diff --git a/codex-rs/network-proxy/src/mitm.rs b/codex-rs/network-proxy/src/mitm.rs index 9056e48f50..af2e0ac066 100644 --- a/codex-rs/network-proxy/src/mitm.rs +++ b/codex-rs/network-proxy/src/mitm.rs @@ -33,6 +33,7 @@ use rama_http::Response; use rama_http::StatusCode; use rama_http::Uri; use rama_http::header::HOST; +use rama_http::header::HeaderName; use rama_http::layer::remove_header::RemoveRequestHeaderLayer; use rama_http::layer::remove_header::RemoveResponseHeaderLayer; use rama_http_backend::server::HttpServer; @@ -85,6 +86,10 @@ enum MitmPolicyDecision { const MITM_INSPECT_BODIES: bool = false; const MITM_MAX_BODY_BYTES: usize = 4096; +const CREDENTIAL_ROUTE_CONNECTOR_ID_HEADER: &str = + "x-openai-internal-credential-route-connector-id"; +const CREDENTIAL_ROUTE_LINK_ID_HEADER: &str = "x-openai-internal-credential-route-link-id"; +const CREDENTIAL_ROUTE_REQUEST_URL_HEADER: &str = "x-openai-internal-credential-route-request-url"; impl std::fmt::Debug for MitmState { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -234,7 +239,12 @@ async fn forward_request(req: Request, request_ctx: &MitmRequestContext) -> Resu && let Some(credentialed_route_proxy) = actions.credentialed_route_proxy.as_ref() { let (proxy_uri, proxy_authority) = - credentialed_route_proxy_request(credentialed_route_proxy, &original_request_uri)?; + credentialed_route_proxy_request(credentialed_route_proxy)?; + apply_credentialed_route_proxy_headers( + &mut parts.headers, + credentialed_route_proxy, + &original_request_uri, + )?; parts.uri = proxy_uri; parts .headers @@ -562,19 +572,33 @@ fn build_https_uri(authority: &str, path: &str) -> Result { fn credentialed_route_proxy_request( action: &crate::mitm_hook::CredentialedRouteProxyAction, - request_uri: &Uri, ) -> Result<(Uri, String)> { - let mut proxy_url = Url::parse(&action.proxy_url) + let proxy_url = Url::parse(&action.proxy_url) .with_context(|| format!("invalid credentialed route proxy URL {}", action.proxy_url))?; - proxy_url - .query_pairs_mut() - .append_pair("link_id", &action.link_id) - .append_pair("connector_id", &action.connector_id) - .append_pair("request_url", &request_uri.to_string()); let proxy_authority = url_authority_header_value(&proxy_url)?; Ok((proxy_url.as_str().parse()?, proxy_authority)) } +fn apply_credentialed_route_proxy_headers( + headers: &mut HeaderMap, + action: &crate::mitm_hook::CredentialedRouteProxyAction, + request_uri: &Uri, +) -> Result<()> { + headers.insert( + HeaderName::from_static(CREDENTIAL_ROUTE_CONNECTOR_ID_HEADER), + HeaderValue::from_str(&action.connector_id)?, + ); + headers.insert( + HeaderName::from_static(CREDENTIAL_ROUTE_LINK_ID_HEADER), + HeaderValue::from_str(&action.link_id)?, + ); + headers.insert( + HeaderName::from_static(CREDENTIAL_ROUTE_REQUEST_URL_HEADER), + HeaderValue::from_str(&request_uri.to_string())?, + ); + Ok(()) +} + fn url_authority_header_value(url: &Url) -> Result { let host = url .host_str() diff --git a/codex-rs/network-proxy/src/mitm_tests.rs b/codex-rs/network-proxy/src/mitm_tests.rs index b069cafbb5..e0e9d0bb2d 100644 --- a/codex-rs/network-proxy/src/mitm_tests.rs +++ b/codex-rs/network-proxy/src/mitm_tests.rs @@ -285,13 +285,32 @@ fn credentialed_route_proxy_request_rewrites_upstream_target() { proxy_url: "http://localhost:8080/api/codex/credential_routes/proxy".to_string(), }; let request_uri: Uri = "https://api.example.com/v1/items?limit=5".parse().unwrap(); + let mut headers = HeaderMap::new(); + headers.insert( + HeaderName::from_static(CREDENTIAL_ROUTE_REQUEST_URL_HEADER), + HeaderValue::from_static("https://spoofed.example.com"), + ); - let (proxy_uri, proxy_authority) = - credentialed_route_proxy_request(&action, &request_uri).unwrap(); + let (proxy_uri, proxy_authority) = credentialed_route_proxy_request(&action).unwrap(); + apply_credentialed_route_proxy_headers(&mut headers, &action, &request_uri).unwrap(); assert_eq!( proxy_uri.to_string(), - "http://localhost:8080/api/codex/credential_routes/proxy?link_id=link_123&connector_id=connector_123&request_url=https%3A%2F%2Fapi.example.com%2Fv1%2Fitems%3Flimit%3D5" + "http://localhost:8080/api/codex/credential_routes/proxy" ); assert_eq!(proxy_authority, "localhost:8080"); + assert_eq!( + headers.get(CREDENTIAL_ROUTE_CONNECTOR_ID_HEADER), + Some(&HeaderValue::from_static("connector_123")) + ); + assert_eq!( + headers.get(CREDENTIAL_ROUTE_LINK_ID_HEADER), + Some(&HeaderValue::from_static("link_123")) + ); + assert_eq!( + headers.get(CREDENTIAL_ROUTE_REQUEST_URL_HEADER), + Some(&HeaderValue::from_static( + "https://api.example.com/v1/items?limit=5" + )) + ); }