1 #[cfg(test)] 2 mod h264_test; 3 4 use crate::{ 5 error::{Error, Result}, 6 packetizer::{Depacketizer, Payloader}, 7 }; 8 9 use bytes::{BufMut, Bytes, BytesMut}; 10 11 /// H264Payloader payloads H264 packets 12 #[derive(Default, Debug, Clone)] 13 pub struct H264Payloader { 14 sps_nalu: Option<Bytes>, 15 pps_nalu: Option<Bytes>, 16 } 17 18 pub const STAPA_NALU_TYPE: u8 = 24; 19 pub const FUA_NALU_TYPE: u8 = 28; 20 pub const FUB_NALU_TYPE: u8 = 29; 21 pub const SPS_NALU_TYPE: u8 = 7; 22 pub const PPS_NALU_TYPE: u8 = 8; 23 pub const AUD_NALU_TYPE: u8 = 9; 24 pub const FILLER_NALU_TYPE: u8 = 12; 25 26 pub const FUA_HEADER_SIZE: usize = 2; 27 pub const STAPA_HEADER_SIZE: usize = 1; 28 pub const STAPA_NALU_LENGTH_SIZE: usize = 2; 29 30 pub const NALU_TYPE_BITMASK: u8 = 0x1F; 31 pub const NALU_REF_IDC_BITMASK: u8 = 0x60; 32 pub const FU_START_BITMASK: u8 = 0x80; 33 pub const FU_END_BITMASK: u8 = 0x40; 34 35 pub const OUTPUT_STAP_AHEADER: u8 = 0x78; 36 37 pub static ANNEXB_NALUSTART_CODE: Bytes = Bytes::from_static(&[0x00, 0x00, 0x00, 0x01]); 38 39 impl H264Payloader { next_ind(nalu: &Bytes, start: usize) -> (isize, isize)40 fn next_ind(nalu: &Bytes, start: usize) -> (isize, isize) { 41 let mut zero_count = 0; 42 43 for (i, &b) in nalu[start..].iter().enumerate() { 44 if b == 0 { 45 zero_count += 1; 46 continue; 47 } else if b == 1 && zero_count >= 2 { 48 return ((start + i - zero_count) as isize, zero_count as isize + 1); 49 } 50 zero_count = 0 51 } 52 (-1, -1) 53 } 54 emit(&mut self, nalu: &Bytes, mtu: usize, payloads: &mut Vec<Bytes>)55 fn emit(&mut self, nalu: &Bytes, mtu: usize, payloads: &mut Vec<Bytes>) { 56 if nalu.is_empty() { 57 return; 58 } 59 60 let nalu_type = nalu[0] & NALU_TYPE_BITMASK; 61 let nalu_ref_idc = nalu[0] & NALU_REF_IDC_BITMASK; 62 63 if nalu_type == AUD_NALU_TYPE || nalu_type == FILLER_NALU_TYPE { 64 return; 65 } else if nalu_type == SPS_NALU_TYPE { 66 self.sps_nalu = Some(nalu.clone()); 67 return; 68 } else if nalu_type == PPS_NALU_TYPE { 69 self.pps_nalu = Some(nalu.clone()); 70 return; 71 } else if let (Some(sps_nalu), Some(pps_nalu)) = (&self.sps_nalu, &self.pps_nalu) { 72 // Pack current NALU with SPS and PPS as STAP-A 73 let sps_len = (sps_nalu.len() as u16).to_be_bytes(); 74 let pps_len = (pps_nalu.len() as u16).to_be_bytes(); 75 76 let mut stap_a_nalu = Vec::with_capacity(1 + 2 + sps_nalu.len() + 2 + pps_nalu.len()); 77 stap_a_nalu.push(OUTPUT_STAP_AHEADER); 78 stap_a_nalu.extend(sps_len); 79 stap_a_nalu.extend_from_slice(sps_nalu); 80 stap_a_nalu.extend(pps_len); 81 stap_a_nalu.extend_from_slice(pps_nalu); 82 if stap_a_nalu.len() <= mtu { 83 payloads.push(Bytes::from(stap_a_nalu)); 84 } 85 } 86 87 if self.sps_nalu.is_some() && self.pps_nalu.is_some() { 88 self.sps_nalu = None; 89 self.pps_nalu = None; 90 } 91 92 // Single NALU 93 if nalu.len() <= mtu { 94 payloads.push(nalu.clone()); 95 return; 96 } 97 98 // FU-A 99 let max_fragment_size = mtu as isize - FUA_HEADER_SIZE as isize; 100 101 // The FU payload consists of fragments of the payload of the fragmented 102 // NAL unit so that if the fragmentation unit payloads of consecutive 103 // FUs are sequentially concatenated, the payload of the fragmented NAL 104 // unit can be reconstructed. The NAL unit type octet of the fragmented 105 // NAL unit is not included as such in the fragmentation unit payload, 106 // but rather the information of the NAL unit type octet of the 107 // fragmented NAL unit is conveyed in the F and NRI fields of the FU 108 // indicator octet of the fragmentation unit and in the type field of 109 // the FU header. An FU payload MAY have any number of octets and MAY 110 // be empty. 111 112 let nalu_data = nalu; 113 // According to the RFC, the first octet is skipped due to redundant information 114 let mut nalu_data_index = 1; 115 let nalu_data_length = nalu.len() as isize - nalu_data_index; 116 let mut nalu_data_remaining = nalu_data_length; 117 118 if std::cmp::min(max_fragment_size, nalu_data_remaining) <= 0 { 119 return; 120 } 121 122 while nalu_data_remaining > 0 { 123 let current_fragment_size = std::cmp::min(max_fragment_size, nalu_data_remaining); 124 //out: = make([]byte, fuaHeaderSize + currentFragmentSize) 125 let mut out = BytesMut::with_capacity(FUA_HEADER_SIZE + current_fragment_size as usize); 126 // +---------------+ 127 // |0|1|2|3|4|5|6|7| 128 // +-+-+-+-+-+-+-+-+ 129 // |F|NRI| Type | 130 // +---------------+ 131 let b0 = FUA_NALU_TYPE | nalu_ref_idc; 132 out.put_u8(b0); 133 134 // +---------------+ 135 //|0|1|2|3|4|5|6|7| 136 //+-+-+-+-+-+-+-+-+ 137 //|S|E|R| Type | 138 //+---------------+ 139 140 let mut b1 = nalu_type; 141 if nalu_data_remaining == nalu_data_length { 142 // Set start bit 143 b1 |= 1 << 7; 144 } else if nalu_data_remaining - current_fragment_size == 0 { 145 // Set end bit 146 b1 |= 1 << 6; 147 } 148 out.put_u8(b1); 149 150 out.put( 151 &nalu_data 152 [nalu_data_index as usize..(nalu_data_index + current_fragment_size) as usize], 153 ); 154 payloads.push(out.freeze()); 155 156 nalu_data_remaining -= current_fragment_size; 157 nalu_data_index += current_fragment_size; 158 } 159 } 160 } 161 162 impl Payloader for H264Payloader { 163 /// Payload fragments a H264 packet across one or more byte arrays payload(&mut self, mtu: usize, payload: &Bytes) -> Result<Vec<Bytes>>164 fn payload(&mut self, mtu: usize, payload: &Bytes) -> Result<Vec<Bytes>> { 165 if payload.is_empty() || mtu == 0 { 166 return Ok(vec![]); 167 } 168 169 let mut payloads = vec![]; 170 171 let (mut next_ind_start, mut next_ind_len) = H264Payloader::next_ind(payload, 0); 172 if next_ind_start == -1 { 173 self.emit(payload, mtu, &mut payloads); 174 } else { 175 while next_ind_start != -1 { 176 let prev_start = (next_ind_start + next_ind_len) as usize; 177 let (next_ind_start2, next_ind_len2) = H264Payloader::next_ind(payload, prev_start); 178 next_ind_start = next_ind_start2; 179 next_ind_len = next_ind_len2; 180 if next_ind_start != -1 { 181 self.emit( 182 &payload.slice(prev_start..next_ind_start as usize), 183 mtu, 184 &mut payloads, 185 ); 186 } else { 187 // Emit until end of stream, no end indicator found 188 self.emit(&payload.slice(prev_start..), mtu, &mut payloads); 189 } 190 } 191 } 192 193 Ok(payloads) 194 } 195 clone_to(&self) -> Box<dyn Payloader + Send + Sync>196 fn clone_to(&self) -> Box<dyn Payloader + Send + Sync> { 197 Box::new(self.clone()) 198 } 199 } 200 201 /// H264Packet represents the H264 header that is stored in the payload of an RTP Packet 202 #[derive(PartialEq, Eq, Debug, Default, Clone)] 203 pub struct H264Packet { 204 pub is_avc: bool, 205 fua_buffer: Option<BytesMut>, 206 } 207 208 impl Depacketizer for H264Packet { 209 /// depacketize parses the passed byte slice and stores the result in the H264Packet this method is called upon depacketize(&mut self, packet: &Bytes) -> Result<Bytes>210 fn depacketize(&mut self, packet: &Bytes) -> Result<Bytes> { 211 if packet.len() <= 2 { 212 return Err(Error::ErrShortPacket); 213 } 214 215 let mut payload = BytesMut::new(); 216 217 // NALU Types 218 // https://tools.ietf.org/html/rfc6184#section-5.4 219 let b0 = packet[0]; 220 let nalu_type = b0 & NALU_TYPE_BITMASK; 221 222 match nalu_type { 223 1..=23 => { 224 if self.is_avc { 225 payload.put_u32(packet.len() as u32); 226 } else { 227 payload.put(&*ANNEXB_NALUSTART_CODE); 228 } 229 payload.put(&*packet.clone()); 230 Ok(payload.freeze()) 231 } 232 STAPA_NALU_TYPE => { 233 let mut curr_offset = STAPA_HEADER_SIZE; 234 while curr_offset < packet.len() { 235 let nalu_size = 236 ((packet[curr_offset] as usize) << 8) | packet[curr_offset + 1] as usize; 237 curr_offset += STAPA_NALU_LENGTH_SIZE; 238 239 if packet.len() < curr_offset + nalu_size { 240 return Err(Error::StapASizeLargerThanBuffer( 241 nalu_size, 242 packet.len() - curr_offset, 243 )); 244 } 245 246 if self.is_avc { 247 payload.put_u32(nalu_size as u32); 248 } else { 249 payload.put(&*ANNEXB_NALUSTART_CODE); 250 } 251 payload.put(&*packet.slice(curr_offset..curr_offset + nalu_size)); 252 curr_offset += nalu_size; 253 } 254 255 Ok(payload.freeze()) 256 } 257 FUA_NALU_TYPE => { 258 if packet.len() < FUA_HEADER_SIZE { 259 return Err(Error::ErrShortPacket); 260 } 261 262 if self.fua_buffer.is_none() { 263 self.fua_buffer = Some(BytesMut::new()); 264 } 265 266 if let Some(fua_buffer) = &mut self.fua_buffer { 267 fua_buffer.put(&*packet.slice(FUA_HEADER_SIZE..)); 268 } 269 270 let b1 = packet[1]; 271 if b1 & FU_END_BITMASK != 0 { 272 let nalu_ref_idc = b0 & NALU_REF_IDC_BITMASK; 273 let fragmented_nalu_type = b1 & NALU_TYPE_BITMASK; 274 275 if let Some(fua_buffer) = self.fua_buffer.take() { 276 if self.is_avc { 277 payload.put_u32((fua_buffer.len() + 1) as u32); 278 } else { 279 payload.put(&*ANNEXB_NALUSTART_CODE); 280 } 281 payload.put_u8(nalu_ref_idc | fragmented_nalu_type); 282 payload.put(fua_buffer); 283 } 284 285 Ok(payload.freeze()) 286 } else { 287 Ok(Bytes::new()) 288 } 289 } 290 _ => Err(Error::NaluTypeIsNotHandled(nalu_type)), 291 } 292 } 293 294 /// is_partition_head checks if this is the head of a packetized nalu stream. is_partition_head(&self, payload: &Bytes) -> bool295 fn is_partition_head(&self, payload: &Bytes) -> bool { 296 if payload.len() < 2 { 297 return false; 298 } 299 300 if payload[0] & NALU_TYPE_BITMASK == FUA_NALU_TYPE 301 || payload[0] & NALU_TYPE_BITMASK == FUB_NALU_TYPE 302 { 303 (payload[1] & FU_START_BITMASK) != 0 304 } else { 305 true 306 } 307 } 308 is_partition_tail(&self, marker: bool, _payload: &Bytes) -> bool309 fn is_partition_tail(&self, marker: bool, _payload: &Bytes) -> bool { 310 marker 311 } 312 } 313