feat: 将rtmp流直接保存成FLV文件，但视频存在一卡一卡的问题

Author: nintha

README.md

@@ -13,12 +13,16 @@ ffplay -fflags nobuffer -analyzeduration 100000 rtmp://localhost:11935/channel/t
 
 使用`-fflags nobuffer -analyzeduration 100000`可以有效降低播放的延迟，目前本地测试延迟大概为1秒
 
-## TODO
-
+## Completed
 - [x] 支持不同分辨率的推流和拉流（之前默认1028x720）
 - [x] 支持音频传输
 
+
+## TODO
+- [x] 输出FLV格式视频
+
 ## 参考资料
 - [RTMP推送AAC ADTS音频流](https://www.jianshu.com/p/1a6f195863c7)
 - [视音频数据处理入门](https://blog.csdn.net/leixiaohua1020/article/details/50534369)
-- [rtmp数据封装](https://blog.csdn.net/Jacob_job/article/details/81880445)
+- [rtmp数据封装](https://blog.csdn.net/Jacob_job/article/details/81880445)
+- [视音频编解码学习工程：FLV封装格式分析器](https://blog.csdn.net/leixiaohua1020/article/details/17934487)

src/protocol/flv.rs

@@ -0,0 +1,71 @@
+use byteorder::{BigEndian, ByteOrder};
+
+use crate::protocol::rtmp::{ChunkMessageType, RtmpMessage};
+
+pub const FLV_HEADER_WITH_TAG0: [u8; 13] = [
+    0x46, 0x4c, 0x56, // signature
+    0x01, // version
+    0x05, // audio and video flag
+    0x00, 0x00, 0x00, 0x09, // header length
+    0x00, 0x00, 0x00, 0x00, // tag0 length
+];
+
+pub struct FlvTag {
+    raw_data: Vec<u8>,
+}
+
+#[allow(unused)]
+impl FlvTag {
+    /// 当RtmpMessage类型不是音频或视频的时候，会返回Error
+    pub fn from_rtmp_message(mut msg: RtmpMessage) -> anyhow::Result<Self> {
+        let mut raw_data = vec![];
+        // data type
+        match msg.header.message_type {
+            ChunkMessageType::AudioMessage => raw_data.push(0x08),
+            ChunkMessageType::VideoMessage => raw_data.push(0x09),
+            _ => Err(anyhow::anyhow!("[FlvTag] invalid message type, {:?}", msg.header.message_type))?,
+        }
+        // data size
+        raw_data.extend_from_slice(&(msg.body.len() as u32).to_be_bytes()[1..4]);
+        // timestamp
+        raw_data.extend_from_slice(&(msg.header.timestamp & 0xFFFFFF).to_be_bytes()[1..4]);
+        // timestamp extended
+        raw_data.push((msg.header.timestamp >> 24) as u8);
+        // stream id
+        raw_data.extend_from_slice(&0u32.to_be_bytes()[1..4]);
+        // body
+        raw_data.append(&mut msg.body);
+
+        Ok(FlvTag { raw_data })
+    }
+
+    /// 0x08=audio, 0x09=video, 0x12=script
+    pub fn tag_type(&self) -> u8 {
+        self.raw_data[0]
+    }
+
+    /// 0x08=audio, 0x09=video, 0x12=script
+    pub fn data_size(&self) -> u32 {
+        BigEndian::read_u24(&self.raw_data[1..4])
+    }
+
+    /// ms timestamp
+    pub fn timestamp(&self) -> u32 {
+        let timestamp_u24 = BigEndian::read_u24(&self.raw_data[4..7]);
+        timestamp_u24 | (self.raw_data[7] as u32) << 24
+    }
+
+    /// audio/video/script data
+    pub fn body(&self) -> &[u8] {
+        &self.raw_data[11..]
+    }
+}
+
+impl AsRef<[u8]> for FlvTag {
+    fn as_ref(&self) -> &[u8] {
+        self.raw_data.as_ref()
+    }
+}
+
+
+

src/protocol/mod.rs

@@ -1,2 +1,3 @@
 pub mod rtmp;
-pub mod h264;
+pub mod h264;
+pub mod flv;

src/protocol/rtmp.rs

@@ -1,19 +1,21 @@
 use std::fmt::{Debug, Formatter};
+use std::sync::{Arc, Weak};
+use std::time::Duration;
 
 use amf::amf0;
 use amf::amf0::Value;
 use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
 use chrono::Local;
 use num::FromPrimitive;
+use smol::channel::Receiver;
 use smol::io::{AsyncReadExt, AsyncWriteExt};
 use smol::net::TcpStream;
+use smol_timeout::TimeoutExt;
 
+use crate::{nalu_eventbus, rtmp_msg_eventbus};
+use crate::protocol::flv::FlvTag;
 use crate::util::{bytes_hex_format, spawn_and_log_error};
-use smol::channel::{Receiver};
-use crate::nalu_eventbus;
-use std::sync::{Arc, Weak};
-use smol_timeout::TimeoutExt;
-use std::time::Duration;
+use crate::protocol::flv;
 
 #[derive(Clone, Debug)]
 pub struct Handshake0 {
@@ -106,16 +108,22 @@ pub struct RtmpContext {
     pub chunk_size: u32,
     pub remain_message_length: u32,
     pub recv_bytes_num: u32,
-    pub arc_receiver: Arc<Receiver<Vec<u8>>>,
+    pub nalu_rx: Arc<Receiver<Vec<u8>>>,
     pub peer_addr: String,
     pub stream_name: String,
+    pub flv_rx: Arc<Receiver<RtmpMessage>>
 }
 
 impl RtmpContext {
     pub fn new(stream: TcpStream) -> Self {
-        let receiver = nalu_eventbus().register_receiver();
-        let receiver = Arc::new(receiver);
-        let nalu_rx_weak = Arc::downgrade(&receiver);
+        let nalu_rx = nalu_eventbus().register_receiver();
+        let nalu_rx = Arc::new(nalu_rx);
+        let nalu_rx_weak = Arc::downgrade(&nalu_rx);
+
+        let flv_rx = rtmp_msg_eventbus().register_receiver();
+        let flv_rx = Arc::new(flv_rx);
+        let flv_rx_weak = Arc::downgrade(&flv_rx);
+
         let peer_addr = stream.peer_addr().map(|a| a.to_string()).unwrap_or_default();
 
         async fn handle_nalu_rx(nalu_rx: Weak<Receiver<Vec<u8>>>, peer_addr: String) -> anyhow::Result<()> {
@@ -144,7 +152,39 @@ impl RtmpContext {
             Ok(())
         }
 
+        async fn handle_flv_rx(flv_rx: Weak<Receiver<RtmpMessage>>, peer_addr: String) -> anyhow::Result<()> {
+            let tmp_dir = "tmp";
+            if smol::fs::read_dir(tmp_dir).await.is_err() {
+                smol::fs::create_dir_all(tmp_dir).await?;
+            }
+
+            let mut file = smol::fs::OpenOptions::new()
+                .create(true)
+                .write(true)
+                .truncate(true)
+                .open("tmp/output.flv")
+                .await?;
+
+            // write header
+            file.write_all(&flv::FLV_HEADER_WITH_TAG0).await?;
+
+            while let Some(rx) = flv_rx.upgrade() {
+                if let Some(rs) = rx.recv().timeout(Duration::from_secs(1)).await {
+                    if let Ok(msg) = rs {
+                        let flv_tag = FlvTag::from_rtmp_message(msg)?;
+                        file.write_all(flv_tag.as_ref()).await?;
+                        file.write_all(&(flv_tag.as_ref().len() as u32).to_be_bytes()).await?;
+                    } else {
+                        break;
+                    }
+                }
+            }
+            log::warn!("[peer={}][handle_flv_rx] closed", peer_addr);
+            Ok(())
+        }
+
         spawn_and_log_error(handle_nalu_rx(nalu_rx_weak, peer_addr.clone()));
+        spawn_and_log_error(handle_flv_rx(flv_rx_weak, peer_addr.clone()));
 
         RtmpContext {
             stream,
@@ -157,7 +197,8 @@ impl RtmpContext {
             chunk_size: 128,
             remain_message_length: 0,
             recv_bytes_num: 0,
-            arc_receiver: receiver,
+            nalu_rx,
+            flv_rx,
             peer_addr,
             stream_name: Default::default(),
         }
@@ -178,7 +219,6 @@ impl RtmpContext {
 #[derive(Debug, Clone)]
 pub struct RtmpMessageHeader {
     /// chunk stream id
-    /// 2 (low level), 3 (high level), 4 (control stream), 5 (video) and 6 (audio).
     pub csid: u8,
     pub timestamp: u32,
     pub message_length: u32,