Skip to main content

rmcp_server_kit/
transport.rs

1use std::{
2    future::Future,
3    net::{IpAddr, SocketAddr},
4    path::{Path, PathBuf},
5    pin::Pin,
6    sync::Arc,
7    time::Duration,
8};
9
10use arc_swap::ArcSwap;
11use axum::{
12    body::Body,
13    extract::{ConnectInfo, Request},
14    middleware::Next,
15    response::IntoResponse,
16};
17use rmcp::{
18    ServerHandler,
19    transport::streamable_http_server::{
20        StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
21    },
22};
23use rustls::RootCertStore;
24use tokio::{
25    net::TcpListener,
26    sync::{Semaphore, mpsc},
27};
28use tokio_util::sync::CancellationToken;
29
30use crate::{
31    auth::{
32        AuthConfig, AuthIdentity, AuthState, MtlsConfig, TlsConnInfo, auth_middleware,
33        build_rate_limiter, extract_mtls_identity,
34    },
35    bounded_limiter::BoundedKeyedLimiter,
36    error::McpxError,
37    mtls_revocation::{self, CrlSet, DynamicClientCertVerifier},
38    rbac::{RbacPolicy, ToolRateLimiter, build_tool_rate_limiter, rbac_middleware},
39};
40
41/// Map an internal `anyhow::Error` chain into a public [`McpxError::Startup`]
42/// at the public API boundary, flattening the chain via the alternate
43/// formatter so callers see the full causal path.
44#[allow(
45    clippy::needless_pass_by_value,
46    reason = "consumed at .map_err(anyhow_to_startup) call sites; by-value matches the closure shape"
47)]
48fn anyhow_to_startup(e: anyhow::Error) -> McpxError {
49    McpxError::Startup(format!("{e:#}"))
50}
51
52/// Map a `std::io::Error` produced during server startup into a public
53/// [`McpxError::Startup`]. We deliberately do not use the [`McpxError::Io`]
54/// `From` impl here because startup-phase IO errors (bind, listener) are
55/// semantically distinct from request-time IO errors and should surface
56/// the originating operation in the message.
57#[allow(
58    clippy::needless_pass_by_value,
59    reason = "consumed at .map_err(|e| io_to_startup(...)) call sites; by-value matches the closure shape"
60)]
61fn io_to_startup(op: &str, e: std::io::Error) -> McpxError {
62    McpxError::Startup(format!("{op}: {e}"))
63}
64
65/// Async readiness check callback for the `/readyz` endpoint.
66///
67/// Returns a JSON object with at least a `"ready"` boolean.
68/// When `ready` is false, the endpoint returns HTTP 503.
69pub type ReadinessCheck =
70    Arc<dyn Fn() -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> + Send + Sync>;
71
72/// Direct socket peer address of the current HTTP/TLS connection.
73///
74/// Inserted as a request extension into every request served by [`serve`] —
75/// on both the plain and the TLS listener — and extractable in any axum
76/// handler, including routes mounted via
77/// [`McpServerConfig::with_extra_router`] (which bypass auth/RBAC and
78/// therefore often need the peer address for their own protection, e.g.
79/// per-IP rate limiting).
80///
81/// The same address is also mirrored into
82/// [`axum::extract::ConnectInfo<SocketAddr>`] on the TLS listener, so
83/// third-party middleware that expects the stock axum extension (e.g.
84/// per-IP rate-limit key extractors) works unmodified under TLS.
85///
86/// # Semantics
87///
88/// - **Direct peer only.** This is the socket's remote address. Behind an
89///   L4/L7 proxy or load balancer it is the proxy's address; the framework
90///   performs **no** `X-Forwarded-For` / `Forwarded` interpretation.
91/// - **Available on HTTP and TLS** transports alike ([`serve`]).
92/// - **Absent under [`serve_stdio`]** — a stdio session has no network
93///   peer (stdio bypasses the HTTP stack entirely).
94/// - The separate Prometheus metrics listener (feature `metrics`) is a
95///   different router and does not carry this extension.
96///
97/// # Privacy
98///
99/// `PeerAddr` exposes raw peer network metadata. The framework deliberately
100/// never logs it on its own; whether to log or persist peer addresses is
101/// application policy.
102///
103/// # Example
104///
105/// ```no_run
106/// use axum::{Router, routing::get};
107/// use rmcp_server_kit::transport::{McpServerConfig, PeerAddr};
108///
109/// async fn whoami(peer: PeerAddr) -> String {
110///     peer.addr.ip().to_string()
111/// }
112///
113/// let _config = McpServerConfig::new("127.0.0.1:8443", "my-server", "1.0.0")
114///     .with_extra_router(Router::new().route("/whoami", get(whoami)));
115/// ```
116#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
117#[non_exhaustive]
118pub struct PeerAddr {
119    /// Direct socket peer of this connection.
120    pub addr: SocketAddr,
121}
122
123impl PeerAddr {
124    /// Construct a new [`PeerAddr`]. Framework-internal: downstream code
125    /// receives `PeerAddr` via request extensions and never constructs it.
126    #[must_use]
127    pub(crate) const fn new(addr: SocketAddr) -> Self {
128        Self { addr }
129    }
130}
131
132/// Extract [`PeerAddr`] from request extensions.
133///
134/// # Rejection
135///
136/// Responds `500 Internal Server Error` when the extension is missing.
137/// A missing `PeerAddr` means the handler is not running under [`serve`]
138/// (e.g. the router was mounted on a hand-rolled listener) — a wiring
139/// bug, not a client error.
140impl<S: Send + Sync> axum::extract::FromRequestParts<S> for PeerAddr {
141    type Rejection = (axum::http::StatusCode, &'static str);
142
143    #[allow(
144        clippy::unused_async_trait_impl,
145        reason = "async is mandated by the axum FromRequestParts trait signature; this impl only reads a request extension synchronously"
146    )]
147    async fn from_request_parts(
148        parts: &mut axum::http::request::Parts,
149        _state: &S,
150    ) -> Result<Self, Self::Rejection> {
151        parts.extensions.get::<Self>().copied().ok_or((
152            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
153            "peer address unavailable: not running under rmcp-server-kit serve()",
154        ))
155    }
156}
157
158/// Resolved client IP of the current request.
159///
160/// Inserted as a request extension on every request served by [`serve`],
161/// right after [`PeerAddr`]. Equals the direct peer's IP unless
162/// **trusted-forwarder mode** is active
163/// ([`McpServerConfig::with_trusted_proxies`]) and the request arrived
164/// through a trusted proxy with a verifiable forwarding chain — in that
165/// case it is the rightmost-untrusted address from `X-Forwarded-For`
166/// (or RFC 7239 `Forwarded`, per
167/// [`McpServerConfig::with_forwarded_header`]).
168///
169/// All built-in per-IP rate limiters key by this value. [`PeerAddr`]
170/// keeps its direct-socket-peer contract unchanged; applications that
171/// need provenance can compare `ClientIp.ip` with `PeerAddr.addr.ip()`.
172///
173/// # Security
174///
175/// Resolution only ever activates when the **direct peer** is inside the
176/// operator's trusted-proxy CIDRs; every ambiguous chain (malformed or
177/// obfuscated entries, all-trusted chains, header bombs) falls back to
178/// the direct peer, never to a header value. The framework never logs
179/// this value outside rate-limit deny paths.
180#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
181#[non_exhaustive]
182pub struct ClientIp {
183    /// Resolved client IP (direct peer unless trusted-forwarder resolution applied).
184    pub ip: IpAddr,
185}
186
187impl ClientIp {
188    /// Construct a new [`ClientIp`]. Framework-internal: downstream code
189    /// receives `ClientIp` via request extensions and never constructs it.
190    #[must_use]
191    pub(crate) const fn new(ip: IpAddr) -> Self {
192        Self { ip }
193    }
194}
195
196/// Which forwarding header trusted-forwarder mode reads.
197///
198/// TOML wire values are kebab-case: `"x-forwarded-for"` (default when
199/// unset) and `"forwarded"`.
200#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Deserialize)]
201#[serde(rename_all = "kebab-case")]
202#[non_exhaustive]
203pub enum ForwardedHeaderMode {
204    /// De-facto standard `X-Forwarded-For` list (nginx, HAProxy, CDNs).
205    XForwardedFor,
206    /// RFC 7239 `Forwarded` header (`for=` parameters).
207    Forwarded,
208}
209
210/// Pre-parsed trusted-forwarder configuration captured by the
211/// peer-normalization middleware.
212struct ForwardResolver {
213    trusted: Vec<ipnet::IpNet>,
214    mode: ForwardedHeaderMode,
215}
216
217/// Per-header overrides for the OWASP security headers emitted by the
218/// global response middleware.
219///
220/// Each field follows a three-state semantic:
221///
222/// | Value         | Behaviour                                                |
223/// |---------------|----------------------------------------------------------|
224/// | `None`        | Use the built-in default (current behaviour).            |
225/// | `Some("")`    | **Omit** the header entirely from responses.             |
226/// | `Some(value)` | Emit `header: value`. Validated at config-load time.     |
227///
228/// All non-empty values are validated via
229/// [`axum::http::HeaderValue::from_str`] inside
230/// [`McpServerConfig::validate`]; invalid values fail fast before the
231/// server starts accepting traffic.
232///
233/// `Strict-Transport-Security` has an additional rule: the substring
234/// `preload` (case-insensitive) is rejected. Operators who want to
235/// commit to the HSTS preload list must do so via a future explicit
236/// builder method, not by smuggling it through this knob.
237#[derive(Debug, Clone, Default)]
238#[non_exhaustive]
239pub struct SecurityHeadersConfig {
240    /// Override for `X-Content-Type-Options`. Default: `nosniff`.
241    pub x_content_type_options: Option<String>,
242    /// Override for `X-Frame-Options`. Default: `deny`.
243    pub x_frame_options: Option<String>,
244    /// Override for `Cache-Control`. Default: `no-store, max-age=0`.
245    pub cache_control: Option<String>,
246    /// Override for `Referrer-Policy`. Default: `no-referrer`.
247    pub referrer_policy: Option<String>,
248    /// Override for `Cross-Origin-Opener-Policy`. Default: `same-origin`.
249    pub cross_origin_opener_policy: Option<String>,
250    /// Override for `Cross-Origin-Resource-Policy`. Default: `same-origin`.
251    pub cross_origin_resource_policy: Option<String>,
252    /// Override for `Cross-Origin-Embedder-Policy`. Default: `require-corp`.
253    pub cross_origin_embedder_policy: Option<String>,
254    /// Override for `Permissions-Policy`. Default:
255    /// `accelerometer=(), camera=(), geolocation=(), microphone=()`.
256    pub permissions_policy: Option<String>,
257    /// Override for `X-Permitted-Cross-Domain-Policies`. Default: `none`.
258    pub x_permitted_cross_domain_policies: Option<String>,
259    /// Override for `Content-Security-Policy`. Default:
260    /// `default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests`.
261    pub content_security_policy: Option<String>,
262    /// Override for `X-DNS-Prefetch-Control`. Default: `off`.
263    pub x_dns_prefetch_control: Option<String>,
264    /// Override for `Strict-Transport-Security`. Default (TLS only):
265    /// `max-age=63072000; includeSubDomains`. Only emitted when TLS is
266    /// active; the override is ignored on plaintext deployments. The
267    /// substring `preload` (any case) is rejected by the validator.
268    pub strict_transport_security: Option<String>,
269}
270
271/// Configuration for the MCP server.
272#[allow(
273    missing_debug_implementations,
274    reason = "contains callback/trait objects that don't impl Debug"
275)]
276#[allow(
277    clippy::struct_excessive_bools,
278    reason = "server configuration naturally has many boolean feature flags"
279)]
280#[non_exhaustive]
281pub struct McpServerConfig {
282    /// Socket address the MCP HTTP server binds to.
283    #[deprecated(
284        since = "0.13.0",
285        note = "use McpServerConfig::new() / with_bind_addr(); direct field access will become pub(crate) in a future major release"
286    )]
287    pub bind_addr: String,
288    /// Server name advertised via MCP `initialize`.
289    #[deprecated(
290        since = "0.13.0",
291        note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
292    )]
293    pub name: String,
294    /// Server version advertised via MCP `initialize`.
295    #[deprecated(
296        since = "0.13.0",
297        note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
298    )]
299    pub version: String,
300    /// Path to the TLS certificate (PEM). Required for TLS/mTLS.
301    #[deprecated(
302        since = "0.13.0",
303        note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
304    )]
305    pub tls_cert_path: Option<PathBuf>,
306    /// Path to the TLS private key (PEM). Required for TLS/mTLS.
307    #[deprecated(
308        since = "0.13.0",
309        note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
310    )]
311    pub tls_key_path: Option<PathBuf>,
312    /// Optional authentication config. When `Some` and `enabled`, auth
313    /// is enforced on `/mcp`. `/healthz` is always open.
314    #[deprecated(
315        since = "0.13.0",
316        note = "use McpServerConfig::with_auth(); direct field access will become pub(crate) in a future major release"
317    )]
318    pub auth: Option<AuthConfig>,
319    /// Optional RBAC policy. When present and enabled, tool calls are
320    /// checked against the policy after authentication.
321    #[deprecated(
322        since = "0.13.0",
323        note = "use McpServerConfig::with_rbac(); direct field access will become pub(crate) in a future major release"
324    )]
325    pub rbac: Option<Arc<RbacPolicy>>,
326    /// Allowed Origin values for DNS rebinding protection (MCP spec MUST).
327    /// When empty and `public_url` is set, the origin is auto-derived from
328    /// the public URL. When both are empty, only requests with no Origin
329    /// header are accepted.
330    /// Example entries: `"http://localhost:3000"`, `"https://myapp.example.com"`.
331    #[deprecated(
332        since = "0.13.0",
333        note = "use McpServerConfig::with_allowed_origins(); direct field access will become pub(crate) in a future major release"
334    )]
335    pub allowed_origins: Vec<String>,
336    /// Maximum tool invocations per source IP per minute.
337    /// When set, enforced on every `tools/call` request.
338    #[deprecated(
339        since = "0.13.0",
340        note = "use McpServerConfig::with_tool_rate_limit(); direct field access will become pub(crate) in a future major release"
341    )]
342    pub tool_rate_limit: Option<u32>,
343    /// Burst capacity for the tool rate limiter: maximum `tools/call`
344    /// requests admitted back-to-back before the sustained
345    /// [`tool_rate_limit`](Self::tool_rate_limit) rate applies. `None`
346    /// (default) keeps governor's default of burst = rate. Requires
347    /// `tool_rate_limit` to be set; must be greater than zero.
348    #[deprecated(
349        since = "1.12.0",
350        note = "use McpServerConfig::with_tool_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
351    )]
352    pub tool_rate_limit_burst: Option<u32>,
353    /// Maximum requests per source IP per minute for routes merged via
354    /// [`with_extra_router`](Self::with_extra_router). Opt-in: `None`
355    /// (the default) installs no limiter. Startup-only (not
356    /// hot-reloadable via [`ReloadHandle`]).
357    ///
358    /// Keyed by the **direct socket peer** ([`PeerAddr`] semantics — no
359    /// `X-Forwarded-For` interpretation): behind a reverse proxy all
360    /// clients share the proxy's bucket, and IPv6 single-host address
361    /// rotation can evade per-IP keying. Treat this as an abuse speed
362    /// bump for unauthenticated application endpoints, not tenant
363    /// isolation. On limit: HTTP 429 with a plain-text body, matching
364    /// the tool/auth limiters.
365    #[deprecated(
366        since = "1.11.0",
367        note = "use McpServerConfig::with_extra_route_rate_limit(); direct field access will become pub(crate) in a future major release"
368    )]
369    pub extra_route_rate_limit: Option<u32>,
370    /// Burst capacity for the extra-route limiter: maximum requests
371    /// admitted back-to-back before the sustained
372    /// [`extra_route_rate_limit`](Self::extra_route_rate_limit) rate
373    /// applies. `None` (default) keeps governor's default of
374    /// burst = rate. Requires `extra_route_rate_limit` to be set; must
375    /// be greater than zero.
376    #[deprecated(
377        since = "1.12.0",
378        note = "use McpServerConfig::with_extra_route_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
379    )]
380    pub extra_route_rate_limit_burst: Option<u32>,
381    /// Exact-match request paths exempt from the extra-route rate
382    /// limiter (e.g. `/.well-known/oauth-authorization-server`, which
383    /// MCP clients fetch on every connect — behind a shared egress the
384    /// limiter would otherwise 429 discovery). Matching is a **raw
385    /// exact string comparison** against `req.uri().path()`: no globs,
386    /// no prefixes, no normalization — trailing slashes,
387    /// percent-encoding, and dot-segments must match byte-for-byte.
388    /// Fail-closed: any path not listed stays rate-limited (a mismatch
389    /// can only mean "still limited", never "accidentally exempt").
390    /// Requires [`extra_route_rate_limit`](Self::extra_route_rate_limit);
391    /// each entry must be non-empty and start with `/` (validated).
392    /// Startup-only.
393    #[deprecated(
394        since = "1.14.0",
395        note = "use McpServerConfig::with_extra_route_rate_limit_exempt_paths(); direct field access will become pub(crate) in a future major release"
396    )]
397    pub extra_route_rate_limit_exempt_paths: Vec<String>,
398    /// Trusted reverse-proxy networks (CIDRs or bare IPs) for
399    /// **trusted-forwarder mode**. Empty (default) = mode off: every
400    /// limiter keys by the direct socket peer. Nonempty = requests whose
401    /// direct peer is inside one of these networks have their client IP
402    /// resolved from the forwarding header (rightmost-untrusted walk);
403    /// see [`ClientIp`]. Only enable when **all** ingress paths traverse
404    /// the listed proxies. Startup-only.
405    #[deprecated(
406        since = "1.13.0",
407        note = "use McpServerConfig::with_trusted_proxies(); direct field access will become pub(crate) in a future major release"
408    )]
409    pub trusted_proxies: Vec<String>,
410    /// Which forwarding header trusted-forwarder mode reads. `None`
411    /// (default) = `X-Forwarded-For`. Setting this requires
412    /// [`trusted_proxies`](Self::trusted_proxies) to be nonempty
413    /// (validated). Startup-only.
414    #[deprecated(
415        since = "1.13.0",
416        note = "use McpServerConfig::with_forwarded_header(); direct field access will become pub(crate) in a future major release"
417    )]
418    pub forwarded_header: Option<ForwardedHeaderMode>,
419    /// Optional readiness probe for `/readyz`.
420    /// When `None`, `/readyz` mirrors `/healthz` (always OK).
421    #[deprecated(
422        since = "0.13.0",
423        note = "use McpServerConfig::with_readiness_check(); direct field access will become pub(crate) in a future major release"
424    )]
425    pub readiness_check: Option<ReadinessCheck>,
426    /// Maximum request body size in bytes. Default: 1 MiB.
427    /// Protects against oversized payloads causing OOM.
428    #[deprecated(
429        since = "0.13.0",
430        note = "use McpServerConfig::with_max_request_body(); direct field access will become pub(crate) in a future major release"
431    )]
432    pub max_request_body: usize,
433    /// Request processing timeout. Default: 120s.
434    /// Requests exceeding this duration receive 408 Request Timeout.
435    #[deprecated(
436        since = "0.13.0",
437        note = "use McpServerConfig::with_request_timeout(); direct field access will become pub(crate) in a future major release"
438    )]
439    pub request_timeout: Duration,
440    /// Graceful shutdown timeout. Default: 30s.
441    /// After the shutdown signal, in-flight requests have this long to finish.
442    #[deprecated(
443        since = "0.13.0",
444        note = "use McpServerConfig::with_shutdown_timeout(); direct field access will become pub(crate) in a future major release"
445    )]
446    pub shutdown_timeout: Duration,
447    /// Idle timeout for MCP sessions. Sessions with no activity for this
448    /// duration are closed automatically. Default: 20 minutes.
449    #[deprecated(
450        since = "0.13.0",
451        note = "use McpServerConfig::with_session_idle_timeout(); direct field access will become pub(crate) in a future major release"
452    )]
453    pub session_idle_timeout: Duration,
454    /// Interval for SSE keep-alive pings. Prevents proxies and load
455    /// balancers from killing idle connections. Default: 15 seconds.
456    #[deprecated(
457        since = "0.13.0",
458        note = "use McpServerConfig::with_sse_keep_alive(); direct field access will become pub(crate) in a future major release"
459    )]
460    pub sse_keep_alive: Duration,
461    /// Callback invoked once the server is built, delivering a
462    /// [`ReloadHandle`] for hot-reloading auth keys and RBAC policy
463    /// at runtime (e.g. on SIGHUP). Only useful when auth/RBAC is enabled.
464    #[deprecated(
465        since = "0.13.0",
466        note = "use McpServerConfig::with_reload_callback(); direct field access will become pub(crate) in a future major release"
467    )]
468    pub on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
469    /// Additional application-specific routes merged into the top-level
470    /// router.  These routes **bypass** the MCP auth and RBAC middleware,
471    /// so the application is responsible for its own auth on them.
472    /// Handlers can extract [`PeerAddr`] (or
473    /// [`axum::extract::ConnectInfo<SocketAddr>`] for third-party
474    /// middleware compatibility) regardless of whether TLS is enabled.
475    #[deprecated(
476        since = "0.13.0",
477        note = "use McpServerConfig::with_extra_router(); direct field access will become pub(crate) in a future major release"
478    )]
479    pub extra_router: Option<axum::Router>,
480    /// Externally reachable base URL (e.g. `https://mcp.example.com`).
481    /// When set, OAuth metadata endpoints advertise this URL instead of
482    /// the listen address. Required when binding `0.0.0.0` behind a
483    /// reverse proxy or inside a container.
484    #[deprecated(
485        since = "0.13.0",
486        note = "use McpServerConfig::with_public_url(); direct field access will become pub(crate) in a future major release"
487    )]
488    pub public_url: Option<String>,
489    /// Log inbound HTTP request headers at DEBUG level.
490    /// Sensitive values remain redacted.
491    #[deprecated(
492        since = "0.13.0",
493        note = "use McpServerConfig::enable_request_header_logging(); direct field access will become pub(crate) in a future major release"
494    )]
495    pub log_request_headers: bool,
496    /// Expose build metadata (`build_git_sha`, `build_timestamp`,
497    /// `rust_version`) on the unauthenticated `/version` endpoint.
498    /// **Default: `false`** -- only `name`, `version`, and `mcpx_version`
499    /// are served otherwise, so build fingerprints are not leaked to
500    /// anonymous callers. Enable via
501    /// [`McpServerConfig::expose_build_metadata`].
502    pub expose_build_metadata: bool,
503    /// Enable gzip/br response compression on MCP responses.
504    /// Defaults to `false` to preserve existing behaviour.
505    #[deprecated(
506        since = "0.13.0",
507        note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
508    )]
509    pub compression_enabled: bool,
510    /// Minimum response body size (in bytes) before compression kicks in.
511    /// Only used when `compression_enabled` is true. Default: 1024.
512    #[deprecated(
513        since = "0.13.0",
514        note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
515    )]
516    pub compression_min_size: u16,
517    /// Global cap on in-flight HTTP requests across the whole server.
518    /// When `Some`, requests over the cap receive 503 Service Unavailable
519    /// via `tower::load_shed`. Default: `None` (unlimited).
520    #[deprecated(
521        since = "0.13.0",
522        note = "use McpServerConfig::with_max_concurrent_requests(); direct field access will become pub(crate) in a future major release"
523    )]
524    pub max_concurrent_requests: Option<usize>,
525    /// Enable `/admin/*` diagnostic endpoints. Requires `auth` to be
526    /// configured and `enabled`. Default: `false`.
527    #[deprecated(
528        since = "0.13.0",
529        note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
530    )]
531    pub admin_enabled: bool,
532    /// RBAC role required to access admin endpoints. Default: `"admin"`.
533    #[deprecated(
534        since = "0.13.0",
535        note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
536    )]
537    pub admin_role: String,
538    /// Enable Prometheus metrics endpoint on a separate listener.
539    /// Requires the `metrics` crate feature.
540    #[cfg(feature = "metrics")]
541    #[deprecated(
542        since = "0.13.0",
543        note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
544    )]
545    pub metrics_enabled: bool,
546    /// Bind address for the Prometheus metrics listener. Default: `127.0.0.1:9090`.
547    #[cfg(feature = "metrics")]
548    #[deprecated(
549        since = "0.13.0",
550        note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
551    )]
552    pub metrics_bind: String,
553    /// Per-header overrides for the OWASP security headers emitted by
554    /// the global response middleware. See [`SecurityHeadersConfig`]
555    /// for the three-state semantic and validation rules.
556    #[deprecated(
557        since = "1.5.0",
558        note = "use McpServerConfig::with_security_headers(); direct field access will become pub(crate) in a future major release"
559    )]
560    pub security_headers: SecurityHeadersConfig,
561    /// Per-handshake deadline on the TLS accept path. Idle or slow-loris
562    /// connections are dropped once it elapses. Default: 10s.
563    ///
564    /// Startup-only: bound at listener construction, NOT hot-reloadable
565    /// via [`ReloadHandle`]. Ignored unless TLS is configured.
566    #[deprecated(
567        since = "1.9.0",
568        note = "use McpServerConfig::with_tls_handshake_timeout(); direct field access will become pub(crate) in a future major release"
569    )]
570    pub tls_handshake_timeout: Duration,
571    /// Cap on concurrently in-flight TLS handshakes. At saturation the
572    /// acceptor stops pulling new connections from the kernel backlog
573    /// (backpressure) instead of accepting and dropping. Default: 256.
574    ///
575    /// Startup-only: bound at listener construction, NOT hot-reloadable
576    /// via [`ReloadHandle`]. Ignored unless TLS is configured.
577    #[deprecated(
578        since = "1.9.0",
579        note = "use McpServerConfig::with_max_concurrent_tls_handshakes(); direct field access will become pub(crate) in a future major release"
580    )]
581    pub max_concurrent_tls_handshakes: usize,
582}
583
584/// Marker that wraps a value proven to satisfy its validation
585/// contract.
586///
587/// The only way to obtain `Validated<McpServerConfig>` is by calling
588/// [`McpServerConfig::validate`], which is the contract enforced at
589/// the type level by [`serve`] and [`serve_with_listener`]. The
590/// inner field is private, so downstream code cannot bypass
591/// validation by hand-constructing the wrapper.
592///
593/// Use [`Validated::as_inner`] for read-only borrowing. To mutate,
594/// recover the raw value with [`Validated::into_inner`] and
595/// re-validate.
596///
597/// # Example
598///
599/// ```no_run
600/// use rmcp_server_kit::transport::{McpServerConfig, Validated, serve};
601/// use rmcp::handler::server::ServerHandler;
602/// use rmcp::model::{ServerCapabilities, ServerInfo};
603///
604/// #[derive(Clone)]
605/// struct H;
606/// impl ServerHandler for H {
607///     fn get_info(&self) -> ServerInfo {
608///         ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
609///     }
610/// }
611///
612/// # async fn example() -> rmcp_server_kit::Result<()> {
613/// let config: Validated<McpServerConfig> =
614///     McpServerConfig::new("127.0.0.1:8080", "my-server", "0.1.0").validate()?;
615/// serve(config, || H).await
616/// # }
617/// ```
618///
619/// Forgetting `.validate()?` is a compile error:
620///
621/// ```compile_fail
622/// use rmcp_server_kit::transport::{McpServerConfig, serve};
623/// use rmcp::handler::server::ServerHandler;
624/// use rmcp::model::{ServerCapabilities, ServerInfo};
625///
626/// #[derive(Clone)]
627/// struct H;
628/// impl ServerHandler for H {
629///     fn get_info(&self) -> ServerInfo {
630///         ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
631///     }
632/// }
633///
634/// # async fn example() -> rmcp_server_kit::Result<()> {
635/// let config = McpServerConfig::new("127.0.0.1:8080", "my-server", "0.1.0");
636/// // Missing `.validate()?` -> mismatched types: expected
637/// // `Validated<McpServerConfig>`, found `McpServerConfig`.
638/// serve(config, || H).await
639/// # }
640/// ```
641#[allow(
642    missing_debug_implementations,
643    reason = "wraps T which may not implement Debug; manual impl below avoids leaking inner contents into logs"
644)]
645pub struct Validated<T>(T);
646
647impl<T> std::fmt::Debug for Validated<T> {
648    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
649        f.debug_struct("Validated").finish_non_exhaustive()
650    }
651}
652
653impl<T> Validated<T> {
654    /// Borrow the inner value.
655    #[must_use]
656    pub fn as_inner(&self) -> &T {
657        &self.0
658    }
659
660    /// Recover the raw value, discarding the validation proof.
661    ///
662    /// Re-validate before re-using the value with [`serve`] or
663    /// [`serve_with_listener`].
664    #[must_use]
665    pub fn into_inner(self) -> T {
666        self.0
667    }
668}
669
670#[allow(
671    deprecated,
672    reason = "internal builders/validators legitimately read/write the deprecated `pub` fields they were designed to manage"
673)]
674impl McpServerConfig {
675    /// Create a new server configuration with the given bind address,
676    /// server name, and version. All other fields use safe defaults.
677    ///
678    /// Use the chainable `with_*` / `enable_*` builder methods to
679    /// customize. Call [`McpServerConfig::validate`] to obtain a
680    /// [`Validated<McpServerConfig>`] proof token, which is required by
681    /// [`serve`] and [`serve_with_listener`].
682    #[must_use]
683    pub fn new(
684        bind_addr: impl Into<String>,
685        name: impl Into<String>,
686        version: impl Into<String>,
687    ) -> Self {
688        Self {
689            bind_addr: bind_addr.into(),
690            name: name.into(),
691            version: version.into(),
692            tls_cert_path: None,
693            tls_key_path: None,
694            auth: None,
695            rbac: None,
696            allowed_origins: Vec::new(),
697            tool_rate_limit: None,
698            readiness_check: None,
699            max_request_body: 1024 * 1024,
700            request_timeout: Duration::from_mins(2),
701            shutdown_timeout: Duration::from_secs(30),
702            session_idle_timeout: Duration::from_mins(20),
703            sse_keep_alive: Duration::from_secs(15),
704            on_reload_ready: None,
705            extra_router: None,
706            public_url: None,
707            log_request_headers: false,
708            expose_build_metadata: false,
709            compression_enabled: false,
710            compression_min_size: 1024,
711            max_concurrent_requests: None,
712            admin_enabled: false,
713            admin_role: "admin".to_owned(),
714            #[cfg(feature = "metrics")]
715            metrics_enabled: false,
716            #[cfg(feature = "metrics")]
717            metrics_bind: "127.0.0.1:9090".into(),
718            security_headers: SecurityHeadersConfig::default(),
719            tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
720            max_concurrent_tls_handshakes: DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES,
721            extra_route_rate_limit: None,
722            tool_rate_limit_burst: None,
723            extra_route_rate_limit_burst: None,
724            extra_route_rate_limit_exempt_paths: Vec::new(),
725            trusted_proxies: Vec::new(),
726            forwarded_header: None,
727        }
728    }
729
730    // ---------------------------------------------------------------
731    // Builder methods (fluent, consume + return self).
732    //
733    // Each method is `#[must_use]` because dropping the returned
734    // `McpServerConfig` discards the configuration change.
735    // ---------------------------------------------------------------
736
737    /// Attach an authentication configuration. Required for
738    /// [`enable_admin`](Self::enable_admin) and any non-public deployment.
739    #[must_use]
740    pub fn with_auth(mut self, auth: AuthConfig) -> Self {
741        self.auth = Some(auth);
742        self
743    }
744
745    /// Override one or more of the OWASP security headers emitted on
746    /// every response. See [`SecurityHeadersConfig`] for the three-state
747    /// semantic (`None` = default, `Some("")` = omit, `Some(v)` =
748    /// override). Values are validated by [`Self::validate`].
749    #[must_use]
750    pub fn with_security_headers(mut self, headers: SecurityHeadersConfig) -> Self {
751        self.security_headers = headers;
752        self
753    }
754
755    /// Override the bind address (e.g. `127.0.0.1:8080`). Useful when the
756    /// final port is only known after pre-binding an ephemeral listener
757    /// (tests, dynamic-port deployments).
758    #[must_use]
759    pub fn with_bind_addr(mut self, addr: impl Into<String>) -> Self {
760        self.bind_addr = addr.into();
761        self
762    }
763
764    /// Attach an RBAC policy. Tool calls are checked against the policy
765    /// after authentication.
766    #[must_use]
767    pub fn with_rbac(mut self, rbac: Arc<RbacPolicy>) -> Self {
768        self.rbac = Some(rbac);
769        self
770    }
771
772    /// Configure TLS by providing the certificate and private key paths
773    /// (PEM). Both must be readable at startup. Without this call, the
774    /// server runs plain HTTP.
775    #[must_use]
776    pub fn with_tls(mut self, cert_path: impl Into<PathBuf>, key_path: impl Into<PathBuf>) -> Self {
777        self.tls_cert_path = Some(cert_path.into());
778        self.tls_key_path = Some(key_path.into());
779        self
780    }
781
782    /// Set the externally reachable base URL (e.g. `https://mcp.example.com`).
783    /// Required when binding `0.0.0.0` behind a reverse proxy or inside
784    /// a container so OAuth metadata and auto-derived origins resolve correctly.
785    #[must_use]
786    pub fn with_public_url(mut self, url: impl Into<String>) -> Self {
787        self.public_url = Some(url.into());
788        self
789    }
790
791    /// Replace the allowed Origin allow-list (DNS-rebinding protection).
792    /// When empty and [`with_public_url`](Self::with_public_url) is set,
793    /// the origin is auto-derived.
794    #[must_use]
795    pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
796    where
797        I: IntoIterator<Item = S>,
798        S: Into<String>,
799    {
800        self.allowed_origins = origins.into_iter().map(Into::into).collect();
801        self
802    }
803
804    /// Merge an additional axum router at the top level. Routes added
805    /// here **bypass** rmcp-server-kit auth and RBAC; the application is responsible
806    /// for its own protection.
807    ///
808    /// To support that protection (e.g. per-IP rate limiting on
809    /// unauthenticated endpoints), every request served by [`serve`]
810    /// carries the client peer address regardless of whether TLS is
811    /// enabled: extract the framework-owned [`PeerAddr`] in your
812    /// handlers, or rely on [`axum::extract::ConnectInfo<SocketAddr>`]
813    /// for stock third-party middleware (e.g. per-IP rate-limit key
814    /// extractors). Neither extension exists under [`serve_stdio`],
815    /// which has no network peer.
816    #[must_use]
817    pub fn with_extra_router(mut self, router: axum::Router) -> Self {
818        self.extra_router = Some(router);
819        self
820    }
821
822    /// Install an async readiness probe for `/readyz`. Without this call,
823    /// `/readyz` mirrors `/healthz` (always 200 OK).
824    #[must_use]
825    pub fn with_readiness_check(mut self, check: ReadinessCheck) -> Self {
826        self.readiness_check = Some(check);
827        self
828    }
829
830    /// Override the maximum request body (bytes). Must be `> 0`.
831    /// Default: 1 MiB.
832    #[must_use]
833    pub fn with_max_request_body(mut self, bytes: usize) -> Self {
834        self.max_request_body = bytes;
835        self
836    }
837
838    /// Override the per-request processing timeout. Default: 2 minutes.
839    #[must_use]
840    pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
841        self.request_timeout = timeout;
842        self
843    }
844
845    /// Override the graceful shutdown grace period. Default: 30 seconds.
846    #[must_use]
847    pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
848        self.shutdown_timeout = timeout;
849        self
850    }
851
852    /// Override the MCP session idle timeout. Default: 20 minutes.
853    #[must_use]
854    pub fn with_session_idle_timeout(mut self, timeout: Duration) -> Self {
855        self.session_idle_timeout = timeout;
856        self
857    }
858
859    /// Override the SSE keep-alive interval. Default: 15 seconds.
860    #[must_use]
861    pub fn with_sse_keep_alive(mut self, interval: Duration) -> Self {
862        self.sse_keep_alive = interval;
863        self
864    }
865
866    /// Cap the global number of in-flight HTTP requests via
867    /// `tower::load_shed`. Excess requests receive 503 Service Unavailable.
868    /// Default: unlimited.
869    #[must_use]
870    pub fn with_max_concurrent_requests(mut self, limit: usize) -> Self {
871        self.max_concurrent_requests = Some(limit);
872        self
873    }
874
875    /// Override the per-handshake deadline on the TLS accept path.
876    /// Idle or slow-loris connections are dropped once it elapses.
877    /// Default: 10s. Must be greater than zero.
878    ///
879    /// Startup-only: the value is bound at listener construction and is
880    /// NOT hot-reloadable via [`ReloadHandle`]. Has no effect unless TLS
881    /// is configured via [`Self::with_tls`].
882    #[must_use]
883    pub fn with_tls_handshake_timeout(mut self, timeout: Duration) -> Self {
884        self.tls_handshake_timeout = timeout;
885        self
886    }
887
888    /// Override the cap on concurrently in-flight TLS handshakes. At
889    /// saturation the acceptor stops pulling new connections from the
890    /// kernel backlog (backpressure) instead of accepting and dropping.
891    /// Default: 256. Must be greater than zero.
892    ///
893    /// Startup-only: the value is bound at listener construction and is
894    /// NOT hot-reloadable via [`ReloadHandle`]. Has no effect unless TLS
895    /// is configured via [`Self::with_tls`].
896    #[must_use]
897    pub fn with_max_concurrent_tls_handshakes(mut self, limit: usize) -> Self {
898        self.max_concurrent_tls_handshakes = limit;
899        self
900    }
901
902    /// Cap tool invocations per source IP per minute. Enforced on every
903    /// `tools/call` request.
904    #[must_use]
905    pub fn with_tool_rate_limit(mut self, per_minute: u32) -> Self {
906        self.tool_rate_limit = Some(per_minute);
907        self
908    }
909
910    /// Cap requests per source IP per minute on routes merged via
911    /// [`with_extra_router`](Self::with_extra_router) — the natural
912    /// protection for unauthenticated application endpoints (OAuth
913    /// callbacks, registration, …) that bypass auth/RBAC by design.
914    ///
915    /// Must be greater than zero (validated by
916    /// [`validate`](Self::validate)). Startup-only. See the
917    /// `extra_route_rate_limit` field docs for keying semantics and
918    /// caveats (direct peer only, IPv6 rotation, proxy collapse,
919    /// bounded-memory shared-fate under key spray).
920    #[must_use]
921    pub fn with_extra_route_rate_limit(mut self, per_minute: u32) -> Self {
922        self.extra_route_rate_limit = Some(per_minute);
923        self
924    }
925
926    /// Set the burst capacity for the tool rate limiter (bucket size;
927    /// the sustained rate stays [`with_tool_rate_limit`](Self::with_tool_rate_limit)).
928    /// Requires the tool rate limit to be set; must be greater than zero
929    /// (both validated by [`validate`](Self::validate)).
930    #[must_use]
931    pub fn with_tool_rate_limit_burst(mut self, burst: u32) -> Self {
932        self.tool_rate_limit_burst = Some(burst);
933        self
934    }
935
936    /// Set the burst capacity for the extra-route rate limiter (bucket
937    /// size; the sustained rate stays
938    /// [`with_extra_route_rate_limit`](Self::with_extra_route_rate_limit)).
939    /// Requires the extra-route rate limit to be set; must be greater
940    /// than zero (both validated by [`validate`](Self::validate)).
941    #[must_use]
942    pub fn with_extra_route_rate_limit_burst(mut self, burst: u32) -> Self {
943        self.extra_route_rate_limit_burst = Some(burst);
944        self
945    }
946
947    /// Exempt specific request paths from the extra-route rate limiter.
948    ///
949    /// Matching is a **raw exact string comparison** against
950    /// `req.uri().path()` — no globs, no prefixes, no normalization
951    /// (trailing slashes, percent-encoding, and dot-segments must match
952    /// byte-for-byte). The check is fail-closed: anything not listed
953    /// stays rate-limited, so a mismatch can only keep a request
954    /// limited, never accidentally exempt it. The exemption is checked
955    /// before key extraction, so exempt requests consume no limiter
956    /// budget and never appear in deny telemetry.
957    ///
958    /// Typical use: the RFC 8414 authorization-server metadata document
959    /// (`/.well-known/oauth-authorization-server`), fetched by MCP
960    /// clients on every connect.
961    ///
962    /// Requires the extra-route rate limit to be set
963    /// ([`with_extra_route_rate_limit`](Self::with_extra_route_rate_limit));
964    /// each entry must be non-empty and start with `/` (both validated
965    /// by [`validate`](Self::validate)). Startup-only.
966    #[must_use]
967    pub fn with_extra_route_rate_limit_exempt_paths<I, S>(mut self, paths: I) -> Self
968    where
969        I: IntoIterator<Item = S>,
970        S: Into<String>,
971    {
972        self.extra_route_rate_limit_exempt_paths = paths.into_iter().map(Into::into).collect();
973        self
974    }
975
976    /// Enable **trusted-forwarder mode**: requests whose direct peer is
977    /// inside one of these networks (CIDRs or bare IPs) have their
978    /// client IP resolved from the forwarding header via the
979    /// rightmost-untrusted walk; all per-IP rate limiters then key by
980    /// the resolved [`ClientIp`]. Headers from peers outside these
981    /// networks are ignored entirely.
982    ///
983    /// Only enable when **all** ingress paths traverse the listed
984    /// proxies — otherwise direct clients keep their own buckets and
985    /// proxied clients collapse into the proxy's. Entries are validated
986    /// at [`validate`](Self::validate) time. Startup-only.
987    #[must_use]
988    pub fn with_trusted_proxies<I, S>(mut self, proxies: I) -> Self
989    where
990        I: IntoIterator<Item = S>,
991        S: Into<String>,
992    {
993        self.trusted_proxies = proxies.into_iter().map(Into::into).collect();
994        self
995    }
996
997    /// Select which forwarding header trusted-forwarder mode reads
998    /// (default: `X-Forwarded-For`). Requires
999    /// [`with_trusted_proxies`](Self::with_trusted_proxies) to be set
1000    /// (validated).
1001    #[must_use]
1002    pub fn with_forwarded_header(mut self, mode: ForwardedHeaderMode) -> Self {
1003        self.forwarded_header = Some(mode);
1004        self
1005    }
1006
1007    /// Register a callback that receives the [`ReloadHandle`] after the
1008    /// server is built. Use it to wire SIGHUP-style hot reloads of API
1009    /// keys and RBAC policy.
1010    #[must_use]
1011    pub fn with_reload_callback<F>(mut self, callback: F) -> Self
1012    where
1013        F: FnOnce(ReloadHandle) + Send + 'static,
1014    {
1015        self.on_reload_ready = Some(Box::new(callback));
1016        self
1017    }
1018
1019    /// Enable gzip/brotli response compression on MCP responses.
1020    /// `min_size` is the smallest body size (bytes) eligible for
1021    /// compression. Default min size: 1024.
1022    #[must_use]
1023    pub fn enable_compression(mut self, min_size: u16) -> Self {
1024        self.compression_enabled = true;
1025        self.compression_min_size = min_size;
1026        self
1027    }
1028
1029    /// Enable `/admin/*` diagnostic endpoints. Requires
1030    /// [`with_auth`](Self::with_auth) to be set and enabled; otherwise
1031    /// [`validate`](Self::validate) returns an error. `role` is the RBAC
1032    /// role gate (default: `"admin"`).
1033    #[must_use]
1034    pub fn enable_admin(mut self, role: impl Into<String>) -> Self {
1035        self.admin_enabled = true;
1036        self.admin_role = role.into();
1037        self
1038    }
1039
1040    /// Log inbound HTTP request headers at DEBUG level. Sensitive
1041    /// values remain redacted by the logging layer.
1042    #[must_use]
1043    pub fn enable_request_header_logging(mut self) -> Self {
1044        self.log_request_headers = true;
1045        self
1046    }
1047
1048    /// Expose build metadata (`build_git_sha`, `build_timestamp`,
1049    /// `rust_version`) on the unauthenticated `/version` endpoint. Off by
1050    /// default so `/version` reveals only `name`, `version`, and
1051    /// `mcpx_version`.
1052    #[must_use]
1053    pub fn expose_build_metadata(mut self) -> Self {
1054        self.expose_build_metadata = true;
1055        self
1056    }
1057
1058    /// Enable the Prometheus metrics listener on `bind` (e.g.
1059    /// `127.0.0.1:9090`). Requires the `metrics` crate feature.
1060    #[cfg(feature = "metrics")]
1061    #[must_use]
1062    pub fn with_metrics(mut self, bind: impl Into<String>) -> Self {
1063        self.metrics_enabled = true;
1064        self.metrics_bind = bind.into();
1065        self
1066    }
1067
1068    /// Validate the configuration and consume `self`, returning a
1069    /// [`Validated<McpServerConfig>`] proof token required by [`serve`]
1070    /// and [`serve_with_listener`]. This is the only way to construct
1071    /// `Validated<McpServerConfig>`, so the type system guarantees
1072    /// validation has run before the server starts.
1073    ///
1074    /// Checks:
1075    ///
1076    /// 1. `admin_enabled` requires `auth` to be configured and enabled.
1077    /// 2. `tls_cert_path` and `tls_key_path` must both be set or both
1078    ///    be unset.
1079    /// 3. `bind_addr` must parse as a [`SocketAddr`].
1080    /// 4. `public_url`, when set, must start with `http://` or `https://`.
1081    /// 5. Each entry in `allowed_origins` must start with `http://` or
1082    ///    `https://`.
1083    /// 6. `max_request_body` must be greater than zero.
1084    /// 7. When the `oauth` feature is enabled and an [`OAuthConfig`] is
1085    ///    present, all OAuth URL fields (`jwks_uri`, `proxy.authorize_url`,
1086    ///    `proxy.token_url`, `proxy.introspection_url`,
1087    ///    `proxy.revocation_url`, `token_exchange.token_url`) must parse
1088    ///    and use the `https` scheme. Set
1089    ///    [`OAuthConfig::allow_http_oauth_urls`] to permit `http://`
1090    ///    targets (strongly discouraged in production - see the field-level
1091    ///    docs for the threat model).
1092    ///
1093    /// [`OAuthConfig`]: crate::oauth::OAuthConfig
1094    /// [`OAuthConfig::allow_http_oauth_urls`]: crate::oauth::OAuthConfig::allow_http_oauth_urls
1095    ///
1096    /// # Errors
1097    ///
1098    /// Returns [`McpxError::Config`] with a human-readable message on
1099    /// the first validation failure.
1100    pub fn validate(self) -> Result<Validated<Self>, McpxError> {
1101        self.check()?;
1102        Ok(Validated(self))
1103    }
1104
1105    /// Validate the burst knobs: every burst must be greater than zero
1106    /// when set, and the two top-level bursts require their base limiter
1107    /// to be configured. The auth bursts
1108    /// (`RateLimitConfig::{burst, pre_auth_burst}`) have no orphan rule:
1109    /// their base rates always resolve (`max_attempts_per_minute` is
1110    /// mandatory; the pre-auth base derives from it when unset).
1111    fn check_burst_knobs(&self) -> Result<(), McpxError> {
1112        if self.tool_rate_limit_burst == Some(0) {
1113            return Err(McpxError::Config(
1114                "tool_rate_limit_burst must be greater than zero".into(),
1115            ));
1116        }
1117        if self.extra_route_rate_limit_burst == Some(0) {
1118            return Err(McpxError::Config(
1119                "extra_route_rate_limit_burst must be greater than zero".into(),
1120            ));
1121        }
1122        if self.tool_rate_limit_burst.is_some() && self.tool_rate_limit.is_none() {
1123            return Err(McpxError::Config(
1124                "tool_rate_limit_burst requires tool_rate_limit to be set".into(),
1125            ));
1126        }
1127        if self.extra_route_rate_limit_burst.is_some() && self.extra_route_rate_limit.is_none() {
1128            return Err(McpxError::Config(
1129                "extra_route_rate_limit_burst requires extra_route_rate_limit to be set".into(),
1130            ));
1131        }
1132        if !self.extra_route_rate_limit_exempt_paths.is_empty()
1133            && self.extra_route_rate_limit.is_none()
1134        {
1135            return Err(McpxError::Config(
1136                "extra_route_rate_limit_exempt_paths requires extra_route_rate_limit to be set"
1137                    .into(),
1138            ));
1139        }
1140        for path in &self.extra_route_rate_limit_exempt_paths {
1141            if path.is_empty() || !path.starts_with('/') {
1142                return Err(McpxError::Config(format!(
1143                    "extra_route_rate_limit_exempt_paths entries must be non-empty and start with '/': {path:?}"
1144                )));
1145            }
1146        }
1147        if let Some(rl) = self.auth.as_ref().and_then(|a| a.rate_limit.as_ref()) {
1148            if rl.burst == Some(0) {
1149                return Err(McpxError::Config(
1150                    "auth rate_limit.burst must be greater than zero".into(),
1151                ));
1152            }
1153            if rl.pre_auth_burst == Some(0) {
1154                return Err(McpxError::Config(
1155                    "auth rate_limit.pre_auth_burst must be greater than zero".into(),
1156                ));
1157            }
1158        }
1159        Ok(())
1160    }
1161
1162    /// Validate the trusted-forwarder knobs: every `trusted_proxies`
1163    /// entry must parse as a CIDR (`ipnet::IpNet`) or a bare IP
1164    /// (normalized to a host network), and `forwarded_header` requires a
1165    /// nonempty proxy list (fail-fast over a silent no-op).
1166    fn check_trusted_forwarder(&self) -> Result<(), McpxError> {
1167        for entry in &self.trusted_proxies {
1168            validate_trusted_proxy_entry(entry).map_err(McpxError::Config)?;
1169        }
1170        if self.forwarded_header.is_some() && self.trusted_proxies.is_empty() {
1171            return Err(McpxError::Config(
1172                "forwarded_header requires trusted_proxies to be nonempty".into(),
1173            ));
1174        }
1175        Ok(())
1176    }
1177
1178    /// Run the validation checks without consuming `self`. Used by
1179    /// internal call sites (e.g. tests) that need to inspect a config
1180    /// without taking ownership.
1181    fn check(&self) -> Result<(), McpxError> {
1182        // 1. admin <-> auth dependency. Mirrors the runtime check in
1183        //    `build_app_router`: admin endpoints require an auth state,
1184        //    which is built only when `auth` is `Some` *and* `enabled`.
1185        if self.admin_enabled {
1186            let auth_enabled = self.auth.as_ref().is_some_and(|a| a.enabled);
1187            if !auth_enabled {
1188                return Err(McpxError::Config(
1189                    "admin_enabled=true requires auth to be configured and enabled".into(),
1190                ));
1191            }
1192        }
1193
1194        // 2. TLS cert / key must be paired
1195        match (&self.tls_cert_path, &self.tls_key_path) {
1196            (Some(_), None) => {
1197                return Err(McpxError::Config(
1198                    "tls_cert_path is set but tls_key_path is missing".into(),
1199                ));
1200            }
1201            (None, Some(_)) => {
1202                return Err(McpxError::Config(
1203                    "tls_key_path is set but tls_cert_path is missing".into(),
1204                ));
1205            }
1206            _ => {}
1207        }
1208
1209        // 3. bind_addr parses
1210        if self.bind_addr.parse::<SocketAddr>().is_err() {
1211            return Err(McpxError::Config(format!(
1212                "bind_addr {:?} is not a valid socket address (expected e.g. 127.0.0.1:8080)",
1213                self.bind_addr
1214            )));
1215        }
1216
1217        // 4. public_url scheme
1218        if let Some(ref url) = self.public_url
1219            && !(url.starts_with("http://") || url.starts_with("https://"))
1220        {
1221            return Err(McpxError::Config(format!(
1222                "public_url {url:?} must start with http:// or https://"
1223            )));
1224        }
1225
1226        // 5. allowed_origins scheme
1227        for origin in &self.allowed_origins {
1228            if !(origin.starts_with("http://") || origin.starts_with("https://")) {
1229                return Err(McpxError::Config(format!(
1230                    "allowed_origins entry {origin:?} must start with http:// or https://"
1231                )));
1232            }
1233        }
1234
1235        // 6. max_request_body > 0
1236        if self.max_request_body == 0 {
1237            return Err(McpxError::Config(
1238                "max_request_body must be greater than zero".into(),
1239            ));
1240        }
1241
1242        // 6b. extra_route_rate_limit, when set, must be > 0. Unlike the
1243        // legacy tool_rate_limit (which clamps 0 to its default at
1244        // construction), new knobs fail fast on nonsensical values.
1245        if self.extra_route_rate_limit == Some(0) {
1246            return Err(McpxError::Config(
1247                "extra_route_rate_limit must be greater than zero".into(),
1248            ));
1249        }
1250
1251        // 6c. Burst knobs (extracted helper).
1252        self.check_burst_knobs()?;
1253
1254        // 6d. Trusted-forwarder knobs (extracted helper).
1255        self.check_trusted_forwarder()?;
1256
1257        // 7. OAuth URL fields enforce HTTPS (unless `allow_http_oauth_urls`)
1258        #[cfg(feature = "oauth")]
1259        if let Some(auth_cfg) = &self.auth
1260            && let Some(oauth_cfg) = &auth_cfg.oauth
1261        {
1262            oauth_cfg.validate()?;
1263        }
1264
1265        // 8. Security-header overrides parse as valid HTTP header values,
1266        //    and HSTS does not smuggle in a `preload` directive.
1267        validate_security_headers(&self.security_headers)?;
1268
1269        // 9. max_concurrent_requests must be > 0 when set. Zero would
1270        //    deadlock the global concurrency limiter and reject every
1271        //    request. Mirrors the TOML-side check in `src/config.rs`.
1272        if self.max_concurrent_requests == Some(0) {
1273            return Err(McpxError::Config(
1274                "max_concurrent_requests must be greater than zero when set".into(),
1275            ));
1276        }
1277
1278        // 10. Auth rate-limit `max_tracked_keys` must be > 0. A zero cap
1279        //     would force `BoundedKeyedLimiter` to evict on every insert
1280        //     and effectively disable rate limiting.
1281        if let Some(auth_cfg) = &self.auth
1282            && let Some(rl) = &auth_cfg.rate_limit
1283            && rl.max_tracked_keys == 0
1284        {
1285            return Err(McpxError::Config(
1286                "auth.rate_limit.max_tracked_keys must be greater than zero".into(),
1287            ));
1288        }
1289
1290        // 11. tls_handshake_timeout must be > 0. A zero deadline would
1291        //     reap every handshake before it could complete, rejecting
1292        //     all TLS connections. Mirrors the TOML-side check in
1293        //     `src/config.rs`.
1294        if self.tls_handshake_timeout == Duration::ZERO {
1295            return Err(McpxError::Config(
1296                "tls_handshake_timeout must be greater than zero".into(),
1297            ));
1298        }
1299
1300        // 12. max_concurrent_tls_handshakes must be > 0. A zero-permit
1301        //     semaphore would never admit a handshake, deadlocking the
1302        //     TLS accept path. Mirrors the TOML-side check in
1303        //     `src/config.rs`.
1304        if self.max_concurrent_tls_handshakes == 0 {
1305            return Err(McpxError::Config(
1306                "max_concurrent_tls_handshakes must be greater than zero".into(),
1307            ));
1308        }
1309
1310        Ok(())
1311    }
1312}
1313
1314/// Handle for hot-reloading server configuration without restart.
1315///
1316/// Obtained via [`McpServerConfig::on_reload_ready`].
1317/// All swap operations are lock-free and wait-free -- in-flight requests
1318/// finish with the old values while new requests see the update immediately.
1319#[allow(
1320    missing_debug_implementations,
1321    reason = "contains Arc<AuthState> with non-Debug fields"
1322)]
1323pub struct ReloadHandle {
1324    auth: Option<Arc<AuthState>>,
1325    rbac: Option<Arc<ArcSwap<RbacPolicy>>>,
1326    crl_set: Option<Arc<CrlSet>>,
1327}
1328
1329impl ReloadHandle {
1330    /// Atomically replace the API key list used by the auth middleware.
1331    pub fn reload_auth_keys(&self, keys: Vec<crate::auth::ApiKeyEntry>) {
1332        if let Some(ref auth) = self.auth {
1333            auth.reload_keys(keys);
1334        }
1335    }
1336
1337    /// Atomically replace the RBAC policy used by the RBAC middleware.
1338    pub fn reload_rbac(&self, policy: RbacPolicy) {
1339        if let Some(ref rbac) = self.rbac {
1340            rbac.store(Arc::new(policy));
1341            tracing::info!("RBAC policy reloaded");
1342        }
1343    }
1344
1345    /// Force an immediate refresh of all cached mTLS CRLs.
1346    ///
1347    /// # Errors
1348    ///
1349    /// Returns an error if CRL refresh is unavailable or verifier rebuild fails.
1350    pub async fn refresh_crls(&self) -> Result<(), McpxError> {
1351        let Some(ref crl_set) = self.crl_set else {
1352            return Err(McpxError::Config(
1353                "CRL refresh requested but mTLS CRL support is not configured".into(),
1354            ));
1355        };
1356
1357        crl_set.force_refresh().await
1358    }
1359}
1360
1361/// Generic MCP HTTP server.
1362///
1363/// Wraps an axum server with `/healthz` and `/mcp` endpoints.
1364/// When `tls_cert_path` and `tls_key_path` are both set, the server binds
1365/// with TLS (rustls). Optionally supports mTLS client certificate auth.
1366///
1367/// # Errors
1368///
1369/// Returns an error if the TCP listener cannot bind, TLS config is invalid,
1370/// or the server fails.
1371// NOTE: cognitive complexity reduced from 111/25 to 83/25 by
1372// extracting `run_server` (serve-loop tail) and `install_oauth_proxy_routes`.
1373// Remaining flow is a linear router builder: middleware layering, feature-
1374// gated auth/RBAC wiring, and PRM/metrics installation. Further extraction
1375// would require threading many `&mut Router` helpers and hurt readability
1376// of the layer order (which is security-relevant and must stay visible).
1377#[allow(
1378    clippy::too_many_lines,
1379    clippy::cognitive_complexity,
1380    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"
1381)]
1382/// Internal bundle of values produced by [`build_app_router`] and
1383/// consumed by [`serve`] / [`serve_with_listener`] when driving the
1384/// HTTP listener.
1385struct AppRunParams {
1386    /// TLS cert/key paths when TLS is configured.
1387    tls_paths: Option<(PathBuf, PathBuf)>,
1388    /// Per-handshake deadline on the TLS accept path.
1389    tls_handshake_timeout: Duration,
1390    /// Cap on concurrently in-flight TLS handshakes.
1391    max_concurrent_tls_handshakes: usize,
1392    /// mTLS configuration when mutual-TLS auth is enabled.
1393    mtls_config: Option<MtlsConfig>,
1394    /// Graceful shutdown drain window.
1395    shutdown_timeout: Duration,
1396    /// Shared auth state used by hot-reload callbacks.
1397    auth_state: Option<Arc<AuthState>>,
1398    /// Hot-reloadable RBAC state used by reload callbacks.
1399    rbac_swap: Arc<ArcSwap<RbacPolicy>>,
1400    /// Optional callback that receives the final [`ReloadHandle`].
1401    on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
1402    /// Server-internal lifecycle cancellation token. Cancelled by
1403    /// [`run_server`] once the shutdown trigger fires, stopping the metrics
1404    /// listener, the CRL refresher, and any external shutdown wiring.
1405    ct: CancellationToken,
1406    /// Cancellation token handed to the MCP service, kept SEPARATE from
1407    /// [`Self::ct`].
1408    ///
1409    /// Cancelling it terminates in-flight MCP sessions and SSE streams, so it
1410    /// must not fire at the *start* of the grace period — that truncates
1411    /// responses a normal SIGTERM rollout is supposed to let finish. It is
1412    /// cancelled only after axum has drained, or when the force-exit timer
1413    /// wins. The force-exit path must still cancel it, or a stuck stream turns
1414    /// a truncation bug into a shutdown hang.
1415    session_ct: CancellationToken,
1416    /// `"http"` or `"https"` -- used only for boot-time logging.
1417    scheme: &'static str,
1418    /// Server name -- used only for boot-time logging.
1419    name: String,
1420}
1421
1422/// Build the full application axum [`axum::Router`] (MCP route +
1423/// middleware stack + admin + OAuth + health endpoints + security
1424/// headers + CORS + compression + concurrency limit + origin check)
1425/// and the [`AppRunParams`] needed to drive it.
1426///
1427/// This is the shared core of [`serve`] and [`serve_with_listener`].
1428/// It performs *no* network I/O: callers are responsible for binding
1429/// (or accepting a pre-bound) [`TcpListener`] and invoking
1430/// [`run_server`].
1431#[allow(
1432    clippy::cognitive_complexity,
1433    reason = "router assembly is intrinsically sequential; splitting harms readability"
1434)]
1435#[allow(
1436    deprecated,
1437    reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1438)]
1439fn build_app_router<H, F>(
1440    mut config: McpServerConfig,
1441    handler_factory: F,
1442) -> anyhow::Result<(axum::Router, AppRunParams)>
1443where
1444    H: ServerHandler + 'static,
1445    F: Fn() -> H + Send + Sync + Clone + 'static,
1446{
1447    let ct = CancellationToken::new();
1448    let session_ct = CancellationToken::new();
1449
1450    let allowed_hosts = derive_allowed_hosts(&config.bind_addr, config.public_url.as_deref());
1451    tracing::info!(allowed_hosts = %allowed_hosts.join(", "), "configured Streamable HTTP allowed hosts");
1452
1453    let mcp_service = StreamableHttpService::new(
1454        move || Ok(handler_factory()),
1455        {
1456            let mut mgr = LocalSessionManager::default();
1457            mgr.session_config.keep_alive = Some(config.session_idle_timeout);
1458            mgr.into()
1459        },
1460        StreamableHttpServerConfig::default()
1461            .with_allowed_hosts(allowed_hosts)
1462            .with_sse_keep_alive(Some(config.sse_keep_alive))
1463            .with_cancellation_token(session_ct.clone()),
1464    );
1465
1466    // Build the MCP route, optionally wrapped with auth and RBAC middleware.
1467    let mut mcp_router = axum::Router::new().nest_service("/mcp", mcp_service);
1468
1469    // Build auth state eagerly when auth is configured so we can wire both
1470    // the auth middleware *and* the optional admin router against the same
1471    // state. The middleware itself is installed further down in layer order.
1472    let auth_state: Option<Arc<AuthState>> = match config.auth {
1473        Some(ref auth_config) if auth_config.enabled => {
1474            let rate_limiter = auth_config.rate_limit.as_ref().map(build_rate_limiter);
1475            let pre_auth_limiter = auth_config
1476                .rate_limit
1477                .as_ref()
1478                .map(crate::auth::build_pre_auth_limiter);
1479
1480            #[cfg(feature = "oauth")]
1481            let jwks_cache = auth_config
1482                .oauth
1483                .as_ref()
1484                .map(|c| crate::oauth::JwksCache::new(c).map(Arc::new))
1485                .transpose()
1486                .map_err(|e| std::io::Error::other(format!("JWKS HTTP client: {e}")))?;
1487
1488            Some(Arc::new(AuthState {
1489                api_keys: ArcSwap::new(Arc::new(auth_config.api_keys.clone())),
1490                rate_limiter,
1491                pre_auth_limiter,
1492                #[cfg(feature = "oauth")]
1493                jwks_cache,
1494                seen_identities: crate::auth::SeenIdentitySet::new(),
1495                counters: crate::auth::AuthCounters::default(),
1496            }))
1497        }
1498        _ => None,
1499    };
1500
1501    // Build the RBAC policy swap early so the admin router and the later
1502    // RBAC middleware layer share the same hot-reloadable state.
1503    let rbac_swap = Arc::new(ArcSwap::new(
1504        config
1505            .rbac
1506            .clone()
1507            .unwrap_or_else(|| Arc::new(RbacPolicy::disabled())),
1508    ));
1509
1510    // Optional /admin/* diagnostic routes. Merged BEFORE the
1511    // body-limit/timeout/RBAC/origin/auth layers so all of them apply.
1512    if config.admin_enabled {
1513        let Some(ref auth_state_ref) = auth_state else {
1514            return Err(anyhow::anyhow!(
1515                "admin_enabled=true requires auth to be configured and enabled"
1516            ));
1517        };
1518        let admin_state = crate::admin::AdminState {
1519            started_at: std::time::Instant::now(),
1520            name: config.name.clone(),
1521            version: config.version.clone(),
1522            auth: Some(Arc::clone(auth_state_ref)),
1523            rbac: Arc::clone(&rbac_swap),
1524        };
1525        let admin_cfg = crate::admin::AdminConfig {
1526            role: config.admin_role.clone(),
1527        };
1528        mcp_router = mcp_router.merge(crate::admin::admin_router(admin_state, &admin_cfg));
1529        tracing::info!(role = %config.admin_role, "/admin/* endpoints enabled");
1530    }
1531
1532    // ----- Middleware order (CRITICAL: read carefully) ------------------
1533    //
1534    // axum/tower applies layers **bottom-up** at runtime: the LAST layer
1535    // added is the OUTERMOST (runs first on a request). To achieve a
1536    // request-time flow of:
1537    //
1538    //   body-limit -> timeout -> auth -> rbac -> handler
1539    //
1540    // we add layers in the REVERSE order:
1541    //
1542    //   1. RBAC               (innermost, runs last before handler)
1543    //   2. auth               (parses identity, sets extension for RBAC)
1544    //   3. timeout            (bounds total request time)
1545    //   4. body-limit         (outermost on /mcp; caps payload before
1546    //                          anything else reads/buffers it)
1547    //
1548    // Origin validation is installed on the OUTER router (after the
1549    // /mcp router is merged in), so it also protects /healthz, /readyz,
1550    // /version, and any OAuth proxy endpoints.
1551    //
1552    // Rationale:
1553    // - Body-limit must be outermost on /mcp so RBAC (which reads the
1554    //   JSON-RPC body) cannot be DoS'd by a 100MB payload.
1555    // - Auth must run before RBAC because RBAC consumes
1556    //   `req.extensions().get::<AuthIdentity>()` to enforce per-role
1557    //   policy.
1558    // - Origin runs before auth so we reject cross-origin requests
1559    //   without spending Argon2 cycles on unauthenticated callers.
1560
1561    // [1] RBAC + tool rate-limit layer (innermost; closest to handler).
1562    // Always installed: even when RBAC is disabled, tool rate limiting may
1563    // be active (MCP spec: servers MUST rate limit tool invocations).
1564    {
1565        let tool_limiter: Option<Arc<ToolRateLimiter>> = config
1566            .tool_rate_limit
1567            .map(|per_minute| build_tool_rate_limiter(per_minute, config.tool_rate_limit_burst));
1568
1569        if rbac_swap.load().is_enabled() {
1570            tracing::info!("RBAC enforcement enabled on /mcp");
1571        }
1572        if let Some(limit) = config.tool_rate_limit {
1573            tracing::info!(limit, "tool rate limiting enabled (calls/min per IP)");
1574        }
1575
1576        let rbac_for_mw = Arc::clone(&rbac_swap);
1577        mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1578            let p = rbac_for_mw.load_full();
1579            let tl = tool_limiter.clone();
1580            rbac_middleware(p, tl, req, next)
1581        }));
1582    }
1583
1584    // [2] Auth layer (runs before RBAC so AuthIdentity is in extensions).
1585    if let Some(ref auth_config) = config.auth
1586        && auth_config.enabled
1587    {
1588        let Some(ref state) = auth_state else {
1589            return Err(anyhow::anyhow!("auth state missing despite enabled config"));
1590        };
1591
1592        let methods: Vec<&str> = [
1593            auth_config.mtls.is_some().then_some("mTLS"),
1594            (!auth_config.api_keys.is_empty()).then_some("bearer"),
1595            #[cfg(feature = "oauth")]
1596            auth_config.oauth.is_some().then_some("oauth-jwt"),
1597        ]
1598        .into_iter()
1599        .flatten()
1600        .collect();
1601
1602        tracing::info!(
1603            methods = %methods.join(", "),
1604            api_keys = auth_config.api_keys.len(),
1605            "auth enabled on /mcp"
1606        );
1607
1608        let state_for_mw = Arc::clone(state);
1609        mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1610            let s = Arc::clone(&state_for_mw);
1611            auth_middleware(s, req, next)
1612        }));
1613    }
1614
1615    // [3] Request timeout (returns 408 on expiry). Bounds total request
1616    // duration including auth + handler.
1617    mcp_router = mcp_router.layer(tower_http::timeout::TimeoutLayer::with_status_code(
1618        axum::http::StatusCode::REQUEST_TIMEOUT,
1619        config.request_timeout,
1620    ));
1621
1622    // [4] Request body size limit (OUTERMOST on /mcp). Prevents OOM /
1623    // DoS from oversized payloads BEFORE any inner layer (auth, RBAC)
1624    // attempts to buffer or parse the body.
1625    mcp_router = mcp_router.layer(tower_http::limit::RequestBodyLimitLayer::new(
1626        config.max_request_body,
1627    ));
1628
1629    // Compute the effective allowed-origins list for the outer
1630    // origin-check layer (installed on the merged router below). When
1631    // `allowed_origins` is empty but `public_url` is set, auto-derive
1632    // the origin from the public URL so MCP clients (e.g. Claude Code)
1633    // that send `Origin: <server-url>` are accepted without explicit
1634    // config.
1635    let mut effective_origins = config.allowed_origins.clone();
1636    if effective_origins.is_empty()
1637        && let Some(ref url) = config.public_url
1638    {
1639        // Origin = scheme + "://" + host (+ ":" + port if non-default).
1640        // Strip any path/query from the public URL. Offsets come from
1641        // `find`, so they are char-boundary-aligned; `get(..)` keeps that
1642        // machine-checked (a violation degrades to an empty slice).
1643        if let Some(scheme_end) = url.find("://") {
1644            let scheme_with_sep = url.get(..scheme_end + 3).unwrap_or_default();
1645            let after_scheme = url.get(scheme_end + 3..).unwrap_or_default();
1646            let host_end = after_scheme.find('/').unwrap_or(after_scheme.len());
1647            let host = after_scheme.get(..host_end).unwrap_or_default();
1648            let origin = format!("{scheme_with_sep}{host}");
1649            tracing::info!(
1650                %origin,
1651                "auto-derived allowed origin from public_url"
1652            );
1653            effective_origins.push(origin);
1654        }
1655    }
1656    let allowed_origins: Arc<[String]> = Arc::from(effective_origins);
1657    let cors_origins = Arc::clone(&allowed_origins);
1658    let log_request_headers = config.log_request_headers;
1659
1660    let readyz_route = if let Some(check) = config.readiness_check.take() {
1661        axum::routing::get(move || readyz(Arc::clone(&check)))
1662    } else {
1663        axum::routing::get(healthz)
1664    };
1665
1666    #[allow(unused_mut)] // mut needed when oauth feature adds PRM route
1667    let mut router = axum::Router::new()
1668        .route("/healthz", axum::routing::get(healthz))
1669        .route("/readyz", readyz_route)
1670        .route(
1671            "/version",
1672            axum::routing::get({
1673                // Pre-serialize the version payload once at router-build
1674                // time. The handler then serves a cheap `Arc::clone` of the
1675                // immutable bytes per request, avoiding `serde_json::Value`
1676                // allocation + serialization on every `/version` hit.
1677                let payload_bytes: Arc<[u8]> = serialize_version_payload(
1678                    &config.name,
1679                    &config.version,
1680                    config.expose_build_metadata,
1681                );
1682                move || {
1683                    let p = Arc::clone(&payload_bytes);
1684                    async move {
1685                        (
1686                            [(axum::http::header::CONTENT_TYPE, "application/json")],
1687                            p.to_vec(),
1688                        )
1689                    }
1690                }
1691            }),
1692        )
1693        .merge(mcp_router);
1694
1695    // Merge application-specific routes (bypass MCP auth/RBAC middleware).
1696    // When configured, wrap them — and only them — in the per-IP rate
1697    // limiter BEFORE merging: axum layers wrap only the routes already
1698    // present on the sub-router, so the limiter can never leak onto
1699    // `/mcp`, health, admin, or OAuth endpoints, while top-level layers
1700    // (origin check, peer-address normalization, ...) still run first.
1701    if let Some(extra) = config.extra_router.take() {
1702        let extra = match config.extra_route_rate_limit {
1703            Some(per_minute) => {
1704                let limiter =
1705                    build_extra_route_rate_limiter(per_minute, config.extra_route_rate_limit_burst);
1706                let exempt: Arc<std::collections::HashSet<String>> = Arc::new(
1707                    config
1708                        .extra_route_rate_limit_exempt_paths
1709                        .iter()
1710                        .cloned()
1711                        .collect(),
1712                );
1713                tracing::info!(
1714                    per_minute,
1715                    exempt_paths = exempt.len(),
1716                    "extra-route per-IP rate limit enabled"
1717                );
1718                extra.layer(axum::middleware::from_fn(move |req, next| {
1719                    let l = Arc::clone(&limiter);
1720                    let e = Arc::clone(&exempt);
1721                    extra_route_rate_limit_middleware(l, e, req, next)
1722                }))
1723            }
1724            None => extra,
1725        };
1726        router = router.merge(extra);
1727    }
1728
1729    // RFC 9728: Protected Resource Metadata endpoint.
1730    // When OAuth is configured, serve full metadata with authorization_servers.
1731    // Otherwise, serve a minimal document with just the resource URL and no
1732    // authorization_servers -- this tells MCP clients (e.g. Claude Code SDK)
1733    // that the server exists but does NOT require OAuth authentication,
1734    // preventing them from gating the connection behind a broken auth flow.
1735    let server_url = if let Some(ref url) = config.public_url {
1736        url.trim_end_matches('/').to_owned()
1737    } else {
1738        let prm_scheme = if config.tls_cert_path.is_some() {
1739            "https"
1740        } else {
1741            "http"
1742        };
1743        format!("{prm_scheme}://{}", config.bind_addr)
1744    };
1745    let resource_url = format!("{server_url}/mcp");
1746
1747    #[cfg(feature = "oauth")]
1748    let prm_metadata = if let Some(ref auth_config) = config.auth
1749        && let Some(ref oauth_config) = auth_config.oauth
1750    {
1751        crate::oauth::protected_resource_metadata(&resource_url, &server_url, oauth_config)
1752    } else {
1753        serde_json::json!({ "resource": resource_url })
1754    };
1755    #[cfg(not(feature = "oauth"))]
1756    let prm_metadata = serde_json::json!({ "resource": resource_url });
1757
1758    router = router.route(
1759        "/.well-known/oauth-protected-resource",
1760        axum::routing::get(move || {
1761            let m = prm_metadata.clone();
1762            async move { axum::Json(m) }
1763        }),
1764    );
1765
1766    // OAuth 2.1 proxy endpoints: when an OAuth proxy is configured, expose
1767    // /authorize, /token, /register, and authorization server metadata so
1768    // MCP clients can perform Authorization Code + PKCE against the upstream
1769    // IdP (e.g. Keycloak) transparently.
1770    #[cfg(feature = "oauth")]
1771    if let Some(ref auth_config) = config.auth
1772        && let Some(ref oauth_config) = auth_config.oauth
1773        && oauth_config.proxy.is_some()
1774    {
1775        router = install_oauth_proxy_routes(
1776            router,
1777            &server_url,
1778            oauth_config,
1779            auth_state.as_ref(),
1780            config.max_request_body,
1781            &config.admin_role,
1782        )?;
1783    }
1784
1785    // OWASP security response headers are installed LAST (after the origin
1786    // layer, below) so they form the OUTERMOST response layer and therefore
1787    // also decorate origin-403, CORS-preflight, overload-503, and 404-fallback
1788    // responses. See the `security_headers_middleware` install site below.
1789
1790    // CORS preflight layer (required for browser-based MCP clients).
1791    // Uses the same effective origins as the origin check middleware
1792    // (including auto-derived origin from public_url).
1793    if !cors_origins.is_empty() {
1794        let cors = tower_http::cors::CorsLayer::new()
1795            .allow_origin(
1796                cors_origins
1797                    .iter()
1798                    .filter_map(|o| o.parse::<axum::http::HeaderValue>().ok())
1799                    .collect::<Vec<_>>(),
1800            )
1801            .allow_methods([
1802                axum::http::Method::GET,
1803                axum::http::Method::POST,
1804                axum::http::Method::OPTIONS,
1805            ])
1806            .allow_headers([
1807                axum::http::header::CONTENT_TYPE,
1808                axum::http::header::AUTHORIZATION,
1809            ]);
1810        router = router.layer(cors);
1811    }
1812
1813    // Optional response compression (gzip + brotli). Skips small bodies
1814    // to avoid overhead. Applied after CORS so preflight responses remain
1815    // uncompressed.
1816    if config.compression_enabled {
1817        use tower_http::compression::Predicate as _;
1818        let predicate = tower_http::compression::DefaultPredicate::new().and(
1819            tower_http::compression::predicate::SizeAbove::new(u64::from(
1820                config.compression_min_size,
1821            )),
1822        );
1823        router = router.layer(
1824            tower_http::compression::CompressionLayer::new()
1825                .gzip(true)
1826                .br(true)
1827                .compress_when(predicate),
1828        );
1829        tracing::info!(
1830            min_size = config.compression_min_size,
1831            "response compression enabled (gzip, br)"
1832        );
1833    }
1834
1835    // Optional global concurrency cap. `load_shed` converts the
1836    // `ConcurrencyLimit` back-pressure error into 503 instead of hanging.
1837    if let Some(max) = config.max_concurrent_requests {
1838        let overload_handler = tower::ServiceBuilder::new()
1839            .layer(axum::error_handling::HandleErrorLayer::new(
1840                |_err: tower::BoxError| async {
1841                    (
1842                        axum::http::StatusCode::SERVICE_UNAVAILABLE,
1843                        axum::Json(serde_json::json!({
1844                            "error": "overloaded",
1845                            "error_description": "server is at capacity, retry later"
1846                        })),
1847                    )
1848                },
1849            ))
1850            .layer(tower::load_shed::LoadShedLayer::new())
1851            .layer(tower::limit::ConcurrencyLimitLayer::new(max));
1852        router = router.layer(overload_handler);
1853        tracing::info!(max, "global concurrency limit enabled");
1854    }
1855
1856    // JSON fallback for unmatched routes. Without this, axum returns
1857    // an empty-body 404 that breaks MCP clients (e.g. Claude Code SDK)
1858    // when they probe OAuth endpoints like /authorize or /token.
1859    router = router.fallback(|| async {
1860        (
1861            axum::http::StatusCode::NOT_FOUND,
1862            axum::Json(serde_json::json!({
1863                "error": "not_found",
1864                "error_description": "The requested endpoint does not exist"
1865            })),
1866        )
1867    });
1868
1869    // Prometheus metrics: recording middleware + separate listener.
1870    #[cfg(feature = "metrics")]
1871    if config.metrics_enabled {
1872        let metrics = Arc::new(
1873            crate::metrics::McpMetrics::new().map_err(|e| anyhow::anyhow!("metrics init: {e}"))?,
1874        );
1875        let m = Arc::clone(&metrics);
1876        router = router.layer(axum::middleware::from_fn(
1877            move |req: Request<Body>, next: Next| {
1878                let m = Arc::clone(&m);
1879                metrics_middleware(m, req, next)
1880            },
1881        ));
1882        let metrics_bind = config.metrics_bind.clone();
1883        let metrics_shutdown = ct.clone();
1884        tokio::spawn(async move {
1885            if let Err(e) =
1886                crate::metrics::serve_metrics(metrics_bind, metrics, metrics_shutdown).await
1887            {
1888                tracing::error!("metrics listener failed: {e}");
1889            }
1890        });
1891    }
1892
1893    // Peer-address normalization. Mirrors the TLS branch's peer address
1894    // into `ConnectInfo<SocketAddr>` and exposes the framework-owned
1895    // `PeerAddr` extension on both listener branches, so ALL routes on
1896    // the merged router (`/mcp`, `/healthz`, OAuth proxy endpoints,
1897    // admin endpoints, extra_router, ...) and all inner middleware see a
1898    // uniform peer-address contract regardless of TLS. Installed just
1899    // inside the origin check, which stays outermost by design.
1900    let forward_resolver: Option<Arc<ForwardResolver>> = if config.trusted_proxies.is_empty() {
1901        None
1902    } else {
1903        // Entries are guaranteed parseable by `check_trusted_forwarder`;
1904        // filter_map is defensive only.
1905        Some(Arc::new(ForwardResolver {
1906            trusted: config
1907                .trusted_proxies
1908                .iter()
1909                .filter_map(|entry| parse_proxy_net(entry))
1910                .collect(),
1911            mode: config
1912                .forwarded_header
1913                .unwrap_or(ForwardedHeaderMode::XForwardedFor),
1914        }))
1915    };
1916    if forward_resolver.is_some() {
1917        tracing::info!(
1918            proxies = config.trusted_proxies.len(),
1919            "trusted-forwarder mode enabled: limiters key by resolved client IP"
1920        );
1921    }
1922    router = router.layer(axum::middleware::from_fn(move |req, next| {
1923        let r = forward_resolver.clone();
1924        normalize_peer_addr_middleware(r, req, next)
1925    }));
1926
1927    // Origin validation layer (MCP spec: servers MUST validate the
1928    // Origin header to prevent DNS rebinding attacks). Installed as the
1929    // outermost REQUEST-side security layer so it protects ALL routes
1930    // (`/mcp`, `/healthz`, `/readyz`, `/version`, OAuth proxy endpoints,
1931    // admin endpoints, extra_router, etc.) and runs BEFORE auth so we
1932    // reject cross-origin attackers without spending Argon2 cycles. Only
1933    // the response-decorating security-headers layer below sits further out.
1934    //
1935    // Origin-less requests (e.g. server-to-server probes, curl, native
1936    // MCP clients) are permitted; only requests with an Origin header
1937    // that does not match `effective_origins` are rejected.
1938    router = router.layer(axum::middleware::from_fn(move |req, next| {
1939        let origins = Arc::clone(&allowed_origins);
1940        origin_check_middleware(origins, log_request_headers, req, next)
1941    }));
1942
1943    // OWASP security response headers. Installed LAST, making this the
1944    // OUTERMOST response layer: every response -- normal handler output,
1945    // origin-403, CORS preflight, the overload-503 (already converted to a
1946    // Response by the HandleErrorLayer nested inside the load-shed stack),
1947    // and the 404 fallback -- flows back out through it and gains the headers.
1948    // This is response-only decoration: on the request path it is a
1949    // pass-through, so origin still runs before auth and the rate limiter
1950    // still sits inside auth.
1951    let is_tls = config.tls_cert_path.is_some();
1952    let security_headers_cfg = Arc::new(config.security_headers.clone());
1953    router = router.layer(axum::middleware::from_fn(move |req, next| {
1954        let cfg = Arc::clone(&security_headers_cfg);
1955        security_headers_middleware(is_tls, cfg, req, next)
1956    }));
1957
1958    let scheme = if config.tls_cert_path.is_some() {
1959        "https"
1960    } else {
1961        "http"
1962    };
1963
1964    let tls_paths = match (&config.tls_cert_path, &config.tls_key_path) {
1965        (Some(cert), Some(key)) => Some((cert.clone(), key.clone())),
1966        _ => None,
1967    };
1968    let tls_handshake_timeout = config.tls_handshake_timeout;
1969    let max_concurrent_tls_handshakes = config.max_concurrent_tls_handshakes;
1970    let mtls_config = config.auth.as_ref().and_then(|a| a.mtls.as_ref()).cloned();
1971
1972    Ok((
1973        router,
1974        AppRunParams {
1975            tls_paths,
1976            tls_handshake_timeout,
1977            max_concurrent_tls_handshakes,
1978            mtls_config,
1979            shutdown_timeout: config.shutdown_timeout,
1980            auth_state,
1981            rbac_swap,
1982            on_reload_ready: config.on_reload_ready.take(),
1983            ct,
1984            session_ct,
1985            scheme,
1986            name: config.name.clone(),
1987        },
1988    ))
1989}
1990
1991/// Run the MCP HTTP server, binding to `config.bind_addr` and serving
1992/// until an OS shutdown signal (Ctrl-C / SIGTERM) is received.
1993///
1994/// This is the standard entry point for production deployments. For
1995/// deterministic shutdown control (e.g. integration tests), see
1996/// [`serve_with_listener`].
1997///
1998/// The configuration must be validated first via
1999/// [`McpServerConfig::validate`], which returns a [`Validated`] proof
2000/// token. This typestate guarantees, at compile time, that the server
2001/// never starts with an invalid configuration.
2002///
2003/// # Errors
2004///
2005/// Returns [`McpxError::Startup`] if binding to `config.bind_addr`
2006/// fails, or if the underlying axum server returns an error.
2007pub async fn serve<H, F>(
2008    config: Validated<McpServerConfig>,
2009    handler_factory: F,
2010) -> Result<(), McpxError>
2011where
2012    H: ServerHandler + 'static,
2013    F: Fn() -> H + Send + Sync + Clone + 'static,
2014{
2015    let config = config.into_inner();
2016    #[allow(
2017        deprecated,
2018        reason = "internal serve() reads `bind_addr` to construct the listener; field becomes pub(crate) in 1.0"
2019    )]
2020    let bind_addr = config.bind_addr.clone();
2021    let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2022
2023    let listener = TcpListener::bind(&bind_addr)
2024        .await
2025        .map_err(|e| io_to_startup(&format!("bind {bind_addr}"), e))?;
2026    log_listening(&params.name, params.scheme, &bind_addr);
2027
2028    run_server(
2029        router,
2030        listener,
2031        params.tls_paths,
2032        params.tls_handshake_timeout,
2033        params.max_concurrent_tls_handshakes,
2034        params.mtls_config,
2035        params.shutdown_timeout,
2036        params.auth_state,
2037        params.rbac_swap,
2038        params.on_reload_ready,
2039        params.ct,
2040        params.session_ct,
2041    )
2042    .await
2043    .map_err(anyhow_to_startup)
2044}
2045
2046/// Run the MCP HTTP server on a pre-bound [`TcpListener`], with optional
2047/// readiness signalling and external shutdown control.
2048///
2049/// This variant is intended for **deterministic integration tests** and
2050/// for embedders that need to bind the listening socket themselves
2051/// (e.g. systemd socket activation). Compared to [`serve`]:
2052///
2053/// * The caller passes a `TcpListener` that is already bound. This
2054///   eliminates the bind race in tests that previously required
2055///   poll-the-`/healthz`-loop start-up detection.
2056/// * `ready_tx`, when `Some`, receives the socket's
2057///   [`SocketAddr`] *after* the router is built and immediately before
2058///   the server starts accepting connections. Tests can `await` the
2059///   matching `oneshot::Receiver` to know exactly when it is safe to
2060///   issue requests.
2061/// * `shutdown`, when `Some`, gives the caller a
2062///   [`CancellationToken`] that triggers the same graceful-shutdown
2063///   path as a real OS signal. This avoids cross-platform issues with
2064///   sending real `SIGTERM` from tests on Windows.
2065///
2066/// All three optional parameters degrade gracefully: if `ready_tx` is
2067/// `None`, no signal is sent; if `shutdown` is `None`, the server only
2068/// stops on an OS signal (just like [`serve`]).
2069///
2070/// # Errors
2071///
2072/// Returns [`McpxError::Startup`] if router construction fails, if reading
2073/// the listener's `local_addr()` fails, or if the underlying axum
2074/// server returns an error.
2075pub async fn serve_with_listener<H, F>(
2076    listener: TcpListener,
2077    config: Validated<McpServerConfig>,
2078    handler_factory: F,
2079    ready_tx: Option<tokio::sync::oneshot::Sender<SocketAddr>>,
2080    shutdown: Option<CancellationToken>,
2081) -> Result<(), McpxError>
2082where
2083    H: ServerHandler + 'static,
2084    F: Fn() -> H + Send + Sync + Clone + 'static,
2085{
2086    let config = config.into_inner();
2087    let local_addr = listener
2088        .local_addr()
2089        .map_err(|e| io_to_startup("listener.local_addr", e))?;
2090    let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2091
2092    log_listening(&params.name, params.scheme, &local_addr.to_string());
2093
2094    // Forward external shutdown into the server-internal cancellation
2095    // token so `run_server`'s shutdown trigger picks it up alongside
2096    // any real OS signal.
2097    if let Some(external) = shutdown {
2098        let internal = params.ct.clone();
2099        tokio::spawn(async move {
2100            external.cancelled().await;
2101            internal.cancel();
2102        });
2103    }
2104
2105    // Signal readiness *after* the router is fully built and external
2106    // shutdown is wired, but *before* run_server takes ownership of
2107    // the listener. The receiver can immediately issue requests.
2108    if let Some(tx) = ready_tx {
2109        // Receiver may have been dropped (test gave up). That's fine.
2110        let _ = tx.send(local_addr);
2111    }
2112
2113    run_server(
2114        router,
2115        listener,
2116        params.tls_paths,
2117        params.tls_handshake_timeout,
2118        params.max_concurrent_tls_handshakes,
2119        params.mtls_config,
2120        params.shutdown_timeout,
2121        params.auth_state,
2122        params.rbac_swap,
2123        params.on_reload_ready,
2124        params.ct,
2125        params.session_ct,
2126    )
2127    .await
2128    .map_err(anyhow_to_startup)
2129}
2130
2131/// Emit the standard "listening on …" log lines used by both
2132/// [`serve`] and [`serve_with_listener`].
2133#[allow(
2134    clippy::cognitive_complexity,
2135    reason = "tracing::info! macro expansions inflate the score; logic is trivial"
2136)]
2137fn log_listening(name: &str, scheme: &str, addr: &str) {
2138    tracing::info!("{name} listening on {addr}");
2139    tracing::info!("  MCP endpoint: {scheme}://{addr}/mcp");
2140    tracing::info!("  Health check: {scheme}://{addr}/healthz");
2141    tracing::info!("  Readiness:   {scheme}://{addr}/readyz");
2142}
2143
2144/// Drive the chosen axum server variant (TLS or plain) with a graceful
2145/// shutdown window. Consumes the router and listener.
2146///
2147/// # Shutdown semantics
2148///
2149/// A single shutdown trigger (the FIRST of: OS signal via
2150/// `shutdown_signal()`, or external cancellation of `ct`) starts BOTH:
2151///
2152/// 1. axum's `.with_graceful_shutdown(...)` future, which stops
2153///    accepting new connections and waits for in-flight requests to
2154///    drain;
2155/// 2. a `tokio::time::sleep(shutdown_timeout)` race that forces exit if
2156///    drainage exceeds `shutdown_timeout`.
2157///
2158/// Previously this function awaited `shutdown_signal()` independently
2159/// in BOTH branches of a `tokio::select!`. Because `shutdown_signal`
2160/// resolves once per future and consumes one signal, the force-exit
2161/// timer was tied to a SECOND signal (a second SIGTERM the operator
2162/// would never send). Under a single SIGTERM the graceful drain could
2163/// hang indefinitely. The current implementation derives both branches
2164/// from a single shared trigger so the timeout race is anchored to the
2165/// FIRST (and only) signal.
2166#[allow(
2167    clippy::too_many_arguments,
2168    clippy::cognitive_complexity,
2169    reason = "server start-up threads TLS, reload state, and graceful shutdown through one flow"
2170)]
2171async fn run_server(
2172    router: axum::Router,
2173    listener: TcpListener,
2174    tls_paths: Option<(PathBuf, PathBuf)>,
2175    tls_handshake_timeout: Duration,
2176    max_concurrent_tls_handshakes: usize,
2177    mtls_config: Option<MtlsConfig>,
2178    shutdown_timeout: Duration,
2179    auth_state: Option<Arc<AuthState>>,
2180    rbac_swap: Arc<ArcSwap<RbacPolicy>>,
2181    mut on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
2182    ct: CancellationToken,
2183    session_ct: CancellationToken,
2184) -> anyhow::Result<()> {
2185    // `shutdown_trigger` fires when the FIRST source resolves: either
2186    // an OS signal (Ctrl-C / SIGTERM) or external cancellation of `ct`
2187    // (which the test harness uses for deterministic shutdown).
2188    let shutdown_trigger = CancellationToken::new();
2189    {
2190        let trigger = shutdown_trigger.clone();
2191        let parent = ct.clone();
2192        tokio::spawn(async move {
2193            // cancel-safe: both arms (signal future, CancellationToken::cancelled)
2194            // are cancel-safe; the losing arm holds no state.
2195            tokio::select! {
2196                () = shutdown_signal() => {}
2197                () = parent.cancelled() => {}
2198            }
2199            trigger.cancel();
2200        });
2201    }
2202
2203    let graceful = {
2204        let trigger = shutdown_trigger.clone();
2205        let ct = ct.clone();
2206        async move {
2207            trigger.cancelled().await;
2208            tracing::info!("shutting down (grace period: {shutdown_timeout:?})");
2209            ct.cancel();
2210        }
2211    };
2212
2213    let force_exit_timer = {
2214        let trigger = shutdown_trigger.clone();
2215        async move {
2216            trigger.cancelled().await;
2217            tokio::time::sleep(shutdown_timeout).await;
2218        }
2219    };
2220
2221    if let Some((cert_path, key_path)) = tls_paths {
2222        let crl_set = if let Some(mtls) = mtls_config.as_ref()
2223            && mtls.crl_enabled
2224        {
2225            let (ca_certs, roots) = load_client_auth_roots(&mtls.ca_cert_path)?;
2226            let (crl_set, discover_rx) =
2227                mtls_revocation::bootstrap_fetch(roots, &ca_certs, mtls.clone())
2228                    .await
2229                    .map_err(|error| anyhow::anyhow!(error.to_string()))?;
2230            tokio::spawn(mtls_revocation::run_crl_refresher(
2231                Arc::clone(&crl_set),
2232                discover_rx,
2233                ct.clone(),
2234            ));
2235            Some(crl_set)
2236        } else {
2237            None
2238        };
2239
2240        if let Some(cb) = on_reload_ready.take() {
2241            cb(ReloadHandle {
2242                auth: auth_state.clone(),
2243                rbac: Some(Arc::clone(&rbac_swap)),
2244                crl_set: crl_set.clone(),
2245            });
2246        }
2247
2248        let tls_listener = TlsListener::new(
2249            listener,
2250            &cert_path,
2251            &key_path,
2252            mtls_config.as_ref(),
2253            crl_set,
2254            tls_handshake_timeout,
2255            max_concurrent_tls_handshakes,
2256        )?;
2257        let make_svc = router.into_make_service_with_connect_info::<TlsConnInfo>();
2258        // cancel-safe: dropping the serve future on force-exit is intentional
2259        // forced-shutdown semantics; force_exit_timer is a Sleep chain.
2260        tokio::select! {
2261            result = axum::serve(tls_listener, make_svc)
2262                .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2263            () = force_exit_timer => {
2264                tracing::warn!("shutdown timeout exceeded, forcing exit");
2265                session_ct.cancel();
2266            }
2267        }
2268    } else {
2269        if let Some(cb) = on_reload_ready.take() {
2270            cb(ReloadHandle {
2271                auth: auth_state,
2272                rbac: Some(rbac_swap),
2273                crl_set: None,
2274            });
2275        }
2276
2277        let make_svc = router.into_make_service_with_connect_info::<SocketAddr>();
2278        // cancel-safe: dropping the serve future on force-exit is intentional
2279        // forced-shutdown semantics; force_exit_timer is a Sleep chain.
2280        tokio::select! {
2281            result = axum::serve(listener, make_svc)
2282                .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2283            () = force_exit_timer => {
2284                tracing::warn!("shutdown timeout exceeded, forcing exit");
2285                session_ct.cancel();
2286            }
2287        }
2288    }
2289
2290    Ok(())
2291}
2292
2293/// Install the OAuth 2.1 proxy endpoints (`/authorize`, `/token`,
2294/// `/register`, and authorization server metadata) on `router`. The
2295/// caller must ensure `oauth_config.proxy` is `Some`.
2296///
2297/// # Errors
2298///
2299/// Returns [`McpxError::Startup`] if the shared
2300/// [`crate::oauth::OauthHttpClient`] cannot be initialized.
2301#[cfg(feature = "oauth")]
2302fn install_oauth_proxy_routes(
2303    router: axum::Router,
2304    server_url: &str,
2305    oauth_config: &crate::oauth::OAuthConfig,
2306    auth_state: Option<&Arc<AuthState>>,
2307    max_request_body: usize,
2308    admin_role: &str,
2309) -> Result<axum::Router, McpxError> {
2310    let Some(ref proxy) = oauth_config.proxy else {
2311        return Ok(router);
2312    };
2313
2314    // Single shared HTTP client for all proxy endpoints. Cloning is
2315    // cheap (refcounted) and shares the underlying connection pool.
2316    let http = crate::oauth::OauthHttpClient::with_config(oauth_config)?;
2317
2318    // Build the proxy endpoints on a DEDICATED sub-router so the request-body
2319    // cap below applies to exactly these routes and cannot leak onto `/mcp`,
2320    // health, or `/version`. Without this, the proxy routes would fall back to
2321    // axum's 2 MB `DefaultBodyLimit` and silently ignore the operator's
2322    // configured `max_request_body` (rust-review MEDIUM finding).
2323    let proxy_router = axum::Router::new();
2324
2325    let asm = crate::oauth::authorization_server_metadata(server_url, oauth_config);
2326    let proxy_router = proxy_router.route(
2327        "/.well-known/oauth-authorization-server",
2328        axum::routing::get(move || {
2329            let m = asm.clone();
2330            async move { axum::Json(m) }
2331        }),
2332    );
2333
2334    let proxy_authorize = proxy.clone();
2335    let proxy_router = proxy_router.route(
2336        "/authorize",
2337        axum::routing::get(
2338            move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
2339                let p = proxy_authorize.clone();
2340                async move { crate::oauth::handle_authorize(&p, &query.unwrap_or_default()) }
2341            },
2342        ),
2343    );
2344
2345    let proxy_token = proxy.clone();
2346    let token_http = http.clone();
2347    let proxy_router = proxy_router.route(
2348        "/token",
2349        axum::routing::post(move |body: String| {
2350            let p = proxy_token.clone();
2351            let h = token_http.clone();
2352            async move { crate::oauth::handle_token(&h, &p, &body).await }
2353        })
2354        .layer(axum::middleware::from_fn(
2355            oauth_token_cache_headers_middleware,
2356        )),
2357    );
2358
2359    let proxy_register = proxy.clone();
2360    let proxy_router = proxy_router.route(
2361        "/register",
2362        axum::routing::post(move |axum::Json(body): axum::Json<serde_json::Value>| {
2363            let p = proxy_register;
2364            async move { axum::Json(crate::oauth::handle_register(&p, &body)) }
2365        })
2366        .layer(axum::middleware::from_fn(
2367            oauth_token_cache_headers_middleware,
2368        )),
2369    );
2370
2371    let admin_routes_enabled = proxy.expose_admin_endpoints
2372        && (proxy.introspection_url.is_some() || proxy.revocation_url.is_some());
2373    if proxy.expose_admin_endpoints
2374        && !proxy.require_auth_on_admin_endpoints
2375        && proxy.allow_unauthenticated_admin_endpoints
2376    {
2377        // M3 escape-hatch in effect: validate() let this through because
2378        // the operator explicitly opted in. Surface it loudly at startup
2379        // so the choice is auditable in logs.
2380        tracing::warn!(
2381            "OAuth introspect/revoke endpoints are unauthenticated by explicit \
2382             allow_unauthenticated_admin_endpoints opt-out; ensure an \
2383             authenticated reverse proxy fronts these routes"
2384        );
2385    }
2386
2387    let admin_router = if admin_routes_enabled {
2388        build_oauth_admin_router(proxy, http, auth_state, admin_role)?
2389    } else {
2390        axum::Router::new()
2391    };
2392
2393    // Merge admin (introspect/revoke) BEFORE applying the body-limit layer so
2394    // those routes inherit the cap too. `.layer` only wraps routes already
2395    // present on `proxy_router`, so this cannot affect the outer router.
2396    let proxy_router =
2397        proxy_router
2398            .merge(admin_router)
2399            .layer(tower_http::limit::RequestBodyLimitLayer::new(
2400                max_request_body,
2401            ));
2402
2403    let router = router.merge(proxy_router);
2404
2405    tracing::info!(
2406        introspect = proxy.expose_admin_endpoints && proxy.introspection_url.is_some(),
2407        revoke = proxy.expose_admin_endpoints && proxy.revocation_url.is_some(),
2408        max_request_body,
2409        "OAuth 2.1 proxy endpoints enabled (/authorize, /token, /register)"
2410    );
2411    Ok(router)
2412}
2413
2414/// Build the optional `/introspect` + `/revoke` admin sub-router.
2415///
2416/// Layered with [`oauth_token_cache_headers_middleware`] so RFC 6749 §5.1
2417/// / RFC 6750 §5.4 cache headers are emitted, and conditionally with the
2418/// auth middleware when `proxy.require_auth_on_admin_endpoints` is set.
2419#[cfg(feature = "oauth")]
2420fn build_oauth_admin_router(
2421    proxy: &crate::oauth::OAuthProxyConfig,
2422    http: crate::oauth::OauthHttpClient,
2423    auth_state: Option<&Arc<AuthState>>,
2424    admin_role: &str,
2425) -> Result<axum::Router, McpxError> {
2426    let mut admin_router = axum::Router::new();
2427    if proxy.introspection_url.is_some() {
2428        let proxy_introspect = proxy.clone();
2429        let introspect_http = http.clone();
2430        admin_router = admin_router.route(
2431            "/introspect",
2432            axum::routing::post(move |body: String| {
2433                let p = proxy_introspect.clone();
2434                let h = introspect_http.clone();
2435                async move { crate::oauth::handle_introspect(&h, &p, &body).await }
2436            }),
2437        );
2438    }
2439    if proxy.revocation_url.is_some() {
2440        let proxy_revoke = proxy.clone();
2441        let revoke_http = http;
2442        admin_router = admin_router.route(
2443            "/revoke",
2444            axum::routing::post(move |body: String| {
2445                let p = proxy_revoke.clone();
2446                let h = revoke_http.clone();
2447                async move { crate::oauth::handle_revoke(&h, &p, &body).await }
2448            }),
2449        );
2450    }
2451
2452    let admin_router = admin_router.layer(axum::middleware::from_fn(
2453        oauth_token_cache_headers_middleware,
2454    ));
2455
2456    if proxy.require_auth_on_admin_endpoints {
2457        let Some(state) = auth_state else {
2458            return Err(McpxError::Startup(
2459                "oauth proxy admin endpoints require auth state".into(),
2460            ));
2461        };
2462        let state_for_mw = Arc::clone(state);
2463        let required_role: Arc<str> = Arc::from(admin_role);
2464        // M6: gate introspect/revoke behind the admin role. Layers are added
2465        // inner-first, so the role check is added BEFORE auth in order to run
2466        // AFTER it at runtime: auth_middleware (outermost) authenticates and
2467        // populates the AuthIdentity, then require_admin_role rejects any
2468        // authenticated-but-non-admin caller with 403.
2469        Ok(admin_router
2470            .layer(axum::middleware::from_fn(move |req, next| {
2471                let r = Arc::clone(&required_role);
2472                crate::admin::require_admin_role(r, req, next)
2473            }))
2474            .layer(axum::middleware::from_fn(move |req, next| {
2475                let s = Arc::clone(&state_for_mw);
2476                auth_middleware(s, req, next)
2477            })))
2478    } else {
2479        Ok(admin_router)
2480    }
2481}
2482
2483/// Build the host allow-list for rmcp's DNS rebinding protection.
2484///
2485/// Includes loopback hosts by default, then augments with host/authority
2486/// derived from `public_url` and the server bind address.
2487fn derive_allowed_hosts(bind_addr: &str, public_url: Option<&str>) -> Vec<String> {
2488    let mut hosts = vec![
2489        "localhost".to_owned(),
2490        "127.0.0.1".to_owned(),
2491        "::1".to_owned(),
2492    ];
2493
2494    if let Some(url) = public_url
2495        && let Ok(uri) = url.parse::<axum::http::Uri>()
2496        && let Some(authority) = uri.authority()
2497    {
2498        let host = authority.host().to_owned();
2499        if !hosts.iter().any(|h| h == &host) {
2500            hosts.push(host);
2501        }
2502
2503        let authority = authority.as_str().to_owned();
2504        if !hosts.iter().any(|h| h == &authority) {
2505            hosts.push(authority);
2506        }
2507    }
2508
2509    if let Ok(uri) = format!("http://{bind_addr}").parse::<axum::http::Uri>()
2510        && let Some(authority) = uri.authority()
2511    {
2512        let host = authority.host().to_owned();
2513        if !hosts.iter().any(|h| h == &host) {
2514            hosts.push(host);
2515        }
2516
2517        let authority = authority.as_str().to_owned();
2518        if !hosts.iter().any(|h| h == &authority) {
2519            hosts.push(authority);
2520        }
2521    }
2522
2523    hosts
2524}
2525
2526// - TLS support -
2527
2528/// Implement axum's `Connected` trait for `TlsConnInfo` so that
2529/// `ConnectInfo<TlsConnInfo>` is available in middleware when serving
2530/// over our custom `TlsListener`.
2531///
2532/// The identity is read directly from the wrapping
2533/// [`AuthenticatedTlsStream`], which guarantees one-to-one correspondence
2534/// between the TLS connection and its mTLS identity. This eliminates the
2535/// previous shared-map approach which was vulnerable to ephemeral-port
2536/// reuse races (an unauthenticated reconnection from the same `(IP, port)`
2537/// pair could alias a stale entry).
2538impl axum::extract::connect_info::Connected<axum::serve::IncomingStream<'_, TlsListener>>
2539    for TlsConnInfo
2540{
2541    fn connect_info(target: axum::serve::IncomingStream<'_, TlsListener>) -> Self {
2542        let addr = *target.remote_addr();
2543        let identity = target.io().identity().cloned();
2544        Self::new(addr, identity)
2545    }
2546}
2547
2548/// Default per-handshake deadline on the TLS accept path. Prevents idle
2549/// or slow-loris connections from pinning handshake worker tasks (and
2550/// their semaphore permits) indefinitely.
2551///
2552/// Configurable since 1.9.0 via
2553/// [`McpServerConfig::with_tls_handshake_timeout`].
2554const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
2555
2556/// Default upper bound on concurrently in-flight TLS handshakes. When
2557/// saturated, the acceptor task stops pulling new connections from the
2558/// kernel backlog (backpressure) instead of accepting and dropping them
2559/// in user space.
2560///
2561/// Configurable since 1.9.0 via
2562/// [`McpServerConfig::with_max_concurrent_tls_handshakes`].
2563const DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256;
2564
2565/// Capacity of the completed-handshake queue between the acceptor task and
2566/// `axum::serve`'s `accept()` loop. Handshake workers block on `send` when
2567/// the queue is full, so a slow accept loop back-pressures handshakes
2568/// rather than buffering completed connections unboundedly.
2569const TLS_ACCEPT_CHANNEL_CAPACITY: usize = 32;
2570
2571/// A TLS-wrapping listener that implements axum's `Listener` trait.
2572///
2573/// TCP accepts and TLS handshakes run on a dedicated background task: each
2574/// accepted connection's handshake is spawned onto its own worker task,
2575/// bounded by a configurable concurrent-handshake cap (default
2576/// [`DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES`]) and a per-handshake timeout
2577/// (default [`DEFAULT_TLS_HANDSHAKE_TIMEOUT`]). A slow or idle client
2578/// therefore cannot stall other connections behind a serialized inline
2579/// handshake.
2580///
2581/// When mTLS is configured, client certificates are verified against the
2582/// configured CA and the client identity is extracted at handshake time.
2583/// The extracted identity is bound to the connection itself via the
2584/// returned [`AuthenticatedTlsStream`], so it is impossible for an
2585/// unrelated connection to observe it.
2586struct TlsListener {
2587    /// Bound address, captured eagerly before the `TcpListener` moves into
2588    /// the acceptor task.
2589    local_addr: SocketAddr,
2590    /// Completed handshakes produced by the acceptor task's workers.
2591    rx: mpsc::Receiver<(AuthenticatedTlsStream, SocketAddr)>,
2592    /// Background task driving TCP accepts and concurrent TLS handshakes.
2593    /// Aborted on drop so the listener releases its port deterministically.
2594    acceptor_task: tokio::task::JoinHandle<()>,
2595}
2596
2597impl TlsListener {
2598    fn new(
2599        inner: TcpListener,
2600        cert_path: &Path,
2601        key_path: &Path,
2602        mtls_config: Option<&MtlsConfig>,
2603        crl_set: Option<Arc<CrlSet>>,
2604        handshake_timeout: Duration,
2605        max_concurrent_handshakes: usize,
2606    ) -> anyhow::Result<Self> {
2607        // Install the ring crypto provider (ok to call multiple times).
2608        rustls::crypto::ring::default_provider()
2609            .install_default()
2610            .ok();
2611
2612        let certs = load_certs(cert_path)?;
2613        let key = load_key(key_path)?;
2614
2615        let mtls_default_role;
2616
2617        let tls_config = if let Some(mtls) = mtls_config {
2618            mtls_default_role = mtls.default_role.clone();
2619            let verifier: Arc<dyn rustls::server::danger::ClientCertVerifier> = if mtls.crl_enabled
2620            {
2621                let Some(crl_set) = crl_set else {
2622                    return Err(anyhow::anyhow!(
2623                        "mTLS CRL verifier requested but CRL state was not initialized"
2624                    ));
2625                };
2626                Arc::new(DynamicClientCertVerifier::new(crl_set))
2627            } else {
2628                let (_, root_store) = load_client_auth_roots(&mtls.ca_cert_path)?;
2629                if mtls.required {
2630                    rustls::server::WebPkiClientVerifier::builder(root_store)
2631                        .build()
2632                        .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2633                } else {
2634                    rustls::server::WebPkiClientVerifier::builder(root_store)
2635                        .allow_unauthenticated()
2636                        .build()
2637                        .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2638                }
2639            };
2640
2641            tracing::info!(
2642                ca = %mtls.ca_cert_path.display(),
2643                required = mtls.required,
2644                crl_enabled = mtls.crl_enabled,
2645                "mTLS client auth configured"
2646            );
2647
2648            rustls::ServerConfig::builder_with_protocol_versions(&[
2649                &rustls::version::TLS12,
2650                &rustls::version::TLS13,
2651            ])
2652            .with_client_cert_verifier(verifier)
2653            .with_single_cert(certs, key)?
2654        } else {
2655            mtls_default_role = "viewer".to_owned();
2656            rustls::ServerConfig::builder_with_protocol_versions(&[
2657                &rustls::version::TLS12,
2658                &rustls::version::TLS13,
2659            ])
2660            .with_no_client_auth()
2661            .with_single_cert(certs, key)?
2662        };
2663
2664        let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
2665        tracing::info!(
2666            "TLS enabled (cert: {}, key: {})",
2667            cert_path.display(),
2668            key_path.display()
2669        );
2670        let local_addr = inner.local_addr()?;
2671        let (tx, rx) = mpsc::channel(TLS_ACCEPT_CHANNEL_CAPACITY);
2672        let acceptor_task = tokio::spawn(run_tls_acceptor(
2673            inner,
2674            acceptor,
2675            mtls_default_role,
2676            tx,
2677            handshake_timeout,
2678            max_concurrent_handshakes,
2679        ));
2680        Ok(Self {
2681            local_addr,
2682            rx,
2683            acceptor_task,
2684        })
2685    }
2686
2687    /// Extract the mTLS client cert identity from a completed TLS handshake.
2688    /// Returns `None` if no client certificate was presented or if the
2689    /// certificate could not be parsed into an [`AuthIdentity`].
2690    fn extract_handshake_identity(
2691        tls_stream: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2692        default_role: &str,
2693        addr: SocketAddr,
2694    ) -> Option<AuthIdentity> {
2695        let (_, server_conn) = tls_stream.get_ref();
2696        let cert_der = server_conn.peer_certificates()?.first()?;
2697        let id = extract_mtls_identity(cert_der.as_ref(), default_role)?;
2698        tracing::debug!(name = %id.name, peer = %addr, "mTLS client cert accepted");
2699        Some(id)
2700    }
2701}
2702
2703/// Drive TCP accepts and concurrent TLS handshakes for [`TlsListener`].
2704///
2705/// Each accepted connection's handshake runs on its own worker task under
2706/// a permit from a `max_concurrent_handshakes`-sized semaphore and a
2707/// `handshake_timeout` deadline. Completed handshakes are pushed to `tx`;
2708/// failures and timeouts are logged at DEBUG and the connection dropped.
2709/// The loop exits when the owning [`TlsListener`] is dropped.
2710async fn run_tls_acceptor(
2711    listener: TcpListener,
2712    acceptor: tokio_rustls::TlsAcceptor,
2713    default_role: String,
2714    tx: mpsc::Sender<(AuthenticatedTlsStream, SocketAddr)>,
2715    handshake_timeout: Duration,
2716    max_concurrent_handshakes: usize,
2717) {
2718    let inflight = Arc::new(Semaphore::new(max_concurrent_handshakes));
2719    loop {
2720        // Acquire the permit BEFORE accepting: at saturation, pending
2721        // connections wait in the kernel backlog instead of being accepted
2722        // and then buffered or dropped in user space.
2723        let Ok(permit) = Arc::clone(&inflight).acquire_owned().await else {
2724            // The semaphore is never closed; defensive exit.
2725            return;
2726        };
2727        let (stream, addr) = match listener.accept().await {
2728            Ok(pair) => pair,
2729            Err(e) => {
2730                tracing::debug!("TCP accept error: {e}");
2731                continue;
2732            }
2733        };
2734        if tx.is_closed() {
2735            // The listener was dropped (shutdown): stop accepting.
2736            return;
2737        }
2738        let acceptor = acceptor.clone();
2739        let default_role = default_role.clone();
2740        let tx = tx.clone();
2741        tokio::spawn(async move {
2742            let _permit = permit;
2743            match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
2744                Ok(Ok(tls_stream)) => {
2745                    let identity =
2746                        TlsListener::extract_handshake_identity(&tls_stream, &default_role, addr);
2747                    let wrapped = AuthenticatedTlsStream {
2748                        inner: tls_stream,
2749                        identity,
2750                    };
2751                    // The receiver only disappears during shutdown; discard
2752                    // the completed connection quietly rather than logging.
2753                    let _ = tx.send((wrapped, addr)).await;
2754                }
2755                Ok(Err(e)) => {
2756                    tracing::debug!("TLS handshake failed from {addr}: {e}");
2757                }
2758                Err(_elapsed) => {
2759                    tracing::debug!(
2760                        "TLS handshake timed out from {addr} after {handshake_timeout:?}"
2761                    );
2762                }
2763            }
2764        });
2765    }
2766}
2767
2768/// A TLS stream paired with the mTLS identity extracted at handshake time.
2769///
2770/// Wraps [`tokio_rustls::server::TlsStream`] so the verified client
2771/// identity travels with the connection itself. This replaces the previous
2772/// shared `MtlsIdentities` map, eliminating the
2773/// `(SocketAddr) -> AuthIdentity` aliasing risk caused by ephemeral-port
2774/// reuse and removing the need for an LRU eviction policy.
2775///
2776/// The wrapper is `Unpin` (its inner stream is `Unpin` because
2777/// [`tokio::net::TcpStream`] is `Unpin`), so `AsyncRead`/`AsyncWrite`
2778/// delegation uses safe pin projection via `Pin::new(&mut self.inner)`.
2779pub(crate) struct AuthenticatedTlsStream {
2780    inner: tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2781    identity: Option<AuthIdentity>,
2782}
2783
2784impl AuthenticatedTlsStream {
2785    /// Returns the verified mTLS client identity, if any.
2786    #[must_use]
2787    pub(crate) const fn identity(&self) -> Option<&AuthIdentity> {
2788        self.identity.as_ref()
2789    }
2790}
2791
2792impl std::fmt::Debug for AuthenticatedTlsStream {
2793    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2794        f.debug_struct("AuthenticatedTlsStream")
2795            .field("identity", &self.identity.as_ref().map(|id| &id.name))
2796            .finish_non_exhaustive()
2797    }
2798}
2799
2800impl tokio::io::AsyncRead for AuthenticatedTlsStream {
2801    fn poll_read(
2802        mut self: Pin<&mut Self>,
2803        cx: &mut std::task::Context<'_>,
2804        buf: &mut tokio::io::ReadBuf<'_>,
2805    ) -> std::task::Poll<std::io::Result<()>> {
2806        Pin::new(&mut self.inner).poll_read(cx, buf)
2807    }
2808}
2809
2810impl tokio::io::AsyncWrite for AuthenticatedTlsStream {
2811    fn poll_write(
2812        mut self: Pin<&mut Self>,
2813        cx: &mut std::task::Context<'_>,
2814        buf: &[u8],
2815    ) -> std::task::Poll<std::io::Result<usize>> {
2816        Pin::new(&mut self.inner).poll_write(cx, buf)
2817    }
2818
2819    fn poll_flush(
2820        mut self: Pin<&mut Self>,
2821        cx: &mut std::task::Context<'_>,
2822    ) -> std::task::Poll<std::io::Result<()>> {
2823        Pin::new(&mut self.inner).poll_flush(cx)
2824    }
2825
2826    fn poll_shutdown(
2827        mut self: Pin<&mut Self>,
2828        cx: &mut std::task::Context<'_>,
2829    ) -> std::task::Poll<std::io::Result<()>> {
2830        Pin::new(&mut self.inner).poll_shutdown(cx)
2831    }
2832
2833    fn poll_write_vectored(
2834        mut self: Pin<&mut Self>,
2835        cx: &mut std::task::Context<'_>,
2836        bufs: &[std::io::IoSlice<'_>],
2837    ) -> std::task::Poll<std::io::Result<usize>> {
2838        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
2839    }
2840
2841    fn is_write_vectored(&self) -> bool {
2842        self.inner.is_write_vectored()
2843    }
2844}
2845
2846impl axum::serve::Listener for TlsListener {
2847    type Io = AuthenticatedTlsStream;
2848    type Addr = SocketAddr;
2849
2850    /// Yield the next fully-handshaken TLS connection.
2851    ///
2852    /// Cancel-safe: this is a plain `mpsc::Receiver::recv`, so cancelling
2853    /// the future (axum selects it against graceful shutdown) never loses
2854    /// a connection.
2855    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
2856        if let Some(pair) = self.rx.recv().await {
2857            return pair;
2858        }
2859        // The channel only closes if the acceptor task terminated, which
2860        // means the TcpListener is gone and the OS already refuses new
2861        // connections. `Listener::accept` is infallible and panicking is
2862        // forbidden, so park forever: existing connections keep being
2863        // served and graceful shutdown still completes.
2864        tracing::error!("TLS acceptor task terminated; no further connections will be accepted");
2865        std::future::pending().await
2866    }
2867
2868    fn local_addr(&self) -> std::io::Result<Self::Addr> {
2869        Ok(self.local_addr)
2870    }
2871}
2872
2873impl Drop for TlsListener {
2874    fn drop(&mut self) {
2875        // Stop accepting immediately and release the bound port. In-flight
2876        // handshake workers notice the closed channel and exit quietly.
2877        self.acceptor_task.abort();
2878    }
2879}
2880
2881fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
2882    use rustls::pki_types::pem::PemObject;
2883    let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_file_iter(path)
2884        .map_err(|e| anyhow::anyhow!("failed to read certs from {}: {e}", path.display()))?
2885        .collect::<Result<_, _>>()
2886        .map_err(|e| anyhow::anyhow!("invalid cert in {}: {e}", path.display()))?;
2887    anyhow::ensure!(
2888        !certs.is_empty(),
2889        "no certificates found in {}",
2890        path.display()
2891    );
2892    Ok(certs)
2893}
2894
2895fn load_client_auth_roots(
2896    path: &Path,
2897) -> anyhow::Result<(
2898    Vec<rustls::pki_types::CertificateDer<'static>>,
2899    Arc<RootCertStore>,
2900)> {
2901    let ca_certs = load_certs(path)?;
2902    let mut root_store = RootCertStore::empty();
2903    for cert in &ca_certs {
2904        root_store
2905            .add(cert.clone())
2906            .map_err(|error| anyhow::anyhow!("invalid CA cert: {error}"))?;
2907    }
2908
2909    Ok((ca_certs, Arc::new(root_store)))
2910}
2911
2912fn load_key(path: &Path) -> anyhow::Result<rustls::pki_types::PrivateKeyDer<'static>> {
2913    use rustls::pki_types::pem::PemObject;
2914    rustls::pki_types::PrivateKeyDer::from_pem_file(path)
2915        .map_err(|e| anyhow::anyhow!("failed to read key from {}: {e}", path.display()))
2916}
2917
2918#[allow(
2919    clippy::unused_async,
2920    reason = "axum route handler signature requires `async fn` even when the body is synchronous"
2921)]
2922async fn healthz() -> impl IntoResponse {
2923    axum::Json(serde_json::json!({
2924        "status": "ok",
2925    }))
2926}
2927
2928/// Build the `/version` JSON payload for a given server name and version.
2929///
2930/// `name`, `version`, and `mcpx_version` are always included. Build
2931/// metadata (`build_git_sha`, `build_timestamp`, `rust_version`) is added
2932/// only when `expose_build_metadata` is true, so anonymous `/version`
2933/// callers do not receive build fingerprints by default. The build values
2934/// are read at compile time from `RMCP_SERVER_KIT_BUILD_SHA`,
2935/// `RMCP_SERVER_KIT_BUILD_TIME`, and `RMCP_SERVER_KIT_RUSTC_VERSION`;
2936/// unset values resolve to `"unknown"`.
2937fn version_payload(name: &str, version: &str, expose_build_metadata: bool) -> serde_json::Value {
2938    let mut map = serde_json::Map::new();
2939    map.insert("name".into(), name.into());
2940    map.insert("version".into(), version.into());
2941    map.insert("mcpx_version".into(), env!("CARGO_PKG_VERSION").into());
2942    if expose_build_metadata {
2943        map.insert(
2944            "build_git_sha".into(),
2945            option_env!("RMCP_SERVER_KIT_BUILD_SHA")
2946                .unwrap_or("unknown")
2947                .into(),
2948        );
2949        map.insert(
2950            "build_timestamp".into(),
2951            option_env!("RMCP_SERVER_KIT_BUILD_TIME")
2952                .unwrap_or("unknown")
2953                .into(),
2954        );
2955        map.insert(
2956            "rust_version".into(),
2957            option_env!("RMCP_SERVER_KIT_RUSTC_VERSION")
2958                .unwrap_or("unknown")
2959                .into(),
2960        );
2961    }
2962    serde_json::Value::Object(map)
2963}
2964
2965/// Pre-serialize the `/version` payload to immutable bytes.
2966///
2967/// This is called once at router-build time so per-request handling can
2968/// reuse a cheap `Arc<[u8]>` clone instead of re-serializing a
2969/// [`serde_json::Value`] on every hit.
2970///
2971/// Serialization of a flat `serde_json::Value` of static-string fields
2972/// cannot fail in practice; the fallback to `b"{}"` exists only to
2973/// satisfy the crate-wide `unwrap_used` / `expect_used` lint policy.
2974fn serialize_version_payload(name: &str, version: &str, expose_build_metadata: bool) -> Arc<[u8]> {
2975    let value = version_payload(name, version, expose_build_metadata);
2976    serde_json::to_vec(&value).map_or_else(|_| Arc::from(&b"{}"[..]), Arc::from)
2977}
2978
2979async fn readyz(check: ReadinessCheck) -> impl IntoResponse {
2980    let status = check().await;
2981    let ready = status
2982        .get("ready")
2983        .and_then(serde_json::Value::as_bool)
2984        .unwrap_or(false);
2985    let code = if ready {
2986        axum::http::StatusCode::OK
2987    } else {
2988        axum::http::StatusCode::SERVICE_UNAVAILABLE
2989    };
2990    (code, axum::Json(status))
2991}
2992
2993/// Wait for SIGINT (ctrl-c) or SIGTERM (container stop).
2994///
2995/// On non-Unix platforms, only SIGINT is handled.
2996async fn shutdown_signal() {
2997    let ctrl_c = tokio::signal::ctrl_c();
2998
2999    #[cfg(unix)]
3000    {
3001        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
3002            Ok(mut term) => {
3003                // cancel-safe: signal-listener futures are cancel-safe per
3004                // tokio docs; no partial state in either arm.
3005                tokio::select! {
3006                    _ = ctrl_c => {}
3007                    _ = term.recv() => {}
3008                }
3009            }
3010            Err(e) => {
3011                tracing::warn!(error = %e, "failed to register SIGTERM handler, using SIGINT only");
3012                ctrl_c.await.ok();
3013            }
3014        }
3015    }
3016
3017    #[cfg(not(unix))]
3018    {
3019        ctrl_c.await.ok();
3020    }
3021}
3022
3023// -- Origin validation (MCP 2025-11-25 spec, section 2.0.1) --
3024
3025/// Middleware that validates the `Origin` header on incoming HTTP requests.
3026///
3027/// Collapse a request into a bounded set of Prometheus label values.
3028///
3029/// Prometheus retains one time series per distinct label set, and this
3030/// middleware runs OUTSIDE the auth layer, so any label derived from raw
3031/// request input is an unauthenticated memory-growth primitive. Both label
3032/// values must therefore come from a closed set.
3033///
3034/// `MatchedPath` yields the route template for ordinary registered routes,
3035/// but it is not available everywhere: `/mcp` is mounted with `nest_service`,
3036/// whose tail match may carry `MatchedNestedPath` instead, and unmatched
3037/// requests that hit the 404 fallback carry neither. The raw URI path is
3038/// never used as a fallback -- that is precisely the unbounded input.
3039#[cfg(feature = "metrics")]
3040fn metrics_labels(req: &Request<Body>) -> (&'static str, String) {
3041    let method = match *req.method() {
3042        axum::http::Method::GET => "GET",
3043        axum::http::Method::POST => "POST",
3044        axum::http::Method::PUT => "PUT",
3045        axum::http::Method::PATCH => "PATCH",
3046        axum::http::Method::DELETE => "DELETE",
3047        axum::http::Method::HEAD => "HEAD",
3048        axum::http::Method::OPTIONS => "OPTIONS",
3049        axum::http::Method::TRACE => "TRACE",
3050        axum::http::Method::CONNECT => "CONNECT",
3051        // HTTP permits extension methods, so anything else collapses to a
3052        // single bucket rather than minting a series per invented verb.
3053        _ => "OTHER",
3054    };
3055
3056    let path = req
3057        .extensions()
3058        .get::<axum::extract::MatchedPath>()
3059        .map_or_else(
3060            || {
3061                let raw = req.uri().path();
3062                if raw == "/mcp" || raw.starts_with("/mcp/") {
3063                    "/mcp".to_owned()
3064                } else {
3065                    "<unmatched>".to_owned()
3066                }
3067            },
3068            |matched| matched.as_str().to_owned(),
3069        );
3070
3071    (method, path)
3072}
3073
3074/// Record HTTP request metrics (method, path, status, duration).
3075///
3076/// Also exposes the shared [`crate::metrics::McpMetrics`] handle to
3077/// inner middleware via a request extension, so the rate limiters can
3078/// increment `rmcp_server_kit_rate_limited_total` at their deny sites
3079/// (see [`crate::metrics::record_rate_limit_deny`]).
3080#[cfg(feature = "metrics")]
3081async fn metrics_middleware(
3082    metrics: Arc<crate::metrics::McpMetrics>,
3083    mut req: Request<Body>,
3084    next: Next,
3085) -> axum::response::Response {
3086    let (method, path) = metrics_labels(&req);
3087    let start = std::time::Instant::now();
3088
3089    req.extensions_mut().insert(Arc::clone(&metrics));
3090    let response = next.run(req).await;
3091
3092    let mut status_buf = core::fmt::NumBuffer::<u16>::new();
3093    let status = response.status().as_u16().format_into(&mut status_buf);
3094    let duration = start.elapsed().as_secs_f64();
3095
3096    metrics
3097        .http_requests_total
3098        .with_label_values(&[method, &path, status])
3099        .inc();
3100    metrics
3101        .http_request_duration_seconds
3102        .with_label_values(&[method, &path])
3103        .observe(duration);
3104
3105    response
3106}
3107
3108/// OWASP security header hardening applied to every response.
3109///
3110/// Sets: `X-Content-Type-Options`, `X-Frame-Options`, `Cache-Control`,
3111/// `Referrer-Policy`, `Cross-Origin-Opener-Policy`, `Cross-Origin-Resource-Policy`,
3112/// `Cross-Origin-Embedder-Policy`, `Permissions-Policy`,
3113/// `X-Permitted-Cross-Domain-Policies`, `Content-Security-Policy`,
3114/// `X-DNS-Prefetch-Control`, and (when TLS is active) `Strict-Transport-Security`.
3115///
3116/// Each header's value can be customised via [`SecurityHeadersConfig`]
3117/// on [`McpServerConfig`]. See that type for the three-state semantic
3118/// (`None` = default, `Some("")` = omit, `Some(v)` = override).
3119async fn security_headers_middleware(
3120    is_tls: bool,
3121    cfg: Arc<SecurityHeadersConfig>,
3122    req: Request<Body>,
3123    next: Next,
3124) -> axum::response::Response {
3125    use axum::http::{HeaderName, header};
3126
3127    let mut resp = next.run(req).await;
3128    let headers = resp.headers_mut();
3129
3130    // Strip server identity headers to reduce information leakage.
3131    headers.remove(header::SERVER);
3132    headers.remove(HeaderName::from_static("x-powered-by"));
3133
3134    apply_security_header(
3135        headers,
3136        header::X_CONTENT_TYPE_OPTIONS,
3137        cfg.x_content_type_options.as_deref(),
3138        "nosniff",
3139    );
3140    apply_security_header(
3141        headers,
3142        header::X_FRAME_OPTIONS,
3143        cfg.x_frame_options.as_deref(),
3144        "deny",
3145    );
3146    apply_security_header(
3147        headers,
3148        header::CACHE_CONTROL,
3149        cfg.cache_control.as_deref(),
3150        "no-store, max-age=0",
3151    );
3152    apply_security_header(
3153        headers,
3154        header::REFERRER_POLICY,
3155        cfg.referrer_policy.as_deref(),
3156        "no-referrer",
3157    );
3158    apply_security_header(
3159        headers,
3160        HeaderName::from_static("cross-origin-opener-policy"),
3161        cfg.cross_origin_opener_policy.as_deref(),
3162        "same-origin",
3163    );
3164    apply_security_header(
3165        headers,
3166        HeaderName::from_static("cross-origin-resource-policy"),
3167        cfg.cross_origin_resource_policy.as_deref(),
3168        "same-origin",
3169    );
3170    apply_security_header(
3171        headers,
3172        HeaderName::from_static("cross-origin-embedder-policy"),
3173        cfg.cross_origin_embedder_policy.as_deref(),
3174        "require-corp",
3175    );
3176    apply_security_header(
3177        headers,
3178        HeaderName::from_static("permissions-policy"),
3179        cfg.permissions_policy.as_deref(),
3180        "accelerometer=(), camera=(), geolocation=(), microphone=()",
3181    );
3182    apply_security_header(
3183        headers,
3184        HeaderName::from_static("x-permitted-cross-domain-policies"),
3185        cfg.x_permitted_cross_domain_policies.as_deref(),
3186        "none",
3187    );
3188    apply_security_header(
3189        headers,
3190        HeaderName::from_static("content-security-policy"),
3191        cfg.content_security_policy.as_deref(),
3192        "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests",
3193    );
3194    apply_security_header(
3195        headers,
3196        HeaderName::from_static("x-dns-prefetch-control"),
3197        cfg.x_dns_prefetch_control.as_deref(),
3198        "off",
3199    );
3200
3201    if is_tls {
3202        apply_security_header(
3203            headers,
3204            header::STRICT_TRANSPORT_SECURITY,
3205            cfg.strict_transport_security.as_deref(),
3206            "max-age=63072000; includeSubDomains",
3207        );
3208    }
3209
3210    resp
3211}
3212
3213/// Set a single security header on the response, honouring the
3214/// three-state override semantic (None = default, Some("") = omit,
3215/// Some(value) = override).
3216///
3217/// Defence-in-depth: if an override value somehow reaches this point
3218/// despite [`validate_security_headers`] having approved it (e.g. a
3219/// runtime mutation on a non-`Validated` field), we log at error level
3220/// and fall back to the static default rather than panicking. The
3221/// `Validated<McpServerConfig>` type makes that path unreachable in
3222/// well-typed code paths.
3223fn apply_security_header(
3224    headers: &mut axum::http::HeaderMap,
3225    name: axum::http::HeaderName,
3226    override_value: Option<&str>,
3227    default: &'static str,
3228) {
3229    use axum::http::HeaderValue;
3230
3231    match override_value {
3232        None => {
3233            headers.insert(name, HeaderValue::from_static(default));
3234        }
3235        Some("") => {
3236            // Operator explicitly opted out of this header.
3237        }
3238        Some(v) => match HeaderValue::from_str(v) {
3239            Ok(hv) => {
3240                headers.insert(name, hv);
3241            }
3242            Err(err) => {
3243                tracing::error!(
3244                    header = %name,
3245                    error = %err,
3246                    "invalid security header override reached middleware; using default"
3247                );
3248                headers.insert(name, HeaderValue::from_static(default));
3249            }
3250        },
3251    }
3252}
3253
3254/// Validate every non-empty entry in a [`SecurityHeadersConfig`].
3255///
3256/// - `None` and `Some("")` are accepted unconditionally (use-default and
3257///   omit, respectively).
3258/// - `Some(v)` is rejected if `axum::http::HeaderValue::from_str(v)` fails.
3259/// - `strict_transport_security` additionally rejects any value
3260///   containing `preload` (case-insensitive). Operators who genuinely
3261///   want to commit to the HSTS preload list must do so via a future
3262///   explicit `with_hsts_preload(true)` builder, not by smuggling
3263///   `preload` through this knob.
3264fn validate_security_headers(cfg: &SecurityHeadersConfig) -> Result<(), McpxError> {
3265    use axum::http::HeaderValue;
3266
3267    let fields: &[(&str, Option<&str>)] = &[
3268        (
3269            "x_content_type_options",
3270            cfg.x_content_type_options.as_deref(),
3271        ),
3272        ("x_frame_options", cfg.x_frame_options.as_deref()),
3273        ("cache_control", cfg.cache_control.as_deref()),
3274        ("referrer_policy", cfg.referrer_policy.as_deref()),
3275        (
3276            "cross_origin_opener_policy",
3277            cfg.cross_origin_opener_policy.as_deref(),
3278        ),
3279        (
3280            "cross_origin_resource_policy",
3281            cfg.cross_origin_resource_policy.as_deref(),
3282        ),
3283        (
3284            "cross_origin_embedder_policy",
3285            cfg.cross_origin_embedder_policy.as_deref(),
3286        ),
3287        ("permissions_policy", cfg.permissions_policy.as_deref()),
3288        (
3289            "x_permitted_cross_domain_policies",
3290            cfg.x_permitted_cross_domain_policies.as_deref(),
3291        ),
3292        (
3293            "content_security_policy",
3294            cfg.content_security_policy.as_deref(),
3295        ),
3296        (
3297            "x_dns_prefetch_control",
3298            cfg.x_dns_prefetch_control.as_deref(),
3299        ),
3300        (
3301            "strict_transport_security",
3302            cfg.strict_transport_security.as_deref(),
3303        ),
3304    ];
3305
3306    for (field, value) in fields {
3307        let Some(v) = value else { continue };
3308        if v.is_empty() {
3309            continue;
3310        }
3311        if let Err(err) = HeaderValue::from_str(v) {
3312            return Err(McpxError::Config(format!(
3313                "invalid security_headers.{field}: {err}"
3314            )));
3315        }
3316    }
3317
3318    if let Some(v) = cfg.strict_transport_security.as_deref()
3319        && !v.is_empty()
3320        && v.to_ascii_lowercase().contains("preload")
3321    {
3322        return Err(McpxError::Config(format!(
3323            "invalid security_headers.strict_transport_security: {v:?} contains the `preload` directive; \
3324             HSTS preload must be opted into explicitly via a dedicated builder, not via this knob"
3325        )));
3326    }
3327
3328    Ok(())
3329}
3330
3331/// Append RFC 6749 §5.1 / RFC 6750 §5.4 cache and `Vary` headers required
3332/// on OAuth token-issuing responses.
3333///
3334/// `Cache-Control: no-store, max-age=0` is already applied globally by
3335/// [`security_headers_middleware`]; this middleware adds:
3336///
3337/// - `Pragma: no-cache` -- mandated by RFC 6749 §5.1 for HTTP/1.0 caches.
3338/// - `Vary: Authorization` -- mandated by RFC 6750 §5.4 for endpoints
3339///   whose response depends on the `Authorization` header.
3340///
3341/// Applied only to the OAuth proxy token-class endpoints (`/token`,
3342/// `/register`, `/introspect`, `/revoke`). `Vary` is appended (not
3343/// inserted) so any `Vary` value already present (e.g. `Accept-Encoding`
3344/// from a compression layer, or `Origin` from a CORS layer) is preserved.
3345#[cfg(feature = "oauth")]
3346async fn oauth_token_cache_headers_middleware(
3347    req: Request<Body>,
3348    next: Next,
3349) -> axum::response::Response {
3350    use axum::http::{HeaderValue, header};
3351
3352    let mut resp = next.run(req).await;
3353    let headers = resp.headers_mut();
3354    headers.insert(header::PRAGMA, HeaderValue::from_static("no-cache"));
3355    headers.append(header::VARY, HeaderValue::from_static("Authorization"));
3356    resp
3357}
3358
3359/// Normalize peer-address request extensions across listener branches.
3360///
3361/// The make-service installs `ConnectInfo<SocketAddr>` on the plain
3362/// listener but `ConnectInfo<TlsConnInfo>` on the TLS listener (the
3363/// latter additionally carries the connection-bound mTLS identity and
3364/// stays `pub(crate)` — see the anti-aliasing rationale on
3365/// [`TlsConnInfo`]). Application routes — in particular those merged via
3366/// [`McpServerConfig::with_extra_router`], which bypass the auth
3367/// middleware and its private fallback — could therefore not read the
3368/// peer address under TLS.
3369///
3370/// This middleware makes both branches look identical to every route and
3371/// inner middleware:
3372///
3373/// 1. mirrors the TLS peer address into `ConnectInfo<SocketAddr>` when
3374///    (and only when) it is absent, so stock axum-ecosystem extractors
3375///    work unmodified,
3376/// 2. inserts the framework-owned [`PeerAddr`] extension on both
3377///    branches, and
3378/// 3. inserts the resolved [`ClientIp`] extension: the direct peer's IP,
3379///    unless trusted-forwarder mode is configured AND the direct peer is
3380///    a trusted proxy AND the forwarding chain resolves — every
3381///    ambiguous chain falls back to the direct peer with only a reason
3382///    code logged at `debug` (never raw header contents).
3383///
3384/// Precedence mirrors the auth middleware: an existing
3385/// `ConnectInfo<SocketAddr>` always wins and is never overwritten. The
3386/// peer address is deliberately not logged here.
3387async fn normalize_peer_addr_middleware(
3388    resolver: Option<Arc<ForwardResolver>>,
3389    mut req: Request<Body>,
3390    next: Next,
3391) -> axum::response::Response {
3392    let direct = req
3393        .extensions()
3394        .get::<ConnectInfo<SocketAddr>>()
3395        .map(|ci| ci.0);
3396    let from_tls = req
3397        .extensions()
3398        .get::<ConnectInfo<TlsConnInfo>>()
3399        .map(|ci| ci.0.addr);
3400    if let Some(addr) = direct.or(from_tls) {
3401        if direct.is_none() {
3402            req.extensions_mut().insert(ConnectInfo(addr));
3403        }
3404        req.extensions_mut().insert(PeerAddr::new(addr));
3405        let client_ip = match &resolver {
3406            Some(r) => {
3407                crate::forwarded::resolve_client_ip(addr.ip(), req.headers(), &r.trusted, r.mode)
3408                    .unwrap_or_else(|reason| {
3409                        tracing::debug!(
3410                            reason = ?reason,
3411                            "forwarded-header resolution fell back to direct peer"
3412                        );
3413                        addr.ip()
3414                    })
3415            }
3416            None => addr.ip(),
3417        };
3418        req.extensions_mut().insert(ClientIp::new(client_ip));
3419    }
3420    next.run(req).await
3421}
3422
3423/// Parse a trusted-proxy entry: a CIDR (`10.0.0.0/8`) or a bare IP
3424/// (normalized to a `/32` / `/128` host network).
3425fn parse_proxy_net(entry: &str) -> Option<ipnet::IpNet> {
3426    if let Ok(net) = entry.parse::<ipnet::IpNet>() {
3427        return Some(net);
3428    }
3429    entry.parse::<IpAddr>().ok().map(ipnet::IpNet::from)
3430}
3431
3432/// Validate one `trusted_proxies` entry. Accepts a CIDR (`ipnet::IpNet`)
3433/// or a bare IP; **rejects a `/0` prefix**, which would mark every peer
3434/// trusted and let any client spoof the resolved client IP via forwarding
3435/// headers. Shared by the builder ([`McpServerConfig::check_trusted_forwarder`])
3436/// and the TOML validator so the two validators cannot drift.
3437///
3438/// # Errors
3439///
3440/// Returns a message when the entry is unparseable or carries a `/0` prefix.
3441pub(crate) fn validate_trusted_proxy_entry(entry: &str) -> Result<(), String> {
3442    match parse_proxy_net(entry) {
3443        None => Err(format!(
3444            "trusted_proxies entry {entry:?} is neither a CIDR nor an IP address"
3445        )),
3446        Some(net) if net.prefix_len() == 0 => Err(format!(
3447            "trusted_proxies entry {entry:?}: prefix length 0 is forbidden (marks every peer trusted, enabling client-IP spoofing)"
3448        )),
3449        Some(_) => Ok(()),
3450    }
3451}
3452
3453/// Rate-limit key for the current request: the resolved [`ClientIp`]
3454/// when present, else the direct peer from either `ConnectInfo` form.
3455/// All four built-in limiters key through this helper.
3456pub(crate) fn limiter_client_ip(extensions: &axum::http::Extensions) -> Option<IpAddr> {
3457    if let Some(client) = extensions.get::<ClientIp>() {
3458        return Some(client.ip);
3459    }
3460    extensions
3461        .get::<ConnectInfo<SocketAddr>>()
3462        .map(|ci| ci.0.ip())
3463        .or_else(|| {
3464            extensions
3465                .get::<ConnectInfo<TlsConnInfo>>()
3466                .map(|ci| ci.0.addr.ip())
3467        })
3468}
3469
3470/// Per-IP rate limiter for `extra_router` routes, keyed by the direct
3471/// socket peer address. Same memory-bounded machinery as the tool
3472/// limiter ([`crate::rbac`]).
3473pub(crate) type ExtraRouteRateLimiter = BoundedKeyedLimiter<IpAddr>;
3474
3475/// Cap on distinct source IPs tracked by the extra-route limiter.
3476/// Mirrors the tool limiter's bound: memory stays bounded at saturation
3477/// via idle-prune + LRU eviction, at the cost of shared-fate fairness
3478/// under key spray (an attacker churning many IPs can reset quieter
3479/// legitimate IPs to fresh buckets).
3480const EXTRA_ROUTE_MAX_TRACKED_KEYS: usize = 10_000;
3481
3482/// Idle-eviction window for the extra-route limiter (15 minutes),
3483/// mirroring the tool limiter.
3484const EXTRA_ROUTE_IDLE_EVICTION: Duration = Duration::from_mins(15);
3485
3486/// Build the per-IP limiter for `extra_router` routes.
3487///
3488/// `per_minute` and `burst` are validated nonzero by
3489/// [`McpServerConfig::validate`]; the `NonZeroU32` fallbacks here are
3490/// defensive only. `burst` overrides governor's default bucket capacity
3491/// (burst = rate).
3492fn build_extra_route_rate_limiter(
3493    per_minute: u32,
3494    burst: Option<u32>,
3495) -> Arc<ExtraRouteRateLimiter> {
3496    let rate = std::num::NonZeroU32::new(per_minute.max(1)).unwrap_or(std::num::NonZeroU32::MIN);
3497    let mut quota = governor::Quota::per_minute(rate);
3498    if let Some(b) = burst.and_then(std::num::NonZeroU32::new) {
3499        quota = quota.allow_burst(b);
3500    }
3501    Arc::new(BoundedKeyedLimiter::new(
3502        quota,
3503        EXTRA_ROUTE_MAX_TRACKED_KEYS,
3504        EXTRA_ROUTE_IDLE_EVICTION,
3505    ))
3506}
3507
3508/// Per-IP rate limit middleware for `extra_router` routes.
3509///
3510/// Applied to the application-supplied router **before** it is merged
3511/// into the top-level router, so it wraps exactly the extra routes
3512/// (and their fallback, if any) and nothing else — `/mcp`, health,
3513/// admin, and OAuth endpoints are never affected. Outer layers (origin
3514/// check, peer-address normalization, security headers, metrics) still
3515/// wrap these routes and run first, so both `ConnectInfo` forms are
3516/// populated by the time this middleware reads them.
3517///
3518/// Semantics mirror the tool/auth limiters exactly: keyed by the
3519/// direct peer `IpAddr` (no `X-Forwarded-For`), fail-open when no peer
3520/// address is present (cannot happen under [`serve`]), and on limit a
3521/// plain-text 429 via [`McpxError::RateLimitedFor`] carrying a
3522/// `Retry-After` header (delta-seconds), consistent with every other
3523/// limiter in the crate.
3524///
3525/// `exempt` holds raw exact-match paths (validated at config time)
3526/// checked against `req.uri().path()` **before** key extraction:
3527/// exempt requests consume no limiter budget and produce no deny
3528/// telemetry. Fail-closed — any non-listed path stays limited.
3529async fn extra_route_rate_limit_middleware(
3530    limiter: Arc<ExtraRouteRateLimiter>,
3531    exempt: Arc<std::collections::HashSet<String>>,
3532    req: Request<Body>,
3533    next: Next,
3534) -> axum::response::Response {
3535    if exempt.contains(req.uri().path()) {
3536        return next.run(req).await;
3537    }
3538    let peer_ip: Option<IpAddr> = limiter_client_ip(req.extensions());
3539    if let Some(ip) = peer_ip
3540        && let Err(wait) = limiter.check_key_wait(&ip)
3541    {
3542        #[cfg(feature = "metrics")]
3543        crate::metrics::record_rate_limit_deny(req.extensions(), "extra_route");
3544        tracing::warn!(%ip, "extra route request rate limited");
3545        return McpxError::RateLimitedFor {
3546            message: "too many requests to application routes from this source".into(),
3547            retry_after: wait,
3548        }
3549        .into_response();
3550    }
3551    next.run(req).await
3552}
3553
3554/// Per the MCP spec: if the Origin header is present and its value is not in
3555/// the allowed list, respond with 403 Forbidden. Requests without an Origin
3556/// header are allowed through (e.g. non-browser clients like curl, SDKs).
3557async fn origin_check_middleware(
3558    allowed: Arc<[String]>,
3559    log_request_headers: bool,
3560    req: Request<Body>,
3561    next: Next,
3562) -> axum::response::Response {
3563    let method = req.method().clone();
3564    let path = req.uri().path().to_owned();
3565
3566    log_incoming_request(&method, &path, req.headers(), log_request_headers);
3567
3568    if let Some(origin) = req.headers().get(axum::http::header::ORIGIN) {
3569        let origin_str = origin.to_str().unwrap_or("");
3570        if !allowed.iter().any(|a| a == origin_str) {
3571            tracing::warn!(
3572                origin = origin_str,
3573                %method,
3574                %path,
3575                allowed = ?&*allowed,
3576                "rejected request: Origin not allowed"
3577            );
3578            return (
3579                axum::http::StatusCode::FORBIDDEN,
3580                "Forbidden: Origin not allowed",
3581            )
3582                .into_response();
3583        }
3584    }
3585    next.run(req).await
3586}
3587
3588/// Emit a DEBUG log for an incoming request, optionally including the full
3589/// (redacted) header set.
3590fn log_incoming_request(
3591    method: &axum::http::Method,
3592    path: &str,
3593    headers: &axum::http::HeaderMap,
3594    log_request_headers: bool,
3595) {
3596    if log_request_headers {
3597        tracing::debug!(
3598            %method,
3599            %path,
3600            headers = %format_request_headers_for_log(headers),
3601            "incoming request"
3602        );
3603    } else {
3604        tracing::debug!(%method, %path, "incoming request");
3605    }
3606}
3607
3608fn format_request_headers_for_log(headers: &axum::http::HeaderMap) -> String {
3609    headers
3610        .iter()
3611        .map(|(k, v)| {
3612            let name = k.as_str();
3613            if name == "authorization" || name == "cookie" || name == "proxy-authorization" {
3614                format!("{name}: [REDACTED]")
3615            } else {
3616                format!("{name}: {}", v.to_str().unwrap_or("<non-utf8>"))
3617            }
3618        })
3619        .collect::<Vec<_>>()
3620        .join(", ")
3621}
3622
3623// -- stdio transport --
3624
3625/// Serve an MCP server over stdin/stdout (stdio transport).
3626///
3627/// # Security warnings
3628///
3629/// - **No authentication**: the parent process has full, unrestricted access.
3630/// - **No RBAC**: all tools are available regardless of policy.
3631/// - **No TLS**: messages travel over OS pipes in plaintext.
3632/// - **Single client**: only the parent process can connect.
3633/// - **No Origin validation**: not applicable to stdio.
3634///
3635/// Use this only when the MCP client spawns the server as a trusted subprocess
3636/// (e.g. Claude Desktop, VS Code Copilot). For network-accessible deployments,
3637/// use `serve()` (Streamable HTTP) instead.
3638///
3639/// # Errors
3640///
3641/// Returns [`McpxError::Startup`] if the handler fails to initialize or the
3642/// transport disconnects unexpectedly.
3643// NOTE: reported complexity 32/25 is driven entirely by `tracing::*!`
3644// macro expansion in this 18-line function (info/warn/info + two matches).
3645// There is nothing meaningful to extract; the allow stays.
3646#[allow(
3647    clippy::cognitive_complexity,
3648    reason = "complexity is purely tracing macro expansion (info/warn + match arms); 18 lines of straight-line code, nothing meaningful to extract"
3649)]
3650pub async fn serve_stdio<H>(handler: H) -> Result<(), McpxError>
3651where
3652    H: ServerHandler + 'static,
3653{
3654    use rmcp::ServiceExt as _;
3655
3656    tracing::info!("stdio transport: serving on stdin/stdout");
3657    tracing::warn!("stdio mode: auth, RBAC, TLS, and Origin checks are DISABLED");
3658
3659    let transport = rmcp::transport::io::stdio();
3660
3661    let service = handler
3662        .serve(transport)
3663        .await
3664        .map_err(|e| McpxError::Startup(format!("stdio initialize failed: {e}")))?;
3665
3666    if let Err(e) = service.waiting().await {
3667        tracing::warn!(error = %e, "stdio session ended with error");
3668    }
3669    tracing::info!("stdio session ended");
3670    Ok(())
3671}
3672
3673#[cfg(test)]
3674mod tests {
3675    #![allow(
3676        clippy::unwrap_used,
3677        clippy::expect_used,
3678        clippy::panic,
3679        clippy::indexing_slicing,
3680        clippy::unwrap_in_result,
3681        clippy::print_stdout,
3682        clippy::print_stderr,
3683        deprecated,
3684        reason = "internal unit tests legitimately read/write the deprecated `pub` fields they were designed to verify"
3685    )]
3686    use std::{sync::Arc, time::Duration};
3687
3688    use axum::{
3689        body::Body,
3690        http::{Request, StatusCode, header},
3691        response::IntoResponse,
3692    };
3693    use http_body_util::BodyExt;
3694    use tower::ServiceExt as _;
3695
3696    use super::*;
3697
3698    // -- McpServerConfig --
3699
3700    #[test]
3701    fn server_config_new_defaults() {
3702        let cfg = McpServerConfig::new("0.0.0.0:8443", "test-server", "1.0.0");
3703        assert_eq!(cfg.bind_addr, "0.0.0.0:8443");
3704        assert_eq!(cfg.name, "test-server");
3705        assert_eq!(cfg.version, "1.0.0");
3706        assert!(cfg.tls_cert_path.is_none());
3707        assert!(cfg.tls_key_path.is_none());
3708        assert!(cfg.auth.is_none());
3709        assert!(cfg.rbac.is_none());
3710        assert!(cfg.allowed_origins.is_empty());
3711        assert!(cfg.tool_rate_limit.is_none());
3712        assert!(cfg.readiness_check.is_none());
3713        assert_eq!(cfg.max_request_body, 1024 * 1024);
3714        assert_eq!(cfg.request_timeout, Duration::from_mins(2));
3715        assert_eq!(cfg.shutdown_timeout, Duration::from_secs(30));
3716        assert!(!cfg.log_request_headers);
3717        assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(10));
3718        assert_eq!(cfg.max_concurrent_tls_handshakes, 256);
3719    }
3720
3721    #[test]
3722    fn tls_handshake_builders_set_fields() {
3723        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3724            .with_tls_handshake_timeout(Duration::from_secs(3))
3725            .with_max_concurrent_tls_handshakes(64);
3726        assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(3));
3727        assert_eq!(cfg.max_concurrent_tls_handshakes, 64);
3728    }
3729
3730    #[test]
3731    fn validate_rejects_zero_tls_handshake_timeout() {
3732        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3733            .with_tls_handshake_timeout(Duration::ZERO);
3734        let err = cfg.validate().expect_err("zero handshake timeout");
3735        assert!(err.to_string().contains("tls_handshake_timeout"));
3736    }
3737
3738    #[test]
3739    fn validate_rejects_zero_max_concurrent_tls_handshakes() {
3740        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3741            .with_max_concurrent_tls_handshakes(0);
3742        let err = cfg.validate().expect_err("zero handshake concurrency");
3743        assert!(err.to_string().contains("max_concurrent_tls_handshakes"));
3744    }
3745
3746    #[test]
3747    fn validate_consumes_and_proves() {
3748        // Valid config -> Validated wrapper, original is consumed.
3749        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
3750        let validated = cfg.validate().expect("valid config");
3751        // as_inner() gives read-only access to inner fields.
3752        assert_eq!(validated.as_inner().name, "test-server");
3753        // into_inner recovers the raw value.
3754        let raw = validated.into_inner();
3755        assert_eq!(raw.name, "test-server");
3756
3757        // Invalid config (zero max_request_body) -> Err.
3758        let mut bad = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
3759        bad.max_request_body = 0;
3760        assert!(bad.validate().is_err(), "zero body cap must fail validate");
3761    }
3762
3763    #[test]
3764    fn validate_rejects_zero_max_concurrent_requests() {
3765        let cfg =
3766            McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_max_concurrent_requests(0);
3767        let err = cfg.validate().expect_err("zero concurrency cap must fail");
3768        assert!(
3769            format!("{err}").contains("max_concurrent_requests"),
3770            "error should mention max_concurrent_requests, got: {err}"
3771        );
3772    }
3773
3774    #[test]
3775    fn validate_rejects_zero_max_tracked_keys() {
3776        // Defaults mirror auth::default_max_attempts / default_idle_eviction
3777        // (module-private in auth.rs); spelled out here for review clarity.
3778        let rl = crate::auth::RateLimitConfig {
3779            max_attempts_per_minute: 30,
3780            pre_auth_max_per_minute: None,
3781            max_tracked_keys: 0,
3782            idle_eviction: Duration::from_secs(15 * 60),
3783            burst: None,
3784            pre_auth_burst: None,
3785        };
3786        let auth_cfg = AuthConfig {
3787            enabled: true,
3788            api_keys: Vec::new(),
3789            mtls: None,
3790            rate_limit: Some(rl),
3791            #[cfg(feature = "oauth")]
3792            oauth: None,
3793        };
3794        let cfg = McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_auth(auth_cfg);
3795        let err = cfg.validate().expect_err("zero max_tracked_keys must fail");
3796        assert!(
3797            format!("{err}").contains("max_tracked_keys"),
3798            "error should mention max_tracked_keys, got: {err}"
3799        );
3800    }
3801
3802    #[test]
3803    fn derive_allowed_hosts_includes_public_host() {
3804        let hosts = derive_allowed_hosts("0.0.0.0:8080", Some("https://mcp.example.com/mcp"));
3805        assert!(
3806            hosts.iter().any(|h| h == "mcp.example.com"),
3807            "public_url host must be allowed"
3808        );
3809    }
3810
3811    #[test]
3812    fn derive_allowed_hosts_includes_bind_authority() {
3813        let hosts = derive_allowed_hosts("127.0.0.1:8080", None);
3814        assert!(
3815            hosts.iter().any(|h| h == "127.0.0.1"),
3816            "bind host must be allowed"
3817        );
3818        assert!(
3819            hosts.iter().any(|h| h == "127.0.0.1:8080"),
3820            "bind authority must be allowed"
3821        );
3822    }
3823
3824    // -- healthz --
3825
3826    #[tokio::test]
3827    async fn healthz_returns_ok_json() {
3828        let resp = healthz().await.into_response();
3829        assert_eq!(resp.status(), StatusCode::OK);
3830        let body = resp.into_body().collect().await.unwrap().to_bytes();
3831        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
3832        assert_eq!(json["status"], "ok");
3833        assert!(
3834            json.get("name").is_none(),
3835            "healthz must not expose server name"
3836        );
3837        assert!(
3838            json.get("version").is_none(),
3839            "healthz must not expose version"
3840        );
3841    }
3842
3843    // -- readyz --
3844
3845    #[tokio::test]
3846    async fn readyz_returns_ok_when_ready() {
3847        let check: ReadinessCheck =
3848            Arc::new(|| Box::pin(async { serde_json::json!({"ready": true, "db": "connected"}) }));
3849        let resp = readyz(check).await.into_response();
3850        assert_eq!(resp.status(), StatusCode::OK);
3851        let body = resp.into_body().collect().await.unwrap().to_bytes();
3852        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
3853        assert_eq!(json["ready"], true);
3854        assert!(
3855            json.get("name").is_none(),
3856            "readyz must not expose server name"
3857        );
3858        assert!(
3859            json.get("version").is_none(),
3860            "readyz must not expose version"
3861        );
3862        assert_eq!(json["db"], "connected");
3863    }
3864
3865    #[tokio::test]
3866    async fn readyz_returns_503_when_not_ready() {
3867        let check: ReadinessCheck =
3868            Arc::new(|| Box::pin(async { serde_json::json!({"ready": false}) }));
3869        let resp = readyz(check).await.into_response();
3870        assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
3871    }
3872
3873    #[tokio::test]
3874    async fn readyz_returns_503_when_ready_missing() {
3875        let check: ReadinessCheck =
3876            Arc::new(|| Box::pin(async { serde_json::json!({"status": "starting"}) }));
3877        let resp = readyz(check).await.into_response();
3878        // Missing "ready" field defaults to false -> 503
3879        assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
3880    }
3881
3882    // -- normalize_peer_addr_middleware / PeerAddr --
3883
3884    /// Build a test router that reports the request's peer-address
3885    /// extensions as `"<ConnectInfo>|<PeerAddr>"` (empty when absent).
3886    fn peer_probe_router() -> axum::Router {
3887        async fn probe(req: Request<Body>) -> String {
3888            let ci = req
3889                .extensions()
3890                .get::<ConnectInfo<SocketAddr>>()
3891                .map(|c| c.0.to_string())
3892                .unwrap_or_default();
3893            let pa = req
3894                .extensions()
3895                .get::<PeerAddr>()
3896                .map(|p| p.addr.to_string())
3897                .unwrap_or_default();
3898            format!("{ci}|{pa}")
3899        }
3900        axum::Router::new()
3901            .route("/probe", axum::routing::get(probe))
3902            .layer(axum::middleware::from_fn(|req, next| {
3903                normalize_peer_addr_middleware(None, req, next)
3904            }))
3905    }
3906
3907    async fn body_string(resp: axum::response::Response) -> String {
3908        let bytes = resp.into_body().collect().await.unwrap().to_bytes();
3909        String::from_utf8(bytes.to_vec()).unwrap()
3910    }
3911
3912    #[tokio::test]
3913    async fn normalize_preserves_existing_connect_info_and_mirrors_peer_addr() {
3914        // Precedence proof: when both extensions exist with DIFFERENT
3915        // addresses, ConnectInfo<SocketAddr> wins and is never overwritten.
3916        let plain: SocketAddr = "10.0.0.1:1111".parse().unwrap();
3917        let tls: SocketAddr = "10.0.0.2:2222".parse().unwrap();
3918        let req = Request::builder()
3919            .uri("/probe")
3920            .extension(ConnectInfo(plain))
3921            .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
3922            .body(Body::empty())
3923            .unwrap();
3924        let resp = peer_probe_router().oneshot(req).await.unwrap();
3925        assert_eq!(resp.status(), StatusCode::OK);
3926        assert_eq!(body_string(resp).await, format!("{plain}|{plain}"));
3927    }
3928
3929    #[tokio::test]
3930    async fn normalize_inserts_connect_info_and_peer_addr_from_tls() {
3931        let tls: SocketAddr = "192.168.1.7:50443".parse().unwrap();
3932        let req = Request::builder()
3933            .uri("/probe")
3934            .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
3935            .body(Body::empty())
3936            .unwrap();
3937        let resp = peer_probe_router().oneshot(req).await.unwrap();
3938        assert_eq!(resp.status(), StatusCode::OK);
3939        assert_eq!(body_string(resp).await, format!("{tls}|{tls}"));
3940    }
3941
3942    #[tokio::test]
3943    async fn normalize_no_op_without_any_connect_info() {
3944        let req = Request::builder()
3945            .uri("/probe")
3946            .body(Body::empty())
3947            .unwrap();
3948        let resp = peer_probe_router().oneshot(req).await.unwrap();
3949        assert_eq!(resp.status(), StatusCode::OK);
3950        assert_eq!(body_string(resp).await, "|");
3951    }
3952
3953    #[tokio::test]
3954    async fn peer_addr_extractor_rejects_when_absent() {
3955        async fn h(peer: PeerAddr) -> String {
3956            peer.addr.to_string()
3957        }
3958        let app = axum::Router::new().route("/p", axum::routing::get(h));
3959        let req = Request::builder().uri("/p").body(Body::empty()).unwrap();
3960        let resp = app.oneshot(req).await.unwrap();
3961        assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
3962    }
3963
3964    #[tokio::test]
3965    async fn peer_addr_extractor_returns_value_when_present() {
3966        async fn h(peer: PeerAddr) -> String {
3967            peer.addr.to_string()
3968        }
3969        let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap();
3970        let app = axum::Router::new().route("/p", axum::routing::get(h));
3971        let req = Request::builder()
3972            .uri("/p")
3973            .extension(PeerAddr::new(addr))
3974            .body(Body::empty())
3975            .unwrap();
3976        let resp = app.oneshot(req).await.unwrap();
3977        assert_eq!(resp.status(), StatusCode::OK);
3978        assert_eq!(body_string(resp).await, addr.to_string());
3979    }
3980
3981    #[tokio::test]
3982    async fn peer_addr_via_extension_extractor() {
3983        async fn h(axum::Extension(peer): axum::Extension<PeerAddr>) -> String {
3984            peer.addr.to_string()
3985        }
3986        let addr: SocketAddr = "127.0.0.1:4242".parse().unwrap();
3987        let app = axum::Router::new().route("/p", axum::routing::get(h));
3988        let req = Request::builder()
3989            .uri("/p")
3990            .extension(PeerAddr::new(addr))
3991            .body(Body::empty())
3992            .unwrap();
3993        let resp = app.oneshot(req).await.unwrap();
3994        assert_eq!(resp.status(), StatusCode::OK);
3995        assert_eq!(body_string(resp).await, addr.to_string());
3996    }
3997
3998    // -- extra_route_rate_limit_middleware --
3999
4000    /// Probe router with the extra-route limiter installed, mirroring
4001    /// the layer-before-merge wiring in `build_app_router`.
4002    fn limited_router(per_minute: u32) -> axum::Router {
4003        limited_router_with_burst(per_minute, None)
4004    }
4005
4006    /// Probe router with an explicit burst capacity.
4007    fn limited_router_with_burst(per_minute: u32, burst: Option<u32>) -> axum::Router {
4008        limited_router_full(per_minute, burst, &[])
4009    }
4010
4011    /// Probe router with explicit burst and exempt paths. `/limited`
4012    /// and `/exempt` are both registered so exemption interplay can be
4013    /// asserted on one limiter instance.
4014    fn limited_router_full(
4015        per_minute: u32,
4016        burst: Option<u32>,
4017        exempt_paths: &[&str],
4018    ) -> axum::Router {
4019        let limiter = build_extra_route_rate_limiter(per_minute, burst);
4020        let exempt: Arc<std::collections::HashSet<String>> =
4021            Arc::new(exempt_paths.iter().map(|s| (*s).to_owned()).collect());
4022        axum::Router::new()
4023            .route("/limited", axum::routing::get(|| async { "ok" }))
4024            .route("/exempt", axum::routing::get(|| async { "ok" }))
4025            .layer(axum::middleware::from_fn(move |req, next| {
4026                let l = Arc::clone(&limiter);
4027                let e = Arc::clone(&exempt);
4028                extra_route_rate_limit_middleware(l, e, req, next)
4029            }))
4030    }
4031
4032    fn limited_req(ip: &str) -> Request<Body> {
4033        limited_req_to(ip, "/limited")
4034    }
4035
4036    fn limited_req_to(ip: &str, path: &str) -> Request<Body> {
4037        let addr: SocketAddr = format!("{ip}:40000").parse().unwrap();
4038        Request::builder()
4039            .uri(path)
4040            .extension(ConnectInfo(addr))
4041            .body(Body::empty())
4042            .unwrap()
4043    }
4044
4045    #[tokio::test]
4046    async fn extra_route_limiter_denies_over_quota() {
4047        let app = limited_router(2);
4048        for i in 0..2 {
4049            let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4050            assert_eq!(resp.status(), StatusCode::OK, "request {i} should pass");
4051        }
4052        let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4053        assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4054        let body = body_string(resp).await;
4055        assert!(
4056            body.contains("too many requests to application routes"),
4057            "deny body should match the limiter message, got: {body}"
4058        );
4059    }
4060
4061    #[tokio::test]
4062    async fn extra_route_limiter_isolates_keys() {
4063        let app = limited_router(2);
4064        for _ in 0..2 {
4065            let resp = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4066            assert_eq!(resp.status(), StatusCode::OK);
4067        }
4068        let exhausted = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4069        assert_eq!(exhausted.status(), StatusCode::TOO_MANY_REQUESTS);
4070        // A different source IP still has a fresh bucket.
4071        let other = app.clone().oneshot(limited_req("10.3.3.3")).await.unwrap();
4072        assert_eq!(other.status(), StatusCode::OK);
4073    }
4074
4075    #[tokio::test]
4076    async fn extra_route_limiter_fails_open_without_peer() {
4077        let app = limited_router(1);
4078        for i in 0..3 {
4079            let req = Request::builder()
4080                .uri("/limited")
4081                .body(Body::empty())
4082                .unwrap();
4083            let resp = app.clone().oneshot(req).await.unwrap();
4084            assert_eq!(
4085                resp.status(),
4086                StatusCode::OK,
4087                "request {i} should fail open"
4088            );
4089        }
4090    }
4091
4092    #[tokio::test]
4093    async fn extra_route_limiter_extracts_tls_conn_info() {
4094        let app = limited_router(2);
4095        let mk = || {
4096            let addr: SocketAddr = "192.168.9.9:55555".parse().unwrap();
4097            Request::builder()
4098                .uri("/limited")
4099                .extension(ConnectInfo(TlsConnInfo::new(addr, None)))
4100                .body(Body::empty())
4101                .unwrap()
4102        };
4103        for _ in 0..2 {
4104            assert_eq!(
4105                app.clone().oneshot(mk()).await.unwrap().status(),
4106                StatusCode::OK
4107            );
4108        }
4109        let resp = app.clone().oneshot(mk()).await.unwrap();
4110        assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4111    }
4112
4113    #[tokio::test]
4114    async fn extra_route_limiter_exempt_path_bypasses_quota() {
4115        // rate=1: a single non-exempt request exhausts the bucket, yet
4116        // repeated exempt-path requests all pass and consume no budget.
4117        let app = limited_router_full(1, None, &["/exempt"]);
4118        for i in 0..5 {
4119            let resp = app
4120                .clone()
4121                .oneshot(limited_req_to("10.6.6.6", "/exempt"))
4122                .await
4123                .unwrap();
4124            assert_eq!(resp.status(), StatusCode::OK, "exempt request {i}");
4125        }
4126        // Budget untouched by exempt traffic: first limited request OK…
4127        let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
4128        assert_eq!(resp.status(), StatusCode::OK);
4129        // …second is denied (exemption did not leak onto /limited).
4130        let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
4131        assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4132    }
4133
4134    #[tokio::test]
4135    async fn extra_route_limiter_exemption_is_raw_exact_match() {
4136        // Trailing-slash and case variants are NOT exempt (fail-closed:
4137        // a mismatch keeps the request limited, never the reverse).
4138        let app = limited_router_full(1, None, &["/exempt"]);
4139        let ok = app
4140            .clone()
4141            .oneshot(limited_req_to("10.7.7.7", "/exempt/"))
4142            .await
4143            .unwrap();
4144        assert_eq!(
4145            ok.status(),
4146            StatusCode::NOT_FOUND,
4147            "variant path routes 404"
4148        );
4149        // The variant consumed limiter budget (it was not exempt):
4150        let denied = app
4151            .clone()
4152            .oneshot(limited_req_to("10.7.7.7", "/limited"))
4153            .await
4154            .unwrap();
4155        assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4156    }
4157
4158    #[cfg(feature = "metrics")]
4159    #[tokio::test]
4160    async fn extra_route_limiter_deny_increments_counter_exempt_does_not() {
4161        let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
4162        let app = limited_router_full(1, None, &["/exempt"]);
4163        let mk = |path: &str| {
4164            let addr: SocketAddr = "10.8.8.8:40000".parse().unwrap();
4165            Request::builder()
4166                .uri(path)
4167                .extension(ConnectInfo(addr))
4168                .extension(Arc::clone(&metrics))
4169                .body(Body::empty())
4170                .unwrap()
4171        };
4172        let counter = || {
4173            metrics
4174                .rate_limited_total
4175                .with_label_values(&["extra_route"])
4176                .get()
4177        };
4178        // Exempt traffic: no budget, no counter.
4179        for _ in 0..3 {
4180            assert_eq!(
4181                app.clone().oneshot(mk("/exempt")).await.unwrap().status(),
4182                StatusCode::OK
4183            );
4184        }
4185        assert_eq!(counter(), 0, "exempt requests must not count as denies");
4186        // Exhaust then deny: counter increments exactly on the deny.
4187        assert_eq!(
4188            app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4189            StatusCode::OK
4190        );
4191        assert_eq!(counter(), 0);
4192        assert_eq!(
4193            app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4194            StatusCode::TOO_MANY_REQUESTS
4195        );
4196        assert_eq!(counter(), 1, "deny must increment the extra_route label");
4197    }
4198
4199    #[test]
4200    fn validate_rejects_exempt_paths_without_base_knob() {
4201        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4202            .with_extra_route_rate_limit_exempt_paths(["/ok"]);
4203        let err = cfg.validate().expect_err("exempt paths without rate limit");
4204        assert!(err.to_string().contains("requires extra_route_rate_limit"));
4205    }
4206
4207    #[test]
4208    fn validate_rejects_malformed_exempt_paths() {
4209        for bad in ["", "no-slash"] {
4210            let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4211                .with_extra_route_rate_limit(10)
4212                .with_extra_route_rate_limit_exempt_paths([bad]);
4213            let err = cfg.validate().expect_err("malformed exempt path");
4214            assert!(
4215                err.to_string()
4216                    .contains("must be non-empty and start with '/'"),
4217                "entry {bad:?}: {err}"
4218            );
4219        }
4220    }
4221
4222    #[test]
4223    fn validate_accepts_wellformed_exempt_paths() {
4224        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4225            .with_extra_route_rate_limit(10)
4226            .with_extra_route_rate_limit_exempt_paths(["/.well-known/oauth-authorization-server"]);
4227        assert!(cfg.validate().is_ok());
4228    }
4229
4230    #[test]
4231    fn validate_rejects_zero_extra_route_rate_limit() {
4232        let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4233            .with_extra_route_rate_limit(0);
4234        let err = cfg.validate().expect_err("zero extra route rate limit");
4235        assert!(err.to_string().contains("extra_route_rate_limit"));
4236    }
4237
4238    #[tokio::test]
4239    async fn extra_route_limiter_burst_allows_initial_spike() {
4240        let app = limited_router_with_burst(1, Some(3));
4241        for i in 0..3 {
4242            let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4243            assert_eq!(resp.status(), StatusCode::OK, "burst request {i}");
4244        }
4245        let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4246        assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4247    }
4248
4249    #[tokio::test]
4250    async fn extra_route_limiter_deny_sets_retry_after() {
4251        let app = limited_router(1);
4252        let ok = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4253        assert_eq!(ok.status(), StatusCode::OK);
4254        let denied = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4255        assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4256        let retry_after = denied
4257            .headers()
4258            .get(header::RETRY_AFTER)
4259            .expect("Retry-After present")
4260            .to_str()
4261            .unwrap()
4262            .parse::<u64>()
4263            .unwrap();
4264        assert!(retry_after >= 1, "delta-seconds must be >= 1");
4265    }
4266
4267    #[test]
4268    fn validate_rejects_zero_burst_knobs() {
4269        let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4270            .with_tool_rate_limit(10)
4271            .with_tool_rate_limit_burst(0)
4272            .validate()
4273            .expect_err("zero tool burst");
4274        assert!(err.to_string().contains("tool_rate_limit_burst"));
4275
4276        let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4277            .with_extra_route_rate_limit(10)
4278            .with_extra_route_rate_limit_burst(0)
4279            .validate()
4280            .expect_err("zero extra route burst");
4281        assert!(err.to_string().contains("extra_route_rate_limit_burst"));
4282    }
4283
4284    #[test]
4285    fn validate_rejects_orphan_burst_knobs() {
4286        let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4287            .with_tool_rate_limit_burst(5)
4288            .validate()
4289            .expect_err("orphan tool burst");
4290        assert!(err.to_string().contains("requires tool_rate_limit"));
4291
4292        let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4293            .with_extra_route_rate_limit_burst(5)
4294            .validate()
4295            .expect_err("orphan extra route burst");
4296        assert!(err.to_string().contains("requires extra_route_rate_limit"));
4297    }
4298
4299    #[test]
4300    fn validate_rejects_zero_auth_bursts() {
4301        let auth = AuthConfig::with_keys(vec![])
4302            .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_burst(0));
4303        let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4304            .with_auth(auth)
4305            .validate()
4306            .expect_err("zero auth burst");
4307        assert!(err.to_string().contains("rate_limit.burst"));
4308
4309        let auth = AuthConfig::with_keys(vec![])
4310            .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(0));
4311        let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4312            .with_auth(auth)
4313            .validate()
4314            .expect_err("zero pre-auth burst");
4315        assert!(err.to_string().contains("pre_auth_burst"));
4316    }
4317
4318    /// `pre_auth_burst` without `pre_auth_max_per_minute` is LEGAL: the
4319    /// pre-auth base rate always resolves (max_attempts_per_minute x 10).
4320    #[test]
4321    fn validate_accepts_pre_auth_burst_without_explicit_pre_auth_rate() {
4322        let auth = AuthConfig::with_keys(vec![])
4323            .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(50));
4324        let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4325        assert!(cfg.validate().is_ok(), "pre_auth_burst has no orphan rule");
4326    }
4327
4328    // -- trusted-forwarder mode (ClientIp / ForwardedHeaderMode) --
4329
4330    fn forward_resolver(trusted: &[&str], mode: ForwardedHeaderMode) -> Arc<ForwardResolver> {
4331        Arc::new(ForwardResolver {
4332            trusted: trusted.iter().map(|s| s.parse().unwrap()).collect(),
4333            mode,
4334        })
4335    }
4336
4337    /// Probe router reporting `"<PeerAddr ip>|<ClientIp>"`.
4338    fn forwarded_probe_router(resolver: Option<Arc<ForwardResolver>>) -> axum::Router {
4339        async fn probe(req: Request<Body>) -> String {
4340            let pa = req
4341                .extensions()
4342                .get::<PeerAddr>()
4343                .map(|p| p.addr.ip().to_string())
4344                .unwrap_or_default();
4345            let ci = req
4346                .extensions()
4347                .get::<ClientIp>()
4348                .map(|c| c.ip.to_string())
4349                .unwrap_or_default();
4350            format!("{pa}|{ci}")
4351        }
4352        axum::Router::new()
4353            .route("/probe", axum::routing::get(probe))
4354            .layer(axum::middleware::from_fn(move |req, next| {
4355                let r = resolver.clone();
4356                normalize_peer_addr_middleware(r, req, next)
4357            }))
4358    }
4359
4360    fn probe_req(peer: &str, header: Option<(&str, &str)>) -> Request<Body> {
4361        let addr: SocketAddr = peer.parse().unwrap();
4362        let mut builder = Request::builder()
4363            .uri("/probe")
4364            .extension(ConnectInfo(addr));
4365        if let Some((name, value)) = header {
4366            builder = builder.header(name, value);
4367        }
4368        builder.body(Body::empty()).unwrap()
4369    }
4370
4371    #[tokio::test]
4372    async fn client_ip_equals_direct_without_resolver() {
4373        let app = forwarded_probe_router(None);
4374        let resp = app
4375            .oneshot(probe_req(
4376                "10.1.2.3:4444",
4377                Some(("x-forwarded-for", "203.0.113.7")),
4378            ))
4379            .await
4380            .unwrap();
4381        assert_eq!(
4382            body_string(resp).await,
4383            "10.1.2.3|10.1.2.3",
4384            "feature off: header ignored, ClientIp == direct"
4385        );
4386    }
4387
4388    #[tokio::test]
4389    async fn client_ip_resolved_for_trusted_peer() {
4390        let app = forwarded_probe_router(Some(forward_resolver(
4391            &["10.0.0.0/8"],
4392            ForwardedHeaderMode::XForwardedFor,
4393        )));
4394        let resp = app
4395            .oneshot(probe_req(
4396                "10.0.0.1:9999",
4397                Some(("x-forwarded-for", "203.0.113.7")),
4398            ))
4399            .await
4400            .unwrap();
4401        assert_eq!(
4402            body_string(resp).await,
4403            "10.0.0.1|203.0.113.7",
4404            "PeerAddr stays direct while ClientIp resolves"
4405        );
4406    }
4407
4408    #[tokio::test]
4409    async fn client_ip_falls_back_to_direct_on_malformed_header() {
4410        let app = forwarded_probe_router(Some(forward_resolver(
4411            &["10.0.0.0/8"],
4412            ForwardedHeaderMode::XForwardedFor,
4413        )));
4414        let resp = app
4415            .oneshot(probe_req(
4416                "10.0.0.1:9999",
4417                Some(("x-forwarded-for", "not-an-ip")),
4418            ))
4419            .await
4420            .unwrap();
4421        assert_eq!(
4422            body_string(resp).await,
4423            "10.0.0.1|10.0.0.1",
4424            "malformed chain falls back to the direct peer"
4425        );
4426    }
4427
4428    #[test]
4429    fn forwarded_header_mode_deserializes_kebab_case() {
4430        #[derive(serde::Deserialize)]
4431        struct Wrapper {
4432            mode: ForwardedHeaderMode,
4433        }
4434        let w: Wrapper = toml::from_str(r#"mode = "x-forwarded-for""#).unwrap();
4435        assert_eq!(w.mode, ForwardedHeaderMode::XForwardedFor);
4436        let w: Wrapper = toml::from_str(r#"mode = "forwarded""#).unwrap();
4437        assert_eq!(w.mode, ForwardedHeaderMode::Forwarded);
4438        assert!(
4439            toml::from_str::<Wrapper>(r#"mode = "XForwardedFor""#).is_err(),
4440            "PascalCase wire value must be rejected"
4441        );
4442    }
4443
4444    #[test]
4445    fn validate_rejects_bad_trusted_proxy_entry() {
4446        let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4447            .with_trusted_proxies(["not-a-cidr"]);
4448        let err = cfg.validate().expect_err("bad CIDR");
4449        assert!(err.to_string().contains("trusted_proxies"));
4450    }
4451
4452    #[test]
4453    fn validate_rejects_zero_prefix_trusted_proxy() {
4454        for entry in ["0.0.0.0/0", "::/0"] {
4455            let cfg =
4456                McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([entry]);
4457            let err = cfg.validate().expect_err("zero-prefix CIDR");
4458            assert!(
4459                err.to_string().contains("prefix length 0"),
4460                "entry {entry}: {err}"
4461            );
4462        }
4463    }
4464
4465    #[test]
4466    fn validate_accepts_cidr_and_bare_ip_proxy_entries() {
4467        let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([
4468            "10.0.0.0/8",
4469            "192.0.2.1",
4470            "2001:db8::1",
4471        ]);
4472        assert!(cfg.validate().is_ok(), "CIDRs and bare IPs are accepted");
4473    }
4474
4475    #[test]
4476    fn validate_rejects_forwarded_header_without_proxies() {
4477        let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4478            .with_forwarded_header(ForwardedHeaderMode::Forwarded);
4479        let err = cfg.validate().expect_err("mode without proxies");
4480        assert!(err.to_string().contains("requires trusted_proxies"));
4481    }
4482
4483    // -- origin_check_middleware --
4484
4485    /// Build a test router with origin check middleware and a simple handler.
4486    fn origin_router(origins: Vec<String>, log_request_headers: bool) -> axum::Router {
4487        let allowed: Arc<[String]> = Arc::from(origins);
4488        axum::Router::new()
4489            .route("/test", axum::routing::get(|| async { "ok" }))
4490            .layer(axum::middleware::from_fn(move |req, next| {
4491                let a = Arc::clone(&allowed);
4492                origin_check_middleware(a, log_request_headers, req, next)
4493            }))
4494    }
4495
4496    #[tokio::test]
4497    async fn origin_allowed_passes() {
4498        let app = origin_router(vec!["http://localhost:3000".into()], false);
4499        let req = Request::builder()
4500            .uri("/test")
4501            .header(header::ORIGIN, "http://localhost:3000")
4502            .body(Body::empty())
4503            .unwrap();
4504        let resp = app.oneshot(req).await.unwrap();
4505        assert_eq!(resp.status(), StatusCode::OK);
4506    }
4507
4508    #[tokio::test]
4509    async fn origin_rejected_returns_403() {
4510        let app = origin_router(vec!["http://localhost:3000".into()], false);
4511        let req = Request::builder()
4512            .uri("/test")
4513            .header(header::ORIGIN, "http://evil.com")
4514            .body(Body::empty())
4515            .unwrap();
4516        let resp = app.oneshot(req).await.unwrap();
4517        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
4518    }
4519
4520    #[tokio::test]
4521    async fn no_origin_header_passes() {
4522        let app = origin_router(vec!["http://localhost:3000".into()], false);
4523        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4524        let resp = app.oneshot(req).await.unwrap();
4525        assert_eq!(resp.status(), StatusCode::OK);
4526    }
4527
4528    #[tokio::test]
4529    async fn empty_allowlist_rejects_any_origin() {
4530        let app = origin_router(vec![], false);
4531        let req = Request::builder()
4532            .uri("/test")
4533            .header(header::ORIGIN, "http://anything.com")
4534            .body(Body::empty())
4535            .unwrap();
4536        let resp = app.oneshot(req).await.unwrap();
4537        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
4538    }
4539
4540    #[tokio::test]
4541    async fn empty_allowlist_passes_without_origin() {
4542        let app = origin_router(vec![], false);
4543        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4544        let resp = app.oneshot(req).await.unwrap();
4545        assert_eq!(resp.status(), StatusCode::OK);
4546    }
4547
4548    #[test]
4549    fn format_request_headers_redacts_sensitive_values() {
4550        let mut headers = axum::http::HeaderMap::new();
4551        headers.insert("authorization", "Bearer secret-token".parse().unwrap());
4552        headers.insert("cookie", "sid=abc".parse().unwrap());
4553        headers.insert("x-request-id", "req-123".parse().unwrap());
4554
4555        let out = format_request_headers_for_log(&headers);
4556        assert!(out.contains("authorization: [REDACTED]"));
4557        assert!(out.contains("cookie: [REDACTED]"));
4558        assert!(out.contains("x-request-id: req-123"));
4559        assert!(!out.contains("secret-token"));
4560    }
4561
4562    // -- security_headers_middleware --
4563
4564    fn security_router(is_tls: bool) -> axum::Router {
4565        security_router_with(is_tls, SecurityHeadersConfig::default())
4566    }
4567
4568    fn security_router_with(is_tls: bool, cfg: SecurityHeadersConfig) -> axum::Router {
4569        let cfg = Arc::new(cfg);
4570        axum::Router::new()
4571            .route("/test", axum::routing::get(|| async { "ok" }))
4572            .layer(axum::middleware::from_fn(move |req, next| {
4573                let c = Arc::clone(&cfg);
4574                security_headers_middleware(is_tls, c, req, next)
4575            }))
4576    }
4577
4578    #[tokio::test]
4579    async fn security_headers_set_on_response() {
4580        let app = security_router(false);
4581        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4582        let resp = app.oneshot(req).await.unwrap();
4583        assert_eq!(resp.status(), StatusCode::OK);
4584
4585        let h = resp.headers();
4586        assert_eq!(h.get("x-content-type-options").unwrap(), "nosniff");
4587        assert_eq!(h.get("x-frame-options").unwrap(), "deny");
4588        assert_eq!(h.get("cache-control").unwrap(), "no-store, max-age=0");
4589        assert_eq!(h.get("referrer-policy").unwrap(), "no-referrer");
4590        assert_eq!(h.get("cross-origin-opener-policy").unwrap(), "same-origin");
4591        assert_eq!(
4592            h.get("cross-origin-resource-policy").unwrap(),
4593            "same-origin"
4594        );
4595        assert_eq!(
4596            h.get("cross-origin-embedder-policy").unwrap(),
4597            "require-corp"
4598        );
4599        assert_eq!(h.get("x-permitted-cross-domain-policies").unwrap(), "none");
4600        assert!(
4601            h.get("permissions-policy")
4602                .unwrap()
4603                .to_str()
4604                .unwrap()
4605                .contains("camera=()"),
4606            "permissions-policy must restrict browser features"
4607        );
4608        assert_eq!(
4609            h.get("content-security-policy").unwrap(),
4610            "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
4611        );
4612        assert_eq!(h.get("x-dns-prefetch-control").unwrap(), "off");
4613        // No HSTS when TLS is off.
4614        assert!(h.get("strict-transport-security").is_none());
4615    }
4616
4617    #[tokio::test]
4618    async fn hsts_set_when_tls_enabled() {
4619        let app = security_router(true);
4620        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4621        let resp = app.oneshot(req).await.unwrap();
4622
4623        let hsts = resp.headers().get("strict-transport-security").unwrap();
4624        assert!(
4625            hsts.to_str().unwrap().contains("max-age=63072000"),
4626            "HSTS must set 2-year max-age"
4627        );
4628    }
4629
4630    #[tokio::test]
4631    async fn default_csp_matches_guideline() {
4632        let app = security_router(false);
4633        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4634        let resp = app.oneshot(req).await.unwrap();
4635        assert_eq!(
4636            resp.headers().get("content-security-policy").unwrap(),
4637            "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
4638        );
4639    }
4640
4641    #[tokio::test]
4642    async fn operator_csp_override_still_wins() {
4643        let cfg = SecurityHeadersConfig {
4644            content_security_policy: Some("default-src 'self'".into()),
4645            ..SecurityHeadersConfig::default()
4646        };
4647        let app = security_router_with(false, cfg);
4648        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4649        let resp = app.oneshot(req).await.unwrap();
4650        assert_eq!(
4651            resp.headers().get("content-security-policy").unwrap(),
4652            "default-src 'self'"
4653        );
4654    }
4655
4656    // -- SecurityHeadersConfig validation + override semantics --
4657
4658    /// Build a minimal config with a custom SecurityHeadersConfig and
4659    /// drive it through `check()`. Returns the result so individual
4660    /// tests can assert on success or specific error messages.
4661    fn check_with_security_headers(headers: SecurityHeadersConfig) -> Result<(), McpxError> {
4662        let cfg =
4663            McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_security_headers(headers);
4664        cfg.check()
4665    }
4666
4667    #[test]
4668    fn security_headers_config_default_validates() {
4669        check_with_security_headers(SecurityHeadersConfig::default())
4670            .expect("default SecurityHeadersConfig must validate");
4671    }
4672
4673    #[test]
4674    fn security_headers_config_validate_accepts_empty_string() {
4675        // All twelve fields explicitly set to "" -> omit-everything mode.
4676        let h = SecurityHeadersConfig {
4677            x_content_type_options: Some(String::new()),
4678            x_frame_options: Some(String::new()),
4679            cache_control: Some(String::new()),
4680            referrer_policy: Some(String::new()),
4681            cross_origin_opener_policy: Some(String::new()),
4682            cross_origin_resource_policy: Some(String::new()),
4683            cross_origin_embedder_policy: Some(String::new()),
4684            permissions_policy: Some(String::new()),
4685            x_permitted_cross_domain_policies: Some(String::new()),
4686            content_security_policy: Some(String::new()),
4687            x_dns_prefetch_control: Some(String::new()),
4688            strict_transport_security: Some(String::new()),
4689        };
4690        check_with_security_headers(h).expect("Some(\"\") on every field must validate (omit-all)");
4691    }
4692
4693    #[test]
4694    fn security_headers_config_validate_rejects_bad_value() {
4695        // 0x07 (BEL) is not a valid HTTP header value char.
4696        let h = SecurityHeadersConfig {
4697            referrer_policy: Some("\u{0007}".into()),
4698            ..SecurityHeadersConfig::default()
4699        };
4700        let err = check_with_security_headers(h)
4701            .expect_err("control char in referrer_policy must reject");
4702        let msg = err.to_string();
4703        assert!(
4704            msg.contains("referrer_policy"),
4705            "error must name the offending field, got: {msg}"
4706        );
4707    }
4708
4709    #[test]
4710    fn security_headers_config_validate_rejects_hsts_preload() {
4711        let h = SecurityHeadersConfig {
4712            strict_transport_security: Some("max-age=63072000; includeSubDomains; preload".into()),
4713            ..SecurityHeadersConfig::default()
4714        };
4715        let err = check_with_security_headers(h).expect_err("HSTS with preload must reject");
4716        let msg = err.to_string();
4717        assert!(
4718            msg.contains("strict_transport_security"),
4719            "error must name the field, got: {msg}"
4720        );
4721        assert!(
4722            msg.to_lowercase().contains("preload"),
4723            "error must mention `preload`, got: {msg}"
4724        );
4725    }
4726
4727    #[test]
4728    fn security_headers_config_validate_rejects_hsts_preload_uppercase() {
4729        // Case-insensitive match.
4730        let h = SecurityHeadersConfig {
4731            strict_transport_security: Some("max-age=600; PRELOAD".into()),
4732            ..SecurityHeadersConfig::default()
4733        };
4734        check_with_security_headers(h).expect_err("HSTS preload check must be case-insensitive");
4735    }
4736
4737    #[tokio::test]
4738    async fn security_headers_override_honored() {
4739        // Override X-Frame-Options to SAMEORIGIN.
4740        let h = SecurityHeadersConfig {
4741            x_frame_options: Some("SAMEORIGIN".into()),
4742            ..SecurityHeadersConfig::default()
4743        };
4744        let app = security_router_with(false, h);
4745        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4746        let resp = app.oneshot(req).await.unwrap();
4747        assert_eq!(resp.status(), StatusCode::OK);
4748
4749        let xfo = resp.headers().get("x-frame-options").unwrap();
4750        assert_eq!(xfo, "SAMEORIGIN");
4751    }
4752
4753    #[tokio::test]
4754    async fn security_headers_empty_string_omits() {
4755        // Empty string on referrer-policy -> header absent.
4756        let h = SecurityHeadersConfig {
4757            referrer_policy: Some(String::new()),
4758            ..SecurityHeadersConfig::default()
4759        };
4760        let app = security_router_with(false, h);
4761        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4762        let resp = app.oneshot(req).await.unwrap();
4763        assert_eq!(resp.status(), StatusCode::OK);
4764
4765        assert!(
4766            resp.headers().get("referrer-policy").is_none(),
4767            "Some(\"\") must omit the header"
4768        );
4769        // Other defaults should still be present.
4770        assert_eq!(
4771            resp.headers().get("x-content-type-options").unwrap(),
4772            "nosniff"
4773        );
4774    }
4775
4776    #[tokio::test]
4777    async fn security_headers_hsts_only_when_tls() {
4778        // HSTS override is irrelevant when TLS is off.
4779        let h = SecurityHeadersConfig {
4780            strict_transport_security: Some("max-age=600".into()),
4781            ..SecurityHeadersConfig::default()
4782        };
4783        let app = security_router_with(false, h);
4784        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4785        let resp = app.oneshot(req).await.unwrap();
4786        assert!(
4787            resp.headers().get("strict-transport-security").is_none(),
4788            "HSTS must remain absent on plaintext deployments even with override"
4789        );
4790    }
4791
4792    // -- oauth_token_cache_headers_middleware --
4793
4794    #[cfg(feature = "oauth")]
4795    #[tokio::test]
4796    async fn oauth_token_cache_headers_set_pragma_and_vary() {
4797        let app = axum::Router::new()
4798            .route("/token", axum::routing::post(|| async { "{}" }))
4799            .layer(axum::middleware::from_fn(
4800                oauth_token_cache_headers_middleware,
4801            ));
4802        let req = Request::builder()
4803            .method("POST")
4804            .uri("/token")
4805            .body(Body::from("{}"))
4806            .unwrap();
4807        let resp = app.oneshot(req).await.unwrap();
4808        assert_eq!(resp.status(), StatusCode::OK);
4809
4810        let h = resp.headers();
4811        assert_eq!(
4812            h.get("pragma").unwrap(),
4813            "no-cache",
4814            "RFC 6749 §5.1: token responses must set Pragma: no-cache"
4815        );
4816        let vary_values: Vec<String> = h
4817            .get_all("vary")
4818            .iter()
4819            .filter_map(|v| v.to_str().ok().map(str::to_owned))
4820            .collect();
4821        assert!(
4822            vary_values
4823                .iter()
4824                .any(|v| v.eq_ignore_ascii_case("Authorization")),
4825            "RFC 6750 §5.4: Vary must include Authorization, got {vary_values:?}"
4826        );
4827    }
4828
4829    #[cfg(feature = "oauth")]
4830    #[tokio::test]
4831    async fn oauth_token_cache_headers_preserve_existing_vary() {
4832        // Simulates a handler/layer that already set `Vary: Accept-Encoding`
4833        // (e.g. compression). Our middleware must APPEND, not REPLACE.
4834        let app = axum::Router::new()
4835            .route(
4836                "/token",
4837                axum::routing::post(|| async {
4838                    axum::response::Response::builder()
4839                        .header("vary", "Accept-Encoding")
4840                        .body(Body::from("{}"))
4841                        .unwrap()
4842                }),
4843            )
4844            .layer(axum::middleware::from_fn(
4845                oauth_token_cache_headers_middleware,
4846            ));
4847        let req = Request::builder()
4848            .method("POST")
4849            .uri("/token")
4850            .body(Body::empty())
4851            .unwrap();
4852        let resp = app.oneshot(req).await.unwrap();
4853
4854        let vary: Vec<String> = resp
4855            .headers()
4856            .get_all("vary")
4857            .iter()
4858            .filter_map(|v| v.to_str().ok().map(str::to_owned))
4859            .collect();
4860        assert!(
4861            vary.iter().any(|v| v.contains("Accept-Encoding")),
4862            "must preserve pre-existing Vary value, got {vary:?}"
4863        );
4864        assert!(
4865            vary.iter().any(|v| v.contains("Authorization")),
4866            "must append Authorization to Vary, got {vary:?}"
4867        );
4868    }
4869
4870    // -- version endpoint --
4871
4872    #[test]
4873    fn version_omits_build_fingerprint_by_default() {
4874        let v = version_payload("my-server", "1.2.3", false);
4875        assert_eq!(v["name"], "my-server");
4876        assert_eq!(v["version"], "1.2.3");
4877        assert!(v["mcpx_version"].is_string());
4878        assert!(
4879            v.get("build_git_sha").is_none(),
4880            "build sha must be hidden by default"
4881        );
4882        assert!(v.get("build_timestamp").is_none());
4883        assert!(v.get("rust_version").is_none());
4884    }
4885
4886    #[test]
4887    fn version_exposes_all_when_enabled() {
4888        let v = version_payload("my-server", "1.2.3", true);
4889        assert!(v["build_git_sha"].is_string());
4890        assert!(v["build_timestamp"].is_string());
4891        assert!(v["rust_version"].is_string());
4892        assert!(v["mcpx_version"].is_string());
4893    }
4894
4895    // -- concurrency limit layer --
4896
4897    #[tokio::test]
4898    async fn concurrency_limit_layer_composes_and_serves() {
4899        // We only assert the layer stack compiles and a single request
4900        // below the cap still succeeds. True back-pressure behaviour
4901        // requires a live HTTP server and is covered by integration tests.
4902        let app = axum::Router::new()
4903            .route("/ok", axum::routing::get(|| async { "ok" }))
4904            .layer(
4905                tower::ServiceBuilder::new()
4906                    .layer(axum::error_handling::HandleErrorLayer::new(
4907                        |_err: tower::BoxError| async { StatusCode::SERVICE_UNAVAILABLE },
4908                    ))
4909                    .layer(tower::load_shed::LoadShedLayer::new())
4910                    .layer(tower::limit::ConcurrencyLimitLayer::new(4)),
4911            );
4912        let resp = app
4913            .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap())
4914            .await
4915            .unwrap();
4916        assert_eq!(resp.status(), StatusCode::OK);
4917    }
4918
4919    // -- compression layer --
4920
4921    #[tokio::test]
4922    async fn compression_layer_gzip_encodes_response() {
4923        use tower_http::compression::Predicate as _;
4924
4925        let big_body = "a".repeat(4096);
4926        let app = axum::Router::new()
4927            .route(
4928                "/big",
4929                axum::routing::get(move || {
4930                    let body = big_body.clone();
4931                    async move { body }
4932                }),
4933            )
4934            .layer(
4935                tower_http::compression::CompressionLayer::new()
4936                    .gzip(true)
4937                    .br(true)
4938                    .compress_when(
4939                        tower_http::compression::DefaultPredicate::new()
4940                            .and(tower_http::compression::predicate::SizeAbove::new(1024)),
4941                    ),
4942            );
4943
4944        let req = Request::builder()
4945            .uri("/big")
4946            .header(header::ACCEPT_ENCODING, "gzip")
4947            .body(Body::empty())
4948            .unwrap();
4949        let resp = app.oneshot(req).await.unwrap();
4950        assert_eq!(resp.status(), StatusCode::OK);
4951        assert_eq!(
4952            resp.headers().get(header::CONTENT_ENCODING).unwrap(),
4953            "gzip"
4954        );
4955    }
4956
4957    // -- TlsListener handshake timeout --
4958
4959    #[tokio::test]
4960    async fn tls_handshake_timeout_reaps_idle_connections() {
4961        use tokio::io::AsyncReadExt as _;
4962
4963        let _ = rustls::crypto::ring::default_provider().install_default();
4964
4965        // Self-signed cert material on disk (TlsListener::new takes paths).
4966        let key = rcgen::KeyPair::generate().expect("generate key");
4967        let cert = rcgen::CertificateParams::new(vec!["localhost".to_owned()])
4968            .expect("cert params")
4969            .self_signed(&key)
4970            .expect("self-signed cert");
4971        let dir = std::env::temp_dir().join(format!(
4972            "rmcp-server-kit-hs-timeout-{}",
4973            std::time::SystemTime::now()
4974                .duration_since(std::time::UNIX_EPOCH)
4975                .expect("clock after epoch")
4976                .as_nanos()
4977        ));
4978        tokio::fs::create_dir_all(&dir).await.expect("temp dir");
4979        let cert_path = dir.join("server.crt");
4980        let key_path = dir.join("server.key");
4981        tokio::fs::write(&cert_path, cert.pem())
4982            .await
4983            .expect("write cert");
4984        tokio::fs::write(&key_path, key.serialize_pem())
4985            .await
4986            .expect("write key");
4987
4988        let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
4989        let tls = TlsListener::new(
4990            listener,
4991            &cert_path,
4992            &key_path,
4993            None,
4994            None,
4995            Duration::from_millis(200),
4996            8, // custom concurrency cap: proves the plumbing end-to-end
4997        )
4998        .expect("tls listener");
4999        let addr = axum::serve::Listener::local_addr(&tls).expect("local addr");
5000
5001        // Connect and send NOTHING: the handshake worker must time out
5002        // after 200ms and drop the stream, which the client observes as
5003        // EOF or a reset well within the 2s deadline.
5004        let mut idle = tokio::net::TcpStream::connect(addr).await.expect("connect");
5005        let mut buf = [0_u8; 16];
5006        let read = tokio::time::timeout(Duration::from_secs(2), idle.read(&mut buf))
5007            .await
5008            .expect("server must reap the idle handshake within its timeout");
5009        match read {
5010            Ok(0) | Err(_) => {} // EOF or reset: connection was dropped.
5011            Ok(n) => panic!("unexpected {n} bytes from server during reaped handshake"),
5012        }
5013
5014        drop(tls);
5015    }
5016
5017    // -- M5: OWASP security headers reach early / fallback responses --
5018
5019    fn assert_owasp_headers(resp: &axum::response::Response, ctx: &str) {
5020        let h = resp.headers();
5021        assert!(
5022            h.contains_key("x-content-type-options"),
5023            "{ctx}: missing X-Content-Type-Options"
5024        );
5025        assert!(
5026            h.contains_key("x-frame-options"),
5027            "{ctx}: missing X-Frame-Options"
5028        );
5029        assert!(
5030            h.contains_key("strict-transport-security"),
5031            "{ctx}: missing Strict-Transport-Security"
5032        );
5033        assert!(
5034            h.contains_key(header::CONTENT_SECURITY_POLICY),
5035            "{ctx}: missing Content-Security-Policy"
5036        );
5037    }
5038
5039    fn m5_router(configure: impl FnOnce(&mut McpServerConfig)) -> axum::Router {
5040        #[derive(Clone)]
5041        struct H;
5042        impl ServerHandler for H {}
5043        // TLS paths make `is_tls` true so HSTS is emitted. The paths are never
5044        // read: these tests drive only the axum router via `oneshot`, not the
5045        // TLS listener.
5046        let mut config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0")
5047            .with_allowed_origins(["http://good.example"])
5048            .with_tls("unused.crt", "unused.key");
5049        configure(&mut config);
5050        let (router, _params) = build_app_router(config, || H).expect("build_app_router");
5051        router
5052    }
5053
5054    #[tokio::test]
5055    async fn headers_on_rejected_origin_403() {
5056        let app = m5_router(|_| {});
5057        let req = Request::builder()
5058            .uri("/healthz")
5059            .header(header::ORIGIN, "http://evil.example")
5060            .body(Body::empty())
5061            .unwrap();
5062        let resp = app.oneshot(req).await.unwrap();
5063        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5064        assert_owasp_headers(&resp, "origin-403");
5065    }
5066
5067    #[tokio::test]
5068    async fn headers_on_cors_preflight() {
5069        let app = m5_router(|_| {});
5070        let req = Request::builder()
5071            .method(axum::http::Method::OPTIONS)
5072            .uri("/mcp")
5073            .header(header::ORIGIN, "http://good.example")
5074            .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
5075            .body(Body::empty())
5076            .unwrap();
5077        let resp = app.oneshot(req).await.unwrap();
5078        assert_owasp_headers(&resp, "cors-preflight");
5079    }
5080
5081    #[tokio::test]
5082    async fn headers_on_404_fallback() {
5083        let app = m5_router(|_| {});
5084        let req = Request::builder()
5085            .uri("/no-such-route")
5086            .body(Body::empty())
5087            .unwrap();
5088        let resp = app.oneshot(req).await.unwrap();
5089        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
5090        assert_owasp_headers(&resp, "404-fallback");
5091    }
5092
5093    #[tokio::test]
5094    async fn headers_on_overload_503() {
5095        // A zero-permit concurrency cap sheds every request, so a single
5096        // oneshot deterministically surfaces the overload 503.
5097        let app = m5_router(|c| c.max_concurrent_requests = Some(0));
5098        let req = Request::builder()
5099            .uri("/healthz")
5100            .body(Body::empty())
5101            .unwrap();
5102        let resp = app.oneshot(req).await.unwrap();
5103        assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
5104        assert_owasp_headers(&resp, "overload-503");
5105    }
5106
5107    // -- M6: OAuth proxy admin endpoints enforce the admin role --
5108
5109    #[cfg(feature = "oauth")]
5110    fn m6_auth_state() -> (Arc<AuthState>, String, String) {
5111        let (admin_token, admin_hash) = crate::auth::generate_api_key().unwrap();
5112        let (viewer_token, viewer_hash) = crate::auth::generate_api_key().unwrap();
5113        let state = Arc::new(AuthState {
5114            api_keys: ArcSwap::from_pointee(vec![
5115                crate::auth::ApiKeyEntry::new("admin-key", admin_hash, "admin"),
5116                crate::auth::ApiKeyEntry::new("viewer-key", viewer_hash, "viewer"),
5117            ]),
5118            rate_limiter: None,
5119            pre_auth_limiter: None,
5120            jwks_cache: None,
5121            seen_identities: crate::auth::SeenIdentitySet::new(),
5122            counters: crate::auth::AuthCounters::default(),
5123        });
5124        (state, admin_token, viewer_token)
5125    }
5126
5127    #[cfg(feature = "oauth")]
5128    fn m6_admin_router(state: &Arc<AuthState>) -> axum::Router {
5129        let proxy = crate::oauth::OAuthProxyConfig::builder(
5130            "https://idp.example/authorize",
5131            "https://idp.example/token",
5132            "client",
5133        )
5134        .introspection_url("http://127.0.0.1:1/introspect")
5135        .revocation_url("http://127.0.0.1:1/revoke")
5136        .expose_admin_endpoints(true)
5137        .require_auth_on_admin_endpoints(true)
5138        .build();
5139        let http = crate::oauth::OauthHttpClient::new().expect("oauth http client");
5140        build_oauth_admin_router(&proxy, http, Some(state), "admin").expect("admin router")
5141    }
5142
5143    #[cfg(feature = "oauth")]
5144    fn m6_req(path: &str, token: &str) -> Request<Body> {
5145        Request::builder()
5146            .method(axum::http::Method::POST)
5147            .uri(path)
5148            .header(header::AUTHORIZATION, format!("Bearer {token}"))
5149            .body(Body::from("token=abc"))
5150            .unwrap()
5151    }
5152
5153    #[cfg(feature = "oauth")]
5154    #[tokio::test]
5155    async fn oauth_proxy_admin_requires_admin_role() {
5156        let (state, _admin, viewer) = m6_auth_state();
5157        for path in ["/introspect", "/revoke"] {
5158            let app = m6_admin_router(&state);
5159            let resp = app.oneshot(m6_req(path, &viewer)).await.unwrap();
5160            assert_eq!(
5161                resp.status(),
5162                StatusCode::FORBIDDEN,
5163                "an authenticated viewer must be rejected with 403 on {path}"
5164            );
5165        }
5166    }
5167
5168    #[cfg(feature = "oauth")]
5169    #[tokio::test]
5170    async fn oauth_proxy_admin_allows_admin_role() {
5171        let (state, admin, _viewer) = m6_auth_state();
5172        for path in ["/introspect", "/revoke"] {
5173            let app = m6_admin_router(&state);
5174            let resp = app.oneshot(m6_req(path, &admin)).await.unwrap();
5175            // The admin identity clears both the auth and role gates; the
5176            // downstream introspection call then fails closed (no upstream),
5177            // so the only guarantee asserted is that it is neither 401 nor 403.
5178            assert_ne!(
5179                resp.status(),
5180                StatusCode::FORBIDDEN,
5181                "an authenticated admin must pass the role gate on {path}"
5182            );
5183            assert_ne!(
5184                resp.status(),
5185                StatusCode::UNAUTHORIZED,
5186                "an authenticated admin must pass the auth gate on {path}"
5187            );
5188        }
5189    }
5190
5191    // -- F3 regression: unbounded Prometheus label cardinality --
5192    //
5193    // `metrics_middleware` runs outside the auth layer, so it observes
5194    // unauthenticated traffic. Labelling with the raw URI path and raw HTTP
5195    // method let any client mint a permanent time series per request, growing
5196    // in-process metric state until OOM. Both labels must now come from a
5197    // closed set.
5198    #[cfg(feature = "metrics")]
5199    mod metrics_labels_bounded {
5200        use super::*;
5201
5202        fn labels_for(method: &str, uri: &str) -> (&'static str, String) {
5203            let req = Request::builder()
5204                .method(method)
5205                .uri(uri)
5206                .body(Body::empty())
5207                .unwrap();
5208            metrics_labels(&req)
5209        }
5210
5211        #[test]
5212        fn many_unmatched_paths_collapse_to_one_label() {
5213            let mut seen = std::collections::HashSet::new();
5214            for i in 0..500 {
5215                let (_, path) = labels_for("GET", &format!("/nonexistent-{i}"));
5216                seen.insert(path);
5217            }
5218            assert_eq!(
5219                seen.len(),
5220                1,
5221                "unmatched paths must collapse to a single label, got {seen:?}"
5222            );
5223            assert!(seen.contains("<unmatched>"));
5224        }
5225
5226        #[test]
5227        fn nested_mcp_paths_collapse_to_the_mount_point() {
5228            let mut seen = std::collections::HashSet::new();
5229            for i in 0..200 {
5230                let (_, path) = labels_for("POST", &format!("/mcp/{i}"));
5231                seen.insert(path);
5232            }
5233            let (_, root) = labels_for("POST", "/mcp");
5234            seen.insert(root);
5235            assert_eq!(
5236                seen.len(),
5237                1,
5238                "nested /mcp paths must collapse to one label, got {seen:?}"
5239            );
5240            assert!(seen.contains("/mcp"));
5241        }
5242
5243        #[test]
5244        fn unusual_methods_collapse_to_one_bucket() {
5245            let mut seen = std::collections::HashSet::new();
5246            for verb in ["FROBNICATE", "WIBBLE", "QUUX", "M-SEARCH"] {
5247                let (method, _) = labels_for(verb, "/healthz");
5248                seen.insert(method);
5249            }
5250            assert_eq!(seen, std::collections::HashSet::from(["OTHER"]));
5251        }
5252
5253        #[test]
5254        fn known_methods_keep_their_identity() {
5255            for verb in ["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"] {
5256                let (method, _) = labels_for(verb, "/healthz");
5257                assert_eq!(method, verb);
5258            }
5259        }
5260
5261        #[test]
5262        fn raw_path_never_leaks_into_a_label() {
5263            let (_, path) = labels_for("GET", "/secret-token-abc123");
5264            assert!(
5265                !path.contains("secret-token"),
5266                "raw request path must never become a label value: {path}"
5267            );
5268        }
5269    }
5270}