diff --git a/crates/entity/src/stream_key.rs b/crates/entity/src/stream_key.rs index 9c86899..beb1d33 100644 --- a/crates/entity/src/stream_key.rs +++ b/crates/entity/src/stream_key.rs @@ -12,6 +12,8 @@ pub struct Model { pub label: String, pub is_active: bool, pub is_unlisted: bool, + pub password: Option, + pub custom_id: Option, pub created_at: DateTimeUtc, } @@ -78,6 +80,16 @@ impl Entity { .all(db) .await } + + pub async fn find_by_custom_id( + db: &DatabaseConnection, + custom_id: String, + ) -> Result, DbErr> { + Entity::find() + .filter(Column::CustomId.eq(custom_id)) + .one(db) + .await + } } impl ActiveModel { @@ -89,4 +101,17 @@ impl ActiveModel { self.label = Set(value); self.update(db).await } + + pub async fn password( + mut self, + db: &DatabaseConnection, + value: String, + ) -> Result { + self.password = if value.is_empty() { + Set(None) + } else { + Set(Some(value)) + }; + self.update(db).await + } } diff --git a/crates/migration/src/lib.rs b/crates/migration/src/lib.rs index 454239f..ccc3161 100644 --- a/crates/migration/src/lib.rs +++ b/crates/migration/src/lib.rs @@ -4,6 +4,8 @@ mod m20260616_000001_create_users; mod m20260616_000002_create_stream_key; mod m20260616_000003_create_stream_session; mod m20260616_000004_create_auth_session; +mod m20260815_000005_add_password_to_stream_key; +mod m20260815_000006_set_stream_keys_unlisted; pub struct Migrator; @@ -15,6 +17,8 @@ impl MigratorTrait for Migrator { Box::new(m20260616_000002_create_stream_key::Migration), Box::new(m20260616_000003_create_stream_session::Migration), Box::new(m20260616_000004_create_auth_session::Migration), + Box::new(m20260815_000005_add_password_to_stream_key::Migration), + Box::new(m20260815_000006_set_stream_keys_unlisted::Migration), ] } } diff --git a/crates/migration/src/m20260616_000002_create_stream_key.rs b/crates/migration/src/m20260616_000002_create_stream_key.rs index b101f8f..684af7b 100644 --- a/crates/migration/src/m20260616_000002_create_stream_key.rs +++ b/crates/migration/src/m20260616_000002_create_stream_key.rs @@ -38,7 +38,7 @@ impl MigrationTrait for Migration { ColumnDef::new(StreamKey::IsUnlisted) .boolean() .not_null() - .default(true), + .default(false), ) .col(ColumnDef::new(StreamKey::CreatedAt).date_time().not_null()) .foreign_key( @@ -67,5 +67,7 @@ pub enum StreamKey { Label, IsActive, IsUnlisted, + Password, + CustomId, CreatedAt, } diff --git a/crates/migration/src/m20260815_000005_add_password_to_stream_key.rs b/crates/migration/src/m20260815_000005_add_password_to_stream_key.rs new file mode 100644 index 0000000..3893b28 --- /dev/null +++ b/crates/migration/src/m20260815_000005_add_password_to_stream_key.rs @@ -0,0 +1,47 @@ +use sea_orm_migration::prelude::*; + +use super::m20260616_000002_create_stream_key::StreamKey; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .alter_table( + Table::alter() + .table(StreamKey::Table) + .add_column(ColumnDef::new(StreamKey::Password).string().null()) + .to_owned(), + ) + .await?; + manager + .alter_table( + Table::alter() + .table(StreamKey::Table) + .add_column(ColumnDef::new(StreamKey::CustomId).string().null()) + .to_owned(), + ) + .await + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .alter_table( + Table::alter() + .table(StreamKey::Table) + .drop_column(StreamKey::Password) + .to_owned(), + ) + .await?; + manager + .alter_table( + Table::alter() + .table(StreamKey::Table) + .drop_column(StreamKey::CustomId) + .to_owned(), + ) + .await + } +} diff --git a/crates/migration/src/m20260815_000006_set_stream_keys_unlisted.rs b/crates/migration/src/m20260815_000006_set_stream_keys_unlisted.rs new file mode 100644 index 0000000..24a0e64 --- /dev/null +++ b/crates/migration/src/m20260815_000006_set_stream_keys_unlisted.rs @@ -0,0 +1,26 @@ +use sea_orm_migration::prelude::*; + +use super::m20260616_000002_create_stream_key::StreamKey; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + // By accident the is_unlisted column defaulted to true; existing stream + // keys were meant to be listed. Reset all rows to false here. + manager + .exec_stmt( + Query::update() + .table(StreamKey::Table) + .values([(StreamKey::IsUnlisted, false.into())]) + .to_owned(), + ) + .await + } + + async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> { + Ok(()) + } +} diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml index 8956638..0c7076d 100644 --- a/crates/server/Cargo.toml +++ b/crates/server/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "server" -version = "0.4.0" +version = "0.5.0" edition = "2024" [target.x86_64-unknown-linux-gnu] diff --git a/crates/server/src/http.rs b/crates/server/src/http.rs index 4d63fb0..c1c0242 100644 --- a/crates/server/src/http.rs +++ b/crates/server/src/http.rs @@ -22,7 +22,7 @@ use axum::{ Json, Router, extract::{FromRequestParts, Path, State}, http::{ - HeaderName, Method, StatusCode, + HeaderMap, HeaderName, Method, StatusCode, header::{ACCEPT, AUTHORIZATION, CONTENT_TYPE}, request::Parts, }, @@ -31,7 +31,7 @@ use axum::{ }; use chrono::{DateTime, Utc}; use entity::{auth_session, stream_key, stream_session, users}; -use sea_orm::{DatabaseConnection, EntityTrait, IntoActiveModel}; +use sea_orm::{ActiveModelTrait, DatabaseConnection, EntityTrait, IntoActiveModel, Set}; use serde::{Deserialize, Serialize}; use sysinfo::System; use tokio::{ @@ -149,7 +149,14 @@ struct StreamListing { #[derive(Deserialize)] struct EditStreamKeyRequest { id: i32, - new: String, + #[serde(default)] + new: Option, + #[serde(default)] + password: Option, + #[serde(default)] + unlisted: Option, + #[serde(default)] + custom_id: Option, } async fn edit_stream_key( @@ -158,14 +165,16 @@ async fn edit_stream_key( Json(payload): Json, ) -> Result { // Trim surrounding whitespace so labels aren't stored with leading/trailing spaces. - let new_label = payload.new.trim(); + let new_label = payload.new.as_deref().map(str::trim); // Length (1..=67) and allowed charset ([A-Za-z0-9 _-], space included). - if !KEY_RE.is_match(new_label) { - return Err(HttpError::BadRequest("invalid label".into())); - } - // Replaces the JS lookahead: label must contain at least one letter. - if !new_label.chars().any(|c| c.is_ascii_alphabetic()) { - return Err(HttpError::BadRequest("label must contain a letter".into())); + if let Some(label) = new_label { + if !KEY_RE.is_match(label) { + return Err(HttpError::BadRequest("invalid label".into())); + } + // Replaces the JS lookahead: label must contain at least one letter. + if !label.chars().any(|c| c.is_ascii_alphabetic()) { + return Err(HttpError::BadRequest("label must contain a letter".into())); + } } let stream_key = stream_key::Entity::find_by_id(payload.id) @@ -176,10 +185,60 @@ async fn edit_stream_key( if stream_key.user_id != auth.0.id { return Err(HttpError::Forbidden); } - stream_key - .into_active_model() - .change_label_value(&state.db, new_label.to_string()) - .await?; + + let lock = state.appstate.lock().await; + let mut ses = { lock.stream_sessions.get_mut(&payload.id) }; + // drop(lock); + + let mut am = stream_key.into_active_model(); + if let Some(custom_id) = payload.custom_id { + am.custom_id = if custom_id.is_empty() { + if let Some(ref mut ses) = ses { + ses.custom_id = None; + } + Set(None) + } else { + // Check for conflict. + if stream_key::Entity::find_by_custom_id(&state.db, custom_id.clone()) + .await? + .is_some() + { + return Err(HttpError::Conflict); + } + if let Some(ref mut ses) = ses { + ses.custom_id = Some(custom_id.clone()); + } + Set(Some(custom_id)) + }; + } + if let Some(label) = new_label { + if let Some(ref mut ses) = ses { + ses.stream_key_label = label.to_string(); + } + am.label = Set(label.to_string()); + } + if let Some(pwd) = payload.password { + am.password = if pwd.is_empty() { + if let Some(ref mut ses) = ses { + ses.password = None; + } + Set(None) + } else { + if let Some(ref mut ses) = ses { + ses.password = Some(pwd.clone()); + } + Set(Some(pwd)) + }; + } + if let Some(unlisted) = payload.unlisted { + if let Some(ref mut ses) = ses { + ses.is_unlisted = unlisted; + } + am.is_unlisted = Set(unlisted); + } + // This hopefully will not fail, if it does, our values for the stream key will be mismatched, + // that would be bad + am.update(&state.db).await?; Ok(StatusCode::OK) } @@ -195,6 +254,7 @@ async fn catalog_handler( .await .stream_sessions .iter() + .filter(|x| !x.is_unlisted) .map(|x| StreamListing { id: x.stream_key_id, label: x.stream_key_label.clone(), @@ -412,16 +472,18 @@ async fn create_user_handler( async fn stream_handler( State(state): State>, Path(slug): Path, + headers: HeaderMap, body: String, ) -> Result { let request_id_clone = state.request_count.fetch_add(1, Ordering::Relaxed); + let app = state.appstate.lock().await; let stream_key_id = { - let app = state.appstate.lock().await; - app.stream_sessions - .iter() - .find(|e| e.value().stream_key_id.to_string() == slug) - .map(|e| *e.key()) + app.stream_sessions.iter().find(|e| { + e.value().stream_key_id.to_string() == slug + || e.value().custom_id.as_deref() == Some(slug.as_str()) + }) + // .map(|e| *e.key()) }; let stream_key_id = stream_key_id.ok_or_else(|| { @@ -429,12 +491,29 @@ async fn stream_handler( HttpError::NotFound })?; - info!(request_id = request_id_clone, slug = %slug, stream_key_id, "WHEP offer received"); + // Dont return stream by its id in the db if its unlisted + if stream_key_id.is_unlisted && stream_key_id.custom_id.clone().ok_or("") != Ok(slug.clone()) { + return Err(HttpError::NotFound); + } + + let auth_header = headers.get("auth"); + + if let Some(password) = &stream_key_id.password { + if let Some(auth) = auth_header { + if auth.to_str().unwrap() != password { + return Err(HttpError::Unauthorized); + } + } else { + return Err(HttpError::Unauthorized); + }; + }; + + info!(request_id = request_id_clone, slug = %slug, stream_key_id.stream_key_id, "WHEP offer received"); let accept_rx = state.accept_rx.activate_cloned(); // The webrtc worker owning the receiver died if this fails. state .offer_tx - .send((request_id_clone, stream_key_id, body)) + .send((request_id_clone, stream_key_id.stream_key_id, body)) .await .map_err(|_| HttpError::Internal)?; debug!( diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 10e4621..1d58acf 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -50,6 +50,9 @@ pub struct StreamSession { pub stream_key_id: i32, pub stream_key_label: String, pub stream_key_user: String, + pub custom_id: Option, + pub is_unlisted: bool, + pub password: Option, pub frame_channel: async_broadcast::Sender>, pub audio_channel: async_broadcast::Sender>, pub codec: Option, @@ -98,6 +101,7 @@ async fn main() -> Result<(), Box> { webrtc_proxy: proxy.clone(), })); let (offer_tx, offer_rx) = tokio::sync::mpsc::channel::<(i32, i32, String)>(64); + // Request_Id, // String_Label, // Offer_body diff --git a/crates/server/src/rtmp.rs b/crates/server/src/rtmp.rs index d3d695c..1b4bf32 100644 --- a/crates/server/src/rtmp.rs +++ b/crates/server/src/rtmp.rs @@ -312,6 +312,9 @@ impl Rtmp { stream_key_id: key.id, stream_key_label: key.label, stream_key_user: user.username, + custom_id: key.custom_id, + is_unlisted: key.is_unlisted, + password: key.password, frame_channel: video_tx.clone(), audio_channel: audio_tx.clone(), codec: None, diff --git a/crates/server/src/webrtc.rs b/crates/server/src/webrtc.rs index ecef992..4ec944e 100644 --- a/crates/server/src/webrtc.rs +++ b/crates/server/src/webrtc.rs @@ -4,7 +4,7 @@ use entity::stream_key; use sea_orm::{DatabaseConnection, EntityTrait}; use std::{net::SocketAddr, sync::Arc, time::Instant}; use tokio::{net::UdpSocket, sync::mpsc::Receiver}; -use tracing::{debug, error, info, warn}; +use tracing::{debug, error, info, trace, warn}; use str0m::{ Candidate, Event, Input, Output, Rtc, @@ -40,7 +40,10 @@ impl Webrtc { request_id, stream_id, "stream key not found in DB, rejecting offer" ); - self.accept_tx.broadcast((request_id, Err(String::new()))).await.unwrap(); + self.accept_tx + .broadcast((request_id, Err(String::new()))) + .await + .unwrap(); continue; } if let Err(ref e) = stream_key { @@ -48,7 +51,10 @@ impl Webrtc { request_id, stream_id, "DB error looking up stream key: {:?}", e ); - self.accept_tx.broadcast((request_id, Err(String::new()))).await.unwrap(); + self.accept_tx + .broadcast((request_id, Err(String::new()))) + .await + .unwrap(); continue; } @@ -134,7 +140,10 @@ impl Webrtc { Ok(sdp) => sdp, Err(e) => { warn!(request_id, stream_id, "malformed SDP offer: {:?}", e); - self.accept_tx.broadcast((request_id, Err(String::new()))).await.unwrap(); + self.accept_tx + .broadcast((request_id, Err(String::new()))) + .await + .unwrap(); continue; } }; @@ -162,7 +171,7 @@ impl Webrtc { } }; let answer_sdp = offer_answer.to_sdp_string(); - info!(request_id, "SDP answer:\n{}", answer_sdp); + trace!(request_id, "SDP answer:\n{}", answer_sdp); // Detect the case where str0m couldn't match any video codec. // str0m serialises the m-line with an empty PT list, which is invalid SDP @@ -179,14 +188,14 @@ impl Webrtc { "no video codec negotiated — browser likely doesn't support {:?}; rejecting offer", stream_codec ); - self.accept_tx - .broadcast((request_id, Err(format!("{:?}", stream_codec)))) - .await - .unwrap(); - continue; - } + self.accept_tx + .broadcast((request_id, Err(format!("{:?}", stream_codec)))) + .await + .unwrap(); + continue; + } - let ufrag = answer_sdp + let ufrag = answer_sdp .lines() .find(|l| l.starts_with("a=ice-ufrag:")) .and_then(|l| l.strip_prefix("a=ice-ufrag:")) diff --git a/crates/server/src/webrtc_ingest.rs b/crates/server/src/webrtc_ingest.rs index ba5cc6e..a73cc3d 100644 --- a/crates/server/src/webrtc_ingest.rs +++ b/crates/server/src/webrtc_ingest.rs @@ -1,4 +1,8 @@ -use std::{net::SocketAddr, sync::Arc, time::{Duration, Instant}}; +use std::{ + net::SocketAddr, + sync::Arc, + time::{Duration, Instant}, +}; use async_broadcast::broadcast; use axum::{ @@ -298,6 +302,9 @@ pub async fn handle_whip_injest( stream_key_id: key.id, stream_key_label: key.label, stream_key_user: user.username, + custom_id: key.custom_id, + is_unlisted: key.is_unlisted, + password: key.password, frame_channel: video_tx, audio_channel: audio_tx, codec: negotiated_codec, @@ -372,6 +379,7 @@ async fn detach_inject_rtc( let sessions_ref = sessions_ref.clone(); let db = db.clone(); async move { + debug!("Cleaning up {:?}", &stream_key_id); sessions_ref.remove(&stream_key_id); if let Ok(Some(s)) = stream_session::Model::get_active_by_stream_key_id(&db, stream_key_id).await @@ -388,6 +396,7 @@ async fn detach_inject_rtc( let mut video_tx: Option>> = None; let mut audio_tx: Option>> = None; let mut _connected = false; + let mut disconnect_timer: Option = None; loop { // Blankly using them so they don't drop (like RTMP). @@ -398,6 +407,7 @@ async fn detach_inject_rtc( match rtc.poll_output() { Ok(Output::Timeout(t)) => break t, Ok(Output::Transmit(t)) => { + // The keep alive loop, by default is every 1 second. trace!( "Whip TX: {} bytes → {}:{}", t.contents.len(), @@ -409,6 +419,22 @@ async fn detach_inject_rtc( cleanup().await; return; } + // 30 Sec time out if disconnected, then we clean up and disconnect. + if let Some(instant_since_disconnet) = disconnect_timer + && Instant::now() + .duration_since(instant_since_disconnet) + .as_secs() + > NO_MEDIA_TIMEOUT.as_secs() + { + info!( + "WHIP connection on stream_key_id: {:?}, has been disconnected for {:?} seconds. Cleaning up and destroying the connection,", + stream_key_id, + NO_MEDIA_TIMEOUT.as_secs() + ); + cleanup().await; + rtc.disconnect(); + return; + } } Ok(Output::Event(e)) => match e { Event::MediaAdded(ma) => { @@ -448,10 +474,19 @@ async fn detach_inject_rtc( } Event::IceConnectionStateChange(state) => { info!(stream_key_id, ?state, "Whip ICE state change"); - if matches!(state, str0m::IceConnectionState::Disconnected) { - info!("Whip ICE disconnected, closing connection"); - cleanup().await; - return; + match state { + str0m::IceConnectionState::Disconnected => { + info!( + "Whip ICE disconnected... (State changed to Disconnected for {:?}) ((This is usually due to network jitter))", + &stream_key_id + ); + disconnect_timer = Some(Instant::now()); + } + str0m::IceConnectionState::Connected + | str0m::IceConnectionState::Completed => { + disconnect_timer = None; + } + _ => {} } } Event::Connected => { diff --git a/flake.lock b/flake.lock index e269dc5..2f4738c 100644 --- a/flake.lock +++ b/flake.lock @@ -53,7 +53,28 @@ "inputs": { "crane": "crane", "flake-utils": "flake-utils", - "nixpkgs": "nixpkgs" + "nixpkgs": "nixpkgs", + "rust-overlay": "rust-overlay" + } + }, + "rust-overlay": { + "inputs": { + "nixpkgs": [ + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1786160316, + "narHash": "sha256-oLoc3ZLg1LX/S5Jb3v6MrF415AzDyC4vgeWy9UcYTQk=", + "owner": "oxalica", + "repo": "rust-overlay", + "rev": "4e1c940c96560ceab7c547f89642231371a66646", + "type": "github" + }, + "original": { + "owner": "oxalica", + "repo": "rust-overlay", + "type": "github" } }, "systems": { diff --git a/flake.nix b/flake.nix index edbf5f6..7160d2d 100644 --- a/flake.nix +++ b/flake.nix @@ -7,6 +7,12 @@ crane.url = "github:ipetkov/crane"; flake-utils.url = "github:numtide/flake-utils"; + + rust-overlay = { + url = "github:oxalica/rust-overlay"; + inputs.nixpkgs.follows = "nixpkgs"; + }; + }; outputs = @@ -16,14 +22,27 @@ crane, flake-utils, ... - }: + }@inputs: flake-utils.lib.eachDefaultSystem ( system: let - pkgs = nixpkgs.legacyPackages.${system}; + pkgs = import inputs.nixpkgs { + inherit system; + overlays = [ (import inputs.rust-overlay) ]; + }; + inherit (pkgs) lib; - craneLib = crane.mkLib pkgs; + craneLib = (inputs.crane.mkLib pkgs).overrideToolchain ( + p: + p.rust-bin.nightly.latest.default.override { + extensions = [ + "rustc-codegen-cranelift-preview" + "rust-analyzer" + "rust-src" + ]; + } + ); # Common arguments can be set here to avoid repeating them later # Note: changes here will rebuild all dependency crates @@ -56,7 +75,7 @@ commonArgs // { pname = "rtmp-to-whip"; - version = "0.1.0"; + version = "0.4.0"; cargoArtifacts = craneLib.buildDepsOnly commonArgs; cargoExtraArgs = "-p server"; src = fileSetForCrate ./crates/server;