a lot 2
This commit is contained in:
+225
-39
@@ -1,52 +1,238 @@
|
||||
#[derive(Default, Debug)]
|
||||
pub struct RtmpSession {
|
||||
client: RtmpClient,
|
||||
use std::{error::Error, sync::Arc};
|
||||
|
||||
use async_broadcast::broadcast;
|
||||
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},
|
||||
};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::{
|
||||
StreamSession,
|
||||
audio::{AACParser, AudioProcesser, OpusAudioFrame},
|
||||
media::{H264Parser, VideoFrame},
|
||||
};
|
||||
|
||||
pub struct Rtmp {
|
||||
pub listener: TcpListener,
|
||||
pub stream_sessions: Arc<DashMap<i32, StreamSession>>,
|
||||
pub db: DatabaseConnection,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct RtmpClient {
|
||||
version: u8,
|
||||
timestamp: u32,
|
||||
magic_bytes: [u8; 1536],
|
||||
}
|
||||
|
||||
impl Default for RtmpClient {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
magic_bytes: [0; 1536],
|
||||
version: 0,
|
||||
timestamp: 0,
|
||||
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 RtmpSession {
|
||||
pub fn consume_handshake(&mut self, bytes: &[u8]) -> Result<(), ()> {
|
||||
let version = bytes[0];
|
||||
let timestamp = u32::from_be_bytes(bytes[1..5].try_into().unwrap());
|
||||
let mut random: [u8; 1536] = [0; 1536];
|
||||
random.copy_from_slice(&bytes[1..1537]);
|
||||
impl Rtmp {
|
||||
async fn handshake(
|
||||
mut socket: TcpStream,
|
||||
) -> Result<(ServerSession, TcpStream), Box<dyn Error>> {
|
||||
let mut server = Handshake::new(PeerType::Server);
|
||||
|
||||
println!("size of rand: {}", random.len());
|
||||
|
||||
let client = RtmpClient {
|
||||
version,
|
||||
timestamp,
|
||||
magic_bytes: random,
|
||||
let mut c0_c1 = [0u8; 1537];
|
||||
socket.read_exact(&mut c0_c1).await.unwrap();
|
||||
let s0_s1_s2 = match server.process_bytes(&c0_c1) {
|
||||
Ok(HandshakeProcessResult::InProgress { response_bytes }) => response_bytes,
|
||||
_ => panic!("handshake failed"),
|
||||
};
|
||||
socket.write_all(&s0_s1_s2).await.unwrap();
|
||||
|
||||
self.client = client;
|
||||
let mut c2 = [0u8; 1536];
|
||||
socket.read_exact(&mut c2).await.unwrap();
|
||||
match server.process_bytes(&c2) {
|
||||
Ok(HandshakeProcessResult::Completed { .. }) => {}
|
||||
Ok(HandshakeProcessResult::InProgress { response_bytes }) => {
|
||||
socket.write_all(&response_bytes).await.unwrap()
|
||||
}
|
||||
x => panic!("Unexpected process_bytes response: {:?}", x),
|
||||
}
|
||||
|
||||
let (rtmp_session, init_bytes) = ServerSession::new(ServerSessionConfig::new()).unwrap();
|
||||
write_outbound(&mut socket, init_bytes).await;
|
||||
|
||||
Ok((rtmp_session, socket))
|
||||
}
|
||||
|
||||
pub fn start(self) -> Result<(), Box<dyn Error>> {
|
||||
let Self {
|
||||
listener,
|
||||
stream_sessions,
|
||||
db,
|
||||
} = self;
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let (mut session, mut socket) = Rtmp::handshake(socket).await.unwrap();
|
||||
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();
|
||||
|
||||
tokio::spawn(async move {
|
||||
// let video_rx = video_rx;
|
||||
loop {
|
||||
let mut buf = [0u8; 4096];
|
||||
let n = socket.read(&mut buf).await.unwrap();
|
||||
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();
|
||||
}
|
||||
ServerSessionResult::RaisedEvent(e) => match e {
|
||||
ServerSessionEvent::ConnectionRequested {
|
||||
request_id, ..
|
||||
} => {
|
||||
let reply = session.accept_request(request_id).unwrap();
|
||||
write_outbound(&mut socket, reply).await;
|
||||
}
|
||||
ServerSessionEvent::PublishStreamRequested {
|
||||
request_id,
|
||||
stream_key,
|
||||
..
|
||||
} => {
|
||||
let key = entity::stream_key::Entity::find_by_key(
|
||||
&db,
|
||||
&stream_key,
|
||||
)
|
||||
.await;
|
||||
|
||||
let key = if let Ok(Some(key)) = key {
|
||||
key
|
||||
} else {
|
||||
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() {
|
||||
let reply = session
|
||||
.reject_request(
|
||||
request_id,
|
||||
"",
|
||||
"You're already streaming...",
|
||||
)
|
||||
.unwrap();
|
||||
write_outbound(&mut socket, reply).await;
|
||||
break;
|
||||
}
|
||||
|
||||
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,
|
||||
..
|
||||
} => {
|
||||
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()
|
||||
.into_active_model()
|
||||
.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;
|
||||
}
|
||||
if let Some(frame) = parser.parse(&data, timestamp.value) {
|
||||
video_tx.broadcast(Arc::new(frame)).await.ok();
|
||||
}
|
||||
}
|
||||
ServerSessionEvent::AudioDataReceived {
|
||||
data,
|
||||
timestamp,
|
||||
..
|
||||
} => {
|
||||
// Consume the non-Send error before any await point.
|
||||
let opus_frames: Vec<_> = aac_parser
|
||||
.parse(&data, timestamp.value)
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|f| audio_proc.encode(f))
|
||||
.unwrap_or_default();
|
||||
for frame in opus_frames {
|
||||
audio_tx.broadcast(Arc::new(frame)).await.ok();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
pub fn response(&self) -> [u8; 1537] {
|
||||
let mut reply: [u8; 1537] = [0; 1537];
|
||||
reply[0] = 3;
|
||||
// reply[1..1537].copy_from_slice(&self.client.magic_bytes);
|
||||
let mut rand: [u8; 1536] = [0; 1536];
|
||||
rand.fill(1);
|
||||
reply[1..1537].copy_from_slice(&rand);
|
||||
|
||||
reply
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user