Enforce network policy throughout HTTP and WebSocket requests (#47389)

## Why Destination restrictions and policy revocation must remain effective during redirects, response body reads, and established WebSocket traffic. Policy denials must also survive error handling so callers do not retry them or report a revoked operation as successful. ## What changed - Route managed HTTP clients through a shared `RequestBuilder` and policy-aware execution, checking each redirect destination before route resolution and retaining a network permit while consuming response bodies. - Guard WebSocket connection setup, reads, and writes with revocable permits, including independent wakeups for split readers and writers. - Preserve `TransportError::Policy` through HTTP, SSE, and realtime error handling, and treat policy denials as non-retryable. - Keep the supplied network policy in plugin startup HTTP requests and propagate response body failures from backend and plugin requests. ## Testing Add regression coverage for rejection before connecting, revocation during redirect routing and streamed body reads, split WebSocket revocation, realtime writer error propagation, and revoked plugin upload responses. GitOrigin-RevId: fba36700444c98a8364c09b4dc5452c4d78a90fb
2026-11-08 20:40:33 +00:00 · 2026-09-22 23:37:39 +00:00
parent 506a328dab
commit 888e02db34
57 changed files with 1661 additions and 924 deletions
@@ -23,6 +23,7 @@ use crate::error::ApiError;
 use crate::provider::Provider;
 use codex_client::backoff;
 use codex_http_client::HttpClientFactory;
+use codex_http_client::NetworkPolicyDenied;
 use codex_http_client::maybe_build_rustls_client_config_with_custom_ca;
 use codex_protocol::protocol::ConversationTextParams;
 use codex_protocol::protocol::ConversationTextRole;
@@ -65,6 +66,7 @@ const TRUNCATED_TRANSCRIPT_PREFIX: &str = "…";
 struct WsStream {
    tx_command: mpsc::Sender<WsCommand>,
    pump_task: tokio::task::JoinHandle<()>,
+    terminal_policy_error: tokio::sync::watch::Receiver<Option<NetworkPolicyDenied>>,
 }

 enum WsCommand {
@@ -84,8 +86,15 @@ impl WsStream {
        let (tx_command, mut rx_command) = mpsc::channel::<WsCommand>(32);
        let (tx_message, rx_message) = async_channel::unbounded::<Result<Message, WsError>>();

+        let (terminal_policy_tx, terminal_policy_error) = tokio::sync::watch::channel(None);
        let pump_task = tokio::spawn(async move {
            let mut inner = inner;
+            let report_write_policy_error = |error: &WsError| {
+                if let Some(denied) = codex_websocket_client::network_policy_denial(error) {
+                    terminal_policy_tx.send_replace(Some(denied));
+                    let _ = tx_message.try_send(Err(Self::policy_error(denied)));
+                }
+            };
            loop {
                tokio::select! {
                    command = rx_command.recv() => {
@@ -99,6 +108,7 @@ impl WsStream {
                                let should_break = result.is_err();
                                if let Err(err) = &result {
                                    error!("realtime websocket send failed: {err}");
+                                    report_write_policy_error(err);
                                }
                                let _ = tx_result.send(result);
                                if should_break {
@@ -110,6 +120,7 @@ impl WsStream {
                                let result = inner.send(Message::Close(None)).await;
                                if let Err(err) = &result {
                                    error!("realtime websocket close failed: {err}");
+                                    report_write_policy_error(err);
                                }
                                let _ = tx_result.send(result);
                                break;
@@ -125,6 +136,7 @@ impl WsStream {
                                trace!(payload_len = payload.len(), "realtime websocket received ping");
                                if let Err(err) = inner.send(Message::Pong(payload)).await {
                                    error!("realtime websocket failed to send pong: {err}");
+                                    terminal_policy_tx.send_replace(codex_websocket_client::network_policy_denial(&err));
                                    let _ = tx_message.send(Err(err)).await;
                                    break;
                                }
@@ -162,6 +174,7 @@ impl WsStream {
                            }
                            Err(err) => {
                                error!("realtime websocket receive failed: {err}");
+                                    terminal_policy_tx.send_replace(codex_websocket_client::network_policy_denial(&err));
                                let _ = tx_message.send(Err(err)).await;
                                break;
                            }
@@ -176,6 +189,7 @@ impl WsStream {
            Self {
                tx_command,
                pump_task,
+                terminal_policy_error,
            },
            rx_message,
        )
@@ -187,9 +201,23 @@ impl WsStream {
    ) -> Result<(), WsError> {
        let (tx_result, rx_result) = oneshot::channel();
        if self.tx_command.send(make_command(tx_result)).await.is_err() {
-            return Err(WsError::ConnectionClosed);
+            return Err(self.closed_error());
        }
-        rx_result.await.unwrap_or(Err(WsError::ConnectionClosed))
+        rx_result.await.unwrap_or_else(|_| Err(self.closed_error()))
+    }
+
+    fn closed_error(&self) -> WsError {
+        match *self.terminal_policy_error.borrow() {
+            Some(denied) => Self::policy_error(denied),
+            None => WsError::ConnectionClosed,
+        }
+    }
+
+    fn policy_error(denied: NetworkPolicyDenied) -> WsError {
+        WsError::Io(std::io::Error::new(
+            std::io::ErrorKind::PermissionDenied,
+            denied,
+        ))
    }

    async fn send(&self, message: Message) -> Result<(), WsError> {
@@ -481,8 +509,8 @@ impl RealtimeWebsocketWriter {

    pub async fn send_payload(&self, payload: String) -> Result<(), ApiError> {
        if self.is_closed.load(Ordering::SeqCst) {
-            return Err(ApiError::Stream(
-                "realtime websocket connection is closed".to_string(),
+            return Err(map_realtime_websocket_connect_error(
+                self.stream.closed_error(),
            ));
        }

@@ -490,7 +518,7 @@ impl RealtimeWebsocketWriter {
        self.stream
            .send(Message::Text(payload.into()))
            .await
-            .map_err(|err| ApiError::Stream(format!("failed to send realtime request: {err}")))?;
+            .map_err(map_realtime_websocket_connect_error)?;
        Ok(())
    }
 }
@@ -519,9 +547,7 @@ impl RealtimeWebsocketEvents {
                Ok(Err(err)) => {
                    self.is_closed.store(true, Ordering::SeqCst);
                    error!("realtime websocket read failed: {err}");
-                    return Err(ApiError::Stream(format!(
-                        "failed to read websocket message: {err}"
-                    )));
+                    return Err(map_realtime_websocket_connect_error(err));
                }
                Err(_) => {
                    self.is_closed.store(true, Ordering::SeqCst);
@@ -918,6 +944,9 @@ impl RealtimeWebsocketClient {
                .await;
            match result {
                Ok(connection) => return Ok(connection),
+                Err(err @ ApiError::Transport(codex_http_client::TransportError::Policy(_))) => {
+                    return Err(err);
+                }
                Err(err) if webrtc_sideband_session_ended(&err) => return Err(err),
                Err(err) if attempt < self.provider.retry.max_attempts => {
                    let delay = backoff(self.provider.retry.base_delay, attempt + 1);
@@ -1081,6 +1110,9 @@ fn webrtc_sideband_session_ended(err: &ApiError) -> bool {
 }

 fn map_realtime_websocket_connect_error(err: WsError) -> ApiError {
+    if let Some(denied) = codex_websocket_client::network_policy_denial(&err) {
+        return ApiError::Transport(codex_http_client::TransportError::Policy(denied));
+    }
    match err {
        WsError::Http(response) => ApiError::Api {
            status: response.status(),
@@ -1261,7 +1293,7 @@ mod tests {
    use super::*;
    use crate::endpoint::realtime_websocket::protocol::RealtimeTranscriptEntry;
    use crate::provider::RetryConfig;
-    use codex_http_client::OutboundProxyPolicy;
+    use codex_http_client::OutboundProxyPolicy::ReqwestDefault;
    use codex_protocol::protocol::RealtimeHandoffRequested;
    use codex_protocol::protocol::RealtimeInputAudioSpeechStarted;
    use codex_protocol::protocol::RealtimeNoopRequested;
@@ -1283,6 +1315,10 @@ mod tests {
    use tokio_tungstenite::tungstenite::Message;
    use tungstenite::protocol::CloseFrame;

+    fn test_client(provider: Provider) -> RealtimeWebsocketClient {
+        RealtimeWebsocketClient::new(provider, HttpClientFactory::new(ReqwestDefault))
+    }
+
    #[test]
    fn parse_session_updated_event() {
        let payload = json!({
@@ -1531,6 +1567,59 @@ mod tests {
        events.next_event().await.is_err()
    }

+    #[tokio::test]
+    async fn writer_policy_failure_reaches_realtime_events() {
+        use codex_http_client::DestinationPolicy::Unrestricted;
+        use codex_http_client::NetworkPolicyController;
+        use codex_http_client::NetworkPolicyDenied::Revoked;
+        use codex_http_client::TransportError::Policy;
+        use codex_websocket_client::network_policy_denial;
+        use tokio::time::timeout;
+
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let url = format!("ws://{}/", listener.local_addr().unwrap());
+        let (writing, started) = oneshot::channel();
+        let server = tokio::spawn(async move {
+            let (socket, _) = listener.accept().await.unwrap();
+            let socket = accept_async(socket).await.unwrap();
+            socket.get_ref().peek(&mut [0]).await.unwrap();
+            writing.send(()).unwrap();
+            std::future::pending::<()>().await;
+        });
+        let controller = NetworkPolicyController::default();
+        let policy = controller.policy();
+        controller.publish(policy.revision(), Unrestricted);
+        let factory = HttpClientFactory::new(ReqwestDefault).with_network_policy(policy.clone());
+        let connector = WebSocketConnector::new(&factory).unwrap();
+        let request = url.into_client_request().unwrap();
+        let connected = connector.connect(request, websocket_config());
+        let (socket, _) = connected.await.unwrap();
+        let (stream, rx_message) = WsStream::new(socket);
+        let parser = RealtimeEventParser::V1;
+        let transcript = RealtimeTranscriptState::default();
+        let connection = RealtimeWebsocketConnection::new(stream, rx_message, parser, transcript);
+        let stream = Arc::clone(&connection.writer.stream);
+        let payload = Message::Binary(vec![0; 8 * 1024 * 1024].into());
+        let writer = tokio::spawn(async move {
+            loop {
+                if let Err(error) = stream.send(payload.clone()).await {
+                    return error;
+                }
+            }
+        });
+        let deadline = Duration::from_secs(/*secs*/ 5);
+        timeout(deadline, started).await.unwrap().unwrap();
+        policy.invalidate();
+        let writer = timeout(deadline, writer).await.unwrap();
+        let error = writer.unwrap();
+        assert_eq!(network_policy_denial(&error), Some(Revoked));
+        let next = timeout(deadline, connection.next_event());
+        let error = next.await.unwrap().unwrap_err();
+        assert!(matches!(error, ApiError::Transport(Policy(Revoked))));
+        assert!(connection.next_event().await.unwrap().is_none());
+        server.abort();
+    }
+
    #[tokio::test]
    async fn websocket_termination_classification_preserves_protocol_behavior() {
        use CloseCode::Away;
@@ -2220,23 +2309,20 @@ mod tests {

    #[test]
    fn webrtc_frameless_sideband_ignores_provider_base_url() {
-        let client = RealtimeWebsocketClient::new(
-            Provider {
-                name: "chatgpt".to_string(),
-                base_url: "https://chatgpt.com/backend-api/codex".to_string(),
-                query_params: None,
-                headers: HeaderMap::new(),
-                retry: RetryConfig {
-                    max_attempts: 0,
-                    base_delay: Duration::ZERO,
-                    retry_429: false,
-                    retry_5xx: false,
-                    retry_transport: false,
-                },
-                stream_idle_timeout: Duration::from_secs(/*secs*/ 5),
+        let client = test_client(Provider {
+            name: "chatgpt".to_string(),
+            base_url: "https://chatgpt.com/backend-api/codex".to_string(),
+            query_params: None,
+            headers: HeaderMap::new(),
+            retry: RetryConfig {
+                max_attempts: 0,
+                base_delay: Duration::ZERO,
+                retry_429: false,
+                retry_5xx: false,
+                retry_transport: false,
            },
-            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
-        );
+            stream_idle_timeout: Duration::from_secs(5),
+        });

        let url = client
            .webrtc_sideband_url(
@@ -2433,10 +2519,7 @@ mod tests {
            },
            stream_idle_timeout: Duration::from_secs(5),
        };
-        let client = RealtimeWebsocketClient::new(
-            provider,
-            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
-        );
+        let client = test_client(provider);
        let connection = client
            .connect(
                RealtimeSessionConfig {
@@ -2762,10 +2845,7 @@ mod tests {
            },
            stream_idle_timeout: Duration::from_secs(5),
        };
-        let client = RealtimeWebsocketClient::new(
-            provider,
-            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
-        );
+        let client = test_client(provider);
        let connection = client
            .connect(
                RealtimeSessionConfig {
@@ -2892,10 +2972,7 @@ mod tests {
            },
            stream_idle_timeout: Duration::from_secs(5),
        };
-        let client = RealtimeWebsocketClient::new(
-            provider,
-            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
-        );
+        let client = test_client(provider);
        let connection = client
            .connect(
                RealtimeSessionConfig {
@@ -3001,10 +3078,7 @@ mod tests {
            },
            stream_idle_timeout: Duration::from_secs(5),
        };
-        let client = RealtimeWebsocketClient::new(
-            provider,
-            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
-        );
+        let client = test_client(provider);
        let connection = client
            .connect(
                RealtimeSessionConfig {
@@ -3096,10 +3170,7 @@ mod tests {
            },
            stream_idle_timeout: Duration::from_secs(5),
        };
-        let client = RealtimeWebsocketClient::new(
-            provider,
-            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
-        );
+        let client = test_client(provider);
        let connection = client
            .connect(
                RealtimeSessionConfig {
@@ -125,17 +125,22 @@ impl WsStream {
    }

    async fn request(
-        &self,
+        &mut self,
        make_command: impl FnOnce(oneshot::Sender<Result<(), WsError>>) -> WsCommand,
    ) -> Result<(), WsError> {
        let (tx_result, rx_result) = oneshot::channel();
-        if self.tx_command.send(make_command(tx_result)).await.is_err() {
-            return Err(WsError::ConnectionClosed);
+        if self.tx_command.send(make_command(tx_result)).await.is_ok()
+            && let Ok(result) = rx_result.await
+        {
+            return result;
        }
-        rx_result.await.unwrap_or(Err(WsError::ConnectionClosed))
+        while let Some(message) = self.rx_message.recv().await {
+            message?;
+        }
+        Err(WsError::ConnectionClosed)
    }

-    async fn send(&self, message: Message) -> Result<(), WsError> {
+    async fn send(&mut self, message: Message) -> Result<(), WsError> {
        self.request(|tx_result| WsCommand::Send { message, tx_result })
            .await
    }
@@ -441,9 +446,7 @@ impl ResponsesWebsocketClient {
            .ok()
            .flatten()
            .transpose()
-            .map_err(|err| {
-                ApiError::Stream(format!("failed to read websocket probe event: {err}"))
-            })?
+            .map_err(map_ws_stream_error)?
            .and_then(immediate_close_from_message);

        Ok(ResponsesWebsocketProbe {
@@ -548,7 +551,17 @@ fn websocket_config() -> WebSocketConfig {
    config
 }

+fn map_ws_stream_error(error: WsError) -> ApiError {
+    match codex_websocket_client::network_policy_denial(&error) {
+        Some(denied) => ApiError::Transport(TransportError::Policy(denied)),
+        None => ApiError::Stream(error.to_string()),
+    }
+}
+
 fn map_ws_error(err: WsError, url: &Url) -> ApiError {
+    if let Some(denied) = codex_websocket_client::network_policy_denial(&err) {
+        return ApiError::Transport(TransportError::Policy(denied));
+    }
    match err {
        WsError::Http(response) => {
            let status = response.status();
@@ -696,7 +709,7 @@ async fn run_websocket_response_stream(
        let message = match response {
            Ok(Some(Ok(msg))) => msg,
            Ok(Some(Err(err))) => {
-                return Err(ApiError::Stream(err.to_string()));
+                return Err(map_ws_stream_error(err));
            }
            Ok(None) => {
                return Err(ApiError::Stream(
@@ -870,7 +883,7 @@ fn safety_buffering_for_event(
 }

 async fn send_websocket_request(
-    ws_stream: &WsStream,
+    ws_stream: &mut WsStream,
    request_text: String,
    idle_timeout: Duration,
    telemetry: Option<&Arc<dyn WebsocketTelemetry>>,
@@ -883,9 +896,7 @@ async fn send_websocket_request(
    )
    .await
    .map_err(|_| ApiError::Stream("idle timeout sending websocket request".into()))
-    .and_then(|result| {
-        result.map_err(|err| ApiError::Stream(format!("failed to send websocket request: {err}")))
-    });
+    .and_then(|result| result.map_err(map_ws_stream_error));

    if let Some(t) = telemetry.as_ref() {
        t.on_ws_request(