From 3aade011c08cb1937c903deb48a711fb0e856958 Mon Sep 17 00:00:00 2001 From: 21pages Date: Mon, 28 Sep 2026 14:14:54 +0800 Subject: [PATCH 1/7] feat: add opt-in WebSocket proxy support Support WS and WSS over SOCKS5, HTTP, and HTTPS proxies, including custom ports. Use the configured proxy only when allow-websocket-proxy is enabled; explicit proxy configuration takes precedence. Propagate proxy failures without retrying a direct connection. Add regression coverage for proxy routing, authentication, TLS, custom ports, direct defaults, and failure handling. Signed-off-by: 21pages fix: flush WebSocket proxy tunnels before native TLS reads Flush buffered HTTP and HTTPS tunnel writes before NativeTls reads so WSS handshakes can complete on macOS. Apply the adapter only to NativeTls WebSocket connections through HTTP or HTTPS proxies. Extend the existing proxy round-trip test with HTTPS and SOCKS5 NativeTls cases. The websocket_proxy integration test passes in the rustdesk workspace. Signed-off-by: 21pages refactor: simplify WebSocket proxy configuration Select the configured proxy inside WsFramedStream::new based on the WebSocket proxy option and remove the explicit proxy argument. Update the socket client and drop the unused Stream::connect_websocket wrapper. Adapt integration tests to configured proxies and cover custom target ports and WebSocket close handling. --- src/config.rs | 6 + src/socket_client.rs | 2 +- src/stream.rs | 16 +- src/websocket.rs | 237 ++++++++++++++----- tests/data/websocket-cert.pem | 12 + tests/data/websocket-key.pem | 5 + tests/websocket_proxy.rs | 419 ++++++++++++++++++++++++++++++++++ 7 files changed, 627 insertions(+), 70 deletions(-) create mode 100644 tests/data/websocket-cert.pem create mode 100644 tests/data/websocket-key.pem create mode 100644 tests/websocket_proxy.rs diff --git a/src/config.rs b/src/config.rs index 6fe100455..1e25a0df5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2842,6 +2842,11 @@ pub fn use_ws() -> bool { option2bool(option, &Config::get_option(option)) } +pub fn ws_use_proxy() -> bool { + let option = keys::OPTION_ALLOW_WEBSOCKET_PROXY; + option2bool(option, &Config::get_option(option)) +} + pub fn allow_insecure_tls_fallback() -> bool { let option = keys::OPTION_ALLOW_INSECURE_TLS_FALLBACK; option2bool(option, &Config::get_option(option)) @@ -2869,6 +2874,7 @@ pub mod keys { pub const OPTION_ALLOW_NUMERNIC_ONE_TIME_PASSWORD: &str = "allow-numeric-one-time-password"; pub const OPTION_DIRECT_SERVER: &str = "direct-server"; pub const OPTION_ALLOW_WEBSOCKET: &str = "allow-websocket"; + pub const OPTION_ALLOW_WEBSOCKET_PROXY: &str = "allow-websocket-proxy"; pub const OPTION_TRACKPAD_SPEED: &str = "trackpad-speed"; pub const OPTION_REGISTER_DEVICE: &str = "register-device"; pub const OPTION_RELAY_SERVER: &str = "relay-server"; diff --git a/src/socket_client.rs b/src/socket_client.rs index d40a0fff4..096a2e5c3 100644 --- a/src/socket_client.rs +++ b/src/socket_client.rs @@ -140,7 +140,7 @@ pub async fn connect_tcp< let target_str = check_ws(&target.to_string()); if is_ws_endpoint(&target_str) { return Ok(Stream::WebSocket( - websocket::WsFramedStream::new(target_str, None, None, ms_timeout).await?, + websocket::WsFramedStream::new(target_str, None, ms_timeout).await?, )); } connect_tcp_local(target, None, ms_timeout).await diff --git a/src/stream.rs b/src/stream.rs index a486302d4..126d08a86 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,4 +1,4 @@ -use crate::{bail, config, tcp, websocket, ResultType}; +use crate::{bail, tcp, websocket, ResultType}; #[cfg(feature = "webrtc")] use crate::webrtc; use sodiumoxide::crypto::secretbox::Key; @@ -245,20 +245,6 @@ impl Stream { } } - /// establish connect from websocket - #[inline] - pub async fn connect_websocket( - url: impl AsRef, - local_addr: Option, - proxy_conf: Option<&config::Socks5Server>, - timeout_ms: u64, - ) -> ResultType { - let ws_stream = - websocket::WsFramedStream::new(url, local_addr, proxy_conf, timeout_ms).await?; - log::debug!("WebSocket connection established"); - Ok(Self::WebSocket(ws_stream)) - } - /// send message #[inline] pub async fn send(&mut self, msg: &impl protobuf::Message) -> ResultType<()> { diff --git a/src/websocket.rs b/src/websocket.rs index 6c506250b..3cf0b0b92 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -1,40 +1,84 @@ use crate::{ config::{ - keys::OPTION_RELAY_SERVER, use_ws, Config, Socks5Server, RELAY_PORT, RENDEZVOUS_PORT, + keys::OPTION_RELAY_SERVER, use_ws, ws_use_proxy, Config, Socks5Server, RELAY_PORT, + RENDEZVOUS_PORT, }, protobuf::Message, socket_client::split_host_port, sodiumoxide::crypto::secretbox::Key, - tcp::{Encrypt, KxTranscript}, + tcp::{DynTcpStream, Encrypt, KxTranscript}, tls::{get_cached_tls_accept_invalid_cert, get_cached_tls_type, upsert_tls_cache, TlsType}, ResultType, }; use anyhow::bail; use async_recursion::async_recursion; use bytes::{Bytes, BytesMut}; -use futures::{SinkExt, StreamExt}; +use futures::{future::Either, SinkExt, StreamExt}; use std::{ - io::{Error, ErrorKind}, + io::{self, Error, ErrorKind}, net::SocketAddr, + pin::Pin, sync::Arc, + task::{ready, Context, Poll}, time::Duration, }; -use tokio::{net::TcpStream, time::timeout}; +use tokio::{ + io::{AsyncRead, AsyncWrite, ReadBuf}, + net::TcpStream, + time::timeout, +}; use tokio_native_tls::native_tls::TlsConnector; use tokio_tungstenite::{ - connect_async_tls_with_config, tungstenite::protocol::Message as WsMessage, Connector, - MaybeTlsStream, WebSocketStream, + client_async_tls_with_config, connect_async_tls_with_config, + tungstenite::protocol::Message as WsMessage, Connector, MaybeTlsStream, WebSocketStream, }; use tungstenite::client::IntoClientRequest; -use tungstenite::protocol::Role; /// tungstenite's own defaults. `set_max_packet_length` clamps to these, so `usize::MAX` puts the /// transport back exactly where it started rather than above it. const DEFAULT_MAX_MESSAGE_SIZE: usize = 64 << 20; const DEFAULT_MAX_FRAME_SIZE: usize = 16 << 20; +// macOS/iOS Native TLS may wait for a response without flushing the proxy tunnel's write buffer. +struct FlushBeforeRead(DynTcpStream); + +impl AsyncRead for FlushBeforeRead { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + ready!(Pin::new(&mut self.0).poll_flush(cx))?; + Pin::new(&mut self.0).poll_read(cx, buf) + } +} + +impl AsyncWrite for FlushBeforeRead { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.0).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.0).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.0).poll_shutdown(cx) + } +} + +// Keep the original non-proxy connection path while supporting proxy streams. +type WsStream = Either< + WebSocketStream>, + WebSocketStream>, +>; + pub struct WsFramedStream { - stream: WebSocketStream>, + stream: WsStream, addr: SocketAddr, encrypt: Option, send_timeout: u64, @@ -72,16 +116,18 @@ impl WsFramedStream { async fn connect( url: &str, + local_addr: Option, + proxy_conf: Option<&Socks5Server>, ms_timeout: u64, - ) -> ResultType>> { - // to-do: websocket proxy. - + ) -> ResultType<(WsStream, SocketAddr)> { let tls_type = get_cached_tls_type(url); let is_tls_type_cached = tls_type.is_some(); let tls_type = tls_type.unwrap_or(TlsType::Rustls); let danger_accept_invalid_cert = get_cached_tls_accept_invalid_cert(&url); Self::try_connect( url, + local_addr, + proxy_conf, ms_timeout, tls_type, is_tls_type_cached, @@ -94,28 +140,83 @@ impl WsFramedStream { #[async_recursion] async fn try_connect( url: &str, + local_addr: Option, + proxy_conf: Option<&Socks5Server>, ms_timeout: u64, tls_type: TlsType, is_tls_type_cached: bool, danger_accept_invalid_cert: Option, original_danger_accept_invalid_certs: Option, - ) -> ResultType>> { + ) -> ResultType<(WsStream, SocketAddr)> { let ws_config = None; - let disable_nagle = false; let request = url .into_client_request() .map_err(|e| Error::new(ErrorKind::Other, e))?; let connector = Self::get_connector(&tls_type, danger_accept_invalid_cert.unwrap_or(false))?; - match timeout( - Duration::from_millis(ms_timeout), - connect_async_tls_with_config(request, ws_config, disable_nagle, connector), - ) - .await? - { - Ok((ws_stream, _)) => { + let result = if let Some(proxy_conf) = proxy_conf { + let host = request.uri().host().ok_or_else(|| { + Error::new(ErrorKind::InvalidInput, "Invalid WebSocket URL: no host") + })?; + let port = request + .uri() + .port_u16() + .or_else(|| match request.uri().scheme_str() { + Some("wss") => Some(443), + Some("ws") => Some(80), + _ => None, + }) + .ok_or_else(|| { + Error::new(ErrorKind::InvalidInput, "Invalid WebSocket URL scheme") + })?; + let target = format!("{host}:{port}"); + timeout(Duration::from_millis(ms_timeout), async { + let stream = crate::tcp::FramedStream::connect( + target.as_str(), + local_addr, + proxy_conf, + ms_timeout, + ) + .await?; + let addr = stream.1; + let mut stream = stream.0.into_inner(); + if cfg!(any(target_os = "macos", target_os = "ios")) + && url.starts_with("wss://") + && matches!(tls_type, TlsType::NativeTls) + && url::Url::parse(&proxy_conf.proxy) + .is_ok_and(|url| matches!(url.scheme(), "http" | "https")) + { + stream = DynTcpStream(Box::new(FlushBeforeRead(stream))); + } + Ok::<_, anyhow::Error>( + client_async_tls_with_config(request, stream, ws_config, connector) + .await + .map(|(stream, _)| (stream, addr)), + ) + }) + .await?? + .map(Either::Right) + } else { + let disable_nagle = false; + timeout( + Duration::from_millis(ms_timeout), + connect_async_tls_with_config(request, ws_config, disable_nagle, connector), + ) + .await? + .map(|(stream, _)| Either::Left(stream)) + }; + match result { + Ok(stream) => { upsert_tls_cache(url, tls_type, danger_accept_invalid_cert.unwrap_or(false)); - Ok(ws_stream) + match stream { + Either::Left(stream) => { + // MaybeTlsStream::get_ref() handles Plain, NativeTls and Rustls, + // returning the underlying TcpStream for peer_addr(). + let addr = stream.get_ref().get_ref().peer_addr()?; + Ok((Either::Left(stream), addr)) + } + Either::Right((stream, addr)) => Ok((Either::Right(stream), addr)), + } } Err(e) => match (tls_type, is_tls_type_cached, danger_accept_invalid_cert) { (TlsType::Rustls, _, None) => { @@ -126,6 +227,8 @@ impl WsFramedStream { ); Self::try_connect( url, + local_addr, + proxy_conf, ms_timeout, tls_type, is_tls_type_cached, @@ -142,6 +245,8 @@ impl WsFramedStream { ); Self::try_connect( url, + local_addr, + proxy_conf, ms_timeout, TlsType::NativeTls, is_tls_type_cached, @@ -158,6 +263,8 @@ impl WsFramedStream { ); Self::try_connect( url, + local_addr, + proxy_conf, ms_timeout, tls_type, is_tls_type_cached, @@ -189,17 +296,16 @@ impl WsFramedStream { pub async fn new>( url: T, - _local_addr: Option, - _proxy_conf: Option<&Socks5Server>, + local_addr: Option, ms_timeout: u64, ) -> ResultType { - let stream = Self::connect(url.as_ref(), ms_timeout).await?; - let addr = match stream.get_ref() { - MaybeTlsStream::Plain(tcp) => tcp.peer_addr()?, - MaybeTlsStream::NativeTls(tls) => tls.get_ref().get_ref().get_ref().peer_addr()?, - MaybeTlsStream::Rustls(tls) => tls.get_ref().0.peer_addr()?, - _ => return Err(Error::new(ErrorKind::Other, "Unsupported stream type").into()), + let proxy_conf = if ws_use_proxy() { + Config::get_socks() + } else { + None }; + let (stream, addr) = + Self::connect(url.as_ref(), local_addr, proxy_conf.as_ref(), ms_timeout).await?; let ws = Self { stream, @@ -222,24 +328,14 @@ impl WsFramedStream { /// fork growing the read buffer a chunk at a time rather than to the length a frame declares. #[inline] pub fn set_max_packet_length(&mut self, n: usize) { - self.stream.set_config(|c| { + let set_config = |c: &mut tungstenite::protocol::WebSocketConfig| { c.max_message_size = Some(n.min(DEFAULT_MAX_MESSAGE_SIZE)); c.max_frame_size = Some(n.min(DEFAULT_MAX_FRAME_SIZE)); - }); - } - - #[inline] - pub async fn from_tcp_stream(stream: TcpStream, addr: SocketAddr) -> ResultType { - let ws_stream = - WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Client, None) - .await; - - Ok(Self { - stream: ws_stream, - addr, - encrypt: None, - send_timeout: 0, - }) + }; + match &mut self.stream { + Either::Left(stream) => stream.set_config(set_config), + Either::Right(stream) => stream.set_config(set_config), + } } #[inline] @@ -439,6 +535,20 @@ mod tests { use super::*; use crate::config::{keys, Config}; use tokio::{io::AsyncWriteExt, net::TcpListener}; + use tungstenite::protocol::Role; + + async fn ws_from_tcp(stream: TcpStream, addr: SocketAddr) -> WsFramedStream { + let ws_stream = + WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Client, None) + .await; + + WsFramedStream { + stream: Either::Left(ws_stream), + addr, + encrypt: None, + send_timeout: 0, + } + } #[tokio::test(flavor = "current_thread")] async fn test_secured_stream_rejects_plaintext_text() { @@ -455,7 +565,7 @@ mod tests { }); let tcp = TcpStream::connect(addr).await.unwrap(); - let mut client = WsFramedStream::from_tcp_stream(tcp, addr).await.unwrap(); + let mut client = ws_from_tcp(tcp, addr).await; client.set_key(Key([0x42; sodiumoxide::crypto::secretbox::KEYBYTES])); let result = client.next().await; @@ -629,7 +739,7 @@ mod tests { let addr = listener.local_addr().unwrap(); let client = TcpStream::connect(addr).await.unwrap(); let (server, _) = listener.accept().await.unwrap(); - let ws = WsFramedStream::from_tcp_stream(client, addr).await.unwrap(); + let ws = ws_from_tcp(client, addr).await; (ws, server) } @@ -640,7 +750,10 @@ mod tests { ws.set_max_packet_length(CAP); // One byte over, header only: refused before there is any payload to buffer. - server.write_all(&ws_binary_frame(CAP + 1, false)).await.unwrap(); + server + .write_all(&ws_binary_frame(CAP + 1, false)) + .await + .unwrap(); match timeout(Duration::from_secs(5), ws.next()).await { Ok(Some(Err(e))) => assert!(e.to_string().contains("Message too long"), "{}", e), Ok(other) => panic!( @@ -657,18 +770,28 @@ mod tests { let (mut ws, mut server) = ws_loopback().await; ws.set_max_packet_length(CAP); - server.write_all(&ws_binary_frame(8 * 1024, true)).await.unwrap(); + server + .write_all(&ws_binary_frame(8 * 1024, true)) + .await + .unwrap(); let got = ws.next().await.unwrap().unwrap(); assert_eq!(got.len(), 8 * 1024, "a message under the cap still arrives"); ws.set_max_packet_length(usize::MAX); - server.write_all(&ws_binary_frame(200_000, true)).await.unwrap(); + server + .write_all(&ws_binary_frame(200_000, true)) + .await + .unwrap(); let got = timeout(Duration::from_secs(5), ws.next()) .await .expect("next() hung after the cap was lifted") .unwrap() .unwrap(); - assert_eq!(got.len(), 200_000, "lifting the cap lets a large message through again"); + assert_eq!( + got.len(), + 200_000, + "lifting the cap lets a large message through again" + ); } // Two frames each under the cap that reassemble to a message over it: the frame bound lets @@ -680,8 +803,14 @@ mod tests { ws.set_max_packet_length(CAP); // Binary with FIN clear, then a continuation with FIN set: 12 KiB each, 24 KiB together. - server.write_all(&ws_frame(0x02, 12 * 1024, true)).await.unwrap(); - server.write_all(&ws_frame(0x80, 12 * 1024, true)).await.unwrap(); + server + .write_all(&ws_frame(0x02, 12 * 1024, true)) + .await + .unwrap(); + server + .write_all(&ws_frame(0x80, 12 * 1024, true)) + .await + .unwrap(); match timeout(Duration::from_secs(5), ws.next()).await { Ok(Some(Err(e))) => assert!(e.to_string().contains("Message too long"), "{}", e), Ok(other) => panic!( diff --git a/tests/data/websocket-cert.pem b/tests/data/websocket-cert.pem new file mode 100644 index 000000000..b289bf58a --- /dev/null +++ b/tests/data/websocket-cert.pem @@ -0,0 +1,12 @@ +-----BEGIN CERTIFICATE----- +MIIB1TCCAXugAwIBAgIUU5BrjQyAn16xm5+tfdJlWneyQA8wCgYIKoZIzj0EAwIw +IzEhMB8GA1UEAwwYd2Vic29ja2V0LXRhcmdldC5pbnZhbGlkMCAXDTI2MDkwODEx +NTg0NloYDzIxMjYwODE1MTE1ODQ2WjAjMSEwHwYDVQQDDBh3ZWJzb2NrZXQtdGFy +Z2V0LmludmFsaWQwWTATBgcqhkjOPQIBBggqhkjOPQMBBwNCAATGz2rNvO/J5IZc +R2tAP18dE/cDqaqhSI/GrG9wsRKRDypx6C1AlhHp67+xsBxfxGX8qP+2IYqxPeEf +OaXTNmRuo4GKMIGHMB0GA1UdDgQWBBSQaINp0ts4Dxa2+7VJqavc7dykejAfBgNV +HSMEGDAWgBSQaINp0ts4Dxa2+7VJqavc7dykejAPBgNVHRMBAf8EBTADAQH/MDQG +A1UdEQQtMCuCGHdlYnNvY2tldC10YXJnZXQuaW52YWxpZIIJbG9jYWxob3N0hwR/ +AAABMAoGCCqGSM49BAMCA0gAMEUCIQCfNhZYsn3MB0781bnGmxY5y+qvq1bl1xgx +6hB2jeCiwQIgeJjyi0sD3kg9IwExk1tvY1CqME8S3vX6B3A6OihqYUg= +-----END CERTIFICATE----- diff --git a/tests/data/websocket-key.pem b/tests/data/websocket-key.pem new file mode 100644 index 000000000..1b2623908 --- /dev/null +++ b/tests/data/websocket-key.pem @@ -0,0 +1,5 @@ +-----BEGIN PRIVATE KEY----- +MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgfMmTjDYIQrnprR5R +3NkH2woEGhwRtkkOSoFqi9yAeN6hRANCAATGz2rNvO/J5IZcR2tAP18dE/cDqaqh +SI/GrG9wsRKRDypx6C1AlhHp67+xsBxfxGX8qP+2IYqxPeEfOaXTNmRu +-----END PRIVATE KEY----- diff --git a/tests/websocket_proxy.rs b/tests/websocket_proxy.rs new file mode 100644 index 000000000..3d6d09edb --- /dev/null +++ b/tests/websocket_proxy.rs @@ -0,0 +1,419 @@ +use futures::{SinkExt, StreamExt}; +use hbb_common::{ + config::{keys, Socks5Server, APP_NAME, OVERWRITE_SETTINGS}, + socket_client, + tcp::DynTcpStream, + tls::{upsert_tls_cache, TlsType}, + websocket::WsFramedStream, + Stream, +}; +use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer}; +use std::{sync::Arc, time::Duration}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + time::timeout, +}; +use tokio_rustls::{rustls, TlsAcceptor}; +use tokio_tungstenite::tungstenite::handshake::server::{Request, Response}; + +const TIMEOUT: u64 = 5_000; +const PAYLOAD: &[u8] = b"RustDesk websocket proxy test"; + +fn tls_acceptor() -> TlsAcceptor { + // This self-signed identity is only used by the local mock servers. + let cert = CertificateDer::from_pem_slice(include_bytes!("data/websocket-cert.pem")).unwrap(); + let key = PrivateKeyDer::from_pem_slice(include_bytes!("data/websocket-key.pem")).unwrap(); + let config = rustls::ServerConfig::builder_with_provider(Arc::new( + rustls::crypto::ring::default_provider(), + )) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert(vec![cert], key) + .unwrap(); + TlsAcceptor::from(Arc::new(config)) +} + +async fn read_headers(stream: &mut DynTcpStream) -> String { + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + assert!(headers.len() < 4096); + headers.push(stream.read_u8().await.unwrap()); + } + String::from_utf8(headers).unwrap() +} + +async fn echo(stream: DynTcpStream, path: &str, host: &str) { + let mut ws = + tokio_tungstenite::accept_hdr_async(stream, |request: &Request, response: Response| { + assert_eq!(request.uri().to_string(), path); + assert_eq!(request.headers()["host"], host); + assert!(!request.headers().contains_key("proxy-authorization")); + Ok(response) + }) + .await + .unwrap(); + let message = ws.next().await.unwrap().unwrap(); + assert!(message.is_binary()); + assert_eq!(message.into_data().as_ref(), PAYLOAD); + ws.send(PAYLOAD.to_vec().into()).await.unwrap(); + ws.close(None).await.unwrap(); +} + +async fn exchange(mut stream: Stream) { + stream.send_bytes(PAYLOAD.to_vec().into()).await.unwrap(); + let reply = stream.next_timeout(TIMEOUT).await.unwrap().unwrap(); + assert_eq!(reply.as_ref(), PAYLOAD); + if matches!(&stream, Stream::WebSocket(_)) { + assert!(timeout(Duration::from_millis(TIMEOUT), stream.next()) + .await + .unwrap() + .is_none()); + } +} + +fn option(key: &str, value: &str) { + OVERWRITE_SETTINGS + .write() + .unwrap() + .insert(key.into(), value.into()); +} + +async fn proxy_round_trip( + scheme: &str, + url: &str, + tls_type: TlsType, + auth: bool, + configured: bool, +) { + option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "Y"); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = Socks5Server { + proxy: format!("{scheme}://{}", listener.local_addr().unwrap()), + username: if auth { "user".into() } else { String::new() }, + password: if auth { "pass".into() } else { String::new() }, + }; + upsert_tls_cache(&proxy.proxy, TlsType::Rustls, true); + upsert_tls_cache(url, tls_type, true); + let target = url::Url::parse(url).unwrap(); + let host = target.host_str().unwrap().to_owned(); + let port = target.port_or_known_default().unwrap(); + let authority = match target.port() { + Some(port) => format!("{host}:{port}"), + None => host.clone(), + }; + let path = format!( + "{}{}", + target.path(), + target.query().map(|q| format!("?{q}")).unwrap_or_default() + ); + let wss = target.scheme() == "wss"; + let scheme = scheme.to_owned(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut stream = DynTcpStream(Box::new(stream)); + if scheme == "https" { + stream = DynTcpStream(Box::new(tls_acceptor().accept(stream).await.unwrap())); + } + if scheme == "socks5" { + let mut greeting = [0; 3]; + stream.read_exact(&mut greeting).await.unwrap(); + assert_eq!(greeting, [5, 1, 0]); + stream.write_all(&[5, 0]).await.unwrap(); + let mut request = [0; 5]; + stream.read_exact(&mut request).await.unwrap(); + assert_eq!(&request[..4], &[5, 1, 0, 3]); + let mut name = vec![0; request[4] as usize]; + stream.read_exact(&mut name).await.unwrap(); + assert_eq!(name, host.as_bytes()); + assert_eq!(stream.read_u16().await.unwrap(), port); + stream + .write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 0]) + .await + .unwrap(); + } else { + let headers = read_headers(&mut stream).await; + assert!(headers.starts_with(&format!("CONNECT {host}:{port} HTTP/1.1\r\n"))); + assert!(headers.contains(&format!("Host: {host}:{port}\r\n"))); + assert_eq!( + headers.contains("Proxy-Authorization: Basic dXNlcjpwYXNz\r\n"), + auth + ); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .unwrap(); + } + stream.flush().await.unwrap(); + if wss { + let tls = tls_acceptor().accept(stream).await.unwrap(); + assert_eq!(tls.get_ref().1.server_name(), Some(host.as_str())); + stream = DynTcpStream(Box::new(tls)); + } + echo(stream, &path, &authority).await; + }); + option(keys::OPTION_PROXY_URL, &proxy.proxy); + option(keys::OPTION_PROXY_USERNAME, &proxy.username); + option(keys::OPTION_PROXY_PASSWORD, &proxy.password); + let stream = if configured { + option(keys::OPTION_ALLOW_WEBSOCKET, "Y"); + option("custom-rendezvous-server", target.host_str().unwrap()); + option( + "api-server", + if wss { + "https://websocket-target.invalid" + } else { + "http://websocket-target.invalid" + }, + ); + socket_client::connect_tcp("websocket-target.invalid:21116", TIMEOUT) + .await + .unwrap() + } else { + WsFramedStream::new(url, None, TIMEOUT) + .await + .map(Stream::WebSocket) + .unwrap() + }; + exchange(stream).await; + timeout(Duration::from_millis(TIMEOUT), server) + .await + .unwrap() + .unwrap(); +} + +async fn rejected_or_stalled_proxy(stalled: bool, configured: bool) { + let direct = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = Socks5Server { + proxy: format!("http://{}", listener.local_addr().unwrap()), + ..Default::default() + }; + option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "Y"); + option(keys::OPTION_PROXY_URL, &proxy.proxy); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut stream = DynTcpStream(Box::new(stream)); + read_headers(&mut stream).await; + if !stalled { + stream + .write_all(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n") + .await + .unwrap(); + } + let mut byte = [0]; + assert_eq!(stream.read(&mut byte).await.unwrap(), 0); + }); + let result = timeout(Duration::from_millis(TIMEOUT), async { + let url = format!("wss://{}/ws/id", direct.local_addr().unwrap()); + if configured { + socket_client::connect_tcp(url, 200).await + } else { + WsFramedStream::new(url, None, 200) + .await + .map(Stream::WebSocket) + } + }) + .await + .unwrap(); + let error = result.err().expect("the proxy must fail"); + if stalled { + assert!(error + .downcast_ref::() + .is_some()); + } else { + assert!(error.to_string().contains("407"), "{}", error); + } + assert!(timeout(Duration::from_millis(50), direct.accept()) + .await + .is_err()); + timeout(Duration::from_millis(TIMEOUT), server) + .await + .unwrap() + .unwrap(); +} + +async fn direct_round_trip(setting: &str, configured: bool, wss: bool) { + option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, setting); + let unused_proxy = TcpListener::bind("127.0.0.1:0").await.unwrap(); + if configured { + option( + keys::OPTION_PROXY_URL, + &format!("http://{}", unused_proxy.local_addr().unwrap()), + ); + } else { + OVERWRITE_SETTINGS + .write() + .unwrap() + .remove(keys::OPTION_PROXY_URL); + } + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let url = format!("{}://{addr}/ws/id", if wss { "wss" } else { "ws" }); + upsert_tls_cache(&url, TlsType::Rustls, true); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut stream = DynTcpStream(Box::new(stream)); + if wss { + stream = DynTcpStream(Box::new(tls_acceptor().accept(stream).await.unwrap())); + } + echo(stream, "/ws/id", &addr.to_string()).await; + }); + let stream = socket_client::connect_tcp(url, TIMEOUT).await.unwrap(); + exchange(stream).await; + timeout(Duration::from_millis(TIMEOUT), server) + .await + .unwrap() + .unwrap(); + assert!(timeout(Duration::from_millis(50), unused_proxy.accept()) + .await + .is_err()); +} + +async fn tcp_proxy_round_trip() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = Socks5Server { + proxy: format!("http://{}", listener.local_addr().unwrap()), + ..Default::default() + }; + option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "N"); + option(keys::OPTION_PROXY_URL, &proxy.proxy); + option(keys::OPTION_PROXY_USERNAME, ""); + option(keys::OPTION_PROXY_PASSWORD, ""); + let server = tokio::spawn(async move { + let (stream, addr) = listener.accept().await.unwrap(); + let mut stream = DynTcpStream(Box::new(stream)); + let headers = read_headers(&mut stream).await; + assert!(headers.starts_with("CONNECT tcp-target.invalid:21116 HTTP/1.1\r\n")); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .unwrap(); + let mut stream = hbb_common::tcp::FramedStream::from(stream, addr); + let message = stream.next().await.unwrap().unwrap(); + assert_eq!(message.as_ref(), PAYLOAD); + stream.send_bytes(message.freeze()).await.unwrap(); + }); + let stream = socket_client::connect_tcp_local("tcp-target.invalid:21116", None, TIMEOUT) + .await + .unwrap(); + exchange(stream).await; + timeout(Duration::from_millis(TIMEOUT), server) + .await + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn websocket_proxy_connections() { + // Keep configuration changes inside this integration test process. + *APP_NAME.write().unwrap() = format!("RustDesk WebSocket Test {}", std::process::id()); + option(keys::OPTION_ALLOW_INSECURE_TLS_FALLBACK, "Y"); + for (scheme, url, tls_type, auth, configured) in [ + ( + "http", + "ws://websocket-target.invalid/ws/id", + TlsType::Plain, + false, + true, + ), + ( + "http", + "wss://websocket-target.invalid/ws/id", + TlsType::Rustls, + true, + true, + ), + ( + "http", + "wss://websocket-target.invalid:8443/ws/relay?test=1", + TlsType::NativeTls, + true, + false, + ), + ( + "https", + "ws://websocket-target.invalid/ws/id", + TlsType::Plain, + true, + false, + ), + ( + "https", + "wss://websocket-target.invalid/ws/relay", + TlsType::Rustls, + true, + false, + ), + ( + "https", + "wss://websocket-target.invalid/ws/relay", + TlsType::NativeTls, + true, + false, + ), + ( + "socks5", + "ws://websocket-target.invalid/ws/id", + TlsType::Plain, + false, + false, + ), + ( + "socks5", + "wss://websocket-target.invalid/ws/relay", + TlsType::Rustls, + false, + false, + ), + ( + "socks5", + "wss://websocket-target.invalid/ws/relay", + TlsType::NativeTls, + false, + false, + ), + ] { + proxy_round_trip(scheme, url, tls_type, auth, configured).await; + } + for scheme in ["http", "https", "socks5"] { + for (url, tls_type) in [ + ( + "ws://websocket-target.invalid:8080/ws/id?test=1", + TlsType::Plain, + ), + ( + "wss://websocket-target.invalid:8443/ws/relay?test=1", + TlsType::Rustls, + ), + ( + "wss://websocket-target.invalid:8443/ws/relay?test=1", + TlsType::NativeTls, + ), + ] { + proxy_round_trip(scheme, url, tls_type, scheme != "socks5", false).await; + } + } + for configured in [false, true] { + rejected_or_stalled_proxy(false, configured).await; + rejected_or_stalled_proxy(true, configured).await; + } + tcp_proxy_round_trip().await; + + for wss in [false, true] { + direct_round_trip("", true, wss).await; + direct_round_trip("N", true, wss).await; + direct_round_trip("Y", false, wss).await; + } + + option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "N"); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + drop(listener); + for scheme in ["wss", "ws"] { + let url = format!("{scheme}://{addr}/ws/id"); + let error = WsFramedStream::new(url, None, TIMEOUT).await.err().unwrap(); + assert!(error.to_string().starts_with("WebSocket error:")); + } +} From 9ee2f3ef9a0676737b023be7487d8415543c60f0 Mon Sep 17 00:00:00 2001 From: 21pages Date: Sat, 10 Oct 2026 19:45:51 +0800 Subject: [PATCH 2/7] fix: bracket IPv6 targets in HTTP CONNECT requests --- src/proxy.rs | 4 ++++ tests/websocket_proxy.rs | 8 +++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/proxy.rs b/src/proxy.rs index d3b8a76d2..538d5eaab 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -663,6 +663,10 @@ impl Proxy { fn get_domain_and_port<'a>(target_addr: &TargetAddr<'a>) -> Result<(String, u16), ProxyError> { match target_addr { + // CONNECT requires brackets around IPv6 literals in both the target and Host header. + tokio_socks::TargetAddr::Ip(SocketAddr::V6(addr)) => { + Ok((format!("[{}]", addr.ip()), addr.port())) + } tokio_socks::TargetAddr::Ip(addr) => Ok((addr.ip().to_string(), addr.port())), tokio_socks::TargetAddr::Domain(name, port) => Ok((name.to_string(), *port)), } diff --git a/tests/websocket_proxy.rs b/tests/websocket_proxy.rs index 3d6d09edb..18c00b619 100644 --- a/tests/websocket_proxy.rs +++ b/tests/websocket_proxy.rs @@ -98,6 +98,7 @@ async fn proxy_round_trip( upsert_tls_cache(url, tls_type, true); let target = url::Url::parse(url).unwrap(); let host = target.host_str().unwrap().to_owned(); + let domain = target.domain().map(str::to_owned); let port = target.port_or_known_default().unwrap(); let authority = match target.port() { Some(port) => format!("{host}:{port}"), @@ -148,7 +149,7 @@ async fn proxy_round_trip( stream.flush().await.unwrap(); if wss { let tls = tls_acceptor().accept(stream).await.unwrap(); - assert_eq!(tls.get_ref().1.server_name(), Some(host.as_str())); + assert_eq!(tls.get_ref().1.server_name(), domain.as_deref()); stream = DynTcpStream(Box::new(tls)); } echo(stream, &path, &authority).await; @@ -310,6 +311,11 @@ async fn websocket_proxy_connections() { // Keep configuration changes inside this integration test process. *APP_NAME.write().unwrap() = format!("RustDesk WebSocket Test {}", std::process::id()); option(keys::OPTION_ALLOW_INSECURE_TLS_FALLBACK, "Y"); + for scheme in ["http", "https"] { + for url in ["ws://[::1]:21119", "wss://[::1]:21119"] { + proxy_round_trip(scheme, url, TlsType::Rustls, false, false).await; + } + } for (scheme, url, tls_type, auth, configured) in [ ( "http", From 19b8b2242d9a4b3f87717f5d81307fbdea7c9eb2 Mon Sep 17 00:00:00 2001 From: 21pages Date: Mon, 12 Oct 2026 18:58:03 +0800 Subject: [PATCH 3/7] fix: always use configured proxy for WebSocket --- src/config.rs | 6 ------ src/websocket.rs | 9 ++------- tests/websocket_proxy.rs | 30 ++++++------------------------ 3 files changed, 8 insertions(+), 37 deletions(-) diff --git a/src/config.rs b/src/config.rs index 1e25a0df5..6fe100455 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2842,11 +2842,6 @@ pub fn use_ws() -> bool { option2bool(option, &Config::get_option(option)) } -pub fn ws_use_proxy() -> bool { - let option = keys::OPTION_ALLOW_WEBSOCKET_PROXY; - option2bool(option, &Config::get_option(option)) -} - pub fn allow_insecure_tls_fallback() -> bool { let option = keys::OPTION_ALLOW_INSECURE_TLS_FALLBACK; option2bool(option, &Config::get_option(option)) @@ -2874,7 +2869,6 @@ pub mod keys { pub const OPTION_ALLOW_NUMERNIC_ONE_TIME_PASSWORD: &str = "allow-numeric-one-time-password"; pub const OPTION_DIRECT_SERVER: &str = "direct-server"; pub const OPTION_ALLOW_WEBSOCKET: &str = "allow-websocket"; - pub const OPTION_ALLOW_WEBSOCKET_PROXY: &str = "allow-websocket-proxy"; pub const OPTION_TRACKPAD_SPEED: &str = "trackpad-speed"; pub const OPTION_REGISTER_DEVICE: &str = "register-device"; pub const OPTION_RELAY_SERVER: &str = "relay-server"; diff --git a/src/websocket.rs b/src/websocket.rs index 3cf0b0b92..2af824fb1 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -1,7 +1,6 @@ use crate::{ config::{ - keys::OPTION_RELAY_SERVER, use_ws, ws_use_proxy, Config, Socks5Server, RELAY_PORT, - RENDEZVOUS_PORT, + keys::OPTION_RELAY_SERVER, use_ws, Config, Socks5Server, RELAY_PORT, RENDEZVOUS_PORT, }, protobuf::Message, socket_client::split_host_port, @@ -299,11 +298,7 @@ impl WsFramedStream { local_addr: Option, ms_timeout: u64, ) -> ResultType { - let proxy_conf = if ws_use_proxy() { - Config::get_socks() - } else { - None - }; + let proxy_conf = Config::get_socks(); let (stream, addr) = Self::connect(url.as_ref(), local_addr, proxy_conf.as_ref(), ms_timeout).await?; diff --git a/tests/websocket_proxy.rs b/tests/websocket_proxy.rs index 18c00b619..5d4a15ab9 100644 --- a/tests/websocket_proxy.rs +++ b/tests/websocket_proxy.rs @@ -87,7 +87,6 @@ async fn proxy_round_trip( auth: bool, configured: bool, ) { - option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "Y"); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let proxy = Socks5Server { proxy: format!("{scheme}://{}", listener.local_addr().unwrap()), @@ -191,7 +190,6 @@ async fn rejected_or_stalled_proxy(stalled: bool, configured: bool) { proxy: format!("http://{}", listener.local_addr().unwrap()), ..Default::default() }; - option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "Y"); option(keys::OPTION_PROXY_URL, &proxy.proxy); let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); @@ -235,20 +233,11 @@ async fn rejected_or_stalled_proxy(stalled: bool, configured: bool) { .unwrap(); } -async fn direct_round_trip(setting: &str, configured: bool, wss: bool) { - option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, setting); - let unused_proxy = TcpListener::bind("127.0.0.1:0").await.unwrap(); - if configured { - option( - keys::OPTION_PROXY_URL, - &format!("http://{}", unused_proxy.local_addr().unwrap()), - ); - } else { - OVERWRITE_SETTINGS - .write() - .unwrap() - .remove(keys::OPTION_PROXY_URL); - } +async fn direct_round_trip(wss: bool) { + OVERWRITE_SETTINGS + .write() + .unwrap() + .remove(keys::OPTION_PROXY_URL); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let url = format!("{}://{addr}/ws/id", if wss { "wss" } else { "ws" }); @@ -267,9 +256,6 @@ async fn direct_round_trip(setting: &str, configured: bool, wss: bool) { .await .unwrap() .unwrap(); - assert!(timeout(Duration::from_millis(50), unused_proxy.accept()) - .await - .is_err()); } async fn tcp_proxy_round_trip() { @@ -278,7 +264,6 @@ async fn tcp_proxy_round_trip() { proxy: format!("http://{}", listener.local_addr().unwrap()), ..Default::default() }; - option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "N"); option(keys::OPTION_PROXY_URL, &proxy.proxy); option(keys::OPTION_PROXY_USERNAME, ""); option(keys::OPTION_PROXY_PASSWORD, ""); @@ -408,12 +393,9 @@ async fn websocket_proxy_connections() { tcp_proxy_round_trip().await; for wss in [false, true] { - direct_round_trip("", true, wss).await; - direct_round_trip("N", true, wss).await; - direct_round_trip("Y", false, wss).await; + direct_round_trip(wss).await; } - option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "N"); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); drop(listener); From 4c04a7a130a7b3796425dfdeb3fef7a1d594d6a7 Mon Sep 17 00:00:00 2001 From: 21pages Date: Mon, 12 Oct 2026 21:48:07 +0800 Subject: [PATCH 4/7] fix: allow proxy reads while flushing is pending --- src/websocket.rs | 10 ++- src/websocket/proxy_tests.rs | 149 +++++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 2 deletions(-) create mode 100644 src/websocket/proxy_tests.rs diff --git a/src/websocket.rs b/src/websocket.rs index 2af824fb1..dc3deedf5 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -18,7 +18,7 @@ use std::{ net::SocketAddr, pin::Pin, sync::Arc, - task::{ready, Context, Poll}, + task::{Context, Poll}, time::Duration, }; use tokio::{ @@ -47,7 +47,10 @@ impl AsyncRead for FlushBeforeRead { cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { - ready!(Pin::new(&mut self.0).poll_flush(cx))?; + // A blocked writer must not prevent reading data that lets the peer make progress. + if let Poll::Ready(Err(err)) = Pin::new(&mut self.0).poll_flush(cx) { + return Poll::Ready(Err(err)); + } Pin::new(&mut self.0).poll_read(cx, buf) } } @@ -525,6 +528,9 @@ pub fn check_ws(endpoint: &str) -> String { format!("{}://{}", protocol, address) } +#[cfg(test)] +mod proxy_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/src/websocket/proxy_tests.rs b/src/websocket/proxy_tests.rs new file mode 100644 index 000000000..06db80965 --- /dev/null +++ b/src/websocket/proxy_tests.rs @@ -0,0 +1,149 @@ +use super::FlushBeforeRead; +use crate::{proxy::Proxy, tcp::DynTcpStream, tls::TlsType}; +use futures::{SinkExt, StreamExt}; +use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer}; +use std::{sync::Arc, time::Duration}; +use tokio::{ + io::{duplex, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream}, + time::timeout, +}; +use tokio_native_tls::native_tls; +use tokio_rustls::{rustls, TlsAcceptor}; +use tokio_socks::IntoTargetAddr; +use tokio_tungstenite::{accept_async, client_async_tls_with_config, Connector, WebSocketStream}; + +fn acceptor() -> TlsAcceptor { + let cert = + CertificateDer::from_pem_slice(include_bytes!("../../tests/data/websocket-cert.pem")) + .unwrap(); + let key = PrivateKeyDer::from_pem_slice(include_bytes!("../../tests/data/websocket-key.pem")) + .unwrap(); + let config = rustls::ServerConfig::builder_with_provider(Arc::new( + rustls::crypto::ring::default_provider(), + )) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert(vec![cert], key) + .unwrap(); + TlsAcceptor::from(Arc::new(config)) +} + +async fn exchange( + ws: WebSocketStream, + sent: u8, + received: u8, +) { + let (mut writer, mut reader) = ws.split(); + let send = async { + for sequence in 0..16 { + writer + .send(vec![sent + sequence; 256 * 1024].into()) + .await + .unwrap(); + } + }; + let receive = async { + for sequence in 0..16 { + let message = reader.next().await.unwrap().unwrap(); + assert!(message.is_binary()); + assert_eq!( + message.into_data().as_ref(), + vec![received + sequence; 256 * 1024] + ); + } + }; + tokio::join!(send, receive); +} + +#[tokio::test] +async fn native_tls_wss_bidirectional_transfer_through_proxy_buffers() { + for outer_tls in [TlsType::Plain, TlsType::Rustls, TlsType::NativeTls] { + // A 4 KiB pipe forces backpressure while both sides send 4 MiB. + let (client, server) = duplex(4096); + let server = async { + let mut stream = DynTcpStream(Box::new(server)); + if !matches!(outer_tls, TlsType::Plain) { + stream = DynTcpStream(Box::new(acceptor().accept(stream).await.unwrap())); + } + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + assert!(headers.len() < 4096); + headers.push(stream.read_u8().await.unwrap()); + } + assert!(headers.starts_with(b"CONNECT target.invalid:443 HTTP/1.1\r\n")); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .unwrap(); + stream.flush().await.unwrap(); + let stream = acceptor().accept(stream).await.unwrap(); + exchange(accept_async(stream).await.unwrap(), 64, 0).await; + }; + let client = async { + let proxy = Proxy::new("https://proxy.invalid:4443", 5_000).unwrap(); + let target = "target.invalid:443".into_target_addr().unwrap(); + let stream = match outer_tls { + TlsType::Plain => { + DynTcpStream(Box::new(proxy.http_connect(client, &target).await.unwrap())) + } + TlsType::Rustls => DynTcpStream(Box::new( + proxy + .https_connect_rustls(client, &target, true) + .await + .unwrap(), + )), + TlsType::NativeTls => DynTcpStream(Box::new( + proxy + .https_connect_nativetls(client, &target, true) + .await + .unwrap(), + )), + }; + let connector = native_tls::TlsConnector::builder() + .danger_accept_invalid_certs(true) + .build() + .unwrap(); + let (ws, _) = client_async_tls_with_config( + "wss://target.invalid/ws/relay", + FlushBeforeRead(stream), + None, + Some(Connector::NativeTls(connector)), + ) + .await + .unwrap(); + exchange(ws, 0, 64).await; + }; + timeout(Duration::from_secs(20), async { + tokio::join!(client, server) + }) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn read_progresses_while_proxy_flush_is_blocked() { + let (client, mut peer) = duplex(64); + let mut client = FlushBeforeRead(DynTcpStream(Box::new(BufStream::new(client)))); + let payload = [0x5a; 4096]; + client.write_all(&payload).await.unwrap(); + peer.write_all(b"pong").await.unwrap(); + + let mut reply = [0; 4]; + timeout(Duration::from_secs(1), client.read_exact(&mut reply)) + .await + .expect("a blocked flush must not prevent receiving queued data") + .unwrap(); + assert_eq!(&reply, b"pong"); + + let mut received = [0; 4096]; + timeout(Duration::from_secs(1), async { + let (flushed, read) = tokio::join!(client.flush(), peer.read_exact(&mut received)); + flushed.unwrap(); + read.unwrap(); + }) + .await + .unwrap(); + assert_eq!(received, payload); +} From 2d998e7b02ee8c83d772b85609e60e98d15ccf3a Mon Sep 17 00:00:00 2001 From: 21pages Date: Mon, 12 Oct 2026 21:40:28 +0800 Subject: [PATCH 5/7] fix: reuse the initial TCP connection for HTTPS proxies --- src/proxy.rs | 15 ++-- tests/proxy_transport.rs | 151 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 158 insertions(+), 8 deletions(-) create mode 100644 tests/proxy_transport.rs diff --git a/src/proxy.rs b/src/proxy.rs index 538d5eaab..2e67e0328 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -428,8 +428,7 @@ impl Proxy { TlsType::NativeTls => { self.https_connect_nativetls_wrap_danger( &url, - local, - proxy, + stream, &target_addr, danger_accept_invalid_cert, ) @@ -480,12 +479,10 @@ impl Proxy { async fn https_connect_nativetls_wrap_danger<'a>( &self, url: &str, - local: SocketAddr, - proxy: SocketAddr, + stream: tokio::net::TcpStream, target_addr: &TargetAddr<'a>, danger_accept_invalid_cert: Option, ) -> ResultType { - let stream = self.new_stream(local, proxy).await?; let s = super::timeout( self.ms_timeout, self.https_connect_nativetls( @@ -535,7 +532,10 @@ impl Proxy { danger_accept_invalid_cert: Option, origin_danger_accept_invalid_cert: Option, ) -> ResultType { - let stream = stream.unwrap_or(self.new_stream(local, proxy).await?); + let stream = match stream { + Some(stream) => stream, + None => self.new_stream(local, proxy).await?, + }; match super::timeout( self.ms_timeout, self.https_connect_rustls( @@ -587,8 +587,7 @@ impl Proxy { log::warn!("Falling back to native-tls for HTTPS proxy server."); self.https_connect_nativetls_wrap_danger( &url, - local, - proxy, + self.new_stream(local, proxy).await?, &target_addr, origin_danger_accept_invalid_cert, ) diff --git a/tests/proxy_transport.rs b/tests/proxy_transport.rs new file mode 100644 index 000000000..67e9f979d --- /dev/null +++ b/tests/proxy_transport.rs @@ -0,0 +1,151 @@ +use hbb_common::{ + config::{APP_NAME, OVERWRITE_SETTINGS}, + proxy::{Proxy, ProxyError}, + tls::{upsert_tls_cache, TlsType}, +}; +use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer}; +use std::{sync::Arc, time::Duration}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + time::timeout, +}; +use tokio_rustls::{rustls, TlsAcceptor}; + +const TIMEOUT: u64 = 3_000; + +fn tls_acceptor() -> TlsAcceptor { + let cert = CertificateDer::from_pem_slice(include_bytes!("data/websocket-cert.pem")).unwrap(); + let key = PrivateKeyDer::from_pem_slice(include_bytes!("data/websocket-key.pem")).unwrap(); + let config = rustls::ServerConfig::builder_with_provider(Arc::new( + rustls::crypto::ring::default_provider(), + )) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert(vec![cert], key) + .unwrap(); + TlsAcceptor::from(Arc::new(config)) +} + +async fn https_connection(tls_type: TlsType, fixed_local_port: bool) { + *APP_NAME.write().unwrap() = format!("RustDesk-proxy-review-{}", std::process::id()); + OVERWRITE_SETTINGS + .write() + .unwrap() + .insert("allow-insecure-tls-fallback".into(), "Y".into()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("https://{}", listener.local_addr().unwrap()); + upsert_tls_cache(&url, tls_type, true); + let local = if fixed_local_port { + let reserved = TcpListener::bind("127.0.0.1:0").await.unwrap(); + Some(reserved.local_addr().unwrap()) + } else { + None + }; + let server = tokio::spawn(async move { + let (tcp, peer) = listener.accept().await.unwrap(); + let mut stream = tls_acceptor().accept(tcp).await.unwrap(); + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + assert!(headers.len() < 4096); + headers.push(stream.read_u8().await.unwrap()); + } + assert!(headers.starts_with(b"CONNECT relay.invalid:21117 HTTP/1.1\r\n")); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .unwrap(); + stream.flush().await.unwrap(); + let mut data = [0; 4]; + stream.read_exact(&mut data).await.unwrap(); + assert_eq!(&data, b"ping"); + stream.write_all(b"pong").await.unwrap(); + stream.flush().await.unwrap(); + assert!( + timeout(Duration::from_millis(100), listener.accept()) + .await + .is_err(), + "an unused second TCP connection was opened" + ); + peer + }); + let proxy = Proxy::new(url, TIMEOUT).unwrap(); + let framed = proxy.connect("relay.invalid:21117", local).await.unwrap(); + let actual_local = framed.1; + if let Some(local) = local { + assert_eq!(actual_local, local); + } + let mut stream = framed.0.into_inner(); + stream.write_all(b"ping").await.unwrap(); + stream.flush().await.unwrap(); + let mut reply = [0; 4]; + timeout( + Duration::from_millis(TIMEOUT), + stream.read_exact(&mut reply), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(&reply, b"pong"); + assert_eq!(server.await.unwrap(), actual_local); +} + +#[tokio::test] +async fn https_rustls_reuses_initial_connection() { + https_connection(TlsType::Rustls, false).await; +} + +#[tokio::test] +async fn https_nativetls_reuses_initial_connection() { + https_connection(TlsType::NativeTls, false).await; +} + +#[tokio::test] +async fn https_rustls_with_fixed_local_port() { + https_connection(TlsType::Rustls, true).await; +} + +#[tokio::test] +async fn https_nativetls_with_fixed_local_port() { + https_connection(TlsType::NativeTls, true).await; +} + +#[tokio::test] +async fn https_reconnects_for_tls_fallback() { + *APP_NAME.write().unwrap() = format!("RustDesk-proxy-review-{}", std::process::id()); + OVERWRITE_SETTINGS + .write() + .unwrap() + .insert("allow-insecure-tls-fallback".into(), "Y".into()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("https://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + // Reject both Rustls attempts so the uncached route must reach NativeTls. + for _ in 0..2 { + let (tcp, _) = listener.accept().await.unwrap(); + drop(tcp); + } + let (tcp, _) = listener.accept().await.unwrap(); + assert!(tls_acceptor().accept(tcp).await.is_err()); + }); + let proxy = Proxy::new(&url, TIMEOUT).unwrap(); + let error = proxy + .connect("relay.invalid:21117", None) + .await + .err() + .unwrap(); + // The third connection reaches NativeTls and still rejects the untrusted identity. + assert!( + matches!( + error.downcast_ref::(), + Some(ProxyError::NativeTlsError(_)) + ), + "{}", + error + ); + timeout(Duration::from_millis(TIMEOUT), server) + .await + .unwrap() + .unwrap(); +} From 0ccaaf2cf6906417c76527a55d3d36b28a0ddf46 Mon Sep 17 00:00:00 2001 From: 21pages Date: Mon, 12 Oct 2026 21:42:08 +0800 Subject: [PATCH 6/7] fix: parse IPv6 HTTP proxy authorities correctly --- src/proxy.rs | 151 +++++++++++++++++++++++++++++++++------ tests/proxy_transport.rs | 69 +++++++++++++++--- 2 files changed, 189 insertions(+), 31 deletions(-) diff --git a/src/proxy.rs b/src/proxy.rs index 2e67e0328..a31ea6f10 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -1,7 +1,4 @@ -use std::{ - io::Error as IoError, - net::{SocketAddr, ToSocketAddrs}, -}; +use std::{io::Error as IoError, net::SocketAddr}; use anyhow::bail; use async_recursion::async_recursion; @@ -234,12 +231,10 @@ impl ProxyScheme { } async fn resolve_host(&self, host: &str, default_port: u16) -> Result { - let (host_str, port) = match host.split_once(':') { - Some((h, p)) => (h, p.parse::().ok()), - None => (host, None), - }; - let addr = (host_str, port.unwrap_or(default_port)) - .to_socket_addrs()? + let addr = self + .http_url(host)? + .socket_addrs(|| Some(default_port))? + .into_iter() .next() .ok_or_else(|| ProxyError::AddressResolutionFailed(host.to_string()))?; Ok(addr) @@ -248,11 +243,11 @@ impl ProxyScheme { pub fn get_domain(&self) -> Result { match self { ProxyScheme::Http { host, .. } | ProxyScheme::Https { host, .. } => { - let domain = host - .split(':') - .next() - .ok_or_else(|| ProxyError::AddressResolutionFailed(host.clone()))?; - Ok(domain.to_string()) + match self.http_url(host)?.host() { + Some(url::Host::Ipv6(ip)) => Ok(ip.to_string()), + Some(host) => Ok(host.to_string()), + None => Err(ProxyError::AddressResolutionFailed(host.clone())), + } } ProxyScheme::Socks5 { addr, .. } => match addr { SocketAddr::V4(addr_v4) => Ok(addr_v4.ip().to_string()), @@ -262,17 +257,30 @@ impl ProxyScheme { } pub fn get_host_and_port(&self) -> Result { match self { - ProxyScheme::Http { host, .. } => Ok(self.append_default_port(host, 80)), - ProxyScheme::Https { host, .. } => Ok(self.append_default_port(host, 443)), + ProxyScheme::Http { host, .. } => self.append_default_port(host, 80), + ProxyScheme::Https { host, .. } => self.append_default_port(host, 443), ProxyScheme::Socks5 { addr, .. } => Ok(format!("{}", addr)), } } - fn append_default_port(&self, host: &str, default_port: u16) -> String { - if host.contains(':') { - host.to_string() + fn append_default_port(&self, host: &str, default_port: u16) -> Result { + let url = self.http_url(host)?; + let host = url + .host_str() + .ok_or_else(|| ProxyError::AddressResolutionFailed(host.to_string()))?; + Ok(format!( + "{}:{}", + host, + url.port_or_known_default().unwrap_or(default_port) + )) + } + + fn http_url(&self, host: &str) -> Result { + let scheme = if matches!(self, Self::Https { .. }) { + "https" } else { - format!("{}:{}", host, default_port) - } + "http" + }; + Ok(Url::parse(&format!("{scheme}://{host}"))?) } } @@ -717,3 +725,102 @@ where None => Err(ProxyError::NoHttpCode), }; } + +#[cfg(test)] +mod tests { + use super::ProxyScheme; + + #[test] + fn http_url_parses_hosts_and_ports() { + for host in ["localhost", "127.0.0.1", "[::1]"] { + for (proxy, scheme, default_port) in [ + (ProxyScheme::http(host).unwrap(), "http", 80), + (ProxyScheme::https(host).unwrap(), "https", 443), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + let url = proxy.http_url(&authority).unwrap(); + assert_eq!(url.scheme(), scheme); + assert_eq!(url.host_str(), Some(host)); + assert_eq!( + url.port_or_known_default(), + Some(port.unwrap_or(default_port)) + ); + } + } + } + } + + #[test] + fn append_default_port_formats_hosts_and_ports() { + for host in ["localhost", "127.0.0.1", "[::1]"] { + for (proxy, default_port) in [ + (ProxyScheme::http(host).unwrap(), 80), + (ProxyScheme::https(host).unwrap(), 443), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + assert_eq!( + proxy.append_default_port(&authority, default_port).unwrap(), + format!("{host}:{}", port.unwrap_or(default_port)) + ); + } + } + } + } + + #[tokio::test] + async fn resolve_host_resolves_hosts_and_ports() { + for (host, expected_ip) in [ + ("localhost", None), + ("127.0.0.1", Some("127.0.0.1")), + ("[::1]", Some("::1")), + ] { + for (proxy, default_port) in [ + (ProxyScheme::http(host).unwrap(), 80), + (ProxyScheme::https(host).unwrap(), 443), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + let addr = proxy.resolve_host(&authority, default_port).await.unwrap(); + assert_eq!(addr.port(), port.unwrap_or(default_port)); + assert!(addr.ip().is_loopback()); + if let Some(ip) = expected_ip { + assert_eq!(addr.ip().to_string(), ip); + } + } + } + } + } + + #[test] + fn get_domain_strips_ports_and_ipv6_brackets() { + for (host, domain) in [ + ("localhost", "localhost"), + ("127.0.0.1", "127.0.0.1"), + ("[::1]", "::1"), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + for proxy in [ + ProxyScheme::http(&authority).unwrap(), + ProxyScheme::https(&authority).unwrap(), + ] { + assert_eq!(proxy.get_domain().unwrap(), domain); + } + } + } + } +} diff --git a/tests/proxy_transport.rs b/tests/proxy_transport.rs index 67e9f979d..2faedf023 100644 --- a/tests/proxy_transport.rs +++ b/tests/proxy_transport.rs @@ -1,6 +1,7 @@ use hbb_common::{ config::{APP_NAME, OVERWRITE_SETTINGS}, proxy::{Proxy, ProxyError}, + tcp::DynTcpStream, tls::{upsert_tls_cache, TlsType}, }; use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer}; @@ -28,24 +29,32 @@ fn tls_acceptor() -> TlsAcceptor { TlsAcceptor::from(Arc::new(config)) } -async fn https_connection(tls_type: TlsType, fixed_local_port: bool) { +async fn proxy_connection(bind: &str, tls_type: TlsType, fixed_local_port: bool) { *APP_NAME.write().unwrap() = format!("RustDesk-proxy-review-{}", std::process::id()); OVERWRITE_SETTINGS .write() .unwrap() .insert("allow-insecure-tls-fallback".into(), "Y".into()); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let url = format!("https://{}", listener.local_addr().unwrap()); + let listener = TcpListener::bind(bind).await.unwrap(); + let scheme = if matches!(tls_type, TlsType::Plain) { + "http" + } else { + "https" + }; + let url = format!("{scheme}://{}", listener.local_addr().unwrap()); upsert_tls_cache(&url, tls_type, true); let local = if fixed_local_port { - let reserved = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let reserved = TcpListener::bind(bind).await.unwrap(); Some(reserved.local_addr().unwrap()) } else { None }; let server = tokio::spawn(async move { let (tcp, peer) = listener.accept().await.unwrap(); - let mut stream = tls_acceptor().accept(tcp).await.unwrap(); + let mut stream = DynTcpStream(Box::new(tcp)); + if !matches!(tls_type, TlsType::Plain) { + stream = DynTcpStream(Box::new(tls_acceptor().accept(stream).await.unwrap())); + } let mut headers = Vec::new(); while !headers.ends_with(b"\r\n\r\n") { assert!(headers.len() < 4096); @@ -93,22 +102,22 @@ async fn https_connection(tls_type: TlsType, fixed_local_port: bool) { #[tokio::test] async fn https_rustls_reuses_initial_connection() { - https_connection(TlsType::Rustls, false).await; + proxy_connection("127.0.0.1:0", TlsType::Rustls, false).await; } #[tokio::test] async fn https_nativetls_reuses_initial_connection() { - https_connection(TlsType::NativeTls, false).await; + proxy_connection("127.0.0.1:0", TlsType::NativeTls, false).await; } #[tokio::test] async fn https_rustls_with_fixed_local_port() { - https_connection(TlsType::Rustls, true).await; + proxy_connection("127.0.0.1:0", TlsType::Rustls, true).await; } #[tokio::test] async fn https_nativetls_with_fixed_local_port() { - https_connection(TlsType::NativeTls, true).await; + proxy_connection("127.0.0.1:0", TlsType::NativeTls, true).await; } #[tokio::test] @@ -149,3 +158,45 @@ async fn https_reconnects_for_tls_fallback() { .unwrap() .unwrap(); } + +#[tokio::test] +async fn proxy_addresses_and_tls_names() { + for (scheme, default_port) in [("http", 80), ("https", 443)] { + for (host, tls_name) in [ + ("[::1]", "::1"), + ("127.0.0.1", "127.0.0.1"), + ("localhost", "localhost"), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + let proxy = Proxy::new(format!("{scheme}://{authority}"), TIMEOUT).unwrap(); + let port = port.unwrap_or(default_port); + assert_eq!( + proxy.intercept.get_domain().unwrap(), + tls_name, + "{scheme}://{authority}" + ); + assert_eq!( + proxy.intercept.get_host_and_port().unwrap(), + format!("{host}:{port}") + ); + let resolved = proxy.proxy_addrs().await.unwrap(); + assert_eq!(resolved.port(), port); + assert!(resolved.ip().is_loopback()); + if host == "[::1]" { + assert!(resolved.is_ipv6()); + } + } + } + } +} + +#[tokio::test] +async fn ipv6_proxy_connect_and_data_exchange() { + for tls_type in [TlsType::Plain, TlsType::Rustls, TlsType::NativeTls] { + proxy_connection("[::1]:0", tls_type, false).await; + } +} From 423fb7108504635b661bb059e3f18105bc0ddf5b Mon Sep 17 00:00:00 2001 From: 21pages Date: Wed, 14 Oct 2026 20:04:48 +0800 Subject: [PATCH 7/7] fix: return local address for direct WebSocket connections --- src/websocket.rs | 4 ++-- tests/websocket_proxy.rs | 7 +++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/websocket.rs b/src/websocket.rs index dc3deedf5..af797acc1 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -213,8 +213,8 @@ impl WsFramedStream { match stream { Either::Left(stream) => { // MaybeTlsStream::get_ref() handles Plain, NativeTls and Rustls, - // returning the underlying TcpStream for peer_addr(). - let addr = stream.get_ref().get_ref().peer_addr()?; + // returning the underlying TcpStream for local_addr(). + let addr = stream.get_ref().get_ref().local_addr()?; Ok((Either::Left(stream), addr)) } Either::Right((stream, addr)) => Ok((Either::Right(stream), addr)), diff --git a/tests/websocket_proxy.rs b/tests/websocket_proxy.rs index 5d4a15ab9..2f3f29763 100644 --- a/tests/websocket_proxy.rs +++ b/tests/websocket_proxy.rs @@ -243,19 +243,22 @@ async fn direct_round_trip(wss: bool) { let url = format!("{}://{addr}/ws/id", if wss { "wss" } else { "ws" }); upsert_tls_cache(&url, TlsType::Rustls, true); let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); + let (stream, client_addr) = listener.accept().await.unwrap(); let mut stream = DynTcpStream(Box::new(stream)); if wss { stream = DynTcpStream(Box::new(tls_acceptor().accept(stream).await.unwrap())); } echo(stream, "/ws/id", &addr.to_string()).await; + client_addr }); let stream = socket_client::connect_tcp(url, TIMEOUT).await.unwrap(); + let local_addr = stream.local_addr(); exchange(stream).await; - timeout(Duration::from_millis(TIMEOUT), server) + let client_addr = timeout(Duration::from_millis(TIMEOUT), server) .await .unwrap() .unwrap(); + assert_eq!(local_addr, client_addr); } async fn tcp_proxy_round_trip() {