From 6cc7eef2181d48b97632175d682bb3025d31a9ef Mon Sep 17 00:00:00 2001 From: Dave Kempe Date: Sun, 18 Oct 2026 10:41:24 +1100 Subject: [PATCH] websocket: terminating a session stops its traffic The proxy raced its two pump tasks against the cancel token with select!, which consumed the JoinHandles. Dropping a handle only detaches its task, so when terminate (or disabling a user) fired the token, or when one direction ended, the other task kept forwarding: traffic still reached the target and the guacd connection stayed open. Race through a helper that aborts both tasks afterwards, whichever way it ended. Tested with tasks that never finish on their own, for cancellation and for one side ending. --- src/websocket.rs | 136 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 115 insertions(+), 21 deletions(-) diff --git a/src/websocket.rs b/src/websocket.rs index a1bbc5e..7b9c854 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -448,27 +448,11 @@ async fn proxy_ws_guacd( let browser_to_guacd = tokio::spawn(async move { ws_to_guacd(ws_read, guacd_write, ws_sink_b).await }); - // Wait for either direction to finish, or cancellation - let result = tokio::select! { - result = guacd_to_browser => { - let err = match result { - Ok(Err(e)) => Some(e.to_string()), - Err(e) => Some(e.to_string()), - _ => None, - }; - ProxyResult::GuacdEnded(err) - } - result = browser_to_guacd => { - let err = match result { - Ok(Err(e)) => Some(e.to_string()), - Err(e) => Some(e.to_string()), - _ => None, - }; - ProxyResult::BrowserEnded(err) - } - _ = cancel.cancelled() => { - ProxyResult::Cancelled - } + // Wait for either direction to finish, or cancellation, then stop both. + let result = match race_and_abort(guacd_to_browser, browser_to_guacd, &cancel).await { + Race::First(result) => ProxyResult::GuacdEnded(task_error(result)), + Race::Second(result) => ProxyResult::BrowserEnded(task_error(result)), + Race::Cancelled => ProxyResult::Cancelled, }; ProxyOutcome { @@ -477,6 +461,45 @@ async fn proxy_ws_guacd( } } +/// Which side of `race_and_abort` finished first. +enum Race { + First(A), + Second(B), + Cancelled, +} + +/// Wait for the first of two tasks to finish, or for `cancel`, then abort +/// both. Dropping a `JoinHandle` only detaches its task, so without the +/// abort a terminated session (or one whose browser went away) kept both +/// pump tasks running: traffic still flowed to and from the target and the +/// guacd connection stayed open until the other side happened to close. +/// Aborting a task that has already finished does nothing. +async fn race_and_abort( + mut first: tokio::task::JoinHandle, + mut second: tokio::task::JoinHandle, + cancel: &CancellationToken, +) -> Race, Result> { + let outcome = tokio::select! { + r = &mut first => Race::First(r), + r = &mut second => Race::Second(r), + _ = cancel.cancelled() => Race::Cancelled, + }; + first.abort(); + second.abort(); + outcome +} + +/// The error text of a pump task that ended, if it ended badly. +fn task_error( + result: Result, tokio::task::JoinError>, +) -> Option { + match result { + Ok(Err(e)) => Some(e.to_string()), + Err(e) => Some(e.to_string()), + _ => None, + } +} + type WsSink = Arc>>; /// Maximum bytes the guacd-side carry buffer is allowed to grow to before @@ -647,6 +670,77 @@ pub(crate) fn origin_host_matches(origin: &str, host: &str) -> bool { mod tests { use super::*; + /// A task that runs until aborted; `alive` turns false when it is dropped. + fn forever( + alive: std::sync::Arc, + ) -> tokio::task::JoinHandle> { + struct Flag(std::sync::Arc); + impl Drop for Flag { + fn drop(&mut self) { + self.0.store(false, std::sync::atomic::Ordering::SeqCst); + } + } + // Captured by the future, so dropping the task (polled or not) + // drops the flag. + let flag = Flag(alive); + tokio::spawn(async move { + let _flag = flag; + std::future::pending::<()>().await; + Ok(()) + }) + } + + async fn gone(alive: &std::sync::atomic::AtomicBool) -> bool { + for _ in 0..100 { + if !alive.load(std::sync::atomic::Ordering::SeqCst) { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + false + } + + #[tokio::test] + async fn terminating_a_session_stops_both_pump_tasks() { + // Terminate fires the cancel token. Before, the select dropped the + // two task handles, which only detaches them: traffic kept flowing + // to the target after the session was "terminated". + use std::sync::atomic::AtomicBool; + use std::sync::Arc; + let (a, b) = ( + Arc::new(AtomicBool::new(true)), + Arc::new(AtomicBool::new(true)), + ); + let cancel = CancellationToken::new(); + let race = tokio::spawn({ + let (a, b, cancel) = (a.clone(), b.clone(), cancel.clone()); + async move { + matches!( + race_and_abort(forever(a), forever(b), &cancel).await, + Race::Cancelled + ) + } + }); + cancel.cancel(); + assert!(race.await.unwrap(), "cancellation wins the race"); + assert!(gone(&a).await, "guacd to browser task stopped"); + assert!(gone(&b).await, "browser to guacd task stopped"); + } + + #[tokio::test] + async fn one_direction_ending_stops_the_other() { + // The browser goes away: its task ends, and the guacd side must not be + // left running with the connection open. + use std::sync::atomic::AtomicBool; + use std::sync::Arc; + let other = Arc::new(AtomicBool::new(true)); + let ended: tokio::task::JoinHandle> = tokio::spawn(async { Ok(()) }); + let cancel = CancellationToken::new(); + let r = race_and_abort(forever(other.clone()), ended, &cancel).await; + assert!(matches!(r, Race::Second(Ok(Ok(()))))); + assert!(gone(&other).await, "the still-running side was stopped"); + } + fn user(email: &str, role: &str) -> AuthIdentity { AuthIdentity::User { email: email.into(),