From f84237b0be12fe5853cf649eace38491175beaba Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Fri, 6 Oct 2023 16:35:03 +0500
Subject: [PATCH] FLV NALU implementation

---
 src/demux/flv.rs | 29 ++++++++++++++++++++++-------
 src/util/mod.rs  |  3 +++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/demux/flv.rs b/src/demux/flv.rs
index 15a2ce2..b84bd23 100644
--- a/src/demux/flv.rs
+++ b/src/demux/flv.rs
@@ -46,7 +46,7 @@ impl FLVReader {
         return Ok(());
     }
 
-    fn read_packet(&mut self) -> Result<(FLVTagType, Vec<u8>), Box<dyn Error>> {
+    fn read_packet(&mut self) -> Result<(FLVTagType, Vec<u8>), Box<dyn Error + Send + Sync>> {
         let mut tag_head = [0u8; 11];
         self.stdin.read_exact(&mut tag_head)?;
         let tag_type = FLVTagType::try_from(tag_head[0])?;
@@ -60,10 +60,13 @@ impl FLVReader {
     }
 
     fn read_metadata(&mut self) -> Result<util::Metadata, Box<dyn Error>> {
-        let (packet_type, packet_data) = self.read_packet()?;
+        let (packet_type, packet_data) = match self.read_packet() {
+            Ok(x) => x,
+            Err(e) => return Err(e)
+        };
         match packet_type {
             FLVTagType::Metadata => (),
-            x => return Err(Box::new(util::DemuxerError::FLVPacketNotMetadata))
+            _ => return Err(Box::new(util::DemuxerError::FLVPacketNotMetadata))
         }
         return process_metadata(packet_data);
     }
@@ -74,8 +77,20 @@ 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>> {
-        todo!();
+    fn read_nalu(&mut self, _metadata: &util::Metadata) -> Result<util::NALUPacket, Box<dyn Error + Send + Sync>> {
+        let (packet_type, data) = self.read_packet()?;
+        dbg!(&data);
+        return match packet_type {
+            FLVTagType::Audio => Ok(util::NALUPacket { 
+                packet_type: util::NALUPacketType::Audio,
+                packet_data: data,
+            }),
+            FLVTagType::Video => Ok(util::NALUPacket {
+                packet_type: util::NALUPacketType::Video,
+                packet_data: data,
+            }),
+            _ => Err(Box::new(util::DemuxerError::FLVUnexpectedTag))
+        };
     }
 }
 
@@ -96,7 +111,7 @@ fn process_metadata(mut data: Vec<u8>) -> Result<util::Metadata, Box<dyn Error>>
     for _ in 0..arr_len {
         let key = String::fetch_amf_val(&mut data);
         let amf_marker = u8::fetch_amf_val(&mut data);
-        if !handle_relevant_keys(key, &mut data, &amf_marker, &mut metadata)? {
+        if !handle_relevant_keys(key, &mut data, &mut metadata)? {
             match amf_marker {
                 0 => {let _ = f64::fetch_amf_val(&mut data);},
                 1 => {let _ = bool::fetch_amf_val(&mut data);},
@@ -108,7 +123,7 @@ fn process_metadata(mut data: Vec<u8>) -> Result<util::Metadata, Box<dyn Error>>
     return if metadata.is_valid() {Ok(metadata)} else {Err(Box::new(util::DemuxerError::MetadataValNotSet))};
 }
 
-fn handle_relevant_keys(key: String, data: &mut Vec<u8>, marker: &u8, metadata: &mut util::Metadata) -> Result<bool, Box<dyn Error>> {
+fn handle_relevant_keys(key: String, data: &mut Vec<u8>, metadata: &mut util::Metadata) -> Result<bool, Box<dyn Error>> {
     match key.as_str() {
         "width" => {metadata.video.width = f64::fetch_amf_val(data) as u32; Ok(true)},
         "height" => {metadata.video.height = f64::fetch_amf_val(data) as u32; Ok(true)},
diff --git a/src/util/mod.rs b/src/util/mod.rs
index 6f11f5e..262af45 100644
--- a/src/util/mod.rs
+++ b/src/util/mod.rs
@@ -83,12 +83,14 @@ pub enum DemuxerError {
     FLVPacketNotMetadata,
     FLVMetadataSyncFail,
     FLVMetadataAMFMarkerError,
+    FLVUnexpectedTag,
     MetadataValNotSet,
     CodecNotImplemented,
 }
 
 impl Error for DemuxerError {}
 
+#[allow(unreachable_patterns)]
 impl fmt::Debug for DemuxerError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
@@ -100,6 +102,7 @@ impl fmt::Debug for DemuxerError {
             DemuxerError::FLVMetadataAMFMarkerError => write!(f, "Marker value not recognized/implemented"),
             DemuxerError::CodecNotImplemented => write!(f, "Input file has unrecognized codec"),
             DemuxerError::MetadataValNotSet => write!(f, "One or more metadata values failed to init"),
+            DemuxerError::FLVUnexpectedTag => write!(f, "Expected data tag, found metadata"),
             _ => write!(f, "Error not described yet")
         }
     }