secrets/crates/secrets-mcp/src/logging.rs

use std::net::SocketAddr;
use std::time::Instant;

use axum::{
    body::{Body, Bytes, to_bytes},
    extract::{ConnectInfo, Request},
    http::{
        HeaderMap, Method,
        header::{CONTENT_LENGTH, CONTENT_TYPE, USER_AGENT},
    },
    middleware::Next,
    response::Response,
};

/// Axum middleware that logs structured info for every HTTP request.
///
/// All requests: method, path, status, latency_ms, client_ip, user_agent.
/// POST /mcp requests: additionally parses JSON-RPC body for jsonrpc_method,
/// tool_name, jsonrpc_id, mcp_session, batch_size.
///
/// Sensitive headers (Authorization, X-Encryption-Key) and secret values
/// are never logged.
pub async fn request_logging_middleware(req: Request, next: Next) -> Response {
    let method = req.method().clone();
    let path = req.uri().path().to_string();
    let ip = client_ip(&req);
    let ua = header_str(req.headers(), USER_AGENT);
    let content_len = header_str(req.headers(), CONTENT_LENGTH).and_then(|v| v.parse::<u64>().ok());
    let mcp_session = req
        .headers()
        .get("mcp-session-id")
        .or_else(|| req.headers().get("x-mcp-session"))
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string());

    let is_mcp_post = path.starts_with("/mcp") && method == Method::POST;
    let is_json = header_str(req.headers(), CONTENT_TYPE)
        .map(|ct| ct.contains("application/json"))
        .unwrap_or(false);

    let start = Instant::now();

    // For MCP JSON-RPC POST requests, buffer body to extract JSON-RPC metadata.
    // We cap at 512 KiB to avoid buffering large payloads.
    if is_mcp_post && is_json {
        let cap = content_len.unwrap_or(0);
        if cap <= 512 * 1024 {
            let (parts, body) = req.into_parts();
            match to_bytes(body, 512 * 1024).await {
                Ok(bytes) => {
                    let rpc = parse_jsonrpc_meta(&bytes);
                    let req = Request::from_parts(parts, Body::from(bytes));
                    let resp = next.run(req).await;
                    let status = resp.status().as_u16();
                    let elapsed = start.elapsed().as_millis();
                    log_mcp_request(
                        &method,
                        &path,
                        status,
                        elapsed,
                        ip.as_deref(),
                        ua.as_deref(),
                        content_len,
                        mcp_session.as_deref(),
                        &rpc,
                    );
                    return resp;
                }
                Err(e) => {
                    tracing::warn!(path, error = %e, "failed to buffer MCP request body for logging");
                    // Reconstruct with empty body; request was consumed — return 500.
                    // This branch is highly unlikely in practice.
                    let resp = next.run(Request::from_parts(parts, Body::empty())).await;
                    return resp;
                }
            }
        }
    }

    let resp = next.run(req).await;
    let status = resp.status().as_u16();
    let elapsed = start.elapsed().as_millis();

    // Known client probe patterns that legitimately 404 — downgrade to debug to
    // avoid noise in production logs.  These are:
    //   • GET /.well-known/* — OAuth/OIDC discovery by MCP clients (RFC 8414 / RFC 9728)
    //   • GET /mcp → 404    — old SSE-transport compatibility probe by clients
    let is_expected_probe_404 = status == 404
        && (path.starts_with("/.well-known/")
            || (method == Method::GET && path.starts_with("/mcp")));

    if is_expected_probe_404 {
        tracing::debug!(
            method = method.as_str(),
            path,
            status,
            elapsed_ms = elapsed,
            client_ip = ip.as_deref(),
            ua = ua.as_deref(),
            "probe request (not found — expected)",
        );
    } else {
        log_http_request(
            &method,
            &path,
            status,
            elapsed,
            ip.as_deref(),
            ua.as_deref(),
            content_len,
        );
    }

    resp
}

// ── Logging helpers ───────────────────────────────────────────────────────────

fn log_http_request(
    method: &Method,
    path: &str,
    status: u16,
    elapsed_ms: u128,
    client_ip: Option<&str>,
    ua: Option<&str>,
    content_length: Option<u64>,
) {
    tracing::info!(
        method = method.as_str(),
        path,
        status,
        elapsed_ms,
        client_ip,
        ua,
        content_length,
        "http request",
    );
}

#[allow(clippy::too_many_arguments)]
fn log_mcp_request(
    method: &Method,
    path: &str,
    status: u16,
    elapsed_ms: u128,
    client_ip: Option<&str>,
    ua: Option<&str>,
    content_length: Option<u64>,
    mcp_session: Option<&str>,
    rpc: &JsonRpcMeta,
) {
    tracing::info!(
        method = method.as_str(),
        path,
        status,
        elapsed_ms,
        client_ip,
        ua,
        content_length,
        mcp_session,
        jsonrpc = rpc.rpc_method.as_deref(),
        tool = rpc.tool_name.as_deref(),
        jsonrpc_id = rpc.request_id.as_deref(),
        batch_size = rpc.batch_size,
        "mcp request",
    );
}

// ── JSON-RPC body parsing ─────────────────────────────────────────────────────

#[derive(Debug, Default)]
struct JsonRpcMeta {
    request_id: Option<String>,
    rpc_method: Option<String>,
    tool_name: Option<String>,
    batch_size: Option<usize>,
}

fn parse_jsonrpc_meta(bytes: &Bytes) -> JsonRpcMeta {
    let Ok(value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
        return JsonRpcMeta::default();
    };

    if let Some(arr) = value.as_array() {
        // Batch request: summarise method(s) from first element only
        let first = arr.first().map(parse_single).unwrap_or_default();
        return JsonRpcMeta {
            batch_size: Some(arr.len()),
            ..first
        };
    }

    parse_single(&value)
}

fn parse_single(value: &serde_json::Value) -> JsonRpcMeta {
    let request_id = value.get("id").and_then(json_to_string);
    let rpc_method = value
        .get("method")
        .and_then(|v| v.as_str())
        .map(|s| s.to_string());
    let tool_name = value
        .pointer("/params/name")
        .and_then(|v| v.as_str())
        .map(|s| s.to_string());

    JsonRpcMeta {
        request_id,
        rpc_method,
        tool_name,
        batch_size: None,
    }
}

fn json_to_string(value: &serde_json::Value) -> Option<String> {
    match value {
        serde_json::Value::Null => None,
        serde_json::Value::String(s) => Some(s.clone()),
        serde_json::Value::Number(n) => Some(n.to_string()),
        serde_json::Value::Bool(b) => Some(b.to_string()),
        other => Some(other.to_string()),
    }
}

// ── Header helpers ────────────────────────────────────────────────────────────

fn header_str(headers: &HeaderMap, name: impl axum::http::header::AsHeaderName) -> Option<String> {
    headers
        .get(name)
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string())
}

fn client_ip(req: &Request) -> Option<String> {
    if let Some(first) = req
        .headers()
        .get("x-forwarded-for")
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.split(',').next())
    {
        let s = first.trim();
        if !s.is_empty() {
            return Some(s.to_string());
        }
    }
    req.extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|c| c.ip().to_string())
}