From 25d1bcce43ab2a1dea72528c115220b30fe31329 Mon Sep 17 00:00:00 2001 From: rustdesk Date: Sat, 5 Sep 2026 10:13:49 +0800 Subject: [PATCH] port forward: closing the tunnel reaches channels parked on their socket A channel whose far end neither reads nor writes has both relays parked on the socket, not on the inbound queue, so `close_all` dropping the queue's sender woke neither: the socket and both tasks lived on until the far end hung up. Both sides now hold a per-tunnel teardown signal that `run_channel` selects on beside its own cancel, and `close_all` sends it after clearing the map. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab --- src/port_forward_mux.rs | 40 +++++++++++++++++++++++++++++++--- src/server/port_forward_mux.rs | 14 +++++++++--- 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/port_forward_mux.rs b/src/port_forward_mux.rs index 1f5319943..5fee3d781 100644 --- a/src/port_forward_mux.rs +++ b/src/port_forward_mux.rs @@ -346,6 +346,9 @@ async fn relay_tunnel_to_socket( /// Runs both halves as independent tasks; whichever ends first cancels the /// other. Sends `close` once, after the last data, and only when the channel /// ended for a local reason — the peer's own `close` is never echoed. +/// `teardown` is the tunnel closing under the channel: it cancels both halves +/// even when they are parked on the socket, where dropping the inbound sender +/// reaches neither. pub async fn run_channel( id: i32, reader: R, @@ -356,6 +359,7 @@ pub async fn run_channel( window: Arc>, inbound: mpsc::UnboundedReceiver, sink: FrameSink, + mut teardown: watch::Receiver<()>, ) where R: AsyncRead + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static, @@ -376,6 +380,10 @@ pub async fn run_channel( let _ = cancel_tx.send(true); (r.unwrap_or(RelayEnd::Cancelled), to_tunnel.await.unwrap_or(RelayEnd::Cancelled)) } + _ = teardown.changed() => { + let _ = cancel_tx.send(true); + (to_tunnel.await.unwrap_or(RelayEnd::Cancelled), to_socket.await.unwrap_or(RelayEnd::Cancelled)) + } }; let peer_closed = first == RelayEnd::PeerClosed || second == RelayEnd::PeerClosed; let tunnel_gone = first == RelayEnd::TunnelGone || second == RelayEnd::TunnelGone; @@ -463,6 +471,7 @@ mod tunnel { channels: Mutex::new(HashMap::new()), next_id: AtomicI32::new(1), reported: Default::default(), + teardown: watch::channel(()).0, }); let state = self.state.clone(); // Publish before spawning: if the loop exits first and resets the @@ -502,6 +511,9 @@ mod tunnel { channels: Mutex>, next_id: AtomicI32, reported: Mutex>, + /// Sent once, by `close_all`, for the channels its `clear` cannot + /// reach: one parked on its local socket is not on the inbound queue. + teardown: watch::Sender<()>, } impl TunnelHandle { @@ -541,10 +553,11 @@ mod tunnel { let open = open_msg(id, host, port, CHANNEL_WINDOW); let (reader, writer) = socket.into_split(); let sink = self.sink.clone(); + let teardown = self.teardown.subscribe(); let handle = self.clone(); tokio::spawn(async move { if sink.send_ordered(open).await.is_ok() { - run_channel(id, reader, writer, prebuf, Vec::new(), credit, window, inbound_rx, sink).await; + run_channel(id, reader, writer, prebuf, Vec::new(), credit, window, inbound_rx, sink, teardown).await; } handle.channels.lock().unwrap().remove(&id); }); @@ -643,6 +656,7 @@ mod tunnel { fn close_all(&self) { self.channels.lock().unwrap().clear(); + self.teardown.send(()).ok(); } #[cfg(test)] @@ -883,6 +897,7 @@ mod tests { credit: Arc, window: Arc>, local: tokio::io::DuplexStream, + teardown: watch::Sender<()>, task: tokio::task::JoinHandle<()>, } @@ -897,10 +912,11 @@ mod tests { let credit = Arc::new(SendCredit::new(INITIAL_WINDOW)); let window = Arc::new(Mutex::new(RecvWindow::new(CHANNEL_WINDOW))); let sink = FrameSink::Queued { data: data_tx, control: control_tx }; + let (teardown, teardown_rx) = watch::channel(()); let task = tokio::spawn(run_channel( - id, r, w, prebuf, initial_out, credit.clone(), window.clone(), inbound_rx, sink, + id, r, w, prebuf, initial_out, credit.clone(), window.clone(), inbound_rx, sink, teardown_rx, )); - Harness { data_rx, control_rx, inbound_tx, credit, window, local, task } + Harness { data_rx, control_rx, inbound_tx, credit, window, local, teardown, task } } // `PortForwardData.data` is generated as `bytes::Bytes` (hbb_common builds @@ -926,6 +942,24 @@ mod tests { } } + #[test] + fn teardown_ends_a_channel_parked_on_a_socket_nobody_reads() { + rt().block_on(async { + // Nothing reads the local side, so once the duplex buffer is full + // the socket relay parks in write_all; nothing writes it either, + // so the tunnel relay parks in read. Neither is on the inbound + // queue, which stays open here: teardown alone must end them. + let mut h = harness(1, vec![], vec![]); + for _ in 0..17 { + h.inbound_tx.send(Inbound::Data(Bytes::from(vec![0u8; MAX_FRAME]))).unwrap(); + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + h.teardown.send(()).ok(); + let ended = tokio::time::timeout(std::time::Duration::from_millis(500), &mut h.task).await; + assert!(ended.is_ok(), "channel task outlived the tunnel"); + }); + } + #[test] fn local_bytes_become_data_frames_capped_at_max_frame() { rt().block_on(async { diff --git a/src/server/port_forward_mux.rs b/src/server/port_forward_mux.rs index d6f22e2d8..312644187 100644 --- a/src/server/port_forward_mux.rs +++ b/src/server/port_forward_mux.rs @@ -8,7 +8,7 @@ use hbb_common::{ log, message_proto::*, timeout, - tokio::{self, net::TcpStream, sync::mpsc}, + tokio::{self, net::TcpStream, sync::{mpsc, watch}}, }; use std::{ collections::HashMap, @@ -35,6 +35,9 @@ pub struct PortForwardMux { channels: HashMap, tx: Sender, login_target: String, + /// Sent once, by `close_all`, for the channels its `clear` cannot reach: + /// one parked on its target socket is not on the inbound queue. + teardown: watch::Sender<()>, } impl PortForwardMux { @@ -43,6 +46,7 @@ impl PortForwardMux { channels: HashMap::new(), tx, login_target, + teardown: watch::channel(()).0, } } @@ -150,6 +154,7 @@ impl PortForwardMux { window, inbound_rx, FrameSink::Direct(self.tx.clone()), + self.teardown.subscribe(), )); } @@ -169,9 +174,11 @@ impl PortForwardMux { self.channels.get(&id).map(|e| e.window.lock().unwrap().remaining()) } - /// Dropping every sender ends every task; each drops its target socket. + /// Every task ends and drops its target socket: the queue's senders go for + /// a task on the queue, `teardown` reaches one parked on the socket. pub fn close_all(&mut self) { self.channels.clear(); + self.teardown.send(()).ok(); } } @@ -186,6 +193,7 @@ async fn run_controlled_channel( window: Arc>, mut inbound: mpsc::UnboundedReceiver, sink: FrameSink, + teardown: watch::Receiver<()>, ) { let mut pending: Vec = Vec::new(); let mut pending_len = 0usize; @@ -235,7 +243,7 @@ async fn run_controlled_channel( return; } let (reader, writer) = socket.into_split(); - run_channel(id, reader, writer, Vec::new(), pending, credit, window, inbound, sink).await; + run_channel(id, reader, writer, Vec::new(), pending, credit, window, inbound, sink, teardown).await; } /// The same words the raw pipe puts in its login error, so one problem reads