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

nautilus_network/websocket/
proxy.rs
