port forward: a channel opened as its tunnel closes still gets the teardown

`open` can straddle `close_all`: the claim passed, the frame receiver was
still alive, and the channel subscribed after the signal had gone out.
`watch::subscribe` marks earlier sends as seen, and the entry sits in a
map that was already cleared, so nothing would ever end it. The signal is
now a level: `close_all` raises it with `send_replace`, which stores even
with no channel live, and `run_channel` waits for the value rather than
for a change.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab
This commit is contained in:
rustdesk
2026-09-05 11:23:55 +08:00
parent 7587cf514a
commit 89b3f88286
2 changed files with 43 additions and 15 deletions

View File

@@ -348,7 +348,8 @@ async fn relay_tunnel_to_socket<W: AsyncWrite + Unpin>(
/// ended for a local reason — the peer's own `close` is never echoed. /// 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 /// `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 /// even when they are parked on the socket, where dropping the inbound sender
/// reaches neither. /// reaches neither. It is a level, so a channel opened as the tunnel closes,
/// subscribing after the signal went out, still sees it.
pub async fn run_channel<R, W>( pub async fn run_channel<R, W>(
id: i32, id: i32,
reader: R, reader: R,
@@ -359,7 +360,7 @@ pub async fn run_channel<R, W>(
window: Arc<Mutex<RecvWindow>>, window: Arc<Mutex<RecvWindow>>,
inbound: mpsc::UnboundedReceiver<Inbound>, inbound: mpsc::UnboundedReceiver<Inbound>,
sink: FrameSink, sink: FrameSink,
mut teardown: watch::Receiver<()>, mut teardown: watch::Receiver<bool>,
) where ) where
R: AsyncRead + Unpin + Send + 'static, R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static,
@@ -380,7 +381,9 @@ pub async fn run_channel<R, W>(
let _ = cancel_tx.send(true); let _ = cancel_tx.send(true);
(r.unwrap_or(RelayEnd::Cancelled), to_tunnel.await.unwrap_or(RelayEnd::Cancelled)) (r.unwrap_or(RelayEnd::Cancelled), to_tunnel.await.unwrap_or(RelayEnd::Cancelled))
} }
_ = teardown.changed() => { // Wrapped so `select!` keeps a `bool`, not the `Ref` (a read guard,
// not `Send`) it would otherwise hold across the joins.
_ = async { teardown.wait_for(|down| *down).await.is_ok() } => {
let _ = cancel_tx.send(true); let _ = cancel_tx.send(true);
(to_tunnel.await.unwrap_or(RelayEnd::Cancelled), to_socket.await.unwrap_or(RelayEnd::Cancelled)) (to_tunnel.await.unwrap_or(RelayEnd::Cancelled), to_socket.await.unwrap_or(RelayEnd::Cancelled))
} }
@@ -471,7 +474,7 @@ mod tunnel {
channels: Mutex::new(HashMap::new()), channels: Mutex::new(HashMap::new()),
next_id: AtomicI32::new(1), next_id: AtomicI32::new(1),
reported: Default::default(), reported: Default::default(),
teardown: watch::channel(()).0, teardown: watch::channel(false).0,
}); });
let state = self.state.clone(); let state = self.state.clone();
// Publish before spawning: if the loop exits first and resets the // Publish before spawning: if the loop exits first and resets the
@@ -511,9 +514,9 @@ mod tunnel {
channels: Mutex<HashMap<i32, ChannelEntry>>, channels: Mutex<HashMap<i32, ChannelEntry>>,
next_id: AtomicI32, next_id: AtomicI32,
reported: Mutex<HashMap<String, Instant>>, reported: Mutex<HashMap<String, Instant>>,
/// Sent once, by `close_all`, for the channels its `clear` cannot /// Raised once, by `close_all`, for the channels its `clear` cannot
/// reach: one parked on its local socket is not on the inbound queue. /// reach: one parked on its local socket is not on the inbound queue.
teardown: watch::Sender<()>, teardown: watch::Sender<bool>,
} }
impl TunnelHandle { impl TunnelHandle {
@@ -656,7 +659,9 @@ mod tunnel {
fn close_all(&self) { fn close_all(&self) {
self.channels.lock().unwrap().clear(); self.channels.lock().unwrap().clear();
self.teardown.send(()).ok(); // Not `send`: with no channel live it stores nothing, and one
// opened as the tunnel closes would never see it.
self.teardown.send_replace(true);
} }
#[cfg(test)] #[cfg(test)]
@@ -897,13 +902,19 @@ mod tests {
credit: Arc<SendCredit>, credit: Arc<SendCredit>,
window: Arc<Mutex<RecvWindow>>, window: Arc<Mutex<RecvWindow>>,
local: tokio::io::DuplexStream, local: tokio::io::DuplexStream,
teardown: watch::Sender<()>, teardown: watch::Sender<bool>,
task: tokio::task::JoinHandle<()>, task: tokio::task::JoinHandle<()>,
} }
/// A channel whose "local socket" is one end of a duplex pipe and whose /// A channel whose "local socket" is one end of a duplex pipe and whose
/// "tunnel" is a pair of queues the test reads directly. /// "tunnel" is a pair of queues the test reads directly.
fn harness(id: i32, prebuf: Vec<u8>, initial_out: Vec<Bytes>) -> Harness { fn harness(id: i32, prebuf: Vec<u8>, initial_out: Vec<Bytes>) -> Harness {
harness_on(id, prebuf, initial_out, watch::channel(false).0)
}
/// The channel subscribes to `teardown` here, so a test can hand in one
/// that has already been raised.
fn harness_on(id: i32, prebuf: Vec<u8>, initial_out: Vec<Bytes>, teardown: watch::Sender<bool>) -> Harness {
let (data_tx, data_rx) = mpsc::channel(DATA_QUEUE_FRAMES); let (data_tx, data_rx) = mpsc::channel(DATA_QUEUE_FRAMES);
let (control_tx, control_rx) = mpsc::unbounded_channel(); let (control_tx, control_rx) = mpsc::unbounded_channel();
let (inbound_tx, inbound_rx) = mpsc::unbounded_channel(); let (inbound_tx, inbound_rx) = mpsc::unbounded_channel();
@@ -912,7 +923,7 @@ mod tests {
let credit = Arc::new(SendCredit::new(INITIAL_WINDOW)); let credit = Arc::new(SendCredit::new(INITIAL_WINDOW));
let window = Arc::new(Mutex::new(RecvWindow::new(CHANNEL_WINDOW))); let window = Arc::new(Mutex::new(RecvWindow::new(CHANNEL_WINDOW)));
let sink = FrameSink::Queued { data: data_tx, control: control_tx }; let sink = FrameSink::Queued { data: data_tx, control: control_tx };
let (teardown, teardown_rx) = watch::channel(()); let teardown_rx = teardown.subscribe();
let task = tokio::spawn(run_channel( let task = tokio::spawn(run_channel(
id, r, w, prebuf, initial_out, credit.clone(), window.clone(), inbound_rx, sink, teardown_rx, id, r, w, prebuf, initial_out, credit.clone(), window.clone(), inbound_rx, sink, teardown_rx,
)); ));
@@ -954,12 +965,27 @@ mod tests {
h.inbound_tx.send(Inbound::Data(Bytes::from(vec![0u8; MAX_FRAME]))).unwrap(); h.inbound_tx.send(Inbound::Data(Bytes::from(vec![0u8; MAX_FRAME]))).unwrap();
} }
tokio::time::sleep(std::time::Duration::from_millis(50)).await; tokio::time::sleep(std::time::Duration::from_millis(50)).await;
h.teardown.send(()).ok(); h.teardown.send_replace(true);
let ended = tokio::time::timeout(std::time::Duration::from_millis(500), &mut h.task).await; let ended = tokio::time::timeout(std::time::Duration::from_millis(500), &mut h.task).await;
assert!(ended.is_ok(), "channel task outlived the tunnel"); assert!(ended.is_ok(), "channel task outlived the tunnel");
}); });
} }
#[test]
fn a_channel_subscribed_after_teardown_ends_at_once() {
rt().block_on(async {
// `open` can race `close_all`: this channel subscribes after the
// signal went out, and its entry sits in a map that was already
// cleared, so nothing will ever drop its inbound sender. No other
// channel was live when the tunnel closed, either.
let teardown = watch::channel(false).0;
teardown.send_replace(true);
let mut h = harness_on(1, vec![], vec![], teardown);
let ended = tokio::time::timeout(std::time::Duration::from_millis(500), &mut h.task).await;
assert!(ended.is_ok(), "late channel outlived the tunnel");
});
}
#[test] #[test]
fn local_bytes_become_data_frames_capped_at_max_frame() { fn local_bytes_become_data_frames_capped_at_max_frame() {
rt().block_on(async { rt().block_on(async {

View File

@@ -35,9 +35,9 @@ pub struct PortForwardMux {
channels: HashMap<i32, Entry>, channels: HashMap<i32, Entry>,
tx: Sender, tx: Sender,
login_target: String, login_target: String,
/// Sent once, by `close_all`, for the channels its `clear` cannot reach: /// Raised once, by `close_all`, for the channels its `clear` cannot reach:
/// one parked on its target socket is not on the inbound queue. /// one parked on its target socket is not on the inbound queue.
teardown: watch::Sender<()>, teardown: watch::Sender<bool>,
} }
impl PortForwardMux { impl PortForwardMux {
@@ -46,7 +46,7 @@ impl PortForwardMux {
channels: HashMap::new(), channels: HashMap::new(),
tx, tx,
login_target, login_target,
teardown: watch::channel(()).0, teardown: watch::channel(false).0,
} }
} }
@@ -178,7 +178,9 @@ impl PortForwardMux {
/// a task on the queue, `teardown` reaches one parked on the socket. /// a task on the queue, `teardown` reaches one parked on the socket.
pub fn close_all(&mut self) { pub fn close_all(&mut self) {
self.channels.clear(); self.channels.clear();
self.teardown.send(()).ok(); // Not `send`: with no channel live it stores nothing, and one opened
// as the tunnel closes would never see it.
self.teardown.send_replace(true);
} }
} }
@@ -193,7 +195,7 @@ async fn run_controlled_channel(
window: Arc<Mutex<RecvWindow>>, window: Arc<Mutex<RecvWindow>>,
mut inbound: mpsc::UnboundedReceiver<Inbound>, mut inbound: mpsc::UnboundedReceiver<Inbound>,
sink: FrameSink, sink: FrameSink,
teardown: watch::Receiver<()>, teardown: watch::Receiver<bool>,
) { ) {
let mut pending: Vec<Bytes> = Vec::new(); let mut pending: Vec<Bytes> = Vec::new();
let mut pending_len = 0usize; let mut pending_len = 0usize;