feat: 实现缓存sps/pps帧信息

Author: nintha

src/main.rs

@@ -13,6 +13,7 @@ use protocol::rtmp::*;
 use crate::util::{bytes_hex_format, print_hex, spawn_and_log_error, gen_random_bytes};
 use crate::eventbus::EventBus;
 use once_cell::sync::OnceCell;
+use dashmap::DashMap;
 
 mod util;
 mod protocol;
@@ -28,6 +29,11 @@ fn rtmp_msg_eventbus() -> &'static EventBus<RtmpMessage> {
     INSTANCE.get_or_init(|| EventBus::with_label("RTMP_MSG"))
 }
 
+fn sps_message_map() -> &'static DashMap<String, RtmpMessage> {
+    static INSTANCE: OnceCell<DashMap<String, RtmpMessage>> = OnceCell::new();
+    INSTANCE.get_or_init(|| DashMap::new())
+}
+
 /// TCP 连接处理
 async fn accept_loop(addr: &str) -> anyhow::Result<()> {
     let listener = TcpListener::bind(addr.clone()).await?;
@@ -65,6 +71,7 @@ async fn connection_loop(stream: TcpStream) -> anyhow::Result<()> {
                     let buffer_length = BigEndian::read_u32(&bytes[6..10]);
                     log::info!("[peer={}] C->S, [{}] set buffer length={}, streamId={}", ctx.peer_addr, message.message_type_desc(), buffer_length, stream_id);
                     response_play(&mut ctx, stream_id).await?;
+                    // TODO 使用动态的值
                     send_meta_data_for_play(&mut ctx, &RtmpMetaData {
                         width: 1280.0,
                         height: 720.0,
@@ -75,24 +82,20 @@ async fn connection_loop(stream: TcpStream) -> anyhow::Result<()> {
                         frame_rate: 30.0,
                         duration: 30.0,
                     }).await?;
-                    // TODO send video stream to peer
-                    fn is_key_frame(msg: &RtmpMessage) -> bool {
-                        if let ChunkMessageType::VideoMessage = msg.header.message_type {
-                            msg.body[0] >> 4 == 1
-                        } else {
-                            false
+
+                    ctx.ctx_begin_timestamp = Local::now().timestamp_millis();
+
+                    // 发送sps/pps帧
+                    if let Some(msg) = sps_message_map().get(&ctx.stream_name) {
+                        let chunks = msg.split_chunks_bytes(ctx.chunk_size);
+                        for chunk in chunks {
+                            ctx.write_to_peer(&chunk).await?;
                         }
                     }
 
-                    let mut wait_key_frame = true;
                     let receiver = rtmp_msg_eventbus().register_receiver();
                     while let Ok(mut msg) = receiver.recv().await {
-                        if wait_key_frame && !is_key_frame(&msg) {
-                            continue;
-                        }
-                        wait_key_frame = false;
-
-                        msg.header.msid = 1;
+                        // log::info!("[peer={}] frame type = {:#04X}, acv_packet_type={:#04X}", &ctx.peer_addr, msg.body[0], msg.body[1]);
                         msg.header.timestamp = (Local::now().timestamp_millis() - ctx.ctx_begin_timestamp) as u32;
                         let chunks = msg.split_chunks_bytes(ctx.chunk_size);
                         for chunk in chunks {
@@ -123,8 +126,14 @@ async fn connection_loop(stream: TcpStream) -> anyhow::Result<()> {
                         response_create_stream(&mut ctx, &values[1]).await?;
                     }
                     "publish" => {
+                        ctx.stream_name = values[3].try_as_str().unwrap_or_default().to_string();
+                        log::info!("[peer={}] stream_name={}", ctx.peer_addr, ctx.stream_name);
                         response_publish(&mut ctx).await?;
                     }
+                    "play" => {
+                        ctx.stream_name = values[3].try_as_str().unwrap_or_default().to_string();
+                        log::info!("[peer={}] stream_name={}", ctx.peer_addr, ctx.stream_name);
+                    }
                     _ => ()
                 }
             }
@@ -145,8 +154,13 @@ async fn connection_loop(stream: TcpStream) -> anyhow::Result<()> {
 
             ChunkMessageType::VideoMessage => {
                 log::debug!("[peer={}] C->S, [{}] header={:?}", ctx.peer_addr, message.message_type_desc(), message.header);
-                rtmp_msg_eventbus().publish(message.clone()).await;
-                handle_video_data(&message.body, &ctx);
+                if message.body[0] == 0x17 && message.body[1] == 0x00 {
+                    sps_message_map().insert(ctx.stream_name.clone(), message.clone());
+                    log::info!("[peer={}] C->S, cache sps/pps message, stream_name={}", ctx.peer_addr, ctx.stream_name);
+                } else {
+                    rtmp_msg_eventbus().publish(message.clone()).await;
+                }
+                // handle_video_data(&message.body, &ctx);
             }
             ChunkMessageType::AudioMessage => {
                 // sender.send(message).await?;
@@ -156,25 +170,6 @@ async fn connection_loop(stream: TcpStream) -> anyhow::Result<()> {
             }
         }
     }
-
-// log::info!("C->S, others:");
-// let mut i = 0;
-// let mut arr: [char; 8] = ['.'; 8];
-// loop {
-//     let byte = stream.read_one_return().await?;
-//     print!("{:#04X}", byte);
-//     if byte.is_ascii_graphic() {
-//         arr[i % 8] = byte as char;
-//     } else {
-//         arr[i % 8] = '.';
-//     }
-//     print!(" ");
-//     std::io::stdout().flush()?;
-//     i += 1;
-//     if i % 8 == 0 {
-//         println!("  {}", arr.iter().collect::<String>());
-//     }
-// }
 }
 
 /// 处理RTMP握手流程
@@ -356,7 +351,7 @@ async fn response_publish(ctx: &mut RtmpContext) -> anyhow::Result<()> {
 
 async fn response_play(ctx: &mut RtmpContext, stream_id: u32) -> anyhow::Result<()> {
     {
-        let mut rs: Vec<u8> = vec![
+        let rs: Vec<u8> = vec![
             0x02, 0x00, 0x00, 0x00,
             0x00, 0x00, 0x02, 0x04,
             0x00, 0x00, 0x00, 0x01,
@@ -503,6 +498,6 @@ async fn send_meta_data_for_play(ctx: &mut RtmpContext, meta_data: &RtmpMetaData
 
 fn main() -> anyhow::Result<()> {
     util::init_logger();
-    let server = "127.0.0.1:1935";
+    let server = "127.0.0.1:11935";
     smol::block_on(accept_loop(server))
 }

src/protocol/rtmp.rs

@@ -108,6 +108,7 @@ pub struct RtmpContext {
     pub recv_bytes_num: u32,
     pub arc_receiver: Arc<Receiver<Vec<u8>>>,
     pub peer_addr: String,
+    pub stream_name: String,
 }
 
 impl RtmpContext {
@@ -158,6 +159,7 @@ impl RtmpContext {
             recv_bytes_num: 0,
             arc_receiver: receiver,
             peer_addr,
+            stream_name: Default::default(),
         }
     }
 
@@ -190,11 +192,20 @@ pub struct RtmpMessageHeader {
 
 impl RtmpMessageHeader {
     pub fn to_bytes(&self) -> Vec<u8> {
+        let enable_extend_timestamp_field = self.timestamp >= 0xFFFFFF;
+
         let mut rs = vec![self.csid];
-        rs.write_u24::<BigEndian>(self.timestamp).unwrap();
+        if enable_extend_timestamp_field {
+            rs.write_u24::<BigEndian>(0xFFFFFF).unwrap();
+        } else {
+            rs.write_u24::<BigEndian>(self.timestamp).unwrap();
+        }
         rs.write_u24::<BigEndian>(self.message_length).unwrap();
         rs.write_u8(self.message_type_id).unwrap();
         rs.write_u32::<BigEndian>(self.msid).unwrap();
+        if enable_extend_timestamp_field {
+            rs.write_u32::<BigEndian>(self.timestamp).unwrap();
+        }
         rs
     }
 }
@@ -231,11 +242,17 @@ impl RtmpMessage {
                 // 时间差值置零
                 ctx.last_timestamp_delta = 0;
                 ctx.last_timestamp = BigEndian::read_u24(&h[0..3]);
+
                 ctx.last_message_length = BigEndian::read_u24(&h[3..6]);
                 ctx.remain_message_length = 0;
                 ctx.last_message_type_id = h[6];
                 ctx.last_message_stream_id = BigEndian::read_u32(&h[7..11]);
                 ctx.recv_bytes_num += 12;
+                if ctx.last_timestamp >= 0xFFFFFF {
+                    let extend = ctx.read_exact_from_peer(4).await?;
+                    ctx.last_timestamp = BigEndian::read_u32(&extend[0..4]);
+                    ctx.recv_bytes_num += 4;
+                }
 
                 (ctx.last_timestamp,
                  ctx.last_message_length,
@@ -510,6 +527,7 @@ pub fn read_all_amf_value(bytes: &[u8]) -> Option<Vec<Value>> {
 ///     0
 /// See ISO 14496-12, 8.15.3 for an explanation of composition
 /// times. The offset in an FLV file is always in milliseconds.
+#[allow(unused)]
 pub fn handle_video_data(bytes: &[u8], ctx: &RtmpContext) {
     let frame_type = bytes[0];
     let mut read_index = 1;
@@ -584,38 +602,13 @@ pub fn handle_video_data(bytes: &[u8], ctx: &RtmpContext) {
     } else {
         unreachable!("unknown acv packet type")
     };
-}
 
-fn handle_nalu(nalu_bytes: Vec<u8>) {
-    smol::block_on(nalu_eventbus().publish(nalu_bytes));
+    fn handle_nalu(nalu_bytes: Vec<u8>) {
+        smol::block_on(nalu_eventbus().publish(nalu_bytes));
+    }
 }
 
-fn nalu_type_desc(nalu_type: &u8) -> String {
-    let priority: String = match (nalu_type & 0x60) >> 5 {
-        0 => "DISPOSABLE".into(),
-        1 => "LOW".into(),
-        2 => "HIGH".into(),
-        3 => "HIGHEST".into(),
-        _ => "UNKNOWN".into(),
-    };
 
-    let t: String = match nalu_type & 0x1F {
-        1 => "SLICE".into(),
-        2 => "DPA".into(),
-        3 => "DPB".into(),
-        4 => "DPC".into(),
-        5 => "IDR".into(),
-        6 => "SEI".into(),
-        7 => "SPS".into(),
-        8 => "PPS".into(),
-        9 => "AUD".into(),
-        10 => "EOSEQ".into(),
-        11 => "EOSTREAM".into(),
-        12 => "FILL".into(),
-        _ => "UNKNOWN".into(),
-    };
 
-    format!("{}::{}", priority, t)
-}