fix(rtmp): replace unwrap() calls with proper error handling and add handshake timeout
This commit is contained in:
+74
-22
@@ -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<DashMap<i32, StreamSession>>,
|
||||
@@ -33,7 +35,10 @@ pub struct Rtmp {
|
||||
async fn write_outbound(socket: &mut TcpStream, results: Vec<ServerSessionResult>) {
|
||||
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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user