add: h265 & AV1 support (with a lot of fixes)
This commit is contained in:
+218
-151
@@ -15,9 +15,14 @@ use tokio::{
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::{
|
||||
StreamSession,
|
||||
StreamCodec, StreamSession,
|
||||
audio::{AACParser, AudioProcesser, OpusAudioFrame},
|
||||
media::{H264Parser, VideoFrame},
|
||||
codec::{
|
||||
CodecParser, VideoFrame,
|
||||
av1::Av1CodecParser,
|
||||
h264::H264CodecParser,
|
||||
h265::H265CodecParser,
|
||||
},
|
||||
};
|
||||
|
||||
pub struct Rtmp {
|
||||
@@ -35,6 +40,30 @@ async fn write_outbound(socket: &mut TcpStream, results: Vec<ServerSessionResult
|
||||
}
|
||||
|
||||
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>> {
|
||||
@@ -64,139 +93,129 @@ impl Rtmp {
|
||||
Ok((rtmp_session, socket))
|
||||
}
|
||||
|
||||
pub fn start(self) -> Result<(), Box<dyn Error>> {
|
||||
pub async fn run(self) {
|
||||
let Self {
|
||||
listener,
|
||||
stream_sessions,
|
||||
db,
|
||||
} = self;
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (socket, peer_addr) = listener.accept().await.unwrap();
|
||||
info!(%peer_addr, "RTMP connection accepted");
|
||||
let (mut session, mut socket) = Rtmp::handshake(socket).await.unwrap();
|
||||
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
|
||||
loop {
|
||||
let (socket, peer_addr) = listener.accept().await.unwrap();
|
||||
info!(%peer_addr, "RTMP connection accepted");
|
||||
let (mut session, mut socket) = Rtmp::handshake(socket).await.unwrap();
|
||||
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 = H264Parser::new();
|
||||
let mut aac_parser = AACParser::new();
|
||||
let mut audio_proc = AudioProcesser::new();
|
||||
let db = db.clone();
|
||||
let stream_sessions = stream_sessions.clone();
|
||||
video_tx.set_overflow(true);
|
||||
audio_tx.set_overflow(true);
|
||||
let mut parser: Option<Box<dyn CodecParser>> = 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 {
|
||||
loop {
|
||||
let mut buf = [0u8; 4096];
|
||||
let n = socket.read(&mut buf).await.unwrap();
|
||||
if n == 0 {
|
||||
debug!("RTMP connection closed by peer");
|
||||
return;
|
||||
}
|
||||
let events = session.handle_input(&buf[..n]).unwrap();
|
||||
// Blankly using it, so it doesnt drop
|
||||
video_rx.is_closed();
|
||||
audio_rx.is_closed();
|
||||
tokio::spawn(async move {
|
||||
let mut current_stream_key_id: Option<i32> = None;
|
||||
let mut codec_stamped = false;
|
||||
loop {
|
||||
let mut buf = [0u8; 4096];
|
||||
let n = socket.read(&mut buf).await.unwrap();
|
||||
if n == 0 {
|
||||
debug!("RTMP connection closed by peer");
|
||||
return;
|
||||
}
|
||||
let events = session.handle_input(&buf[..n]).unwrap();
|
||||
// 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();
|
||||
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;
|
||||
}
|
||||
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;
|
||||
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
|
||||
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();
|
||||
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");
|
||||
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(),
|
||||
},
|
||||
);
|
||||
|
||||
let reply = session.accept_request(request_id).unwrap();
|
||||
stream_session::Model::create_stream_session(
|
||||
&db,
|
||||
key.id,
|
||||
Local::now().into(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
write_outbound(&mut socket, reply).await;
|
||||
}
|
||||
ServerSessionEvent::PublishStreamFinished {
|
||||
stream_key,
|
||||
..
|
||||
} => {
|
||||
info!(stream_key = %stream_key, "publish stream finished");
|
||||
let key = entity::stream_key::Entity::find_by_key(
|
||||
&db,
|
||||
&stream_key,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
stream_sessions.remove(&key.id);
|
||||
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,
|
||||
},
|
||||
);
|
||||
|
||||
let reply = session.accept_request(request_id).unwrap();
|
||||
stream_session::Model::create_stream_session(
|
||||
&db,
|
||||
key.id,
|
||||
Local::now().into(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
write_outbound(&mut socket, reply).await;
|
||||
}
|
||||
ServerSessionEvent::PublishStreamFinished {
|
||||
stream_key, ..
|
||||
} => {
|
||||
info!(stream_key = %stream_key, "publish stream finished");
|
||||
let key =
|
||||
entity::stream_key::Entity::find_by_key(&db, &stream_key)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
stream_sessions.remove(&key.id);
|
||||
stream_session::Model::get_active_by_stream_key_id(&db, key.id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
@@ -204,44 +223,92 @@ impl Rtmp {
|
||||
.finish_stream_session(&db, Local::now().into())
|
||||
.await
|
||||
.unwrap();
|
||||
stream_sessions.get(&key.id).unwrap().frame_channel.close();
|
||||
}
|
||||
ServerSessionEvent::VideoDataReceived {
|
||||
data,
|
||||
timestamp,
|
||||
..
|
||||
} => {
|
||||
if data.len() >= 5 && &data[1..5] == b"hvc1" {
|
||||
warn!("HEVC/H.265 not supported, closing connection");
|
||||
return;
|
||||
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;
|
||||
}
|
||||
if let Some(frame) = parser.parse(&data, timestamp.value) {
|
||||
video_tx.broadcast(Arc::new(frame)).await.ok();
|
||||
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)");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
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();
|
||||
}
|
||||
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();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user