rmcp_server_kit/

transport.rs

use std::{
    future::Future,
    net::SocketAddr,
    path::{Path, PathBuf},
    pin::Pin,
    sync::Arc,
    time::Duration,
};

use arc_swap::ArcSwap;
use axum::{
    body::Body,
    extract::{ConnectInfo, Request},
    middleware::Next,
    response::IntoResponse,
};
use rmcp::{
    ServerHandler,
    transport::streamable_http_server::{
        StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
    },
};
use rustls::RootCertStore;
use tokio::{
    net::TcpListener,
    sync::{Semaphore, mpsc},
};
use tokio_util::sync::CancellationToken;

use crate::{
    auth::{
        AuthConfig, AuthIdentity, AuthState, MtlsConfig, TlsConnInfo, auth_middleware,
        build_rate_limiter, extract_mtls_identity,
    },
    error::McpxError,
    mtls_revocation::{self, CrlSet, DynamicClientCertVerifier},
    rbac::{RbacPolicy, ToolRateLimiter, build_tool_rate_limiter, rbac_middleware},
};

/// Map an internal `anyhow::Error` chain into a public [`McpxError::Startup`]
/// at the public API boundary, flattening the chain via the alternate
/// formatter so callers see the full causal path.
#[allow(
    clippy::needless_pass_by_value,
    reason = "consumed at .map_err(anyhow_to_startup) call sites; by-value matches the closure shape"
)]
fn anyhow_to_startup(e: anyhow::Error) -> McpxError {
    McpxError::Startup(format!("{e:#}"))
}

/// Map a `std::io::Error` produced during server startup into a public
/// [`McpxError::Startup`]. We deliberately do not use the [`McpxError::Io`]
/// `From` impl here because startup-phase IO errors (bind, listener) are
/// semantically distinct from request-time IO errors and should surface
/// the originating operation in the message.
#[allow(
    clippy::needless_pass_by_value,
    reason = "consumed at .map_err(|e| io_to_startup(...)) call sites; by-value matches the closure shape"
)]
fn io_to_startup(op: &str, e: std::io::Error) -> McpxError {
    McpxError::Startup(format!("{op}: {e}"))
}

/// Async readiness check callback for the `/readyz` endpoint.
///
/// Returns a JSON object with at least a `"ready"` boolean.
/// When `ready` is false, the endpoint returns HTTP 503.
pub type ReadinessCheck =
    Arc<dyn Fn() -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> + Send + Sync>;

/// Direct socket peer address of the current HTTP/TLS connection.
///
/// Inserted as a request extension into every request served by [`serve`] —
/// on both the plain and the TLS listener — and extractable in any axum
/// handler, including routes mounted via
/// [`McpServerConfig::with_extra_router`] (which bypass auth/RBAC and
/// therefore often need the peer address for their own protection, e.g.
/// per-IP rate limiting).
///
/// The same address is also mirrored into
/// [`axum::extract::ConnectInfo`]`<SocketAddr>` on the TLS listener, so
/// third-party middleware that expects the stock axum extension (e.g.
/// per-IP rate-limit key extractors) works unmodified under TLS.
///
/// # Semantics
///
/// - **Direct peer only.** This is the socket's remote address. Behind an
///   L4/L7 proxy or load balancer it is the proxy's address; the framework
///   performs **no** `X-Forwarded-For` / `Forwarded` interpretation.
/// - **Available on HTTP and TLS** transports alike ([`serve`]).
/// - **Absent under [`serve_stdio`]** — a stdio session has no network
///   peer (stdio bypasses the HTTP stack entirely).
/// - The separate Prometheus metrics listener (feature `metrics`) is a
///   different router and does not carry this extension.
///
/// # Privacy
///
/// `PeerAddr` exposes raw peer network metadata. The framework deliberately
/// never logs it on its own; whether to log or persist peer addresses is
/// application policy.
///
/// # Example
///
/// ```no_run
/// use axum::{Router, routing::get};
/// use rmcp_server_kit::transport::{McpServerConfig, PeerAddr};
///
/// async fn whoami(peer: PeerAddr) -> String {
///     peer.addr.ip().to_string()
/// }
///
/// let _config = McpServerConfig::new("127.0.0.1:8443", "my-server", "1.0.0")
///     .with_extra_router(Router::new().route("/whoami", get(whoami)));
/// ```
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub struct PeerAddr {
    /// Direct socket peer of this connection.
    pub addr: SocketAddr,
}

impl PeerAddr {
    /// Construct a new [`PeerAddr`]. Framework-internal: downstream code
    /// receives `PeerAddr` via request extensions and never constructs it.
    #[must_use]
    pub(crate) const fn new(addr: SocketAddr) -> Self {
        Self { addr }
    }
}

/// Extract [`PeerAddr`] from request extensions.
///
/// # Rejection
///
/// Responds `500 Internal Server Error` when the extension is missing.
/// A missing `PeerAddr` means the handler is not running under [`serve`]
/// (e.g. the router was mounted on a hand-rolled listener) — a wiring
/// bug, not a client error.
impl<S: Send + Sync> axum::extract::FromRequestParts<S> for PeerAddr {
    type Rejection = (axum::http::StatusCode, &'static str);

    async fn from_request_parts(
        parts: &mut axum::http::request::Parts,
        _state: &S,
    ) -> Result<Self, Self::Rejection> {
        parts.extensions.get::<Self>().copied().ok_or((
            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
            "peer address unavailable: not running under rmcp-server-kit serve()",
        ))
    }
}

/// Per-header overrides for the OWASP security headers emitted by the
/// global response middleware.
///
/// Each field follows a three-state semantic:
///
/// | Value         | Behaviour                                                |
/// |---------------|----------------------------------------------------------|
/// | `None`        | Use the built-in default (current behaviour).            |
/// | `Some("")`    | **Omit** the header entirely from responses.             |
/// | `Some(value)` | Emit `header: value`. Validated at config-load time.     |
///
/// All non-empty values are validated via
/// [`axum::http::HeaderValue::from_str`] inside
/// [`McpServerConfig::validate`]; invalid values fail fast before the
/// server starts accepting traffic.
///
/// `Strict-Transport-Security` has an additional rule: the substring
/// `preload` (case-insensitive) is rejected. Operators who want to
/// commit to the HSTS preload list must do so via a future explicit
/// builder method, not by smuggling it through this knob.
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct SecurityHeadersConfig {
    /// Override for `X-Content-Type-Options`. Default: `nosniff`.
    pub x_content_type_options: Option<String>,
    /// Override for `X-Frame-Options`. Default: `deny`.
    pub x_frame_options: Option<String>,
    /// Override for `Cache-Control`. Default: `no-store, max-age=0`.
    pub cache_control: Option<String>,
    /// Override for `Referrer-Policy`. Default: `no-referrer`.
    pub referrer_policy: Option<String>,
    /// Override for `Cross-Origin-Opener-Policy`. Default: `same-origin`.
    pub cross_origin_opener_policy: Option<String>,
    /// Override for `Cross-Origin-Resource-Policy`. Default: `same-origin`.
    pub cross_origin_resource_policy: Option<String>,
    /// Override for `Cross-Origin-Embedder-Policy`. Default: `require-corp`.
    pub cross_origin_embedder_policy: Option<String>,
    /// Override for `Permissions-Policy`. Default:
    /// `accelerometer=(), camera=(), geolocation=(), microphone=()`.
    pub permissions_policy: Option<String>,
    /// Override for `X-Permitted-Cross-Domain-Policies`. Default: `none`.
    pub x_permitted_cross_domain_policies: Option<String>,
    /// Override for `Content-Security-Policy`. Default:
    /// `default-src 'none'; frame-ancestors 'none'`.
    pub content_security_policy: Option<String>,
    /// Override for `X-DNS-Prefetch-Control`. Default: `off`.
    pub x_dns_prefetch_control: Option<String>,
    /// Override for `Strict-Transport-Security`. Default (TLS only):
    /// `max-age=63072000; includeSubDomains`. Only emitted when TLS is
    /// active; the override is ignored on plaintext deployments. The
    /// substring `preload` (any case) is rejected by the validator.
    pub strict_transport_security: Option<String>,
}

/// Configuration for the MCP server.
#[allow(
    missing_debug_implementations,
    reason = "contains callback/trait objects that don't impl Debug"
)]
#[allow(
    clippy::struct_excessive_bools,
    reason = "server configuration naturally has many boolean feature flags"
)]
#[non_exhaustive]
pub struct McpServerConfig {
    /// Socket address the MCP HTTP server binds to.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::new() / with_bind_addr(); direct field access will become pub(crate) in a future major release"
    )]
    pub bind_addr: String,
    /// Server name advertised via MCP `initialize`.
    #[deprecated(
        since = "0.13.0",
        note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
    )]
    pub name: String,
    /// Server version advertised via MCP `initialize`.
    #[deprecated(
        since = "0.13.0",
        note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
    )]
    pub version: String,
    /// Path to the TLS certificate (PEM). Required for TLS/mTLS.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
    )]
    pub tls_cert_path: Option<PathBuf>,
    /// Path to the TLS private key (PEM). Required for TLS/mTLS.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
    )]
    pub tls_key_path: Option<PathBuf>,
    /// Optional authentication config. When `Some` and `enabled`, auth
    /// is enforced on `/mcp`. `/healthz` is always open.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_auth(); direct field access will become pub(crate) in a future major release"
    )]
    pub auth: Option<AuthConfig>,
    /// Optional RBAC policy. When present and enabled, tool calls are
    /// checked against the policy after authentication.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_rbac(); direct field access will become pub(crate) in a future major release"
    )]
    pub rbac: Option<Arc<RbacPolicy>>,
    /// Allowed Origin values for DNS rebinding protection (MCP spec MUST).
    /// When empty and `public_url` is set, the origin is auto-derived from
    /// the public URL. When both are empty, only requests with no Origin
    /// header are accepted.
    /// Example entries: `"http://localhost:3000"`, `"https://myapp.example.com"`.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_allowed_origins(); direct field access will become pub(crate) in a future major release"
    )]
    pub allowed_origins: Vec<String>,
    /// Maximum tool invocations per source IP per minute.
    /// When set, enforced on every `tools/call` request.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_tool_rate_limit(); direct field access will become pub(crate) in a future major release"
    )]
    pub tool_rate_limit: Option<u32>,
    /// Optional readiness probe for `/readyz`.
    /// When `None`, `/readyz` mirrors `/healthz` (always OK).
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_readiness_check(); direct field access will become pub(crate) in a future major release"
    )]
    pub readiness_check: Option<ReadinessCheck>,
    /// Maximum request body size in bytes. Default: 1 MiB.
    /// Protects against oversized payloads causing OOM.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_max_request_body(); direct field access will become pub(crate) in a future major release"
    )]
    pub max_request_body: usize,
    /// Request processing timeout. Default: 120s.
    /// Requests exceeding this duration receive 408 Request Timeout.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_request_timeout(); direct field access will become pub(crate) in a future major release"
    )]
    pub request_timeout: Duration,
    /// Graceful shutdown timeout. Default: 30s.
    /// After the shutdown signal, in-flight requests have this long to finish.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_shutdown_timeout(); direct field access will become pub(crate) in a future major release"
    )]
    pub shutdown_timeout: Duration,
    /// Idle timeout for MCP sessions. Sessions with no activity for this
    /// duration are closed automatically. Default: 20 minutes.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_session_idle_timeout(); direct field access will become pub(crate) in a future major release"
    )]
    pub session_idle_timeout: Duration,
    /// Interval for SSE keep-alive pings. Prevents proxies and load
    /// balancers from killing idle connections. Default: 15 seconds.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_sse_keep_alive(); direct field access will become pub(crate) in a future major release"
    )]
    pub sse_keep_alive: Duration,
    /// Callback invoked once the server is built, delivering a
    /// [`ReloadHandle`] for hot-reloading auth keys and RBAC policy
    /// at runtime (e.g. on SIGHUP). Only useful when auth/RBAC is enabled.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_reload_callback(); direct field access will become pub(crate) in a future major release"
    )]
    pub on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
    /// Additional application-specific routes merged into the top-level
    /// router.  These routes **bypass** the MCP auth and RBAC middleware,
    /// so the application is responsible for its own auth on them.
    /// Handlers can extract [`PeerAddr`] (or
    /// [`axum::extract::ConnectInfo`]`<SocketAddr>` for third-party
    /// middleware compatibility) regardless of whether TLS is enabled.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_extra_router(); direct field access will become pub(crate) in a future major release"
    )]
    pub extra_router: Option<axum::Router>,
    /// Externally reachable base URL (e.g. `https://mcp.example.com`).
    /// When set, OAuth metadata endpoints advertise this URL instead of
    /// the listen address. Required when binding `0.0.0.0` behind a
    /// reverse proxy or inside a container.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_public_url(); direct field access will become pub(crate) in a future major release"
    )]
    pub public_url: Option<String>,
    /// Log inbound HTTP request headers at DEBUG level.
    /// Sensitive values remain redacted.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::enable_request_header_logging(); direct field access will become pub(crate) in a future major release"
    )]
    pub log_request_headers: bool,
    /// Enable gzip/br response compression on MCP responses.
    /// Defaults to `false` to preserve existing behaviour.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
    )]
    pub compression_enabled: bool,
    /// Minimum response body size (in bytes) before compression kicks in.
    /// Only used when `compression_enabled` is true. Default: 1024.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
    )]
    pub compression_min_size: u16,
    /// Global cap on in-flight HTTP requests across the whole server.
    /// When `Some`, requests over the cap receive 503 Service Unavailable
    /// via `tower::load_shed`. Default: `None` (unlimited).
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_max_concurrent_requests(); direct field access will become pub(crate) in a future major release"
    )]
    pub max_concurrent_requests: Option<usize>,
    /// Enable `/admin/*` diagnostic endpoints. Requires `auth` to be
    /// configured and `enabled`. Default: `false`.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
    )]
    pub admin_enabled: bool,
    /// RBAC role required to access admin endpoints. Default: `"admin"`.
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
    )]
    pub admin_role: String,
    /// Enable Prometheus metrics endpoint on a separate listener.
    /// Requires the `metrics` crate feature.
    #[cfg(feature = "metrics")]
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
    )]
    pub metrics_enabled: bool,
    /// Bind address for the Prometheus metrics listener. Default: `127.0.0.1:9090`.
    #[cfg(feature = "metrics")]
    #[deprecated(
        since = "0.13.0",
        note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
    )]
    pub metrics_bind: String,
    /// Per-header overrides for the OWASP security headers emitted by
    /// the global response middleware. See [`SecurityHeadersConfig`]
    /// for the three-state semantic and validation rules.
    #[deprecated(
        since = "1.5.0",
        note = "use McpServerConfig::with_security_headers(); direct field access will become pub(crate) in a future major release"
    )]
    pub security_headers: SecurityHeadersConfig,
    /// Per-handshake deadline on the TLS accept path. Idle or slow-loris
    /// connections are dropped once it elapses. Default: 10s.
    ///
    /// Startup-only: bound at listener construction, NOT hot-reloadable
    /// via [`ReloadHandle`]. Ignored unless TLS is configured.
    #[deprecated(
        since = "1.9.0",
        note = "use McpServerConfig::with_tls_handshake_timeout(); direct field access will become pub(crate) in a future major release"
    )]
    pub tls_handshake_timeout: Duration,
    /// Cap on concurrently in-flight TLS handshakes. At saturation the
    /// acceptor stops pulling new connections from the kernel backlog
    /// (backpressure) instead of accepting and dropping. Default: 256.
    ///
    /// Startup-only: bound at listener construction, NOT hot-reloadable
    /// via [`ReloadHandle`]. Ignored unless TLS is configured.
    #[deprecated(
        since = "1.9.0",
        note = "use McpServerConfig::with_max_concurrent_tls_handshakes(); direct field access will become pub(crate) in a future major release"
    )]
    pub max_concurrent_tls_handshakes: usize,
}

/// Marker that wraps a value proven to satisfy its validation
/// contract.
///
/// The only way to obtain `Validated<McpServerConfig>` is by calling
/// [`McpServerConfig::validate`], which is the contract enforced at
/// the type level by [`serve`] and [`serve_with_listener`]. The
/// inner field is private, so downstream code cannot bypass
/// validation by hand-constructing the wrapper.
///
/// Use [`Validated::as_inner`] for read-only borrowing. To mutate,
/// recover the raw value with [`Validated::into_inner`] and
/// re-validate.
///
/// # Example
///
/// ```no_run
/// use rmcp_server_kit::transport::{McpServerConfig, Validated, serve};
/// use rmcp::handler::server::ServerHandler;
/// use rmcp::model::{ServerCapabilities, ServerInfo};
///
/// #[derive(Clone)]
/// struct H;
/// impl ServerHandler for H {
///     fn get_info(&self) -> ServerInfo {
///         ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
///     }
/// }
///
/// # async fn example() -> rmcp_server_kit::Result<()> {
/// let config: Validated<McpServerConfig> =
///     McpServerConfig::new("127.0.0.1:8080", "my-server", "0.1.0").validate()?;
/// serve(config, || H).await
/// # }
/// ```
///
/// Forgetting `.validate()?` is a compile error:
///
/// ```compile_fail
/// use rmcp_server_kit::transport::{McpServerConfig, serve};
/// use rmcp::handler::server::ServerHandler;
/// use rmcp::model::{ServerCapabilities, ServerInfo};
///
/// #[derive(Clone)]
/// struct H;
/// impl ServerHandler for H {
///     fn get_info(&self) -> ServerInfo {
///         ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
///     }
/// }
///
/// # async fn example() -> rmcp_server_kit::Result<()> {
/// let config = McpServerConfig::new("127.0.0.1:8080", "my-server", "0.1.0");
/// // Missing `.validate()?` -> mismatched types: expected
/// // `Validated<McpServerConfig>`, found `McpServerConfig`.
/// serve(config, || H).await
/// # }
/// ```
#[allow(
    missing_debug_implementations,
    reason = "wraps T which may not implement Debug; manual impl below avoids leaking inner contents into logs"
)]
pub struct Validated<T>(T);

impl<T> std::fmt::Debug for Validated<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Validated").finish_non_exhaustive()
    }
}

impl<T> Validated<T> {
    /// Borrow the inner value.
    #[must_use]
    pub fn as_inner(&self) -> &T {
        &self.0
    }

    /// Recover the raw value, discarding the validation proof.
    ///
    /// Re-validate before re-using the value with [`serve`] or
    /// [`serve_with_listener`].
    #[must_use]
    pub fn into_inner(self) -> T {
        self.0
    }
}

#[allow(
    deprecated,
    reason = "internal builders/validators legitimately read/write the deprecated `pub` fields they were designed to manage"
)]
impl McpServerConfig {
    /// Create a new server configuration with the given bind address,
    /// server name, and version. All other fields use safe defaults.
    ///
    /// Use the chainable `with_*` / `enable_*` builder methods to
    /// customize. Call [`McpServerConfig::validate`] to obtain a
    /// [`Validated<McpServerConfig>`] proof token, which is required by
    /// [`serve`] and [`serve_with_listener`].
    #[must_use]
    pub fn new(
        bind_addr: impl Into<String>,
        name: impl Into<String>,
        version: impl Into<String>,
    ) -> Self {
        Self {
            bind_addr: bind_addr.into(),
            name: name.into(),
            version: version.into(),
            tls_cert_path: None,
            tls_key_path: None,
            auth: None,
            rbac: None,
            allowed_origins: Vec::new(),
            tool_rate_limit: None,
            readiness_check: None,
            max_request_body: 1024 * 1024,
            request_timeout: Duration::from_mins(2),
            shutdown_timeout: Duration::from_secs(30),
            session_idle_timeout: Duration::from_mins(20),
            sse_keep_alive: Duration::from_secs(15),
            on_reload_ready: None,
            extra_router: None,
            public_url: None,
            log_request_headers: false,
            compression_enabled: false,
            compression_min_size: 1024,
            max_concurrent_requests: None,
            admin_enabled: false,
            admin_role: "admin".to_owned(),
            #[cfg(feature = "metrics")]
            metrics_enabled: false,
            #[cfg(feature = "metrics")]
            metrics_bind: "127.0.0.1:9090".into(),
            security_headers: SecurityHeadersConfig::default(),
            tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
            max_concurrent_tls_handshakes: DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES,
        }
    }

    // ---------------------------------------------------------------
    // Builder methods (fluent, consume + return self).
    //
    // Each method is `#[must_use]` because dropping the returned
    // `McpServerConfig` discards the configuration change.
    // ---------------------------------------------------------------

    /// Attach an authentication configuration. Required for
    /// [`enable_admin`](Self::enable_admin) and any non-public deployment.
    #[must_use]
    pub fn with_auth(mut self, auth: AuthConfig) -> Self {
        self.auth = Some(auth);
        self
    }

    /// Override one or more of the OWASP security headers emitted on
    /// every response. See [`SecurityHeadersConfig`] for the three-state
    /// semantic (`None` = default, `Some("")` = omit, `Some(v)` =
    /// override). Values are validated by [`Self::validate`].
    #[must_use]
    pub fn with_security_headers(mut self, headers: SecurityHeadersConfig) -> Self {
        self.security_headers = headers;
        self
    }

    /// Override the bind address (e.g. `127.0.0.1:8080`). Useful when the
    /// final port is only known after pre-binding an ephemeral listener
    /// (tests, dynamic-port deployments).
    #[must_use]
    pub fn with_bind_addr(mut self, addr: impl Into<String>) -> Self {
        self.bind_addr = addr.into();
        self
    }

    /// Attach an RBAC policy. Tool calls are checked against the policy
    /// after authentication.
    #[must_use]
    pub fn with_rbac(mut self, rbac: Arc<RbacPolicy>) -> Self {
        self.rbac = Some(rbac);
        self
    }

    /// Configure TLS by providing the certificate and private key paths
    /// (PEM). Both must be readable at startup. Without this call, the
    /// server runs plain HTTP.
    #[must_use]
    pub fn with_tls(mut self, cert_path: impl Into<PathBuf>, key_path: impl Into<PathBuf>) -> Self {
        self.tls_cert_path = Some(cert_path.into());
        self.tls_key_path = Some(key_path.into());
        self
    }

    /// Set the externally reachable base URL (e.g. `https://mcp.example.com`).
    /// Required when binding `0.0.0.0` behind a reverse proxy or inside
    /// a container so OAuth metadata and auto-derived origins resolve correctly.
    #[must_use]
    pub fn with_public_url(mut self, url: impl Into<String>) -> Self {
        self.public_url = Some(url.into());
        self
    }

    /// Replace the allowed Origin allow-list (DNS-rebinding protection).
    /// When empty and [`with_public_url`](Self::with_public_url) is set,
    /// the origin is auto-derived.
    #[must_use]
    pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.allowed_origins = origins.into_iter().map(Into::into).collect();
        self
    }

    /// Merge an additional axum router at the top level. Routes added
    /// here **bypass** rmcp-server-kit auth and RBAC; the application is responsible
    /// for its own protection.
    ///
    /// To support that protection (e.g. per-IP rate limiting on
    /// unauthenticated endpoints), every request served by [`serve`]
    /// carries the client peer address regardless of whether TLS is
    /// enabled: extract the framework-owned [`PeerAddr`] in your
    /// handlers, or rely on [`axum::extract::ConnectInfo`]`<SocketAddr>`
    /// for stock third-party middleware (e.g. per-IP rate-limit key
    /// extractors). Neither extension exists under [`serve_stdio`],
    /// which has no network peer.
    #[must_use]
    pub fn with_extra_router(mut self, router: axum::Router) -> Self {
        self.extra_router = Some(router);
        self
    }

    /// Install an async readiness probe for `/readyz`. Without this call,
    /// `/readyz` mirrors `/healthz` (always 200 OK).
    #[must_use]
    pub fn with_readiness_check(mut self, check: ReadinessCheck) -> Self {
        self.readiness_check = Some(check);
        self
    }

    /// Override the maximum request body (bytes). Must be `> 0`.
    /// Default: 1 MiB.
    #[must_use]
    pub fn with_max_request_body(mut self, bytes: usize) -> Self {
        self.max_request_body = bytes;
        self
    }

    /// Override the per-request processing timeout. Default: 2 minutes.
    #[must_use]
    pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
        self.request_timeout = timeout;
        self
    }

    /// Override the graceful shutdown grace period. Default: 30 seconds.
    #[must_use]
    pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
        self.shutdown_timeout = timeout;
        self
    }

    /// Override the MCP session idle timeout. Default: 20 minutes.
    #[must_use]
    pub fn with_session_idle_timeout(mut self, timeout: Duration) -> Self {
        self.session_idle_timeout = timeout;
        self
    }

    /// Override the SSE keep-alive interval. Default: 15 seconds.
    #[must_use]
    pub fn with_sse_keep_alive(mut self, interval: Duration) -> Self {
        self.sse_keep_alive = interval;
        self
    }

    /// Cap the global number of in-flight HTTP requests via
    /// `tower::load_shed`. Excess requests receive 503 Service Unavailable.
    /// Default: unlimited.
    #[must_use]
    pub fn with_max_concurrent_requests(mut self, limit: usize) -> Self {
        self.max_concurrent_requests = Some(limit);
        self
    }

    /// Override the per-handshake deadline on the TLS accept path.
    /// Idle or slow-loris connections are dropped once it elapses.
    /// Default: 10s. Must be greater than zero.
    ///
    /// Startup-only: the value is bound at listener construction and is
    /// NOT hot-reloadable via [`ReloadHandle`]. Has no effect unless TLS
    /// is configured via [`Self::with_tls`].
    #[must_use]
    pub fn with_tls_handshake_timeout(mut self, timeout: Duration) -> Self {
        self.tls_handshake_timeout = timeout;
        self
    }

    /// Override the cap on concurrently in-flight TLS handshakes. At
    /// saturation the acceptor stops pulling new connections from the
    /// kernel backlog (backpressure) instead of accepting and dropping.
    /// Default: 256. Must be greater than zero.
    ///
    /// Startup-only: the value is bound at listener construction and is
    /// NOT hot-reloadable via [`ReloadHandle`]. Has no effect unless TLS
    /// is configured via [`Self::with_tls`].
    #[must_use]
    pub fn with_max_concurrent_tls_handshakes(mut self, limit: usize) -> Self {
        self.max_concurrent_tls_handshakes = limit;
        self
    }

    /// Cap tool invocations per source IP per minute. Enforced on every
    /// `tools/call` request.
    #[must_use]
    pub fn with_tool_rate_limit(mut self, per_minute: u32) -> Self {
        self.tool_rate_limit = Some(per_minute);
        self
    }

    /// Register a callback that receives the [`ReloadHandle`] after the
    /// server is built. Use it to wire SIGHUP-style hot reloads of API
    /// keys and RBAC policy.
    #[must_use]
    pub fn with_reload_callback<F>(mut self, callback: F) -> Self
    where
        F: FnOnce(ReloadHandle) + Send + 'static,
    {
        self.on_reload_ready = Some(Box::new(callback));
        self
    }

    /// Enable gzip/brotli response compression on MCP responses.
    /// `min_size` is the smallest body size (bytes) eligible for
    /// compression. Default min size: 1024.
    #[must_use]
    pub fn enable_compression(mut self, min_size: u16) -> Self {
        self.compression_enabled = true;
        self.compression_min_size = min_size;
        self
    }

    /// Enable `/admin/*` diagnostic endpoints. Requires
    /// [`with_auth`](Self::with_auth) to be set and enabled; otherwise
    /// [`validate`](Self::validate) returns an error. `role` is the RBAC
    /// role gate (default: `"admin"`).
    #[must_use]
    pub fn enable_admin(mut self, role: impl Into<String>) -> Self {
        self.admin_enabled = true;
        self.admin_role = role.into();
        self
    }

    /// Log inbound HTTP request headers at DEBUG level. Sensitive
    /// values remain redacted by the logging layer.
    #[must_use]
    pub fn enable_request_header_logging(mut self) -> Self {
        self.log_request_headers = true;
        self
    }

    /// Enable the Prometheus metrics listener on `bind` (e.g.
    /// `127.0.0.1:9090`). Requires the `metrics` crate feature.
    #[cfg(feature = "metrics")]
    #[must_use]
    pub fn with_metrics(mut self, bind: impl Into<String>) -> Self {
        self.metrics_enabled = true;
        self.metrics_bind = bind.into();
        self
    }

    /// Validate the configuration and consume `self`, returning a
    /// [`Validated<McpServerConfig>`] proof token required by [`serve`]
    /// and [`serve_with_listener`]. This is the only way to construct
    /// `Validated<McpServerConfig>`, so the type system guarantees
    /// validation has run before the server starts.
    ///
    /// Checks:
    ///
    /// 1. `admin_enabled` requires `auth` to be configured and enabled.
    /// 2. `tls_cert_path` and `tls_key_path` must both be set or both
    ///    be unset.
    /// 3. `bind_addr` must parse as a [`SocketAddr`].
    /// 4. `public_url`, when set, must start with `http://` or `https://`.
    /// 5. Each entry in `allowed_origins` must start with `http://` or
    ///    `https://`.
    /// 6. `max_request_body` must be greater than zero.
    /// 7. When the `oauth` feature is enabled and an [`OAuthConfig`] is
    ///    present, all OAuth URL fields (`jwks_uri`, `proxy.authorize_url`,
    ///    `proxy.token_url`, `proxy.introspection_url`,
    ///    `proxy.revocation_url`, `token_exchange.token_url`) must parse
    ///    and use the `https` scheme. Set
    ///    [`OAuthConfig::allow_http_oauth_urls`] to permit `http://`
    ///    targets (strongly discouraged in production - see the field-level
    ///    docs for the threat model).
    ///
    /// [`OAuthConfig`]: crate::oauth::OAuthConfig
    /// [`OAuthConfig::allow_http_oauth_urls`]: crate::oauth::OAuthConfig::allow_http_oauth_urls
    ///
    /// # Errors
    ///
    /// Returns [`McpxError::Config`] with a human-readable message on
    /// the first validation failure.
    pub fn validate(self) -> Result<Validated<Self>, McpxError> {
        self.check()?;
        Ok(Validated(self))
    }

    /// Run the validation checks without consuming `self`. Used by
    /// internal call sites (e.g. tests) that need to inspect a config
    /// without taking ownership.
    fn check(&self) -> Result<(), McpxError> {
        // 1. admin <-> auth dependency. Mirrors the runtime check in
        //    `build_app_router`: admin endpoints require an auth state,
        //    which is built only when `auth` is `Some` *and* `enabled`.
        if self.admin_enabled {
            let auth_enabled = self.auth.as_ref().is_some_and(|a| a.enabled);
            if !auth_enabled {
                return Err(McpxError::Config(
                    "admin_enabled=true requires auth to be configured and enabled".into(),
                ));
            }
        }

        // 2. TLS cert / key must be paired
        match (&self.tls_cert_path, &self.tls_key_path) {
            (Some(_), None) => {
                return Err(McpxError::Config(
                    "tls_cert_path is set but tls_key_path is missing".into(),
                ));
            }
            (None, Some(_)) => {
                return Err(McpxError::Config(
                    "tls_key_path is set but tls_cert_path is missing".into(),
                ));
            }
            _ => {}
        }

        // 3. bind_addr parses
        if self.bind_addr.parse::<SocketAddr>().is_err() {
            return Err(McpxError::Config(format!(
                "bind_addr {:?} is not a valid socket address (expected e.g. 127.0.0.1:8080)",
                self.bind_addr
            )));
        }

        // 4. public_url scheme
        if let Some(ref url) = self.public_url
            && !(url.starts_with("http://") || url.starts_with("https://"))
        {
            return Err(McpxError::Config(format!(
                "public_url {url:?} must start with http:// or https://"
            )));
        }

        // 5. allowed_origins scheme
        for origin in &self.allowed_origins {
            if !(origin.starts_with("http://") || origin.starts_with("https://")) {
                return Err(McpxError::Config(format!(
                    "allowed_origins entry {origin:?} must start with http:// or https://"
                )));
            }
        }

        // 6. max_request_body > 0
        if self.max_request_body == 0 {
            return Err(McpxError::Config(
                "max_request_body must be greater than zero".into(),
            ));
        }

        // 7. OAuth URL fields enforce HTTPS (unless `allow_http_oauth_urls`)
        #[cfg(feature = "oauth")]
        if let Some(auth_cfg) = &self.auth
            && let Some(oauth_cfg) = &auth_cfg.oauth
        {
            oauth_cfg.validate()?;
        }

        // 8. Security-header overrides parse as valid HTTP header values,
        //    and HSTS does not smuggle in a `preload` directive.
        validate_security_headers(&self.security_headers)?;

        // 9. max_concurrent_requests must be > 0 when set. Zero would
        //    deadlock the global concurrency limiter and reject every
        //    request. Mirrors the TOML-side check in `src/config.rs`.
        if let Some(0) = self.max_concurrent_requests {
            return Err(McpxError::Config(
                "max_concurrent_requests must be greater than zero when set".into(),
            ));
        }

        // 10. Auth rate-limit `max_tracked_keys` must be > 0. A zero cap
        //     would force `BoundedKeyedLimiter` to evict on every insert
        //     and effectively disable rate limiting.
        if let Some(auth_cfg) = &self.auth
            && let Some(rl) = &auth_cfg.rate_limit
            && rl.max_tracked_keys == 0
        {
            return Err(McpxError::Config(
                "auth.rate_limit.max_tracked_keys must be greater than zero".into(),
            ));
        }

        // 11. tls_handshake_timeout must be > 0. A zero deadline would
        //     reap every handshake before it could complete, rejecting
        //     all TLS connections. Mirrors the TOML-side check in
        //     `src/config.rs`.
        if self.tls_handshake_timeout == Duration::ZERO {
            return Err(McpxError::Config(
                "tls_handshake_timeout must be greater than zero".into(),
            ));
        }

        // 12. max_concurrent_tls_handshakes must be > 0. A zero-permit
        //     semaphore would never admit a handshake, deadlocking the
        //     TLS accept path. Mirrors the TOML-side check in
        //     `src/config.rs`.
        if self.max_concurrent_tls_handshakes == 0 {
            return Err(McpxError::Config(
                "max_concurrent_tls_handshakes must be greater than zero".into(),
            ));
        }

        Ok(())
    }
}

/// Handle for hot-reloading server configuration without restart.
///
/// Obtained via [`McpServerConfig::on_reload_ready`].
/// All swap operations are lock-free and wait-free -- in-flight requests
/// finish with the old values while new requests see the update immediately.
#[allow(
    missing_debug_implementations,
    reason = "contains Arc<AuthState> with non-Debug fields"
)]
pub struct ReloadHandle {
    auth: Option<Arc<AuthState>>,
    rbac: Option<Arc<ArcSwap<RbacPolicy>>>,
    crl_set: Option<Arc<CrlSet>>,
}

impl ReloadHandle {
    /// Atomically replace the API key list used by the auth middleware.
    pub fn reload_auth_keys(&self, keys: Vec<crate::auth::ApiKeyEntry>) {
        if let Some(ref auth) = self.auth {
            auth.reload_keys(keys);
        }
    }

    /// Atomically replace the RBAC policy used by the RBAC middleware.
    pub fn reload_rbac(&self, policy: RbacPolicy) {
        if let Some(ref rbac) = self.rbac {
            rbac.store(Arc::new(policy));
            tracing::info!("RBAC policy reloaded");
        }
    }

    /// Force an immediate refresh of all cached mTLS CRLs.
    ///
    /// # Errors
    ///
    /// Returns an error if CRL refresh is unavailable or verifier rebuild fails.
    pub async fn refresh_crls(&self) -> Result<(), McpxError> {
        let Some(ref crl_set) = self.crl_set else {
            return Err(McpxError::Config(
                "CRL refresh requested but mTLS CRL support is not configured".into(),
            ));
        };

        crl_set.force_refresh().await
    }
}

/// Generic MCP HTTP server.
///
/// Wraps an axum server with `/healthz` and `/mcp` endpoints.
/// When `tls_cert_path` and `tls_key_path` are both set, the server binds
/// with TLS (rustls). Optionally supports mTLS client certificate auth.
///
/// # Errors
///
/// Returns an error if the TCP listener cannot bind, TLS config is invalid,
/// or the server fails.
// NOTE: cognitive complexity reduced from 111/25 to 83/25 by
// extracting `run_server` (serve-loop tail) and `install_oauth_proxy_routes`.
// Remaining flow is a linear router builder: middleware layering, feature-
// gated auth/RBAC wiring, and PRM/metrics installation. Further extraction
// would require threading many `&mut Router` helpers and hurt readability
// of the layer order (which is security-relevant and must stay visible).
#[allow(
    clippy::too_many_lines,
    clippy::cognitive_complexity,
    reason = "middleware layer order is security-critical and must remain visible at one glance; extracting `&mut Router` helpers would obscure the auth/RBAC/origin/rate-limit ordering"
)]
/// Internal bundle of values produced by [`build_app_router`] and
/// consumed by [`serve`] / [`serve_with_listener`] when driving the
/// HTTP listener.
struct AppRunParams {
    /// TLS cert/key paths when TLS is configured.
    tls_paths: Option<(PathBuf, PathBuf)>,
    /// Per-handshake deadline on the TLS accept path.
    tls_handshake_timeout: Duration,
    /// Cap on concurrently in-flight TLS handshakes.
    max_concurrent_tls_handshakes: usize,
    /// mTLS configuration when mutual-TLS auth is enabled.
    mtls_config: Option<MtlsConfig>,
    /// Graceful shutdown drain window.
    shutdown_timeout: Duration,
    /// Shared auth state used by hot-reload callbacks.
    auth_state: Option<Arc<AuthState>>,
    /// Hot-reloadable RBAC state used by reload callbacks.
    rbac_swap: Arc<ArcSwap<RbacPolicy>>,
    /// Optional callback that receives the final [`ReloadHandle`].
    on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
    /// Server-internal cancellation token. Cancelled by [`run_server`]
    /// once the shutdown trigger fires (so rmcp's child token also
    /// fires, terminating in-flight MCP sessions).
    ct: CancellationToken,
    /// `"http"` or `"https"` -- used only for boot-time logging.
    scheme: &'static str,
    /// Server name -- used only for boot-time logging.
    name: String,
}

/// Build the full application axum [`axum::Router`] (MCP route +
/// middleware stack + admin + OAuth + health endpoints + security
/// headers + CORS + compression + concurrency limit + origin check)
/// and the [`AppRunParams`] needed to drive it.
///
/// This is the shared core of [`serve`] and [`serve_with_listener`].
/// It performs *no* network I/O: callers are responsible for binding
/// (or accepting a pre-bound) [`TcpListener`] and invoking
/// [`run_server`].
#[allow(
    clippy::cognitive_complexity,
    reason = "router assembly is intrinsically sequential; splitting harms readability"
)]
#[allow(
    deprecated,
    reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
)]
fn build_app_router<H, F>(
    mut config: McpServerConfig,
    handler_factory: F,
) -> anyhow::Result<(axum::Router, AppRunParams)>
where
    H: ServerHandler + 'static,
    F: Fn() -> H + Send + Sync + Clone + 'static,
{
    let ct = CancellationToken::new();

    let allowed_hosts = derive_allowed_hosts(&config.bind_addr, config.public_url.as_deref());
    tracing::info!(allowed_hosts = ?allowed_hosts, "configured Streamable HTTP allowed hosts");

    let mcp_service = StreamableHttpService::new(
        move || Ok(handler_factory()),
        {
            let mut mgr = LocalSessionManager::default();
            mgr.session_config.keep_alive = Some(config.session_idle_timeout);
            mgr.into()
        },
        StreamableHttpServerConfig::default()
            .with_allowed_hosts(allowed_hosts)
            .with_sse_keep_alive(Some(config.sse_keep_alive))
            .with_cancellation_token(ct.child_token()),
    );

    // Build the MCP route, optionally wrapped with auth and RBAC middleware.
    let mut mcp_router = axum::Router::new().nest_service("/mcp", mcp_service);

    // Build auth state eagerly when auth is configured so we can wire both
    // the auth middleware *and* the optional admin router against the same
    // state. The middleware itself is installed further down in layer order.
    let auth_state: Option<Arc<AuthState>> = match config.auth {
        Some(ref auth_config) if auth_config.enabled => {
            let rate_limiter = auth_config.rate_limit.as_ref().map(build_rate_limiter);
            let pre_auth_limiter = auth_config
                .rate_limit
                .as_ref()
                .map(crate::auth::build_pre_auth_limiter);

            #[cfg(feature = "oauth")]
            let jwks_cache = auth_config
                .oauth
                .as_ref()
                .map(|c| crate::oauth::JwksCache::new(c).map(Arc::new))
                .transpose()
                .map_err(|e| std::io::Error::other(format!("JWKS HTTP client: {e}")))?;

            Some(Arc::new(AuthState {
                api_keys: ArcSwap::new(Arc::new(auth_config.api_keys.clone())),
                rate_limiter,
                pre_auth_limiter,
                #[cfg(feature = "oauth")]
                jwks_cache,
                seen_identities: crate::auth::SeenIdentitySet::new(),
                counters: crate::auth::AuthCounters::default(),
            }))
        }
        _ => None,
    };

    // Build the RBAC policy swap early so the admin router and the later
    // RBAC middleware layer share the same hot-reloadable state.
    let rbac_swap = Arc::new(ArcSwap::new(
        config
            .rbac
            .clone()
            .unwrap_or_else(|| Arc::new(RbacPolicy::disabled())),
    ));

    // Optional /admin/* diagnostic routes. Merged BEFORE the
    // body-limit/timeout/RBAC/origin/auth layers so all of them apply.
    if config.admin_enabled {
        let Some(ref auth_state_ref) = auth_state else {
            return Err(anyhow::anyhow!(
                "admin_enabled=true requires auth to be configured and enabled"
            ));
        };
        let admin_state = crate::admin::AdminState {
            started_at: std::time::Instant::now(),
            name: config.name.clone(),
            version: config.version.clone(),
            auth: Some(Arc::clone(auth_state_ref)),
            rbac: Arc::clone(&rbac_swap),
        };
        let admin_cfg = crate::admin::AdminConfig {
            role: config.admin_role.clone(),
        };
        mcp_router = mcp_router.merge(crate::admin::admin_router(admin_state, &admin_cfg));
        tracing::info!(role = %config.admin_role, "/admin/* endpoints enabled");
    }

    // ----- Middleware order (CRITICAL: read carefully) ------------------
    //
    // axum/tower applies layers **bottom-up** at runtime: the LAST layer
    // added is the OUTERMOST (runs first on a request). To achieve a
    // request-time flow of:
    //
    //   body-limit -> timeout -> auth -> rbac -> handler
    //
    // we add layers in the REVERSE order:
    //
    //   1. RBAC               (innermost, runs last before handler)
    //   2. auth               (parses identity, sets extension for RBAC)
    //   3. timeout            (bounds total request time)
    //   4. body-limit         (outermost on /mcp; caps payload before
    //                          anything else reads/buffers it)
    //
    // Origin validation is installed on the OUTER router (after the
    // /mcp router is merged in), so it also protects /healthz, /readyz,
    // /version, and any OAuth proxy endpoints.
    //
    // Rationale:
    // - Body-limit must be outermost on /mcp so RBAC (which reads the
    //   JSON-RPC body) cannot be DoS'd by a 100MB payload.
    // - Auth must run before RBAC because RBAC consumes
    //   `req.extensions().get::<AuthIdentity>()` to enforce per-role
    //   policy.
    // - Origin runs before auth so we reject cross-origin requests
    //   without spending Argon2 cycles on unauthenticated callers.

    // [1] RBAC + tool rate-limit layer (innermost; closest to handler).
    // Always installed: even when RBAC is disabled, tool rate limiting may
    // be active (MCP spec: servers MUST rate limit tool invocations).
    {
        let tool_limiter: Option<Arc<ToolRateLimiter>> =
            config.tool_rate_limit.map(build_tool_rate_limiter);

        if rbac_swap.load().is_enabled() {
            tracing::info!("RBAC enforcement enabled on /mcp");
        }
        if let Some(limit) = config.tool_rate_limit {
            tracing::info!(limit, "tool rate limiting enabled (calls/min per IP)");
        }

        let rbac_for_mw = Arc::clone(&rbac_swap);
        mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
            let p = rbac_for_mw.load_full();
            let tl = tool_limiter.clone();
            rbac_middleware(p, tl, req, next)
        }));
    }

    // [2] Auth layer (runs before RBAC so AuthIdentity is in extensions).
    if let Some(ref auth_config) = config.auth
        && auth_config.enabled
    {
        let Some(ref state) = auth_state else {
            return Err(anyhow::anyhow!("auth state missing despite enabled config"));
        };

        let methods: Vec<&str> = [
            auth_config.mtls.is_some().then_some("mTLS"),
            (!auth_config.api_keys.is_empty()).then_some("bearer"),
            #[cfg(feature = "oauth")]
            auth_config.oauth.is_some().then_some("oauth-jwt"),
        ]
        .into_iter()
        .flatten()
        .collect();

        tracing::info!(
            methods = %methods.join(", "),
            api_keys = auth_config.api_keys.len(),
            "auth enabled on /mcp"
        );

        let state_for_mw = Arc::clone(state);
        mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
            let s = Arc::clone(&state_for_mw);
            auth_middleware(s, req, next)
        }));
    }

    // [3] Request timeout (returns 408 on expiry). Bounds total request
    // duration including auth + handler.
    mcp_router = mcp_router.layer(tower_http::timeout::TimeoutLayer::with_status_code(
        axum::http::StatusCode::REQUEST_TIMEOUT,
        config.request_timeout,
    ));

    // [4] Request body size limit (OUTERMOST on /mcp). Prevents OOM /
    // DoS from oversized payloads BEFORE any inner layer (auth, RBAC)
    // attempts to buffer or parse the body.
    mcp_router = mcp_router.layer(tower_http::limit::RequestBodyLimitLayer::new(
        config.max_request_body,
    ));

    // Compute the effective allowed-origins list for the outer
    // origin-check layer (installed on the merged router below). When
    // `allowed_origins` is empty but `public_url` is set, auto-derive
    // the origin from the public URL so MCP clients (e.g. Claude Code)
    // that send `Origin: <server-url>` are accepted without explicit
    // config.
    let mut effective_origins = config.allowed_origins.clone();
    if effective_origins.is_empty()
        && let Some(ref url) = config.public_url
    {
        // Origin = scheme + "://" + host (+ ":" + port if non-default).
        // Strip any path/query from the public URL. Offsets come from
        // `find`, so they are char-boundary-aligned; `get(..)` keeps that
        // machine-checked (a violation degrades to an empty slice).
        if let Some(scheme_end) = url.find("://") {
            let scheme_with_sep = url.get(..scheme_end + 3).unwrap_or_default();
            let after_scheme = url.get(scheme_end + 3..).unwrap_or_default();
            let host_end = after_scheme.find('/').unwrap_or(after_scheme.len());
            let host = after_scheme.get(..host_end).unwrap_or_default();
            let origin = format!("{scheme_with_sep}{host}");
            tracing::info!(
                %origin,
                "auto-derived allowed origin from public_url"
            );
            effective_origins.push(origin);
        }
    }
    let allowed_origins: Arc<[String]> = Arc::from(effective_origins);
    let cors_origins = Arc::clone(&allowed_origins);
    let log_request_headers = config.log_request_headers;

    let readyz_route = if let Some(check) = config.readiness_check.take() {
        axum::routing::get(move || readyz(Arc::clone(&check)))
    } else {
        axum::routing::get(healthz)
    };

    #[allow(unused_mut)] // mut needed when oauth feature adds PRM route
    let mut router = axum::Router::new()
        .route("/healthz", axum::routing::get(healthz))
        .route("/readyz", readyz_route)
        .route(
            "/version",
            axum::routing::get({
                // Pre-serialize the version payload once at router-build
                // time. The handler then serves a cheap `Arc::clone` of the
                // immutable bytes per request, avoiding `serde_json::Value`
                // allocation + serialization on every `/version` hit.
                let payload_bytes: Arc<[u8]> =
                    serialize_version_payload(&config.name, &config.version);
                move || {
                    let p = Arc::clone(&payload_bytes);
                    async move {
                        (
                            [(axum::http::header::CONTENT_TYPE, "application/json")],
                            p.to_vec(),
                        )
                    }
                }
            }),
        )
        .merge(mcp_router);

    // Merge application-specific routes (bypass MCP auth/RBAC middleware).
    if let Some(extra) = config.extra_router.take() {
        router = router.merge(extra);
    }

    // RFC 9728: Protected Resource Metadata endpoint.
    // When OAuth is configured, serve full metadata with authorization_servers.
    // Otherwise, serve a minimal document with just the resource URL and no
    // authorization_servers -- this tells MCP clients (e.g. Claude Code SDK)
    // that the server exists but does NOT require OAuth authentication,
    // preventing them from gating the connection behind a broken auth flow.
    let server_url = if let Some(ref url) = config.public_url {
        url.trim_end_matches('/').to_owned()
    } else {
        let prm_scheme = if config.tls_cert_path.is_some() {
            "https"
        } else {
            "http"
        };
        format!("{prm_scheme}://{}", config.bind_addr)
    };
    let resource_url = format!("{server_url}/mcp");

    #[cfg(feature = "oauth")]
    let prm_metadata = if let Some(ref auth_config) = config.auth
        && let Some(ref oauth_config) = auth_config.oauth
    {
        crate::oauth::protected_resource_metadata(&resource_url, &server_url, oauth_config)
    } else {
        serde_json::json!({ "resource": resource_url })
    };
    #[cfg(not(feature = "oauth"))]
    let prm_metadata = serde_json::json!({ "resource": resource_url });

    router = router.route(
        "/.well-known/oauth-protected-resource",
        axum::routing::get(move || {
            let m = prm_metadata.clone();
            async move { axum::Json(m) }
        }),
    );

    // OAuth 2.1 proxy endpoints: when an OAuth proxy is configured, expose
    // /authorize, /token, /register, and authorization server metadata so
    // MCP clients can perform Authorization Code + PKCE against the upstream
    // IdP (e.g. Keycloak) transparently.
    #[cfg(feature = "oauth")]
    if let Some(ref auth_config) = config.auth
        && let Some(ref oauth_config) = auth_config.oauth
        && oauth_config.proxy.is_some()
    {
        router =
            install_oauth_proxy_routes(router, &server_url, oauth_config, auth_state.as_ref())?;
    }

    // OWASP security response headers (applied to all responses).
    // HSTS is conditional on TLS being configured.
    let is_tls = config.tls_cert_path.is_some();
    let security_headers_cfg = Arc::new(config.security_headers.clone());
    router = router.layer(axum::middleware::from_fn(move |req, next| {
        let cfg = Arc::clone(&security_headers_cfg);
        security_headers_middleware(is_tls, cfg, req, next)
    }));

    // CORS preflight layer (required for browser-based MCP clients).
    // Uses the same effective origins as the origin check middleware
    // (including auto-derived origin from public_url).
    if !cors_origins.is_empty() {
        let cors = tower_http::cors::CorsLayer::new()
            .allow_origin(
                cors_origins
                    .iter()
                    .filter_map(|o| o.parse::<axum::http::HeaderValue>().ok())
                    .collect::<Vec<_>>(),
            )
            .allow_methods([
                axum::http::Method::GET,
                axum::http::Method::POST,
                axum::http::Method::OPTIONS,
            ])
            .allow_headers([
                axum::http::header::CONTENT_TYPE,
                axum::http::header::AUTHORIZATION,
            ]);
        router = router.layer(cors);
    }

    // Optional response compression (gzip + brotli). Skips small bodies
    // to avoid overhead. Applied after CORS so preflight responses remain
    // uncompressed.
    if config.compression_enabled {
        use tower_http::compression::Predicate as _;
        let predicate = tower_http::compression::DefaultPredicate::new().and(
            tower_http::compression::predicate::SizeAbove::new(config.compression_min_size),
        );
        router = router.layer(
            tower_http::compression::CompressionLayer::new()
                .gzip(true)
                .br(true)
                .compress_when(predicate),
        );
        tracing::info!(
            min_size = config.compression_min_size,
            "response compression enabled (gzip, br)"
        );
    }

    // Optional global concurrency cap. `load_shed` converts the
    // `ConcurrencyLimit` back-pressure error into 503 instead of hanging.
    if let Some(max) = config.max_concurrent_requests {
        let overload_handler = tower::ServiceBuilder::new()
            .layer(axum::error_handling::HandleErrorLayer::new(
                |_err: tower::BoxError| async {
                    (
                        axum::http::StatusCode::SERVICE_UNAVAILABLE,
                        axum::Json(serde_json::json!({
                            "error": "overloaded",
                            "error_description": "server is at capacity, retry later"
                        })),
                    )
                },
            ))
            .layer(tower::load_shed::LoadShedLayer::new())
            .layer(tower::limit::ConcurrencyLimitLayer::new(max));
        router = router.layer(overload_handler);
        tracing::info!(max, "global concurrency limit enabled");
    }

    // JSON fallback for unmatched routes. Without this, axum returns
    // an empty-body 404 that breaks MCP clients (e.g. Claude Code SDK)
    // when they probe OAuth endpoints like /authorize or /token.
    router = router.fallback(|| async {
        (
            axum::http::StatusCode::NOT_FOUND,
            axum::Json(serde_json::json!({
                "error": "not_found",
                "error_description": "The requested endpoint does not exist"
            })),
        )
    });

    // Prometheus metrics: recording middleware + separate listener.
    #[cfg(feature = "metrics")]
    if config.metrics_enabled {
        let metrics = Arc::new(
            crate::metrics::McpMetrics::new().map_err(|e| anyhow::anyhow!("metrics init: {e}"))?,
        );
        let m = Arc::clone(&metrics);
        router = router.layer(axum::middleware::from_fn(
            move |req: Request<Body>, next: Next| {
                let m = Arc::clone(&m);
                metrics_middleware(m, req, next)
            },
        ));
        let metrics_bind = config.metrics_bind.clone();
        let metrics_shutdown = ct.clone();
        tokio::spawn(async move {
            if let Err(e) =
                crate::metrics::serve_metrics(metrics_bind, metrics, metrics_shutdown).await
            {
                tracing::error!("metrics listener failed: {e}");
            }
        });
    }

    // Peer-address normalization. Mirrors the TLS branch's peer address
    // into `ConnectInfo<SocketAddr>` and exposes the framework-owned
    // `PeerAddr` extension on both listener branches, so ALL routes on
    // the merged router (`/mcp`, `/healthz`, OAuth proxy endpoints,
    // admin endpoints, extra_router, ...) and all inner middleware see a
    // uniform peer-address contract regardless of TLS. Installed just
    // inside the origin check, which stays outermost by design.
    router = router.layer(axum::middleware::from_fn(normalize_peer_addr_middleware));

    // Origin validation layer (MCP spec: servers MUST validate the
    // Origin header to prevent DNS rebinding attacks). Installed as the
    // OUTERMOST layer on the OUTER router so it protects ALL routes
    // (`/mcp`, `/healthz`, `/readyz`, `/version`, OAuth proxy endpoints,
    // admin endpoints, extra_router, etc.) and runs BEFORE auth so we
    // reject cross-origin attackers without spending Argon2 cycles.
    //
    // Origin-less requests (e.g. server-to-server probes, curl, native
    // MCP clients) are permitted; only requests with an Origin header
    // that does not match `effective_origins` are rejected.
    router = router.layer(axum::middleware::from_fn(move |req, next| {
        let origins = Arc::clone(&allowed_origins);
        origin_check_middleware(origins, log_request_headers, req, next)
    }));

    let scheme = if config.tls_cert_path.is_some() {
        "https"
    } else {
        "http"
    };

    let tls_paths = match (&config.tls_cert_path, &config.tls_key_path) {
        (Some(cert), Some(key)) => Some((cert.clone(), key.clone())),
        _ => None,
    };
    let tls_handshake_timeout = config.tls_handshake_timeout;
    let max_concurrent_tls_handshakes = config.max_concurrent_tls_handshakes;
    let mtls_config = config.auth.as_ref().and_then(|a| a.mtls.as_ref()).cloned();

    Ok((
        router,
        AppRunParams {
            tls_paths,
            tls_handshake_timeout,
            max_concurrent_tls_handshakes,
            mtls_config,
            shutdown_timeout: config.shutdown_timeout,
            auth_state,
            rbac_swap,
            on_reload_ready: config.on_reload_ready.take(),
            ct,
            scheme,
            name: config.name.clone(),
        },
    ))
}

/// Run the MCP HTTP server, binding to `config.bind_addr` and serving
/// until an OS shutdown signal (Ctrl-C / SIGTERM) is received.
///
/// This is the standard entry point for production deployments. For
/// deterministic shutdown control (e.g. integration tests), see
/// [`serve_with_listener`].
///
/// The configuration must be validated first via
/// [`McpServerConfig::validate`], which returns a [`Validated`] proof
/// token. This typestate guarantees, at compile time, that the server
/// never starts with an invalid configuration.
///
/// # Errors
///
/// Returns [`McpxError::Startup`] if binding to `config.bind_addr`
/// fails, or if the underlying axum server returns an error.
pub async fn serve<H, F>(
    config: Validated<McpServerConfig>,
    handler_factory: F,
) -> Result<(), McpxError>
where
    H: ServerHandler + 'static,
    F: Fn() -> H + Send + Sync + Clone + 'static,
{
    let config = config.into_inner();
    #[allow(
        deprecated,
        reason = "internal serve() reads `bind_addr` to construct the listener; field becomes pub(crate) in 1.0"
    )]
    let bind_addr = config.bind_addr.clone();
    let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;

    let listener = TcpListener::bind(&bind_addr)
        .await
        .map_err(|e| io_to_startup(&format!("bind {bind_addr}"), e))?;
    log_listening(&params.name, params.scheme, &bind_addr);

    run_server(
        router,
        listener,
        params.tls_paths,
        params.tls_handshake_timeout,
        params.max_concurrent_tls_handshakes,
        params.mtls_config,
        params.shutdown_timeout,
        params.auth_state,
        params.rbac_swap,
        params.on_reload_ready,
        params.ct,
    )
    .await
    .map_err(anyhow_to_startup)
}

/// Run the MCP HTTP server on a pre-bound [`TcpListener`], with optional
/// readiness signalling and external shutdown control.
///
/// This variant is intended for **deterministic integration tests** and
/// for embedders that need to bind the listening socket themselves
/// (e.g. systemd socket activation). Compared to [`serve`]:
///
/// * The caller passes a `TcpListener` that is already bound. This
///   eliminates the bind race in tests that previously required
///   poll-the-`/healthz`-loop start-up detection.
/// * `ready_tx`, when `Some`, receives the socket's
///   [`SocketAddr`] *after* the router is built and immediately before
///   the server starts accepting connections. Tests can `await` the
///   matching `oneshot::Receiver` to know exactly when it is safe to
///   issue requests.
/// * `shutdown`, when `Some`, gives the caller a
///   [`CancellationToken`] that triggers the same graceful-shutdown
///   path as a real OS signal. This avoids cross-platform issues with
///   sending real `SIGTERM` from tests on Windows.
///
/// All three optional parameters degrade gracefully: if `ready_tx` is
/// `None`, no signal is sent; if `shutdown` is `None`, the server only
/// stops on an OS signal (just like [`serve`]).
///
/// # Errors
///
/// Returns [`McpxError::Startup`] if router construction fails, if reading
/// the listener's `local_addr()` fails, or if the underlying axum
/// server returns an error.
pub async fn serve_with_listener<H, F>(
    listener: TcpListener,
    config: Validated<McpServerConfig>,
    handler_factory: F,
    ready_tx: Option<tokio::sync::oneshot::Sender<SocketAddr>>,
    shutdown: Option<CancellationToken>,
) -> Result<(), McpxError>
where
    H: ServerHandler + 'static,
    F: Fn() -> H + Send + Sync + Clone + 'static,
{
    let config = config.into_inner();
    let local_addr = listener
        .local_addr()
        .map_err(|e| io_to_startup("listener.local_addr", e))?;
    let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;

    log_listening(&params.name, params.scheme, &local_addr.to_string());

    // Forward external shutdown into the server-internal cancellation
    // token so `run_server`'s shutdown trigger picks it up alongside
    // any real OS signal.
    if let Some(external) = shutdown {
        let internal = params.ct.clone();
        tokio::spawn(async move {
            external.cancelled().await;
            internal.cancel();
        });
    }

    // Signal readiness *after* the router is fully built and external
    // shutdown is wired, but *before* run_server takes ownership of
    // the listener. The receiver can immediately issue requests.
    if let Some(tx) = ready_tx {
        // Receiver may have been dropped (test gave up). That's fine.
        let _ = tx.send(local_addr);
    }

    run_server(
        router,
        listener,
        params.tls_paths,
        params.tls_handshake_timeout,
        params.max_concurrent_tls_handshakes,
        params.mtls_config,
        params.shutdown_timeout,
        params.auth_state,
        params.rbac_swap,
        params.on_reload_ready,
        params.ct,
    )
    .await
    .map_err(anyhow_to_startup)
}

/// Emit the standard "listening on …" log lines used by both
/// [`serve`] and [`serve_with_listener`].
#[allow(
    clippy::cognitive_complexity,
    reason = "tracing::info! macro expansions inflate the score; logic is trivial"
)]
fn log_listening(name: &str, scheme: &str, addr: &str) {
    tracing::info!("{name} listening on {addr}");
    tracing::info!("  MCP endpoint: {scheme}://{addr}/mcp");
    tracing::info!("  Health check: {scheme}://{addr}/healthz");
    tracing::info!("  Readiness:   {scheme}://{addr}/readyz");
}

/// Drive the chosen axum server variant (TLS or plain) with a graceful
/// shutdown window. Consumes the router and listener.
///
/// # Shutdown semantics
///
/// A single shutdown trigger (the FIRST of: OS signal via
/// `shutdown_signal()`, or external cancellation of `ct`) starts BOTH:
///
/// 1. axum's `.with_graceful_shutdown(...)` future, which stops
///    accepting new connections and waits for in-flight requests to
///    drain;
/// 2. a `tokio::time::sleep(shutdown_timeout)` race that forces exit if
///    drainage exceeds `shutdown_timeout`.
///
/// Previously this function awaited `shutdown_signal()` independently
/// in BOTH branches of a `tokio::select!`. Because `shutdown_signal`
/// resolves once per future and consumes one signal, the force-exit
/// timer was tied to a SECOND signal (a second SIGTERM the operator
/// would never send). Under a single SIGTERM the graceful drain could
/// hang indefinitely. The current implementation derives both branches
/// from a single shared trigger so the timeout race is anchored to the
/// FIRST (and only) signal.
#[allow(
    clippy::too_many_arguments,
    clippy::cognitive_complexity,
    reason = "server start-up threads TLS, reload state, and graceful shutdown through one flow"
)]
async fn run_server(
    router: axum::Router,
    listener: TcpListener,
    tls_paths: Option<(PathBuf, PathBuf)>,
    tls_handshake_timeout: Duration,
    max_concurrent_tls_handshakes: usize,
    mtls_config: Option<MtlsConfig>,
    shutdown_timeout: Duration,
    auth_state: Option<Arc<AuthState>>,
    rbac_swap: Arc<ArcSwap<RbacPolicy>>,
    mut on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
    ct: CancellationToken,
) -> anyhow::Result<()> {
    // `shutdown_trigger` fires when the FIRST source resolves: either
    // an OS signal (Ctrl-C / SIGTERM) or external cancellation of `ct`
    // (which the test harness uses for deterministic shutdown).
    let shutdown_trigger = CancellationToken::new();
    {
        let trigger = shutdown_trigger.clone();
        let parent = ct.clone();
        tokio::spawn(async move {
            tokio::select! {
                () = shutdown_signal() => {}
                () = parent.cancelled() => {}
            }
            trigger.cancel();
        });
    }

    let graceful = {
        let trigger = shutdown_trigger.clone();
        let ct = ct.clone();
        async move {
            trigger.cancelled().await;
            tracing::info!("shutting down (grace period: {shutdown_timeout:?})");
            ct.cancel();
        }
    };

    let force_exit_timer = {
        let trigger = shutdown_trigger.clone();
        async move {
            trigger.cancelled().await;
            tokio::time::sleep(shutdown_timeout).await;
        }
    };

    if let Some((cert_path, key_path)) = tls_paths {
        let crl_set = if let Some(mtls) = mtls_config.as_ref()
            && mtls.crl_enabled
        {
            let (ca_certs, roots) = load_client_auth_roots(&mtls.ca_cert_path)?;
            let (crl_set, discover_rx) =
                mtls_revocation::bootstrap_fetch(roots, &ca_certs, mtls.clone())
                    .await
                    .map_err(|error| anyhow::anyhow!(error.to_string()))?;
            tokio::spawn(mtls_revocation::run_crl_refresher(
                Arc::clone(&crl_set),
                discover_rx,
                ct.clone(),
            ));
            Some(crl_set)
        } else {
            None
        };

        if let Some(cb) = on_reload_ready.take() {
            cb(ReloadHandle {
                auth: auth_state.clone(),
                rbac: Some(Arc::clone(&rbac_swap)),
                crl_set: crl_set.clone(),
            });
        }

        let tls_listener = TlsListener::new(
            listener,
            &cert_path,
            &key_path,
            mtls_config.as_ref(),
            crl_set,
            tls_handshake_timeout,
            max_concurrent_tls_handshakes,
        )?;
        let make_svc = router.into_make_service_with_connect_info::<TlsConnInfo>();
        tokio::select! {
            result = axum::serve(tls_listener, make_svc)
                .with_graceful_shutdown(graceful) => { result?; }
            () = force_exit_timer => {
                tracing::warn!("shutdown timeout exceeded, forcing exit");
            }
        }
    } else {
        if let Some(cb) = on_reload_ready.take() {
            cb(ReloadHandle {
                auth: auth_state,
                rbac: Some(rbac_swap),
                crl_set: None,
            });
        }

        let make_svc = router.into_make_service_with_connect_info::<SocketAddr>();
        tokio::select! {
            result = axum::serve(listener, make_svc)
                .with_graceful_shutdown(graceful) => { result?; }
            () = force_exit_timer => {
                tracing::warn!("shutdown timeout exceeded, forcing exit");
            }
        }
    }

    Ok(())
}

/// Install the OAuth 2.1 proxy endpoints (`/authorize`, `/token`,
/// `/register`, and authorization server metadata) on `router`. The
/// caller must ensure `oauth_config.proxy` is `Some`.
///
/// # Errors
///
/// Returns [`McpxError::Startup`] if the shared
/// [`crate::oauth::OauthHttpClient`] cannot be initialized.
#[cfg(feature = "oauth")]
fn install_oauth_proxy_routes(
    router: axum::Router,
    server_url: &str,
    oauth_config: &crate::oauth::OAuthConfig,
    auth_state: Option<&Arc<AuthState>>,
) -> Result<axum::Router, McpxError> {
    let Some(ref proxy) = oauth_config.proxy else {
        return Ok(router);
    };

    // Single shared HTTP client for all proxy endpoints. Cloning is
    // cheap (refcounted) and shares the underlying connection pool.
    let http = crate::oauth::OauthHttpClient::with_config(oauth_config)?;

    let asm = crate::oauth::authorization_server_metadata(server_url, oauth_config);
    let router = router.route(
        "/.well-known/oauth-authorization-server",
        axum::routing::get(move || {
            let m = asm.clone();
            async move { axum::Json(m) }
        }),
    );

    let proxy_authorize = proxy.clone();
    let router = router.route(
        "/authorize",
        axum::routing::get(
            move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
                let p = proxy_authorize.clone();
                async move { crate::oauth::handle_authorize(&p, &query.unwrap_or_default()) }
            },
        ),
    );

    let proxy_token = proxy.clone();
    let token_http = http.clone();
    let router = router.route(
        "/token",
        axum::routing::post(move |body: String| {
            let p = proxy_token.clone();
            let h = token_http.clone();
            async move { crate::oauth::handle_token(&h, &p, &body).await }
        })
        .layer(axum::middleware::from_fn(
            oauth_token_cache_headers_middleware,
        )),
    );

    let proxy_register = proxy.clone();
    let router = router.route(
        "/register",
        axum::routing::post(move |axum::Json(body): axum::Json<serde_json::Value>| {
            let p = proxy_register;
            async move { axum::Json(crate::oauth::handle_register(&p, &body)) }
        })
        .layer(axum::middleware::from_fn(
            oauth_token_cache_headers_middleware,
        )),
    );

    let admin_routes_enabled = proxy.expose_admin_endpoints
        && (proxy.introspection_url.is_some() || proxy.revocation_url.is_some());
    if proxy.expose_admin_endpoints
        && !proxy.require_auth_on_admin_endpoints
        && proxy.allow_unauthenticated_admin_endpoints
    {
        // M3 escape-hatch in effect: validate() let this through because
        // the operator explicitly opted in. Surface it loudly at startup
        // so the choice is auditable in logs.
        tracing::warn!(
            "OAuth introspect/revoke endpoints are unauthenticated by explicit \
             allow_unauthenticated_admin_endpoints opt-out; ensure an \
             authenticated reverse proxy fronts these routes"
        );
    }

    let admin_router = if admin_routes_enabled {
        build_oauth_admin_router(proxy, http, auth_state)?
    } else {
        axum::Router::new()
    };

    let router = router.merge(admin_router);

    tracing::info!(
        introspect = proxy.expose_admin_endpoints && proxy.introspection_url.is_some(),
        revoke = proxy.expose_admin_endpoints && proxy.revocation_url.is_some(),
        "OAuth 2.1 proxy endpoints enabled (/authorize, /token, /register)"
    );
    Ok(router)
}

/// Build the optional `/introspect` + `/revoke` admin sub-router.
///
/// Layered with [`oauth_token_cache_headers_middleware`] so RFC 6749 §5.1
/// / RFC 6750 §5.4 cache headers are emitted, and conditionally with the
/// auth middleware when `proxy.require_auth_on_admin_endpoints` is set.
#[cfg(feature = "oauth")]
fn build_oauth_admin_router(
    proxy: &crate::oauth::OAuthProxyConfig,
    http: crate::oauth::OauthHttpClient,
    auth_state: Option<&Arc<AuthState>>,
) -> Result<axum::Router, McpxError> {
    let mut admin_router = axum::Router::new();
    if proxy.introspection_url.is_some() {
        let proxy_introspect = proxy.clone();
        let introspect_http = http.clone();
        admin_router = admin_router.route(
            "/introspect",
            axum::routing::post(move |body: String| {
                let p = proxy_introspect.clone();
                let h = introspect_http.clone();
                async move { crate::oauth::handle_introspect(&h, &p, &body).await }
            }),
        );
    }
    if proxy.revocation_url.is_some() {
        let proxy_revoke = proxy.clone();
        let revoke_http = http;
        admin_router = admin_router.route(
            "/revoke",
            axum::routing::post(move |body: String| {
                let p = proxy_revoke.clone();
                let h = revoke_http.clone();
                async move { crate::oauth::handle_revoke(&h, &p, &body).await }
            }),
        );
    }

    let admin_router = admin_router.layer(axum::middleware::from_fn(
        oauth_token_cache_headers_middleware,
    ));

    if proxy.require_auth_on_admin_endpoints {
        let Some(state) = auth_state else {
            return Err(McpxError::Startup(
                "oauth proxy admin endpoints require auth state".into(),
            ));
        };
        let state_for_mw = Arc::clone(state);
        Ok(
            admin_router.layer(axum::middleware::from_fn(move |req, next| {
                let s = Arc::clone(&state_for_mw);
                auth_middleware(s, req, next)
            })),
        )
    } else {
        Ok(admin_router)
    }
}

/// Build the host allow-list for rmcp's DNS rebinding protection.
///
/// Includes loopback hosts by default, then augments with host/authority
/// derived from `public_url` and the server bind address.
fn derive_allowed_hosts(bind_addr: &str, public_url: Option<&str>) -> Vec<String> {
    let mut hosts = vec![
        "localhost".to_owned(),
        "127.0.0.1".to_owned(),
        "::1".to_owned(),
    ];

    if let Some(url) = public_url
        && let Ok(uri) = url.parse::<axum::http::Uri>()
        && let Some(authority) = uri.authority()
    {
        let host = authority.host().to_owned();
        if !hosts.iter().any(|h| h == &host) {
            hosts.push(host);
        }

        let authority = authority.as_str().to_owned();
        if !hosts.iter().any(|h| h == &authority) {
            hosts.push(authority);
        }
    }

    if let Ok(uri) = format!("http://{bind_addr}").parse::<axum::http::Uri>()
        && let Some(authority) = uri.authority()
    {
        let host = authority.host().to_owned();
        if !hosts.iter().any(|h| h == &host) {
            hosts.push(host);
        }

        let authority = authority.as_str().to_owned();
        if !hosts.iter().any(|h| h == &authority) {
            hosts.push(authority);
        }
    }

    hosts
}

// - TLS support -

/// Implement axum's `Connected` trait for `TlsConnInfo` so that
/// `ConnectInfo<TlsConnInfo>` is available in middleware when serving
/// over our custom `TlsListener`.
///
/// The identity is read directly from the wrapping
/// [`AuthenticatedTlsStream`], which guarantees one-to-one correspondence
/// between the TLS connection and its mTLS identity. This eliminates the
/// previous shared-map approach which was vulnerable to ephemeral-port
/// reuse races (an unauthenticated reconnection from the same `(IP, port)`
/// pair could alias a stale entry).
impl axum::extract::connect_info::Connected<axum::serve::IncomingStream<'_, TlsListener>>
    for TlsConnInfo
{
    fn connect_info(target: axum::serve::IncomingStream<'_, TlsListener>) -> Self {
        let addr = *target.remote_addr();
        let identity = target.io().identity().cloned();
        TlsConnInfo::new(addr, identity)
    }
}

/// Default per-handshake deadline on the TLS accept path. Prevents idle
/// or slow-loris connections from pinning handshake worker tasks (and
/// their semaphore permits) indefinitely.
///
/// Configurable since 1.9.0 via
/// [`McpServerConfig::with_tls_handshake_timeout`].
const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);

/// Default upper bound on concurrently in-flight TLS handshakes. When
/// saturated, the acceptor task stops pulling new connections from the
/// kernel backlog (backpressure) instead of accepting and dropping them
/// in user space.
///
/// Configurable since 1.9.0 via
/// [`McpServerConfig::with_max_concurrent_tls_handshakes`].
const DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256;

/// Capacity of the completed-handshake queue between the acceptor task and
/// `axum::serve`'s `accept()` loop. Handshake workers block on `send` when
/// the queue is full, so a slow accept loop back-pressures handshakes
/// rather than buffering completed connections unboundedly.
const TLS_ACCEPT_CHANNEL_CAPACITY: usize = 32;

/// A TLS-wrapping listener that implements axum's `Listener` trait.
///
/// TCP accepts and TLS handshakes run on a dedicated background task: each
/// accepted connection's handshake is spawned onto its own worker task,
/// bounded by a configurable concurrent-handshake cap (default
/// [`DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES`]) and a per-handshake timeout
/// (default [`DEFAULT_TLS_HANDSHAKE_TIMEOUT`]). A slow or idle client
/// therefore cannot stall other connections behind a serialized inline
/// handshake.
///
/// When mTLS is configured, client certificates are verified against the
/// configured CA and the client identity is extracted at handshake time.
/// The extracted identity is bound to the connection itself via the
/// returned [`AuthenticatedTlsStream`], so it is impossible for an
/// unrelated connection to observe it.
struct TlsListener {
    /// Bound address, captured eagerly before the `TcpListener` moves into
    /// the acceptor task.
    local_addr: SocketAddr,
    /// Completed handshakes produced by the acceptor task's workers.
    rx: mpsc::Receiver<(AuthenticatedTlsStream, SocketAddr)>,
    /// Background task driving TCP accepts and concurrent TLS handshakes.
    /// Aborted on drop so the listener releases its port deterministically.
    acceptor_task: tokio::task::JoinHandle<()>,
}

impl TlsListener {
    fn new(
        inner: TcpListener,
        cert_path: &Path,
        key_path: &Path,
        mtls_config: Option<&MtlsConfig>,
        crl_set: Option<Arc<CrlSet>>,
        handshake_timeout: Duration,
        max_concurrent_handshakes: usize,
    ) -> anyhow::Result<Self> {
        // Install the ring crypto provider (ok to call multiple times).
        rustls::crypto::ring::default_provider()
            .install_default()
            .ok();

        let certs = load_certs(cert_path)?;
        let key = load_key(key_path)?;

        let mtls_default_role;

        let tls_config = if let Some(mtls) = mtls_config {
            mtls_default_role = mtls.default_role.clone();
            let verifier: Arc<dyn rustls::server::danger::ClientCertVerifier> = if mtls.crl_enabled
            {
                let Some(crl_set) = crl_set else {
                    return Err(anyhow::anyhow!(
                        "mTLS CRL verifier requested but CRL state was not initialized"
                    ));
                };
                Arc::new(DynamicClientCertVerifier::new(crl_set))
            } else {
                let (_, root_store) = load_client_auth_roots(&mtls.ca_cert_path)?;
                if mtls.required {
                    rustls::server::WebPkiClientVerifier::builder(root_store)
                        .build()
                        .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
                } else {
                    rustls::server::WebPkiClientVerifier::builder(root_store)
                        .allow_unauthenticated()
                        .build()
                        .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
                }
            };

            tracing::info!(
                ca = %mtls.ca_cert_path.display(),
                required = mtls.required,
                crl_enabled = mtls.crl_enabled,
                "mTLS client auth configured"
            );

            rustls::ServerConfig::builder_with_protocol_versions(&[
                &rustls::version::TLS12,
                &rustls::version::TLS13,
            ])
            .with_client_cert_verifier(verifier)
            .with_single_cert(certs, key)?
        } else {
            mtls_default_role = "viewer".to_owned();
            rustls::ServerConfig::builder_with_protocol_versions(&[
                &rustls::version::TLS12,
                &rustls::version::TLS13,
            ])
            .with_no_client_auth()
            .with_single_cert(certs, key)?
        };

        let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
        tracing::info!(
            "TLS enabled (cert: {}, key: {})",
            cert_path.display(),
            key_path.display()
        );
        let local_addr = inner.local_addr()?;
        let (tx, rx) = mpsc::channel(TLS_ACCEPT_CHANNEL_CAPACITY);
        let acceptor_task = tokio::spawn(run_tls_acceptor(
            inner,
            acceptor,
            mtls_default_role,
            tx,
            handshake_timeout,
            max_concurrent_handshakes,
        ));
        Ok(Self {
            local_addr,
            rx,
            acceptor_task,
        })
    }

    /// Extract the mTLS client cert identity from a completed TLS handshake.
    /// Returns `None` if no client certificate was presented or if the
    /// certificate could not be parsed into an [`AuthIdentity`].
    fn extract_handshake_identity(
        tls_stream: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
        default_role: &str,
        addr: SocketAddr,
    ) -> Option<AuthIdentity> {
        let (_, server_conn) = tls_stream.get_ref();
        let cert_der = server_conn.peer_certificates()?.first()?;
        let id = extract_mtls_identity(cert_der.as_ref(), default_role)?;
        tracing::debug!(name = %id.name, peer = %addr, "mTLS client cert accepted");
        Some(id)
    }
}

/// Drive TCP accepts and concurrent TLS handshakes for [`TlsListener`].
///
/// Each accepted connection's handshake runs on its own worker task under
/// a permit from a `max_concurrent_handshakes`-sized semaphore and a
/// `handshake_timeout` deadline. Completed handshakes are pushed to `tx`;
/// failures and timeouts are logged at DEBUG and the connection dropped.
/// The loop exits when the owning [`TlsListener`] is dropped.
async fn run_tls_acceptor(
    listener: TcpListener,
    acceptor: tokio_rustls::TlsAcceptor,
    default_role: String,
    tx: mpsc::Sender<(AuthenticatedTlsStream, SocketAddr)>,
    handshake_timeout: Duration,
    max_concurrent_handshakes: usize,
) {
    let inflight = Arc::new(Semaphore::new(max_concurrent_handshakes));
    loop {
        // Acquire the permit BEFORE accepting: at saturation, pending
        // connections wait in the kernel backlog instead of being accepted
        // and then buffered or dropped in user space.
        let Ok(permit) = Arc::clone(&inflight).acquire_owned().await else {
            // The semaphore is never closed; defensive exit.
            return;
        };
        let (stream, addr) = match listener.accept().await {
            Ok(pair) => pair,
            Err(e) => {
                tracing::debug!("TCP accept error: {e}");
                continue;
            }
        };
        if tx.is_closed() {
            // The listener was dropped (shutdown): stop accepting.
            return;
        }
        let acceptor = acceptor.clone();
        let default_role = default_role.clone();
        let tx = tx.clone();
        tokio::spawn(async move {
            let _permit = permit;
            match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
                Ok(Ok(tls_stream)) => {
                    let identity =
                        TlsListener::extract_handshake_identity(&tls_stream, &default_role, addr);
                    let wrapped = AuthenticatedTlsStream {
                        inner: tls_stream,
                        identity,
                    };
                    // The receiver only disappears during shutdown; discard
                    // the completed connection quietly rather than logging.
                    let _ = tx.send((wrapped, addr)).await;
                }
                Ok(Err(e)) => {
                    tracing::debug!("TLS handshake failed from {addr}: {e}");
                }
                Err(_elapsed) => {
                    tracing::debug!(
                        "TLS handshake timed out from {addr} after {handshake_timeout:?}"
                    );
                }
            }
        });
    }
}

/// A TLS stream paired with the mTLS identity extracted at handshake time.
///
/// Wraps [`tokio_rustls::server::TlsStream`] so the verified client
/// identity travels with the connection itself. This replaces the previous
/// shared `MtlsIdentities` map, eliminating the
/// `(SocketAddr) -> AuthIdentity` aliasing risk caused by ephemeral-port
/// reuse and removing the need for an LRU eviction policy.
///
/// The wrapper is `Unpin` (its inner stream is `Unpin` because
/// [`tokio::net::TcpStream`] is `Unpin`), so `AsyncRead`/`AsyncWrite`
/// delegation uses safe pin projection via `Pin::new(&mut self.inner)`.
pub(crate) struct AuthenticatedTlsStream {
    inner: tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
    identity: Option<AuthIdentity>,
}

impl AuthenticatedTlsStream {
    /// Returns the verified mTLS client identity, if any.
    #[must_use]
    pub(crate) const fn identity(&self) -> Option<&AuthIdentity> {
        self.identity.as_ref()
    }
}

impl std::fmt::Debug for AuthenticatedTlsStream {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AuthenticatedTlsStream")
            .field("identity", &self.identity.as_ref().map(|id| &id.name))
            .finish_non_exhaustive()
    }
}

impl tokio::io::AsyncRead for AuthenticatedTlsStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        Pin::new(&mut self.inner).poll_read(cx, buf)
    }
}

impl tokio::io::AsyncWrite for AuthenticatedTlsStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write(cx, buf)
    }

    fn poll_flush(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }

    fn poll_write_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        bufs: &[std::io::IoSlice<'_>],
    ) -> std::task::Poll<std::io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
    }

    fn is_write_vectored(&self) -> bool {
        self.inner.is_write_vectored()
    }
}

impl axum::serve::Listener for TlsListener {
    type Io = AuthenticatedTlsStream;
    type Addr = SocketAddr;

    /// Yield the next fully-handshaken TLS connection.
    ///
    /// Cancel-safe: this is a plain `mpsc::Receiver::recv`, so cancelling
    /// the future (axum selects it against graceful shutdown) never loses
    /// a connection.
    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
        if let Some(pair) = self.rx.recv().await {
            return pair;
        }
        // The channel only closes if the acceptor task terminated, which
        // means the TcpListener is gone and the OS already refuses new
        // connections. `Listener::accept` is infallible and panicking is
        // forbidden, so park forever: existing connections keep being
        // served and graceful shutdown still completes.
        tracing::error!("TLS acceptor task terminated; no further connections will be accepted");
        std::future::pending().await
    }

    fn local_addr(&self) -> std::io::Result<Self::Addr> {
        Ok(self.local_addr)
    }
}

impl Drop for TlsListener {
    fn drop(&mut self) {
        // Stop accepting immediately and release the bound port. In-flight
        // handshake workers notice the closed channel and exit quietly.
        self.acceptor_task.abort();
    }
}

fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
    use rustls::pki_types::pem::PemObject;
    let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_file_iter(path)
        .map_err(|e| anyhow::anyhow!("failed to read certs from {}: {e}", path.display()))?
        .collect::<Result<_, _>>()
        .map_err(|e| anyhow::anyhow!("invalid cert in {}: {e}", path.display()))?;
    anyhow::ensure!(
        !certs.is_empty(),
        "no certificates found in {}",
        path.display()
    );
    Ok(certs)
}

fn load_client_auth_roots(
    path: &Path,
) -> anyhow::Result<(
    Vec<rustls::pki_types::CertificateDer<'static>>,
    Arc<RootCertStore>,
)> {
    let ca_certs = load_certs(path)?;
    let mut root_store = RootCertStore::empty();
    for cert in &ca_certs {
        root_store
            .add(cert.clone())
            .map_err(|error| anyhow::anyhow!("invalid CA cert: {error}"))?;
    }

    Ok((ca_certs, Arc::new(root_store)))
}

fn load_key(path: &Path) -> anyhow::Result<rustls::pki_types::PrivateKeyDer<'static>> {
    use rustls::pki_types::pem::PemObject;
    rustls::pki_types::PrivateKeyDer::from_pem_file(path)
        .map_err(|e| anyhow::anyhow!("failed to read key from {}: {e}", path.display()))
}

#[allow(
    clippy::unused_async,
    reason = "axum route handler signature requires `async fn` even when the body is synchronous"
)]
async fn healthz() -> impl IntoResponse {
    axum::Json(serde_json::json!({
        "status": "ok",
    }))
}

/// Build the `/version` JSON payload for a given server name and version.
///
/// Build metadata (`build_git_sha`, `build_timestamp`, `rust_version`) is
/// read at compile time from the `RMCP_SERVER_KIT_BUILD_SHA`,
/// `RMCP_SERVER_KIT_BUILD_TIME`, and `RMCP_SERVER_KIT_RUSTC_VERSION` env
/// vars. Unset values resolve to `"unknown"`.
fn version_payload(name: &str, version: &str) -> serde_json::Value {
    serde_json::json!({
        "name": name,
        "version": version,
        "build_git_sha": option_env!("RMCP_SERVER_KIT_BUILD_SHA").unwrap_or("unknown"),
        "build_timestamp": option_env!("RMCP_SERVER_KIT_BUILD_TIME").unwrap_or("unknown"),
        "rust_version": option_env!("RMCP_SERVER_KIT_RUSTC_VERSION").unwrap_or("unknown"),
        "mcpx_version": env!("CARGO_PKG_VERSION"),
    })
}

/// Pre-serialize the `/version` payload to immutable bytes.
///
/// This is called once at router-build time so per-request handling can
/// reuse a cheap `Arc<[u8]>` clone instead of re-serializing a
/// [`serde_json::Value`] on every hit.
///
/// Serialization of a flat `serde_json::Value` of static-string fields
/// cannot fail in practice; the fallback to `b"{}"` exists only to
/// satisfy the crate-wide `unwrap_used` / `expect_used` lint policy.
fn serialize_version_payload(name: &str, version: &str) -> Arc<[u8]> {
    let value = version_payload(name, version);
    serde_json::to_vec(&value).map_or_else(|_| Arc::from(&b"{}"[..]), Arc::from)
}

async fn readyz(check: ReadinessCheck) -> impl IntoResponse {
    let status = check().await;
    let ready = status
        .get("ready")
        .and_then(serde_json::Value::as_bool)
        .unwrap_or(false);
    let code = if ready {
        axum::http::StatusCode::OK
    } else {
        axum::http::StatusCode::SERVICE_UNAVAILABLE
    };
    (code, axum::Json(status))
}

/// Wait for SIGINT (ctrl-c) or SIGTERM (container stop).
///
/// On non-Unix platforms, only SIGINT is handled.
async fn shutdown_signal() {
    let ctrl_c = tokio::signal::ctrl_c();

    #[cfg(unix)]
    {
        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
            Ok(mut term) => {
                tokio::select! {
                    _ = ctrl_c => {}
                    _ = term.recv() => {}
                }
            }
            Err(e) => {
                tracing::warn!(error = %e, "failed to register SIGTERM handler, using SIGINT only");
                ctrl_c.await.ok();
            }
        }
    }

    #[cfg(not(unix))]
    {
        ctrl_c.await.ok();
    }
}

// -- Origin validation (MCP 2025-11-25 spec, section 2.0.1) --

/// Middleware that validates the `Origin` header on incoming HTTP requests.
///
/// Record HTTP request metrics (method, path, status, duration).
#[cfg(feature = "metrics")]
async fn metrics_middleware(
    metrics: Arc<crate::metrics::McpMetrics>,
    req: Request<Body>,
    next: Next,
) -> axum::response::Response {
    let method = req.method().to_string();
    let path = req.uri().path().to_owned();
    let start = std::time::Instant::now();

    let response = next.run(req).await;

    let status = response.status().as_u16().to_string();
    let duration = start.elapsed().as_secs_f64();

    metrics
        .http_requests_total
        .with_label_values(&[&method, &path, &status])
        .inc();
    metrics
        .http_request_duration_seconds
        .with_label_values(&[&method, &path])
        .observe(duration);

    response
}

/// OWASP security header hardening applied to every response.
///
/// Sets: `X-Content-Type-Options`, `X-Frame-Options`, `Cache-Control`,
/// `Referrer-Policy`, `Cross-Origin-Opener-Policy`, `Cross-Origin-Resource-Policy`,
/// `Cross-Origin-Embedder-Policy`, `Permissions-Policy`,
/// `X-Permitted-Cross-Domain-Policies`, `Content-Security-Policy`,
/// `X-DNS-Prefetch-Control`, and (when TLS is active) `Strict-Transport-Security`.
///
/// Each header's value can be customised via [`SecurityHeadersConfig`]
/// on [`McpServerConfig`]. See that type for the three-state semantic
/// (`None` = default, `Some("")` = omit, `Some(v)` = override).
async fn security_headers_middleware(
    is_tls: bool,
    cfg: Arc<SecurityHeadersConfig>,
    req: Request<Body>,
    next: Next,
) -> axum::response::Response {
    use axum::http::{HeaderName, header};

    let mut resp = next.run(req).await;
    let headers = resp.headers_mut();

    // Strip server identity headers to reduce information leakage.
    headers.remove(header::SERVER);
    headers.remove(HeaderName::from_static("x-powered-by"));

    apply_security_header(
        headers,
        header::X_CONTENT_TYPE_OPTIONS,
        cfg.x_content_type_options.as_deref(),
        "nosniff",
    );
    apply_security_header(
        headers,
        header::X_FRAME_OPTIONS,
        cfg.x_frame_options.as_deref(),
        "deny",
    );
    apply_security_header(
        headers,
        header::CACHE_CONTROL,
        cfg.cache_control.as_deref(),
        "no-store, max-age=0",
    );
    apply_security_header(
        headers,
        header::REFERRER_POLICY,
        cfg.referrer_policy.as_deref(),
        "no-referrer",
    );
    apply_security_header(
        headers,
        HeaderName::from_static("cross-origin-opener-policy"),
        cfg.cross_origin_opener_policy.as_deref(),
        "same-origin",
    );
    apply_security_header(
        headers,
        HeaderName::from_static("cross-origin-resource-policy"),
        cfg.cross_origin_resource_policy.as_deref(),
        "same-origin",
    );
    apply_security_header(
        headers,
        HeaderName::from_static("cross-origin-embedder-policy"),
        cfg.cross_origin_embedder_policy.as_deref(),
        "require-corp",
    );
    apply_security_header(
        headers,
        HeaderName::from_static("permissions-policy"),
        cfg.permissions_policy.as_deref(),
        "accelerometer=(), camera=(), geolocation=(), microphone=()",
    );
    apply_security_header(
        headers,
        HeaderName::from_static("x-permitted-cross-domain-policies"),
        cfg.x_permitted_cross_domain_policies.as_deref(),
        "none",
    );
    apply_security_header(
        headers,
        HeaderName::from_static("content-security-policy"),
        cfg.content_security_policy.as_deref(),
        "default-src 'none'; frame-ancestors 'none'",
    );
    apply_security_header(
        headers,
        HeaderName::from_static("x-dns-prefetch-control"),
        cfg.x_dns_prefetch_control.as_deref(),
        "off",
    );

    if is_tls {
        apply_security_header(
            headers,
            header::STRICT_TRANSPORT_SECURITY,
            cfg.strict_transport_security.as_deref(),
            "max-age=63072000; includeSubDomains",
        );
    }

    resp
}

/// Set a single security header on the response, honouring the
/// three-state override semantic (None = default, Some("") = omit,
/// Some(value) = override).
///
/// Defence-in-depth: if an override value somehow reaches this point
/// despite [`validate_security_headers`] having approved it (e.g. a
/// runtime mutation on a non-`Validated` field), we log at error level
/// and fall back to the static default rather than panicking. The
/// `Validated<McpServerConfig>` type makes that path unreachable in
/// well-typed code paths.
fn apply_security_header(
    headers: &mut axum::http::HeaderMap,
    name: axum::http::HeaderName,
    override_value: Option<&str>,
    default: &'static str,
) {
    use axum::http::HeaderValue;

    match override_value {
        None => {
            headers.insert(name, HeaderValue::from_static(default));
        }
        Some("") => {
            // Operator explicitly opted out of this header.
        }
        Some(v) => match HeaderValue::from_str(v) {
            Ok(hv) => {
                headers.insert(name, hv);
            }
            Err(err) => {
                tracing::error!(
                    header = %name,
                    error = %err,
                    "invalid security header override reached middleware; using default"
                );
                headers.insert(name, HeaderValue::from_static(default));
            }
        },
    }
}

/// Validate every non-empty entry in a [`SecurityHeadersConfig`].
///
/// - `None` and `Some("")` are accepted unconditionally (use-default and
///   omit, respectively).
/// - `Some(v)` is rejected if `axum::http::HeaderValue::from_str(v)` fails.
/// - `strict_transport_security` additionally rejects any value
///   containing `preload` (case-insensitive). Operators who genuinely
///   want to commit to the HSTS preload list must do so via a future
///   explicit `with_hsts_preload(true)` builder, not by smuggling
///   `preload` through this knob.
fn validate_security_headers(cfg: &SecurityHeadersConfig) -> Result<(), McpxError> {
    use axum::http::HeaderValue;

    let fields: &[(&str, Option<&str>)] = &[
        (
            "x_content_type_options",
            cfg.x_content_type_options.as_deref(),
        ),
        ("x_frame_options", cfg.x_frame_options.as_deref()),
        ("cache_control", cfg.cache_control.as_deref()),
        ("referrer_policy", cfg.referrer_policy.as_deref()),
        (
            "cross_origin_opener_policy",
            cfg.cross_origin_opener_policy.as_deref(),
        ),
        (
            "cross_origin_resource_policy",
            cfg.cross_origin_resource_policy.as_deref(),
        ),
        (
            "cross_origin_embedder_policy",
            cfg.cross_origin_embedder_policy.as_deref(),
        ),
        ("permissions_policy", cfg.permissions_policy.as_deref()),
        (
            "x_permitted_cross_domain_policies",
            cfg.x_permitted_cross_domain_policies.as_deref(),
        ),
        (
            "content_security_policy",
            cfg.content_security_policy.as_deref(),
        ),
        (
            "x_dns_prefetch_control",
            cfg.x_dns_prefetch_control.as_deref(),
        ),
        (
            "strict_transport_security",
            cfg.strict_transport_security.as_deref(),
        ),
    ];

    for (field, value) in fields {
        let Some(v) = value else { continue };
        if v.is_empty() {
            continue;
        }
        if let Err(err) = HeaderValue::from_str(v) {
            return Err(McpxError::Config(format!(
                "invalid security_headers.{field}: {err}"
            )));
        }
    }

    if let Some(v) = cfg.strict_transport_security.as_deref()
        && !v.is_empty()
        && v.to_ascii_lowercase().contains("preload")
    {
        return Err(McpxError::Config(format!(
            "invalid security_headers.strict_transport_security: {v:?} contains the `preload` directive; \
             HSTS preload must be opted into explicitly via a dedicated builder, not via this knob"
        )));
    }

    Ok(())
}

/// Append RFC 6749 §5.1 / RFC 6750 §5.4 cache and `Vary` headers required
/// on OAuth token-issuing responses.
///
/// `Cache-Control: no-store, max-age=0` is already applied globally by
/// [`security_headers_middleware`]; this middleware adds:
///
/// - `Pragma: no-cache` -- mandated by RFC 6749 §5.1 for HTTP/1.0 caches.
/// - `Vary: Authorization` -- mandated by RFC 6750 §5.4 for endpoints
///   whose response depends on the `Authorization` header.
///
/// Applied only to the OAuth proxy token-class endpoints (`/token`,
/// `/register`, `/introspect`, `/revoke`). `Vary` is appended (not
/// inserted) so any `Vary` value already present (e.g. `Accept-Encoding`
/// from a compression layer, or `Origin` from a CORS layer) is preserved.
#[cfg(feature = "oauth")]
async fn oauth_token_cache_headers_middleware(
    req: Request<Body>,
    next: Next,
) -> axum::response::Response {
    use axum::http::{HeaderValue, header};

    let mut resp = next.run(req).await;
    let headers = resp.headers_mut();
    headers.insert(header::PRAGMA, HeaderValue::from_static("no-cache"));
    headers.append(header::VARY, HeaderValue::from_static("Authorization"));
    resp
}

/// Normalize peer-address request extensions across listener branches.
///
/// The make-service installs `ConnectInfo<SocketAddr>` on the plain
/// listener but `ConnectInfo<TlsConnInfo>` on the TLS listener (the
/// latter additionally carries the connection-bound mTLS identity and
/// stays `pub(crate)` — see the anti-aliasing rationale on
/// [`TlsConnInfo`]). Application routes — in particular those merged via
/// [`McpServerConfig::with_extra_router`], which bypass the auth
/// middleware and its private fallback — could therefore not read the
/// peer address under TLS.
///
/// This middleware makes both branches look identical to every route and
/// inner middleware:
///
/// 1. mirrors the TLS peer address into `ConnectInfo<SocketAddr>` when
///    (and only when) it is absent, so stock axum-ecosystem extractors
///    work unmodified, and
/// 2. inserts the framework-owned [`PeerAddr`] extension on both
///    branches.
///
/// Precedence mirrors the auth middleware: an existing
/// `ConnectInfo<SocketAddr>` always wins and is never overwritten. The
/// peer address is deliberately not logged here.
async fn normalize_peer_addr_middleware(
    mut req: Request<Body>,
    next: Next,
) -> axum::response::Response {
    let direct = req
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|ci| ci.0);
    let from_tls = req
        .extensions()
        .get::<ConnectInfo<TlsConnInfo>>()
        .map(|ci| ci.0.addr);
    if let Some(addr) = direct.or(from_tls) {
        if direct.is_none() {
            req.extensions_mut().insert(ConnectInfo(addr));
        }
        req.extensions_mut().insert(PeerAddr::new(addr));
    }
    next.run(req).await
}

/// Per the MCP spec: if the Origin header is present and its value is not in
/// the allowed list, respond with 403 Forbidden. Requests without an Origin
/// header are allowed through (e.g. non-browser clients like curl, SDKs).
async fn origin_check_middleware(
    allowed: Arc<[String]>,
    log_request_headers: bool,
    req: Request<Body>,
    next: Next,
) -> axum::response::Response {
    let method = req.method().clone();
    let path = req.uri().path().to_owned();

    log_incoming_request(&method, &path, req.headers(), log_request_headers);

    if let Some(origin) = req.headers().get(axum::http::header::ORIGIN) {
        let origin_str = origin.to_str().unwrap_or("");
        if !allowed.iter().any(|a| a == origin_str) {
            tracing::warn!(
                origin = origin_str,
                %method,
                %path,
                allowed = ?&*allowed,
                "rejected request: Origin not allowed"
            );
            return (
                axum::http::StatusCode::FORBIDDEN,
                "Forbidden: Origin not allowed",
            )
                .into_response();
        }
    }
    next.run(req).await
}

/// Emit a DEBUG log for an incoming request, optionally including the full
/// (redacted) header set.
fn log_incoming_request(
    method: &axum::http::Method,
    path: &str,
    headers: &axum::http::HeaderMap,
    log_request_headers: bool,
) {
    if log_request_headers {
        tracing::debug!(
            %method,
            %path,
            headers = %format_request_headers_for_log(headers),
            "incoming request"
        );
    } else {
        tracing::debug!(%method, %path, "incoming request");
    }
}

fn format_request_headers_for_log(headers: &axum::http::HeaderMap) -> String {
    headers
        .iter()
        .map(|(k, v)| {
            let name = k.as_str();
            if name == "authorization" || name == "cookie" || name == "proxy-authorization" {
                format!("{name}: [REDACTED]")
            } else {
                format!("{name}: {}", v.to_str().unwrap_or("<non-utf8>"))
            }
        })
        .collect::<Vec<_>>()
        .join(", ")
}

// -- stdio transport --

/// Serve an MCP server over stdin/stdout (stdio transport).
///
/// # Security warnings
///
/// - **No authentication**: the parent process has full, unrestricted access.
/// - **No RBAC**: all tools are available regardless of policy.
/// - **No TLS**: messages travel over OS pipes in plaintext.
/// - **Single client**: only the parent process can connect.
/// - **No Origin validation**: not applicable to stdio.
///
/// Use this only when the MCP client spawns the server as a trusted subprocess
/// (e.g. Claude Desktop, VS Code Copilot). For network-accessible deployments,
/// use `serve()` (Streamable HTTP) instead.
///
/// # Errors
///
/// Returns [`McpxError::Startup`] if the handler fails to initialize or the
/// transport disconnects unexpectedly.
// NOTE: reported complexity 32/25 is driven entirely by `tracing::*!`
// macro expansion in this 18-line function (info/warn/info + two matches).
// There is nothing meaningful to extract; the allow stays.
#[allow(
    clippy::cognitive_complexity,
    reason = "complexity is purely tracing macro expansion (info/warn + match arms); 18 lines of straight-line code, nothing meaningful to extract"
)]
pub async fn serve_stdio<H>(handler: H) -> Result<(), McpxError>
where
    H: ServerHandler + 'static,
{
    use rmcp::ServiceExt as _;

    tracing::info!("stdio transport: serving on stdin/stdout");
    tracing::warn!("stdio mode: auth, RBAC, TLS, and Origin checks are DISABLED");

    let transport = rmcp::transport::io::stdio();

    let service = handler
        .serve(transport)
        .await
        .map_err(|e| McpxError::Startup(format!("stdio initialize failed: {e}")))?;

    if let Err(e) = service.waiting().await {
        tracing::warn!(error = %e, "stdio session ended with error");
    }
    tracing::info!("stdio session ended");
    Ok(())
}

#[cfg(test)]
mod tests {
    #![allow(
        clippy::unwrap_used,
        clippy::expect_used,
        clippy::panic,
        clippy::indexing_slicing,
        clippy::unwrap_in_result,
        clippy::print_stdout,
        clippy::print_stderr,
        deprecated,
        reason = "internal unit tests legitimately read/write the deprecated `pub` fields they were designed to verify"
    )]
    use std::{sync::Arc, time::Duration};

    use axum::{
        body::Body,
        http::{Request, StatusCode, header},
        response::IntoResponse,
    };
    use http_body_util::BodyExt;
    use tower::ServiceExt as _;

    use super::*;

    // -- McpServerConfig --

    #[test]
    fn server_config_new_defaults() {
        let cfg = McpServerConfig::new("0.0.0.0:8443", "test-server", "1.0.0");
        assert_eq!(cfg.bind_addr, "0.0.0.0:8443");
        assert_eq!(cfg.name, "test-server");
        assert_eq!(cfg.version, "1.0.0");
        assert!(cfg.tls_cert_path.is_none());
        assert!(cfg.tls_key_path.is_none());
        assert!(cfg.auth.is_none());
        assert!(cfg.rbac.is_none());
        assert!(cfg.allowed_origins.is_empty());
        assert!(cfg.tool_rate_limit.is_none());
        assert!(cfg.readiness_check.is_none());
        assert_eq!(cfg.max_request_body, 1024 * 1024);
        assert_eq!(cfg.request_timeout, Duration::from_mins(2));
        assert_eq!(cfg.shutdown_timeout, Duration::from_secs(30));
        assert!(!cfg.log_request_headers);
        assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(10));
        assert_eq!(cfg.max_concurrent_tls_handshakes, 256);
    }

    #[test]
    fn tls_handshake_builders_set_fields() {
        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
            .with_tls_handshake_timeout(Duration::from_secs(3))
            .with_max_concurrent_tls_handshakes(64);
        assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(3));
        assert_eq!(cfg.max_concurrent_tls_handshakes, 64);
    }

    #[test]
    fn validate_rejects_zero_tls_handshake_timeout() {
        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
            .with_tls_handshake_timeout(Duration::ZERO);
        let err = cfg.validate().expect_err("zero handshake timeout");
        assert!(err.to_string().contains("tls_handshake_timeout"));
    }

    #[test]
    fn validate_rejects_zero_max_concurrent_tls_handshakes() {
        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
            .with_max_concurrent_tls_handshakes(0);
        let err = cfg.validate().expect_err("zero handshake concurrency");
        assert!(err.to_string().contains("max_concurrent_tls_handshakes"));
    }

    #[test]
    fn validate_consumes_and_proves() {
        // Valid config -> Validated wrapper, original is consumed.
        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
        let validated = cfg.validate().expect("valid config");
        // as_inner() gives read-only access to inner fields.
        assert_eq!(validated.as_inner().name, "test-server");
        // into_inner recovers the raw value.
        let raw = validated.into_inner();
        assert_eq!(raw.name, "test-server");

        // Invalid config (zero max_request_body) -> Err.
        let mut bad = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
        bad.max_request_body = 0;
        assert!(bad.validate().is_err(), "zero body cap must fail validate");
    }

    #[test]
    fn validate_rejects_zero_max_concurrent_requests() {
        let cfg =
            McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_max_concurrent_requests(0);
        let err = cfg.validate().expect_err("zero concurrency cap must fail");
        assert!(
            format!("{err}").contains("max_concurrent_requests"),
            "error should mention max_concurrent_requests, got: {err}"
        );
    }

    #[test]
    fn validate_rejects_zero_max_tracked_keys() {
        // Defaults mirror auth::default_max_attempts / default_idle_eviction
        // (module-private in auth.rs); spelled out here for review clarity.
        let rl = crate::auth::RateLimitConfig {
            max_attempts_per_minute: 30,
            pre_auth_max_per_minute: None,
            max_tracked_keys: 0,
            idle_eviction: Duration::from_secs(15 * 60),
        };
        let auth_cfg = AuthConfig {
            enabled: true,
            api_keys: Vec::new(),
            mtls: None,
            rate_limit: Some(rl),
            #[cfg(feature = "oauth")]
            oauth: None,
        };
        let cfg = McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_auth(auth_cfg);
        let err = cfg.validate().expect_err("zero max_tracked_keys must fail");
        assert!(
            format!("{err}").contains("max_tracked_keys"),
            "error should mention max_tracked_keys, got: {err}"
        );
    }

    #[test]
    fn derive_allowed_hosts_includes_public_host() {
        let hosts = derive_allowed_hosts("0.0.0.0:8080", Some("https://mcp.example.com/mcp"));
        assert!(
            hosts.iter().any(|h| h == "mcp.example.com"),
            "public_url host must be allowed"
        );
    }

    #[test]
    fn derive_allowed_hosts_includes_bind_authority() {
        let hosts = derive_allowed_hosts("127.0.0.1:8080", None);
        assert!(
            hosts.iter().any(|h| h == "127.0.0.1"),
            "bind host must be allowed"
        );
        assert!(
            hosts.iter().any(|h| h == "127.0.0.1:8080"),
            "bind authority must be allowed"
        );
    }

    // -- healthz --

    #[tokio::test]
    async fn healthz_returns_ok_json() {
        let resp = healthz().await.into_response();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["status"], "ok");
        assert!(
            json.get("name").is_none(),
            "healthz must not expose server name"
        );
        assert!(
            json.get("version").is_none(),
            "healthz must not expose version"
        );
    }

    // -- readyz --

    #[tokio::test]
    async fn readyz_returns_ok_when_ready() {
        let check: ReadinessCheck =
            Arc::new(|| Box::pin(async { serde_json::json!({"ready": true, "db": "connected"}) }));
        let resp = readyz(check).await.into_response();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["ready"], true);
        assert!(
            json.get("name").is_none(),
            "readyz must not expose server name"
        );
        assert!(
            json.get("version").is_none(),
            "readyz must not expose version"
        );
        assert_eq!(json["db"], "connected");
    }

    #[tokio::test]
    async fn readyz_returns_503_when_not_ready() {
        let check: ReadinessCheck =
            Arc::new(|| Box::pin(async { serde_json::json!({"ready": false}) }));
        let resp = readyz(check).await.into_response();
        assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
    }

    #[tokio::test]
    async fn readyz_returns_503_when_ready_missing() {
        let check: ReadinessCheck =
            Arc::new(|| Box::pin(async { serde_json::json!({"status": "starting"}) }));
        let resp = readyz(check).await.into_response();
        // Missing "ready" field defaults to false -> 503
        assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
    }

    // -- normalize_peer_addr_middleware / PeerAddr --

    /// Build a test router that reports the request's peer-address
    /// extensions as `"<ConnectInfo>|<PeerAddr>"` (empty when absent).
    fn peer_probe_router() -> axum::Router {
        async fn probe(req: Request<Body>) -> String {
            let ci = req
                .extensions()
                .get::<ConnectInfo<SocketAddr>>()
                .map(|c| c.0.to_string())
                .unwrap_or_default();
            let pa = req
                .extensions()
                .get::<PeerAddr>()
                .map(|p| p.addr.to_string())
                .unwrap_or_default();
            format!("{ci}|{pa}")
        }
        axum::Router::new()
            .route("/probe", axum::routing::get(probe))
            .layer(axum::middleware::from_fn(normalize_peer_addr_middleware))
    }

    async fn body_string(resp: axum::response::Response) -> String {
        let bytes = resp.into_body().collect().await.unwrap().to_bytes();
        String::from_utf8(bytes.to_vec()).unwrap()
    }

    #[tokio::test]
    async fn normalize_preserves_existing_connect_info_and_mirrors_peer_addr() {
        // Precedence proof: when both extensions exist with DIFFERENT
        // addresses, ConnectInfo<SocketAddr> wins and is never overwritten.
        let plain: SocketAddr = "10.0.0.1:1111".parse().unwrap();
        let tls: SocketAddr = "10.0.0.2:2222".parse().unwrap();
        let req = Request::builder()
            .uri("/probe")
            .extension(ConnectInfo(plain))
            .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
            .body(Body::empty())
            .unwrap();
        let resp = peer_probe_router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(body_string(resp).await, format!("{plain}|{plain}"));
    }

    #[tokio::test]
    async fn normalize_inserts_connect_info_and_peer_addr_from_tls() {
        let tls: SocketAddr = "192.168.1.7:50443".parse().unwrap();
        let req = Request::builder()
            .uri("/probe")
            .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
            .body(Body::empty())
            .unwrap();
        let resp = peer_probe_router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(body_string(resp).await, format!("{tls}|{tls}"));
    }

    #[tokio::test]
    async fn normalize_no_op_without_any_connect_info() {
        let req = Request::builder()
            .uri("/probe")
            .body(Body::empty())
            .unwrap();
        let resp = peer_probe_router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(body_string(resp).await, "|");
    }

    #[tokio::test]
    async fn peer_addr_extractor_rejects_when_absent() {
        async fn h(peer: PeerAddr) -> String {
            peer.addr.to_string()
        }
        let app = axum::Router::new().route("/p", axum::routing::get(h));
        let req = Request::builder().uri("/p").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
    }

    #[tokio::test]
    async fn peer_addr_extractor_returns_value_when_present() {
        async fn h(peer: PeerAddr) -> String {
            peer.addr.to_string()
        }
        let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap();
        let app = axum::Router::new().route("/p", axum::routing::get(h));
        let req = Request::builder()
            .uri("/p")
            .extension(PeerAddr::new(addr))
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(body_string(resp).await, addr.to_string());
    }

    #[tokio::test]
    async fn peer_addr_via_extension_extractor() {
        async fn h(axum::Extension(peer): axum::Extension<PeerAddr>) -> String {
            peer.addr.to_string()
        }
        let addr: SocketAddr = "127.0.0.1:4242".parse().unwrap();
        let app = axum::Router::new().route("/p", axum::routing::get(h));
        let req = Request::builder()
            .uri("/p")
            .extension(PeerAddr::new(addr))
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(body_string(resp).await, addr.to_string());
    }

    // -- origin_check_middleware --

    /// Build a test router with origin check middleware and a simple handler.
    fn origin_router(origins: Vec<String>, log_request_headers: bool) -> axum::Router {
        let allowed: Arc<[String]> = Arc::from(origins);
        axum::Router::new()
            .route("/test", axum::routing::get(|| async { "ok" }))
            .layer(axum::middleware::from_fn(move |req, next| {
                let a = Arc::clone(&allowed);
                origin_check_middleware(a, log_request_headers, req, next)
            }))
    }

    #[tokio::test]
    async fn origin_allowed_passes() {
        let app = origin_router(vec!["http://localhost:3000".into()], false);
        let req = Request::builder()
            .uri("/test")
            .header(header::ORIGIN, "http://localhost:3000")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn origin_rejected_returns_403() {
        let app = origin_router(vec!["http://localhost:3000".into()], false);
        let req = Request::builder()
            .uri("/test")
            .header(header::ORIGIN, "http://evil.com")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn no_origin_header_passes() {
        let app = origin_router(vec!["http://localhost:3000".into()], false);
        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn empty_allowlist_rejects_any_origin() {
        let app = origin_router(vec![], false);
        let req = Request::builder()
            .uri("/test")
            .header(header::ORIGIN, "http://anything.com")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn empty_allowlist_passes_without_origin() {
        let app = origin_router(vec![], false);
        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[test]
    fn format_request_headers_redacts_sensitive_values() {
        let mut headers = axum::http::HeaderMap::new();
        headers.insert("authorization", "Bearer secret-token".parse().unwrap());
        headers.insert("cookie", "sid=abc".parse().unwrap());
        headers.insert("x-request-id", "req-123".parse().unwrap());

        let out = format_request_headers_for_log(&headers);
        assert!(out.contains("authorization: [REDACTED]"));
        assert!(out.contains("cookie: [REDACTED]"));
        assert!(out.contains("x-request-id: req-123"));
        assert!(!out.contains("secret-token"));
    }

    // -- security_headers_middleware --

    fn security_router(is_tls: bool) -> axum::Router {
        security_router_with(is_tls, SecurityHeadersConfig::default())
    }

    fn security_router_with(is_tls: bool, cfg: SecurityHeadersConfig) -> axum::Router {
        let cfg = Arc::new(cfg);
        axum::Router::new()
            .route("/test", axum::routing::get(|| async { "ok" }))
            .layer(axum::middleware::from_fn(move |req, next| {
                let c = Arc::clone(&cfg);
                security_headers_middleware(is_tls, c, req, next)
            }))
    }

    #[tokio::test]
    async fn security_headers_set_on_response() {
        let app = security_router(false);
        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let h = resp.headers();
        assert_eq!(h.get("x-content-type-options").unwrap(), "nosniff");
        assert_eq!(h.get("x-frame-options").unwrap(), "deny");
        assert_eq!(h.get("cache-control").unwrap(), "no-store, max-age=0");
        assert_eq!(h.get("referrer-policy").unwrap(), "no-referrer");
        assert_eq!(h.get("cross-origin-opener-policy").unwrap(), "same-origin");
        assert_eq!(
            h.get("cross-origin-resource-policy").unwrap(),
            "same-origin"
        );
        assert_eq!(
            h.get("cross-origin-embedder-policy").unwrap(),
            "require-corp"
        );
        assert_eq!(h.get("x-permitted-cross-domain-policies").unwrap(), "none");
        assert!(
            h.get("permissions-policy")
                .unwrap()
                .to_str()
                .unwrap()
                .contains("camera=()"),
            "permissions-policy must restrict browser features"
        );
        assert_eq!(
            h.get("content-security-policy").unwrap(),
            "default-src 'none'; frame-ancestors 'none'"
        );
        assert_eq!(h.get("x-dns-prefetch-control").unwrap(), "off");
        // No HSTS when TLS is off.
        assert!(h.get("strict-transport-security").is_none());
    }

    #[tokio::test]
    async fn hsts_set_when_tls_enabled() {
        let app = security_router(true);
        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();

        let hsts = resp.headers().get("strict-transport-security").unwrap();
        assert!(
            hsts.to_str().unwrap().contains("max-age=63072000"),
            "HSTS must set 2-year max-age"
        );
    }

    // -- SecurityHeadersConfig validation + override semantics --

    /// Build a minimal config with a custom SecurityHeadersConfig and
    /// drive it through `check()`. Returns the result so individual
    /// tests can assert on success or specific error messages.
    fn check_with_security_headers(headers: SecurityHeadersConfig) -> Result<(), McpxError> {
        let cfg =
            McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_security_headers(headers);
        cfg.check()
    }

    #[test]
    fn security_headers_config_default_validates() {
        check_with_security_headers(SecurityHeadersConfig::default())
            .expect("default SecurityHeadersConfig must validate");
    }

    #[test]
    fn security_headers_config_validate_accepts_empty_string() {
        // All twelve fields explicitly set to "" -> omit-everything mode.
        let h = SecurityHeadersConfig {
            x_content_type_options: Some(String::new()),
            x_frame_options: Some(String::new()),
            cache_control: Some(String::new()),
            referrer_policy: Some(String::new()),
            cross_origin_opener_policy: Some(String::new()),
            cross_origin_resource_policy: Some(String::new()),
            cross_origin_embedder_policy: Some(String::new()),
            permissions_policy: Some(String::new()),
            x_permitted_cross_domain_policies: Some(String::new()),
            content_security_policy: Some(String::new()),
            x_dns_prefetch_control: Some(String::new()),
            strict_transport_security: Some(String::new()),
        };
        check_with_security_headers(h).expect("Some(\"\") on every field must validate (omit-all)");
    }

    #[test]
    fn security_headers_config_validate_rejects_bad_value() {
        // 0x07 (BEL) is not a valid HTTP header value char.
        let h = SecurityHeadersConfig {
            referrer_policy: Some("\u{0007}".into()),
            ..SecurityHeadersConfig::default()
        };
        let err = check_with_security_headers(h)
            .expect_err("control char in referrer_policy must reject");
        let msg = err.to_string();
        assert!(
            msg.contains("referrer_policy"),
            "error must name the offending field, got: {msg}"
        );
    }

    #[test]
    fn security_headers_config_validate_rejects_hsts_preload() {
        let h = SecurityHeadersConfig {
            strict_transport_security: Some("max-age=63072000; includeSubDomains; preload".into()),
            ..SecurityHeadersConfig::default()
        };
        let err = check_with_security_headers(h).expect_err("HSTS with preload must reject");
        let msg = err.to_string();
        assert!(
            msg.contains("strict_transport_security"),
            "error must name the field, got: {msg}"
        );
        assert!(
            msg.to_lowercase().contains("preload"),
            "error must mention `preload`, got: {msg}"
        );
    }

    #[test]
    fn security_headers_config_validate_rejects_hsts_preload_uppercase() {
        // Case-insensitive match.
        let h = SecurityHeadersConfig {
            strict_transport_security: Some("max-age=600; PRELOAD".into()),
            ..SecurityHeadersConfig::default()
        };
        check_with_security_headers(h).expect_err("HSTS preload check must be case-insensitive");
    }

    #[tokio::test]
    async fn security_headers_override_honored() {
        // Override X-Frame-Options to SAMEORIGIN.
        let h = SecurityHeadersConfig {
            x_frame_options: Some("SAMEORIGIN".into()),
            ..SecurityHeadersConfig::default()
        };
        let app = security_router_with(false, h);
        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let xfo = resp.headers().get("x-frame-options").unwrap();
        assert_eq!(xfo, "SAMEORIGIN");
    }

    #[tokio::test]
    async fn security_headers_empty_string_omits() {
        // Empty string on referrer-policy -> header absent.
        let h = SecurityHeadersConfig {
            referrer_policy: Some(String::new()),
            ..SecurityHeadersConfig::default()
        };
        let app = security_router_with(false, h);
        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        assert!(
            resp.headers().get("referrer-policy").is_none(),
            "Some(\"\") must omit the header"
        );
        // Other defaults should still be present.
        assert_eq!(
            resp.headers().get("x-content-type-options").unwrap(),
            "nosniff"
        );
    }

    #[tokio::test]
    async fn security_headers_hsts_only_when_tls() {
        // HSTS override is irrelevant when TLS is off.
        let h = SecurityHeadersConfig {
            strict_transport_security: Some("max-age=600".into()),
            ..SecurityHeadersConfig::default()
        };
        let app = security_router_with(false, h);
        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert!(
            resp.headers().get("strict-transport-security").is_none(),
            "HSTS must remain absent on plaintext deployments even with override"
        );
    }

    // -- oauth_token_cache_headers_middleware --

    #[cfg(feature = "oauth")]
    #[tokio::test]
    async fn oauth_token_cache_headers_set_pragma_and_vary() {
        let app = axum::Router::new()
            .route("/token", axum::routing::post(|| async { "{}" }))
            .layer(axum::middleware::from_fn(
                oauth_token_cache_headers_middleware,
            ));
        let req = Request::builder()
            .method("POST")
            .uri("/token")
            .body(Body::from("{}"))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let h = resp.headers();
        assert_eq!(
            h.get("pragma").unwrap(),
            "no-cache",
            "RFC 6749 §5.1: token responses must set Pragma: no-cache"
        );
        let vary_values: Vec<String> = h
            .get_all("vary")
            .iter()
            .filter_map(|v| v.to_str().ok().map(str::to_owned))
            .collect();
        assert!(
            vary_values
                .iter()
                .any(|v| v.eq_ignore_ascii_case("Authorization")),
            "RFC 6750 §5.4: Vary must include Authorization, got {vary_values:?}"
        );
    }

    #[cfg(feature = "oauth")]
    #[tokio::test]
    async fn oauth_token_cache_headers_preserve_existing_vary() {
        // Simulates a handler/layer that already set `Vary: Accept-Encoding`
        // (e.g. compression). Our middleware must APPEND, not REPLACE.
        let app = axum::Router::new()
            .route(
                "/token",
                axum::routing::post(|| async {
                    axum::response::Response::builder()
                        .header("vary", "Accept-Encoding")
                        .body(axum::body::Body::from("{}"))
                        .unwrap()
                }),
            )
            .layer(axum::middleware::from_fn(
                oauth_token_cache_headers_middleware,
            ));
        let req = Request::builder()
            .method("POST")
            .uri("/token")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();

        let vary: Vec<String> = resp
            .headers()
            .get_all("vary")
            .iter()
            .filter_map(|v| v.to_str().ok().map(str::to_owned))
            .collect();
        assert!(
            vary.iter().any(|v| v.contains("Accept-Encoding")),
            "must preserve pre-existing Vary value, got {vary:?}"
        );
        assert!(
            vary.iter().any(|v| v.contains("Authorization")),
            "must append Authorization to Vary, got {vary:?}"
        );
    }

    // -- version endpoint --

    #[test]
    fn version_payload_contains_expected_fields() {
        let v = version_payload("my-server", "1.2.3");
        assert_eq!(v["name"], "my-server");
        assert_eq!(v["version"], "1.2.3");
        assert!(v["build_git_sha"].is_string());
        assert!(v["build_timestamp"].is_string());
        assert!(v["rust_version"].is_string());
        assert!(v["mcpx_version"].is_string());
    }

    // -- concurrency limit layer --

    #[tokio::test]
    async fn concurrency_limit_layer_composes_and_serves() {
        // We only assert the layer stack compiles and a single request
        // below the cap still succeeds. True back-pressure behaviour
        // requires a live HTTP server and is covered by integration tests.
        let app = axum::Router::new()
            .route("/ok", axum::routing::get(|| async { "ok" }))
            .layer(
                tower::ServiceBuilder::new()
                    .layer(axum::error_handling::HandleErrorLayer::new(
                        |_err: tower::BoxError| async { StatusCode::SERVICE_UNAVAILABLE },
                    ))
                    .layer(tower::load_shed::LoadShedLayer::new())
                    .layer(tower::limit::ConcurrencyLimitLayer::new(4)),
            );
        let resp = app
            .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap())
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    // -- compression layer --

    #[tokio::test]
    async fn compression_layer_gzip_encodes_response() {
        use tower_http::compression::Predicate as _;

        let big_body = "a".repeat(4096);
        let app = axum::Router::new()
            .route(
                "/big",
                axum::routing::get(move || {
                    let body = big_body.clone();
                    async move { body }
                }),
            )
            .layer(
                tower_http::compression::CompressionLayer::new()
                    .gzip(true)
                    .br(true)
                    .compress_when(
                        tower_http::compression::DefaultPredicate::new()
                            .and(tower_http::compression::predicate::SizeAbove::new(1024)),
                    ),
            );

        let req = Request::builder()
            .uri("/big")
            .header(header::ACCEPT_ENCODING, "gzip")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(header::CONTENT_ENCODING).unwrap(),
            "gzip"
        );
    }

    // -- TlsListener handshake timeout --

    #[tokio::test]
    async fn tls_handshake_timeout_reaps_idle_connections() {
        use tokio::io::AsyncReadExt as _;

        let _ = rustls::crypto::ring::default_provider().install_default();

        // Self-signed cert material on disk (TlsListener::new takes paths).
        let key = rcgen::KeyPair::generate().expect("generate key");
        let cert = rcgen::CertificateParams::new(vec!["localhost".to_owned()])
            .expect("cert params")
            .self_signed(&key)
            .expect("self-signed cert");
        let dir = std::env::temp_dir().join(format!(
            "rmcp-server-kit-hs-timeout-{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .expect("clock after epoch")
                .as_nanos()
        ));
        tokio::fs::create_dir_all(&dir).await.expect("temp dir");
        let cert_path = dir.join("server.crt");
        let key_path = dir.join("server.key");
        tokio::fs::write(&cert_path, cert.pem())
            .await
            .expect("write cert");
        tokio::fs::write(&key_path, key.serialize_pem())
            .await
            .expect("write key");

        let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
        let tls = TlsListener::new(
            listener,
            &cert_path,
            &key_path,
            None,
            None,
            Duration::from_millis(200),
            8, // custom concurrency cap: proves the plumbing end-to-end
        )
        .expect("tls listener");
        let addr = axum::serve::Listener::local_addr(&tls).expect("local addr");

        // Connect and send NOTHING: the handshake worker must time out
        // after 200ms and drop the stream, which the client observes as
        // EOF or a reset well within the 2s deadline.
        let mut idle = tokio::net::TcpStream::connect(addr).await.expect("connect");
        let mut buf = [0_u8; 16];
        let read = tokio::time::timeout(Duration::from_secs(2), idle.read(&mut buf))
            .await
            .expect("server must reap the idle handshake within its timeout");
        match read {
            Ok(0) | Err(_) => {} // EOF or reset: connection was dropped.
            Ok(n) => panic!("unexpected {n} bytes from server during reaped handshake"),
        }

        drop(tls);
    }
}