From a2c74a6e81f836047e26833e536c0695021a62ed Mon Sep 17 00:00:00 2001 From: Doloro1978 Date: Sun, 19 Jul 2026 01:02:40 +0100 Subject: [PATCH] fix(rtmp): replace unwrap() calls with proper error handling and add handshake timeout --- crates/server/src/rtmp.rs | 96 ++++++++++++++++++++++++++++++--------- 1 file changed, 74 insertions(+), 22 deletions(-) diff --git a/crates/server/src/rtmp.rs b/crates/server/src/rtmp.rs index 87529b0..81f9ee3 100644 --- a/crates/server/src/rtmp.rs +++ b/crates/server/src/rtmp.rs @@ -12,7 +12,7 @@ use sea_orm::{DatabaseConnection, IntoActiveModel, sqlx::types::chrono::Local}; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}, - time::{Instant, timeout}, + time::timeout, }; use tracing::{debug, info, warn}; @@ -24,6 +24,8 @@ use crate::{ }, }; +const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); + pub struct Rtmp { pub listener: TcpListener, pub stream_sessions: Arc>, @@ -33,7 +35,10 @@ pub struct Rtmp { async fn write_outbound(socket: &mut TcpStream, results: Vec) { for r in results { if let ServerSessionResult::OutboundResponse(p) = r { - socket.write_all(&p.bytes).await.unwrap(); + if let Err(e) = socket.write_all(&p.bytes).await { + warn!("RTMP write error: {e}"); + return; + } } } } @@ -69,7 +74,10 @@ impl Rtmp { let mut server = Handshake::new(PeerType::Server); let mut c0_c1 = [0u8; 1537]; - socket.read_exact(&mut c0_c1).await?; + let n = timeout(HANDSHAKE_TIMEOUT, socket.read_exact(&mut c0_c1)) + .await + .map_err(|_| "handshake C0+C1 timeout")??; + debug!("handshake C0+C1 read {} bytes", n); let s0_s1_s2 = match server.process_bytes(&c0_c1) { Ok(HandshakeProcessResult::InProgress { response_bytes }) => response_bytes, x => return Err(format!("unexpected handshake state: {:?}", x).into()), @@ -77,7 +85,9 @@ impl Rtmp { socket.write_all(&s0_s1_s2).await?; let mut c2 = [0u8; 1536]; - socket.read_exact(&mut c2).await?; + timeout(HANDSHAKE_TIMEOUT, socket.read_exact(&mut c2)) + .await + .map_err(|_| "handshake C2 timeout")??; match server.process_bytes(&c2) { Ok(HandshakeProcessResult::Completed { .. }) => {} Ok(HandshakeProcessResult::InProgress { response_bytes }) => { @@ -86,7 +96,8 @@ impl Rtmp { x => return Err(format!("unexpected handshake state: {:?}", x).into()), } - let (rtmp_session, init_bytes) = ServerSession::new(ServerSessionConfig::new()).unwrap(); + let (rtmp_session, init_bytes) = + ServerSession::new(ServerSessionConfig::new()).map_err(|e| format!("ServerSession::new failed: {e}"))?; write_outbound(&mut socket, init_bytes).await; Ok((rtmp_session, socket)) @@ -193,12 +204,21 @@ impl Rtmp { for event in events { match event { ServerSessionResult::OutboundResponse(p) => { - socket.write_all(&p.bytes).await.unwrap(); + if let Err(e) = socket.write_all(&p.bytes).await { + warn!("RTMP outbound write error: {e}"); + return; + } } ServerSessionResult::RaisedEvent(e) => match e { ServerSessionEvent::ConnectionRequested { request_id, .. } => { debug!("RTMP ConnectionRequested, accepting"); - let reply = session.accept_request(request_id).unwrap(); + let reply = match session.accept_request(request_id) { + Ok(r) => r, + Err(e) => { + warn!("Failed to accept connection request {request_id}: {e}"); + return; + } + }; write_outbound(&mut socket, reply).await; } ServerSessionEvent::PublishStreamRequested { @@ -215,28 +235,50 @@ impl Rtmp { key } else { warn!(stream_key = %stream_key, "stream key not found, rejecting"); - let reply = session - .reject_request(request_id, "", "Stream key invalid") - .unwrap(); + let reply = match session.reject_request(request_id, "", "Stream key invalid") { + Ok(r) => r, + Err(e) => { + warn!("Failed to reject stream key request: {e}"); + return; + } + }; write_outbound(&mut socket, reply).await; break; }; let already_live = - stream_session::Model::get_active_by_stream_key_id( + match stream_session::Model::get_active_by_stream_key_id( &db, key.id, ) .await - .unwrap(); + { + Ok(s) => s, + Err(e) => { + warn!("DB error checking active stream: {e}"); + let reply = match session.reject_request(request_id, "", "Internal error") { + Ok(r) => r, + Err(e2) => { + warn!("Failed to reject request on DB error: {e2}"); + return; + } + }; + write_outbound(&mut socket, reply).await; + return; + } + }; if already_live.is_some() { warn!(stream_key_id = key.id, label = %key.label, "stream key already live, rejecting duplicate publish"); - let reply = session - .reject_request( - request_id, - "", - "You're already streaming...", - ) - .unwrap(); + let reply = match session.reject_request( + request_id, + "", + "You're already streaming...", + ) { + Ok(r) => r, + Err(e) => { + warn!("Failed to reject duplicate publish: {e}"); + return; + } + }; write_outbound(&mut socket, reply).await; break; } @@ -255,14 +297,24 @@ impl Rtmp { }, ); - let reply = session.accept_request(request_id).unwrap(); - stream_session::Model::create_stream_session( + let reply = match session.accept_request(request_id) { + Ok(r) => r, + Err(e) => { + warn!("Failed to accept publish request: {e}"); + return; + } + }; + if let Err(e) = stream_session::Model::create_stream_session( &db, key.id, Local::now().into(), ) .await - .unwrap(); + { + warn!("Failed to create stream session record: {e}"); + let _ = session.reject_request(request_id, "", "Internal error"); + return; + } stream_id = Some(key.id); write_outbound(&mut socket, reply).await; }