enhancement for zip chunk headers

harlanc · Oct 28, 2023 · 23b4072 · 23b4072
1 parent b2fe9af
commit 23b4072
Showing 1 changed file with 29 additions and 17 deletions.
diff --git a/protocol/rtmp/src/chunk/packetizer.rs b/protocol/rtmp/src/chunk/packetizer.rs
@@ -38,27 +38,39 @@ impl ChunkPacketizer {
         chunk_info.basic_header.format = 0;
 
         let pre_header = self
-            .csid_2_chunk_header
-            .get_mut(&chunk_info.basic_header.chunk_stream_id);
-
-        if let Some(val) = pre_header {
-            let cur_msg_header = &mut chunk_info.message_header;
-            let pre_msg_header = &val.message_header;
-
-            if cur_msg_header.msg_streamd_id == pre_msg_header.msg_streamd_id {
-                chunk_info.basic_header.format = 1;
-                cur_msg_header.timestamp -= pre_msg_header.timestamp;
-
-                if cur_msg_header.msg_type_id == pre_msg_header.msg_type_id
-                    && cur_msg_header.msg_length == pre_msg_header.msg_length
-                {
-                    chunk_info.basic_header.format = 2;
-                    if chunk_info.message_header.timestamp == pre_msg_header.timestamp {
-                        chunk_info.basic_header.format = 3;
+        .csid_2_chunk_header
+        .get_mut(&chunk_info.basic_header.chunk_stream_id);
+
+        match pre_header {
+            None => {
+                self.csid_2_chunk_header.insert(
+                    chunk_info.basic_header.chunk_stream_id,
+                    ChunkHeader {
+                        basic_header: chunk_info.basic_header.clone(),
+                        message_header: chunk_info.message_header.clone(),
+                    },
+                );
+            }
+            Some(val) => {
+                let cur_msg_header = &mut chunk_info.message_header;
+                let pre_msg_header = &val.message_header;
+
+                if cur_msg_header.msg_streamd_id == pre_msg_header.msg_streamd_id {
+                    chunk_info.basic_header.format = 1;
+                    cur_msg_header.timestamp -= pre_msg_header.timestamp;
+
+                    if cur_msg_header.msg_type_id == pre_msg_header.msg_type_id
+                        && cur_msg_header.msg_length == pre_msg_header.msg_length
+                    {
+                        chunk_info.basic_header.format = 2;
+                        if chunk_info.message_header.timestamp == pre_msg_header.timestamp {
+                            chunk_info.basic_header.format = 3;
+                        }
                     }
                 }
             }
         }
+
         Ok(PackResult::Success)
     }