Skip to main content

rmcp_server_kit/
transport.rs

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