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//! HTTP `CONNECT` tunneling for outbound WebSocket connections.
17//!
18//! HTTP and HTTPS proxy URLs are supported. An HTTPS proxy adds TLS to the proxy hop; a `wss`
19//! target adds a separate TLS session after the tunnel is established. URL user information
20//! becomes Basic proxy authentication, and credential-bearing values are redacted from `Debug`
21//! output.
22//!
23//! The tunnel accepts only a `2xx` response and bounds response headers before parsing. It returns
24//! a stream positioned for the WebSocket handshake rather than performing that handshake itself.
25//!
26//! SOCKS URLs are recognized but not tunneled: the client logs a warning and connects directly.
27//! Both transport backends tunnel through the same [`ProxiedStream`], which implements the IO
28//! traits over every tunnel shape so each backend runs its own handshake over the finished stream.
29
30use std::{
31    fmt::Debug,
32    io,
33    pin::Pin,
34    task::{Context, Poll},
35};
36
37use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
38use nautilus_core::string::secret::REDACTED;
39use rustls::{ClientConfig, RootCertStore, pki_types::ServerName};
40use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
41use tokio_rustls::{TlsConnector, client::TlsStream};
42use url::Url;
43
44use crate::{net::TcpStream, transport::TransportError};
45
46/// Maximum size of a `CONNECT` proxy response we are willing to read.
47///
48/// Bounds the buffer so a malicious or broken proxy cannot make us allocate
49/// indefinitely while we wait for the header terminator.
50const MAX_PROXY_RESPONSE_BYTES: usize = 16 * 1024;
51
52/// Validated HTTP or HTTPS proxy URL.
53///
54/// The underlying URL is intentionally redacted from [`Debug`] output because
55/// URL user-info can contain proxy credentials.
56#[derive(Clone, PartialEq, Eq)]
57pub struct ProxyUrl(String);
58
59impl ProxyUrl {
60    /// Parses and validates an HTTP or HTTPS proxy URL.
61    ///
62    /// # Errors
63    ///
64    /// Returns [`TransportError::InvalidUrl`] when the URL is malformed, has no host, or uses an
65    /// unsupported scheme.
66    pub fn parse(value: impl Into<String>) -> Result<Self, TransportError> {
67        let value = value.into();
68        ProxyTarget::parse(&value)?;
69        Ok(Self(value))
70    }
71
72    /// Returns the validated URL for transport configuration.
73    #[must_use]
74    pub fn expose(&self) -> &str {
75        &self.0
76    }
77}
78
79impl Debug for ProxyUrl {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        f.debug_tuple(stringify!(ProxyUrl))
82            .field(&REDACTED)
83            .finish()
84    }
85}
86
87/// Stream produced by `tunnel_via_proxy` when the upstream is `ws://`
88/// (no upstream TLS, but the proxy hop itself may have been TLS-protected).
89///
90/// The TLS-bearing variants are boxed because [`tokio_rustls::client::TlsStream`]
91/// is large enough that a flat enum trips `clippy::large_enum_variant`. Boxing
92/// keeps the discriminant cheap to move while leaving the rare TLS path on the
93/// heap.
94#[derive(Debug)]
95pub enum ProxiedStream {
96    /// Plain TCP after a plain proxy hop.
97    Plain(TcpStream),
98    /// Plain TCP after a TLS proxy hop.
99    PlainOverTlsProxy(Box<TlsStream<TcpStream>>),
100    /// Upstream TLS over a plain proxy hop.
101    Tls(Box<TlsStream<TcpStream>>),
102    /// Upstream TLS over a TLS proxy hop.
103    TlsOverTlsProxy(Box<TlsStream<TlsStream<TcpStream>>>),
104}
105
106/// Combines the two IO traits so [`ProxiedStream`] resolves its variant in one place.
107trait ProxiedIo: AsyncRead + AsyncWrite + Send + Unpin {}
108
109impl<T: AsyncRead + AsyncWrite + Send + Unpin> ProxiedIo for T {}
110
111impl ProxiedStream {
112    fn inner_mut(&mut self) -> &mut dyn ProxiedIo {
113        match self {
114            Self::Plain(s) => s,
115            Self::PlainOverTlsProxy(s) | Self::Tls(s) => s.as_mut(),
116            Self::TlsOverTlsProxy(s) => s.as_mut(),
117        }
118    }
119}
120
121impl AsyncRead for ProxiedStream {
122    fn poll_read(
123        self: Pin<&mut Self>,
124        cx: &mut Context<'_>,
125        buf: &mut ReadBuf<'_>,
126    ) -> Poll<io::Result<()>> {
127        Pin::new(self.get_mut().inner_mut()).poll_read(cx, buf)
128    }
129}
130
131impl AsyncWrite for ProxiedStream {
132    fn poll_write(
133        self: Pin<&mut Self>,
134        cx: &mut Context<'_>,
135        buf: &[u8],
136    ) -> Poll<io::Result<usize>> {
137        Pin::new(self.get_mut().inner_mut()).poll_write(cx, buf)
138    }
139
140    fn poll_write_vectored(
141        self: Pin<&mut Self>,
142        cx: &mut Context<'_>,
143        bufs: &[io::IoSlice<'_>],
144    ) -> Poll<io::Result<usize>> {
145        Pin::new(self.get_mut().inner_mut()).poll_write_vectored(cx, bufs)
146    }
147
148    fn is_write_vectored(&self) -> bool {
149        match self {
150            Self::Plain(s) => s.is_write_vectored(),
151            Self::PlainOverTlsProxy(s) | Self::Tls(s) => s.is_write_vectored(),
152            Self::TlsOverTlsProxy(s) => s.is_write_vectored(),
153        }
154    }
155
156    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
157        Pin::new(self.get_mut().inner_mut()).poll_flush(cx)
158    }
159
160    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
161        Pin::new(self.get_mut().inner_mut()).poll_shutdown(cx)
162    }
163}
164
165/// Parsed components of a target WebSocket URL needed by the proxy hop.
166#[derive(Debug, Clone, PartialEq, Eq)]
167pub struct WsTarget {
168    /// Host name for DNS / SNI / `CONNECT` request line.
169    pub host: String,
170    /// TCP port of the WebSocket origin.
171    pub port: u16,
172    /// `true` when the WebSocket scheme is `wss://`.
173    pub is_tls: bool,
174}
175
176impl WsTarget {
177    /// Parses a `ws://` or `wss://` URL into the host, port, and TLS components.
178    ///
179    /// # Errors
180    ///
181    /// Returns [`TransportError::InvalidUrl`] when the URL fails to parse,
182    /// is missing a hostname, or uses a scheme other than `ws`/`wss`.
183    pub fn parse(url: &str) -> Result<Self, TransportError> {
184        let parsed = Url::parse(url)
185            .map_err(|e| TransportError::InvalidUrl(format!("invalid WebSocket URL: {e}")))?;
186
187        let is_tls = match parsed.scheme() {
188            "ws" => false,
189            "wss" => true,
190            other => {
191                return Err(TransportError::InvalidUrl(format!(
192                    "expected ws:// or wss:// scheme, was {other}"
193                )));
194            }
195        };
196
197        let raw_host = parsed
198            .host_str()
199            .ok_or_else(|| TransportError::InvalidUrl("missing hostname".to_string()))?;
200
201        // url::Url stores IPv6 literals in bracketed form (`[::1]`); the
202        // `CONNECT` request line and TLS SNI both want the unbracketed form.
203        let host = if raw_host.starts_with('[') && raw_host.ends_with(']') {
204            raw_host[1..raw_host.len() - 1].to_string()
205        } else {
206            raw_host.to_string()
207        };
208
209        let port = parsed.port().unwrap_or(if is_tls { 443 } else { 80 });
210
211        Ok(Self { host, port, is_tls })
212    }
213}
214
215/// Outcome of parsing a proxy URL prior to opening a tunnel.
216///
217/// SOCKS schemes are recognized but not implemented for the WebSocket path
218/// yet. They are surfaced as [`ProxyKind::Unsupported`] so callers can log
219/// a warning and fall back to a direct connection, preserving compatibility
220/// with REST configs that already pointed at a SOCKS proxy.
221#[derive(Debug, Clone, PartialEq, Eq)]
222pub enum ProxyKind {
223    /// HTTP / HTTPS forward proxy reachable via `CONNECT` tunneling.
224    Http(ProxyTarget),
225    /// Recognized scheme without a working tunnel (currently SOCKS).
226    Unsupported {
227        /// Original URL scheme (e.g. `socks5`).
228        scheme: String,
229    },
230}
231
232impl ProxyKind {
233    /// Parses a proxy URL into a [`ProxyKind`]. Returns
234    /// [`TransportError::InvalidUrl`] for malformed input or non-proxy
235    /// schemes (`ftp://`, `ws://`, etc.).
236    ///
237    /// # Errors
238    ///
239    /// See [`ProxyTarget::parse`] for the underlying validation.
240    pub fn parse(url: &str) -> Result<Self, TransportError> {
241        let parsed = Url::parse(url)
242            .map_err(|e| TransportError::InvalidUrl(format!("invalid proxy URL: {e}")))?;
243
244        match parsed.scheme() {
245            "http" | "https" => ProxyTarget::parse(url).map(ProxyKind::Http),
246            scheme @ ("socks5" | "socks5h" | "socks4" | "socks4a") => {
247                // Reject malformed inputs like `socks5:host:port` that parse as
248                // scheme + opaque path with no authority: surfacing them as
249                // Unsupported would silently fall back to a direct connection
250                // and hide the typo.
251                if parsed.host_str().is_none_or(str::is_empty) {
252                    return Err(TransportError::InvalidUrl(format!(
253                        "proxy URL is missing a host (did you mean {scheme}://...)?"
254                    )));
255                }
256                Ok(Self::Unsupported {
257                    scheme: scheme.to_string(),
258                })
259            }
260            other => Err(TransportError::InvalidUrl(format!(
261                "unsupported proxy scheme '{other}'; expected http:// or https://"
262            ))),
263        }
264    }
265}
266
267/// Parsed components of a forward proxy URL.
268#[derive(Clone, PartialEq, Eq)]
269pub struct ProxyTarget {
270    /// Host name of the proxy (used for both DNS and TLS SNI when
271    /// [`ProxyTarget::is_tls`] is `true`).
272    pub host: String,
273    /// TCP port of the proxy.
274    pub port: u16,
275    /// `true` when the proxy URL scheme is `https`.
276    pub is_tls: bool,
277    /// Pre-computed `Proxy-Authorization` header value, if the URL embeds
278    /// `user:pass@`.
279    pub auth_header: Option<String>,
280}
281
282impl Debug for ProxyTarget {
283    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
284        f.debug_struct(stringify!(ProxyTarget))
285            .field("host", &self.host)
286            .field("port", &self.port)
287            .field("is_tls", &self.is_tls)
288            .field("auth_header", &self.auth_header.as_ref().map(|_| REDACTED))
289            .finish()
290    }
291}
292
293impl ProxyTarget {
294    /// Parses a proxy URL into the components needed to establish the tunnel.
295    ///
296    /// Only `http://` and `https://` schemes are accepted here. Use
297    /// [`ProxyKind::parse`] when callers need to distinguish recognised but
298    /// unsupported schemes (currently SOCKS) from malformed input.
299    ///
300    /// # Errors
301    ///
302    /// Returns [`TransportError::InvalidUrl`] for malformed URLs, missing
303    /// hosts, or any scheme other than `http`/`https`.
304    pub fn parse(url: &str) -> Result<Self, TransportError> {
305        let parsed = Url::parse(url)
306            .map_err(|e| TransportError::InvalidUrl(format!("invalid proxy URL: {e}")))?;
307
308        let is_tls = match parsed.scheme() {
309            "http" => false,
310            "https" => true,
311            "socks5" | "socks5h" | "socks4" | "socks4a" => {
312                return Err(TransportError::InvalidUrl(format!(
313                    "SOCKS proxy scheme '{}' is not yet supported for WebSocket connections; \
314                    use an http:// or https:// proxy",
315                    parsed.scheme()
316                )));
317            }
318            other => {
319                return Err(TransportError::InvalidUrl(format!(
320                    "unsupported proxy scheme '{other}'; expected http:// or https://"
321                )));
322            }
323        };
324
325        let raw_host = parsed
326            .host_str()
327            .ok_or_else(|| TransportError::InvalidUrl("proxy URL missing hostname".to_string()))?;
328
329        // url::Url stores IPv6 literals bracketed (`[::1]`); the bracketed
330        // form is only valid in the HTTP `Host:` header, not for DNS or
331        // TLS SNI, so we keep both representations.
332        let host = if raw_host.starts_with('[') && raw_host.ends_with(']') {
333            raw_host[1..raw_host.len() - 1].to_string()
334        } else {
335            raw_host.to_string()
336        };
337
338        let port = parsed.port().unwrap_or(if is_tls { 443 } else { 80 });
339
340        let auth_header = if parsed.username().is_empty() && parsed.password().is_none() {
341            None
342        } else {
343            let username = decode_userinfo(parsed.username());
344            let password = decode_userinfo(parsed.password().unwrap_or(""));
345            let credentials = format!("{username}:{password}");
346            Some(format!("Basic {}", BASE64.encode(credentials)))
347        };
348
349        Ok(Self {
350            host,
351            port,
352            is_tls,
353            auth_header,
354        })
355    }
356}
357
358/// Percent-decode a userinfo field from a proxy URL. `url::Url` keeps the
359/// raw percent-encoded form, so we decode it here before assembling the
360/// `Basic` credentials.
361fn decode_userinfo(value: &str) -> String {
362    let bytes = nautilus_core::string::urlencoding::decode_bytes(value.as_bytes());
363    String::from_utf8_lossy(&bytes).into_owned()
364}
365
366/// Establish a tunneled connection through `proxy` to the WebSocket `target`.
367///
368/// On success the returned stream is positioned right after the proxy's
369/// `200`/`2xx` response, ready for the WebSocket handshake. The function does
370/// not perform the WebSocket handshake itself; callers wrap the stream in
371/// `tokio-tungstenite::client_async`.
372///
373/// # Errors
374///
375/// Returns a [`TransportError`] when:
376/// - The TCP connection to the proxy fails ([`TransportError::Io`]).
377/// - The TLS layer to the proxy or upstream cannot be established
378///   ([`TransportError::Tls`]).
379/// - The proxy returns a non-success status, malformed headers, or closes the
380///   stream before completing the response.
381pub async fn tunnel_via_proxy(
382    target: &WsTarget,
383    proxy: &ProxyTarget,
384) -> Result<ProxiedStream, TransportError> {
385    let tcp = TcpStream::connect((proxy.host.as_str(), proxy.port))
386        .await
387        .map_err(TransportError::Io)?;
388
389    crate::net::apply_socket_options(&tcp);
390
391    if proxy.is_tls {
392        let proxy_tls = wrap_tls(tcp, &proxy.host).await?;
393        let tunneled = send_connect(proxy_tls, target, proxy).await?;
394        if target.is_tls {
395            let upstream = wrap_tls(tunneled, &target.host).await?;
396            Ok(ProxiedStream::TlsOverTlsProxy(Box::new(upstream)))
397        } else {
398            Ok(ProxiedStream::PlainOverTlsProxy(Box::new(tunneled)))
399        }
400    } else {
401        let tunneled = send_connect(tcp, target, proxy).await?;
402        if target.is_tls {
403            let upstream = wrap_tls(tunneled, &target.host).await?;
404            Ok(ProxiedStream::Tls(Box::new(upstream)))
405        } else {
406            Ok(ProxiedStream::Plain(tunneled))
407        }
408    }
409}
410
411/// Sends a `CONNECT` request and returns the underlying stream once a `2xx`
412/// status is received. The returned stream is positioned after the empty line
413/// terminating the proxy response headers.
414async fn send_connect<S>(
415    mut stream: S,
416    target: &WsTarget,
417    proxy: &ProxyTarget,
418) -> Result<S, TransportError>
419where
420    S: AsyncRead + AsyncWrite + Unpin,
421{
422    let host_header = format_host_header(&target.host, target.port);
423    let mut request = format!(
424        "CONNECT {host_header} HTTP/1.1\r\n\
425         Host: {host_header}\r\n\
426         Proxy-Connection: Keep-Alive\r\n"
427    );
428
429    if let Some(auth) = &proxy.auth_header {
430        request.push_str("Proxy-Authorization: ");
431        request.push_str(auth);
432        request.push_str("\r\n");
433    }
434    request.push_str("\r\n");
435
436    stream
437        .write_all(request.as_bytes())
438        .await
439        .map_err(TransportError::Io)?;
440    stream.flush().await.map_err(TransportError::Io)?;
441
442    read_connect_response(&mut stream).await?;
443    Ok(stream)
444}
445
446fn format_host_header(host: &str, port: u16) -> String {
447    if host.contains(':') && !(host.starts_with('[') && host.ends_with(']')) {
448        format!("[{host}]:{port}")
449    } else {
450        format!("{host}:{port}")
451    }
452}
453
454/// Reads the proxy's response up to the empty line that terminates the
455/// headers, validating the status line.
456async fn read_connect_response<S>(stream: &mut S) -> Result<(), TransportError>
457where
458    S: AsyncRead + Unpin,
459{
460    let mut buf = Vec::with_capacity(512);
461    let mut byte = [0u8; 1];
462
463    loop {
464        let n = stream.read(&mut byte).await.map_err(TransportError::Io)?;
465        if n == 0 {
466            return Err(TransportError::ConnectionClosed);
467        }
468
469        buf.push(byte[0]);
470
471        if buf.ends_with(b"\r\n\r\n") {
472            break;
473        }
474
475        if buf.len() > MAX_PROXY_RESPONSE_BYTES {
476            return Err(TransportError::Handshake(format!(
477                "proxy CONNECT response exceeded {MAX_PROXY_RESPONSE_BYTES} bytes without terminator"
478            )));
479        }
480    }
481
482    let text = std::str::from_utf8(&buf).map_err(|_| {
483        TransportError::Handshake("proxy CONNECT response was not valid UTF-8".to_string())
484    })?;
485
486    let status_line = text.lines().next().ok_or_else(|| {
487        TransportError::Handshake("proxy CONNECT response missing status line".to_string())
488    })?;
489
490    // Expect: `HTTP/1.1 200 Connection established` (or any 2xx).
491    let mut parts = status_line.splitn(3, ' ');
492    let version = parts.next().ok_or_else(|| {
493        TransportError::Handshake("proxy CONNECT response has a malformed status line".to_string())
494    })?;
495
496    // The version gates the status branch below: without this check a malformed line such as
497    // `NOT-HTTP 503 ...` would yield a retryable `ProxyConnectRejected` rather than staying a
498    // permanent handshake failure.
499    if !matches!(version, "HTTP/1.0" | "HTTP/1.1") {
500        return Err(TransportError::Handshake(
501            "proxy CONNECT response has a malformed status line".to_string(),
502        ));
503    }
504
505    // Parsed as a `StatusCode` rather than a bare `u16` for the same reason the version is
506    // validated above: the number now selects retry behavior, and `u16` parsing accepts tokens
507    // HTTP does not, normalizing `0503` to `503`.
508    let status_code = parts
509        .next()
510        .ok_or_else(|| {
511            TransportError::Handshake(
512                "proxy CONNECT response has a malformed status line".to_string(),
513            )
514        })?
515        .parse::<http::StatusCode>()
516        .map_err(|_| {
517            TransportError::Handshake(
518                "proxy CONNECT response has a invalid status code".to_string(),
519            )
520        })?
521        .as_u16();
522
523    if !(200..300).contains(&status_code) {
524        return Err(TransportError::ProxyConnectRejected(status_code));
525    }
526
527    Ok(())
528}
529
530/// Wraps a stream in a `rustls`-backed TLS session using `webpki_roots`.
531async fn wrap_tls<S>(stream: S, server_name: &str) -> Result<TlsStream<S>, TransportError>
532where
533    S: AsyncRead + AsyncWrite + Unpin,
534{
535    let mut root_store = RootCertStore::empty();
536    root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
537
538    let config = ClientConfig::builder()
539        .with_root_certificates(root_store)
540        .with_no_client_auth();
541
542    let connector = TlsConnector::from(std::sync::Arc::new(config));
543    let domain = ServerName::try_from(server_name.to_string())
544        .map_err(|e| TransportError::Tls(format!("invalid DNS name '{server_name}': {e}")))?;
545
546    connector
547        .connect(domain, stream)
548        .await
549        .map_err(TransportError::Io)
550}
551
552#[cfg(test)]
553#[cfg(not(feature = "turmoil"))] // proxy hop is not modelled under the turmoil simulator
554mod tests {
555    use std::net::SocketAddr;
556
557    use rstest::rstest;
558    use tokio::net::TcpListener;
559
560    use super::*;
561
562    #[rstest]
563    fn ws_target_parses_wss() {
564        let target = WsTarget::parse("wss://stream.binance.com:9443/ws/btcusdt@trade").unwrap();
565        assert_eq!(target.host, "stream.binance.com");
566        assert_eq!(target.port, 9443);
567        assert!(target.is_tls);
568    }
569
570    #[rstest]
571    fn ws_target_default_ports() {
572        let plain = WsTarget::parse("ws://example.com/path").unwrap();
573        assert_eq!(plain.port, 80);
574        assert!(!plain.is_tls);
575
576        let tls = WsTarget::parse("wss://example.com/path").unwrap();
577        assert_eq!(tls.port, 443);
578        assert!(tls.is_tls);
579    }
580
581    #[rstest]
582    fn ws_target_strips_ipv6_brackets() {
583        let target = WsTarget::parse("wss://[::1]:9443/ws").unwrap();
584        assert_eq!(target.host, "::1");
585        assert_eq!(target.port, 9443);
586    }
587
588    #[rstest]
589    fn ws_target_rejects_non_ws_scheme() {
590        let err = WsTarget::parse("https://example.com").unwrap_err();
591        assert!(matches!(err, TransportError::InvalidUrl(_)));
592    }
593
594    #[rstest]
595    fn proxy_target_parses_http() {
596        let proxy = ProxyTarget::parse("http://127.0.0.1:9999").unwrap();
597        assert_eq!(proxy.host, "127.0.0.1");
598        assert_eq!(proxy.port, 9999);
599        assert!(!proxy.is_tls);
600        assert!(proxy.auth_header.is_none());
601    }
602
603    #[rstest]
604    fn proxy_target_default_ports() {
605        let plain = ProxyTarget::parse("http://proxy.example.com").unwrap();
606        assert_eq!(plain.port, 80);
607        let tls = ProxyTarget::parse("https://proxy.example.com").unwrap();
608        assert_eq!(tls.port, 443);
609        assert!(tls.is_tls);
610    }
611
612    #[rstest]
613    fn proxy_target_basic_auth() {
614        let proxy =
615            ProxyTarget::parse("http://proxytest:fixture42@proxy.example.com:8080").unwrap();
616        // base64("proxytest:fixture42") == "cHJveHl0ZXN0OmZpeHR1cmU0Mg=="
617        assert_eq!(
618            proxy.auth_header.unwrap(),
619            "Basic cHJveHl0ZXN0OmZpeHR1cmU0Mg=="
620        );
621    }
622
623    #[rstest]
624    fn proxy_target_basic_auth_decodes_percent_encoded() {
625        // `p%40ss` should decode to `p@ss` before assembling Basic credentials
626        let proxy = ProxyTarget::parse("http://us%2Fer:p%40ss@proxy.example.com:8080").unwrap();
627        let header = proxy.auth_header.unwrap();
628        // base64("us/er:p@ss") == "dXMvZXI6cEBzcw=="
629        assert_eq!(header, "Basic dXMvZXI6cEBzcw==");
630    }
631
632    #[rstest]
633    fn proxy_target_basic_auth_with_empty_username() {
634        let proxy = ProxyTarget::parse("http://:fixture42@proxy.example.com:8080").unwrap();
635
636        assert_eq!(proxy.auth_header.unwrap(), "Basic OmZpeHR1cmU0Mg==");
637    }
638
639    #[rstest]
640    fn proxy_debug_redacts_credentials() {
641        const SECRET: &str = "unique-proxy-secret";
642        let url = format!("http://proxytest:{SECRET}@proxy.example.com:8080");
643        let proxy_url = ProxyUrl::parse(url.clone()).unwrap();
644        let target = ProxyTarget::parse(&url).unwrap();
645        let proxy_url_debug = format!("{proxy_url:?}");
646        let target_debug = format!("{target:?}");
647        let encoded_credentials = BASE64.encode(format!("proxytest:{SECRET}"));
648
649        assert_eq!(proxy_url_debug, "ProxyUrl(\"<redacted>\")");
650        assert!(!proxy_url_debug.contains(SECRET));
651        assert!(!target_debug.contains(SECRET));
652        assert!(!target_debug.contains(&encoded_credentials));
653        assert!(target_debug.contains(REDACTED));
654    }
655
656    #[rstest]
657    fn proxy_parse_error_redacts_credentials() {
658        const SECRET: &str = "unique-proxy-secret";
659        let err = ProxyUrl::parse(format!("http://proxytest:{SECRET}@[::1"))
660            .expect_err("malformed proxy URL should fail");
661
662        assert!(!err.to_string().contains(SECRET));
663    }
664
665    #[tokio::test]
666    async fn send_connect_includes_proxy_authorization() {
667        let (client, mut server) = tokio::io::duplex(1024);
668        let target = WsTarget::parse("ws://example.com:80/path").unwrap();
669        let proxy =
670            ProxyTarget::parse("http://proxytest:fixture42@proxy.example.com:8080").unwrap();
671
672        let server_task = tokio::spawn(async move {
673            let mut request = Vec::new();
674            let mut chunk = [0; 64];
675            loop {
676                let n = server.read(&mut chunk).await.unwrap();
677                request.extend_from_slice(&chunk[..n]);
678                if request.ends_with(b"\r\n\r\n") {
679                    break;
680                }
681            }
682
683            let request = String::from_utf8(request).unwrap();
684            assert_eq!(
685                request,
686                "CONNECT example.com:80 HTTP/1.1\r\n\
687                 Host: example.com:80\r\n\
688                 Proxy-Connection: Keep-Alive\r\n\
689                 Proxy-Authorization: Basic cHJveHl0ZXN0OmZpeHR1cmU0Mg==\r\n\r\n"
690            );
691            server
692                .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
693                .await
694        });
695
696        let _stream = send_connect(client, &target, &proxy).await.unwrap();
697        server_task.await.unwrap().unwrap();
698    }
699
700    #[rstest]
701    fn proxy_target_strips_ipv6_brackets() {
702        let proxy = ProxyTarget::parse("http://[::1]:8080").unwrap();
703        assert_eq!(proxy.host, "::1");
704        assert_eq!(proxy.port, 8080);
705    }
706
707    #[rstest]
708    fn proxy_target_rejects_socks() {
709        let err = ProxyTarget::parse("socks5://127.0.0.1:1080").unwrap_err();
710        let TransportError::InvalidUrl(msg) = err else {
711            panic!("expected InvalidUrl");
712        };
713        assert!(msg.contains("SOCKS"));
714    }
715
716    #[rstest]
717    fn proxy_kind_classifies_http() {
718        let kind = ProxyKind::parse("http://127.0.0.1:9999").unwrap();
719        assert!(matches!(kind, ProxyKind::Http(_)));
720    }
721
722    #[rstest]
723    fn proxy_kind_classifies_socks_as_unsupported() {
724        let kind = ProxyKind::parse("socks5://127.0.0.1:1080").unwrap();
725        let ProxyKind::Unsupported { scheme } = kind else {
726            panic!("expected Unsupported");
727        };
728        assert_eq!(scheme, "socks5");
729    }
730
731    #[rstest]
732    fn proxy_kind_rejects_garbage() {
733        assert!(ProxyKind::parse("ftp://x").is_err());
734        assert!(ProxyKind::parse("").is_err());
735    }
736
737    #[rstest]
738    fn proxy_kind_rejects_socks_without_authority() {
739        // `socks5:host:port` (no `//`) parses as scheme + opaque path; surface
740        // as a real error instead of a silent direct-fallback.
741        let err = ProxyKind::parse("socks5:127.0.0.1:1080").unwrap_err();
742        assert!(matches!(err, TransportError::InvalidUrl(_)));
743    }
744
745    #[rstest]
746    fn proxy_target_rejects_unknown_scheme() {
747        let err = ProxyTarget::parse("ftp://proxy.example.com").unwrap_err();
748        assert!(matches!(err, TransportError::InvalidUrl(_)));
749    }
750
751    #[rstest]
752    fn proxy_target_rejects_empty() {
753        let err = ProxyTarget::parse("").unwrap_err();
754        assert!(matches!(err, TransportError::InvalidUrl(_)));
755    }
756
757    #[rstest]
758    fn host_header_brackets_ipv6() {
759        assert_eq!(format_host_header("example.com", 443), "example.com:443");
760        assert_eq!(format_host_header("::1", 443), "[::1]:443");
761        assert_eq!(format_host_header("[::1]", 443), "[::1]:443");
762    }
763
764    /// Spawn a fake HTTP proxy that returns the configured response after
765    /// reading one CONNECT request line. Returns the bound address.
766    async fn spawn_fake_proxy(response: &'static [u8]) -> SocketAddr {
767        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
768        let addr = listener.local_addr().unwrap();
769        tokio::spawn(async move {
770            let (mut stream, _) = listener.accept().await.unwrap();
771            let mut buf = vec![0u8; 1024];
772            // Read until we see the CONNECT terminator.
773            loop {
774                let n = AsyncReadExt::read(&mut stream, &mut buf).await.unwrap();
775                if n == 0 {
776                    break;
777                }
778
779                if buf[..n].windows(4).any(|w| w == b"\r\n\r\n") {
780                    break;
781                }
782            }
783            stream.write_all(response).await.unwrap();
784            stream.flush().await.unwrap();
785        });
786        addr
787    }
788
789    #[tokio::test]
790    async fn read_connect_response_accepts_2xx() {
791        let addr = spawn_fake_proxy(b"HTTP/1.1 200 Connection established\r\n\r\n").await;
792        let mut stream = TcpStream::connect(addr).await.unwrap();
793        stream
794            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
795            .await
796            .unwrap();
797        stream.flush().await.unwrap();
798        read_connect_response(&mut stream).await.unwrap();
799    }
800
801    #[tokio::test]
802    async fn read_connect_response_rejects_403() {
803        let addr = spawn_fake_proxy(b"HTTP/1.1 403 Forbidden\r\n\r\n").await;
804        let mut stream = TcpStream::connect(addr).await.unwrap();
805        stream
806            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
807            .await
808            .unwrap();
809        stream.flush().await.unwrap();
810        let err = read_connect_response(&mut stream).await.unwrap_err();
811        assert!(matches!(err, TransportError::ProxyConnectRejected(403)));
812    }
813
814    /// 300 sits on the upper boundary of the accepted `200..300` range; if
815    /// the check is ever loosened to `200..=300` this test fails. 407 is the
816    /// classic "Proxy Authentication Required" response. Non-numeric status
817    /// probes the parse path.
818    #[rstest]
819    #[case::status_300(&b"HTTP/1.1 300 Multiple Choices\r\n\r\n"[..], Some(300), None)]
820    #[case::status_407(
821        &b"HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic\r\n\r\n"[..],
822        Some(407),
823        None,
824    )]
825    #[case::malformed_status(
826        &b"HTTP/1.1 abc Boom\r\n\r\n"[..],
827        None,
828        Some("invalid status code"),
829    )]
830    #[tokio::test]
831    async fn read_connect_response_rejects_non_2xx(
832        #[case] response: &'static [u8],
833        #[case] expected_status: Option<u16>,
834        #[case] expected_msg_substring: Option<&'static str>,
835    ) {
836        let addr = spawn_fake_proxy(response).await;
837        let mut stream = TcpStream::connect(addr).await.unwrap();
838        stream
839            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
840            .await
841            .unwrap();
842        stream.flush().await.unwrap();
843        let err = read_connect_response(&mut stream).await.unwrap_err();
844
845        match (expected_status, expected_msg_substring) {
846            (Some(status), None) => {
847                assert!(
848                    matches!(err, TransportError::ProxyConnectRejected(actual) if actual == status)
849                );
850            }
851            (None, Some(expected_msg_substring)) => {
852                let TransportError::Handshake(msg) = err else {
853                    panic!("expected Handshake error, was {err:?}");
854                };
855                assert!(
856                    msg.contains(expected_msg_substring),
857                    "expected error message to contain {expected_msg_substring:?}, was {msg:?}"
858                );
859            }
860            _ => unreachable!(),
861        }
862    }
863
864    #[tokio::test]
865    async fn read_connect_response_does_not_expose_reason_phrase() {
866        const SECRET: &str = "unique-proxy-secret";
867        let response = format!("HTTP/1.1 407 {SECRET}\r\n\r\n").into_bytes();
868        let response = Box::leak(response.into_boxed_slice());
869        let addr = spawn_fake_proxy(response).await;
870        let mut stream = TcpStream::connect(addr).await.unwrap();
871        stream
872            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
873            .await
874            .unwrap();
875        stream.flush().await.unwrap();
876
877        let err = read_connect_response(&mut stream).await.unwrap_err();
878
879        assert_eq!(err.to_string(), "proxy CONNECT rejected with status 407");
880        assert!(!err.to_string().contains(SECRET));
881    }
882
883    /// Closing the connection mid-response should produce a clear handshake
884    /// error rather than spinning on a zero-byte read.
885    #[tokio::test]
886    async fn read_connect_response_rejects_eof_before_terminator() {
887        // Truncated response: missing the empty line that ends the headers
888        let addr = spawn_fake_proxy(b"HTTP/1.1 200 OK\r\n").await;
889        let mut stream = TcpStream::connect(addr).await.unwrap();
890        stream
891            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
892            .await
893            .unwrap();
894        stream.flush().await.unwrap();
895        let err = read_connect_response(&mut stream).await.unwrap_err();
896        assert!(matches!(err, TransportError::ConnectionClosed));
897    }
898
899    /// A malformed version with an otherwise retryable status must stay a permanent
900    /// `Handshake` failure. The status is 503 deliberately: were the version left
901    /// unvalidated, this would parse as `ProxyConnectRejected(503)` and be retried.
902    #[tokio::test]
903    async fn read_connect_response_rejects_malformed_version_with_retryable_status() {
904        let addr = spawn_fake_proxy(b"NOT-HTTP 503 Service Unavailable\r\n\r\n").await;
905        let mut stream = TcpStream::connect(addr).await.unwrap();
906        stream
907            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
908            .await
909            .unwrap();
910        stream.flush().await.unwrap();
911        let err = read_connect_response(&mut stream).await.unwrap_err();
912        let TransportError::Handshake(msg) = err else {
913            panic!("expected Handshake error, was {err:?}");
914        };
915        assert!(msg.contains("malformed status line"), "was {msg}");
916    }
917
918    /// `0503` is not a valid status token but parses as `503` under bare `u16`
919    /// parsing, which would make a malformed line retryable. 503 is used because
920    /// a permanent status would pass whether or not the token is validated.
921    #[tokio::test]
922    async fn read_connect_response_rejects_malformed_status_token_with_retryable_value() {
923        let addr = spawn_fake_proxy(b"HTTP/1.1 0503 Service Unavailable\r\n\r\n").await;
924        let mut stream = TcpStream::connect(addr).await.unwrap();
925        stream
926            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
927            .await
928            .unwrap();
929        stream.flush().await.unwrap();
930        let err = read_connect_response(&mut stream).await.unwrap_err();
931        let TransportError::Handshake(msg) = err else {
932            panic!("expected Handshake error, was {err:?}");
933        };
934        assert!(msg.contains("invalid status code"), "was {msg}");
935    }
936
937    /// A proxy that streams headers without ever emitting `\r\n\r\n` should
938    /// trip the size cap rather than allocating without bound.
939    #[tokio::test]
940    async fn read_connect_response_rejects_oversize_headers() {
941        let mut response = b"HTTP/1.1 200 OK\r\n".to_vec();
942        while response.len() <= MAX_PROXY_RESPONSE_BYTES {
943            response.extend_from_slice(b"X-Pad: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\r\n");
944        }
945        let leaked: &'static [u8] = response.leak();
946        let addr = spawn_fake_proxy(leaked).await;
947        let mut stream = TcpStream::connect(addr).await.unwrap();
948        stream
949            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
950            .await
951            .unwrap();
952        stream.flush().await.unwrap();
953        let err = read_connect_response(&mut stream).await.unwrap_err();
954        let TransportError::Handshake(msg) = err else {
955            panic!("expected Handshake error, was {err:?}");
956        };
957        assert!(
958            msg.contains("exceeded"),
959            "unexpected handshake error: {msg}"
960        );
961    }
962
963    /// After accepting the 2xx response, the stream cursor must sit immediately
964    /// after the terminating `\r\n\r\n` so the WebSocket handshake can read its
965    /// own response. Regression guard against over-reading the terminator.
966    #[tokio::test]
967    async fn read_connect_response_preserves_trailing_bytes() {
968        let addr = spawn_fake_proxy(b"HTTP/1.1 200 Connection established\r\n\r\nLEFTOVER").await;
969        let mut stream = TcpStream::connect(addr).await.unwrap();
970        stream
971            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
972            .await
973            .unwrap();
974        stream.flush().await.unwrap();
975        read_connect_response(&mut stream).await.unwrap();
976
977        let mut tail = [0u8; b"LEFTOVER".len()];
978        AsyncReadExt::read_exact(&mut stream, &mut tail)
979            .await
980            .unwrap();
981        assert_eq!(&tail, b"LEFTOVER");
982    }
983}
nautilus_network/websocket/proxy.rs

nautilus_network/websocket/
proxy.rs
