doloro/simple-instant-stream
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(rtmp): replace unwrap() calls with proper error handling and add handshake timeout

Browse Source
This commit is contained in:
Doloro1978
2026-07-19 01:02:40 +01:00
parent ad997d3b65
commit a2c74a6e81
1 changed files with 74 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/server/src/rtmp.rs
+74 -22
View File
@@ -12,7 +12,7 @@ use sea_orm::{DatabaseConnection, IntoActiveModel, sqlx::types::chrono::Local};
use tokio::{ use tokio::{
io::{AsyncReadExt, AsyncWriteExt}, io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream}, net::{TcpListener, TcpStream},
time::{Instant, timeout}, time::timeout,
}; };
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
@@ -24,6 +24,8 @@ use crate::{
}, },
}; };
const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
pub struct Rtmp { pub struct Rtmp {
pub listener: TcpListener, pub listener: TcpListener,
pub stream_sessions: Arc<DashMap<i32, StreamSession>>, pub stream_sessions: Arc<DashMap<i32, StreamSession>>,
@@ -33,7 +35,10 @@ pub struct Rtmp {
async fn write_outbound(socket: &mut TcpStream, results: Vec<ServerSessionResult>) { async fn write_outbound(socket: &mut TcpStream, results: Vec<ServerSessionResult>) {
for r in results { for r in results {
if let ServerSessionResult::OutboundResponse(p) = r { 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 server = Handshake::new(PeerType::Server);
let mut c0_c1 = [0u8; 1537]; 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) { let s0_s1_s2 = match server.process_bytes(&c0_c1) {
Ok(HandshakeProcessResult::InProgress { response_bytes }) => response_bytes, Ok(HandshakeProcessResult::InProgress { response_bytes }) => response_bytes,
x => return Err(format!("unexpected handshake state: {:?}", x).into()), x => return Err(format!("unexpected handshake state: {:?}", x).into()),
@@ -77,7 +85,9 @@ impl Rtmp {
socket.write_all(&s0_s1_s2).await?; socket.write_all(&s0_s1_s2).await?;
let mut c2 = [0u8; 1536]; 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) { match server.process_bytes(&c2) {
Ok(HandshakeProcessResult::Completed { .. }) => {} Ok(HandshakeProcessResult::Completed { .. }) => {}
Ok(HandshakeProcessResult::InProgress { response_bytes }) => { Ok(HandshakeProcessResult::InProgress { response_bytes }) => {
@@ -86,7 +96,8 @@ impl Rtmp {
x => return Err(format!("unexpected handshake state: {:?}", x).into()), 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; write_outbound(&mut socket, init_bytes).await;
Ok((rtmp_session, socket)) Ok((rtmp_session, socket))
@@ -193,12 +204,21 @@ impl Rtmp {
for event in events { for event in events {
match event { match event {
ServerSessionResult::OutboundResponse(p) => { 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 { ServerSessionResult::RaisedEvent(e) => match e {
ServerSessionEvent::ConnectionRequested { request_id, .. } => { ServerSessionEvent::ConnectionRequested { request_id, .. } => {
debug!("RTMP ConnectionRequested, accepting"); 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; write_outbound(&mut socket, reply).await;
} }
ServerSessionEvent::PublishStreamRequested { ServerSessionEvent::PublishStreamRequested {
@@ -215,28 +235,50 @@ impl Rtmp {
key key
} else { } else {
warn!(stream_key = %stream_key, "stream key not found, rejecting"); warn!(stream_key = %stream_key, "stream key not found, rejecting");
let reply = session let reply = match session.reject_request(request_id, "", "Stream key invalid") {
.reject_request(request_id, "", "Stream key invalid") Ok(r) => r,
.unwrap(); Err(e) => {
warn!("Failed to reject stream key request: {e}");
return;
}
};
write_outbound(&mut socket, reply).await; write_outbound(&mut socket, reply).await;
break; break;
}; };
let already_live = 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, &db, key.id,
) )
.await .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() { if already_live.is_some() {
warn!(stream_key_id = key.id, label = %key.label, "stream key already live, rejecting duplicate publish"); warn!(stream_key_id = key.id, label = %key.label, "stream key already live, rejecting duplicate publish");
let reply = session let reply = match session.reject_request(
.reject_request( request_id,
request_id, "",
"", "You're already streaming...",
"You're already streaming...", ) {
) Ok(r) => r,
.unwrap(); Err(e) => {
warn!("Failed to reject duplicate publish: {e}");
return;
}
};
write_outbound(&mut socket, reply).await; write_outbound(&mut socket, reply).await;
break; break;
} }
@@ -255,14 +297,24 @@ impl Rtmp {
}, },
); );
let reply = session.accept_request(request_id).unwrap(); let reply = match session.accept_request(request_id) {
stream_session::Model::create_stream_session( Ok(r) => r,
Err(e) => {
warn!("Failed to accept publish request: {e}");
return;
}
};
if let Err(e) = stream_session::Model::create_stream_session(
&db, &db,
key.id, key.id,
Local::now().into(), Local::now().into(),
) )
.await .await
.unwrap(); {
warn!("Failed to create stream session record: {e}");
let _ = session.reject_request(request_id, "", "Internal error");
return;
}
stream_id = Some(key.id); stream_id = Some(key.id);
write_outbound(&mut socket, reply).await; write_outbound(&mut socket, reply).await;
} }