Merge pull request #600 from 21pages/websocket-proxy

websocket proxy support

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<SocketAddr, ProxyError> {
-        let (host_str, port) = match host.split_once(':') {
-            Some((h, p)) => (h, p.parse::<u16>().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<String, ProxyError> {
         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<String, ProxyError> {
         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<String, ProxyError> {
+        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<Url, ProxyError> {
+        let scheme = if matches!(self, Self::Https { .. }) {
+            "https"
         } else {
-            format!("{}:{}", host, default_port)
-        }
+            "http"
+        };
+        Ok(Url::parse(&format!("{scheme}://{host}"))?)
     }
 }
 
@@ -428,8 +436,7 @@ impl Proxy {
                     TlsType::NativeTls => {
                         self.https_connect_nativetls_wrap_danger(
                             &url,
-                            local,
-                            proxy,
+                            stream,
                             &target_addr,
                             danger_accept_invalid_cert,
                         )
@@ -480,12 +487,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<bool>,
     ) -> ResultType<DynTcpStream> {
-        let stream = self.new_stream(local, proxy).await?;
         let s = super::timeout(
             self.ms_timeout,
             self.https_connect_nativetls(
@@ -535,7 +540,10 @@ impl Proxy {
         danger_accept_invalid_cert: Option<bool>,
         origin_danger_accept_invalid_cert: Option<bool>,
     ) -> ResultType<DynTcpStream> {
-        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 +595,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,
                     )
@@ -663,6 +670,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)),
     }
@@ -714,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);
+                }
+            }
+        }
+    }
+}

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

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<str>,
-        local_addr: Option<SocketAddr>,
-        proxy_conf: Option<&config::Socks5Server>,
-        timeout_ms: u64,
-    ) -> ResultType<Self> {
-        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<()> {

src/websocket.rs

@@ -5,36 +5,82 @@ use crate::{
     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::{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<io::Result<()>> {
+        // 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)
+    }
+}
+
+impl AsyncWrite for FlushBeforeRead {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut self.0).poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_shutdown(cx)
+    }
+}
+
+// Keep the original non-proxy connection path while supporting proxy streams.
+type WsStream = Either<
+    WebSocketStream<MaybeTlsStream<TcpStream>>,
+    WebSocketStream<MaybeTlsStream<DynTcpStream>>,
+>;
+
 pub struct WsFramedStream {
-    stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
+    stream: WsStream,
     addr: SocketAddr,
     encrypt: Option<Encrypt>,
     send_timeout: u64,
@@ -72,16 +118,18 @@ impl WsFramedStream {
 
     async fn connect(
         url: &str,
+        local_addr: Option<SocketAddr>,
+        proxy_conf: Option<&Socks5Server>,
         ms_timeout: u64,
-    ) -> ResultType<WebSocketStream<MaybeTlsStream<TcpStream>>> {
-        // 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 +142,83 @@ impl WsFramedStream {
     #[async_recursion]
     async fn try_connect(
         url: &str,
+        local_addr: Option<SocketAddr>,
+        proxy_conf: Option<&Socks5Server>,
         ms_timeout: u64,
         tls_type: TlsType,
         is_tls_type_cached: bool,
         danger_accept_invalid_cert: Option<bool>,
         original_danger_accept_invalid_certs: Option<bool>,
-    ) -> ResultType<WebSocketStream<MaybeTlsStream<TcpStream>>> {
+    ) -> 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 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)),
+                }
             }
             Err(e) => match (tls_type, is_tls_type_cached, danger_accept_invalid_cert) {
                 (TlsType::Rustls, _, None) => {
@@ -126,6 +229,8 @@ impl WsFramedStream {
                         );
                     Self::try_connect(
                         url,
+                        local_addr,
+                        proxy_conf,
                         ms_timeout,
                         tls_type,
                         is_tls_type_cached,
@@ -142,6 +247,8 @@ impl WsFramedStream {
                     );
                     Self::try_connect(
                         url,
+                        local_addr,
+                        proxy_conf,
                         ms_timeout,
                         TlsType::NativeTls,
                         is_tls_type_cached,
@@ -158,6 +265,8 @@ impl WsFramedStream {
                         );
                     Self::try_connect(
                         url,
+                        local_addr,
+                        proxy_conf,
                         ms_timeout,
                         tls_type,
                         is_tls_type_cached,
@@ -189,17 +298,12 @@ impl WsFramedStream {
 
     pub async fn new<T: AsRef<str>>(
         url: T,
-        _local_addr: Option<SocketAddr>,
-        _proxy_conf: Option<&Socks5Server>,
+        local_addr: Option<SocketAddr>,
         ms_timeout: u64,
     ) -> ResultType<Self> {
-        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 = Config::get_socks();
+        let (stream, addr) =
+            Self::connect(url.as_ref(), local_addr, proxy_conf.as_ref(), ms_timeout).await?;
 
         let ws = Self {
             stream,
@@ -222,24 +326,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<Self> {
-        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]
@@ -434,11 +528,28 @@ pub fn check_ws(endpoint: &str) -> String {
     format!("{}://{}", protocol, address)
 }
 
+#[cfg(test)]
+mod proxy_tests;
+
 #[cfg(test)]
 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 +566,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 +740,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 +751,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 +771,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 +804,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!(

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<S: AsyncRead + AsyncWrite + Unpin>(
+    ws: WebSocketStream<S>,
+    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);
+}

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-----

tests/data/websocket-key.pem

@@ -0,0 +1,5 @@
+-----BEGIN PRIVATE KEY-----
+MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgfMmTjDYIQrnprR5R
+3NkH2woEGhwRtkkOSoFqi9yAeN6hRANCAATGz2rNvO/J5IZcR2tAP18dE/cDqaqh
+SI/GrG9wsRKRDypx6C1AlhHp67+xsBxfxGX8qP+2IYqxPeEfOaXTNmRu
+-----END PRIVATE KEY-----

tests/proxy_transport.rs

@@ -0,0 +1,202 @@
+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};
+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 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(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(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 = 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);
+            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() {
+    proxy_connection("127.0.0.1:0", TlsType::Rustls, false).await;
+}
+
+#[tokio::test]
+async fn https_nativetls_reuses_initial_connection() {
+    proxy_connection("127.0.0.1:0", TlsType::NativeTls, false).await;
+}
+
+#[tokio::test]
+async fn https_rustls_with_fixed_local_port() {
+    proxy_connection("127.0.0.1:0", TlsType::Rustls, true).await;
+}
+
+#[tokio::test]
+async fn https_nativetls_with_fixed_local_port() {
+    proxy_connection("127.0.0.1:0", 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::<ProxyError>(),
+            Some(ProxyError::NativeTlsError(_))
+        ),
+        "{}",
+        error
+    );
+    timeout(Duration::from_millis(TIMEOUT), server)
+        .await
+        .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;
+    }
+}

tests/websocket_proxy.rs

@@ -0,0 +1,410 @@
+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,
+) {
+    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 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}"),
+        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(), domain.as_deref());
+            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_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::<tokio::time::error::Elapsed>()
+            .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(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" });
+    upsert_tls_cache(&url, TlsType::Rustls, true);
+    let server = tokio::spawn(async move {
+        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;
+    let client_addr = timeout(Duration::from_millis(TIMEOUT), server)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(local_addr, client_addr);
+}
+
+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_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 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",
+            "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(wss).await;
+    }
+
+    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:"));
+    }
+}