Skip to main content

rmcp_server_kit/
transport.rs

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