nautilus_network/websocket/
proxy.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Proxy support for outbound WebSocket connections.
17//!
18//! Implements HTTP `CONNECT` tunneling so a `WebSocketClient` can be reached
19//! through an HTTP or HTTPS forward proxy. The same `proxy_url` field is used
20//! by the HTTP client (via `reqwest::Proxy::all`), keeping a single config
21//! field for both transports.
22//!
23//! `socks5://` / `socks5h://` URLs are recognized but not yet implemented
24//! for the WebSocket path. The dispatcher logs a warning and falls back to
25//! a direct connection so that REST configs that already point at a SOCKS
26//! proxy keep working unchanged. SOCKS support requires the optional
27//! `tokio-socks` crate, which is not yet a workspace dependency.
28//!
29//! The tunnel is established as follows:
30//! 1. TCP connect to the proxy host / port.
31//! 2. If the proxy URL scheme is `https`, layer TLS using the proxy host as
32//!    the SNI and certificate domain.
33//! 3. Send `CONNECT target_host:target_port HTTP/1.1` plus the matching
34//!    `Host:` header (and optional `Proxy-Authorization:` derived from the
35//!    proxy URL user-info).
36//! 4. Read the response line and headers; require a `2xx` status.
37//! 5. If the upstream WebSocket scheme is `wss`, layer a second TLS session
38//!    using the upstream host name.
39//! 6. Hand the resulting stream to `tokio-tungstenite`'s `client_async` so the
40//!    WebSocket handshake completes over the tunnel.
41
42use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
43use rustls::{ClientConfig, RootCertStore, pki_types::ServerName};
44use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
45use tokio_rustls::{TlsConnector, client::TlsStream};
46use url::Url;
47
48use crate::{net::TcpStream, transport::TransportError};
49
50/// Maximum size of a `CONNECT` proxy response we are willing to read.
51///
52/// Bounds the buffer so a malicious or broken proxy cannot make us allocate
53/// indefinitely while we wait for the header terminator.
54const MAX_PROXY_RESPONSE_BYTES: usize = 16 * 1024;
55
56/// Stream produced by `tunnel_via_proxy` when the upstream is `ws://`
57/// (no upstream TLS, but the proxy hop itself may have been TLS-protected).
58///
59/// The TLS-bearing variants are boxed because [`tokio_rustls::client::TlsStream`]
60/// is large enough that a flat enum trips `clippy::large_enum_variant`. Boxing
61/// keeps the discriminant cheap to move while leaving the rare TLS path on the
62/// heap.
63#[derive(Debug)]
64pub enum ProxiedStream {
65    /// Plain TCP after a plain proxy hop.
66    Plain(TcpStream),
67    /// Plain TCP after a TLS proxy hop.
68    PlainOverTlsProxy(Box<TlsStream<TcpStream>>),
69    /// Upstream TLS over a plain proxy hop.
70    Tls(Box<TlsStream<TcpStream>>),
71    /// Upstream TLS over a TLS proxy hop.
72    TlsOverTlsProxy(Box<TlsStream<TlsStream<TcpStream>>>),
73}
74
75/// Parsed components of a target WebSocket URL needed by the proxy hop.
76#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct WsTarget {
78    /// Host name for DNS / SNI / `CONNECT` request line.
79    pub host: String,
80    /// TCP port of the WebSocket origin.
81    pub port: u16,
82    /// `true` when the WebSocket scheme is `wss://`.
83    pub is_tls: bool,
84}
85
86impl WsTarget {
87    /// Parse a `ws://` or `wss://` URL into the host/port/TLS components.
88    ///
89    /// # Errors
90    ///
91    /// Returns [`TransportError::InvalidUrl`] when the URL fails to parse,
92    /// is missing a hostname, or uses a scheme other than `ws`/`wss`.
93    pub fn parse(url: &str) -> Result<Self, TransportError> {
94        let parsed =
95            Url::parse(url).map_err(|e| TransportError::InvalidUrl(format!("{url}: {e}")))?;
96
97        let is_tls = match parsed.scheme() {
98            "ws" => false,
99            "wss" => true,
100            other => {
101                return Err(TransportError::InvalidUrl(format!(
102                    "expected ws:// or wss:// scheme, was {other}"
103                )));
104            }
105        };
106
107        let raw_host = parsed
108            .host_str()
109            .ok_or_else(|| TransportError::InvalidUrl("missing hostname".to_string()))?;
110
111        // url::Url stores IPv6 literals in bracketed form (`[::1]`); the
112        // `CONNECT` request line and TLS SNI both want the unbracketed form.
113        let host = if raw_host.starts_with('[') && raw_host.ends_with(']') {
114            raw_host[1..raw_host.len() - 1].to_string()
115        } else {
116            raw_host.to_string()
117        };
118
119        let port = parsed.port().unwrap_or(if is_tls { 443 } else { 80 });
120
121        Ok(Self { host, port, is_tls })
122    }
123}
124
125/// Outcome of parsing a proxy URL prior to opening a tunnel.
126///
127/// SOCKS schemes are recognized but not implemented for the WebSocket path
128/// yet. They are surfaced as [`ProxyKind::Unsupported`] so callers can log
129/// a warning and fall back to a direct connection, preserving compatibility
130/// with REST configs that already pointed at a SOCKS proxy.
131#[derive(Debug, Clone, PartialEq, Eq)]
132pub enum ProxyKind {
133    /// HTTP / HTTPS forward proxy reachable via `CONNECT` tunneling.
134    Http(ProxyTarget),
135    /// Recognized scheme without a working tunnel (currently SOCKS).
136    Unsupported {
137        /// Original URL scheme (e.g. `socks5`).
138        scheme: String,
139    },
140}
141
142impl ProxyKind {
143    /// Parse a proxy URL into a [`ProxyKind`]. Returns
144    /// [`TransportError::InvalidUrl`] for malformed input or non-proxy
145    /// schemes (`ftp://`, `ws://`, etc.).
146    ///
147    /// # Errors
148    ///
149    /// See [`ProxyTarget::parse`] for the underlying validation.
150    pub fn parse(url: &str) -> Result<Self, TransportError> {
151        let parsed =
152            Url::parse(url).map_err(|e| TransportError::InvalidUrl(format!("{url}: {e}")))?;
153
154        match parsed.scheme() {
155            "http" | "https" => ProxyTarget::parse(url).map(ProxyKind::Http),
156            scheme @ ("socks5" | "socks5h" | "socks4" | "socks4a") => {
157                // Reject malformed inputs like `socks5:host:port` that parse as
158                // scheme + opaque path with no authority: surfacing them as
159                // Unsupported would silently fall back to a direct connection
160                // and hide the typo.
161                if parsed.host_str().is_none_or(str::is_empty) {
162                    return Err(TransportError::InvalidUrl(format!(
163                        "proxy URL '{url}' is missing a host (did you mean {scheme}://...)?"
164                    )));
165                }
166                Ok(Self::Unsupported {
167                    scheme: scheme.to_string(),
168                })
169            }
170            other => Err(TransportError::InvalidUrl(format!(
171                "unsupported proxy scheme '{other}'; expected http:// or https://"
172            ))),
173        }
174    }
175}
176
177/// Parsed components of a forward proxy URL.
178#[derive(Debug, Clone, PartialEq, Eq)]
179pub struct ProxyTarget {
180    /// Host name of the proxy (used for both DNS and TLS SNI when
181    /// [`ProxyTarget::is_tls`] is `true`).
182    pub host: String,
183    /// TCP port of the proxy.
184    pub port: u16,
185    /// `true` when the proxy URL scheme is `https`.
186    pub is_tls: bool,
187    /// Pre-computed `Proxy-Authorization` header value, if the URL embeds
188    /// `user:pass@`.
189    pub auth_header: Option<String>,
190}
191
192impl ProxyTarget {
193    /// Parse a proxy URL into the components needed to establish the tunnel.
194    ///
195    /// Only `http://` and `https://` schemes are accepted here. Use
196    /// [`ProxyKind::parse`] when callers need to distinguish recognised but
197    /// unsupported schemes (currently SOCKS) from malformed input.
198    ///
199    /// # Errors
200    ///
201    /// Returns [`TransportError::InvalidUrl`] for malformed URLs, missing
202    /// hosts, or any scheme other than `http`/`https`.
203    pub fn parse(url: &str) -> Result<Self, TransportError> {
204        let parsed =
205            Url::parse(url).map_err(|e| TransportError::InvalidUrl(format!("{url}: {e}")))?;
206
207        let is_tls = match parsed.scheme() {
208            "http" => false,
209            "https" => true,
210            "socks5" | "socks5h" | "socks4" | "socks4a" => {
211                return Err(TransportError::InvalidUrl(format!(
212                    "SOCKS proxy scheme '{}' is not yet supported for WebSocket connections; \
213                    use an http:// or https:// proxy",
214                    parsed.scheme()
215                )));
216            }
217            other => {
218                return Err(TransportError::InvalidUrl(format!(
219                    "unsupported proxy scheme '{other}'; expected http:// or https://"
220                )));
221            }
222        };
223
224        let raw_host = parsed
225            .host_str()
226            .ok_or_else(|| TransportError::InvalidUrl("proxy URL missing hostname".to_string()))?;
227
228        // url::Url stores IPv6 literals bracketed (`[::1]`); the bracketed
229        // form is only valid in the HTTP `Host:` header, not for DNS or
230        // TLS SNI, so we keep both representations.
231        let host = if raw_host.starts_with('[') && raw_host.ends_with(']') {
232            raw_host[1..raw_host.len() - 1].to_string()
233        } else {
234            raw_host.to_string()
235        };
236
237        let port = parsed.port().unwrap_or(if is_tls { 443 } else { 80 });
238
239        let auth_header = if parsed.username().is_empty() {
240            None
241        } else {
242            let username = decode_userinfo(parsed.username());
243            let password = decode_userinfo(parsed.password().unwrap_or(""));
244            let credentials = format!("{username}:{password}");
245            Some(format!("Basic {}", BASE64.encode(credentials)))
246        };
247
248        Ok(Self {
249            host,
250            port,
251            is_tls,
252            auth_header,
253        })
254    }
255}
256
257/// Percent-decode a userinfo field from a proxy URL. `url::Url` keeps the
258/// raw percent-encoded form, so we decode it here before assembling the
259/// `Basic` credentials.
260fn decode_userinfo(value: &str) -> String {
261    let bytes = nautilus_core::string::urlencoding::decode_bytes(value.as_bytes());
262    String::from_utf8_lossy(&bytes).into_owned()
263}
264
265/// Establish a tunneled connection through `proxy` to the WebSocket `target`.
266///
267/// On success the returned stream is positioned right after the proxy's
268/// `200`/`2xx` response, ready for the WebSocket handshake. The function does
269/// not perform the WebSocket handshake itself; callers wrap the stream in
270/// `tokio-tungstenite::client_async`.
271///
272/// # Errors
273///
274/// Returns a [`TransportError`] when:
275/// - The TCP connection to the proxy fails ([`TransportError::Io`]).
276/// - The TLS layer to the proxy or upstream cannot be established
277///   ([`TransportError::Tls`]).
278/// - The proxy returns a non-success status, malformed headers, or closes the
279///   stream before completing the response ([`TransportError::Handshake`]).
280pub async fn tunnel_via_proxy(
281    target: &WsTarget,
282    proxy: &ProxyTarget,
283) -> Result<ProxiedStream, TransportError> {
284    let tcp = TcpStream::connect((proxy.host.as_str(), proxy.port))
285        .await
286        .map_err(TransportError::Io)?;
287
288    if let Err(e) = tcp.set_nodelay(true) {
289        log::warn!("Failed to enable TCP_NODELAY on proxy connection: {e:?}");
290    }
291
292    if proxy.is_tls {
293        let proxy_tls = wrap_tls(tcp, &proxy.host).await?;
294        let tunneled = send_connect(proxy_tls, target, proxy).await?;
295        if target.is_tls {
296            let upstream = wrap_tls(tunneled, &target.host).await?;
297            Ok(ProxiedStream::TlsOverTlsProxy(Box::new(upstream)))
298        } else {
299            Ok(ProxiedStream::PlainOverTlsProxy(Box::new(tunneled)))
300        }
301    } else {
302        let tunneled = send_connect(tcp, target, proxy).await?;
303        if target.is_tls {
304            let upstream = wrap_tls(tunneled, &target.host).await?;
305            Ok(ProxiedStream::Tls(Box::new(upstream)))
306        } else {
307            Ok(ProxiedStream::Plain(tunneled))
308        }
309    }
310}
311
312/// Send a `CONNECT` request and return the underlying stream once a `2xx`
313/// status is received. The returned stream is positioned after the empty line
314/// terminating the proxy response headers.
315async fn send_connect<S>(
316    mut stream: S,
317    target: &WsTarget,
318    proxy: &ProxyTarget,
319) -> Result<S, TransportError>
320where
321    S: AsyncRead + AsyncWrite + Unpin,
322{
323    let host_header = format_host_header(&target.host, target.port);
324    let mut request = format!(
325        "CONNECT {host_header} HTTP/1.1\r\n\
326         Host: {host_header}\r\n\
327         Proxy-Connection: Keep-Alive\r\n"
328    );
329
330    if let Some(auth) = &proxy.auth_header {
331        request.push_str("Proxy-Authorization: ");
332        request.push_str(auth);
333        request.push_str("\r\n");
334    }
335    request.push_str("\r\n");
336
337    stream
338        .write_all(request.as_bytes())
339        .await
340        .map_err(TransportError::Io)?;
341    stream.flush().await.map_err(TransportError::Io)?;
342
343    read_connect_response(&mut stream).await?;
344    Ok(stream)
345}
346
347fn format_host_header(host: &str, port: u16) -> String {
348    if host.contains(':') && !(host.starts_with('[') && host.ends_with(']')) {
349        format!("[{host}]:{port}")
350    } else {
351        format!("{host}:{port}")
352    }
353}
354
355/// Read the proxy's response up to the empty line that terminates the
356/// headers, validating the status line.
357async fn read_connect_response<S>(stream: &mut S) -> Result<(), TransportError>
358where
359    S: AsyncRead + Unpin,
360{
361    let mut buf = Vec::with_capacity(512);
362    let mut byte = [0u8; 1];
363
364    loop {
365        let n = stream.read(&mut byte).await.map_err(TransportError::Io)?;
366        if n == 0 {
367            return Err(TransportError::Handshake(
368                "proxy closed connection before sending CONNECT response".to_string(),
369            ));
370        }
371
372        buf.push(byte[0]);
373
374        if buf.ends_with(b"\r\n\r\n") {
375            break;
376        }
377
378        if buf.len() > MAX_PROXY_RESPONSE_BYTES {
379            return Err(TransportError::Handshake(format!(
380                "proxy CONNECT response exceeded {MAX_PROXY_RESPONSE_BYTES} bytes without terminator"
381            )));
382        }
383    }
384
385    let text = std::str::from_utf8(&buf).map_err(|_| {
386        TransportError::Handshake("proxy CONNECT response was not valid UTF-8".to_string())
387    })?;
388
389    let status_line = text.lines().next().ok_or_else(|| {
390        TransportError::Handshake("proxy CONNECT response missing status line".to_string())
391    })?;
392
393    // Expect: `HTTP/1.1 200 Connection established` (or any 2xx).
394    let mut parts = status_line.splitn(3, ' ');
395    let _version = parts.next().ok_or_else(|| {
396        TransportError::Handshake(format!("malformed status line: {status_line}"))
397    })?;
398    let status_code = parts
399        .next()
400        .ok_or_else(|| TransportError::Handshake(format!("malformed status line: {status_line}")))?
401        .parse::<u16>()
402        .map_err(|_| TransportError::Handshake(format!("non-numeric status: {status_line}")))?;
403
404    if !(200..300).contains(&status_code) {
405        return Err(TransportError::Handshake(format!(
406            "proxy refused CONNECT: {status_line}"
407        )));
408    }
409
410    Ok(())
411}
412
413/// Wrap a stream in a `rustls`-backed TLS session using `webpki_roots`.
414async fn wrap_tls<S>(stream: S, server_name: &str) -> Result<TlsStream<S>, TransportError>
415where
416    S: AsyncRead + AsyncWrite + Unpin,
417{
418    let mut root_store = RootCertStore::empty();
419    root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
420
421    let config = ClientConfig::builder()
422        .with_root_certificates(root_store)
423        .with_no_client_auth();
424
425    let connector = TlsConnector::from(std::sync::Arc::new(config));
426    let domain = ServerName::try_from(server_name.to_string())
427        .map_err(|e| TransportError::Tls(format!("invalid DNS name '{server_name}': {e}")))?;
428
429    connector
430        .connect(domain, stream)
431        .await
432        .map_err(TransportError::Io)
433}
434
435#[cfg(test)]
436#[cfg(not(feature = "turmoil"))] // proxy hop is not modelled under the turmoil simulator
437mod tests {
438    use std::net::SocketAddr;
439
440    use rstest::rstest;
441    use tokio::net::TcpListener;
442
443    use super::*;
444
445    #[rstest]
446    fn ws_target_parses_wss() {
447        let target = WsTarget::parse("wss://stream.binance.com:9443/ws/btcusdt@trade").unwrap();
448        assert_eq!(target.host, "stream.binance.com");
449        assert_eq!(target.port, 9443);
450        assert!(target.is_tls);
451    }
452
453    #[rstest]
454    fn ws_target_default_ports() {
455        let plain = WsTarget::parse("ws://example.com/path").unwrap();
456        assert_eq!(plain.port, 80);
457        assert!(!plain.is_tls);
458
459        let tls = WsTarget::parse("wss://example.com/path").unwrap();
460        assert_eq!(tls.port, 443);
461        assert!(tls.is_tls);
462    }
463
464    #[rstest]
465    fn ws_target_strips_ipv6_brackets() {
466        let target = WsTarget::parse("wss://[::1]:9443/ws").unwrap();
467        assert_eq!(target.host, "::1");
468        assert_eq!(target.port, 9443);
469    }
470
471    #[rstest]
472    fn ws_target_rejects_non_ws_scheme() {
473        let err = WsTarget::parse("https://example.com").unwrap_err();
474        assert!(matches!(err, TransportError::InvalidUrl(_)));
475    }
476
477    #[rstest]
478    fn proxy_target_parses_http() {
479        let proxy = ProxyTarget::parse("http://127.0.0.1:9999").unwrap();
480        assert_eq!(proxy.host, "127.0.0.1");
481        assert_eq!(proxy.port, 9999);
482        assert!(!proxy.is_tls);
483        assert!(proxy.auth_header.is_none());
484    }
485
486    #[rstest]
487    fn proxy_target_default_ports() {
488        let plain = ProxyTarget::parse("http://proxy.example.com").unwrap();
489        assert_eq!(plain.port, 80);
490        let tls = ProxyTarget::parse("https://proxy.example.com").unwrap();
491        assert_eq!(tls.port, 443);
492        assert!(tls.is_tls);
493    }
494
495    #[rstest]
496    fn proxy_target_basic_auth() {
497        let proxy =
498            ProxyTarget::parse("http://proxytest:fixture42@proxy.example.com:8080").unwrap();
499        // base64("proxytest:fixture42") == "cHJveHl0ZXN0OmZpeHR1cmU0Mg=="
500        assert_eq!(
501            proxy.auth_header.unwrap(),
502            "Basic cHJveHl0ZXN0OmZpeHR1cmU0Mg=="
503        );
504    }
505
506    #[rstest]
507    fn proxy_target_basic_auth_decodes_percent_encoded() {
508        // `p%40ss` should decode to `p@ss` before assembling Basic credentials
509        let proxy = ProxyTarget::parse("http://us%2Fer:p%40ss@proxy.example.com:8080").unwrap();
510        let header = proxy.auth_header.unwrap();
511        // base64("us/er:p@ss") == "dXMvZXI6cEBzcw=="
512        assert_eq!(header, "Basic dXMvZXI6cEBzcw==");
513    }
514
515    #[tokio::test]
516    async fn send_connect_includes_proxy_authorization() {
517        let (client, mut server) = tokio::io::duplex(1024);
518        let target = WsTarget::parse("ws://example.com:80/path").unwrap();
519        let proxy =
520            ProxyTarget::parse("http://proxytest:fixture42@proxy.example.com:8080").unwrap();
521
522        let server_task = tokio::spawn(async move {
523            let mut request = Vec::new();
524            let mut chunk = [0; 64];
525            loop {
526                let n = server.read(&mut chunk).await.unwrap();
527                request.extend_from_slice(&chunk[..n]);
528                if request.ends_with(b"\r\n\r\n") {
529                    break;
530                }
531            }
532
533            let request = String::from_utf8(request).unwrap();
534            assert_eq!(
535                request,
536                "CONNECT example.com:80 HTTP/1.1\r\n\
537                 Host: example.com:80\r\n\
538                 Proxy-Connection: Keep-Alive\r\n\
539                 Proxy-Authorization: Basic cHJveHl0ZXN0OmZpeHR1cmU0Mg==\r\n\r\n"
540            );
541            server
542                .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
543                .await
544        });
545
546        let _stream = send_connect(client, &target, &proxy).await.unwrap();
547        server_task.await.unwrap().unwrap();
548    }
549
550    #[rstest]
551    fn proxy_target_strips_ipv6_brackets() {
552        let proxy = ProxyTarget::parse("http://[::1]:8080").unwrap();
553        assert_eq!(proxy.host, "::1");
554        assert_eq!(proxy.port, 8080);
555    }
556
557    #[rstest]
558    fn proxy_target_rejects_socks() {
559        let err = ProxyTarget::parse("socks5://127.0.0.1:1080").unwrap_err();
560        let TransportError::InvalidUrl(msg) = err else {
561            panic!("expected InvalidUrl");
562        };
563        assert!(msg.contains("SOCKS"));
564    }
565
566    #[rstest]
567    fn proxy_kind_classifies_http() {
568        let kind = ProxyKind::parse("http://127.0.0.1:9999").unwrap();
569        assert!(matches!(kind, ProxyKind::Http(_)));
570    }
571
572    #[rstest]
573    fn proxy_kind_classifies_socks_as_unsupported() {
574        let kind = ProxyKind::parse("socks5://127.0.0.1:1080").unwrap();
575        let ProxyKind::Unsupported { scheme } = kind else {
576            panic!("expected Unsupported");
577        };
578        assert_eq!(scheme, "socks5");
579    }
580
581    #[rstest]
582    fn proxy_kind_rejects_garbage() {
583        assert!(ProxyKind::parse("ftp://x").is_err());
584        assert!(ProxyKind::parse("").is_err());
585    }
586
587    #[rstest]
588    fn proxy_kind_rejects_socks_without_authority() {
589        // `socks5:host:port` (no `//`) parses as scheme + opaque path; surface
590        // as a real error instead of a silent direct-fallback.
591        let err = ProxyKind::parse("socks5:127.0.0.1:1080").unwrap_err();
592        assert!(matches!(err, TransportError::InvalidUrl(_)));
593    }
594
595    #[rstest]
596    fn proxy_target_rejects_unknown_scheme() {
597        let err = ProxyTarget::parse("ftp://proxy.example.com").unwrap_err();
598        assert!(matches!(err, TransportError::InvalidUrl(_)));
599    }
600
601    #[rstest]
602    fn proxy_target_rejects_empty() {
603        let err = ProxyTarget::parse("").unwrap_err();
604        assert!(matches!(err, TransportError::InvalidUrl(_)));
605    }
606
607    #[rstest]
608    fn host_header_brackets_ipv6() {
609        assert_eq!(format_host_header("example.com", 443), "example.com:443");
610        assert_eq!(format_host_header("::1", 443), "[::1]:443");
611        assert_eq!(format_host_header("[::1]", 443), "[::1]:443");
612    }
613
614    /// Spawn a fake HTTP proxy that returns the configured response after
615    /// reading one CONNECT request line. Returns the bound address.
616    async fn spawn_fake_proxy(response: &'static [u8]) -> SocketAddr {
617        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
618        let addr = listener.local_addr().unwrap();
619        tokio::spawn(async move {
620            let (mut stream, _) = listener.accept().await.unwrap();
621            let mut buf = vec![0u8; 1024];
622            // Read until we see the CONNECT terminator.
623            loop {
624                let n = AsyncReadExt::read(&mut stream, &mut buf).await.unwrap();
625                if n == 0 {
626                    break;
627                }
628
629                if buf[..n].windows(4).any(|w| w == b"\r\n\r\n") {
630                    break;
631                }
632            }
633            stream.write_all(response).await.unwrap();
634            stream.flush().await.unwrap();
635        });
636        addr
637    }
638
639    #[tokio::test]
640    async fn read_connect_response_accepts_2xx() {
641        let addr = spawn_fake_proxy(b"HTTP/1.1 200 Connection established\r\n\r\n").await;
642        let mut stream = TcpStream::connect(addr).await.unwrap();
643        stream
644            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
645            .await
646            .unwrap();
647        stream.flush().await.unwrap();
648        read_connect_response(&mut stream).await.unwrap();
649    }
650
651    #[tokio::test]
652    async fn read_connect_response_rejects_403() {
653        let addr = spawn_fake_proxy(b"HTTP/1.1 403 Forbidden\r\n\r\n").await;
654        let mut stream = TcpStream::connect(addr).await.unwrap();
655        stream
656            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
657            .await
658            .unwrap();
659        stream.flush().await.unwrap();
660        let err = read_connect_response(&mut stream).await.unwrap_err();
661        let TransportError::Handshake(msg) = err else {
662            panic!("expected Handshake error");
663        };
664        assert!(msg.contains("403"));
665    }
666
667    /// 300 sits on the upper boundary of the accepted `200..300` range; if
668    /// the check is ever loosened to `200..=300` this test fails. 407 is the
669    /// classic "Proxy Authentication Required" response. Non-numeric status
670    /// probes the parse path.
671    #[rstest]
672    #[case::status_300(&b"HTTP/1.1 300 Multiple Choices\r\n\r\n"[..], "300")]
673    #[case::status_407(
674        &b"HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic\r\n\r\n"[..],
675        "407",
676    )]
677    #[case::malformed_status(&b"HTTP/1.1 abc Boom\r\n\r\n"[..], "non-numeric")]
678    #[tokio::test]
679    async fn read_connect_response_rejects_non_2xx(
680        #[case] response: &'static [u8],
681        #[case] expected_msg_substring: &'static str,
682    ) {
683        let addr = spawn_fake_proxy(response).await;
684        let mut stream = TcpStream::connect(addr).await.unwrap();
685        stream
686            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
687            .await
688            .unwrap();
689        stream.flush().await.unwrap();
690        let err = read_connect_response(&mut stream).await.unwrap_err();
691        let TransportError::Handshake(msg) = err else {
692            panic!("expected Handshake error, was {err:?}");
693        };
694        assert!(
695            msg.contains(expected_msg_substring),
696            "expected error message to contain {expected_msg_substring:?}, was {msg:?}"
697        );
698    }
699
700    /// Closing the connection mid-response should produce a clear handshake
701    /// error rather than spinning on a zero-byte read.
702    #[tokio::test]
703    async fn read_connect_response_rejects_eof_before_terminator() {
704        // Truncated response: missing the empty line that ends the headers
705        let addr = spawn_fake_proxy(b"HTTP/1.1 200 OK\r\n").await;
706        let mut stream = TcpStream::connect(addr).await.unwrap();
707        stream
708            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
709            .await
710            .unwrap();
711        stream.flush().await.unwrap();
712        let err = read_connect_response(&mut stream).await.unwrap_err();
713        let TransportError::Handshake(msg) = err else {
714            panic!("expected Handshake error, was {err:?}");
715        };
716        assert!(
717            msg.contains("closed connection"),
718            "unexpected handshake error: {msg}"
719        );
720    }
721
722    /// A proxy that streams headers without ever emitting `\r\n\r\n` should
723    /// trip the size cap rather than allocating without bound.
724    #[tokio::test]
725    async fn read_connect_response_rejects_oversize_headers() {
726        let mut response = b"HTTP/1.1 200 OK\r\n".to_vec();
727        while response.len() <= MAX_PROXY_RESPONSE_BYTES {
728            response.extend_from_slice(b"X-Pad: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\r\n");
729        }
730        let leaked: &'static [u8] = response.leak();
731        let addr = spawn_fake_proxy(leaked).await;
732        let mut stream = TcpStream::connect(addr).await.unwrap();
733        stream
734            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
735            .await
736            .unwrap();
737        stream.flush().await.unwrap();
738        let err = read_connect_response(&mut stream).await.unwrap_err();
739        let TransportError::Handshake(msg) = err else {
740            panic!("expected Handshake error, was {err:?}");
741        };
742        assert!(
743            msg.contains("exceeded"),
744            "unexpected handshake error: {msg}"
745        );
746    }
747
748    /// After accepting the 2xx response, the stream cursor must sit immediately
749    /// after the terminating `\r\n\r\n` so the WebSocket handshake can read its
750    /// own response. Regression guard against over-reading the terminator.
751    #[tokio::test]
752    async fn read_connect_response_preserves_trailing_bytes() {
753        let addr = spawn_fake_proxy(b"HTTP/1.1 200 Connection established\r\n\r\nLEFTOVER").await;
754        let mut stream = TcpStream::connect(addr).await.unwrap();
755        stream
756            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
757            .await
758            .unwrap();
759        stream.flush().await.unwrap();
760        read_connect_response(&mut stream).await.unwrap();
761
762        let mut tail = [0u8; b"LEFTOVER".len()];
763        AsyncReadExt::read_exact(&mut stream, &mut tail)
764            .await
765            .unwrap();
766        assert_eq!(&tail, b"LEFTOVER");
767    }
768}
nautilus_network/websocket/proxy.rs

nautilus_network/websocket/
proxy.rs
