T4.1: wzp-video crate scaffold + H.264 NAL framer + depacketizer

crates/wzp-video/src/framer.rs

@@ -0,0 +1,218 @@
+//! H.264 NAL framer — splits access units into MTU-sized packets.
+//!
+//! Supports Single-NAL and FU-A (Fragmentation Unit type A) per RFC 6184.
+
+/// One framed packet emitted by [`H264Framer`].
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct FramedPacket {
+    pub payload: Vec<u8>,
+    /// True when this is the last packet of the access unit.
+    pub is_frame_end: bool,
+}
+
+/// H.264 access-unit framer.
+///
+/// Parses NAL units from a raw access unit and emits either Single-NAL
+/// packets or FU-A fragments so that every payload fits in `max_payload_size`.
+pub struct H264Framer {
+    max_payload_size: usize,
+}
+
+impl H264Framer {
+    /// Create a framer with the given maximum payload size per packet.
+    ///
+    /// Typical value: `MTU - MediaHeader::WIRE_SIZE - AEAD_TAG_SIZE`.
+    pub fn new(max_payload_size: usize) -> Self {
+        Self { max_payload_size }
+    }
+
+    /// Frame one access unit into a sequence of packets.
+    ///
+    /// The input may contain one or more NAL units separated by H.264 start
+    /// codes (`0x000001` or `0x00000001`).  The last emitted packet has
+    /// `is_frame_end = true`.
+    pub fn frame(&self, access_unit: &[u8]) -> Vec<FramedPacket> {
+        let nals = split_nals(access_unit);
+        if nals.is_empty() {
+            return Vec::new();
+        }
+
+        let mut packets = Vec::new();
+        let nal_count = nals.len();
+
+        for (idx, nal) in nals.iter().enumerate() {
+            let is_last_nal = idx + 1 == nal_count;
+
+            if nal.len() <= self.max_payload_size {
+                // Single-NAL packet.
+                packets.push(FramedPacket {
+                    payload: nal.to_vec(),
+                    is_frame_end: is_last_nal,
+                });
+            } else {
+                // FU-A fragmentation.
+                let original_header = nal[0];
+                let nal_type = original_header & 0x1F;
+                let nri = original_header & 0x60;
+
+                // FU indicator: same as original header but with type = 28.
+                let fu_indicator = nri | 28;
+
+                let payload = &nal[1..];
+                let mut offset = 0;
+                let mut frag_idx = 0;
+                let total_frags = payload.len().div_ceil(self.max_payload_size - 2);
+
+                while offset < payload.len() {
+                    let remaining = payload.len() - offset;
+                    let frag_data_len = remaining.min(self.max_payload_size.saturating_sub(2));
+                    let is_first = frag_idx == 0;
+                    let is_last = frag_idx + 1 == total_frags;
+
+                    let fu_header = (if is_first { 0x80 } else { 0 })
+                        | (if is_last { 0x40 } else { 0 })
+                        | nal_type;
+
+                    let mut pkt = Vec::with_capacity(2 + frag_data_len);
+                    pkt.push(fu_indicator);
+                    pkt.push(fu_header);
+                    pkt.extend_from_slice(&payload[offset..offset + frag_data_len]);
+
+                    packets.push(FramedPacket {
+                        payload: pkt,
+                        is_frame_end: is_last_nal && is_last,
+                    });
+
+                    offset += frag_data_len;
+                    frag_idx += 1;
+                }
+            }
+        }
+
+        packets
+    }
+}
+
+/// Split a byte slice into individual NAL units.
+///
+/// NAL units are separated by start codes (`0x000001` or `0x00000001`).
+/// Each returned slice starts with the NAL header byte and contains no
+/// start-code prefix.
+fn split_nals(data: &[u8]) -> Vec<&[u8]> {
+    let mut nals = Vec::new();
+    let mut i = 0;
+
+    while i < data.len() {
+        // Skip leading zeros.
+        while i < data.len() && data[i] == 0 {
+            i += 1;
+        }
+        // Need at least one more byte for the 0x01 marker.
+        if i >= data.len() || data[i] != 1 {
+            break;
+        }
+        i += 1; // skip the 0x01
+
+        let start = i;
+        // Find the next start code or end of data.
+        while i + 3 < data.len() {
+            if data[i] == 0
+                && data[i + 1] == 0
+                && (data[i + 2] == 1
+                    || (data[i + 2] == 0 && i + 4 < data.len() && data[i + 3] == 1))
+            {
+                break;
+            }
+            i += 1;
+        }
+        // If no more start codes were found, consume to the end.
+        if i + 3 >= data.len() {
+            i = data.len();
+        }
+        let end = i;
+        if start < end {
+            nals.push(&data[start..end]);
+        }
+    }
+
+    nals
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Build a synthetic access unit with two NAL units.
+    fn make_access_unit() -> Vec<u8> {
+        let mut au = Vec::new();
+        // Start code + NAL 1 (IDR slice, type 5)
+        au.extend_from_slice(&[0x00, 0x00, 0x00, 0x01, 0x65, 0x01, 0x02, 0x03]);
+        // Start code + NAL 2 (non-IDR slice, type 1)
+        au.extend_from_slice(&[0x00, 0x00, 0x01, 0x41, 0x04, 0x05]);
+        au
+    }
+
+    #[test]
+    fn frame_single_nal_roundtrip() {
+        let framer = H264Framer::new(100);
+        let au = make_access_unit();
+        let packets = framer.frame(&au);
+        assert_eq!(packets.len(), 2);
+        assert_eq!(packets[0].payload, vec![0x65, 0x01, 0x02, 0x03]);
+        assert!(!packets[0].is_frame_end);
+        assert_eq!(packets[1].payload, vec![0x41, 0x04, 0x05]);
+        assert!(packets[1].is_frame_end);
+    }
+
+    #[test]
+    fn frame_empty_input() {
+        let framer = H264Framer::new(100);
+        let packets = framer.frame(&[]);
+        assert!(packets.is_empty());
+    }
+
+    #[test]
+    fn frame_fu_a_fragmentation() {
+        let framer = H264Framer::new(10);
+        // One NAL unit: header 0x65 (IDR) + 20 bytes payload.
+        let mut au = vec![0x00, 0x00, 0x01];
+        au.push(0x65);
+        au.extend_from_slice(&[0xAA; 20]);
+
+        let packets = framer.frame(&au);
+        // max_payload_size = 10, so each fragment can carry 8 bytes of data
+        // (2 bytes FU-A header + 8 data = 10).
+        // 20 bytes payload → 3 fragments (8 + 8 + 4).
+        assert_eq!(packets.len(), 3);
+
+        // First fragment.
+        assert_eq!(packets[0].payload[0], 0x65 & 0x60 | 28); // FU indicator
+        assert_eq!(packets[0].payload[1], 0x80 | 0x05); // S=1, E=0, type=5
+        assert_eq!(packets[0].payload.len(), 10);
+        assert!(!packets[0].is_frame_end);
+
+        // Middle fragment.
+        assert_eq!(packets[1].payload[1], 0x05); // S=0, E=0, type=5
+        assert_eq!(packets[1].payload.len(), 10);
+        assert!(!packets[1].is_frame_end);
+
+        // Last fragment.
+        assert_eq!(packets[2].payload[1], 0x40 | 0x05); // S=0, E=1, type=5
+        assert_eq!(packets[2].payload.len(), 6); // 2 header + 4 data
+        assert!(packets[2].is_frame_end);
+    }
+
+    #[test]
+    fn frame_fu_a_exact_fit() {
+        let framer = H264Framer::new(12);
+        // NAL: 1 header + 10 payload = 11 bytes total → fits in 12, no FU-A.
+        let mut au = vec![0x00, 0x00, 0x01];
+        au.push(0x41);
+        au.extend_from_slice(&[0xBB; 10]);
+
+        let packets = framer.frame(&au);
+        assert_eq!(packets.len(), 1);
+        assert_eq!(packets[0].payload.len(), 11);
+        assert!(packets[0].is_frame_end);
+    }
+}