diff --git a/codex-rs/network-proxy/src/brokered_tunnel.rs b/codex-rs/network-proxy/src/brokered_tunnel.rs new file mode 100644 index 0000000000..9dd6ce6520 --- /dev/null +++ b/codex-rs/network-proxy/src/brokered_tunnel.rs @@ -0,0 +1,66 @@ +//! Detect broker-required TLS without consuming bytes needed for opaque forwarding. + +use anyhow::Context; +use anyhow::Result; +use rama_core::extensions::ExtensionsMut; +use rama_core::stream::PeekStream; +use rama_core::stream::Stream; +use std::io::Cursor; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::time::timeout; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct BrokeredProtocols { + pub(crate) tls: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum TunnelProtocol { + Tls, + Opaque, +} + +const PREFIX_TIMEOUT: Duration = Duration::from_millis(250); + +/// Detect only configured protocols, replaying every byte for opaque forwarding. +pub(crate) async fn peek_protocol( + mut stream: S, + protocols: BrokeredProtocols, +) -> Result<(TunnelProtocol, PeekStream>, S>)> +where + S: Stream + Unpin + ExtensionsMut, +{ + let mut prefix = Vec::new(); + let mut buf = [0_u8; 5]; + if let Ok(result) = timeout(PREFIX_TIMEOUT, stream.read(&mut buf)).await { + prefix.extend_from_slice(&buf[..result.context("read tunnel prefix")?]) + } + let protocol = if prefix.first() == Some(&0x16) && protocols.tls { + // Do not time out a fragmented TLS record once the client has started its handshake. + while prefix.len() < 5 + && matches!( + prefix.as_slice(), + [0x16] | [0x16, 0x03] | [0x16, 0x03, 0x00..=0x04, ..] + ) + { + let read = stream.read(&mut buf[..5 - prefix.len()]).await?; + if read == 0 { + break; + } + prefix.extend_from_slice(&buf[..read]); + } + if matches!(prefix.as_slice(), [0x16, 0x03, 0x00..=0x04, _, _]) { + TunnelProtocol::Tls + } else { + TunnelProtocol::Opaque + } + } else { + TunnelProtocol::Opaque + }; + Ok((protocol, PeekStream::new(Cursor::new(prefix), stream))) +} + +#[cfg(test)] +#[path = "brokered_tunnel_tests.rs"] +mod tests; diff --git a/codex-rs/network-proxy/src/brokered_tunnel_tests.rs b/codex-rs/network-proxy/src/brokered_tunnel_tests.rs new file mode 100644 index 0000000000..9fbc768860 --- /dev/null +++ b/codex-rs/network-proxy/src/brokered_tunnel_tests.rs @@ -0,0 +1,94 @@ +//! Regression coverage for fragmented TLS detection and lossless prefix replay. + +use pretty_assertions::assert_eq; +use rama_core::extensions::Extensions; +use rama_core::extensions::ExtensionsMut; +use rama_core::extensions::ExtensionsRef; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; +use tokio::io::DuplexStream; +use tokio::io::ReadBuf; +use tokio::time::Duration; + +struct TestStream { + inner: DuplexStream, + extensions: Extensions, +} + +impl TestStream { + fn new(inner: DuplexStream) -> Self { + Self { + inner, + extensions: Extensions::new(), + } + } +} + +impl AsyncRead for TestStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for TestStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } +} + +impl ExtensionsRef for TestStream { + fn extensions(&self) -> &Extensions { + &self.extensions + } +} + +impl ExtensionsMut for TestStream { + fn extensions_mut(&mut self) -> &mut Extensions { + &mut self.extensions + } +} + +#[tokio::test] +async fn tls_prefix_detection_accumulates_fragmented_reads() { + let tls_prefix = [0x16, 0x03, 0x03, 0x00, 0x80]; + let (mut writer, reader) = tokio::io::duplex(16); + let writer_task = tokio::spawn(async move { + writer.write_all(&tls_prefix[..1]).await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + writer.write_all(&tls_prefix[1..]).await.unwrap(); + }); + + let (protocol, mut stream) = crate::brokered_tunnel::peek_protocol( + TestStream::new(reader), + crate::brokered_tunnel::BrokeredProtocols { tls: true }, + ) + .await + .unwrap(); + let mut replayed = [0_u8; 5]; + stream.read_exact(&mut replayed).await.unwrap(); + + assert_eq!(protocol, crate::brokered_tunnel::TunnelProtocol::Tls); + assert_eq!(replayed, tls_prefix); + writer_task.await.unwrap(); +} diff --git a/codex-rs/network-proxy/src/credential_broker.rs b/codex-rs/network-proxy/src/credential_broker.rs index 5b598c34d8..e36b1deb66 100644 --- a/codex-rs/network-proxy/src/credential_broker.rs +++ b/codex-rs/network-proxy/src/credential_broker.rs @@ -1082,24 +1082,27 @@ impl CredentialBroker { #[cfg(test)] pub(crate) fn host_requires_mitm(&self, host: &str, port: u16) -> bool { - self.host_requires_mitm_for_environment(host, port, /*environment_id*/ None) + self.host_protocols_for_environment(host, port, /*environment_id*/ None) + .tls } - pub(crate) fn host_requires_mitm_for_environment( + pub(crate) fn host_protocols_for_environment( &self, host: &str, port: u16, environment_id: Option<&str>, - ) -> bool { + ) -> crate::brokered_tunnel::BrokeredProtocols { let normalized_host = normalize_host(host); let state = self.read_state(); - state.enabled - && state.credentials.iter().any(|credential| { - credential.belongs_to_environment(environment_id) - && credential - .host_bindings() - .any(|binding| binding.requires_mitm(&normalized_host, port)) - }) + crate::brokered_tunnel::BrokeredProtocols { + tls: state.enabled + && state.credentials.iter().any(|credential| { + credential.belongs_to_environment(environment_id) + && credential + .host_bindings() + .any(|binding| binding.requires_mitm(&normalized_host, port)) + }), + } } pub(crate) fn virtualize_text(&self, text: &mut String, env: &HashMap) -> bool { diff --git a/codex-rs/network-proxy/src/credential_broker/configured_tests.rs b/codex-rs/network-proxy/src/credential_broker/configured_tests.rs index 9831fd4b71..4578e55d27 100644 --- a/codex-rs/network-proxy/src/credential_broker/configured_tests.rs +++ b/codex-rs/network-proxy/src/credential_broker/configured_tests.rs @@ -812,11 +812,15 @@ fn configured_owner_keeps_prefixed_builtin_shaped_aliases_scoped() { retain_source.then_some(value) ); assert_eq!(value != token, brokered); - assert!(!broker.host_requires_mitm_for_environment( - "api.openai.com", - /*port*/ 443, - environment_id, - )); + assert!( + !broker + .host_protocols_for_environment( + "api.openai.com", + /*port*/ 443, + environment_id, + ) + .tls + ); let headers_for = |value: &str| { HeaderMap::from_iter([( AUTHORIZATION, @@ -2623,11 +2627,11 @@ fn configured_provider_destination_history_is_scoped_to_the_environment() { translated("second-environment", host, &second_dummy), format!("Bearer {second_dummy}") ); - assert!(!broker.host_requires_mitm_for_environment( - host, - /*port*/ 443, - Some("second-environment"), - )); + assert!( + !broker + .host_protocols_for_environment(host, /*port*/ 443, Some("second-environment"),) + .tls + ); } } @@ -2718,11 +2722,11 @@ fn configured_provider_preserves_filtered_destinations_but_honors_explicit_overr .and_then(|value| value.to_str().ok()), Some(format!("Bearer {dummy}").as_str()) ); - assert!(!broker.host_requires_mitm_for_environment( - "first.example", - /*port*/ 443, - Some("environment") - )); + assert!( + !broker + .host_protocols_for_environment("first.example", /*port*/ 443, Some("environment")) + .tls + ); } #[test] diff --git a/codex-rs/network-proxy/src/credential_broker_tests.rs b/codex-rs/network-proxy/src/credential_broker_tests.rs index 14e748efd6..a368d2258e 100644 --- a/codex-rs/network-proxy/src/credential_broker_tests.rs +++ b/codex-rs/network-proxy/src/credential_broker_tests.rs @@ -1097,11 +1097,13 @@ fn concurrent_commands_preserve_discovered_credential_destinations() { "{key}, {environment_id}, {host}" ); assert_eq!( - broker.host_requires_mitm_for_environment( - host, - /*port*/ 443, - Some(environment_id), - ), + broker + .host_protocols_for_environment( + host, + /*port*/ 443, + Some(environment_id), + ) + .tls, injected, "{key}, {environment_id}, {host}" ); diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index d460310b32..0dd011b784 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/codex-rs/network-proxy/src/http_proxy.rs @@ -89,7 +89,7 @@ use tracing::warn; enum ConnectMitmMode { Disabled, Enabled, - DetectTls, + DetectProtocol(crate::brokered_tunnel::BrokeredProtocols), } pub async fn run_http_proxy( @@ -312,7 +312,9 @@ async fn http_connect_accept( } else { match host_mitm_requirement { HostMitmRequirement::None => ConnectMitmMode::Disabled, - HostMitmRequirement::Tls => ConnectMitmMode::DetectTls, + HostMitmRequirement::Credential(protocols) => { + ConnectMitmMode::DetectProtocol(protocols) + } HostMitmRequirement::Always => ConnectMitmMode::Enabled, } }; @@ -387,11 +389,17 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { let result: Result<(), OpaqueError> = match connect_mitm_mode { ConnectMitmMode::Disabled => forward_connect_tunnel(upgraded).await, ConnectMitmMode::Enabled => mitm_connect_tunnel(upgraded).await, - ConnectMitmMode::DetectTls => match mitm::peek_tls_prefix(upgraded).await { - Ok((true, stream)) => mitm_connect_tunnel(stream).await, - Ok((false, stream)) => forward_connect_tunnel(stream).await, - Err(err) => Err(OpaqueError::from_display(format!("detect TLS: {err:#}"))), - }, + ConnectMitmMode::DetectProtocol(protocols) => { + match crate::brokered_tunnel::peek_protocol(upgraded, protocols).await { + Ok((crate::brokered_tunnel::TunnelProtocol::Tls, stream)) => { + mitm_connect_tunnel(stream).await + } + Ok((crate::brokered_tunnel::TunnelProtocol::Opaque, stream)) => { + forward_connect_tunnel(stream).await + } + Err(err) => Err(OpaqueError::from_display(format!("detect TLS: {err:#}"))), + } + } }; if let Err(err) = result { warn!("CONNECT tunnel error: {err}"); @@ -1324,7 +1332,9 @@ mod tests { assert_eq!(response.status(), StatusCode::OK); assert_eq!( request.extensions().get::().copied(), - Some(ConnectMitmMode::DetectTls) + Some(ConnectMitmMode::DetectProtocol( + crate::brokered_tunnel::BrokeredProtocols { tls: true } + )) ); } diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index 565c99e70c..7f9a52bf37 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -2,6 +2,7 @@ mod attribution; mod authorization_path; +mod brokered_tunnel; mod certs; mod config; mod connect_policy; diff --git a/codex-rs/network-proxy/src/mitm.rs b/codex-rs/network-proxy/src/mitm.rs index def083b036..16eb36bef6 100644 --- a/codex-rs/network-proxy/src/mitm.rs +++ b/codex-rs/network-proxy/src/mitm.rs @@ -26,8 +26,6 @@ use rama_core::extensions::ExtensionsRef; use rama_core::futures::stream::Stream as FuturesStream; use rama_core::rt::Executor; use rama_core::service::service_fn; -use rama_core::stream::PeekStream; -use rama_core::stream::StackReader; use rama_core::stream::Stream; use rama_http::Body; use rama_http::BodyDataStream; @@ -43,7 +41,6 @@ use rama_http::layer::remove_header::RemoveResponseHeaderLayer; use rama_http_backend::server::HttpServer; use rama_net::proxy::ProxyTarget; use rama_net::stream::SocketInfo; -use rama_net::tls::server::TlsPeekStream; use rama_tls_rustls::dep::rustls; use rama_tls_rustls::server::TlsAcceptorData; use rama_tls_rustls::server::TlsAcceptorLayer; @@ -51,9 +48,6 @@ use std::pin::Pin; use std::sync::Arc; use std::task::Context as TaskContext; use std::task::Poll; -use std::time::Duration; -use tokio::io::AsyncReadExt; -use tokio::time::timeout; use tracing::info; use tracing::warn; @@ -94,52 +88,6 @@ enum MitmPolicyDecision { const MITM_INSPECT_BODIES: bool = false; const MITM_MAX_BODY_BYTES: usize = 4096; -const TLS_PREFIX_LEN: usize = 5; -const TLS_PREFIX_FIRST_BYTE_TIMEOUT: Duration = Duration::from_millis(250); - -/// Peeks enough bytes to distinguish a TLS handshake from an opaque CONNECT stream. -/// -/// The first-byte timeout preserves server-first protocols. Once the client starts a possible TLS -/// prefix, all five record-header bytes are accumulated so fragmented handshakes cannot bypass -/// interception. Every byte read is replayed through `TlsPeekStream`. -pub(crate) async fn peek_tls_prefix(mut stream: S) -> Result<(bool, TlsPeekStream)> -where - S: Stream + Unpin + ExtensionsMut, -{ - let mut peek_buf = [0_u8; TLS_PREFIX_LEN]; - let mut bytes_read = - match timeout(TLS_PREFIX_FIRST_BYTE_TIMEOUT, stream.read(&mut peek_buf)).await { - Ok(result) => result.context("read TLS prefix")?, - Err(_) => 0, - }; - while bytes_read > 0 && bytes_read < TLS_PREFIX_LEN { - let possible_tls_prefix = matches!( - &peek_buf[..bytes_read], - [0x16] | [0x16, 0x03] | [0x16, 0x03, 0x00..=0x04, ..] - ); - if !possible_tls_prefix { - break; - } - let read = stream - .read(&mut peek_buf[bytes_read..]) - .await - .context("read TLS prefix")?; - if read == 0 { - break; - } - bytes_read += read; - } - - let is_tls = bytes_read == TLS_PREFIX_LEN && matches!(peek_buf, [0x16, 0x03, 0x00..=0x04, ..]); - let offset = TLS_PREFIX_LEN - bytes_read; - if offset > 0 { - peek_buf.copy_within(0..bytes_read, offset); - } - let mut peek = StackReader::new(peek_buf); - peek.skip(offset); - Ok((is_tls, PeekStream::new(peek, stream))) -} - impl std::fmt::Debug for MitmState { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { // Avoid dumping internal state (CA material, connectors, etc.) to logs. diff --git a/codex-rs/network-proxy/src/mitm_tests.rs b/codex-rs/network-proxy/src/mitm_tests.rs index a449ce32a2..2e423737ae 100644 --- a/codex-rs/network-proxy/src/mitm_tests.rs +++ b/codex-rs/network-proxy/src/mitm_tests.rs @@ -7,9 +7,6 @@ use crate::reasons::REASON_NOT_ALLOWED_LOCAL; use crate::runtime::network_proxy_state_for_policy; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; -use rama_core::extensions::Extensions; -use rama_core::extensions::ExtensionsMut; -use rama_core::extensions::ExtensionsRef; use rama_http::Body; use rama_http::HeaderMap; use rama_http::HeaderValue; @@ -17,90 +14,7 @@ use rama_http::Method; use rama_http::Request; use rama_http::StatusCode; use rama_http::header::HeaderName; -use std::pin::Pin; -use std::task::Context; -use std::task::Poll; use tempfile::NamedTempFile; -use tokio::io::AsyncRead; -use tokio::io::AsyncReadExt; -use tokio::io::AsyncWrite; -use tokio::io::AsyncWriteExt; -use tokio::io::DuplexStream; -use tokio::io::ReadBuf; -use tokio::time::Duration; - -struct TestStream { - inner: DuplexStream, - extensions: Extensions, -} - -impl TestStream { - fn new(inner: DuplexStream) -> Self { - Self { - inner, - extensions: Extensions::new(), - } - } -} - -impl AsyncRead for TestStream { - fn poll_read( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_read(cx, buf) - } -} - -impl AsyncWrite for TestStream { - fn poll_write( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) - } - - fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_flush(cx) - } - - fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) - } -} - -impl ExtensionsRef for TestStream { - fn extensions(&self) -> &Extensions { - &self.extensions - } -} - -impl ExtensionsMut for TestStream { - fn extensions_mut(&mut self) -> &mut Extensions { - &mut self.extensions - } -} - -#[tokio::test] -async fn tls_prefix_detection_accumulates_fragmented_reads() { - let tls_prefix = [0x16, 0x03, 0x03, 0x00, 0x80]; - let (mut writer, reader) = tokio::io::duplex(16); - let writer_task = tokio::spawn(async move { - writer.write_all(&tls_prefix[..1]).await.unwrap(); - tokio::time::sleep(TLS_PREFIX_FIRST_BYTE_TIMEOUT + Duration::from_millis(50)).await; - writer.write_all(&tls_prefix[1..]).await.unwrap(); - }); - - let (is_tls, mut stream) = peek_tls_prefix(TestStream::new(reader)).await.unwrap(); - let mut replayed = [0_u8; 5]; - stream.read_exact(&mut replayed).await.unwrap(); - - assert!(is_tls); - assert_eq!(replayed, tls_prefix); - writer_task.await.unwrap(); -} fn github_write_hook() -> crate::mitm_hook::MitmHookConfig { crate::mitm_hook::MitmHookConfig { diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index 0f417e97ca..fd3131ad21 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -247,7 +247,7 @@ pub struct NetworkProxyState { #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum HostMitmRequirement { None, - Tls, + Credential(crate::brokered_tunnel::BrokeredProtocols), Always, } @@ -900,14 +900,15 @@ impl NetworkProxyState { let guard = self.state.read().await; guard.mitm_hooks.contains_key(&normalized_host) }; - Ok(if host_has_mitm_hooks { - HostMitmRequirement::Always - } else if self.credential_broker.host_requires_mitm_for_environment( + let protocols = self.credential_broker.host_protocols_for_environment( &normalized_host, port, self.environment_id(), - ) { - HostMitmRequirement::Tls + ); + Ok(if host_has_mitm_hooks { + HostMitmRequirement::Always + } else if protocols.tls { + HostMitmRequirement::Credential(protocols) } else { HostMitmRequirement::None }) diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index 610f7ba1b4..642305fa88 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/codex-rs/network-proxy/src/socks5.rs @@ -383,7 +383,7 @@ async fn handle_socks5_tcp( } else { match host_mitm_requirement { HostMitmRequirement::None => SocksMitmMode::Disabled, - HostMitmRequirement::Tls => SocksMitmMode::DetectTls, + HostMitmRequirement::Credential(protocols) => SocksMitmMode::DetectProtocol(protocols), HostMitmRequirement::Always => SocksMitmMode::Enabled, } }; @@ -439,7 +439,8 @@ async fn handle_socks5_tcp( mitm: mitm_state, extensions: Extensions::new(), }), - SocksMitmMode::DetectTls => Some(Socks5TcpConnection::DetectTls { + SocksMitmMode::DetectProtocol(protocols) => Some(Socks5TcpConnection::DetectProtocol { + protocols, target, mode, mitm: mitm_state, @@ -483,7 +484,7 @@ async fn handle_socks5_tcp( enum SocksMitmMode { Disabled, Enabled, - DetectTls, + DetectProtocol(crate::brokered_tunnel::BrokeredProtocols), } #[derive(Debug)] @@ -495,7 +496,8 @@ enum Socks5TcpConnection { mitm: Arc, extensions: Extensions, }, - DetectTls { + DetectProtocol { + protocols: crate::brokered_tunnel::BrokeredProtocols, target: HostWithPort, mode: NetworkMode, mitm: Arc, @@ -512,7 +514,7 @@ impl AsyncRead for Socks5TcpConnection { ) -> Poll> { match self.get_mut() { Self::Direct(stream) => Pin::new(stream).poll_read(cx, buf), - Self::Mitm { .. } | Self::DetectTls { .. } => Poll::Ready(Ok(())), + Self::Mitm { .. } | Self::DetectProtocol { .. } => Poll::Ready(Ok(())), } } } @@ -525,21 +527,21 @@ impl AsyncWrite for Socks5TcpConnection { ) -> Poll> { match self.get_mut() { Self::Direct(stream) => Pin::new(stream).poll_write(cx, buf), - Self::Mitm { .. } | Self::DetectTls { .. } => Poll::Ready(Ok(buf.len())), + Self::Mitm { .. } | Self::DetectProtocol { .. } => Poll::Ready(Ok(buf.len())), } } fn poll_flush(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { match self.get_mut() { Self::Direct(stream) => Pin::new(stream).poll_flush(cx), - Self::Mitm { .. } | Self::DetectTls { .. } => Poll::Ready(Ok(())), + Self::Mitm { .. } | Self::DetectProtocol { .. } => Poll::Ready(Ok(())), } } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { match self.get_mut() { Self::Direct(stream) => Pin::new(stream).poll_shutdown(cx), - Self::Mitm { .. } | Self::DetectTls { .. } => Poll::Ready(Ok(())), + Self::Mitm { .. } | Self::DetectProtocol { .. } => Poll::Ready(Ok(())), } } } @@ -548,14 +550,18 @@ impl Socket for Socks5TcpConnection { fn local_addr(&self) -> io::Result { match self { Self::Direct(stream) => stream.local_addr(), - Self::Mitm { .. } | Self::DetectTls { .. } => Ok(SocketAddr::from(([0, 0, 0, 0], 0))), + Self::Mitm { .. } | Self::DetectProtocol { .. } => { + Ok(SocketAddr::from(([0, 0, 0, 0], 0))) + } } } fn peer_addr(&self) -> io::Result { match self { Self::Direct(stream) => stream.peer_addr(), - Self::Mitm { .. } | Self::DetectTls { .. } => Ok(SocketAddr::from(([0, 0, 0, 0], 0))), + Self::Mitm { .. } | Self::DetectProtocol { .. } => { + Ok(SocketAddr::from(([0, 0, 0, 0], 0))) + } } } } @@ -564,7 +570,7 @@ impl ExtensionsRef for Socks5TcpConnection { fn extensions(&self) -> &Extensions { match self { Self::Direct(stream) => stream.extensions(), - Self::Mitm { extensions, .. } | Self::DetectTls { extensions, .. } => extensions, + Self::Mitm { extensions, .. } | Self::DetectProtocol { extensions, .. } => extensions, } } } @@ -573,7 +579,7 @@ impl ExtensionsMut for Socks5TcpConnection { fn extensions_mut(&mut self) -> &mut Extensions { match self { Self::Direct(stream) => stream.extensions_mut(), - Self::Mitm { extensions, .. } | Self::DetectTls { extensions, .. } => extensions, + Self::Mitm { extensions, .. } | Self::DetectProtocol { extensions, .. } => extensions, } } } @@ -595,7 +601,8 @@ async fn proxy_socks5_tcp( source.extensions_mut().insert(mitm); mitm::mitm_stream(source).await.map_err(Into::into) } - Socks5TcpConnection::DetectTls { + Socks5TcpConnection::DetectProtocol { + protocols, target, mode, mitm, @@ -605,29 +612,32 @@ async fn proxy_socks5_tcp( source.extensions_mut().insert(ProxyTarget(target.clone())); source.extensions_mut().insert(mode); source.extensions_mut().insert(mitm); - let (is_tls, source) = mitm::peek_tls_prefix(source) + let (protocol, source) = crate::brokered_tunnel::peek_protocol(source, protocols) .await .map_err(|err| -> BoxError { err.into() })?; - if is_tls { - mitm::mitm_stream(source).await.map_err(Into::into) - } else { - info!("SOCKS opaque upstream dial started (target={target})"); - let connect_started_at = Instant::now(); - let EstablishedClientConnection { conn: upstream, .. } = - TargetCheckedTcpConnector::new(state) - .serve(TcpRequest::new(target.clone())) - .await?; - info!( - "SOCKS opaque upstream dial established (target={target}, elapsed_ms={})", - connect_started_at.elapsed().as_millis() - ); - StreamForwardService::default() - .serve(ProxyRequest { - source, - target: upstream, - }) - .await - .map_err(Into::into) + match protocol { + crate::brokered_tunnel::TunnelProtocol::Tls => { + mitm::mitm_stream(source).await.map_err(Into::into) + } + crate::brokered_tunnel::TunnelProtocol::Opaque => { + info!("SOCKS opaque upstream dial started (target={target})"); + let connect_started_at = Instant::now(); + let EstablishedClientConnection { conn: upstream, .. } = + TargetCheckedTcpConnector::new(state) + .serve(TcpRequest::new(target.clone())) + .await?; + info!( + "SOCKS opaque upstream dial established (target={target}, elapsed_ms={})", + connect_started_at.elapsed().as_millis() + ); + StreamForwardService::default() + .serve(ProxyRequest { + source, + target: upstream, + }) + .await + .map_err(Into::into) + } } } } @@ -1025,7 +1035,10 @@ mod tests { .await .expect("brokered TLS should defer MITM until protocol detection"); - assert!(matches!(result.conn, Socks5TcpConnection::DetectTls { .. })); + assert!(matches!( + result.conn, + Socks5TcpConnection::DetectProtocol { .. } + )); } #[tokio::test(flavor = "current_thread")]