webrtc_media/io/h264_reader/
mod.rs

1#[cfg(test)]
2mod h264_reader_test;
3
4use std::fmt;
5use std::io::Read;
6
7use bytes::{BufMut, BytesMut};
8
9use crate::error::{Error, Result};
10
11/// NalUnitType is the type of a NAL
12/// Enums for NalUnitTypes
13#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
14pub enum NalUnitType {
15    /// Unspecified
16    #[default]
17    Unspecified = 0,
18    /// Coded slice of a non-IDR picture
19    CodedSliceNonIdr = 1,
20    /// Coded slice data partition A
21    CodedSliceDataPartitionA = 2,
22    /// Coded slice data partition B
23    CodedSliceDataPartitionB = 3,
24    /// Coded slice data partition C
25    CodedSliceDataPartitionC = 4,
26    /// Coded slice of an IDR picture
27    CodedSliceIdr = 5,
28    /// Supplemental enhancement information (SEI)
29    SEI = 6,
30    /// Sequence parameter set
31    SPS = 7,
32    /// Picture parameter set
33    PPS = 8,
34    /// Access unit delimiter
35    AUD = 9,
36    /// End of sequence
37    EndOfSequence = 10,
38    /// End of stream
39    EndOfStream = 11,
40    /// Filler data
41    Filler = 12,
42    /// Sequence parameter set extension
43    SpsExt = 13,
44    /// Coded slice of an auxiliary coded picture without partitioning
45    CodedSliceAux = 19,
46    ///Reserved
47    Reserved,
48    // 14..18                                            // Reserved
49    // 20..23                                            // Reserved
50    // 24..31                                            // Unspecified
51}
52
53impl fmt::Display for NalUnitType {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        let s = match *self {
56            NalUnitType::Unspecified => "Unspecified",
57            NalUnitType::CodedSliceNonIdr => "CodedSliceNonIdr",
58            NalUnitType::CodedSliceDataPartitionA => "CodedSliceDataPartitionA",
59            NalUnitType::CodedSliceDataPartitionB => "CodedSliceDataPartitionB",
60            NalUnitType::CodedSliceDataPartitionC => "CodedSliceDataPartitionC",
61            NalUnitType::CodedSliceIdr => "CodedSliceIdr",
62            NalUnitType::SEI => "SEI",
63            NalUnitType::SPS => "SPS",
64            NalUnitType::PPS => "PPS",
65            NalUnitType::AUD => "AUD",
66            NalUnitType::EndOfSequence => "EndOfSequence",
67            NalUnitType::EndOfStream => "EndOfStream",
68            NalUnitType::Filler => "Filler",
69            NalUnitType::SpsExt => "SpsExt",
70            NalUnitType::CodedSliceAux => "NalUnitTypeCodedSliceAux",
71            _ => "Reserved",
72        };
73        write!(f, "{}({})", s, *self as u8)
74    }
75}
76
77impl From<u8> for NalUnitType {
78    fn from(v: u8) -> Self {
79        match v {
80            0 => NalUnitType::Unspecified,
81            1 => NalUnitType::CodedSliceNonIdr,
82            2 => NalUnitType::CodedSliceDataPartitionA,
83            3 => NalUnitType::CodedSliceDataPartitionB,
84            4 => NalUnitType::CodedSliceDataPartitionC,
85            5 => NalUnitType::CodedSliceIdr,
86            6 => NalUnitType::SEI,
87            7 => NalUnitType::SPS,
88            8 => NalUnitType::PPS,
89            9 => NalUnitType::AUD,
90            10 => NalUnitType::EndOfSequence,
91            11 => NalUnitType::EndOfStream,
92            12 => NalUnitType::Filler,
93            13 => NalUnitType::SpsExt,
94            19 => NalUnitType::CodedSliceAux,
95            _ => NalUnitType::Reserved,
96        }
97    }
98}
99
100/// NAL H.264 Network Abstraction Layer
101pub struct NAL {
102    pub picture_order_count: u32,
103
104    /// NAL header
105    pub forbidden_zero_bit: bool,
106    pub ref_idc: u8,
107    pub unit_type: NalUnitType,
108
109    /// header byte + rbsp
110    pub data: BytesMut,
111}
112
113impl NAL {
114    fn new(data: BytesMut) -> Self {
115        NAL {
116            picture_order_count: 0,
117            forbidden_zero_bit: false,
118            ref_idc: 0,
119            unit_type: NalUnitType::Unspecified,
120            data,
121        }
122    }
123
124    fn parse_header(&mut self) {
125        let first_byte = self.data[0];
126        self.forbidden_zero_bit = ((first_byte & 0x80) >> 7) == 1; // 0x80 = 0b10000000
127        self.ref_idc = (first_byte & 0x60) >> 5; // 0x60 = 0b01100000
128        self.unit_type = NalUnitType::from(first_byte & 0x1F); // 0x1F = 0b00011111
129    }
130}
131
132const NAL_PREFIX_3BYTES: [u8; 3] = [0, 0, 1];
133const NAL_PREFIX_4BYTES: [u8; 4] = [0, 0, 0, 1];
134
135/// Wrapper class around reading buffer
136struct ReadBuffer {
137    buffer: Box<[u8]>,
138    read_end: usize,
139    filled_end: usize,
140}
141
142impl ReadBuffer {
143    fn new(capacity: usize) -> ReadBuffer {
144        Self {
145            buffer: vec![0u8; capacity].into_boxed_slice(),
146            read_end: 0,
147            filled_end: 0,
148        }
149    }
150
151    #[inline]
152    fn in_buffer(&self) -> usize {
153        self.filled_end - self.read_end
154    }
155
156    fn consume(&mut self, consume: usize) -> &[u8] {
157        debug_assert!(self.read_end + consume <= self.filled_end);
158        let result = &self.buffer[self.read_end..][..consume];
159        self.read_end += consume;
160        result
161    }
162
163    pub(crate) fn fill_buffer(&mut self, reader: &mut impl Read) -> Result<()> {
164        debug_assert_eq!(self.read_end, self.filled_end);
165
166        self.read_end = 0;
167        self.filled_end = reader.read(&mut self.buffer)?;
168
169        Ok(())
170    }
171}
172
173/// H264Reader reads data from stream and constructs h264 nal units
174pub struct H264Reader<R: Read> {
175    reader: R,
176    // reading buffers
177    buffer: ReadBuffer,
178    // for reading
179    nal_prefix_parsed: bool,
180    count_of_consecutive_zero_bytes: usize,
181    nal_buffer: BytesMut,
182}
183
184impl<R: Read> H264Reader<R> {
185    /// new creates new `H264Reader` with `capacity` sized read buffer.
186    pub fn new(reader: R, capacity: usize) -> H264Reader<R> {
187        H264Reader {
188            reader,
189            nal_prefix_parsed: false,
190            buffer: ReadBuffer::new(capacity),
191            count_of_consecutive_zero_bytes: 0,
192            nal_buffer: BytesMut::new(),
193        }
194    }
195
196    fn read4(&mut self) -> Result<([u8; 4], usize)> {
197        let mut result = [0u8; 4];
198        let mut result_filled = 0;
199        loop {
200            let in_buffer = self.buffer.in_buffer();
201
202            if in_buffer + result_filled >= 4 {
203                let consume = 4 - result_filled;
204                result[result_filled..].copy_from_slice(self.buffer.consume(consume));
205                return Ok((result, 4));
206            }
207
208            result[result_filled..][..in_buffer].copy_from_slice(self.buffer.consume(in_buffer));
209            result_filled += in_buffer;
210
211            self.buffer.fill_buffer(&mut self.reader)?;
212
213            if self.buffer.in_buffer() == 0 {
214                return Ok((result, result_filled));
215            }
216        }
217    }
218
219    fn read1(&mut self) -> Result<Option<u8>> {
220        if self.buffer.in_buffer() == 0 {
221            self.buffer.fill_buffer(&mut self.reader)?;
222
223            if self.buffer.in_buffer() == 0 {
224                return Ok(None);
225            }
226        }
227
228        Ok(Some(self.buffer.consume(1)[0]))
229    }
230
231    fn bit_stream_starts_with_h264prefix(&mut self) -> Result<usize> {
232        let (prefix_buffer, n) = self.read4()?;
233
234        if n == 0 {
235            return Err(Error::ErrIoEOF);
236        }
237
238        if n < 3 {
239            return Err(Error::ErrDataIsNotH264Stream);
240        }
241
242        let nal_prefix3bytes_found = NAL_PREFIX_3BYTES[..] == prefix_buffer[..3];
243        if n == 3 {
244            if nal_prefix3bytes_found {
245                return Err(Error::ErrIoEOF);
246            }
247            return Err(Error::ErrDataIsNotH264Stream);
248        }
249
250        // n == 4
251        if nal_prefix3bytes_found {
252            self.nal_buffer.put_u8(prefix_buffer[3]);
253            return Ok(3);
254        }
255
256        let nal_prefix4bytes_found = NAL_PREFIX_4BYTES[..] == prefix_buffer;
257        if nal_prefix4bytes_found {
258            Ok(4)
259        } else {
260            Err(Error::ErrDataIsNotH264Stream)
261        }
262    }
263
264    /// next_nal reads from stream and returns then next NAL,
265    /// and an error if there is incomplete frame data.
266    /// Returns all nil values when no more NALs are available.
267    pub fn next_nal(&mut self) -> Result<NAL> {
268        if !self.nal_prefix_parsed {
269            self.bit_stream_starts_with_h264prefix()?;
270
271            self.nal_prefix_parsed = true;
272        }
273
274        loop {
275            let Some(read_byte) = self.read1()? else {
276                break;
277            };
278
279            let nal_found = self.process_byte(read_byte);
280            if nal_found {
281                let nal_unit_type = NalUnitType::from(self.nal_buffer[0] & 0x1F);
282                if nal_unit_type == NalUnitType::SEI {
283                    self.nal_buffer.clear();
284                    continue;
285                } else {
286                    break;
287                }
288            }
289
290            self.nal_buffer.put_u8(read_byte);
291        }
292
293        if self.nal_buffer.is_empty() {
294            return Err(Error::ErrIoEOF);
295        }
296
297        let mut nal = NAL::new(self.nal_buffer.split());
298        nal.parse_header();
299
300        Ok(nal)
301    }
302
303    fn process_byte(&mut self, read_byte: u8) -> bool {
304        let mut nal_found = false;
305
306        match read_byte {
307            0 => {
308                self.count_of_consecutive_zero_bytes += 1;
309            }
310            1 => {
311                if self.count_of_consecutive_zero_bytes >= 2 {
312                    let count_of_consecutive_zero_bytes_in_prefix =
313                        if self.count_of_consecutive_zero_bytes > 2 {
314                            3
315                        } else {
316                            2
317                        };
318                    let nal_unit_length =
319                        self.nal_buffer.len() - count_of_consecutive_zero_bytes_in_prefix;
320                    if nal_unit_length > 0 {
321                        let _ = self.nal_buffer.split_off(nal_unit_length);
322                        nal_found = true;
323                    }
324                }
325                self.count_of_consecutive_zero_bytes = 0;
326            }
327            _ => {
328                self.count_of_consecutive_zero_bytes = 0;
329            }
330        }
331
332        nal_found
333    }
334}
webrtc_media/io/h264_reader/mod.rs

webrtc_media/io/h264_reader/
mod.rs
