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 recognized 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 modeled 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        let expected = format!("Basic {}", BASE64.encode("proxytest:fixture42"));
617        assert_eq!(proxy.auth_header.unwrap(), expected);
618    }
619
620    #[rstest]
621    fn proxy_target_basic_auth_decodes_percent_encoded() {
622        // `p%40ss` should decode to `p@ss` before assembling Basic credentials
623        let proxy = ProxyTarget::parse("http://us%2Fer:p%40ss@proxy.example.com:8080").unwrap();
624        let header = proxy.auth_header.unwrap();
625        let expected = format!("Basic {}", BASE64.encode("us/er:p@ss"));
626        assert_eq!(header, expected);
627    }
628
629    #[rstest]
630    fn proxy_target_basic_auth_with_empty_username() {
631        let proxy = ProxyTarget::parse("http://:fixture42@proxy.example.com:8080").unwrap();
632        let expected = format!("Basic {}", BASE64.encode(":fixture42"));
633
634        assert_eq!(proxy.auth_header.unwrap(), expected);
635    }
636
637    #[rstest]
638    fn proxy_debug_redacts_credentials() {
639        const SECRET: &str = "unique-proxy-secret";
640        let url = format!("http://proxytest:{SECRET}@proxy.example.com:8080");
641        let proxy_url = ProxyUrl::parse(url.clone()).unwrap();
642        let target = ProxyTarget::parse(&url).unwrap();
643        let proxy_url_debug = format!("{proxy_url:?}");
644        let target_debug = format!("{target:?}");
645        let encoded_credentials = BASE64.encode(format!("proxytest:{SECRET}"));
646
647        assert_eq!(proxy_url_debug, "ProxyUrl(\"<redacted>\")");
648        assert!(!proxy_url_debug.contains(SECRET));
649        assert!(!target_debug.contains(SECRET));
650        assert!(!target_debug.contains(&encoded_credentials));
651        assert!(target_debug.contains(REDACTED));
652    }
653
654    #[rstest]
655    fn proxy_parse_error_redacts_credentials() {
656        const SECRET: &str = "unique-proxy-secret";
657        let err = ProxyUrl::parse(format!("http://proxytest:{SECRET}@[::1"))
658            .expect_err("malformed proxy URL should fail");
659
660        assert!(!err.to_string().contains(SECRET));
661    }
662
663    #[tokio::test]
664    async fn send_connect_includes_proxy_authorization() {
665        let (client, mut server) = tokio::io::duplex(1024);
666        let target = WsTarget::parse("ws://example.com:80/path").unwrap();
667        let proxy =
668            ProxyTarget::parse("http://proxytest:fixture42@proxy.example.com:8080").unwrap();
669
670        let server_task = tokio::spawn(async move {
671            let mut request = Vec::new();
672            let mut chunk = [0; 64];
673            loop {
674                let n = server.read(&mut chunk).await.unwrap();
675                request.extend_from_slice(&chunk[..n]);
676                if request.ends_with(b"\r\n\r\n") {
677                    break;
678                }
679            }
680
681            let request = String::from_utf8(request).unwrap();
682            let expected = format!(
683                "CONNECT example.com:80 HTTP/1.1\r\n\
684                 Host: example.com:80\r\n\
685                 Proxy-Connection: Keep-Alive\r\n\
686                 Proxy-Authorization: Basic {}\r\n\r\n",
687                BASE64.encode("proxytest:fixture42")
688            );
689            assert_eq!(request, expected);
690            server
691                .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
692                .await
693        });
694
695        let _stream = send_connect(client, &target, &proxy).await.unwrap();
696        server_task.await.unwrap().unwrap();
697    }
698
699    #[rstest]
700    fn proxy_target_strips_ipv6_brackets() {
701        let proxy = ProxyTarget::parse("http://[::1]:8080").unwrap();
702        assert_eq!(proxy.host, "::1");
703        assert_eq!(proxy.port, 8080);
704    }
705
706    #[rstest]
707    fn proxy_target_rejects_socks() {
708        let err = ProxyTarget::parse("socks5://127.0.0.1:1080").unwrap_err();
709        let TransportError::InvalidUrl(msg) = err else {
710            panic!("expected InvalidUrl");
711        };
712        assert!(msg.contains("SOCKS"));
713    }
714
715    #[rstest]
716    fn proxy_kind_classifies_http() {
717        let kind = ProxyKind::parse("http://127.0.0.1:9999").unwrap();
718        assert!(matches!(kind, ProxyKind::Http(_)));
719    }
720
721    #[rstest]
722    fn proxy_kind_classifies_socks_as_unsupported() {
723        let kind = ProxyKind::parse("socks5://127.0.0.1:1080").unwrap();
724        let ProxyKind::Unsupported { scheme } = kind else {
725            panic!("expected Unsupported");
726        };
727        assert_eq!(scheme, "socks5");
728    }
729
730    #[rstest]
731    fn proxy_kind_rejects_garbage() {
732        assert!(ProxyKind::parse("ftp://x").is_err());
733        assert!(ProxyKind::parse("").is_err());
734    }
735
736    #[rstest]
737    fn proxy_kind_rejects_socks_without_authority() {
738        // `socks5:host:port` (no `//`) parses as scheme + opaque path; surface
739        // as a real error instead of a silent direct-fallback.
740        let err = ProxyKind::parse("socks5:127.0.0.1:1080").unwrap_err();
741        assert!(matches!(err, TransportError::InvalidUrl(_)));
742    }
743
744    #[rstest]
745    fn proxy_target_rejects_unknown_scheme() {
746        let err = ProxyTarget::parse("ftp://proxy.example.com").unwrap_err();
747        assert!(matches!(err, TransportError::InvalidUrl(_)));
748    }
749
750    #[rstest]
751    fn proxy_target_rejects_empty() {
752        let err = ProxyTarget::parse("").unwrap_err();
753        assert!(matches!(err, TransportError::InvalidUrl(_)));
754    }
755
756    #[rstest]
757    fn host_header_brackets_ipv6() {
758        assert_eq!(format_host_header("example.com", 443), "example.com:443");
759        assert_eq!(format_host_header("::1", 443), "[::1]:443");
760        assert_eq!(format_host_header("[::1]", 443), "[::1]:443");
761    }
762
763    /// Spawn a fake HTTP proxy that returns the configured response after
764    /// reading one CONNECT request line. Returns the bound address.
765    async fn spawn_fake_proxy(response: &'static [u8]) -> SocketAddr {
766        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
767        let addr = listener.local_addr().unwrap();
768        tokio::spawn(async move {
769            let (mut stream, _) = listener.accept().await.unwrap();
770            let mut buf = vec![0u8; 1024];
771            // Read until we see the CONNECT terminator.
772            loop {
773                let n = AsyncReadExt::read(&mut stream, &mut buf).await.unwrap();
774                if n == 0 {
775                    break;
776                }
777
778                if buf[..n].windows(4).any(|w| w == b"\r\n\r\n") {
779                    break;
780                }
781            }
782            stream.write_all(response).await.unwrap();
783            stream.flush().await.unwrap();
784        });
785        addr
786    }
787
788    #[tokio::test]
789    async fn read_connect_response_accepts_2xx() {
790        let addr = spawn_fake_proxy(b"HTTP/1.1 200 Connection established\r\n\r\n").await;
791        let mut stream = TcpStream::connect(addr).await.unwrap();
792        stream
793            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
794            .await
795            .unwrap();
796        stream.flush().await.unwrap();
797        read_connect_response(&mut stream).await.unwrap();
798    }
799
800    #[tokio::test]
801    async fn read_connect_response_rejects_403() {
802        let addr = spawn_fake_proxy(b"HTTP/1.1 403 Forbidden\r\n\r\n").await;
803        let mut stream = TcpStream::connect(addr).await.unwrap();
804        stream
805            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
806            .await
807            .unwrap();
808        stream.flush().await.unwrap();
809        let err = read_connect_response(&mut stream).await.unwrap_err();
810        assert!(matches!(err, TransportError::ProxyConnectRejected(403)));
811    }
812
813    /// 300 sits on the upper boundary of the accepted `200..300` range; if
814    /// the check is ever loosened to `200..=300` this test fails. 407 is the
815    /// classic "Proxy Authentication Required" response. Non-numeric status
816    /// probes the parse path.
817    #[rstest]
818    #[case::status_300(&b"HTTP/1.1 300 Multiple Choices\r\n\r\n"[..], Some(300), None)]
819    #[case::status_407(
820        &b"HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic\r\n\r\n"[..],
821        Some(407),
822        None,
823    )]
824    #[case::malformed_status(
825        &b"HTTP/1.1 abc Boom\r\n\r\n"[..],
826        None,
827        Some("invalid status code"),
828    )]
829    #[tokio::test]
830    async fn read_connect_response_rejects_non_2xx(
831        #[case] response: &'static [u8],
832        #[case] expected_status: Option<u16>,
833        #[case] expected_msg_substring: Option<&'static str>,
834    ) {
835        let addr = spawn_fake_proxy(response).await;
836        let mut stream = TcpStream::connect(addr).await.unwrap();
837        stream
838            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
839            .await
840            .unwrap();
841        stream.flush().await.unwrap();
842        let err = read_connect_response(&mut stream).await.unwrap_err();
843
844        match (expected_status, expected_msg_substring) {
845            (Some(status), None) => {
846                assert!(
847                    matches!(err, TransportError::ProxyConnectRejected(actual) if actual == status)
848                );
849            }
850            (None, Some(expected_msg_substring)) => {
851                let TransportError::Handshake(msg) = err else {
852                    panic!("expected Handshake error, was {err:?}");
853                };
854                assert!(
855                    msg.contains(expected_msg_substring),
856                    "expected error message to contain {expected_msg_substring:?}, was {msg:?}"
857                );
858            }
859            _ => unreachable!(),
860        }
861    }
862
863    #[tokio::test]
864    async fn read_connect_response_does_not_expose_reason_phrase() {
865        const SECRET: &str = "unique-proxy-secret";
866        let response = format!("HTTP/1.1 407 {SECRET}\r\n\r\n").into_bytes();
867        let response = Box::leak(response.into_boxed_slice());
868        let addr = spawn_fake_proxy(response).await;
869        let mut stream = TcpStream::connect(addr).await.unwrap();
870        stream
871            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
872            .await
873            .unwrap();
874        stream.flush().await.unwrap();
875
876        let err = read_connect_response(&mut stream).await.unwrap_err();
877
878        assert_eq!(err.to_string(), "proxy CONNECT rejected with status 407");
879        assert!(!err.to_string().contains(SECRET));
880    }
881
882    /// Closing the connection mid-response should produce a clear handshake
883    /// error rather than spinning on a zero-byte read.
884    #[tokio::test]
885    async fn read_connect_response_rejects_eof_before_terminator() {
886        // Truncated response: missing the empty line that ends the headers
887        let addr = spawn_fake_proxy(b"HTTP/1.1 200 OK\r\n").await;
888        let mut stream = TcpStream::connect(addr).await.unwrap();
889        stream
890            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
891            .await
892            .unwrap();
893        stream.flush().await.unwrap();
894        let err = read_connect_response(&mut stream).await.unwrap_err();
895        assert!(matches!(err, TransportError::ConnectionClosed));
896    }
897
898    /// A malformed version with an otherwise retryable status must stay a permanent
899    /// `Handshake` failure. The status is 503 deliberately: were the version left
900    /// unvalidated, this would parse as `ProxyConnectRejected(503)` and be retried.
901    #[tokio::test]
902    async fn read_connect_response_rejects_malformed_version_with_retryable_status() {
903        let addr = spawn_fake_proxy(b"NOT-HTTP 503 Service Unavailable\r\n\r\n").await;
904        let mut stream = TcpStream::connect(addr).await.unwrap();
905        stream
906            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
907            .await
908            .unwrap();
909        stream.flush().await.unwrap();
910        let err = read_connect_response(&mut stream).await.unwrap_err();
911        let TransportError::Handshake(msg) = err else {
912            panic!("expected Handshake error, was {err:?}");
913        };
914        assert!(msg.contains("malformed status line"), "was {msg}");
915    }
916
917    /// `0503` is not a valid status token but parses as `503` under bare `u16`
918    /// parsing, which would make a malformed line retryable. 503 is used because
919    /// a permanent status would pass whether or not the token is validated.
920    #[tokio::test]
921    async fn read_connect_response_rejects_malformed_status_token_with_retryable_value() {
922        let addr = spawn_fake_proxy(b"HTTP/1.1 0503 Service Unavailable\r\n\r\n").await;
923        let mut stream = TcpStream::connect(addr).await.unwrap();
924        stream
925            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
926            .await
927            .unwrap();
928        stream.flush().await.unwrap();
929        let err = read_connect_response(&mut stream).await.unwrap_err();
930        let TransportError::Handshake(msg) = err else {
931            panic!("expected Handshake error, was {err:?}");
932        };
933        assert!(msg.contains("invalid status code"), "was {msg}");
934    }
935
936    /// A proxy that streams headers without ever emitting `\r\n\r\n` should
937    /// trip the size cap rather than allocating without bound.
938    #[tokio::test]
939    async fn read_connect_response_rejects_oversize_headers() {
940        let mut response = b"HTTP/1.1 200 OK\r\n".to_vec();
941        while response.len() <= MAX_PROXY_RESPONSE_BYTES {
942            response.extend_from_slice(b"X-Pad: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\r\n");
943        }
944        let leaked: &'static [u8] = response.leak();
945        let addr = spawn_fake_proxy(leaked).await;
946        let mut stream = TcpStream::connect(addr).await.unwrap();
947        stream
948            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
949            .await
950            .unwrap();
951        stream.flush().await.unwrap();
952        let err = read_connect_response(&mut stream).await.unwrap_err();
953        let TransportError::Handshake(msg) = err else {
954            panic!("expected Handshake error, was {err:?}");
955        };
956        assert!(
957            msg.contains("exceeded"),
958            "unexpected handshake error: {msg}"
959        );
960    }
961
962    /// After accepting the 2xx response, the stream cursor must sit immediately
963    /// after the terminating `\r\n\r\n` so the WebSocket handshake can read its
964    /// own response. Regression guard against over-reading the terminator.
965    #[tokio::test]
966    async fn read_connect_response_preserves_trailing_bytes() {
967        let addr = spawn_fake_proxy(b"HTTP/1.1 200 Connection established\r\n\r\nLEFTOVER").await;
968        let mut stream = TcpStream::connect(addr).await.unwrap();
969        stream
970            .write_all(b"CONNECT host:443 HTTP/1.1\r\nHost: host:443\r\n\r\n")
971            .await
972            .unwrap();
973        stream.flush().await.unwrap();
974        read_connect_response(&mut stream).await.unwrap();
975
976        let mut tail = [0u8; b"LEFTOVER".len()];
977        AsyncReadExt::read_exact(&mut stream, &mut tail)
978            .await
979            .unwrap();
980        assert_eq!(&tail, b"LEFTOVER");
981    }
982}
nautilus_network/websocket/proxy.rs

nautilus_network/websocket/
proxy.rs
