From 67c00703247a81e22fbc8773bfc894cd9424f2a8 Mon Sep 17 00:00:00 2001 From: Shehriyar Qureshi Date: Wed, 9 Nov 2022 14:47:29 +0500 Subject: [PATCH] feat(h264): handle annex b boundary b/w pkts If an Annex B boundary cuts off b/w pkts, `break_apart_nals` will handle and split up nals just as if the boundary was inside single packet --- src/codec/h264.rs | 210 ++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 204 insertions(+), 6 deletions(-) diff --git a/src/codec/h264.rs b/src/codec/h264.rs index badd4d2..ef32aa8 100644 --- a/src/codec/h264.rs +++ b/src/codec/h264.rs @@ -86,15 +86,54 @@ impl NalParser { /// boundaries (i.e. three consecutive 0x00). If it finds a boundary, it splits on it /// to break apart NALs from the Annex B stream. fn break_apart_nals(&mut self, data: Bytes) -> Result { + // Should `&mut Bytes` be passed instead of doing this? + let mut data_copy = data; + let mut nal_start_idx = 0; let mut did_find_boundary = false; - for (idx, byte) in data.iter().enumerate() { + // Check if current packet starts with either a single or two 0x00 bytes. + // If so, check if previous piece ended with either one or two 0x00 bytes. + // Case 1: [ [.., 0x00], [pkt_header, nal_header, 0x00, 0x00, ..] ] + // Case 2: [ [.., 0x00, 0x00], [pkt_header, nal_header, 0x00, ..] ] + // If we match any of the above cases, we remove the 0x00 that form the boundary + // and start off a new nal after it. + if !self.pieces.is_empty() && data_copy.slice(..2)[..] == [0x00; 2] { + let last_piece = self + .pieces + .last_mut() + .expect("pieces should be non-empty because break_apart_nals checked for it to be non-empty"); + if last_piece.ends_with(&[0x00]) { + let last_byte_index = last_piece.len() - 1; + last_piece.truncate(last_byte_index); + self.nals.last_mut().unwrap().len -= 1; + data_copy.advance(3); + let nal_header = NalHeader::new(data_copy[0]).expect("Header w/o F bit is valid"); + self.start_rtp_nal(nal_header)?; + data_copy.advance(1); + } + } else if !self.pieces.is_empty() && data_copy.first().unwrap() == &0x00 { + let last_piece = self + .pieces + .last_mut() + .expect("pieces should be non-empty because break_apart_nals checked for it to be non-empty"); + if last_piece.ends_with(&[0x00, 0x00]) { + let last_byte_index = last_piece.len() - 1; + last_piece.truncate(last_byte_index - 1); + self.nals.last_mut().unwrap().len -= 1; + data_copy.advance(2); + let nal_header = NalHeader::new(data_copy[0]).expect("Header w/o F bit is valid"); + self.start_rtp_nal(nal_header)?; + data_copy.advance(1); + } + } + + for (idx, byte) in data_copy.iter().enumerate() { // TODO: Handle boundaries split b/w packets. // If the current FU-A has a boundary that splits at end, ignore the last or // last two zeros because this boundary will be handled when the start of // next packet is being read, and these zeros will be removed on walk-back. - if byte == &0x00 && idx + 2 < data.len() && data[idx..idx + 3] == [0x00; 3] { + if byte == &0x00 && idx + 2 < data_copy.len() && data_copy[idx..idx + 3] == [0x00; 3] { debug!("Found boundary with index range: {} - {}.", idx, idx + 2); // we found a boundary, let NalParser know that it should now keep adding // to last NAL even if the next FU-A frag header byte does not match the @@ -102,7 +141,7 @@ impl NalParser { did_find_boundary = true; let nal_end_idx = idx; - let nal = data.slice( + let nal = data_copy.slice( if nal_start_idx == 0 { // this is only for the first boundary, since `data` passed already has advanced 2 bytes // to skip the packet type and NAL header @@ -119,7 +158,7 @@ impl NalParser { nal_start_idx = idx + 3; // create new nal which'll get updated - let nal_header = data[idx + 4]; + let nal_header = data_copy[idx + 4]; self.nals.push(Nal { hdr: NalHeader::new(nal_header).expect("header w/o F bit set is valid"), next_piece_idx: u32::MAX, @@ -130,7 +169,7 @@ impl NalParser { // if we had found a boundary, we need to add the last NAL to pieces now if did_find_boundary { - self.push(data.slice(nal_start_idx + 2..))?; + self.push(data_copy.slice(nal_start_idx + 2..))?; } Ok(did_find_boundary) @@ -1125,11 +1164,13 @@ enum PacketizerState { mod tests { use std::num::NonZeroU32; - use h264_reader::nal::UnitType; + use h264_reader::nal::{NalHeader, UnitType}; use crate::testutil::init_logging; use crate::{codec::CodecItem, rtp::ReceivedPacketBuilder}; + use super::NalParser; + /* * This test requires * 1. a hacked version of the "mp4" crate to fix a couple bugs @@ -1817,4 +1858,161 @@ mod tests { ); assert!(push_result.is_err()); } + + #[test] + fn split_annex_b_from_single_packet() { + init_logging(); + let timestamp = crate::Timestamp { + timestamp: 0, + clock_rate: NonZeroU32::new(90_000).unwrap(), + start: 0, + }; + let rtp_pkt_1 = ReceivedPacketBuilder { + // FU-A start fragment + ctx: crate::PacketContext::dummy(), + stream_id: 0, + timestamp, + ssrc: 0, + sequence_number: 0, + loss: 0, + mark: false, + payload_type: 0, + } + .build(*b"\x3c\x07previous-nal\x00\x00\x00\x3c\x08new-nal") + .unwrap(); + + let rtp_payload = rtp_pkt_1.into_payload_bytes(); + let nal_header = NalHeader::new(rtp_payload[1]).unwrap(); + + let mut nal_parser = NalParser::new(); + nal_parser.start_rtp_nal(nal_header).unwrap(); + nal_parser.append_rtp_nal(rtp_payload.slice(2..)).unwrap(); + + assert_eq!(nal_parser.nals.len(), 2); + assert_eq!(nal_parser.pieces.len(), 2); + assert_eq!( + nal_parser.nals.first().unwrap().hdr.nal_unit_type(), + UnitType::SeqParameterSet + ); + assert_eq!( + nal_parser.nals.last().unwrap().hdr.nal_unit_type(), + UnitType::PicParameterSet + ); + } + + #[test] + fn split_annex_b_boundary_cut_off_between_packets_variation_one() { + init_logging(); + let timestamp = crate::Timestamp { + timestamp: 0, + clock_rate: NonZeroU32::new(90_000).unwrap(), + start: 0, + }; + let rtp_pkt_1 = ReceivedPacketBuilder { + // FU-A start fragment + ctx: crate::PacketContext::dummy(), + stream_id: 0, + timestamp, + ssrc: 0, + sequence_number: 0, + loss: 0, + mark: false, + payload_type: 0, + } + .build(*b"\x3c\x07previous-nal\x00\x00") + .unwrap(); + let rtp_pkt_2 = ReceivedPacketBuilder { + // FU-A start fragment + ctx: crate::PacketContext::dummy(), + stream_id: 0, + timestamp, + ssrc: 0, + sequence_number: 0, + loss: 0, + mark: false, + payload_type: 0, + } + .build(*b"\x3c\x07\x00\x3c\x08new-nal") + .unwrap(); + + let rtp_1_payload = rtp_pkt_1.into_payload_bytes(); + let rtp_2_payload = rtp_pkt_2.into_payload_bytes(); + let nal_header_1 = NalHeader::new(rtp_1_payload[1]).unwrap(); + + let mut nal_parser = NalParser::new(); + nal_parser.start_rtp_nal(nal_header_1).unwrap(); + nal_parser.append_rtp_nal(rtp_1_payload.slice(2..)).unwrap(); + assert_eq!(nal_parser.nals.len(), 1); + assert_eq!(nal_parser.pieces.len(), 1); + nal_parser.append_rtp_nal(rtp_2_payload.slice(2..)).unwrap(); + + assert_eq!(nal_parser.nals.len(), 2); + assert_eq!(nal_parser.pieces.len(), 2); + assert_eq!( + nal_parser.nals.first().unwrap().hdr.nal_unit_type(), + UnitType::SeqParameterSet + ); + assert_eq!( + nal_parser.nals.last().unwrap().hdr.nal_unit_type(), + UnitType::PicParameterSet + ); + } + + #[test] + fn split_annex_b_boundary_cut_off_between_packets_variation_two() { + init_logging(); + let timestamp = crate::Timestamp { + timestamp: 0, + clock_rate: NonZeroU32::new(90_000).unwrap(), + start: 0, + }; + let rtp_pkt_1 = ReceivedPacketBuilder { + // FU-A start fragment + ctx: crate::PacketContext::dummy(), + stream_id: 0, + timestamp, + ssrc: 0, + sequence_number: 0, + loss: 0, + mark: false, + payload_type: 0, + } + .build(*b"\x3c\x07previous-nal\x00") + .unwrap(); + let rtp_pkt_2 = ReceivedPacketBuilder { + // FU-A start fragment + ctx: crate::PacketContext::dummy(), + stream_id: 0, + timestamp, + ssrc: 0, + sequence_number: 0, + loss: 0, + mark: false, + payload_type: 0, + } + .build(*b"\x3c\x07\x00\x00\x3c\x08new-nal") + .unwrap(); + + let rtp_1_payload = rtp_pkt_1.into_payload_bytes(); + let rtp_2_payload = rtp_pkt_2.into_payload_bytes(); + let nal_header_1 = NalHeader::new(rtp_1_payload[1]).unwrap(); + + let mut nal_parser = NalParser::new(); + nal_parser.start_rtp_nal(nal_header_1).unwrap(); + nal_parser.append_rtp_nal(rtp_1_payload.slice(2..)).unwrap(); + assert_eq!(nal_parser.nals.len(), 1); + assert_eq!(nal_parser.pieces.len(), 1); + nal_parser.append_rtp_nal(rtp_2_payload.slice(2..)).unwrap(); + + assert_eq!(nal_parser.nals.len(), 2); + assert_eq!(nal_parser.pieces.len(), 2); + assert_eq!( + nal_parser.nals.first().unwrap().hdr.nal_unit_type(), + UnitType::SeqParameterSet + ); + assert_eq!( + nal_parser.nals.last().unwrap().hdr.nal_unit_type(), + UnitType::PicParameterSet + ); + } }