From 5a65fd87d84215b24ea79100f8d47238f578a5f2 Mon Sep 17 00:00:00 2001 From: jelson-oai Date: Tue, 8 Sep 2026 18:28:03 +0000 Subject: [PATCH] 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 | 99 +++++++ .../connection_lifecycle/lifecycle_tests.rs | 266 +++++++++++++++++ .../src/connection_lifecycle/listeners.rs | 47 +++ .../src/connection_lifecycle/mod.rs | 12 + .../src/connection_lifecycle/scope.rs | 35 +++ .../src/connection_lifecycle/service.rs | 43 +++ codex-rs/network-proxy/src/http_proxy.rs | 25 +- codex-rs/network-proxy/src/lib.rs | 1 + codex-rs/network-proxy/src/proxy.rs | 278 ++++++++---------- codex-rs/network-proxy/src/socks5.rs | 22 +- 10 files changed, 654 insertions(+), 174 deletions(-) create mode 100644 codex-rs/network-proxy/src/connection_lifecycle/lifecycle_tests.rs create mode 100644 codex-rs/network-proxy/src/connection_lifecycle/listeners.rs create mode 100644 codex-rs/network-proxy/src/connection_lifecycle/mod.rs create mode 100644 codex-rs/network-proxy/src/connection_lifecycle/scope.rs create mode 100644 codex-rs/network-proxy/src/connection_lifecycle/service.rs diff --git a/codex-rs/core/tests/suite/network_approval.rs b/codex-rs/core/tests/suite/network_approval.rs index 5651c6f825..4d6f92e966 100644 --- a/codex-rs/core/tests/suite/network_approval.rs +++ b/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"), diff --git a/codex-rs/network-proxy/src/connection_lifecycle/lifecycle_tests.rs b/codex-rs/network-proxy/src/connection_lifecycle/lifecycle_tests.rs new file mode 100644 index 0000000000..60be13c30f --- /dev/null +++ b/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, +} + +#[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 { + 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 { + 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)?) +} diff --git a/codex-rs/network-proxy/src/connection_lifecycle/listeners.rs b/codex-rs/network-proxy/src/connection_lifecycle/listeners.rs new file mode 100644 index 0000000000..5c414bcd58 --- /dev/null +++ b/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>, +} + +impl ProxyListeners { + pub(crate) fn new() -> Self { + Self { + connections: ConnectionLifecycle::new(), + listeners: JoinSet::new(), + } + } + + pub(crate) fn spawn(&mut self, listener: F) + where + F: FnOnce(ShutdownGuard) -> Fut, + Fut: Future> + 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; + } +} diff --git a/codex-rs/network-proxy/src/connection_lifecycle/mod.rs b/codex-rs/network-proxy/src/connection_lifecycle/mod.rs new file mode 100644 index 0000000000..b0219bbcd4 --- /dev/null +++ b/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; diff --git a/codex-rs/network-proxy/src/connection_lifecycle/scope.rs b/codex-rs/network-proxy/src/connection_lifecycle/scope.rs new file mode 100644 index 0000000000..2af4b500c8 --- /dev/null +++ b/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>, + 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; + } +} diff --git a/codex-rs/network-proxy/src/connection_lifecycle/service.rs b/codex-rs/network-proxy/src/connection_lifecycle/service.rs new file mode 100644 index 0000000000..2444b9cb67 --- /dev/null +++ b/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 { + inner: S, +} + +impl CancelOnShutdown { + pub(crate) fn new(inner: S) -> Self { + Self { inner } + } +} + +impl Service for CancelOnShutdown +where + S: Service, + Request: ExtensionsRef + Send + 'static, +{ + type Output = (); + type Error = S::Error; + + async fn serve(&self, request: Request) -> Result<(), Self::Error> { + let guard = request + .extensions() + .get::() + .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, + } + } +} diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index 8d9670a41d..da9bc494a8 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/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>, environment_id: Option, + 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>, environment_id: Option, + 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>, environment_id: Option, + 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) diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index b2affe0b8c..38ac654d46 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/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; diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index 8172645e13..42a2a8135b 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/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>, - socks_task: Option>>, - }, +enum ProxyRuntime { + Listeners(ProxyListeners), #[cfg(target_os = "windows")] - SharedIngress { _route: Arc }, + SharedIngress { + route: Arc, + }, } -impl EnvironmentProxyRuntime { +impl ProxyRuntime { #[cfg(target_os = "windows")] fn network_proxy_restricting_sid(&self) -> Option { 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 { + 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>>, - socks_task: Option>>, - environment_proxies: Arc>>, - completed: bool, + runtime: Option, + environment_proxies: Option>>>, #[cfg(target_os = "windows")] windows_active_route: Option>>>>, } @@ -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, active_route: Arc>>>, environment_proxies: Arc>>, ) -> 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 { + #[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>>) { - if let Some(task) = task { - task.abort(); - let _ = task.await; - } -} - -async fn abort_tasks( - http_task: Option>>, - socks_task: Option>>, -) { - abort_task(http_task).await; - abort_task(socks_task).await; -} - -async fn abort_environment_proxies( - environment_proxies: Arc>>, -) { - let proxies = { - let mut guard = environment_proxies - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - guard.drain().map(|(_, proxy)| proxy).collect::>() - }; - 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>>, -) { - 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()); } } diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index 9817e880d9..098d3feb27 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/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>, environment_id: Option, 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>, environment_id: Option, 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>, environment_id: Option, 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(()) }