From 4bcf6c797d2304916206ed41b548b4b2cc0e6229 Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Tue, 10 Oct 2023 14:42:08 +0500
Subject: [PATCH] preprocessing nalu to apply bytestream formatting needed by
 decoder

---
 src/demux/flv.rs  |  9 +++++----
 src/demux/mod.rs  |  1 +
 src/demux/nalu.rs | 50 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 src/demux/nalu.rs

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<util::NALUPacket, Box<dyn Error + Send + Sync>> {
-        let (packet_type, data) = self.read_packet()?;
+    fn read_nalu(&mut self, metadata: &util::Metadata) -> Result<util::NALUPacket, Box<dyn Error + Send + Sync>> {
+        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<u8>, metadata: &util::Metadata) -> Vec<u8> {
+    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<u8>, metadata: &util::Metadata) -> Vec<u8> {
+    return data;
+}
+
+fn parse_avcc(mut data: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
+    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::<Vec<u8>>();
+    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::<Vec<u8>>();
+    return (sps, pps);
+}
+
+fn preprocess_h264(mut data: Vec<u8>) -> Vec<u8> {
+    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]);
+        }
+    }
+}