mirror of
https://github.com/rustdesk/rustdesk.git
synced 2026-09-11 23:11:01 +03:00
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab
This commit is contained in:
@@ -346,6 +346,9 @@ async fn relay_tunnel_to_socket<W: AsyncWrite + Unpin>(
|
|||||||
/// Runs both halves as independent tasks; whichever ends first cancels the
|
/// Runs both halves as independent tasks; whichever ends first cancels the
|
||||||
/// other. Sends `close` once, after the last data, and only when the channel
|
/// 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.
|
/// 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<R, W>(
|
pub async fn run_channel<R, W>(
|
||||||
id: i32,
|
id: i32,
|
||||||
reader: R,
|
reader: R,
|
||||||
@@ -356,6 +359,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<()>,
|
||||||
) where
|
) where
|
||||||
R: AsyncRead + Unpin + Send + 'static,
|
R: AsyncRead + Unpin + Send + 'static,
|
||||||
W: AsyncWrite + Unpin + Send + 'static,
|
W: AsyncWrite + Unpin + Send + 'static,
|
||||||
@@ -376,6 +380,10 @@ 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() => {
|
||||||
|
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 peer_closed = first == RelayEnd::PeerClosed || second == RelayEnd::PeerClosed;
|
||||||
let tunnel_gone = first == RelayEnd::TunnelGone || second == RelayEnd::TunnelGone;
|
let tunnel_gone = first == RelayEnd::TunnelGone || second == RelayEnd::TunnelGone;
|
||||||
@@ -463,6 +471,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,
|
||||||
});
|
});
|
||||||
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
|
||||||
@@ -502,6 +511,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
|
||||||
|
/// reach: one parked on its local socket is not on the inbound queue.
|
||||||
|
teardown: watch::Sender<()>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TunnelHandle {
|
impl TunnelHandle {
|
||||||
@@ -541,10 +553,11 @@ mod tunnel {
|
|||||||
let open = open_msg(id, host, port, CHANNEL_WINDOW);
|
let open = open_msg(id, host, port, CHANNEL_WINDOW);
|
||||||
let (reader, writer) = socket.into_split();
|
let (reader, writer) = socket.into_split();
|
||||||
let sink = self.sink.clone();
|
let sink = self.sink.clone();
|
||||||
|
let teardown = self.teardown.subscribe();
|
||||||
let handle = self.clone();
|
let handle = self.clone();
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
if sink.send_ordered(open).await.is_ok() {
|
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);
|
handle.channels.lock().unwrap().remove(&id);
|
||||||
});
|
});
|
||||||
@@ -643,6 +656,7 @@ mod tunnel {
|
|||||||
|
|
||||||
fn close_all(&self) {
|
fn close_all(&self) {
|
||||||
self.channels.lock().unwrap().clear();
|
self.channels.lock().unwrap().clear();
|
||||||
|
self.teardown.send(()).ok();
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -883,6 +897,7 @@ 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<()>,
|
||||||
task: tokio::task::JoinHandle<()>,
|
task: tokio::task::JoinHandle<()>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -897,10 +912,11 @@ 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 task = tokio::spawn(run_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
|
// `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]
|
#[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 {
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ use hbb_common::{
|
|||||||
log,
|
log,
|
||||||
message_proto::*,
|
message_proto::*,
|
||||||
timeout,
|
timeout,
|
||||||
tokio::{self, net::TcpStream, sync::mpsc},
|
tokio::{self, net::TcpStream, sync::{mpsc, watch}},
|
||||||
};
|
};
|
||||||
use std::{
|
use std::{
|
||||||
collections::HashMap,
|
collections::HashMap,
|
||||||
@@ -35,6 +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:
|
||||||
|
/// one parked on its target socket is not on the inbound queue.
|
||||||
|
teardown: watch::Sender<()>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl PortForwardMux {
|
impl PortForwardMux {
|
||||||
@@ -43,6 +46,7 @@ impl PortForwardMux {
|
|||||||
channels: HashMap::new(),
|
channels: HashMap::new(),
|
||||||
tx,
|
tx,
|
||||||
login_target,
|
login_target,
|
||||||
|
teardown: watch::channel(()).0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -150,6 +154,7 @@ impl PortForwardMux {
|
|||||||
window,
|
window,
|
||||||
inbound_rx,
|
inbound_rx,
|
||||||
FrameSink::Direct(self.tx.clone()),
|
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())
|
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) {
|
pub fn close_all(&mut self) {
|
||||||
self.channels.clear();
|
self.channels.clear();
|
||||||
|
self.teardown.send(()).ok();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -186,6 +193,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<()>,
|
||||||
) {
|
) {
|
||||||
let mut pending: Vec<Bytes> = Vec::new();
|
let mut pending: Vec<Bytes> = Vec::new();
|
||||||
let mut pending_len = 0usize;
|
let mut pending_len = 0usize;
|
||||||
@@ -235,7 +243,7 @@ async fn run_controlled_channel(
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
let (reader, writer) = socket.into_split();
|
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
|
/// The same words the raw pipe puts in its login error, so one problem reads
|
||||||
|
|||||||
Reference in New Issue
Block a user