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:")); + } +}