diff --git a/src/io/h264_writer/mod.rs b/src/io/h264_writer/mod.rs index 1510d3f..ed7158f 100644 --- a/src/io/h264_writer/mod.rs +++ b/src/io/h264_writer/mod.rs @@ -65,9 +65,9 @@ impl Writer for H264Writer { } if let Some(cached_packet) = &mut self.cached_packet { - cached_packet.depacketize(&packet.payload)?; + let payload = cached_packet.depacketize(&packet.payload)?; - self.writer.write_all(&cached_packet.payload)?; + self.writer.write_all(&payload)?; } Ok(()) diff --git a/src/io/ivf_writer/ivf_writer_test.rs b/src/io/ivf_writer/ivf_writer_test.rs index 8751cbb..41c663c 100644 --- a/src/io/ivf_writer/ivf_writer_test.rs +++ b/src/io/ivf_writer/ivf_writer_test.rs @@ -87,33 +87,29 @@ fn test_ivf_writer_add_packet_and_close() -> Result<()> { // Check valid packet parameters let mut vp8packet = rtp::codecs::vp8::Vp8Packet::default(); - vp8packet.depacketize(&valid_packet.payload)?; + let payload = vp8packet.depacketize(&valid_packet.payload)?; assert_eq!(1, vp8packet.s, "Start packet S value should be 1"); assert_eq!( 1, - vp8packet.payload[0] & 0x01, + payload[0] & 0x01, "Non Keyframe packet P value should be 1" ); // Check mid partition packet parameters let mut vp8packet = rtp::codecs::vp8::Vp8Packet::default(); - vp8packet.depacketize(&mid_part_packet.payload)?; + let payload = vp8packet.depacketize(&mid_part_packet.payload)?; assert_eq!(0, vp8packet.s, "Mid Partition packet S value should be 0"); assert_eq!( 1, - vp8packet.payload[0] & 0x01, + payload[0] & 0x01, "Non Keyframe packet P value should be 1" ); // Check keyframe packet parameters let mut vp8packet = rtp::codecs::vp8::Vp8Packet::default(); - vp8packet.depacketize(&keyframe_packet.payload)?; + let payload = vp8packet.depacketize(&keyframe_packet.payload)?; assert_eq!(1, vp8packet.s, "Start packet S value should be 1"); - assert_eq!( - 0, - vp8packet.payload[0] & 0x01, - "Keyframe packet P value should be 0" - ); + assert_eq!(0, payload[0] & 0x01, "Keyframe packet P value should be 0"); let add_packet_test_case = vec![ ( @@ -170,7 +166,7 @@ fn test_ivf_writer_add_packet_and_close() -> Result<()> { assert!(result.is_ok(), "{}", msg1); } - assert_eq!(seen_key_frame, writer.seen_key_frame); + assert_eq!(seen_key_frame, writer.seen_key_frame, "{} failed", msg1); if count == 1 { assert_eq!(0, writer.count); } else if count == 2 { diff --git a/src/io/ivf_writer/mod.rs b/src/io/ivf_writer/mod.rs index c01ca29..da312a7 100644 --- a/src/io/ivf_writer/mod.rs +++ b/src/io/ivf_writer/mod.rs @@ -16,6 +16,7 @@ pub struct IVFWriter { count: u64, seen_key_frame: bool, current_frame: Option, + is_vp9: bool, } impl IVFWriter { @@ -26,6 +27,7 @@ impl IVFWriter { count: 0, seen_key_frame: false, current_frame: None, + is_vp9: &header.four_cc != b"VP80", }; w.write_header(header)?; @@ -54,24 +56,29 @@ impl IVFWriter { impl Writer for IVFWriter { /// write_rtp adds a new packet and writes the appropriate headers for it fn write_rtp(&mut self, packet: &rtp::packet::Packet) -> Result<()> { - let mut vp8packet = rtp::codecs::vp8::Vp8Packet::default(); - vp8packet.depacketize(&packet.payload)?; + let mut depacketizer: Box = if self.is_vp9 { + Box::new(rtp::codecs::vp9::Vp9Packet::default()) + } else { + Box::new(rtp::codecs::vp8::Vp8Packet::default()) + }; + + let payload = depacketizer.depacketize(&packet.payload)?; - let is_key_frame = vp8packet.payload[0] & 0x01; + let is_key_frame = payload[0] & 0x01; if (!self.seen_key_frame && is_key_frame == 1) - || (self.current_frame.is_none() && vp8packet.s != 1) + || (self.current_frame.is_none() && !depacketizer.is_partition_head(&packet.payload)) { return Ok(()); } self.seen_key_frame = true; let frame_length = if let Some(current_frame) = &mut self.current_frame { - current_frame.extend(vp8packet.payload); + current_frame.extend(payload); current_frame.len() } else { let mut current_frame = BytesMut::new(); - current_frame.extend(vp8packet.payload); + current_frame.extend(payload); let frame_length = current_frame.len(); self.current_frame = Some(current_frame); frame_length diff --git a/src/io/ogg_writer/mod.rs b/src/io/ogg_writer/mod.rs index 649718f..65b2659 100644 --- a/src/io/ogg_writer/mod.rs +++ b/src/io/ogg_writer/mod.rs @@ -159,7 +159,7 @@ impl Writer for OggWriter { /// write_rtp adds a new packet and writes the appropriate headers for it fn write_rtp(&mut self, packet: &rtp::packet::Packet) -> Result<()> { let mut opus_packet = rtp::codecs::opus::OpusPacket::default(); - opus_packet.depacketize(&packet.payload)?; + let payload = opus_packet.depacketize(&packet.payload)?; // Should be equivalent to sample_rate * duration if self.previous_timestamp != 1 { @@ -169,7 +169,7 @@ impl Writer for OggWriter { self.previous_timestamp = packet.header.timestamp; self.write_page( - &opus_packet.payload, + &payload, PAGE_HEADER_TYPE_CONTINUATION_OF_STREAM, self.previous_granule_position, self.page_index,