punktfunk/crates/punktfunk-host/src/gamestream/video.rs

//! GameStream video wire packetization: an encoded access unit → UDP datagrams a stock
//! Moonlight client decodes (and recovers under loss). Each datagram is
//!   `RTP_PACKET(12, big-endian) + reserved[4] + NV_VIDEO_PACKET(16, little-endian) + payload`
//! and the frame's bitstream is prefixed with an 8-byte `video_short_frame_header_t`, then
//! striped into ≤4 FEC blocks of ≤255 shards. Byte-exact spec:
//! `design/research/gamestream-protocol-research.json` (video plane).
//!
//! FEC (P1.5): each block carries `m = ⌈k·pct/100⌉` Reed–Solomon parity shards generated by
//! `punktfunk_core::fec::Gf8Coder` (the nanors-compatible Cauchy GF(2⁸) coder). Crucially, RS runs
//! over the **whole `blocksize` shard** — Moonlight decodes over `packetSize + 16` bytes from
//! the datagram start (`RtpVideoQueue.c`), and rejects a recovered shard whose reconstructed
//! `flags` byte isn't valid — so the NV header fields RS must reproduce (streamPacketIndex,
//! frameIndex, flags, multiFec*) are written into the data shards **before** encoding, and only
//! the transport fields (RTP header/seq/timestamp + fecInfo) are stamped **after**, matching
//! Sunshine `stream.cpp`. `pct = 0` falls back to data-shards-only. Plaintext (AES-GCM video
//! encryption is negotiated off for now).

use punktfunk_core::fec::{ErasureCoder, Gf8Coder};

/// RTP `header` byte: version 2 (0x80) | extension (0x10) — Moonlight keys on the extension.
const RTP_HEADER_BYTE: u8 = 0x80 | 0x10;
const FLAG_PIC: u8 = 0x1;
const FLAG_EOF: u8 = 0x2;
const FLAG_SOF: u8 = 0x4;
const MULTI_FEC_FLAGS: u8 = 0x10;
const MAX_DATA_SHARDS_PER_BLOCK: usize = 255;
const MAX_FEC_BLOCKS: usize = 4;
/// Per-shard header: RTP(12) + reserved(4) + NV_VIDEO_PACKET(16).
const SHARD_HEADER: usize = 32;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FrameType {
    Idr,
    P,
}

/// Splits encoded access units into GameStream video datagrams (data + FEC parity shards).
pub struct VideoPacketizer {
    /// Negotiated `packetSize` (ANNOUNCE `x-nv-video[0].packetSize`).
    packet_size: usize,
    /// Per-shard payload bytes = `blocksize - SHARD_HEADER`, `blocksize = packetSize + 16`.
    payload_per_shard: usize,
    /// Requested FEC overhead percent (0 = data shards only). The wire carries the recomputed
    /// per-block `(100·m)/k` so Moonlight derives the same parity count.
    fec_percentage: usize,
    /// Minimum parity shards per block (the client's `fec.minRequiredFecPackets`) — protects
    /// small frames whose `⌈k·pct/100⌉` would otherwise be just 1.
    min_fec: usize,
    frame_index: u32,
    /// Monotonic per-stream packet counter (the RTP sequence / streamPacketIndex source).
    seq: u32,
}

impl VideoPacketizer {
    pub fn new(packet_size: usize, fec_percentage: u8, min_fec: u8) -> Self {
        VideoPacketizer {
            packet_size,
            // Defense in depth: `pps` is a divisor in `packetize` (`% pps`, `div_ceil(pps)`), so it
            // must never be 0. `blocksize = packet_size + 16`; a tiny attacker-supplied packet_size
            // (≤ SHARD_HEADER-16 = 16) would otherwise underflow (panic) or yield pps==0 (div-by-zero).
            // `stream_config` already rejects out-of-range packetSize; this saturating `.max(1)` makes
            // a degenerate value structurally unable to panic, without affecting any valid size.
            payload_per_shard: (packet_size + 16).saturating_sub(SHARD_HEADER).max(1),
            fec_percentage: fec_percentage as usize,
            min_fec: min_fec as usize,
            frame_index: 0,
            seq: 0,
        }
    }

    /// Packetize one encoded AU into wire datagrams (data shards + Cauchy RS parity shards).
    pub fn packetize(
        &mut self,
        au: &[u8],
        frame_type: FrameType,
        timestamp_90k: u32,
    ) -> Vec<Vec<u8>> {
        let frame_index = self.frame_index;
        self.frame_index = self.frame_index.wrapping_add(1);
        let pps = self.payload_per_shard;
        let blocksize = SHARD_HEADER + pps; // = packet_size + 16
        let pct = self.fec_percentage;

        // frame payload = 8-byte short frame header + the AU bitstream.
        let total_len = 8 + au.len();
        let last_payload_len = match total_len % pps {
            0 => pps,
            r => r,
        };
        let mut fp = Vec::with_capacity(total_len);
        fp.extend_from_slice(&short_frame_header(frame_type, last_payload_len as u16));
        fp.extend_from_slice(au);

        let total_data = total_len.div_ceil(pps).max(1);
        // With parity, cap per-block data so k + m ≤ 255 (the GF(2⁸) ceiling): parity for k
        // data shards is ⌈k·pct/100⌉, so k ≤ 255·100/(100+pct).
        let max_data = if pct > 0 {
            (255 * 100) / (100 + pct)
        } else {
            MAX_DATA_SHARDS_PER_BLOCK
        };
        let n_blocks = total_data.div_ceil(max_data).clamp(1, MAX_FEC_BLOCKS);
        let per_block = total_data.div_ceil(n_blocks);

        let mut packets = Vec::with_capacity(total_data + total_data * pct / 100 + n_blocks);
        for b in 0..n_blocks {
            let first = b * per_block;
            let last = ((b + 1) * per_block).min(total_data);
            if first >= last {
                break;
            }
            let k = last - first;
            let block_seq_base = self.seq;
            let multi_fec_blocks = ((b as u8) << 4) | (((n_blocks - 1) as u8) << 6);

            // 1. Build this block's k data-shard datagrams (full `blocksize`), writing the NV
            //    header fields RS must reproduce on recovery (streamPacketIndex, frameIndex,
            //    flags, multiFec*). The RTP header + fecInfo are left zero (stamped post-RS).
            let mut shards: Vec<Vec<u8>> = Vec::with_capacity(k);
            for i in 0..k {
                let global = first + i;
                let seq = block_seq_base + i as u32;
                let mut buf = vec![0u8; blocksize];
                let mut flags = FLAG_PIC;
                if global == 0 {
                    flags |= FLAG_SOF;
                }
                if global == total_data - 1 {
                    flags |= FLAG_EOF;
                }
                buf[16..20].copy_from_slice(&(seq << 8).to_le_bytes()); // streamPacketIndex
                buf[20..24].copy_from_slice(&frame_index.to_le_bytes()); // frameIndex
                buf[24] = flags;
                buf[26] = MULTI_FEC_FLAGS;
                buf[27] = multi_fec_blocks;
                let ps = global * pps;
                let pe = (ps + pps).min(fp.len());
                buf[SHARD_HEADER..SHARD_HEADER + (pe - ps)].copy_from_slice(&fp[ps..pe]);
                shards.push(buf);
            }

            // 2. m = ⌈k·pct/100⌉ parity shards (floored at the client's min, capped so k+m≤255)
            //    over the full datagrams. The wire percentage is recomputed from m so the client
            //    derives the same count.
            let m = if pct > 0 {
                (k * pct).div_ceil(100).max(self.min_fec).min(255 - k)
            } else {
                0
            };
            let wire_pct = if m > 0 { (100 * m) / k } else { 0 };
            let parity = if m > 0 {
                Gf8Coder.encode(&shards, m).unwrap_or_default()
            } else {
                Vec::new()
            };

            // 3. Stamp transport headers (RTP + fecInfo) on every shard. We do NOT touch the
            //    flags/streamPacketIndex bytes, so a recovered data shard's RS-reconstructed
            //    NV header stays valid.
            self.seq = block_seq_base + k as u32;
            for (i, mut buf) in shards.into_iter().enumerate() {
                let seq = block_seq_base + i as u32;
                finalize(
                    &mut buf,
                    seq,
                    timestamp_90k,
                    frame_index,
                    multi_fec_blocks,
                    fec_info(k, i, wire_pct),
                );
                packets.push(buf);
            }
            for (j, mut buf) in parity.into_iter().enumerate() {
                let seq = self.seq;
                self.seq = self.seq.wrapping_add(1);
                finalize(
                    &mut buf,
                    seq,
                    timestamp_90k,
                    frame_index,
                    multi_fec_blocks,
                    fec_info(k, k + j, wire_pct),
                );
                packets.push(buf);
            }
        }
        packets
    }
}

/// `fecInfo` (u32, little-endian): `dataShards<<22 | fecIndex<<12 | fecPercentage<<4`.
fn fec_info(k: usize, fec_index: usize, pct: usize) -> u32 {
    ((k as u32) << 22) | ((fec_index as u32) << 12) | ((pct as u32) << 4)
}

/// Stamp the post-RS transport fields into a shard datagram (in place). Leaves the NV
/// `flags`/`streamPacketIndex`/`multiFecFlags` bytes untouched (RS-covered).
fn finalize(
    buf: &mut [u8],
    seq: u32,
    ts_90k: u32,
    frame_index: u32,
    multi_fec_blocks: u8,
    fec_info: u32,
) {
    buf[0] = RTP_HEADER_BYTE; // header (version 2 + extension)
    buf[2..4].copy_from_slice(&(seq as u16).to_be_bytes()); // sequenceNumber (BE)
    buf[4..8].copy_from_slice(&ts_90k.to_be_bytes()); // timestamp (90 kHz, BE)
    buf[20..24].copy_from_slice(&frame_index.to_le_bytes()); // frameIndex (re-affirm for parity)
    buf[27] = multi_fec_blocks; // re-affirm for parity
    buf[28..32].copy_from_slice(&fec_info.to_le_bytes()); // fecInfo (LE)
}

/// 8-byte `video_short_frame_header_t` (little-endian), prefixed to the AU bitstream.
fn short_frame_header(frame_type: FrameType, last_payload_len: u16) -> [u8; 8] {
    let mut h = [0u8; 8];
    h[0] = 0x01; // headerType
    h[1..3].copy_from_slice(&0u16.to_le_bytes()); // frame_processing_latency
    h[3] = match frame_type {
        FrameType::Idr => 2,
        FrameType::P => 1,
    };
    h[4..6].copy_from_slice(&last_payload_len.to_le_bytes());
    // h[6..8] unknown = 0
    h
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn single_block_layout() {
        let mut pk = VideoPacketizer::new(1392, 0, 0); // data-only; pps = 1392+16-32 = 1376
        assert_eq!(pk.payload_per_shard, 1376);
        let au = vec![0xABu8; 4000]; // 8+4000 = 4008 → ceil(4008/1376) = 3 data shards
        let pkts = pk.packetize(&au, FrameType::Idr, 90_000);
        assert_eq!(pkts.len(), 3);
        for p in &pkts {
            assert_eq!(p.len(), SHARD_HEADER + 1376);
            assert_eq!(p[0], 0x90); // RTP header byte
        }
        let first = &pkts[0];
        assert_eq!(first[24] & FLAG_SOF, FLAG_SOF);
        assert_eq!(first[24] & FLAG_PIC, FLAG_PIC);
        let frame_index = u32::from_le_bytes(first[20..24].try_into().unwrap());
        assert_eq!(frame_index, 0);
        let fec_info = u32::from_le_bytes(first[28..32].try_into().unwrap());
        assert_eq!(fec_info >> 22, 3); // dataShards = 3
        assert_eq!((fec_info >> 12) & 0x3ff, 0); // fecIndex 0
        let last = &pkts[2];
        assert_eq!(last[24] & FLAG_EOF, FLAG_EOF);
        let fec_info_last = u32::from_le_bytes(last[28..32].try_into().unwrap());
        assert_eq!((fec_info_last >> 12) & 0x3ff, 2);
        for (i, p) in pkts.iter().enumerate() {
            assert_eq!(u16::from_be_bytes(p[2..4].try_into().unwrap()), i as u16);
        }
    }

    #[test]
    fn degenerate_packet_size_does_not_panic() {
        // A pre-auth attacker drives packetSize via the RTSP ANNOUNCE. `stream_config` rejects
        // out-of-range values, but the packetizer must ALSO never panic (div-by-zero on `% pps` /
        // `div_ceil(pps)`, or usize underflow) for ANY input — pps is clamped to >= 1.
        for ps in [0usize, 15, 16, 17, 32] {
            let mut pk = VideoPacketizer::new(ps, 20, 2);
            assert!(pk.payload_per_shard >= 1, "pps must never be 0 (ps={ps})");
            let _ = pk.packetize(&[0xCDu8; 200], FrameType::Idr, 0); // must not panic
        }
    }

    #[test]
    fn multi_block_split() {
        let mut pk = VideoPacketizer::new(1392, 0, 0); // data-only
        let au = vec![0u8; 600_000];
        let pkts = pk.packetize(&au, FrameType::P, 0);
        let total = (8 + au.len()).div_ceil(1376);
        assert_eq!(pkts.len(), total);
        let n_blocks = total.div_ceil(255).clamp(1, 4);
        let last_block = ((pkts.last().unwrap()[27]) >> 6) & 0x3;
        assert_eq!(last_block as usize, n_blocks - 1);
    }

    #[test]
    fn emits_parity_shards() {
        let mut pk = VideoPacketizer::new(1392, 20, 0); // pps = 1376, 20% FEC
        let au = vec![0xABu8; 4000]; // 8+4000 = 4008 → 3 data shards (k=3)
        let pkts = pk.packetize(&au, FrameType::Idr, 0);
        // m = ceil(3*20/100) = 1 parity shard → 4 packets; wire_pct = 100*1/3 = 33.
        assert_eq!(pkts.len(), 4);
        for p in &pkts {
            let fec_info = u32::from_le_bytes(p[28..32].try_into().unwrap());
            assert_eq!(fec_info >> 22, 3); // dataShards = k = 3
            assert_eq!((fec_info >> 4) & 0xff, 33); // wire fecPercentage
        }
        // The parity shard is last: fecIndex = k = 3.
        let parity = &pkts[3];
        let fec_info = u32::from_le_bytes(parity[28..32].try_into().unwrap());
        assert_eq!((fec_info >> 12) & 0x3ff, 3);
        // Data shards keep SOF (first) / EOF (last data shard) / PIC.
        assert_eq!(pkts[0][24] & FLAG_SOF, FLAG_SOF);
        assert_eq!(pkts[2][24] & FLAG_EOF, FLAG_EOF);
        // RTP sequence numbers are contiguous across data + parity (0,1,2,3).
        for (i, p) in pkts.iter().enumerate() {
            assert_eq!(u16::from_be_bytes(p[2..4].try_into().unwrap()), i as u16);
        }
    }

    /// End-to-end recovery: parity over the full datagram reconstructs a dropped data shard's
    /// payload AND its NV `flags` byte (the byte Moonlight validates), proving the layout.
    #[test]
    fn parity_recovers_full_datagram_incl_flags() {
        let mut pk = VideoPacketizer::new(1392, 50, 0); // high pct → plenty of parity
        let au = vec![0x5Au8; 4000]; // k = 3
        let pkts = pk.packetize(&au, FrameType::Idr, 0);
        let k = 3usize;
        let m = pkts.len() - k;
        assert!(m >= 1);
        // Drop data shard 1; reconstruct from the rest via the same Cauchy coder.
        let mut received: Vec<Option<Vec<u8>>> = pkts.iter().map(|p| Some(p.clone())).collect();
        received[1] = None;
        let recovered = Gf8Coder.reconstruct(k, m, &mut received).unwrap();
        // The recovered shard equals the original data shard's RS-covered bytes: its flags
        // byte (offset 24) is PIC (middle shard), proving the NV header recovers correctly.
        assert_eq!(recovered[1][24], FLAG_PIC);
        // ...and the payload region matches the original.
        assert_eq!(recovered[1][SHARD_HEADER..], pkts[1][SHARD_HEADER..]);
    }
}