Merge aac057d136 into sapling-pr-archive-bolinfest

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

@@ -11,6 +11,8 @@ use clap::Parser;
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::sync::atomic::AtomicU16;
+use std::sync::atomic::Ordering;
 use tokio::task::JoinHandle;
 use tracing::warn;
 
@@ -18,14 +20,32 @@ use tracing::warn;
 #[command(name = "codex-network-proxy", about = "Codex network sandbox proxy")]
 pub struct Args {}
 
-#[derive(Clone, Default)]
+const FALLBACK_EPHEMERAL_PORT_START: u16 = 40_000;
+static FALLBACK_EPHEMERAL_PORT: AtomicU16 = AtomicU16::new(FALLBACK_EPHEMERAL_PORT_START);
+
+#[derive(Clone)]
 pub struct NetworkProxyBuilder {
     state: Option<Arc<NetworkProxyState>>,
     http_addr: Option<SocketAddr>,
+    socks_addr: Option<SocketAddr>,
     admin_addr: Option<SocketAddr>,
+    managed_by_codex: bool,
     policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
 }
 
+impl Default for NetworkProxyBuilder {
+    fn default() -> Self {
+        Self {
+            state: None,
+            http_addr: None,
+            socks_addr: None,
+            admin_addr: None,
+            managed_by_codex: true,
+            policy_decider: None,
+        }
+    }
+}
+
 impl NetworkProxyBuilder {
     pub fn state(mut self, state: Arc<NetworkProxyState>) -> Self {
         self.state = Some(state);
@@ -37,11 +57,21 @@ impl NetworkProxyBuilder {
         self
     }
 
+    pub fn socks_addr(mut self, addr: SocketAddr) -> Self {
+        self.socks_addr = Some(addr);
+        self
+    }
+
     pub fn admin_addr(mut self, addr: SocketAddr) -> Self {
         self.admin_addr = Some(addr);
         self
     }
 
+    pub fn managed_by_codex(mut self, managed_by_codex: bool) -> Self {
+        self.managed_by_codex = managed_by_codex;
+        self
+    }
+
     pub fn policy_decider<D>(mut self, decider: D) -> Self
     where
         D: NetworkPolicyDecider,
@@ -62,12 +92,27 @@ impl NetworkProxyBuilder {
             )
         })?;
         let current_cfg = state.current_cfg().await?;
-        let runtime = config::resolve_runtime(&current_cfg)?;
+        let (requested_http_addr, requested_socks_addr, requested_admin_addr) =
+            if self.managed_by_codex {
+                (
+                    reserve_loopback_ephemeral_addr()?,
+                    reserve_loopback_ephemeral_addr()?,
+                    reserve_loopback_ephemeral_addr()?,
+                )
+            } else {
+                let runtime = config::resolve_runtime(&current_cfg)?;
+                (
+                    self.http_addr.unwrap_or(runtime.http_addr),
+                    self.socks_addr.unwrap_or(runtime.socks_addr),
+                    self.admin_addr.unwrap_or(runtime.admin_addr),
+                )
+            };
+
         // Reapply bind clamping for caller overrides so unix-socket proxying stays loopback-only.
         let (http_addr, socks_addr, admin_addr) = config::clamp_bind_addrs(
-            self.http_addr.unwrap_or(runtime.http_addr),
-            runtime.socks_addr,
-            self.admin_addr.unwrap_or(runtime.admin_addr),
+            requested_http_addr,
+            requested_socks_addr,
+            requested_admin_addr,
             &current_cfg.network,
         );
 
@@ -81,6 +126,33 @@ impl NetworkProxyBuilder {
     }
 }
 
+fn reserve_loopback_ephemeral_addr() -> Result<SocketAddr> {
+    match std::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))) {
+        Ok(listener) => listener
+            .local_addr()
+            .context("failed to get loopback ephemeral port"),
+        Err(err) => {
+            let port = next_fallback_ephemeral_port();
+            warn!(
+                "failed to reserve loopback ephemeral port ({err}); falling back to 127.0.0.1:{port}"
+            );
+            Ok(SocketAddr::from(([127, 0, 0, 1], port)))
+        }
+    }
+}
+
+fn next_fallback_ephemeral_port() -> u16 {
+    // Keep ports nonzero and in a high range to reduce accidental collisions in test environments
+    // where loopback binds are restricted and we cannot reserve an ephemeral port directly.
+    let next = FALLBACK_EPHEMERAL_PORT.fetch_add(1, Ordering::Relaxed);
+    if next == u16::MAX {
+        FALLBACK_EPHEMERAL_PORT.store(FALLBACK_EPHEMERAL_PORT_START + 1, Ordering::Relaxed);
+        FALLBACK_EPHEMERAL_PORT_START
+    } else {
+        next
+    }
+}
+
 #[derive(Clone)]
 pub struct NetworkProxy {
     state: Arc<NetworkProxyState>,
@@ -241,3 +313,55 @@ impl Drop for NetworkProxyHandle {
         });
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::NetworkProxySettings;
+    use crate::state::network_proxy_state_for_policy;
+
+    #[tokio::test]
+    async fn managed_proxy_builder_uses_loopback_ephemeral_ports() {
+        let state = Arc::new(network_proxy_state_for_policy(
+            NetworkProxySettings::default(),
+        ));
+        let proxy = NetworkProxy::builder().state(state).build().await.unwrap();
+
+        assert!(proxy.http_addr.ip().is_loopback());
+        assert!(proxy.socks_addr.ip().is_loopback());
+        assert!(proxy.admin_addr.ip().is_loopback());
+        assert_ne!(proxy.http_addr.port(), 0);
+        assert_ne!(proxy.socks_addr.port(), 0);
+        assert_ne!(proxy.admin_addr.port(), 0);
+    }
+
+    #[tokio::test]
+    async fn non_codex_managed_proxy_builder_uses_configured_ports() {
+        let settings = NetworkProxySettings {
+            proxy_url: "http://127.0.0.1:43128".to_string(),
+            socks_url: "http://127.0.0.1:48081".to_string(),
+            admin_url: "http://127.0.0.1:48080".to_string(),
+            ..NetworkProxySettings::default()
+        };
+        let state = Arc::new(network_proxy_state_for_policy(settings));
+        let proxy = NetworkProxy::builder()
+            .state(state)
+            .managed_by_codex(false)
+            .build()
+            .await
+            .unwrap();
+
+        assert_eq!(
+            proxy.http_addr,
+            "127.0.0.1:43128".parse::<SocketAddr>().unwrap()
+        );
+        assert_eq!(
+            proxy.socks_addr,
+            "127.0.0.1:48081".parse::<SocketAddr>().unwrap()
+        );
+        assert_eq!(
+            proxy.admin_addr,
+            "127.0.0.1:48080".parse::<SocketAddr>().unwrap()
+        );
+    }
+}