1use 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
46const MAX_PROXY_RESPONSE_BYTES: usize = 16 * 1024;
51
52#[derive(Clone, PartialEq, Eq)]
57pub struct ProxyUrl(String);
58
59impl ProxyUrl {
60 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 #[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#[derive(Debug)]
94pub enum ProxiedStream {
95 Plain(TcpStream),
97 PlainOverTlsProxy(Box<TlsStream<TcpStream>>),
99 Tls(Box<TlsStream<TcpStream>>),
101 TlsOverTlsProxy(Box<TlsStream<TlsStream<TcpStream>>>),
103}
104
105trait 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#[derive(Debug, Clone, PartialEq, Eq)]
166pub struct WsTarget {
167 pub host: String,
169 pub port: u16,
171 pub is_tls: bool,
173}
174
175impl WsTarget {
176 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 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#[derive(Debug, Clone, PartialEq, Eq)]
221pub enum ProxyKind {
222 Http(ProxyTarget),
224 Unsupported {
226 scheme: String,
228 },
229}
230
231impl ProxyKind {
232 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 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#[derive(Clone, PartialEq, Eq)]
269pub struct ProxyTarget {
270 pub host: String,
273 pub port: u16,
275 pub is_tls: bool,
277 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 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 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
358fn 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
366pub 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
409async 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
453async 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 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 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 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
529async 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"))] mod 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 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 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 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 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 #[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 #[tokio::test]
892 async fn read_connect_response_rejects_eof_before_terminator() {
893 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 #[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 #[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 #[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 #[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}