typeduck-codex-async-utils 0.59.0

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;
use crate::runtime::HostBlockReason;
use crate::state::BlockedRequest;
use crate::state::BlockedRequestArgs;
use crate::state::NetworkProxyState;
use crate::upstream::UpstreamClient;
use anyhow::Context as _;
use anyhow::Result;
use anyhow::anyhow;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use rama_core::Layer;
use rama_core::Service;
use rama_core::bytes::Bytes;
use rama_core::error::BoxError;
use rama_core::extensions::ExtensionsMut;
use rama_core::extensions::ExtensionsRef;
use rama_core::futures::stream::Stream as FuturesStream;
use rama_core::rt::Executor;
use rama_core::service::service_fn;
use rama_core::stream::PeekStream;
use rama_core::stream::StackReader;
use rama_core::stream::Stream;
use rama_http::Body;
use rama_http::BodyDataStream;
use rama_http::HeaderMap;
use rama_http::HeaderValue;
use rama_http::Request;
use rama_http::Response;
use rama_http::StatusCode;
use rama_http::Uri;
use rama_http::header::HOST;
use rama_http::layer::remove_header::RemoveRequestHeaderLayer;
use rama_http::layer::remove_header::RemoveResponseHeaderLayer;
use rama_http_backend::server::HttpServer;
use rama_net::proxy::ProxyTarget;
use rama_net::stream::SocketInfo;
use rama_net::tls::server::TlsPeekStream;
use rama_tls_rustls::dep::rustls;
use rama_tls_rustls::server::TlsAcceptorData;
use rama_tls_rustls::server::TlsAcceptorLayer;
use std::pin::Pin;
use std::sync::Arc;
use std::task::Context as TaskContext;
use std::task::Poll;
use std::time::Duration;
use tokio::io::AsyncReadExt;
use tokio::time::timeout;
use tracing::info;
use tracing::warn;

/// State needed to terminate a CONNECT tunnel and enforce policy on inner HTTPS requests.
pub struct MitmState {
    ca: Arc<ManagedMitmCa>,
    allow_upstream_proxy: bool,
    upstream_tls_root_store: Arc<rustls::RootCertStore>,
    inspect: bool,
    max_body_bytes: usize,
}

pub(crate) struct MitmUpstreamConfig {
    pub(crate) allow_upstream_proxy: bool,
}

#[derive(Clone)]
struct MitmPolicyContext {
    target_host: String,
    target_port: u16,
    mode: NetworkMode,
    app_state: Arc<NetworkProxyState>,
}

#[derive(Clone)]
struct MitmRequestContext {
    policy: MitmPolicyContext,
    mitm: Arc<MitmState>,
    upstream: UpstreamClient,
}

enum MitmPolicyDecision {
    Allow {
        hook_actions: Option<MitmHookActions>,
    },
    Block(Response),
}

const MITM_INSPECT_BODIES: bool = false;
const MITM_MAX_BODY_BYTES: usize = 4096;
const TLS_PREFIX_LEN: usize = 5;
const TLS_PREFIX_FIRST_BYTE_TIMEOUT: Duration = Duration::from_millis(250);

/// Peeks enough bytes to distinguish a TLS handshake from an opaque CONNECT stream.
///
/// The first-byte timeout preserves server-first protocols. Once the client starts a possible TLS
/// prefix, all five record-header bytes are accumulated so fragmented handshakes cannot bypass
/// interception. Every byte read is replayed through `TlsPeekStream`.
pub(crate) async fn peek_tls_prefix<S>(mut stream: S) -> Result<(bool, TlsPeekStream<S>)>
where
    S: Stream + Unpin + ExtensionsMut,
{
    let mut peek_buf = [0_u8; TLS_PREFIX_LEN];
    let mut bytes_read =
        match timeout(TLS_PREFIX_FIRST_BYTE_TIMEOUT, stream.read(&mut peek_buf)).await {
            Ok(result) => result.context("read TLS prefix")?,
            Err(_) => 0,
        };
    while bytes_read > 0 && bytes_read < TLS_PREFIX_LEN {
        let possible_tls_prefix = matches!(
            &peek_buf[..bytes_read],
            [0x16] | [0x16, 0x03] | [0x16, 0x03, 0x00..=0x04, ..]
        );
        if !possible_tls_prefix {
            break;
        }
        let read = stream
            .read(&mut peek_buf[bytes_read..])
            .await
            .context("read TLS prefix")?;
        if read == 0 {
            break;
        }
        bytes_read += read;
    }

    let is_tls = bytes_read == TLS_PREFIX_LEN && matches!(peek_buf, [0x16, 0x03, 0x00..=0x04, ..]);
    let offset = TLS_PREFIX_LEN - bytes_read;
    if offset > 0 {
        peek_buf.copy_within(0..bytes_read, offset);
    }
    let mut peek = StackReader::new(peek_buf);
    peek.skip(offset);
    Ok((is_tls, PeekStream::new(peek, stream)))
}

impl std::fmt::Debug for MitmState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Avoid dumping internal state (CA material, connectors, etc.) to logs.
        f.debug_struct("MitmState")
            .field("inspect", &self.inspect)
            .field("max_body_bytes", &self.max_body_bytes)
            .finish_non_exhaustive()
    }
}

impl MitmState {
    pub(crate) fn new(config: MitmUpstreamConfig) -> Result<Self> {
        ensure_rustls_crypto_provider();

        // 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 a process-local CA and issue per-host leaf certs so we can terminate TLS and
        // apply policy.
        let ca = ManagedMitmCa::load_or_create()?;
        let upstream_tls_root_store =
            crate::certs::upstream_tls_root_store(&crate::certs::ca_env_from_process())?;

        Ok(Self {
            ca,
            allow_upstream_proxy: config.allow_upstream_proxy,
            upstream_tls_root_store,
            inspect: MITM_INSPECT_BODIES,
            max_body_bytes: MITM_MAX_BODY_BYTES,
        })
    }

    fn tls_acceptor_data_for_host(&self, host: &str) -> Result<TlsAcceptorData> {
        self.ca.tls_acceptor_data_for_host(host)
    }

    pub(crate) fn inspect_enabled(&self) -> bool {
        self.inspect
    }

    pub(crate) fn max_body_bytes(&self) -> usize {
        self.max_body_bytes
    }
}

/// Terminate a raw client stream with a generated leaf cert and proxy inner HTTPS traffic.
pub(crate) async fn mitm_stream<S>(stream: S) -> Result<()>
where
    S: Stream + Unpin + ExtensionsMut,
{
    let mitm = stream
        .extensions()
        .get::<Arc<MitmState>>()
        .cloned()
        .context("missing MITM state")?;
    let app_state = stream
        .extensions()
        .get::<Arc<NetworkProxyState>>()
        .cloned()
        .context("missing app state")?;
    let target = stream
        .extensions()
        .get::<ProxyTarget>()
        .context("missing proxy target")?
        .0
        .clone();
    let target_host = normalize_host(&target.host.to_string());
    let target_port = target.port;
    let acceptor_data = mitm.tls_acceptor_data_for_host(&target_host)?;
    let mode = stream
        .extensions()
        .get::<NetworkMode>()
        .copied()
        .unwrap_or(NetworkMode::Full);
    let upstream = if mitm.allow_upstream_proxy {
        UpstreamClient::from_env_proxy_with_tls_root_store(
            app_state.clone(),
            mitm.upstream_tls_root_store.clone(),
        )
    } else {
        UpstreamClient::direct_with_tls_root_store(
            app_state.clone(),
            mitm.upstream_tls_root_store.clone(),
        )
    };
    let request_ctx = Arc::new(MitmRequestContext {
        policy: MitmPolicyContext {
            target_host,
            target_port,
            mode,
            app_state,
        },
        mitm,
        upstream,
    });

    let executor = stream
        .extensions()
        .get::<Executor>()
        .cloned()
        .unwrap_or_default();

    let http_service = HttpServer::auto(executor).service(
        (
            RemoveResponseHeaderLayer::hop_by_hop(),
            RemoveRequestHeaderLayer::hop_by_hop(),
        )
            .into_layer(service_fn({
                let request_ctx = request_ctx.clone();
                move |req| {
                    let request_ctx = request_ctx.clone();
                    async move { handle_mitm_request(req, request_ctx).await }
                }
            })),
    );

    let https_service = TlsAcceptorLayer::new(acceptor_data)
        .with_store_client_hello(true)
        .into_layer(http_service);

    https_service
        .serve(stream)
        .await
        .map_err(|err| anyhow!("MITM serve error: {err}"))?;
    Ok(())
}

async fn handle_mitm_request(
    req: Request,
    request_ctx: Arc<MitmRequestContext>,
) -> Result<Response, std::convert::Infallible> {
    let response = match forward_request(req, &request_ctx).await {
        Ok(resp) => resp,
        Err(err) => {
            warn!("MITM request handling failed: {err}");
            text_response(StatusCode::BAD_GATEWAY, "mitm upstream error")
        }
    };
    Ok(response)
}

async fn forward_request(req: Request, request_ctx: &MitmRequestContext) -> Result<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;
    let mitm = request_ctx.mitm.clone();

    let method = req.method().as_str().to_string();
    let path = path_and_query(req.uri());
    let log_path = path_for_log(req.uri());

    let (mut parts, body) = req.into_parts();
    request_ctx
        .policy
        .app_state
        .inject_request_credentials(&target_host, &mut parts.headers);
    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
        .headers
        .insert(HOST, HeaderValue::from_str(&authority)?);

    let inspect = mitm.inspect_enabled();
    let max_body_bytes = mitm.max_body_bytes();
    let body = if inspect {
        inspect_body(
            body,
            max_body_bytes,
            RequestLogContext {
                host: authority.clone(),
                method: method.clone(),
                path: log_path.clone(),
            },
        )
    } else {
        body
    };

    let upstream_req = Request::from_parts(parts, body);
    let upstream_resp = request_ctx.upstream.serve(upstream_req).await?;
    respond_with_inspection(
        upstream_resp,
        inspect,
        max_body_bytes,
        &method,
        &log_path,
        &authority,
    )
}

#[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(MitmPolicyDecision::Block(text_response(
            StatusCode::METHOD_NOT_ALLOWED,
            "CONNECT not supported inside MITM",
        )));
    }

    let method = req.method().as_str().to_string();
    let log_path = path_for_log(req.uri());
    let client = req
        .extensions()
        .get::<SocketInfo>()
        .map(|info| info.peer_addr().to_string());

    if let Some(request_host) = extract_request_host(req) {
        let normalized = normalize_host(&request_host);
        if !normalized.is_empty() && normalized != policy.target_host {
            warn!(
                "MITM host mismatch (target={}, request_host={normalized})",
                policy.target_host
            );
            return Ok(MitmPolicyDecision::Block(text_response(
                StatusCode::BAD_REQUEST,
                "host mismatch",
            )));
        }
    }

    // CONNECT already handled allowlist/denylist + decider policy. Re-check local/private
    // resolution here to defend against DNS rebinding between CONNECT and inner HTTPS requests.
    if matches!(
        policy
            .app_state
            .host_blocked(&policy.target_host, policy.target_port)
            .await?,
        HostBlockDecision::Blocked(HostBlockReason::NotAllowedLocal)
    ) {
        let reason = HostBlockReason::NotAllowedLocal.as_str();
        let _ = policy
            .app_state
            .record_blocked(BlockedRequest::new(BlockedRequestArgs {
                host: policy.target_host.clone(),
                reason: reason.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 local/private target after CONNECT (host={}, port={}, method={method}, path={log_path})",
            policy.target_host, policy.target_port
        );
        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
            .record_blocked(BlockedRequest::new(BlockedRequestArgs {
                host: policy.target_host.clone(),
                reason: REASON_METHOD_NOT_ALLOWED.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 method policy (host={}, method={method}, path={log_path}, mode={:?}, allowed_methods=GET, HEAD, OPTIONS)",
            policy.target_host, policy.mode
        );
        return Ok(MitmPolicyDecision::Block(blocked_text_response(
            REASON_METHOD_NOT_ALLOWED,
        )));
    }

    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(
    resp: Response,
    inspect: bool,
    max_body_bytes: usize,
    method: &str,
    log_path: &str,
    authority: &str,
) -> Result<Response> {
    if !inspect {
        return Ok(resp);
    }

    let (parts, body) = resp.into_parts();
    let body = inspect_body(
        body,
        max_body_bytes,
        ResponseLogContext {
            host: authority.to_string(),
            method: method.to_string(),
            path: log_path.to_string(),
            status: parts.status,
        },
    );
    Ok(Response::from_parts(parts, body))
}

fn inspect_body<T: BodyLoggable + Send + 'static>(
    body: Body,
    max_body_bytes: usize,
    ctx: T,
) -> Body {
    Body::from_stream(InspectStream {
        inner: Box::pin(body.into_data_stream()),
        ctx: Some(Box::new(ctx)),
        len: 0,
        max_body_bytes,
    })
}

struct InspectStream<T> {
    inner: Pin<Box<BodyDataStream>>,
    ctx: Option<Box<T>>,
    len: usize,
    max_body_bytes: usize,
}

impl<T: BodyLoggable> FuturesStream for InspectStream<T> {
    type Item = Result<Bytes, BoxError>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        match this.inner.as_mut().poll_next(cx) {
            Poll::Ready(Some(Ok(bytes))) => {
                this.len = this.len.saturating_add(bytes.len());
                Poll::Ready(Some(Ok(bytes)))
            }
            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
            Poll::Ready(None) => {
                if let Some(ctx) = this.ctx.take() {
                    ctx.log(this.len, this.len > this.max_body_bytes);
                }
                Poll::Ready(None)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

struct RequestLogContext {
    host: String,
    method: String,
    path: String,
}

struct ResponseLogContext {
    host: String,
    method: String,
    path: String,
    status: StatusCode,
}

trait BodyLoggable {
    fn log(self, len: usize, truncated: bool);
}

impl BodyLoggable for RequestLogContext {
    fn log(self, len: usize, truncated: bool) {
        let host = self.host;
        let method = self.method;
        let path = self.path;
        info!(
            "MITM inspected request body (host={host}, method={method}, path={path}, body_len={len}, truncated={truncated})"
        );
    }
}

impl BodyLoggable for ResponseLogContext {
    fn log(self, len: usize, truncated: bool) {
        let host = self.host;
        let method = self.method;
        let path = self.path;
        let status = self.status;
        info!(
            "MITM inspected response body (host={host}, method={method}, path={path}, status={status}, body_len={len}, truncated={truncated})"
        );
    }
}

fn extract_request_host(req: &Request) -> Option<String> {
    req.headers()
        .get(HOST)
        .and_then(|v| v.to_str().ok())
        .map(ToString::to_string)
        .or_else(|| req.uri().authority().map(|a| a.as_str().to_string()))
}

fn authority_header_value(host: &str, port: u16) -> String {
    // Host header / URI authority formatting.
    if host.contains(':') {
        if port == 443 {
            format!("[{host}]")
        } else {
            format!("[{host}]:{port}")
        }
    } else if port == 443 {
        host.to_string()
    } else {
        format!("{host}:{port}")
    }
}

fn build_https_uri(authority: &str, path: &str) -> Result<Uri> {
    let target = format!("https://{authority}{path}");
    Ok(target.parse()?)
}

fn path_and_query(uri: &Uri) -> String {
    uri.path_and_query()
        .map(rama_http::uri::PathAndQuery::as_str)
        .unwrap_or("/")
        .to_string()
}

fn path_for_log(uri: &Uri) -> String {
    uri.path().to_string()
}

#[cfg(test)]
#[path = "mitm_tests.rs"]
mod tests;