diff --git a/src/demux/flv.rs b/src/demux/flv.rs index f9c2785..b077b66 100644 --- a/src/demux/flv.rs +++ b/src/demux/flv.rs @@ -4,6 +4,7 @@ use std::io::Read; use crate::util; use crate::demux::input; +use crate::demux::nalu; enum FLVTagType { Metadata, @@ -81,16 +82,16 @@ impl input::FileReader for FLVReader { self.skip_init_header()?; return self.read_metadata() } - fn read_nalu(&mut self, _metadata: &util::Metadata) -> Result> { - let (packet_type, data) = self.read_packet()?; + fn read_nalu(&mut self, metadata: &util::Metadata) -> Result> { + let (packet_type, mut data) = self.read_packet()?; return match packet_type { FLVTagType::Audio => Ok(util::NALUPacket { packet_type: util::NALUPacketType::Audio, - packet_data: data, + packet_data: nalu::prepend_a_flv(data, &metadata), }), FLVTagType::Video => Ok(util::NALUPacket { packet_type: util::NALUPacketType::Video, - packet_data: data, + packet_data: nalu::prepend_v_flv(data, &metadata), }), _ => Err(Box::new(util::DemuxerError::FLVUnexpectedTag)) }; diff --git a/src/demux/mod.rs b/src/demux/mod.rs index 1334195..604d10f 100644 --- a/src/demux/mod.rs +++ b/src/demux/mod.rs @@ -1,5 +1,6 @@ mod input; mod flv; +mod nalu; use std::sync::mpsc; use std::error::Error; diff --git a/src/demux/nalu.rs b/src/demux/nalu.rs new file mode 100644 index 0000000..5a0c143 --- /dev/null +++ b/src/demux/nalu.rs @@ -0,0 +1,50 @@ +use crate::util; + +pub fn prepend_v_flv(data: Vec, metadata: &util::Metadata) -> Vec { + match metadata.video.codec { + Some(util::VideoCodec::H264) => preprocess_h264(data), + None => panic!("not possible for codec to be None by this point") + } +} + +pub fn prepend_a_flv(data: Vec, metadata: &util::Metadata) -> Vec { + return data; +} + +fn parse_avcc(mut data: Vec) -> (Vec, Vec) { + data.drain(..6); + let mut buff = [0u8; 2]; + buff.copy_from_slice(data.drain(..2).as_slice()); + let sps_len = u16::from_be_bytes(buff) as usize; + let sps = data.drain(..sps_len).collect::>(); + data.drain(..1); + let mut buff = [0u8; 2]; + buff.copy_from_slice(data.drain(..2).as_slice()); + let pps_len = u16::from_be_bytes(buff) as usize; + let pps = data.drain(..pps_len).collect::>(); + return (sps, pps); +} + +fn preprocess_h264(mut data: Vec) -> Vec { + let mut nalu = vec![0u8, 0u8, 0u8, 1u8]; + if data[1] == 0 { + data.drain(..5); + let (mut sps, mut pps) = parse_avcc(data); + nalu.append(&mut sps); + nalu.append(&mut vec![0u8, 0u8, 0u8, 1u8]); + nalu.append(&mut pps); + return nalu; + } else if data[1] == 2 { + return vec![0u8; 0]; + } else { + data.drain(..5); + loop { + let mut buff = [0u8; 4]; + buff.copy_from_slice(data.drain(..4).as_slice()); + let nalu_len = u32::from_be_bytes(buff) as usize; + nalu.extend_from_slice(data.drain(..nalu_len).as_slice()); + if data.len() <= 4 {return nalu} + nalu.append(&mut vec![0u8, 0u8, 0u8, 1u8]); + } + } +}