Close active network proxy connections on teardown (#43884)

## Why

Stopping proxy listeners could leave accepted connections and half-closed tunnels alive after their owning thread was unloaded.

## What changed

Tie HTTP and SOCKS5 connection work, including HTTP CONNECT upgrades, to the lifetime of the proxy listeners. Cancel connections across the main and environment proxies when the handle is dropped, shutdown is requested, or `wait()` is canceled. Explicit shutdown waits for connection cleanup.

## Testing

Add live TCP regression tests for HTTP keep-alive connections and open or half-closed HTTP CONNECT and SOCKS5 tunnels across shutdown, handle drop, and canceled waits. Verify that half-closed tunnels still carry return traffic while the proxy is running, and that repeated thread creation and unloading closes managed proxy tunnels.

GitOrigin-RevId: b523f3b08d73e9ae368f6932952cc930c7b038c9

codex-rs/core/tests/suite/network_approval.rs

@@ -3,6 +3,7 @@ use anyhow::Result;
 use codex_config::types::ApprovalsReviewer;
 use codex_core::EnvironmentConfig;
 use codex_core::EnvironmentNetworkPolicy;
+use codex_core::StartThreadOptions;
 use codex_core::TurnInputRequest;
 use codex_core::config::Constrained;
 use codex_core::config::NetworkProxySpec;
@@ -1506,6 +1507,104 @@ async fn ambiguous_unattributed_network_request_is_not_assigned_to_active_calls(
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn thread_turnover_closes_managed_proxy_tunnels() -> Result<()> {
+    skip_if_host_windows!(Ok(()));
+    skip_if_no_network!(Ok(()));
+    skip_if_sandbox!(Ok(()));
+
+    let server = start_mock_server().await;
+    let test = managed_network_unified_exec_test(&server).await?;
+    let mut config = test.config.clone();
+    let mut network = NetworkProxyConfig {
+        enabled: true,
+        mode: codex_network_proxy::NetworkMode::Full,
+        allow_local_binding: true,
+        allow_upstream_proxy: false,
+        ..NetworkProxyConfig::default()
+    };
+    network.set_allowed_domains(vec!["127.0.0.1".to_string()]);
+    config.permissions.network = Some(NetworkProxySpec::from_config_and_constraints(
+        network,
+        /*requirements*/ None,
+        config.permissions.permission_profile(),
+    )?);
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let target = listener.local_addr()?;
+
+    for cycle in 0..3 {
+        let started = test
+            .thread_manager
+            .start_thread(StartThreadOptions::new(config.clone()))
+            .await?;
+        let proxy = started
+            .session_configured
+            .network_proxy
+            .as_ref()
+            .context("expected managed network proxy for new thread")?;
+        let (mut client, upstream) = tokio::time::timeout(Duration::from_secs(5), async {
+            let mut client = tokio::net::TcpStream::connect(&proxy.http_addr).await?;
+            client
+                .write_all(
+                    format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n").as_bytes(),
+                )
+                .await?;
+            let (mut upstream, _) = listener.accept().await?;
+            let mut headers = Vec::new();
+            while !headers.ends_with(b"\r\n\r\n") {
+                headers.push(client.read_u8().await?);
+            }
+            let headers = String::from_utf8(headers)?;
+            assert!(headers.starts_with("HTTP/1.1 200 "), "{headers:?}");
+            client.write_all(b"request").await?;
+            client.shutdown().await?;
+            let mut request = Vec::new();
+            upstream.read_to_end(&mut request).await?;
+            assert_eq!(request, b"request");
+            // Preserve valid half-close behavior while the thread is alive; its
+            // upstream peer deliberately keeps the write half open through unload.
+            upstream.write_all(b"response").await?;
+            let mut response = [0_u8; 8];
+            client.read_exact(&mut response).await?;
+            assert_eq!(&response, b"response");
+            Ok::<_, anyhow::Error>((client, upstream))
+        })
+        .await
+        .with_context(|| {
+            format!("thread cycle {cycle} did not establish a half-closed tunnel")
+        })??;
+
+        tokio::time::timeout(Duration::from_secs(5), started.thread.shutdown_and_wait())
+            .await
+            .with_context(|| format!("thread cycle {cycle} did not shut down"))??;
+        // Unloading releases both the manager's and caller's strong session owners.
+        drop(
+            test.thread_manager
+                .remove_thread(&started.thread_id)
+                .await
+                .context("expected stopped thread in manager")?,
+        );
+        drop(started);
+
+        let result = tokio::time::timeout(Duration::from_secs(2), client.read(&mut [0_u8; 1]))
+            .await
+            .with_context(|| format!("thread cycle {cycle} retained its tunnel after unload"))?;
+        match result {
+            Ok(bytes) => assert_eq!(bytes, 0),
+            Err(error) => assert!(
+                matches!(
+                    error.kind(),
+                    std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted
+                ),
+                "unexpected tunnel error after thread unload: {error}"
+            ),
+        }
+        drop(upstream);
+    }
+    test.codex.shutdown_and_wait().await?;
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[cfg_attr(
     not(target_os = "linux"),

codex-rs/network-proxy/src/connection_lifecycle/lifecycle_tests.rs

@@ -0,0 +1,266 @@
+//! Exercises proxy teardown through live TCP connections, including half-closed tunnels.
+
+use crate::NetworkProxy;
+use crate::NetworkProxyConfig;
+use crate::state::network_proxy_state_for_policy;
+use anyhow::Context;
+use anyhow::Result;
+use pretty_assertions::assert_eq;
+use std::collections::HashMap;
+use std::future::Future;
+use std::future::poll_fn;
+use std::io::ErrorKind;
+use std::net::Ipv4Addr;
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::task::Poll;
+use std::time::Duration;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::net::TcpListener;
+use tokio::net::TcpStream;
+use tokio::time::timeout;
+
+#[derive(Clone, Copy, Debug)]
+enum ConnectionKind {
+    HttpConnect(TunnelState),
+    Socks5(TunnelState),
+    HttpKeepAlive,
+}
+
+#[derive(Clone, Copy, Debug)]
+enum TunnelState {
+    Open,
+    ClientWriteClosed,
+    UpstreamWriteClosed,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum PeerSide {
+    Client,
+    Upstream,
+}
+
+struct Connection {
+    client: TcpStream,
+    upstream: TcpStream,
+    eof_received_by: Option<PeerSide>,
+}
+
+#[derive(Clone, Copy, Debug)]
+enum StopProxy {
+    Shutdown,
+    Drop,
+    CancelWait,
+}
+
+#[tokio::test]
+async fn shutdown_closes_proxy_connections() -> Result<()> {
+    assert_stop_closes_connections(StopProxy::Shutdown).await
+}
+
+#[tokio::test]
+async fn dropping_handle_closes_proxy_connections() -> Result<()> {
+    assert_stop_closes_connections(StopProxy::Drop).await
+}
+
+#[tokio::test]
+async fn canceling_wait_closes_proxy_connections() -> Result<()> {
+    assert_stop_closes_connections(StopProxy::CancelWait).await
+}
+
+async fn assert_stop_closes_connections(stop: StopProxy) -> Result<()> {
+    let mut config = NetworkProxyConfig {
+        allow_local_binding: true,
+        allow_upstream_proxy: false,
+        enable_socks5_udp: false,
+        ..NetworkProxyConfig::default()
+    };
+    config.set_allowed_domains(vec![Ipv4Addr::LOCALHOST.to_string()]);
+    let proxy = NetworkProxy::builder()
+        .state(Arc::new(network_proxy_state_for_policy(config)))
+        .managed_by_codex(!cfg!(target_os = "windows"))
+        .http_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
+        .socks_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
+        .build()
+        .await?;
+    let handle = proxy.run().await?;
+    // Environment listeners reserve ephemeral ports on every platform. Managed main
+    // listeners do so off Windows, whose shared ingress has separate route tests.
+    let prepared = proxy.prepare_for_remote_environment(HashMap::new(), "lifecycle-test")?;
+    let scopes = [
+        ("environment", prepared.env),
+        #[cfg(not(target_os = "windows"))]
+        (
+            "main",
+            proxy
+                .prepare_for_optional_environment(HashMap::new(), /*environment_id*/ None)?
+                .env,
+        ),
+    ];
+    let mut connections = Vec::new();
+    for (scope, env) in scopes {
+        for kind in [
+            ConnectionKind::HttpConnect(TunnelState::Open),
+            ConnectionKind::HttpConnect(TunnelState::ClientWriteClosed),
+            ConnectionKind::HttpConnect(TunnelState::UpstreamWriteClosed),
+            ConnectionKind::Socks5(TunnelState::Open),
+            ConnectionKind::Socks5(TunnelState::ClientWriteClosed),
+            ConnectionKind::Socks5(TunnelState::UpstreamWriteClosed),
+            ConnectionKind::HttpKeepAlive,
+        ] {
+            let (key, prefix) = match kind {
+                ConnectionKind::HttpConnect(_) | ConnectionKind::HttpKeepAlive => {
+                    ("HTTP_PROXY", "http://")
+                }
+                ConnectionKind::Socks5(_) => ("ALL_PROXY", "socks5h://"),
+            };
+            let addr = env[key]
+                .strip_prefix(prefix)
+                .context("proxy URL scheme")?
+                .parse()?;
+            let connection = timeout(Duration::from_secs(5), open_connection(kind, addr))
+                .await
+                .with_context(|| format!("{scope} {kind:?} connection did not become ready"))??;
+            connections.push((scope, kind, connection));
+        }
+    }
+
+    match stop {
+        StopProxy::Shutdown => timeout(Duration::from_secs(2), handle.shutdown())
+            .await
+            .context("proxy shutdown did not finish with active connections")??,
+        StopProxy::Drop => drop(handle),
+        StopProxy::CancelWait => {
+            let mut wait = Box::pin(handle.wait());
+            let state = poll_fn(|cx| Poll::Ready(wait.as_mut().poll(cx))).await;
+            assert!(
+                state.is_pending(),
+                "running proxy unexpectedly stopped: {state:?}"
+            );
+            drop(wait);
+        }
+    }
+
+    // Retain every endpoint until all checks finish: dropping a test peer could
+    // otherwise cause the EOF that a later assertion attributes to proxy teardown.
+    for (scope, kind, connection) in &mut connections {
+        for (side, stream) in [
+            (PeerSide::Client, &mut connection.client),
+            (PeerSide::Upstream, &mut connection.upstream),
+        ] {
+            if connection.eof_received_by == Some(side) {
+                continue;
+            }
+            let mut byte = [0_u8; 1];
+            let result = timeout(Duration::from_secs(2), stream.read(&mut byte))
+                .await
+                .with_context(|| {
+                    format!("{scope} {kind:?} {side:?} connection remained open after {stop:?}")
+                })?;
+            match result {
+                Ok(0) => {}
+                Err(error)
+                    if matches!(
+                        error.kind(),
+                        ErrorKind::ConnectionReset | ErrorKind::ConnectionAborted
+                    ) => {}
+                result => anyhow::bail!(
+                    "{scope} {kind:?} {side:?} did not close after {stop:?}: {result:?}"
+                ),
+            }
+        }
+    }
+    Ok(())
+}
+
+async fn open_connection(kind: ConnectionKind, proxy_addr: SocketAddr) -> Result<Connection> {
+    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
+    let target = listener.local_addr()?;
+    let mut client = TcpStream::connect(proxy_addr).await?;
+    let request = match kind {
+        ConnectionKind::HttpConnect(_) => {
+            format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n").into_bytes()
+        }
+        ConnectionKind::HttpKeepAlive => {
+            format!("GET http://{target}/ HTTP/1.1\r\nHost: {target}\r\n\r\n").into_bytes()
+        }
+        ConnectionKind::Socks5(_) => {
+            client.write_all(&[5, 1, 0]).await?;
+            let mut greeting = [0_u8; 2];
+            client.read_exact(&mut greeting).await?;
+            assert_eq!(greeting, [5, 0]);
+            let [port_high, port_low] = target.port().to_be_bytes();
+            vec![5, 1, 0, 1, 127, 0, 0, 1, port_high, port_low]
+        }
+    };
+    client.write_all(&request).await?;
+    let (upstream, _) = listener.accept().await?;
+    let mut connection = Connection {
+        client,
+        upstream,
+        eof_received_by: None,
+    };
+    let tunnel_state = match kind {
+        ConnectionKind::HttpConnect(state) => {
+            let response = read_http_headers(&mut connection.client).await?;
+            assert!(response.starts_with("HTTP/1.1 200 "), "{response:?}");
+            state
+        }
+        ConnectionKind::Socks5(state) => {
+            let mut response = [0_u8; 10];
+            connection.client.read_exact(&mut response).await?;
+            assert_eq!(&response[..4], &[5, 0, 0, 1]);
+            state
+        }
+        ConnectionKind::HttpKeepAlive => {
+            read_http_headers(&mut connection.upstream).await?;
+            connection
+                .upstream
+                .write_all(b"HTTP/1.1 204 No Content\r\n\r\n")
+                .await?;
+            let response = read_http_headers(&mut connection.client).await?;
+            assert!(response.starts_with("HTTP/1.1 204 "), "{response:?}");
+            return Ok(connection);
+        }
+    };
+
+    let (sender, receiver, receiver_side) = match tunnel_state {
+        TunnelState::Open | TunnelState::ClientWriteClosed => (
+            &mut connection.client,
+            &mut connection.upstream,
+            PeerSide::Upstream,
+        ),
+        TunnelState::UpstreamWriteClosed => (
+            &mut connection.upstream,
+            &mut connection.client,
+            PeerSide::Client,
+        ),
+    };
+    sender.write_all(b"request").await?;
+    let mut request = [0_u8; 7];
+    receiver.read_exact(&mut request).await?;
+    assert_eq!(&request, b"request");
+    match tunnel_state {
+        TunnelState::Open => {}
+        TunnelState::ClientWriteClosed | TunnelState::UpstreamWriteClosed => {
+            sender.shutdown().await?;
+            assert_eq!(receiver.read(&mut [0_u8; 1]).await?, 0);
+            connection.eof_received_by = Some(receiver_side);
+        }
+    }
+    // Traffic in the opposite direction remains valid after a write half closes.
+    receiver.write_all(b"response").await?;
+    let mut response = [0_u8; 8];
+    sender.read_exact(&mut response).await?;
+    assert_eq!(&response, b"response");
+    Ok(connection)
+}
+
+async fn read_http_headers(stream: &mut TcpStream) -> Result<String> {
+    let mut headers = Vec::new();
+    while !headers.ends_with(b"\r\n\r\n") {
+        headers.push(stream.read_u8().await?);
+    }
+    Ok(String::from_utf8(headers)?)
+}

codex-rs/network-proxy/src/connection_lifecycle/listeners.rs

@@ -0,0 +1,47 @@
+//! Keeps listener tasks and their connection scope alive under the same runtime owner.
+
+use super::scope::ConnectionLifecycle;
+use anyhow::Result;
+use rama_core::graceful::ShutdownGuard;
+use std::future::Future;
+use tokio::task::JoinSet;
+
+pub(crate) struct ProxyListeners {
+    connections: ConnectionLifecycle,
+    listeners: JoinSet<Result<()>>,
+}
+
+impl ProxyListeners {
+    pub(crate) fn new() -> Self {
+        Self {
+            connections: ConnectionLifecycle::new(),
+            listeners: JoinSet::new(),
+        }
+    }
+
+    pub(crate) fn spawn<F, Fut>(&mut self, listener: F)
+    where
+        F: FnOnce(ShutdownGuard) -> Fut,
+        Fut: Future<Output = Result<()>> + Send + 'static,
+    {
+        self.listeners.spawn(listener(self.connections.guard()));
+    }
+
+    pub(crate) fn cancel(&mut self) {
+        self.connections.cancel();
+        self.listeners.abort_all();
+    }
+
+    pub(crate) async fn wait(&mut self) -> Result<()> {
+        while let Some(result) = self.listeners.join_next().await {
+            result??;
+        }
+        Ok(())
+    }
+
+    pub(crate) async fn shutdown(mut self) {
+        self.cancel();
+        self.listeners.shutdown().await;
+        self.connections.shutdown().await;
+    }
+}

codex-rs/network-proxy/src/connection_lifecycle/mod.rs

@@ -0,0 +1,12 @@
+mod listeners;
+mod scope;
+mod service;
+
+pub(crate) use listeners::ProxyListeners;
+#[cfg(test)]
+pub(crate) use scope::ConnectionLifecycle;
+pub(crate) use service::CancelOnShutdown;
+
+#[cfg(test)]
+#[path = "lifecycle_tests.rs"]
+mod tests;

codex-rs/network-proxy/src/connection_lifecycle/scope.rs

@@ -0,0 +1,35 @@
+//! Owns cancellation and completion of a proxy's accepted connections and descendant tasks.
+
+use rama_core::graceful::Shutdown;
+use rama_core::graceful::ShutdownGuard;
+use tokio::sync::oneshot;
+
+pub(crate) struct ConnectionLifecycle {
+    // Only the runtime owner holds the sender. Dropping it also signals shutdown when an
+    // enclosing wait/shutdown future is cancelled, without spawning another cleanup task.
+    cancel: Option<oneshot::Sender<()>>,
+    shutdown: Shutdown,
+}
+
+impl ConnectionLifecycle {
+    pub(crate) fn new() -> Self {
+        let (cancel, signal) = oneshot::channel();
+        Self {
+            cancel: Some(cancel),
+            shutdown: Shutdown::new(signal),
+        }
+    }
+
+    pub(crate) fn guard(&self) -> ShutdownGuard {
+        self.shutdown.guard()
+    }
+
+    pub(crate) fn cancel(&mut self) {
+        self.cancel.take();
+    }
+
+    pub(crate) async fn shutdown(mut self) {
+        self.cancel();
+        self.shutdown.shutdown().await;
+    }
+}

codex-rs/network-proxy/src/connection_lifecycle/service.rs

@@ -0,0 +1,43 @@
+//! Cancels connection work when its executor's proxy scope ends, including HTTP upgrades.
+
+use rama_core::Service;
+use rama_core::extensions::ExtensionsRef;
+use rama_core::rt::Executor;
+
+#[derive(Clone)]
+pub(crate) struct CancelOnShutdown<S> {
+    inner: S,
+}
+
+impl<S> CancelOnShutdown<S> {
+    pub(crate) fn new(inner: S) -> Self {
+        Self { inner }
+    }
+}
+
+impl<S, Request> Service<Request> for CancelOnShutdown<S>
+where
+    S: Service<Request, Output = ()>,
+    Request: ExtensionsRef + Send + 'static,
+{
+    type Output = ();
+    type Error = S::Error;
+
+    async fn serve(&self, request: Request) -> Result<(), Self::Error> {
+        let guard = request
+            .extensions()
+            .get::<Executor>()
+            .and_then(Executor::guard)
+            .cloned();
+        match guard {
+            Some(guard) => {
+                tokio::select! {
+                    biased;
+                    _ = guard.cancelled() => Ok(()),
+                    result = self.inner.serve(request) => result,
+                }
+            }
+            None => self.inner.serve(request).await,
+        }
+    }
+}

codex-rs/network-proxy/src/http_proxy.rs

@@ -1,6 +1,7 @@
 use crate::attribution::BindConnectionAttribution;
 use crate::config::NetworkMode;
 use crate::connect_policy::TargetCheckedTcpConnector;
+use crate::connection_lifecycle::CancelOnShutdown;
 use crate::mitm;
 use crate::network_policy::BlockDecisionAuditEventArgs;
 use crate::network_policy::NetworkDecision;
@@ -41,6 +42,7 @@ use rama_core::error::ErrorExt as _;
 use rama_core::error::OpaqueError;
 use rama_core::extensions::ExtensionsMut;
 use rama_core::extensions::ExtensionsRef;
+use rama_core::graceful::ShutdownGuard;
 use rama_core::service::BoxService;
 use rama_core::service::service_fn;
 use rama_core::stream::Stream;
@@ -95,6 +97,7 @@ pub async fn run_http_proxy(
     addr: SocketAddr,
     policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
     environment_id: Option<String>,
+    guard: ShutdownGuard,
 ) -> Result<()> {
     let listener = TcpListener::build()
         .bind(addr)
@@ -107,7 +110,7 @@ pub async fn run_http_proxy(
         .map_err(anyhow::Error::from)
         .with_context(|| format!("bind HTTP proxy: {addr}"))?;
 
-    run_http_proxy_with_listener(state, listener, policy_decider, environment_id).await
+    run_http_proxy_with_listener(state, listener, policy_decider, environment_id, guard).await
 }
 
 pub async fn run_http_proxy_with_std_listener(
@@ -115,10 +118,11 @@ pub async fn run_http_proxy_with_std_listener(
     listener: StdTcpListener,
     policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
     environment_id: Option<String>,
+    guard: ShutdownGuard,
 ) -> Result<()> {
     let listener =
         TcpListener::try_from(listener).context("convert std listener to HTTP proxy listener")?;
-    run_http_proxy_with_listener(state, listener, policy_decider, environment_id).await
+    run_http_proxy_with_listener(state, listener, policy_decider, environment_id, guard).await
 }
 
 async fn run_http_proxy_with_listener(
@@ -126,6 +130,7 @@ async fn run_http_proxy_with_listener(
     listener: TcpListener,
     policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
     environment_id: Option<String>,
+    guard: ShutdownGuard,
 ) -> Result<()> {
     let addr = listener
         .local_addr()
@@ -134,7 +139,10 @@ async fn run_http_proxy_with_listener(
     info!("HTTP proxy listening on {addr}");
 
     listener
-        .serve(http_proxy_service(state, policy_decider, environment_id))
+        .serve_graceful(
+            guard,
+            CancelOnShutdown::new(http_proxy_service(state, policy_decider, environment_id)),
+        )
         .await;
     Ok(())
 }
@@ -161,7 +169,7 @@ pub(crate) fn http_proxy_service(
                         http_connect_accept(policy_decider.clone(), environment_id.clone(), req)
                     }
                 }),
-                service_fn(http_connect_proxy),
+                CancelOnShutdown::new(service_fn(http_connect_proxy)),
             ),
             RemoveResponseHeaderLayer::hop_by_hop(),
         )
@@ -1445,8 +1453,13 @@ mod tests {
         let proxy_addr = listener
             .local_addr()
             .expect("proxy listener should expose local addr");
+        let lifecycle = crate::connection_lifecycle::ConnectionLifecycle::new();
         let proxy_task = tokio::spawn(run_http_proxy_with_std_listener(
-            state, listener, /*policy_decider*/ None, /*environment_id*/ None,
+            state,
+            listener,
+            /*policy_decider*/ None,
+            /*environment_id*/ None,
+            lifecycle.guard(),
         ));
 
         let mut stream = tokio::net::TcpStream::connect(proxy_addr)
@@ -1522,11 +1535,13 @@ mod tests {
         let proxy_addr = listener
             .local_addr()
             .expect("proxy listener should expose local addr");
+        let lifecycle = crate::connection_lifecycle::ConnectionLifecycle::new();
         let proxy_task = tokio::spawn(run_http_proxy_with_std_listener(
             state.clone(),
             listener,
             /*policy_decider*/ None,
             /*environment_id*/ None,
+            lifecycle.guard(),
         ));
 
         let mut stream = tokio::net::TcpStream::connect(proxy_addr)

codex-rs/network-proxy/src/lib.rs

@@ -5,6 +5,7 @@ mod authorization_path;
 mod certs;
 mod config;
 mod connect_policy;
+mod connection_lifecycle;
 mod credential_broker;
 mod environment_policy;
 mod http_proxy;

codex-rs/network-proxy/src/proxy.rs

@@ -2,6 +2,7 @@ mod execution_scope;
 
 use crate::attribution::PROXY_ATTRIBUTION_TOKEN_ENV_KEY;
 use crate::config;
+use crate::connection_lifecycle::ProxyListeners;
 use crate::credential_broker::BROKERED_CREDENTIALS_ENV_KEY;
 use crate::credential_broker::CREDENTIAL_BROKER_ACTIVE_ENV_KEY;
 use crate::http_proxy;
@@ -35,7 +36,6 @@ use std::ops::RangeInclusive;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::RwLock;
-use tokio::task::JoinHandle;
 use tracing::warn;
 
 use self::execution_scope::ExecutionScope;
@@ -470,7 +470,7 @@ pub struct PreparedManagedNetwork {
 
 struct EnvironmentProxy {
     addrs: EnvironmentProxyAddrs,
-    runtime: EnvironmentProxyRuntime,
+    runtime: ProxyRuntime,
 }
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
@@ -479,21 +479,40 @@ enum EnvironmentProxyClient {
     TrustedBridge,
 }
 
-enum EnvironmentProxyRuntime {
-    ListenerTasks {
-        http_task: JoinHandle<Result<()>>,
-        socks_task: Option<JoinHandle<Result<()>>>,
-    },
+enum ProxyRuntime {
+    Listeners(ProxyListeners),
     #[cfg(target_os = "windows")]
-    SharedIngress { _route: Arc<WindowsProxyRoute> },
+    SharedIngress {
+        route: Arc<WindowsProxyRoute>,
+    },
 }
 
-impl EnvironmentProxyRuntime {
+impl ProxyRuntime {
     #[cfg(target_os = "windows")]
     fn network_proxy_restricting_sid(&self) -> Option<String> {
         match self {
-            Self::ListenerTasks { .. } => None,
-            Self::SharedIngress { _route: route } => Some(route.sid().to_string()),
+            Self::Listeners(_) => None,
+            Self::SharedIngress { route } => Some(route.sid().to_string()),
+        }
+    }
+
+    async fn wait(&mut self) -> Result<()> {
+        match self {
+            Self::Listeners(listeners) => listeners.wait().await,
+            #[cfg(target_os = "windows")]
+            Self::SharedIngress { .. } => std::future::pending().await,
+        }
+    }
+
+    fn stop(self) -> Option<ProxyListeners> {
+        match self {
+            Self::Listeners(mut listeners) => {
+                listeners.cancel();
+                Some(listeners)
+            }
+            // Shared ingress owns its listeners; dropping the route unregisters it.
+            #[cfg(target_os = "windows")]
+            Self::SharedIngress { .. } => None,
         }
     }
 }
@@ -1201,8 +1220,8 @@ impl NetworkProxy {
             anyhow::ensure!(
                 matches!(
                     (&proxy.runtime, uses_shared_ingress),
-                    (EnvironmentProxyRuntime::SharedIngress { .. }, true)
-                        | (EnvironmentProxyRuntime::ListenerTasks { .. }, false)
+                    (ProxyRuntime::SharedIngress { .. }, true)
+                        | (ProxyRuntime::Listeners(_), false)
                 ),
                 "network proxy for environment `{environment_id}` was prepared for a different client type"
             );
@@ -1244,13 +1263,13 @@ impl NetworkProxy {
                 environment_id,
                 EnvironmentProxy {
                     addrs,
-                    runtime: EnvironmentProxyRuntime::SharedIngress { _route: route },
+                    runtime: ProxyRuntime::SharedIngress { route },
                 },
             );
             return Ok(addrs);
         }
 
-        let runtime = tokio::runtime::Handle::try_current().with_context(|| {
+        tokio::runtime::Handle::try_current().with_context(|| {
             format!("failed to create network proxy for environment `{environment_id}`")
         })?;
         let listeners =
@@ -1272,49 +1291,47 @@ impl NetworkProxy {
             socks_listener,
         } = listeners;
 
+        let mut listeners = ProxyListeners::new();
         let environment_id = environment_id.to_string();
         let http_state = self.state.clone();
         let http_decider = self.policy_decider.clone();
         let http_environment_id = Some(environment_id.clone());
-        let http_task = runtime.spawn(async move {
+        listeners.spawn(move |guard| async move {
             http_proxy::run_http_proxy_with_std_listener(
                 http_state,
                 http_listener,
                 http_decider,
                 http_environment_id,
+                guard,
             )
             .await
         });
 
-        let socks_task = if self.socks_enabled {
+        if self.socks_enabled
+            && let Some(listener) = socks_listener
+        {
             let socks_state = self.state.clone();
             let socks_decider = self.policy_decider.clone();
             let socks_environment_id = Some(environment_id.clone());
             let socks5_udp_enabled = self.socks5_udp_enabled;
-            socks_listener.map(|listener| {
-                runtime.spawn(async move {
-                    socks5::run_socks5_with_std_listener(
-                        socks_state,
-                        listener,
-                        socks_decider,
-                        socks_environment_id,
-                        socks5_udp_enabled,
-                    )
-                    .await
-                })
-            })
-        } else {
-            None
-        };
+            listeners.spawn(move |guard| async move {
+                socks5::run_socks5_with_std_listener(
+                    socks_state,
+                    listener,
+                    socks_decider,
+                    socks_environment_id,
+                    socks5_udp_enabled,
+                    guard,
+                )
+                .await
+            });
+        }
 
         proxies.insert(
             environment_id,
             EnvironmentProxy {
                 addrs,
-                runtime: EnvironmentProxyRuntime::ListenerTasks {
-                    http_task,
-                    socks_task,
-                },
+                runtime: ProxyRuntime::Listeners(listeners),
             },
         );
         Ok(addrs)
@@ -1386,12 +1403,14 @@ impl NetworkProxy {
                 active_route.is_none(),
                 "shared managed Windows proxy route is already running"
             );
-            *active_route = Some(Arc::new(windows_runtime.ingress.register_route(
+            let route = Arc::new(windows_runtime.ingress.register_route(
                 windows_runtime.http_service.clone(),
                 windows_runtime.socks_service.clone(),
-            )));
+            ));
+            *active_route = Some(Arc::clone(&route));
             drop(active_route);
             return Ok(NetworkProxyHandle::windows_shared(
+                route,
                 Arc::clone(&windows_runtime.active_route),
                 Arc::clone(&self.environment_proxies),
             ));
@@ -1401,10 +1420,11 @@ impl NetworkProxy {
         let http_listener = reserved_listeners.and_then(|listeners| listeners.take_http());
         let socks_listener = reserved_listeners.and_then(|listeners| listeners.take_socks());
 
+        let mut listeners = ProxyListeners::new();
         let http_state = self.state.clone();
         let http_decider = self.policy_decider.clone();
         let http_addr = self.http_addr;
-        let http_task = tokio::spawn(async move {
+        listeners.spawn(move |guard| async move {
             match http_listener {
                 Some(listener) => {
                     http_proxy::run_http_proxy_with_std_listener(
@@ -1412,6 +1432,7 @@ impl NetworkProxy {
                         listener,
                         http_decider,
                         /*environment_id*/ None,
+                        guard,
                     )
                     .await
                 }
@@ -1421,18 +1442,19 @@ impl NetworkProxy {
                         http_addr,
                         http_decider,
                         /*environment_id*/ None,
+                        guard,
                     )
                     .await
                 }
             }
         });
 
-        let socks_task = if current_cfg.enable_socks5 {
+        if current_cfg.enable_socks5 {
             let socks_state = self.state.clone();
             let socks_decider = self.policy_decider.clone();
             let socks_addr = self.socks_addr;
             let enable_socks5_udp = current_cfg.enable_socks5_udp;
-            Some(tokio::spawn(async move {
+            listeners.spawn(move |guard| async move {
                 match socks_listener {
                     Some(listener) => {
                         socks5::run_socks5_with_std_listener(
@@ -1441,6 +1463,7 @@ impl NetworkProxy {
                             socks_decider,
                             /*environment_id*/ None,
                             enable_socks5_udp,
+                            guard,
                         )
                         .await
                     }
@@ -1451,20 +1474,17 @@ impl NetworkProxy {
                             socks_decider,
                             /*environment_id*/ None,
                             enable_socks5_udp,
+                            guard,
                         )
                         .await
                     }
                 }
-            }))
-        } else {
-            None
-        };
+            });
+        }
 
         Ok(NetworkProxyHandle {
-            http_task: Some(http_task),
-            socks_task,
-            environment_proxies: self.environment_proxies.clone(),
-            completed: false,
+            runtime: Some(ProxyRuntime::Listeners(listeners)),
+            environment_proxies: Some(Arc::clone(&self.environment_proxies)),
             #[cfg(target_os = "windows")]
             windows_active_route: None,
         })
@@ -1472,10 +1492,8 @@ impl NetworkProxy {
 }
 
 pub struct NetworkProxyHandle {
-    http_task: Option<JoinHandle<Result<()>>>,
-    socks_task: Option<JoinHandle<Result<()>>>,
-    environment_proxies: Arc<Mutex<HashMap<String, EnvironmentProxy>>>,
-    completed: bool,
+    runtime: Option<ProxyRuntime>,
+    environment_proxies: Option<Arc<Mutex<HashMap<String, EnvironmentProxy>>>>,
     #[cfg(target_os = "windows")]
     windows_active_route: Option<Arc<Mutex<Option<Arc<WindowsProxyRoute>>>>>,
 }
@@ -1483,10 +1501,8 @@ pub struct NetworkProxyHandle {
 impl NetworkProxyHandle {
     fn noop() -> Self {
         Self {
-            http_task: Some(tokio::spawn(async { Ok(()) })),
-            socks_task: None,
-            environment_proxies: Arc::new(Mutex::new(HashMap::new())),
-            completed: true,
+            runtime: None,
+            environment_proxies: None,
             #[cfg(target_os = "windows")]
             windows_active_route: None,
         }
@@ -1494,131 +1510,67 @@ impl NetworkProxyHandle {
 
     #[cfg(target_os = "windows")]
     fn windows_shared(
+        route: Arc<WindowsProxyRoute>,
         active_route: Arc<Mutex<Option<Arc<WindowsProxyRoute>>>>,
         environment_proxies: Arc<Mutex<HashMap<String, EnvironmentProxy>>>,
     ) -> Self {
         Self {
-            http_task: Some(tokio::spawn(async {
-                std::future::pending::<()>().await;
-                Ok(())
-            })),
-            socks_task: None,
-            environment_proxies,
-            completed: false,
+            runtime: Some(ProxyRuntime::SharedIngress { route }),
+            environment_proxies: Some(environment_proxies),
             windows_active_route: Some(active_route),
         }
     }
 
-    #[cfg(target_os = "windows")]
-    fn deactivate_windows_route(&mut self) {
+    pub async fn wait(mut self) -> Result<()> {
+        let result = match self.runtime.as_mut() {
+            Some(runtime) => runtime.wait().await,
+            None => Ok(()),
+        };
+        self.shutdown().await?;
+        result
+    }
+
+    pub async fn shutdown(mut self) -> Result<()> {
+        for listeners in self.stop_runtimes() {
+            listeners.shutdown().await;
+        }
+        Ok(())
+    }
+
+    fn stop_runtimes(&mut self) -> Vec<ProxyListeners> {
+        #[cfg(target_os = "windows")]
         if let Some(active_route) = self.windows_active_route.take() {
             active_route
                 .lock()
                 .unwrap_or_else(std::sync::PoisonError::into_inner)
                 .take();
         }
+
+        let environments = self
+            .environment_proxies
+            .take()
+            .map(|proxies| {
+                std::mem::take(
+                    &mut *proxies
+                        .lock()
+                        .unwrap_or_else(std::sync::PoisonError::into_inner),
+                )
+            })
+            .unwrap_or_default();
+        // Stop every runtime before awaiting any one of them. Cancellation of this
+        // cleanup future must not leave later environments accepting connections.
+        self.runtime
+            .take()
+            .into_iter()
+            .chain(environments.into_values().map(|proxy| proxy.runtime))
+            .filter_map(ProxyRuntime::stop)
+            .collect()
     }
-
-    pub async fn wait(mut self) -> Result<()> {
-        let http_task = self.http_task.take().context("missing http proxy task")?;
-        let socks_task = self.socks_task.take();
-        let http_result = http_task.await;
-        let socks_result = match socks_task {
-            Some(task) => Some(task.await),
-            None => None,
-        };
-        #[cfg(target_os = "windows")]
-        self.deactivate_windows_route();
-        self.completed = true;
-        abort_environment_proxies(self.environment_proxies.clone()).await;
-        http_result??;
-        if let Some(socks_result) = socks_result {
-            socks_result??;
-        }
-        Ok(())
-    }
-
-    pub async fn shutdown(mut self) -> Result<()> {
-        #[cfg(target_os = "windows")]
-        self.deactivate_windows_route();
-        abort_tasks(self.http_task.take(), self.socks_task.take()).await;
-        abort_environment_proxies(self.environment_proxies.clone()).await;
-        self.completed = true;
-        Ok(())
-    }
-}
-
-async fn abort_task(task: Option<JoinHandle<Result<()>>>) {
-    if let Some(task) = task {
-        task.abort();
-        let _ = task.await;
-    }
-}
-
-async fn abort_tasks(
-    http_task: Option<JoinHandle<Result<()>>>,
-    socks_task: Option<JoinHandle<Result<()>>>,
-) {
-    abort_task(http_task).await;
-    abort_task(socks_task).await;
-}
-
-async fn abort_environment_proxies(
-    environment_proxies: Arc<Mutex<HashMap<String, EnvironmentProxy>>>,
-) {
-    let proxies = {
-        let mut guard = environment_proxies
-            .lock()
-            .unwrap_or_else(std::sync::PoisonError::into_inner);
-        guard.drain().map(|(_, proxy)| proxy).collect::<Vec<_>>()
-    };
-    for proxy in proxies {
-        match proxy.runtime {
-            EnvironmentProxyRuntime::ListenerTasks {
-                http_task,
-                socks_task,
-            } => {
-                abort_task(Some(http_task)).await;
-                abort_task(socks_task).await;
-            }
-            #[cfg(target_os = "windows")]
-            EnvironmentProxyRuntime::SharedIngress { .. } => {}
-        }
-    }
-}
-
-#[cfg(target_os = "windows")]
-fn unregister_windows_ingress_environment_routes(
-    environment_proxies: &Arc<Mutex<HashMap<String, EnvironmentProxy>>>,
-) {
-    environment_proxies
-        .lock()
-        .unwrap_or_else(std::sync::PoisonError::into_inner)
-        .retain(|_, proxy| {
-            matches!(
-                &proxy.runtime,
-                EnvironmentProxyRuntime::ListenerTasks { .. }
-            )
-        });
 }
 
 impl Drop for NetworkProxyHandle {
     fn drop(&mut self) {
-        if self.completed {
-            return;
-        }
-        let http_task = self.http_task.take();
-        let socks_task = self.socks_task.take();
-        let environment_proxies = self.environment_proxies.clone();
-        #[cfg(target_os = "windows")]
-        {
-            self.deactivate_windows_route();
-            unregister_windows_ingress_environment_routes(&environment_proxies);
-        }
-        tokio::spawn(async move {
-            abort_tasks(http_task, socks_task).await;
-            abort_environment_proxies(environment_proxies).await;
-        });
+        drop(self.stop_runtimes());
     }
 }
 

codex-rs/network-proxy/src/socks5.rs

@@ -1,6 +1,7 @@
 use crate::attribution::BindConnectionAttribution;
 use crate::config::NetworkMode;
 use crate::connect_policy::TargetCheckedTcpConnector;
+use crate::connection_lifecycle::CancelOnShutdown;
 use crate::mitm;
 use crate::network_policy::BlockDecisionAuditEventArgs;
 use crate::network_policy::NetworkDecision;
@@ -29,6 +30,7 @@ use rama_core::error::BoxError;
 use rama_core::extensions::Extensions;
 use rama_core::extensions::ExtensionsMut;
 use rama_core::extensions::ExtensionsRef;
+use rama_core::graceful::ShutdownGuard;
 use rama_core::service::BoxService;
 use rama_core::service::service_fn;
 use rama_net::address::HostWithPort;
@@ -67,6 +69,7 @@ pub async fn run_socks5(
     policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
     environment_id: Option<String>,
     enable_socks5_udp: bool,
+    guard: ShutdownGuard,
 ) -> Result<()> {
     let listener = TcpListener::build()
         .bind(addr)
@@ -82,6 +85,7 @@ pub async fn run_socks5(
         policy_decider,
         environment_id,
         enable_socks5_udp,
+        guard,
     )
     .await
 }
@@ -92,6 +96,7 @@ pub async fn run_socks5_with_std_listener(
     policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
     environment_id: Option<String>,
     enable_socks5_udp: bool,
+    guard: ShutdownGuard,
 ) -> Result<()> {
     let listener =
         TcpListener::try_from(listener).context("convert std listener to SOCKS5 proxy listener")?;
@@ -101,6 +106,7 @@ pub async fn run_socks5_with_std_listener(
         policy_decider,
         environment_id,
         enable_socks5_udp,
+        guard,
     )
     .await
 }
@@ -111,6 +117,7 @@ async fn run_socks5_with_listener(
     policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
     environment_id: Option<String>,
     enable_socks5_udp: bool,
+    guard: ShutdownGuard,
 ) -> Result<()> {
     let addr = listener
         .local_addr()
@@ -131,12 +138,15 @@ async fn run_socks5_with_listener(
     }
 
     listener
-        .serve(socks5_proxy_service(
-            state,
-            policy_decider,
-            environment_id,
-            enable_socks5_udp,
-        ))
+        .serve_graceful(
+            guard,
+            CancelOnShutdown::new(socks5_proxy_service(
+                state,
+                policy_decider,
+                environment_id,
+                enable_socks5_udp,
+            )),
+        )
         .await;
     Ok(())
 }