383 lines
18 KiB
Rust
383 lines
18 KiB
Rust
use std::{error::Error, sync::Arc, time::Duration};
|
|
|
|
use async_broadcast::broadcast;
|
|
use chrono::Utc;
|
|
use dashmap::DashMap;
|
|
use entity::stream_session;
|
|
use rml_rtmp::{
|
|
handshake::{Handshake, HandshakeProcessResult, PeerType},
|
|
sessions::{ServerSession, ServerSessionConfig, ServerSessionEvent, ServerSessionResult},
|
|
};
|
|
use sea_orm::{DatabaseConnection, IntoActiveModel, sqlx::types::chrono::Local};
|
|
use tokio::{
|
|
io::{AsyncReadExt, AsyncWriteExt},
|
|
net::{TcpListener, TcpStream},
|
|
time::{Instant, timeout},
|
|
};
|
|
use tracing::{debug, info, warn};
|
|
|
|
use crate::{
|
|
StreamCodec, StreamSession,
|
|
audio::{AACParser, AudioProcesser, OpusAudioFrame},
|
|
codec::{
|
|
CodecParser, VideoFrame, av1::Av1CodecParser, h264::H264CodecParser, h265::H265CodecParser,
|
|
},
|
|
};
|
|
|
|
pub struct Rtmp {
|
|
pub listener: TcpListener,
|
|
pub stream_sessions: Arc<DashMap<i32, StreamSession>>,
|
|
pub db: DatabaseConnection,
|
|
}
|
|
|
|
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();
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Rtmp {
|
|
fn parse_video_codec(payload: &[u8]) -> Result<StreamCodec, Box<dyn Error + Send + Sync>> {
|
|
if payload.is_empty() {
|
|
return Err("empty video payload".into());
|
|
}
|
|
|
|
let is_ex = payload[0] & 0x80 != 0;
|
|
|
|
if is_ex {
|
|
if payload.len() < 5 {
|
|
return Err("enhanced RTMP payload too short".into());
|
|
}
|
|
match &payload[1..5] {
|
|
b"hvc1" => Ok(StreamCodec::H265),
|
|
b"avc1" => Ok(StreamCodec::H264),
|
|
b"av01" => Ok(StreamCodec::AV1),
|
|
_ => Err("unsupported FourCC".into()),
|
|
}
|
|
} else {
|
|
match payload[0] & 0x0F {
|
|
7 => Ok(StreamCodec::H264),
|
|
_ => Err("unsupported legacy codec ID".into()),
|
|
}
|
|
}
|
|
}
|
|
async fn handshake(
|
|
mut socket: TcpStream,
|
|
) -> Result<(ServerSession, TcpStream), Box<dyn Error>> {
|
|
let mut server = Handshake::new(PeerType::Server);
|
|
|
|
let mut c0_c1 = [0u8; 1537];
|
|
socket.read_exact(&mut c0_c1).await?;
|
|
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()),
|
|
};
|
|
socket.write_all(&s0_s1_s2).await?;
|
|
|
|
let mut c2 = [0u8; 1536];
|
|
socket.read_exact(&mut c2).await?;
|
|
match server.process_bytes(&c2) {
|
|
Ok(HandshakeProcessResult::Completed { .. }) => {}
|
|
Ok(HandshakeProcessResult::InProgress { response_bytes }) => {
|
|
socket.write_all(&response_bytes).await?
|
|
}
|
|
x => return Err(format!("unexpected handshake state: {:?}", x).into()),
|
|
}
|
|
|
|
let (rtmp_session, init_bytes) = ServerSession::new(ServerSessionConfig::new()).unwrap();
|
|
write_outbound(&mut socket, init_bytes).await;
|
|
|
|
Ok((rtmp_session, socket))
|
|
}
|
|
|
|
pub async fn run(self) {
|
|
let Self {
|
|
listener,
|
|
stream_sessions,
|
|
db,
|
|
} = self;
|
|
loop {
|
|
let connection = listener.accept().await;
|
|
let (socket, peer_addr) = match connection {
|
|
Ok(e) => e,
|
|
Err(e) => {
|
|
tracing::error!("Error on connection accept: {}", e);
|
|
continue;
|
|
}
|
|
};
|
|
info!(%peer_addr, "RTMP connection accepted");
|
|
let (mut session, mut socket) = match Rtmp::handshake(socket).await {
|
|
Ok(v) => v,
|
|
Err(e) => {
|
|
warn!(%peer_addr, "RTMP handshake failed: {e}");
|
|
continue;
|
|
}
|
|
};
|
|
info!(%peer_addr, "RTMP handshake complete");
|
|
let (mut video_tx, mut video_rx) = broadcast::<Arc<VideoFrame>>(32);
|
|
let (mut audio_tx, mut audio_rx) = broadcast::<Arc<OpusAudioFrame>>(32);
|
|
// video_rx.cycle
|
|
|
|
video_tx.set_overflow(true);
|
|
audio_tx.set_overflow(true);
|
|
let mut parser: Option<Box<dyn CodecParser>> = None;
|
|
let mut stream_id: Option<i32> = None;
|
|
let mut aac_parser = AACParser::new();
|
|
let mut audio_proc = AudioProcesser::new();
|
|
let db = db.clone();
|
|
let stream_sessions = stream_sessions.clone();
|
|
|
|
tokio::spawn(async move {
|
|
let mut current_stream_key_id: Option<i32> = None;
|
|
let mut codec_stamped = false;
|
|
|
|
let cleanup = |id: i32| {
|
|
let stream_sessions = stream_sessions.clone();
|
|
let db = db.clone();
|
|
async move {
|
|
stream_sessions.remove(&id);
|
|
if let Ok(Some(s)) =
|
|
stream_session::Model::get_active_by_stream_key_id(&db, id).await
|
|
{
|
|
s.into_active_model()
|
|
.finish_stream_session(&db, Local::now().into())
|
|
.await
|
|
.ok();
|
|
}
|
|
}
|
|
};
|
|
|
|
loop {
|
|
let mut buf = [0u8; 4096];
|
|
const IDLE_TIMEOUT: Duration = Duration::from_secs(15);
|
|
let n = match timeout(IDLE_TIMEOUT, socket.read(&mut buf)).await {
|
|
Ok(Ok(0)) => {
|
|
debug!("RTMP connection closed by peer");
|
|
if let Some(id) = current_stream_key_id {
|
|
cleanup(id).await;
|
|
}
|
|
return;
|
|
}
|
|
Ok(Ok(n)) => n,
|
|
Ok(Err(e)) => {
|
|
warn!("RTMP read error: {e}");
|
|
if let Some(id) = current_stream_key_id {
|
|
cleanup(id).await;
|
|
}
|
|
return;
|
|
}
|
|
Err(_) => {
|
|
warn!("RTMP timed out (15 sec no packets) id: {:?}", stream_id);
|
|
if let Some(id) = current_stream_key_id {
|
|
cleanup(id).await;
|
|
}
|
|
return;
|
|
}
|
|
};
|
|
let events = match session.handle_input(&buf[..n]) {
|
|
Ok(e) => e,
|
|
Err(e) => {
|
|
warn!("RTMP session error: {e}");
|
|
if let Some(id) = current_stream_key_id {
|
|
cleanup(id).await;
|
|
}
|
|
return;
|
|
}
|
|
};
|
|
// Blankly using it, so it doesnt drop
|
|
video_rx.is_closed();
|
|
audio_rx.is_closed();
|
|
|
|
for event in events {
|
|
match event {
|
|
ServerSessionResult::OutboundResponse(p) => {
|
|
socket.write_all(&p.bytes).await.unwrap();
|
|
}
|
|
ServerSessionResult::RaisedEvent(e) => match e {
|
|
ServerSessionEvent::ConnectionRequested { request_id, .. } => {
|
|
debug!("RTMP ConnectionRequested, accepting");
|
|
let reply = session.accept_request(request_id).unwrap();
|
|
write_outbound(&mut socket, reply).await;
|
|
}
|
|
ServerSessionEvent::PublishStreamRequested {
|
|
request_id,
|
|
stream_key,
|
|
..
|
|
} => {
|
|
info!(stream_key = %stream_key, "publish stream requested");
|
|
let key =
|
|
entity::stream_key::Entity::find_by_key(&db, &stream_key)
|
|
.await;
|
|
|
|
let key = if let Ok(Some(key)) = key {
|
|
key
|
|
} else {
|
|
warn!(stream_key = %stream_key, "stream key not found, rejecting");
|
|
let reply = session
|
|
.reject_request(request_id, "", "Stream key invalid")
|
|
.unwrap();
|
|
write_outbound(&mut socket, reply).await;
|
|
break;
|
|
};
|
|
|
|
let already_live =
|
|
stream_session::Model::get_active_by_stream_key_id(
|
|
&db, key.id,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
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();
|
|
write_outbound(&mut socket, reply).await;
|
|
break;
|
|
}
|
|
|
|
info!(stream_key_id = key.id, label = %key.label, "stream started");
|
|
current_stream_key_id = Some(key.id);
|
|
stream_sessions.insert(
|
|
key.id,
|
|
StreamSession {
|
|
stream_key_id: key.id,
|
|
stream_key_label: key.label,
|
|
frame_channel: video_tx.clone(),
|
|
audio_channel: audio_tx.clone(),
|
|
codec: None,
|
|
started_at: Utc::now(),
|
|
},
|
|
);
|
|
|
|
let reply = session.accept_request(request_id).unwrap();
|
|
stream_session::Model::create_stream_session(
|
|
&db,
|
|
key.id,
|
|
Local::now().into(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
stream_id = Some(key.id);
|
|
write_outbound(&mut socket, reply).await;
|
|
}
|
|
ServerSessionEvent::PublishStreamFinished {
|
|
stream_key, ..
|
|
} => {
|
|
info!(stream_key = %stream_key, "publish stream finished");
|
|
match entity::stream_key::Entity::find_by_key(&db, &stream_key)
|
|
.await
|
|
{
|
|
Ok(Some(key)) => cleanup(key.id).await,
|
|
Ok(None) => {
|
|
warn!(stream_key = %stream_key, "publish finished for unknown stream key");
|
|
}
|
|
Err(e) => {
|
|
warn!(stream_key = %stream_key, "failed to look up stream key on finish: {e}");
|
|
}
|
|
}
|
|
current_stream_key_id = None;
|
|
}
|
|
// TODO: We can totally replace the broadcast with a
|
|
// circular_buff
|
|
// Arc<Vec<ArcSwap<Frame>>>
|
|
ServerSessionEvent::VideoDataReceived {
|
|
data, timestamp, ..
|
|
} => match Self::parse_video_codec(&data) {
|
|
Ok(codec) => {
|
|
if !codec_stamped {
|
|
if let Some(id) = current_stream_key_id {
|
|
if let Some(mut session) =
|
|
stream_sessions.get_mut(&id)
|
|
{
|
|
session.codec = Some(codec.clone());
|
|
}
|
|
}
|
|
codec_stamped = true;
|
|
}
|
|
match codec {
|
|
StreamCodec::H264 => {
|
|
let p = parser.get_or_insert_with(|| {
|
|
Box::new(H264CodecParser::new())
|
|
});
|
|
if let Some(frame) = p.parse(&data, timestamp.value)
|
|
{
|
|
video_tx.broadcast(Arc::new(frame)).await.ok();
|
|
}
|
|
}
|
|
StreamCodec::H265 => {
|
|
let p = parser.get_or_insert_with(|| {
|
|
Box::new(H265CodecParser::new())
|
|
});
|
|
if let Some(frame) = p.parse(&data, timestamp.value)
|
|
{
|
|
video_tx.broadcast(Arc::new(frame)).await.ok();
|
|
}
|
|
}
|
|
StreamCodec::AV1 => {
|
|
let p = parser.get_or_insert_with(|| {
|
|
Box::new(Av1CodecParser::new())
|
|
});
|
|
let pkt_type = data[0] & 0x0F;
|
|
match p.parse(&data, timestamp.value) {
|
|
Some(frame) => {
|
|
debug!(
|
|
pkt_type,
|
|
is_keyframe = frame.is_keyframe,
|
|
ts = frame.timestamp_ms,
|
|
bytes = frame.data.len(),
|
|
"AV1 frame → broadcast"
|
|
);
|
|
video_tx
|
|
.broadcast(Arc::new(frame))
|
|
.await
|
|
.ok();
|
|
}
|
|
None => {
|
|
debug!(
|
|
pkt_type,
|
|
"AV1 packet produced no frame (seq header or unknown type)"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Err(_err) => {
|
|
warn!("");
|
|
}
|
|
},
|
|
ServerSessionEvent::AudioDataReceived {
|
|
data, timestamp, ..
|
|
} => {
|
|
// Consume the non-Send error before any await point.
|
|
let opus_frames: Vec<_> =
|
|
match aac_parser.parse(&data, timestamp.value) {
|
|
Err(e) => {
|
|
warn!("AAC parse error: {}", e);
|
|
vec![]
|
|
}
|
|
Ok(frame) => frame
|
|
.map(|f| audio_proc.encode(f))
|
|
.unwrap_or_default(),
|
|
};
|
|
for frame in opus_frames {
|
|
audio_tx.broadcast(Arc::new(frame)).await.ok();
|
|
}
|
|
}
|
|
_ => {}
|
|
},
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|