TheLanc3/rustdesk
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk.git synced 2026-09-05 15:41:23 +03:00
Code Issues Packages Projects Releases Wiki Activity

port forward: share one multiplexed tunnel across a window's listeners

Browse Source 
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab
This commit is contained in:
rustdesk
2026-09-04 01:36:24 +08:00
parent a05247a07a
commit 48fe740ad8
5 changed files with 191 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
libs/hbb_common

Submodule libs/hbb_common updated: 6918d35935...26da70418c

1
src/client.rs
View File

@@ -2769,6 +2769,7 @@ impl LoginConfigHandler {
ConnType::PORT_FORWARD | ConnType::RDP => lr.set_port_forward(PortForward {
host: self.port_forward.0.clone(),
port: self.port_forward.1,
multiplex: crate::common::get_port_forward_mux_enabled(),
..Default::default()
}),
ConnType::TERMINAL => {

7
src/common.rs
View File

@@ -1167,6 +1167,13 @@ pub fn get_ipv6_punch_enabled() -> bool {
)
}
pub fn get_port_forward_mux_enabled() -> bool {
config::option2bool(
keys::OPTION_ENABLE_PORT_FORWARD_MUX,
&get_local_option(keys::OPTION_ENABLE_PORT_FORWARD_MUX),
)
}
pub fn get_local_option(key: &str) -> String {
let v = LocalConfig::get_option(key);
if key == keys::OPTION_ENABLE_UDP_PUNCH || key == keys::OPTION_ENABLE_IPV6_PUNCH {

359
src/port_forward.rs
View File

@@ -1,6 +1,7 @@
use std::sync::{Arc, RwLock};
use crate::client::*;
use crate::port_forward_mux::{Claim, Tunnel, CHANNEL_WINDOW};
use hbb_common::{
allow_err, bail,
config::READ_TIMEOUT,
@@ -79,6 +80,7 @@ pub async fn listen(
lc: Arc<RwLock<LoginConfigHandler>>,
remote_host: String,
remote_port: i32,
tunnel: Arc<Tunnel>,
) -> ResultType<()> {
let listener = tcp::new_listener(format!("127.0.0.1:{}", port), true).await?;
let addr = listener.local_addr()?;
@@ -90,29 +92,94 @@ pub async fn listen(
let mut ui_receiver = ui_receiver;
loop {
tokio::select! {
Ok((forward, addr)) = listener.accept() => {
log::info!("new connection from {:?}", addr);
let id = id.clone();
let password = password.clone();
let mut forward = Framed::new(forward, BytesCodec::new());
let mut close_port_forward = false;
match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward, &remote_host, remote_port).await {
Ok(Some(stream)) => {
let interface = interface.clone();
tokio::spawn(async move {
if let Err(err) = run_forward(forward, stream).await {
interface.msgbox("error", "Error", &err.to_string(), "");
// `addr` above the loop is the listener's own address; `run_rdp` needs
// that port. The accepted peer's address gets its own name so it can
// never shadow it.
Ok((forward, peer_addr)) = listener.accept() => {
log::debug!("new connection from {:?}", peer_addr);
let claim = tunnel.try_claim();
let claim = match claim {
Claim::Wait => {
// Keep servicing the UI while the establishing accept holds the
// prompt. `None` means establishment failed: this accept is dropped.
let resolved: Option<Claim> = loop {
tokio::select! {
// `tokio::select!` brings its own `Poll::{Ready, Pending}` into
// scope, so `Ready` must be qualified here or it resolves to that.
r = tunnel.wait_ready() => break match r {
crate::port_forward_mux::Ready::Muxed(h) => Some(Claim::Muxed(h)),
crate::port_forward_mux::Ready::Legacy => Some(Claim::Legacy),
crate::port_forward_mux::Ready::Failed => None,
},
d = ui_receiver.recv() => match d {
Some(Data::Close) => return Ok(()),
Some(Data::NewRDP) => run_rdp(addr.port(), &rdp_display_name(&lc, &id)),
_ => {}
},
}
log::info!("connection from {:?} closed", addr);
});
};
match resolved {
Some(c) => c,
None => {
log::debug!("tunnel failed while {:?} waited; dropping it", peer_addr);
continue;
}
}
}
_ if close_port_forward => {
break;
other => other,
};
match claim {
Claim::Muxed(handle) => {
if let Err(e) = handle.open(&remote_host, remote_port, forward, Vec::new()) {
log::debug!("cannot open channel for {:?}: {}", peer_addr, e);
}
}
Err(err) => {
interface.on_establish_connection_error(err.to_string());
Claim::Legacy => {
lc.write().unwrap().port_forward = (remote_host.clone(), remote_port);
let mut forward = Framed::new(forward, BytesCodec::new());
let mut close_port_forward = false;
match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward).await {
Ok(Some(outcome)) => run_legacy(outcome, forward, peer_addr, interface.clone()),
_ if close_port_forward => break,
Err(err) => interface.on_establish_connection_error(err.to_string()),
_ => {}
}
}
_ => {}
Claim::Claimed => {
lc.write().unwrap().port_forward = (remote_host.clone(), remote_port);
let mut forward = Framed::new(forward, BytesCodec::new());
let mut close_port_forward = false;
match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward).await {
Ok(Some(outcome)) if outcome.mux => {
let handle = tunnel.set_muxed(outcome.stream, interface.clone());
if !outcome.local_eof {
let (socket, prebuf) = take_socket(forward, outcome.prebuf);
if let Err(e) = handle.open(&remote_host, remote_port, socket, prebuf) {
log::debug!("cannot open channel for {:?}: {}", peer_addr, e);
}
}
}
Ok(Some(outcome)) => {
tunnel.set_legacy();
if outcome.local_eof {
log::debug!("legacy peer and local {:?} already gone", peer_addr);
} else {
run_legacy(outcome, forward, peer_addr, interface.clone());
}
}
_ if close_port_forward => {
tunnel.set_failed();
break;
}
Err(err) => {
tunnel.set_failed();
interface.on_establish_connection_error(err.to_string());
}
_ => tunnel.set_failed(),
}
}
// Resolved above; a stray `Wait` just drops this accept.
Claim::Wait => continue,
}
}
d = ui_receiver.recv() => {
@@ -132,6 +199,45 @@ pub async fn listen(
Ok(())
}
/// Today's raw pipe, for peers without multiplexing.
fn run_legacy(
outcome: LoginOutcome,
forward: Framed<TcpStream, BytesCodec>,
addr: std::net::SocketAddr,
interface: impl Interface,
) {
let mut stream = outcome.stream;
let prebuf = outcome.prebuf;
tokio::spawn(async move {
stream.set_raw();
if !prebuf.is_empty() {
allow_err!(stream.send_bytes(prebuf.into()).await);
}
if let Err(err) = run_forward(forward, stream).await {
interface.msgbox("error", "Error", &err.to_string(), "");
}
log::info!("connection from {:?} closed", addr);
});
}
pub(crate) struct LoginOutcome {
stream: Stream,
mux: bool,
prebuf: Vec<u8>,
local_eof: bool,
}
fn peer_supports_mux(pi: &PeerInfo) -> bool {
pi.features.as_ref().map(|f| f.port_forward_mux).unwrap_or(false)
}
/// `into_inner()` would drop bytes the codec pulled but never yielded.
fn take_socket(forward: Framed<TcpStream, BytesCodec>, mut prebuf: Vec<u8>) -> (TcpStream, Vec<u8>) {
let parts = forward.into_parts();
prebuf.extend_from_slice(&parts.read_buf);
(parts.io, prebuf)
}
async fn connect_and_login(
id: &str,
password: &str,
@@ -142,9 +248,7 @@ async fn connect_and_login(
token: &str,
is_rdp: bool,
close_port_forward: &mut bool,
remote_host: &str,
remote_port: i32,
) -> ResultType<Option<Stream>> {
) -> ResultType<Option<LoginOutcome>> {
let conn_type = if is_rdp {
ConnType::RDP
} else {
@@ -160,6 +264,8 @@ async fn connect_and_login(
}
}
let mut buffer = Vec::new();
let mut local_eof = false;
let mut mux = false;
let mut received = false;
let mut challenge = None;
let mut pending_login = None;
@@ -192,6 +298,7 @@ async fn connect_and_login(
}
}
Some(login_response::Union::PeerInfo(pi)) => {
mux = peer_supports_mux(&pi);
interface.handle_peer_info(pi);
break;
}
@@ -222,20 +329,24 @@ async fn connect_and_login(
_ => {}
}
},
res = forward.next() => {
// Stop pulling once the pre-read buffer is a window deep; the
// rest waits in the kernel until the channel opens. A local EOF
// no longer aborts the login: the tunnel may still be wanted.
res = forward.next(), if !local_eof && buffer.len() < CHANNEL_WINDOW as usize => {
if let Some(Ok(bytes)) = res {
buffer.extend(bytes);
} else {
return Ok(None);
local_eof = true;
}
},
}
}
stream.set_raw();
if !buffer.is_empty() {
allow_err!(stream.send_bytes(buffer.into()).await);
}
Ok(Some(stream))
Ok(Some(LoginOutcome {
stream,
mux,
prebuf: buffer,
local_eof,
}))
}
@@ -334,170 +445,50 @@ async fn run_forward(forward: Framed<TcpStream, BytesCodec>, stream: Stream) ->
}
#[cfg(test)]
mod login_tests {
mod tests {
use super::*;
use async_trait::async_trait;
use hbb_common::{
tcp::FramedStream,
tokio::time::{sleep, Duration},
};
use sha2::{Digest, Sha256};
/// A window's interface over its shared handler. `handle_hash` can pause
/// before building the login, where the real one looks passwords up.
#[derive(Clone)]
struct Ui {
lc: Arc<RwLock<LoginConfigHandler>>,
pause: Duration,
}
#[async_trait]
impl Interface for Ui {
fn send(&self, _data: Data) {}
fn msgbox(&self, _msgtype: &str, _title: &str, _text: &str, _link: &str) {}
fn handle_login_error(&self, _err: &str) -> bool {
false
}
fn handle_peer_info(&self, _pi: PeerInfo) {}
fn set_multiple_windows_session(&self, _sessions: Vec<WindowsSession>) {}
async fn handle_hash(&self, pass: &str, hash: Hash, peer: &mut Stream) -> bool {
sleep(self.pause).await;
crate::client::handle_hash(self.lc.clone(), pass, hash, self, peer).await
}
async fn handle_login_from_ui(
&self,
os_username: String,
os_password: String,
password: String,
remember: bool,
peer: &mut Stream,
) {
crate::client::handle_login_from_ui(
self.lc.clone(),
os_username,
os_password,
password,
remember,
peer,
)
.await
}
async fn handle_test_delay(&self, _t: TestDelay, _peer: &mut Stream) {}
fn get_lch(&self) -> Arc<RwLock<LoginConfigHandler>> {
self.lc.clone()
}
}
fn window() -> Ui {
let mut lc = LoginConfigHandler::default();
lc.conn_type = ConnType::PORT_FORWARD;
Ui {
lc: Arc::new(RwLock::new(lc)),
pause: Duration::ZERO,
}
}
/// (our end, the peer's end) of one connection.
async fn loopback() -> (Stream, Stream) {
let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = l.local_addr().unwrap();
let client = tokio::net::TcpStream::connect(addr).await.unwrap();
let (server, _) = l.accept().await.unwrap();
(
Stream::Tcp(FramedStream::from(client, addr)),
Stream::Tcp(FramedStream::from(server, addr)),
)
}
async fn login_at(peer: &mut Stream) -> LoginRequest {
let bytes = peer.next().await.unwrap().unwrap();
Message::parse_from_bytes(&bytes)
.unwrap()
.login_request()
.clone()
}
fn target(lr: &LoginRequest) -> (String, i32) {
(lr.port_forward().host.clone(), lr.port_forward().port)
}
fn hash(challenge: &str) -> Hash {
Hash {
salt: "salt".to_owned(),
challenge: challenge.to_owned(),
..Default::default()
}
}
/// What the peer expects for password `pw` under `hash(challenge)`.
fn digest(challenge: &str) -> Vec<u8> {
let mut h = Sha256::new();
h.update("pw");
h.update("salt");
let salted = h.finalize();
let mut h2 = Sha256::new();
h2.update(&salted[..]);
h2.update(challenge);
h2.finalize()[..].to_vec()
#[test]
fn peer_supports_mux_reads_the_features_bit() {
let mut pi = PeerInfo::new();
assert!(!peer_supports_mux(&pi));
pi.features = Some(Features { port_forward_mux: false, ..Default::default() }).into();
assert!(!peer_supports_mux(&pi));
pi.features = Some(Features { port_forward_mux: true, ..Default::default() }).into();
assert!(peer_supports_mux(&pi));
}
#[test]
fn mappings_logging_in_at_once_each_carry_their_own_target() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let mut ui = window();
ui.pause = Duration::from_millis(50);
let (mut a, mut a_peer) = loopback().await;
let (mut b, mut b_peer) = loopback().await;
tokio::join!(
login_with_hash(&ui, "pw", hash("a"), "a", 1, &mut a),
login_with_hash(&ui, "pw", hash("b"), "b", 2, &mut b),
);
assert_eq!(target(&login_at(&mut a_peer).await), ("a".to_owned(), 1));
assert_eq!(target(&login_at(&mut b_peer).await), ("b".to_owned(), 2));
});
fn port_forward_mux_defaults_to_on() {
use hbb_common::config::{keys, option2bool};
// The `enable-` prefix is what makes an unset value mean on. Getting the
// key's name wrong silently flips the default, so pin it here.
assert!(option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, ""));
assert!(option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, "Y"));
assert!(!option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, "N"));
}
#[test]
fn a_mapping_answers_the_prompt_with_its_own_challenge() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
fn take_socket_hands_back_a_working_socket_and_the_prebuf() {
use hbb_common::tokio::io::{AsyncReadExt, AsyncWriteExt};
let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
rt.block_on(async {
let ui = window();
let (mut a, mut a_peer) = loopback().await;
let (mut b, mut b_peer) = loopback().await;
// A's hash arrived last, so it is the one the handler holds.
assert!(login_with_hash(&ui, "pw", hash("a"), "a", 1, &mut a).await);
login_at(&mut a_peer).await;
let typed = (String::new(), String::new(), "pw".to_owned(), false);
login_from_ui(&ui, &hash("b"), typed, "b", 2, &mut b).await;
let lr = login_at(&mut b_peer).await;
assert_eq!(lr.password, digest("b"));
assert_eq!(target(&lr), ("b".to_owned(), 2));
});
}
#[test]
fn a_password_typed_before_this_connections_hash_answers_it_when_it_comes() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let ui = window();
let (mut b, mut b_peer) = loopback().await;
// The prompt's password reached B before its hash, and no other
// mapping has stored it in the handler yet.
let typed = (String::new(), String::new(), "pw".to_owned(), false);
assert!(hash_arrived(&ui, "", hash("b"), Some(typed), "b", 2, &mut b).await);
let lr = login_at(&mut b_peer).await;
assert_eq!(lr.password, digest("b"));
assert_eq!(target(&lr), ("b".to_owned(), 2));
let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = l.local_addr().unwrap();
let mut client = TcpStream::connect(addr).await.unwrap();
let (server, _) = l.accept().await.unwrap();
let mut framed = Framed::new(server, BytesCodec::new());
client.write_all(b"abc").await.unwrap();
// Read through the codec, as connect_and_login does during login.
let pulled = framed.next().await.unwrap().unwrap();
assert_eq!(&pulled[..], b"abc");
let (mut sock, prebuf) = take_socket(framed, pulled.to_vec());
assert_eq!(prebuf, b"abc".to_vec());
// Bytes written after the handoff arrive on the bare socket.
client.write_all(b"def").await.unwrap();
let mut buf = [0u8; 3];
sock.read_exact(&mut buf).await.unwrap();
assert_eq!(&buf, b"def");
});
}
}

8
src/ui_session_interface.rs
View File

@@ -1944,6 +1944,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
let key = crate::get_key(false).await;
#[cfg(not(any(target_os = "android", target_os = "ios")))]
if handler.is_port_forward() {
let tunnel = std::sync::Arc::new(crate::port_forward_mux::Tunnel::new());
if handler.is_rdp() {
let port = handler
.get_option("rdp_port".to_owned())
@@ -1958,7 +1959,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
handler.get_option("rdp_password".to_owned()),
);
log::info!("Remote rdp port: {}", port);
start_one_port_forward(handler, 0, "".to_owned(), port, receiver, &key, &token).await;
start_one_port_forward(handler, 0, "".to_owned(), port, receiver, &key, &token, tunnel).await;
} else if handler.args.len() == 0 {
let pfs = handler.lc.read().unwrap().port_forwards.clone();
let mut queues = HashMap::<i32, mpsc::UnboundedSender<Data>>::new();
@@ -1976,6 +1977,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
let handler = handler.clone();
let key = key.clone();
let token = token.clone();
let tunnel = tunnel.clone();
tokio::spawn(async move {
start_one_port_forward(
handler,
@@ -1985,6 +1987,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
receiver,
&key,
&token,
tunnel,
)
.await;
});
@@ -2023,6 +2026,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
receiver,
&key,
&token,
tunnel,
)
.await;
}
@@ -2042,6 +2046,7 @@ async fn start_one_port_forward<T: InvokeUiSession>(
receiver: mpsc::UnboundedReceiver<Data>,
key: &str,
token: &str,
tunnel: std::sync::Arc<crate::port_forward_mux::Tunnel>,
) {
if let Err(err) = crate::port_forward::listen(
handler.get_id(),
@@ -2054,6 +2059,7 @@ async fn start_one_port_forward<T: InvokeUiSession>(
handler.lc.clone(),
remote_host,
remote_port,
tunnel,
)
.await
{