str0m/packet/
h264.rs

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