Bump : 0.5.4 : WHIP cleanup guarded by session id, stale detach task can't kill a newer stream on the same key

Cargo.lock

@@ -2906,7 +2906,7 @@ dependencies = [
 
 [[package]]
 name = "server"
-version = "0.5.2"
+version = "0.5.4"
 dependencies = [
  "argon2",
  "async-broadcast",

crates/server/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "server"
-version = "0.5.3"
+version = "0.5.4"
 edition = "2024"
 
 [target.x86_64-unknown-linux-gnu]

crates/server/src/main.rs

@@ -13,6 +13,7 @@ use migration::{Migrator, MigratorTrait};
 use sea_orm::Database;
 use tokio::{net::TcpListener, sync::Mutex, task::JoinSet};
 use tracing_subscriber::EnvFilter;
+use uuid::Uuid;
 
 use crate::{
     audio::OpusAudioFrame,
@@ -64,6 +65,9 @@ pub struct StreamSession {
     /// profile instead of all three H.264 profiles.
     pub video_profile_level_id: Option<u32>,
     // ----
+    /// Unique id per publish. Lets a session's cleanup verify it is still the
+    /// live session for its key instead of clobbering a newer one.
+    pub session_id: Uuid,
     pub started_at: DateTime<Utc>,
     pub active_clients: AtomicU32,
 }

crates/server/src/rtmp.rs

@@ -15,6 +15,7 @@ use tokio::{
     time::timeout,
 };
 use tracing::{debug, info, trace, warn};
+use uuid::Uuid;
 
 use crate::{
     StreamCodec, StreamSession,
@@ -320,6 +321,7 @@ impl Rtmp {
                                             codec: None,
                                             started_at: Utc::now(),
                                             active_clients: 0.into(),
+                                            session_id: Uuid::new_v4(),
                                             video_pt: None,
                                             video_profile_level_id: None,
                                         },

crates/server/src/webrtc_ingest.rs

@@ -23,6 +23,7 @@ use str0m::{
 };
 use tokio::{net::UdpSocket, sync::mpsc::Receiver};
 use tracing::{debug, error, info, trace, warn};
+use uuid::Uuid;
 
 use crate::{
     StreamSession, audio::OpusAudioFrame, codec::VideoFrame, http::HttpServer,
@@ -65,6 +66,9 @@ pub async fn handle_whip_injest_delete(
     }
 
     state.appstate.lock().await.stream_sessions.remove(&key.id);
+    // Drop the trickle channel too, so PATCHes to a deleted session 404 and
+    // the lingering detach task's own cleanup can't reach a newer session.
+    state.appstate.lock().await.webrtc_proxy.trickle_tx.remove(&key.id);
 
     Ok(StatusCode::OK)
 }
@@ -86,7 +90,10 @@ pub async fn handle_whip_injest_patch(
 
     // Look up the trickle sender for this session.
     let trickle_map = state.appstate.lock().await.webrtc_proxy.trickle_tx.clone();
-    let tx = trickle_map.get(&stream_key_id).ok_or_else(|| {
+    let tx = trickle_map
+        .get(&stream_key_id)
+        .map(|e| e.value().1.clone())
+        .ok_or_else(|| {
             warn!(stream_key_id, "Whip PATCH: no trickle channel for session");
             HttpError::NotFound
         })?;
@@ -283,6 +290,9 @@ pub async fn handle_whip_injest(
     let (socket, rx) = state.appstate.lock().await.webrtc_proxy.add_client(ufrag);
 
     let stream_sessions = state.appstate.lock().await.stream_sessions.clone();
+    // Unique per publish: cleanup only removes the session this task created,
+    // never a newer one that re-published on the same stream key.
+    let session_id = Uuid::new_v4();
     let (mut video_tx, video_rx) = broadcast::<Arc<VideoFrame>>(32);
     let (mut audio_tx, audio_rx) = broadcast::<Arc<OpusAudioFrame>>(32);
     // Never block the ingest loop on slow/missing viewers: overwrite the
@@ -308,6 +318,7 @@ pub async fn handle_whip_injest(
             frame_channel: video_tx,
             audio_channel: audio_tx,
             codec: negotiated_codec,
+            session_id,
             started_at: Utc::now(),
             active_clients: 0.into(),
             video_pt,
@@ -323,7 +334,7 @@ pub async fn handle_whip_injest(
     // initial offer. We forward them to the Rtc task for add_remote_candidate.
     let (trickle_tx, trickle_rx) = tokio::sync::mpsc::unbounded_channel();
     let trickle_map = state.appstate.lock().await.webrtc_proxy.trickle_tx.clone();
-    trickle_map.insert(key.id, trickle_tx);
+    trickle_map.insert(key.id, (session_id, trickle_tx));
 
     let db = state.db.clone();
     tokio::spawn(async move {
@@ -338,10 +349,12 @@ pub async fn handle_whip_injest(
             public_addr,
             video_rx,
             audio_rx,
+            session_id,
         )
         .await;
-        // Clean up trickle channel when done.
+        // Drop only our own trickle channel; a newer session on the same key
-        trickle_map.remove(&key.id);
+        // must keep its own.
+        trickle_map.remove_if(&key.id, |_, (sid, _)| *sid == session_id);
     });
 
     stream_session::Model::create_stream_session(&state.db, key.id, Local::now().into())
@@ -374,11 +387,25 @@ async fn detach_inject_rtc(
     local_addr: SocketAddr,
     video_rx: async_broadcast::Receiver<Arc<VideoFrame>>,
     audio_rx: async_broadcast::Receiver<Arc<OpusAudioFrame>>,
+    session_id: Uuid,
 ) {
     let cleanup = || {
         let sessions_ref = sessions_ref.clone();
         let db = db.clone();
         async move {
+            // Only clean up the session this task created. If the user has
+            // already opened another stream on the same key (e.g. DELETE then
+            // immediate republish), that newer session must not be touched.
+            let is_ours = sessions_ref
+                .get(&stream_key_id)
+                .is_some_and(|s| s.session_id == session_id);
+            if !is_ours {
+                debug!(
+                    stream_key_id,
+                    "skipping cleanup: session replaced by a newer stream on this key"
+                );
+                return;
+            }
             debug!("Cleaning up {:?}", &stream_key_id);
             sessions_ref.remove(&stream_key_id);
             if let Ok(Some(s)) =

crates/server/src/webrtc_proxy.rs

@@ -12,6 +12,7 @@ use tokio::{
     sync::mpsc::{self, Receiver},
 };
 use tracing::{debug, error, info, trace, warn};
+use uuid::Uuid;
 
 pub struct WebRtcProxyConfig {
     pub proxy_port: i32,
@@ -24,9 +25,10 @@ pub struct WebrtcProxy {
     socket: Arc<UdpSocket>,
     public_addr: SocketAddr,
     /// Trickle-ICE candidate channels for WHIP ingest.
-    /// Keyed by stream_key_id; sender stored here so PATCH handler
+    /// Keyed by stream_key_id; each entry is tagged with the owning session's
-    /// can forward candidates to the detach task.
+    /// id so cleanup can remove only its own entry and never a newer session
-    pub trickle_tx: Arc<DashMap<i32, tokio::sync::mpsc::UnboundedSender<String>>>,
+    /// that re-published on the same key.
+    pub trickle_tx: Arc<DashMap<i32, (Uuid, tokio::sync::mpsc::UnboundedSender<String>)>>,
 }
 
 const STUN_MAGIC: u32 = 0x2112A442;