Wire MITM hooks into runtime enforcement

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

@@ -80,6 +80,9 @@ use tracing::error;
 use tracing::info;
 use tracing::warn;
 
+#[derive(Clone, Copy, Debug)]
+struct ConnectMitmEnabled(bool);
+
 pub async fn run_http_proxy(
     state: Arc<NetworkProxyState>,
     addr: SocketAddr,
@@ -256,10 +259,18 @@ async fn http_connect_accept(
             return Err(text_response(StatusCode::INTERNAL_SERVER_ERROR, "error"));
         }
     };
+    let host_has_mitm_hooks = match app_state.host_has_mitm_hooks(&host).await {
+        Ok(has_hooks) => has_hooks,
+        Err(err) => {
+            error!("failed to inspect MITM hooks for {host}: {err}");
+            return Err(text_response(StatusCode::INTERNAL_SERVER_ERROR, "error"));
+        }
+    };
+    let connect_needs_mitm = mode == NetworkMode::Limited || host_has_mitm_hooks;
 
-    if mode == NetworkMode::Limited && mitm_state.is_none() {
-        // Limited mode is designed to be read-only. Without MITM, a CONNECT tunnel would hide the
-        // inner HTTP method/headers from the proxy, effectively bypassing method policy.
+    if connect_needs_mitm && mitm_state.is_none() {
+        // CONNECT needs MITM whenever HTTPS policy depends on inner-request inspection, either for
+        // limited-mode method enforcement or for host-specific MITM hooks.
         emit_http_block_decision_audit_event(
             &app_state,
             BlockDecisionAuditEventArgs {
@@ -286,7 +297,7 @@ async fn http_connect_accept(
                 reason: REASON_MITM_REQUIRED.to_string(),
                 client: client.clone(),
                 method: Some("CONNECT".to_string()),
-                mode: Some(NetworkMode::Limited),
+                mode: Some(mode),
                 protocol: "http-connect".to_string(),
                 decision: Some(details.decision.as_str().to_string()),
                 source: Some(details.source.as_str().to_string()),
@@ -295,14 +306,16 @@ async fn http_connect_accept(
             .await;
         let client = client.as_deref().unwrap_or_default();
         warn!(
-            "CONNECT blocked; MITM required for read-only HTTPS in limited mode (client={client}, host={host}, mode=limited, allowed_methods=GET, HEAD, OPTIONS)"
+            "CONNECT blocked; MITM required to enforce HTTPS policy (client={client}, host={host}, mode={mode:?}, hooked_host={host_has_mitm_hooks})"
         );
         return Err(blocked_text_with_details(REASON_MITM_REQUIRED, &details));
     }
 
     req.extensions_mut().insert(ProxyTarget(authority));
+    req.extensions_mut()
+        .insert(ConnectMitmEnabled(connect_needs_mitm));
     req.extensions_mut().insert(mode);
-    if let Some(mitm_state) = mitm_state {
+    if connect_needs_mitm && let Some(mitm_state) = mitm_state {
         req.extensions_mut().insert(mitm_state);
     }
 
@@ -331,7 +344,10 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> {
         return Ok(());
     };
 
-    if mode == NetworkMode::Limited
+    if upgraded
+        .extensions()
+        .get::<ConnectMitmEnabled>()
+        .is_some_and(|enabled| enabled.0)
         && upgraded
             .extensions()
             .get::<Arc<mitm::MitmState>>()
@@ -1094,6 +1110,42 @@ mod tests {
         assert_eq!(response.status(), StatusCode::OK);
     }
 
+    #[tokio::test]
+    async fn http_connect_accept_blocks_hooked_host_in_full_mode_without_mitm_state() {
+        let mut policy = NetworkProxySettings {
+            mitm: true,
+            mitm_hooks: vec![crate::mitm_hook::MitmHookConfig {
+                host: "api.github.com".to_string(),
+                matcher: crate::mitm_hook::MitmHookMatchConfig {
+                    methods: vec!["POST".to_string()],
+                    path_prefixes: vec!["/repos/openai/".to_string()],
+                    ..crate::mitm_hook::MitmHookMatchConfig::default()
+                },
+                actions: crate::mitm_hook::MitmHookActionsConfig::default(),
+            }],
+            ..Default::default()
+        };
+        policy.set_allowed_domains(vec!["api.github.com".to_string()]);
+        let state = Arc::new(network_proxy_state_for_policy(policy));
+
+        let mut req = Request::builder()
+            .method(Method::CONNECT)
+            .uri("https://api.github.com:443")
+            .header("host", "api.github.com:443")
+            .body(Body::empty())
+            .unwrap();
+        req.extensions_mut().insert(state);
+
+        let response = http_connect_accept(/*policy_decider*/ None, req)
+            .await
+            .unwrap_err();
+        assert_eq!(response.status(), StatusCode::FORBIDDEN);
+        assert_eq!(
+            response.headers().get("x-proxy-error").unwrap(),
+            "blocked-by-mitm-required"
+        );
+    }
+
     #[tokio::test]
     async fn http_proxy_listener_accepts_plain_http1_connect_requests() {
         let target_listener = TokioTcpListener::bind((Ipv4Addr::LOCALHOST, 0))

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

@@ -1,7 +1,10 @@
 use crate::certs::ManagedMitmCa;
 use crate::config::NetworkMode;
+use crate::mitm_hook::HookEvaluation;
+use crate::mitm_hook::MitmHookActions;
 use crate::policy::normalize_host;
 use crate::reasons::REASON_METHOD_NOT_ALLOWED;
+use crate::reasons::REASON_MITM_HOOK_DENIED;
 use crate::responses::blocked_text_response;
 use crate::responses::text_response;
 use crate::runtime::HostBlockDecision;
@@ -23,6 +26,7 @@ use rama_core::rt::Executor;
 use rama_core::service::service_fn;
 use rama_http::Body;
 use rama_http::BodyDataStream;
+use rama_http::HeaderMap;
 use rama_http::HeaderValue;
 use rama_http::Request;
 use rama_http::Response;
@@ -71,6 +75,13 @@ struct MitmRequestContext {
     mitm: Arc<MitmState>,
 }
 
+enum MitmPolicyDecision {
+    Allow {
+        hook_actions: Option<MitmHookActions>,
+    },
+    Block(Response),
+}
+
 const MITM_INSPECT_BODIES: bool = false;
 const MITM_MAX_BODY_BYTES: usize = 4096;
 
@@ -86,9 +97,10 @@ impl std::fmt::Debug for MitmState {
 
 impl MitmState {
     pub(crate) fn new(config: MitmUpstreamConfig) -> Result<Self> {
-        // MITM exists to make limited-mode HTTPS enforceable: once CONNECT is established, plain
-        // proxying would lose visibility into the inner HTTP request. We generate/load a local CA
-        // and issue per-host leaf certs so we can terminate TLS and apply policy.
+        // MITM exists when HTTPS policy depends on the inner request: limited-mode method clamps
+        // and host-specific hooks both need visibility after CONNECT is established. We
+        // generate/load a local CA and issue per-host leaf certs so we can terminate TLS and
+        // apply policy.
         let ca = ManagedMitmCa::load_or_create()?;
 
         let upstream = if config.allow_upstream_proxy {
@@ -200,9 +212,10 @@ async fn handle_mitm_request(
 }
 
 async fn forward_request(req: Request, request_ctx: &MitmRequestContext) -> Result<Response> {
-    if let Some(response) = mitm_blocking_response(&req, &request_ctx.policy).await? {
-        return Ok(response);
-    }
+    let hook_actions = match evaluate_mitm_policy(&req, &request_ctx.policy).await? {
+        MitmPolicyDecision::Allow { hook_actions } => hook_actions,
+        MitmPolicyDecision::Block(response) => return Ok(response),
+    };
 
     let target_host = request_ctx.policy.target_host.clone();
     let target_port = request_ctx.policy.target_port;
@@ -213,6 +226,7 @@ async fn forward_request(req: Request, request_ctx: &MitmRequestContext) -> Resu
     let log_path = path_for_log(req.uri());
 
     let (mut parts, body) = req.into_parts();
+    apply_mitm_hook_actions(&mut parts.headers, hook_actions.as_ref());
     let authority = authority_header_value(&target_host, target_port);
     parts.uri = build_https_uri(&authority, &path)?;
     parts
@@ -247,12 +261,23 @@ async fn forward_request(req: Request, request_ctx: &MitmRequestContext) -> Resu
     )
 }
 
+#[cfg_attr(not(test), allow(dead_code))]
 async fn mitm_blocking_response(
     req: &Request,
     policy: &MitmPolicyContext,
 ) -> Result<Option<Response>> {
+    match evaluate_mitm_policy(req, policy).await? {
+        MitmPolicyDecision::Allow { .. } => Ok(None),
+        MitmPolicyDecision::Block(response) => Ok(Some(response)),
+    }
+}
+
+async fn evaluate_mitm_policy(
+    req: &Request,
+    policy: &MitmPolicyContext,
+) -> Result<MitmPolicyDecision> {
     if req.method().as_str() == "CONNECT" {
-        return Ok(Some(text_response(
+        return Ok(MitmPolicyDecision::Block(text_response(
             StatusCode::METHOD_NOT_ALLOWED,
             "CONNECT not supported inside MITM",
         )));
@@ -272,7 +297,7 @@ async fn mitm_blocking_response(
                 "MITM host mismatch (target={}, request_host={normalized})",
                 policy.target_host
             );
-            return Ok(Some(text_response(
+            return Ok(MitmPolicyDecision::Block(text_response(
                 StatusCode::BAD_REQUEST,
                 "host mismatch",
             )));
@@ -307,9 +332,41 @@ async fn mitm_blocking_response(
             "MITM blocked local/private target after CONNECT (host={}, port={}, method={method}, path={log_path})",
             policy.target_host, policy.target_port
         );
-        return Ok(Some(blocked_text_response(reason)));
+        return Ok(MitmPolicyDecision::Block(blocked_text_response(reason)));
     }
 
+    let hook_actions = match policy
+        .app_state
+        .evaluate_mitm_hook_request(&policy.target_host, req)
+        .await?
+    {
+        HookEvaluation::Matched { actions } => Some(actions),
+        HookEvaluation::HookedHostNoMatch => {
+            let _ = policy
+                .app_state
+                .record_blocked(BlockedRequest::new(BlockedRequestArgs {
+                    host: policy.target_host.clone(),
+                    reason: REASON_MITM_HOOK_DENIED.to_string(),
+                    client: client.clone(),
+                    method: Some(method.clone()),
+                    mode: Some(policy.mode),
+                    protocol: "https".to_string(),
+                    decision: None,
+                    source: None,
+                    port: Some(policy.target_port),
+                }))
+                .await;
+            warn!(
+                "MITM blocked by hook policy (host={}, method={method}, mode={:?})",
+                policy.target_host, policy.mode
+            );
+            return Ok(MitmPolicyDecision::Block(blocked_text_response(
+                REASON_MITM_HOOK_DENIED,
+            )));
+        }
+        HookEvaluation::NoHooksForHost => None,
+    };
+
     if !policy.mode.allows_method(&method) {
         let _ = policy
             .app_state
@@ -329,10 +386,25 @@ async fn mitm_blocking_response(
             "MITM blocked by method policy (host={}, method={method}, path={log_path}, mode={:?}, allowed_methods=GET, HEAD, OPTIONS)",
             policy.target_host, policy.mode
         );
-        return Ok(Some(blocked_text_response(REASON_METHOD_NOT_ALLOWED)));
+        return Ok(MitmPolicyDecision::Block(blocked_text_response(
+            REASON_METHOD_NOT_ALLOWED,
+        )));
     }
 
-    Ok(None)
+    Ok(MitmPolicyDecision::Allow { hook_actions })
+}
+
+fn apply_mitm_hook_actions(headers: &mut HeaderMap, actions: Option<&MitmHookActions>) {
+    let Some(actions) = actions else {
+        return;
+    };
+
+    for header_name in &actions.strip_request_headers {
+        headers.remove(header_name);
+    }
+    for injected_header in &actions.inject_request_headers {
+        headers.insert(injected_header.name.clone(), injected_header.value.clone());
+    }
 }
 
 fn respond_with_inspection(