Refactor credential-broker tunnel protocol detection (#44077)

## What changed

Move TLS prefix detection into `brokered_tunnel` and carry explicit `BrokeredProtocols` requirements from the credential broker through HTTP `CONNECT` and SOCKS5 tunnel handling. Dispatch on `TunnelProtocol::Tls` or `TunnelProtocol::Opaque`, preserving the initial-read timeout, fragmented TLS header detection, and replay of all inspected bytes.

## Testing

Move the fragmented TLS regression test alongside the detector, retaining coverage for delayed header fragments and lossless prefix replay. Update broker and proxy assertions for the protocol-aware types.

GitOrigin-RevId: 6e5a92d8bdf6db5b59cff7f210c17b01d93ee23b

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<S>(
+    mut stream: S,
+    protocols: BrokeredProtocols,
+) -> Result<(TunnelProtocol, PeekStream<Cursor<Vec<u8>>, 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;

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<std::io::Result<()>> {
+        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<std::io::Result<usize>> {
+        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        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();
+}

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<String, String>) -> bool {

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]

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}"
                 );

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::<ConnectMitmMode>().copied(),
-            Some(ConnectMitmMode::DetectTls)
+            Some(ConnectMitmMode::DetectProtocol(
+                crate::brokered_tunnel::BrokeredProtocols { tls: true }
+            ))
         );
     }
 

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;

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<S>(mut stream: S) -> Result<(bool, TlsPeekStream<S>)>
-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.

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<std::io::Result<()>> {
-        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<std::io::Result<usize>> {
-        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
-    }
-
-    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
-    }
-
-    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        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 {

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
         })

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<mitm::MitmState>,
         extensions: Extensions,
     },
-    DetectTls {
+    DetectProtocol {
+        protocols: crate::brokered_tunnel::BrokeredProtocols,
         target: HostWithPort,
         mode: NetworkMode,
         mitm: Arc<mitm::MitmState>,
@@ -512,7 +514,7 @@ impl AsyncRead for Socks5TcpConnection {
     ) -> Poll<io::Result<()>> {
         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<io::Result<usize>> {
         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<io::Result<()>> {
         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<io::Result<()>> {
         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<SocketAddr> {
         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<SocketAddr> {
         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")]