Skip to content
rachel/esp32-radio
File

Blob: firmware/vendor/str0m/src/packet/h264.rs

rust928 lines
1#![allow(clippy::all)]
2 
3use super::{CodecExtra, Depacketizer, PacketError, Packetizer};
4 
5/// H264 information describing the depacketized / packetized data
6#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
7pub struct H264CodecExtra {
8    /// Flag which indicates that within [`MediaData`], there is an individual frame
9    /// containing complete and independent visual information. This frame serves
10    /// as a reference point for other frames in the video sequence.
11    ///
12    /// [`MediaData`]: crate::media::MediaData
13    pub is_keyframe: bool,
14}
15 
16/// Packetizes H264 RTP packets.
17///
18/// ## Unversioned API surface
19///
20/// This struct is not currently versioned according to semver rules.
21/// Breaking changes may be made in minor or patch releases.
22#[derive(Default, Debug, Clone)]
23pub struct H264Packetizer {
24    sps_nalu: Option<Vec<u8>>,
25    pps_nalu: Option<Vec<u8>>,
26}
27 
28pub const STAPA_NALU_TYPE: u8 = 24;
29pub const FUA_NALU_TYPE: u8 = 28;
30pub const FUB_NALU_TYPE: u8 = 29;
31pub const IDR_NALU_TYPE: u8 = 5;
32pub const SPS_NALU_TYPE: u8 = 7;
33pub const PPS_NALU_TYPE: u8 = 8;
34pub const AUD_NALU_TYPE: u8 = 9;
35pub const FILLER_NALU_TYPE: u8 = 12;
36 
37pub const FUA_HEADER_SIZE: usize = 2;
38pub const STAPA_HEADER_SIZE: usize = 1;
39pub const STAPA_NALU_LENGTH_SIZE: usize = 2;
40 
41pub const NALU_TYPE_BITMASK: u8 = 0x1F;
42pub const NALU_REF_IDC_BITMASK: u8 = 0x60;
43pub const FU_START_BITMASK: u8 = 0x80;
44pub const FU_END_BITMASK: u8 = 0x40;
45 
46pub const OUTPUT_STAP_AHEADER: u8 = 0x78;
47 
48pub static ANNEXB_NALUSTART_CODE: &[u8] = &[0x00, 0x00, 0x00, 0x01];
49 
50/// Detect whether an H264 RTP payload contains a keyframe.
51///
52/// Checks for IDR (Instantaneous Decoding Refresh) NAL units (type 5)
53/// in the RTP payload. Handles single NAL units, STAP-A aggregation
54/// packets, and FU-A fragmentation units.
55///
56/// For FU-A packets, only the start fragment (S=1) is detected as a
57/// keyframe since the NAL type is only reliably available there.
58pub fn detect_h264_keyframe(payload: &[u8]) -> bool {
59    if payload.is_empty() {
60        return false;
61    }
62 
63    let nalu_type = payload[0] & NALU_TYPE_BITMASK;
64 
65    match nalu_type {
66        // Single NAL unit (types 1-23)
67        1..=23 => nalu_type == IDR_NALU_TYPE,
68 
69        // STAP-A: check all aggregated NALUs
70        STAPA_NALU_TYPE => {
71            let mut offset = STAPA_HEADER_SIZE;
72            while offset + STAPA_NALU_LENGTH_SIZE <= payload.len() {
73                let nalu_size = ((payload[offset] as usize) << 8) | payload[offset + 1] as usize;
74                offset += STAPA_NALU_LENGTH_SIZE;
75                if offset + nalu_size > payload.len() {
76                    break;
77                }
78                if let Some(&b0) = payload.get(offset) {
79                    if b0 & NALU_TYPE_BITMASK == IDR_NALU_TYPE {
80                        return true;
81                    }
82                }
83                offset += nalu_size;
84            }
85            false
86        }
87 
88        // FU-A: check start fragment for IDR type
89        FUA_NALU_TYPE => {
90            if payload.len() < FUA_HEADER_SIZE {
91                return false;
92            }
93            let b1 = payload[1];
94            // Only the start fragment (S=1) reliably carries the original NAL type
95            if b1 & FU_START_BITMASK == 0 {
96                return false;
97            }
98            b1 & NALU_TYPE_BITMASK == IDR_NALU_TYPE
99        }
100 
101        _ => false,
102    }
103}
104 
105impl H264Packetizer {
106    fn next_ind(nalu: &[u8], start: usize) -> (isize, isize) {
107        let mut zero_count = 0;
108 
109        for (i, &b) in nalu[start..].iter().enumerate() {
110            if b == 0 {
111                zero_count += 1;
112                continue;
113            } else if b == 1 && zero_count >= 2 {
114                return ((start + i - zero_count) as isize, zero_count as isize + 1);
115            }
116            zero_count = 0
117        }
118        (-1, -1)
119    }
120 
121    fn emit(&mut self, nalu: &[u8], mtu: usize, payloads: &mut Vec<Vec<u8>>) {
122        if nalu.is_empty() {
123            return;
124        }
125 
126        let nalu_type = nalu[0] & NALU_TYPE_BITMASK;
127        let nalu_ref_idc = nalu[0] & NALU_REF_IDC_BITMASK;
128 
129        if nalu_type == AUD_NALU_TYPE || nalu_type == FILLER_NALU_TYPE {
130            return;
131        } else if nalu_type == SPS_NALU_TYPE {
132            self.sps_nalu = Some(nalu.to_vec());
133            return;
134        } else if nalu_type == PPS_NALU_TYPE {
135            self.pps_nalu = Some(nalu.to_vec());
136            return;
137        } else if let (Some(sps_nalu), Some(pps_nalu)) = (&self.sps_nalu, &self.pps_nalu) {
138            // Pack current NALU with SPS and PPS as STAP-A
139            let sps_len = (sps_nalu.len() as u16).to_be_bytes();
140            let pps_len = (pps_nalu.len() as u16).to_be_bytes();
141 
142            let mut stap_a_nalu = Vec::with_capacity(1 + 2 + sps_nalu.len() + 2 + pps_nalu.len());
143            stap_a_nalu.push(OUTPUT_STAP_AHEADER);
144            stap_a_nalu.extend(sps_len);
145            stap_a_nalu.extend_from_slice(sps_nalu);
146            stap_a_nalu.extend(pps_len);
147            stap_a_nalu.extend_from_slice(pps_nalu);
148            if stap_a_nalu.len() <= mtu {
149                payloads.push(stap_a_nalu);
150            }
151        }
152 
153        if self.sps_nalu.is_some() && self.pps_nalu.is_some() {
154            self.sps_nalu = None;
155            self.pps_nalu = None;
156        }
157 
158        // Single NALU
159        if nalu.len() <= mtu {
160            payloads.push(nalu.to_vec());
161            return;
162        }
163 
164        // FU-A
165        let max_fragment_size = mtu as isize - FUA_HEADER_SIZE as isize;
166 
167        // The FU payload consists of fragments of the payload of the fragmented
168        // NAL unit so that if the fragmentation unit payloads of consecutive
169        // FUs are sequentially concatenated, the payload of the fragmented NAL
170        // unit can be reconstructed.  The NAL unit type octet of the fragmented
171        // NAL unit is not included as such in the fragmentation unit payload,
172        // 	but rather the information of the NAL unit type octet of the
173        // fragmented NAL unit is conveyed in the F and NRI fields of the FU
174        // indicator octet of the fragmentation unit and in the type field of
175        // the FU header.  An FU payload MAY have any number of octets and MAY
176        // be empty.
177 
178        let nalu_data = nalu;
179        // According to the RFC, the first octet is skipped due to redundant information
180        let mut nalu_data_index = 1;
181        let nalu_data_length = nalu.len() as isize - nalu_data_index;
182        let mut nalu_data_remaining = nalu_data_length;
183 
184        if std::cmp::min(max_fragment_size, nalu_data_remaining) <= 0 {
185            return;
186        }
187 
188        while nalu_data_remaining > 0 {
189            let current_fragment_size = std::cmp::min(max_fragment_size, nalu_data_remaining);
190            //out: = make([]byte, fuaHeaderSize + currentFragmentSize)
191            let mut out = Vec::with_capacity(FUA_HEADER_SIZE + current_fragment_size as usize);
192            // +---------------+
193            // |0|1|2|3|4|5|6|7|
194            // +-+-+-+-+-+-+-+-+
195            // |F|NRI|  Type   |
196            // +---------------+
197            let b0 = FUA_NALU_TYPE | nalu_ref_idc;
198            out.push(b0);
199 
200            // +---------------+
201            //|0|1|2|3|4|5|6|7|
202            //+-+-+-+-+-+-+-+-+
203            //|S|E|R|  Type   |
204            //+---------------+
205 
206            let mut b1 = nalu_type;
207            if nalu_data_remaining == nalu_data_length {
208                // Set start bit
209                b1 |= 1 << 7;
210            } else if nalu_data_remaining - current_fragment_size == 0 {
211                // Set end bit
212                b1 |= 1 << 6;
213            }
214            out.push(b1);
215 
216            out.extend_from_slice(
217                &nalu_data
218                    [nalu_data_index as usize..(nalu_data_index + current_fragment_size) as usize],
219            );
220            payloads.push(out);
221 
222            nalu_data_remaining -= current_fragment_size;
223            nalu_data_index += current_fragment_size;
224        }
225    }
226}
227 
228impl Packetizer for H264Packetizer {
229    /// Payload fragments a H264 packet across one or more byte arrays
230    fn packetize(&mut self, mtu: usize, payload: &[u8]) -> Result<Vec<Vec<u8>>, PacketError> {
231        if payload.is_empty() || mtu == 0 {
232            return Ok(vec![]);
233        }
234 
235        let mut payloads = vec![];
236 
237        let (mut next_ind_start, mut next_ind_len) = H264Packetizer::next_ind(payload, 0);
238        if next_ind_start == -1 {
239            self.emit(payload, mtu, &mut payloads);
240        } else {
241            while next_ind_start != -1 {
242                let prev_start = (next_ind_start + next_ind_len) as usize;
243                let (next_ind_start2, next_ind_len2) =
244                    H264Packetizer::next_ind(payload, prev_start);
245                next_ind_start = next_ind_start2;
246                next_ind_len = next_ind_len2;
247                if next_ind_start != -1 {
248                    self.emit(
249                        &payload[prev_start..next_ind_start as usize],
250                        mtu,
251                        &mut payloads,
252                    );
253                } else {
254                    // Emit until end of stream, no end indicator found
255                    self.emit(&payload[prev_start..], mtu, &mut payloads);
256                }
257            }
258        }
259 
260        Ok(payloads)
261    }
262 
263    fn is_marker(&mut self, _data: &[u8], _previous: Option<&[u8]>, last: bool) -> bool {
264        last
265    }
266}
267 
268/// Depacketizes H264 RTP packets.
269///
270/// ## Unversioned API surface
271///
272/// This struct is not currently versioned according to semver rules.
273/// Breaking changes may be made in minor or patch releases.
274#[derive(PartialEq, Eq, Debug, Default, Clone)]
275pub struct H264Depacketizer {
276    /// Whether to output in AVC format (length-prefixed NALUs) instead of Annex B format
277    pub is_avc: bool,
278    fua_buffer: Option<Vec<u8>>,
279}
280 
281impl Depacketizer for H264Depacketizer {
282    fn out_size_hint(&self, packets_size: usize) -> Option<usize> {
283        // Roughly account for Annex B start codes or AVC length prefixes.
284        let estimated_packets = (packets_size / 1200).saturating_add(1);
285        Some(packets_size.saturating_add(4usize.saturating_mul(estimated_packets)))
286    }
287 
288    /// depacketize parses the passed byte slice and stores the result in the
289    /// H264Packet this method is called upon
290    fn depacketize(
291        &mut self,
292        packet: &[u8],
293        out: &mut Vec<u8>,
294        extra: &mut CodecExtra,
295    ) -> Result<(), PacketError> {
296        if packet.len() == 0 {
297            return Err(PacketError::ErrShortPacket);
298        }
299 
300        // NALU Types
301        // https://tools.ietf.org/html/rfc6184#section-5.4
302        let b0 = packet[0];
303        let nalu_type = b0 & NALU_TYPE_BITMASK;
304 
305        match nalu_type {
306            t @ 1..=23 => {
307                let is_keyframe = if let CodecExtra::H264(e) = extra {
308                    (t == IDR_NALU_TYPE) | e.is_keyframe
309                } else {
310                    t == IDR_NALU_TYPE
311                };
312                *extra = CodecExtra::H264(H264CodecExtra { is_keyframe });
313 
314                if self.is_avc {
315                    out.extend_from_slice(&(packet.len() as u32).to_be_bytes());
316                } else {
317                    out.extend_from_slice(ANNEXB_NALUSTART_CODE);
318                }
319                out.extend_from_slice(packet);
320                Ok(())
321            }
322            STAPA_NALU_TYPE => {
323                let mut curr_offset = STAPA_HEADER_SIZE;
324                while curr_offset + 1 < packet.len() {
325                    let nalu_size =
326                        ((packet[curr_offset] as usize) << 8) | packet[curr_offset + 1] as usize;
327                    curr_offset += STAPA_NALU_LENGTH_SIZE;
328 
329                    if curr_offset + nalu_size > packet.len() {
330                        return Err(PacketError::StapASizeLargerThanBuffer(
331                            nalu_size,
332                            packet.len() - curr_offset,
333                        ));
334                    }
335 
336                    let Some(b0) = packet.get(curr_offset) else {
337                        continue;
338                    };
339                    let t = b0 & NALU_TYPE_BITMASK;
340                    let is_keyframe = if let CodecExtra::H264(e) = extra {
341                        (t == IDR_NALU_TYPE) | e.is_keyframe
342                    } else {
343                        t == IDR_NALU_TYPE
344                    };
345                    *extra = CodecExtra::H264(H264CodecExtra { is_keyframe });
346 
347                    if self.is_avc {
348                        out.extend_from_slice(&(nalu_size as u32).to_be_bytes());
349                    } else {
350                        out.extend_from_slice(ANNEXB_NALUSTART_CODE);
351                    }
352                    out.extend_from_slice(&packet[curr_offset..curr_offset + nalu_size]);
353                    curr_offset += nalu_size;
354                }
355 
356                Ok(())
357            }
358            FUA_NALU_TYPE => {
359                if packet.len() < FUA_HEADER_SIZE as usize {
360                    return Err(PacketError::ErrShortPacket);
361                }
362 
363                if self.fua_buffer.is_none() {
364                    self.fua_buffer = Some(Vec::new());
365                }
366 
367                if let Some(fua_buffer) = &mut self.fua_buffer {
368                    fua_buffer.extend_from_slice(&packet[FUA_HEADER_SIZE as usize..]);
369                }
370 
371                let b1 = packet[1];
372                if b1 & FU_END_BITMASK != 0 {
373                    let nalu_ref_idc = b0 & NALU_REF_IDC_BITMASK;
374                    let fragmented_nalu_type = b1 & NALU_TYPE_BITMASK;
375 
376                    let is_keyframe = if let CodecExtra::H264(e) = extra {
377                        (fragmented_nalu_type == IDR_NALU_TYPE) | e.is_keyframe
378                    } else {
379                        fragmented_nalu_type == IDR_NALU_TYPE
380                    };
381                    *extra = CodecExtra::H264(H264CodecExtra { is_keyframe });
382 
383                    if let Some(fua_buffer) = self.fua_buffer.take() {
384                        if self.is_avc {
385                            out.extend_from_slice(&((fua_buffer.len() + 1) as u32).to_be_bytes());
386                        } else {
387                            out.extend_from_slice(ANNEXB_NALUSTART_CODE);
388                        }
389                        out.push(nalu_ref_idc | fragmented_nalu_type);
390                        out.extend_from_slice(&fua_buffer);
391                    }
392 
393                    Ok(())
394                } else {
395                    Ok(())
396                }
397            }
398            _ => Err(PacketError::NaluTypeIsNotHandled(nalu_type)),
399        }
400    }
401 
402    /// is_partition_head checks if this is the head of a packetized nalu stream.
403    fn is_partition_head(&self, packet: &[u8]) -> bool {
404        if packet.len() < 2 {
405            return false;
406        }
407 
408        if packet[0] & NALU_TYPE_BITMASK == FUA_NALU_TYPE
409            || packet[0] & NALU_TYPE_BITMASK == FUB_NALU_TYPE
410        {
411            (packet[1] & FU_START_BITMASK) != 0
412        } else {
413            true
414        }
415    }
416 
417    fn is_partition_tail(&self, marker: bool, _packet: &[u8]) -> bool {
418        marker
419    }
420}
421 
422#[cfg(test)]
423mod test {
424    use super::*;
425 
426    #[test]
427    fn test_h264_payload() -> Result<(), PacketError> {
428        let empty = &[];
429        let small_payload = &[0x90, 0x90, 0x90];
430        let multiple_payload = &[0x00, 0x00, 0x01, 0x90, 0x00, 0x00, 0x01, 0x90];
431        let large_payload = &[
432            0x00, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x10,
433            0x11, 0x12, 0x13, 0x14, 0x15,
434        ];
435        let large_payload_packetized = vec![
436            &[0x1c, 0x80, 0x01, 0x02, 0x03],
437            &[0x1c, 0x00, 0x04, 0x05, 0x06],
438            &[0x1c, 0x00, 0x07, 0x08, 0x09],
439            &[0x1c, 0x00, 0x10, 0x11, 0x12],
440            &[0x1c, 0x40, 0x13, 0x14, 0x15],
441        ];
442 
443        let mut pck = H264Packetizer::default();
444 
445        // Positive MTU, empty payload
446        let result = pck.packetize(1, empty)?;
447        assert!(result.is_empty(), "Generated payload should be empty");
448 
449        // 0 MTU, small payload
450        let result = pck.packetize(0, small_payload)?;
451        assert_eq!(result.len(), 0, "Generated payload should be empty");
452 
453        // Positive MTU, small payload
454        let result = pck.packetize(1, small_payload)?;
455        assert_eq!(result.len(), 0, "Generated payload should be empty");
456 
457        // Positive MTU, small payload
458        let result = pck.packetize(5, small_payload)?;
459        assert_eq!(result.len(), 1, "Generated payload should be the 1");
460        assert_eq!(
461            result[0].len(),
462            small_payload.len(),
463            "Generated payload should be the same size as original payload size"
464        );
465 
466        // Multiple NALU in a single payload
467        let result = pck.packetize(5, multiple_payload)?;
468        assert_eq!(result.len(), 2, "2 nal units should be broken out");
469        for i in 0..2 {
470            assert_eq!(
471                result[i].len(),
472                1,
473                "Payload {} of 2 is packed incorrectly",
474                i + 1,
475            );
476        }
477 
478        // Large Payload split across multiple RTP Packets
479        let result = pck.packetize(5, large_payload)?;
480        assert_eq!(
481            result, large_payload_packetized,
482            "FU-A packetization failed"
483        );
484 
485        // Nalu type 9 or 12
486        let small_payload2 = &[0x09, 0x00, 0x00];
487        let result = pck.packetize(5, small_payload2)?;
488        assert_eq!(result.len(), 0, "Generated payload should be empty");
489 
490        Ok(())
491    }
492 
493    macro_rules! test_h264 {
494        ($name:tt, $is_avc:expr, $is_ok: expr, $payload:expr, $err:tt) => {
495            #[test]
496            fn $name() -> Result<(), PacketError> {
497                let mut pkt = H264Depacketizer::default();
498                pkt.is_avc = $is_avc;
499                let mut extra = CodecExtra::None;
500                let mut out: Vec<u8> = Vec::new();
501                let result = pkt.depacketize($payload, &mut out, &mut extra);
502                if $is_ok {
503                    assert!(result.is_ok(), $err);
504                } else {
505                    assert!(result.is_err(), $err);
506                }
507                Ok(())
508            }
509        };
510    }
511 
512    test_h264!(
513        nil_payload,
514        false,
515        false,
516        &[],
517        "Unmarshal did not fail on nil payload"
518    );
519    test_h264!(
520        unit_delimiter,
521        false,
522        true,
523        &[0x09, 0x30],
524        "Unmarshal should accept minimal h.264 access unit delimiter"
525    );
526    test_h264!(
527        end_of_sequence_nalu,
528        false,
529        true,
530        &[0x0A],
531        "Unmarshal should accept end of sequence NALU"
532    );
533    test_h264!(
534        not_handled,
535        false,
536        false,
537        &[0xFF, 0x00, 0x00],
538        "Unmarshal accepted a packet with a NALU Type we don't handle"
539    );
540    test_h264!(
541        incomplete_single_payload_multi_nalu,
542        false,
543        false,
544        &[
545            0x78, 0x00, 0x0f, 0x67, 0x42, 0xc0, 0x1f, 0x1a, 0x32, 0x35, 0x01, 0x40, 0x7a, 0x40,
546            0x3c, 0x22, 0x11,
547        ],
548        "Unmarshal accepted a STAP-A packet with insufficient data"
549    );
550 
551    #[test]
552    fn single_payload() -> Result<(), PacketError> {
553        let mut pkt = H264Depacketizer::default();
554        let mut extra = CodecExtra::None;
555        let mut out: Vec<u8> = Vec::new();
556        let single_payload = &[0x90, 0x90, 0x90];
557        let _ = pkt.depacketize(single_payload, &mut out, &mut extra);
558        let single_payload_unmarshaled = &[0x00, 0x00, 0x00, 0x01, 0x90, 0x90, 0x90];
559        assert_eq!(
560            out, single_payload_unmarshaled,
561            "Unmarshaling a single payload shouldn't modify the payload"
562        );
563        Ok(())
564    }
565 
566    #[test]
567    fn single_payload_avc() -> Result<(), PacketError> {
568        let mut pkt = H264Depacketizer::default();
569        pkt.is_avc = true;
570        let mut extra = CodecExtra::None;
571        let mut out: Vec<u8> = Vec::new();
572        let single_payload = &[0x90, 0x90, 0x90];
573        let _ = pkt.depacketize(single_payload, &mut out, &mut extra);
574        let single_payload_unmarshaled_avc = &[0x00, 0x00, 0x00, 0x03, 0x90, 0x90, 0x90];
575        assert_eq!(
576            out, single_payload_unmarshaled_avc,
577            "Unmarshaling a single payload into avc stream shouldn't modify the payload"
578        );
579        Ok(())
580    }
581 
582    #[test]
583    fn h264_large_out() -> Result<(), PacketError> {
584        let large_payload_packetized = vec![
585            &[0x1c, 0x80, 0x01, 0x02, 0x03],
586            &[0x1c, 0x00, 0x04, 0x05, 0x06],
587            &[0x1c, 0x00, 0x07, 0x08, 0x09],
588            &[0x1c, 0x00, 0x10, 0x11, 0x12],
589            &[0x1c, 0x40, 0x13, 0x14, 0x15],
590        ];
591 
592        let large_payload = &[
593            0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
594            0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
595        ];
596 
597        let mut pkt = H264Depacketizer::default();
598        let mut extra = CodecExtra::None;
599 
600        let mut large_out = Vec::new();
601        for p in &large_payload_packetized {
602            pkt.depacketize(*p, &mut large_out, &mut extra)?;
603        }
604        assert_eq!(
605            large_out, large_payload,
606            "Failed to unmarshal a large payload"
607        );
608 
609        Ok(())
610    }
611 
612    #[test]
613    fn h264_large_out_avc() -> Result<(), PacketError> {
614        let large_payload_packetized = vec![
615            &[0x1c, 0x80, 0x01, 0x02, 0x03],
616            &[0x1c, 0x00, 0x04, 0x05, 0x06],
617            &[0x1c, 0x00, 0x07, 0x08, 0x09],
618            &[0x1c, 0x00, 0x10, 0x11, 0x12],
619            &[0x1c, 0x40, 0x13, 0x14, 0x15],
620        ];
621 
622        let large_payload_avc = &[
623            0x00, 0x00, 0x00, 0x10, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
624            0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
625        ];
626 
627        let mut avc_pkt = H264Depacketizer {
628            is_avc: true,
629            ..Default::default()
630        };
631 
632        let mut extra = CodecExtra::None;
633 
634        let mut large_out_avc = Vec::new();
635        for p in &large_payload_packetized {
636            avc_pkt.depacketize(*p, &mut large_out_avc, &mut extra)?;
637        }
638        assert_eq!(
639            large_out_avc, large_payload_avc,
640            "Failed to unmarshal a large payload into avc stream"
641        );
642 
643        Ok(())
644    }
645 
646    #[test]
647    fn single_payload_multi_nalu() -> Result<(), PacketError> {
648        let single_payload_multi_nalu = &[
649            0x78, 0x00, 0x0f, 0x67, 0x42, 0xc0, 0x1f, 0x1a, 0x32, 0x35, 0x01, 0x40, 0x7a, 0x40,
650            0x3c, 0x22, 0x11, 0xa8, 0x00, 0x05, 0x68, 0x1a, 0x34, 0xe3, 0xc8, 0x00,
651        ];
652        let single_payload_multi_nalu_unmarshaled = &[
653            0x00, 0x00, 0x00, 0x01, 0x67, 0x42, 0xc0, 0x1f, 0x1a, 0x32, 0x35, 0x01, 0x40, 0x7a,
654            0x40, 0x3c, 0x22, 0x11, 0xa8, 0x00, 0x00, 0x00, 0x01, 0x68, 0x1a, 0x34, 0xe3, 0xc8,
655        ];
656 
657        let mut pkt = H264Depacketizer::default();
658 
659        let mut extra = CodecExtra::None;
660 
661        let mut out = Vec::new();
662        pkt.depacketize(single_payload_multi_nalu, &mut out, &mut extra)?;
663        assert_eq!(
664            out, single_payload_multi_nalu_unmarshaled,
665            "Failed to unmarshal a single packet with multiple NALUs"
666        );
667 
668        Ok(())
669    }
670 
671    #[test]
672    fn single_payload_multi_nalu_avc() -> Result<(), PacketError> {
673        let single_payload_multi_nalu = &[
674            0x78, 0x00, 0x0f, 0x67, 0x42, 0xc0, 0x1f, 0x1a, 0x32, 0x35, 0x01, 0x40, 0x7a, 0x40,
675            0x3c, 0x22, 0x11, 0xa8, 0x00, 0x05, 0x68, 0x1a, 0x34, 0xe3, 0xc8, 0x00,
676        ];
677        let single_payload_multi_nalu_unmarshaled_avc = &[
678            0x00, 0x00, 0x00, 0x0f, 0x67, 0x42, 0xc0, 0x1f, 0x1a, 0x32, 0x35, 0x01, 0x40, 0x7a,
679            0x40, 0x3c, 0x22, 0x11, 0xa8, 0x00, 0x00, 0x00, 0x05, 0x68, 0x1a, 0x34, 0xe3, 0xc8,
680        ];
681 
682        let mut avc_pkt = H264Depacketizer::default();
683        avc_pkt.is_avc = true;
684 
685        let mut extra = CodecExtra::None;
686 
687        let mut out = Vec::new();
688        avc_pkt.depacketize(single_payload_multi_nalu, &mut out, &mut extra)?;
689        assert_eq!(
690            out, single_payload_multi_nalu_unmarshaled_avc,
691            "Failed to unmarshal a single packet with multiple NALUs into avc stream"
692        );
693 
694        Ok(())
695    }
696 
697    #[test]
698    fn test_h264_partition_head_checker_is_partition_head() -> Result<(), PacketError> {
699        let h264 = H264Depacketizer::default();
700        let empty_nalu = &[];
701        assert!(
702            !h264.is_partition_head(empty_nalu),
703            "empty nalu must not be a partition head"
704        );
705 
706        let single_nalu = &[1, 0];
707        assert!(
708            h264.is_partition_head(single_nalu),
709            "single nalu must be a partition head"
710        );
711 
712        let stapa_nalu = &[STAPA_NALU_TYPE, 0];
713        assert!(
714            h264.is_partition_head(stapa_nalu),
715            "stapa nalu must be a partition head"
716        );
717 
718        let fua_start_nalu = &[FUA_NALU_TYPE, FU_START_BITMASK];
719        assert!(
720            h264.is_partition_head(fua_start_nalu),
721            "fua start nalu must be a partition head"
722        );
723 
724        let fua_end_nalu = &[FUA_NALU_TYPE, FU_END_BITMASK];
725        assert!(
726            !h264.is_partition_head(fua_end_nalu),
727            "fua end nalu must not be a partition head"
728        );
729 
730        let fub_start_nalu = &[FUB_NALU_TYPE, FU_START_BITMASK];
731        assert!(
732            h264.is_partition_head(fub_start_nalu),
733            "fub start nalu must be a partition head"
734        );
735 
736        let fub_end_nalu = &[FUB_NALU_TYPE, FU_END_BITMASK];
737        assert!(
738            !h264.is_partition_head(fub_end_nalu),
739            "fub end nalu must not be a partition head"
740        );
741 
742        Ok(())
743    }
744 
745    #[test]
746    fn test_h264_packetizer_payload_sps_and_pps_handling() -> Result<(), PacketError> {
747        let mut pck = H264Packetizer::default();
748        let expected: Vec<&[u8]> = vec![
749            &[
750                0x78, 0x00, 0x03, 0x07, 0x00, 0x01, 0x00, 0x03, 0x08, 0x02, 0x03,
751            ],
752            &[0x05, 0x04, 0x05],
753        ];
754 
755        // When packetizing SPS and PPS are emitted with following NALU
756        let res = pck.packetize(1500, &[0x07, 0x00, 0x01])?;
757        assert!(res.is_empty(), "Generated payload should be empty");
758 
759        let res = pck.packetize(1500, &[0x08, 0x02, 0x03])?;
760        assert!(res.is_empty(), "Generated payload should be empty");
761 
762        let actual = pck.packetize(1500, &[0x05, 0x04, 0x05])?;
763        assert_eq!(actual, expected, "SPS and PPS aren't packed together");
764 
765        Ok(())
766    }
767 
768    #[test]
769    fn test_h264_depacketizer_idr_handling() -> Result<(), PacketError> {
770        let mut pck = H264Depacketizer::default();
771        let mut extra = CodecExtra::None;
772        let mut out = vec![];
773 
774        // First byte is NALU type
775        let packet = [0x85];
776        pck.depacketize(&packet, &mut out, &mut extra)?;
777        let CodecExtra::H264(e) = extra else {
778            panic!("Expected CodecExtra::H264");
779        };
780        assert!(e.is_keyframe);
781 
782        // First byte is STAPA NALU type
783        let packet = vec![
784            vec![
785                120, 0, 15, 103, 66, 192, 21, 140, 141, 64, 160, 203, 207, 0, 240, 136, 70, 160, 0,
786                4, 104, 206, 60, 128, 1, 20, 101,
787            ],
788            vec![0; 276],
789        ]
790        .into_iter()
791        .flatten()
792        .collect::<Vec<_>>();
793        pck.depacketize(packet.as_slice(), &mut out, &mut extra)?;
794        let CodecExtra::H264(e) = extra else {
795            panic!("Expected CodecExtra::H264");
796        };
797        assert!(e.is_keyframe);
798 
799        // First byte is FUA NALU type
800        let packet = [124, 69];
801        pck.depacketize(&packet, &mut out, &mut extra)?;
802        let CodecExtra::H264(e) = extra else {
803            panic!("Expected CodecExtra::H264");
804        };
805        assert!(e.is_keyframe);
806        Ok(())
807    }
808 
809    #[test]
810    fn parse_first_packet() {
811        const PACKET: &[u8] = &[
812            120, 000, 015, 103, 066, 192, 021, 140, 141, 064, 160, 203, 207, 000, 240, 136, 070,
813            160, 000, 004, 104, 206, 060, 128, 000, 204, 101, 184, 000, 004, 000, 000, 005, 057,
814            049, 064, 000, 064, 222, 078, 078, 078, 078, 078, 078, 078, 078, 078, 078, 078, 078,
815            078, 078, 078, 078, 078, 078, 078, 186, 235, 174, 186, 235, 174, 186, 235, 174, 186,
816            235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174,
817            186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 173, 223, 039, 125, 247, 223,
818            125, 245, 215, 093, 117, 215, 093, 117, 214, 239, 174, 187, 235, 174, 186, 235, 174,
819            186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235,
820            174, 186, 235, 174, 183, 093, 117, 215, 093, 117, 215, 093, 117, 215, 093, 117, 215,
821            093, 117, 215, 093, 117, 215, 092, 189, 117, 215, 093, 117, 215, 093, 117, 215, 093,
822            117, 215, 093, 117, 215, 093, 117, 215, 093, 117, 214, 239, 190, 251, 239, 190, 186,
823            235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174, 186, 235, 174,
824            186, 235, 174, 186, 235, 175, 227, 255, 240, 247, 021, 223, 125, 247, 223, 125, 247,
825            223, 125, 247, 223, 125, 247, 223, 125, 248,
826        ];
827 
828        let mut pck = H264Depacketizer::default();
829        let mut extra = CodecExtra::None;
830        let mut out = vec![];
831        pck.depacketize(PACKET, &mut out, &mut extra).unwrap();
832    }
833 
834    #[test]
835    fn test_out_of_bounds_access() {
836        const PACKET: &[u8] = &[STAPA_NALU_TYPE, 0x00, 0x00];
837 
838        let mut pck = H264Depacketizer::default();
839        let mut extra = CodecExtra::None;
840        let mut out = vec![];
841        pck.depacketize(PACKET, &mut out, &mut extra).unwrap();
842    }
843 
844    #[test]
845    fn test_detect_h264_keyframe() {
846        // Empty payload
847        assert!(!detect_h264_keyframe(&[]));
848 
849        // Single IDR NAL unit (type 5)
850        assert!(detect_h264_keyframe(&[0x65, 0x00, 0x00])); // 0x65 & 0x1F = 5
851 
852        // Single non-IDR NAL unit (type 1 = coded slice)
853        assert!(!detect_h264_keyframe(&[0x41, 0x00, 0x00])); // 0x41 & 0x1F = 1
854 
855        // SPS (type 7) - not a keyframe
856        assert!(!detect_h264_keyframe(&[0x67, 0x00, 0x00]));
857 
858        // PPS (type 8) - not a keyframe
859        assert!(!detect_h264_keyframe(&[0x68, 0x00, 0x00]));
860 
861        // STAP-A containing IDR
862        // Header: type 24 (STAP-A)
863        // NALU 1: 2 bytes, type 7 (SPS)
864        // NALU 2: 2 bytes, type 5 (IDR)
865        let stapa_with_idr = [
866            0x18, // STAP-A (24)
867            0x00, 0x02, 0x67, 0xAA, // SPS: size=2, type=7
868            0x00, 0x02, 0x65, 0xBB, // IDR: size=2, type=5
869        ];
870        assert!(detect_h264_keyframe(&stapa_with_idr));
871 
872        // STAP-A without IDR
873        let stapa_no_idr = [
874            0x18, // STAP-A (24)
875            0x00, 0x02, 0x67, 0xAA, // SPS: size=2, type=7
876            0x00, 0x02, 0x68, 0xBB, // PPS: size=2, type=8
877        ];
878        assert!(!detect_h264_keyframe(&stapa_no_idr));
879 
880        // FU-A start fragment with IDR type
881        let fua_start_idr = [
882            0x7C, // FU indicator: type=28 (FU-A)
883            0x85, // FU header: S=1, type=5 (IDR)
884            0x00, 0x00,
885        ];
886        assert!(detect_h264_keyframe(&fua_start_idr));
887 
888        // FU-A start fragment with non-IDR type
889        let fua_start_non_idr = [
890            0x7C, // FU indicator: type=28 (FU-A)
891            0x81, // FU header: S=1, type=1 (non-IDR)
892            0x00, 0x00,
893        ];
894        assert!(!detect_h264_keyframe(&fua_start_non_idr));
895 
896        // FU-A continuation fragment (S=0) - cannot detect
897        let fua_continuation = [
898            0x7C, // FU indicator: type=28 (FU-A)
899            0x05, // FU header: S=0, type=5 (IDR) but S=0
900            0x00, 0x00,
901        ];
902        assert!(!detect_h264_keyframe(&fua_continuation));
903 
904        // FU-A too short
905        assert!(!detect_h264_keyframe(&[0x7C]));
906    }
907 
908    #[test]
909    fn packetize_respects_mtu() -> Result<(), PacketError> {
910        // NAL header byte (type=5, IDR slice) + payload.
911        let mut payload = vec![0x05u8];
912        payload.extend(std::iter::repeat(0xABu8).take(2000));
913        for &mtu in &[100usize, 300, 600, 1200] {
914            let mut pck = H264Packetizer::default();
915            let pkts = pck.packetize(mtu, &payload)?;
916            assert!(!pkts.is_empty(), "H264 produced no packets at mtu {mtu}");
917            for (i, pkt) in pkts.iter().enumerate() {
918                assert!(
919                    pkt.len() <= mtu,
920                    "H264 packet {i} size {} > mtu {mtu}",
921                    pkt.len()
922                );
923            }
924        }
925        Ok(())
926    }
927}