[codex] Bind network attribution at shared proxy ingress

codex-rs/core/src/network_policy_decision_tests.rs

@@ -161,6 +161,7 @@ fn denied_network_policy_message_requires_deny_decision() {
         method: Some("GET".to_string()),
         mode: None,
         protocol: "http".to_string(),
+        execution_id: None,
         decision: Some("ask".to_string()),
         source: Some("decider".to_string()),
         port: Some(80),
@@ -178,6 +179,7 @@ fn denied_network_policy_message_for_denylist_block_is_explicit() {
         method: Some("GET".to_string()),
         mode: None,
         protocol: "http".to_string(),
+        execution_id: None,
         decision: Some("deny".to_string()),
         source: Some("baseline_policy".to_string()),
         port: Some(80),

codex-rs/core/src/tools/network_approval_tests.rs

@@ -280,6 +280,7 @@ fn denied_blocked_request(host: &str) -> BlockedRequest {
         method: None,
         mode: None,
         protocol: "http".to_string(),
+        execution_id: None,
         decision: Some("deny".to_string()),
         source: Some("decider".to_string()),
         port: Some(80),

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

@@ -0,0 +1,166 @@
+use crate::network_policy::NetworkRequestContext;
+use rama_core::Service;
+use rama_core::error::BoxError;
+use rama_core::extensions::ExtensionsMut;
+use rama_tcp::TcpStream;
+use std::collections::HashMap;
+use std::io;
+use std::io::Write;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::Duration;
+use tokio::io::AsyncReadExt;
+
+/// Internal handoff to a trusted proxy bridge.
+#[doc(hidden)]
+pub const PROXY_ATTRIBUTION_TOKEN_ENV_KEY: &str = "CODEX_NETWORK_PROXY_ATTRIBUTION";
+
+const ATTRIBUTION_FRAME_MAGIC: &[u8; 8] = b"\0CDXPXY1";
+const MAX_ATTRIBUTION_TOKEN_LEN: usize = 128;
+const ATTRIBUTION_FRAME_TIMEOUT: Duration = Duration::from_secs(3);
+
+#[derive(Clone, Default)]
+pub(crate) struct AttributionRegistry {
+    contexts: Arc<Mutex<HashMap<String, NetworkRequestContext>>>,
+}
+
+impl AttributionRegistry {
+    pub(crate) fn context_for_token(&self, token: &str) -> Option<NetworkRequestContext> {
+        self.contexts
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .get(token)
+            .cloned()
+    }
+
+    pub(crate) fn clear(&self) {
+        self.contexts
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clear();
+    }
+
+    #[cfg(test)]
+    fn register(&self, token: impl Into<String>, context: NetworkRequestContext) {
+        self.contexts
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .insert(token.into(), context);
+    }
+}
+
+#[derive(Clone)]
+pub(crate) enum ConnectionAttribution {
+    Fixed(NetworkRequestContext),
+    Registry(AttributionRegistry),
+}
+
+pub(crate) struct BindConnectionAttribution<S> {
+    inner: S,
+    attribution: ConnectionAttribution,
+}
+
+impl<S> BindConnectionAttribution<S> {
+    pub(crate) fn new(inner: S, attribution: ConnectionAttribution) -> Self {
+        Self { inner, attribution }
+    }
+}
+
+impl<S> Service<TcpStream> for BindConnectionAttribution<S>
+where
+    S: Service<TcpStream>,
+    S::Error: Into<BoxError>,
+{
+    type Output = S::Output;
+    type Error = BoxError;
+
+    async fn serve(&self, mut stream: TcpStream) -> Result<Self::Output, Self::Error> {
+        let context = match &self.attribution {
+            ConnectionAttribution::Fixed(context) => context.clone(),
+            ConnectionAttribution::Registry(registry) => {
+                match read_attribution_token(&mut stream).await? {
+                    Some(token) => registry.context_for_token(&token).ok_or_else(|| {
+                        io::Error::new(
+                            io::ErrorKind::PermissionDenied,
+                            "unknown network proxy attribution token",
+                        )
+                    })?,
+                    None => NetworkRequestContext::default(),
+                }
+            }
+        };
+        stream.extensions_mut().insert(context);
+        self.inner.serve(stream).await.map_err(Into::into)
+    }
+}
+
+async fn read_attribution_token(stream: &mut TcpStream) -> Result<Option<String>, BoxError> {
+    let mut marker = [0_u8; 1];
+    let read = stream.stream.peek(&mut marker).await?;
+    if read == 0 {
+        return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "empty proxy connection").into());
+    }
+    if marker[0] != ATTRIBUTION_FRAME_MAGIC[0] {
+        return Ok(None);
+    }
+
+    let token = tokio::time::timeout(ATTRIBUTION_FRAME_TIMEOUT, async {
+        let mut magic = [0_u8; ATTRIBUTION_FRAME_MAGIC.len()];
+        stream.read_exact(&mut magic).await?;
+        if &magic != ATTRIBUTION_FRAME_MAGIC {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "invalid network proxy attribution frame",
+            ));
+        }
+
+        let token_len = stream.read_u16().await? as usize;
+        if token_len == 0 || token_len > MAX_ATTRIBUTION_TOKEN_LEN {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "invalid network proxy attribution token length",
+            ));
+        }
+        let mut token = vec![0_u8; token_len];
+        stream.read_exact(&mut token).await?;
+        String::from_utf8(token).map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                "network proxy attribution token is not UTF-8",
+            )
+        })
+    })
+    .await
+    .map_err(|_| {
+        io::Error::new(
+            io::ErrorKind::TimedOut,
+            "network proxy attribution frame timed out",
+        )
+    })??;
+
+    Ok(Some(token))
+}
+
+/// Writes the trusted bridge preface consumed by the shared proxy ingress.
+#[doc(hidden)]
+pub fn write_attribution_frame(writer: &mut impl Write, token: &str) -> io::Result<()> {
+    if token.is_empty() || token.len() > MAX_ATTRIBUTION_TOKEN_LEN {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidInput,
+            "invalid network proxy attribution token length",
+        ));
+    }
+    let token_len = u16::try_from(token.len()).map_err(|_| {
+        io::Error::new(
+            io::ErrorKind::InvalidInput,
+            "network proxy attribution token is too long",
+        )
+    })?;
+    writer.write_all(ATTRIBUTION_FRAME_MAGIC)?;
+    writer.write_all(&token_len.to_be_bytes())?;
+    writer.write_all(token.as_bytes())
+}
+
+#[cfg(test)]
+#[path = "attribution_tests.rs"]
+mod tests;

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

@@ -0,0 +1,58 @@
+use super::AttributionRegistry;
+use super::BindConnectionAttribution;
+use super::ConnectionAttribution;
+use super::write_attribution_frame;
+use crate::network_policy::NetworkRequestContext;
+use pretty_assertions::assert_eq;
+use rama_core::Service;
+use rama_core::error::BoxError;
+use rama_core::extensions::ExtensionsRef;
+use rama_core::service::service_fn;
+use rama_tcp::TcpStream as RamaTcpStream;
+use std::io;
+use tokio::io::AsyncWriteExt;
+use tokio::net::TcpListener;
+use tokio::net::TcpStream;
+
+#[test]
+fn attribution_frame_has_bounded_binary_prefix() -> io::Result<()> {
+    let mut frame = Vec::new();
+    write_attribution_frame(&mut frame, "token-1")?;
+
+    assert_eq!(&frame[..8], b"\0CDXPXY1");
+    assert_eq!(u16::from_be_bytes([frame[8], frame[9]]), 7);
+    assert_eq!(&frame[10..], b"token-1");
+    Ok(())
+}
+
+#[tokio::test]
+async fn framed_connection_receives_registered_attribution() -> Result<(), BoxError> {
+    let registry = AttributionRegistry::default();
+    let expected = NetworkRequestContext {
+        environment_id: Some("local".to_string()),
+        execution_id: Some("execution-1".to_string()),
+    };
+    registry.register("token-1", expected.clone());
+
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let addr = listener.local_addr()?;
+    let client = tokio::spawn(async move {
+        let mut stream = TcpStream::connect(addr).await?;
+        let mut frame = Vec::new();
+        write_attribution_frame(&mut frame, "token-1")?;
+        stream.write_all(&frame).await
+    });
+
+    let (stream, _) = listener.accept().await?;
+    let service = BindConnectionAttribution::new(
+        service_fn(|stream: RamaTcpStream| async move {
+            Ok::<_, io::Error>(stream.extensions().get::<NetworkRequestContext>().cloned())
+        }),
+        ConnectionAttribution::Registry(registry),
+    );
+    let actual = service.serve(RamaTcpStream::new(stream)).await?;
+    client.await??;
+
+    assert_eq!(actual, Some(expected));
+    Ok(())
+}