From 3aade011c08cb1937c903deb48a711fb0e856958 Mon Sep 17 00:00:00 2001 From: 21pages Date: Mon, 28 Sep 2026 14:14:54 +0800 Subject: [PATCH] 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:")); + } +}