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