1use std::{
2 future::Future,
3 net::{IpAddr, SocketAddr},
4 num::NonZeroUsize,
5 path::{Path, PathBuf},
6 pin::Pin,
7 sync::Arc,
8 time::Duration,
9};
10
11use arc_swap::ArcSwap;
12use axum::{
13 body::Body,
14 extract::{ConnectInfo, Request},
15 middleware::Next,
16 response::IntoResponse,
17};
18use rmcp::{
19 ServerHandler,
20 transport::streamable_http_server::{
21 StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
22 },
23};
24use rustls::RootCertStore;
25use tokio::{
26 net::TcpListener,
27 sync::{Semaphore, mpsc},
28};
29use tokio_util::sync::CancellationToken;
30
31use crate::{
32 auth::{
33 AuthConfig, AuthIdentity, AuthState, MtlsConfig, TlsConnInfo, auth_middleware,
34 build_rate_limiter, extract_mtls_identity,
35 },
36 bounded_limiter::{BoundedKeyedLimiter, BoundedLimiterDeny, KeyEvictionPolicy},
37 error::RmcpServerKitError,
38 mtls_revocation::{self, CrlSet, DynamicClientCertVerifier},
39 rbac::{RbacPolicy, ToolRateLimiter, build_tool_rate_limiter_with_policy, rbac_middleware},
40};
41
42#[allow(
46 clippy::needless_pass_by_value,
47 reason = "consumed at .map_err(anyhow_to_startup) call sites; by-value matches the closure shape"
48)]
49fn anyhow_to_startup(e: anyhow::Error) -> RmcpServerKitError {
50 RmcpServerKitError::Startup(format!("{e:#}"))
51}
52
53#[allow(
59 clippy::needless_pass_by_value,
60 reason = "consumed at .map_err(|e| io_to_startup(...)) call sites; by-value matches the closure shape"
61)]
62fn io_to_startup(op: &str, e: std::io::Error) -> RmcpServerKitError {
63 RmcpServerKitError::Startup(format!("{op}: {e}"))
64}
65
66pub type ReadinessCheck =
71 Arc<dyn Fn() -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> + Send + Sync>;
72
73#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
118#[non_exhaustive]
119pub struct PeerAddr {
120 pub addr: SocketAddr,
122}
123
124impl PeerAddr {
125 #[must_use]
128 pub(crate) const fn new(addr: SocketAddr) -> Self {
129 Self { addr }
130 }
131}
132
133impl<S: Send + Sync> axum::extract::FromRequestParts<S> for PeerAddr {
142 type Rejection = (axum::http::StatusCode, &'static str);
143
144 #[allow(
145 clippy::unused_async_trait_impl,
146 reason = "async is mandated by the axum FromRequestParts trait signature; this impl only reads a request extension synchronously"
147 )]
148 async fn from_request_parts(
149 parts: &mut axum::http::request::Parts,
150 _state: &S,
151 ) -> Result<Self, Self::Rejection> {
152 parts.extensions.get::<Self>().copied().ok_or((
153 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
154 "peer address unavailable: not running under rmcp-server-kit serve()",
155 ))
156 }
157}
158
159#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
182#[non_exhaustive]
183pub struct ClientIp {
184 pub ip: IpAddr,
186}
187
188impl ClientIp {
189 #[must_use]
192 pub(crate) const fn new(ip: IpAddr) -> Self {
193 Self { ip }
194 }
195}
196
197#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Deserialize)]
202#[serde(rename_all = "kebab-case")]
203#[non_exhaustive]
204pub enum ForwardedHeaderMode {
205 XForwardedFor,
207 Forwarded,
209}
210
211struct ForwardResolver {
214 trusted: Vec<ipnet::IpNet>,
215 mode: ForwardedHeaderMode,
216 max_scanned_entries: usize,
217}
218
219#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
240#[serde(default)]
241#[serde(deny_unknown_fields)]
242#[non_exhaustive]
243pub struct SecurityHeadersConfig {
244 pub x_content_type_options: Option<String>,
246 pub x_frame_options: Option<String>,
248 pub cache_control: Option<String>,
250 pub referrer_policy: Option<String>,
252 pub cross_origin_opener_policy: Option<String>,
254 pub cross_origin_resource_policy: Option<String>,
256 pub cross_origin_embedder_policy: Option<String>,
258 pub permissions_policy: Option<String>,
261 pub x_permitted_cross_domain_policies: Option<String>,
263 pub content_security_policy: Option<String>,
266 pub x_dns_prefetch_control: Option<String>,
268 pub strict_transport_security: Option<String>,
273}
274
275#[allow(
277 missing_debug_implementations,
278 reason = "contains callback/trait objects that don't impl Debug"
279)]
280#[allow(
281 clippy::struct_excessive_bools,
282 reason = "server configuration naturally has many boolean feature flags"
283)]
284#[non_exhaustive]
285pub struct McpServerConfig {
286 #[deprecated(
288 since = "0.13.0",
289 note = "use McpServerConfig::new() / with_bind_addr(); direct field access will become pub(crate) in a future major release"
290 )]
291 pub bind_addr: String,
292 #[deprecated(
294 since = "0.13.0",
295 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
296 )]
297 pub name: String,
298 #[deprecated(
300 since = "0.13.0",
301 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
302 )]
303 pub version: String,
304 #[deprecated(
306 since = "0.13.0",
307 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
308 )]
309 pub tls_cert_path: Option<PathBuf>,
310 #[deprecated(
312 since = "0.13.0",
313 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
314 )]
315 pub tls_key_path: Option<PathBuf>,
316 #[deprecated(
319 since = "0.13.0",
320 note = "use McpServerConfig::with_auth(); direct field access will become pub(crate) in a future major release"
321 )]
322 pub auth: Option<AuthConfig>,
323 #[deprecated(
326 since = "0.13.0",
327 note = "use McpServerConfig::with_rbac(); direct field access will become pub(crate) in a future major release"
328 )]
329 pub rbac: Option<Arc<RbacPolicy>>,
330 #[deprecated(
336 since = "0.13.0",
337 note = "use McpServerConfig::with_allowed_origins(); direct field access will become pub(crate) in a future major release"
338 )]
339 pub allowed_origins: Vec<String>,
340 #[deprecated(
343 since = "0.13.0",
344 note = "use McpServerConfig::with_tool_rate_limit(); direct field access will become pub(crate) in a future major release"
345 )]
346 pub tool_rate_limit: Option<u32>,
347 #[deprecated(
353 since = "1.12.0",
354 note = "use McpServerConfig::with_tool_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
355 )]
356 pub tool_rate_limit_burst: Option<u32>,
357 #[deprecated(
370 since = "1.11.0",
371 note = "use McpServerConfig::with_extra_route_rate_limit(); direct field access will become pub(crate) in a future major release"
372 )]
373 pub extra_route_rate_limit: Option<u32>,
374 #[deprecated(
381 since = "1.12.0",
382 note = "use McpServerConfig::with_extra_route_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
383 )]
384 pub extra_route_rate_limit_burst: Option<u32>,
385 #[deprecated(
398 since = "1.14.0",
399 note = "use McpServerConfig::with_extra_route_rate_limit_exempt_paths(); direct field access will become pub(crate) in a future major release"
400 )]
401 pub extra_route_rate_limit_exempt_paths: Vec<String>,
402
403 pub key_eviction_policy: KeyEvictionPolicy,
405
406 pub trusted_forwarder_max_entries: usize,
413 #[deprecated(
421 since = "1.13.0",
422 note = "use McpServerConfig::with_trusted_proxies(); direct field access will become pub(crate) in a future major release"
423 )]
424 pub trusted_proxies: Vec<String>,
425 #[deprecated(
430 since = "1.13.0",
431 note = "use McpServerConfig::with_forwarded_header(); direct field access will become pub(crate) in a future major release"
432 )]
433 pub forwarded_header: Option<ForwardedHeaderMode>,
434 #[deprecated(
437 since = "0.13.0",
438 note = "use McpServerConfig::with_readiness_check(); direct field access will become pub(crate) in a future major release"
439 )]
440 pub readiness_check: Option<ReadinessCheck>,
441 #[deprecated(
444 since = "0.13.0",
445 note = "use McpServerConfig::with_max_request_body(); direct field access will become pub(crate) in a future major release"
446 )]
447 pub max_request_body: usize,
448 #[deprecated(
451 since = "0.13.0",
452 note = "use McpServerConfig::with_request_timeout(); direct field access will become pub(crate) in a future major release"
453 )]
454 pub request_timeout: Duration,
455 #[deprecated(
458 since = "0.13.0",
459 note = "use McpServerConfig::with_shutdown_timeout(); direct field access will become pub(crate) in a future major release"
460 )]
461 pub shutdown_timeout: Duration,
462 #[deprecated(
465 since = "0.13.0",
466 note = "use McpServerConfig::with_session_idle_timeout(); direct field access will become pub(crate) in a future major release"
467 )]
468 pub session_idle_timeout: Duration,
469 #[deprecated(
472 since = "0.13.0",
473 note = "use McpServerConfig::with_sse_keep_alive(); direct field access will become pub(crate) in a future major release"
474 )]
475 pub sse_keep_alive: Duration,
476 #[deprecated(
480 since = "0.13.0",
481 note = "use McpServerConfig::with_reload_callback(); direct field access will become pub(crate) in a future major release"
482 )]
483 pub on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
484 #[deprecated(
491 since = "0.13.0",
492 note = "use McpServerConfig::with_extra_router(); direct field access will become pub(crate) in a future major release"
493 )]
494 pub extra_router: Option<axum::Router>,
495 #[deprecated(
500 since = "0.13.0",
501 note = "use McpServerConfig::with_public_url(); direct field access will become pub(crate) in a future major release"
502 )]
503 pub public_url: Option<String>,
504 #[deprecated(
507 since = "0.13.0",
508 note = "use McpServerConfig::enable_request_header_logging(); direct field access will become pub(crate) in a future major release"
509 )]
510 pub log_request_headers: bool,
511 pub expose_build_metadata: bool,
518 #[deprecated(
521 since = "0.13.0",
522 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
523 )]
524 pub compression_enabled: bool,
525 #[deprecated(
528 since = "0.13.0",
529 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
530 )]
531 pub compression_min_size: u16,
532 #[deprecated(
536 since = "0.13.0",
537 note = "use McpServerConfig::with_max_concurrent_requests(); direct field access will become pub(crate) in a future major release"
538 )]
539 pub max_concurrent_requests: Option<usize>,
540 #[deprecated(
543 since = "0.13.0",
544 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
545 )]
546 pub admin_enabled: bool,
547 #[deprecated(
549 since = "0.13.0",
550 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
551 )]
552 pub admin_role: String,
553 #[cfg(feature = "metrics")]
556 #[deprecated(
557 since = "0.13.0",
558 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
559 )]
560 pub metrics_enabled: bool,
561 #[cfg(feature = "metrics")]
563 #[deprecated(
564 since = "0.13.0",
565 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
566 )]
567 pub metrics_bind: String,
568 #[deprecated(
572 since = "1.5.0",
573 note = "use McpServerConfig::with_security_headers(); direct field access will become pub(crate) in a future major release"
574 )]
575 pub security_headers: SecurityHeadersConfig,
576 #[deprecated(
582 since = "1.9.0",
583 note = "use McpServerConfig::with_tls_handshake_timeout(); direct field access will become pub(crate) in a future major release"
584 )]
585 pub tls_handshake_timeout: Duration,
586 #[deprecated(
593 since = "1.9.0",
594 note = "use McpServerConfig::with_max_concurrent_tls_handshakes(); direct field access will become pub(crate) in a future major release"
595 )]
596 pub max_concurrent_tls_handshakes: usize,
597}
598
599#[allow(
657 missing_debug_implementations,
658 reason = "wraps T which may not implement Debug; manual impl below avoids leaking inner contents into logs"
659)]
660pub struct Validated<T>(T);
661
662impl<T> std::fmt::Debug for Validated<T> {
663 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
664 f.debug_struct("Validated").finish_non_exhaustive()
665 }
666}
667
668impl<T> Validated<T> {
669 #[must_use]
671 pub fn as_inner(&self) -> &T {
672 &self.0
673 }
674
675 #[must_use]
680 pub fn into_inner(self) -> T {
681 self.0
682 }
683}
684
685#[allow(
686 deprecated,
687 reason = "internal builders/validators legitimately read/write the deprecated `pub` fields they were designed to manage"
688)]
689impl McpServerConfig {
690 #[must_use]
698 pub fn new(
699 bind_addr: impl Into<String>,
700 name: impl Into<String>,
701 version: impl Into<String>,
702 ) -> Self {
703 Self {
704 bind_addr: bind_addr.into(),
705 name: name.into(),
706 version: version.into(),
707 tls_cert_path: None,
708 tls_key_path: None,
709 auth: None,
710 rbac: None,
711 allowed_origins: Vec::new(),
712 tool_rate_limit: None,
713 readiness_check: None,
714 max_request_body: 1024 * 1024,
715 request_timeout: Duration::from_mins(2),
716 shutdown_timeout: Duration::from_secs(30),
717 session_idle_timeout: Duration::from_mins(20),
718 sse_keep_alive: Duration::from_secs(15),
719 on_reload_ready: None,
720 extra_router: None,
721 public_url: None,
722 log_request_headers: false,
723 expose_build_metadata: false,
724 compression_enabled: false,
725 compression_min_size: 1024,
726 max_concurrent_requests: None,
727 admin_enabled: false,
728 admin_role: "admin".to_owned(),
729 #[cfg(feature = "metrics")]
730 metrics_enabled: false,
731 #[cfg(feature = "metrics")]
732 metrics_bind: "127.0.0.1:9090".into(),
733 security_headers: SecurityHeadersConfig::default(),
734 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
735 max_concurrent_tls_handshakes: DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES,
736 extra_route_rate_limit: None,
737 tool_rate_limit_burst: None,
738 extra_route_rate_limit_burst: None,
739 extra_route_rate_limit_exempt_paths: Vec::new(),
740 key_eviction_policy: KeyEvictionPolicy::default(),
741 trusted_forwarder_max_entries: crate::forwarded::MAX_SCANNED_ENTRIES,
742 trusted_proxies: Vec::new(),
743 forwarded_header: None,
744 }
745 }
746
747 #[must_use]
757 pub fn with_auth(mut self, auth: AuthConfig) -> Self {
758 self.auth = Some(auth);
759 self
760 }
761
762 #[must_use]
767 pub fn with_security_headers(mut self, headers: SecurityHeadersConfig) -> Self {
768 self.security_headers = headers;
769 self
770 }
771
772 #[must_use]
776 pub fn with_bind_addr(mut self, addr: impl Into<String>) -> Self {
777 self.bind_addr = addr.into();
778 self
779 }
780
781 #[must_use]
784 pub fn with_rbac(mut self, rbac: Arc<RbacPolicy>) -> Self {
785 self.rbac = Some(rbac);
786 self
787 }
788
789 #[must_use]
793 pub fn with_tls(mut self, cert_path: impl Into<PathBuf>, key_path: impl Into<PathBuf>) -> Self {
794 self.tls_cert_path = Some(cert_path.into());
795 self.tls_key_path = Some(key_path.into());
796 self
797 }
798
799 #[must_use]
803 pub fn with_public_url(mut self, url: impl Into<String>) -> Self {
804 self.public_url = Some(url.into());
805 self
806 }
807
808 #[must_use]
812 pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
813 where
814 I: IntoIterator<Item = S>,
815 S: Into<String>,
816 {
817 self.allowed_origins = origins.into_iter().map(Into::into).collect();
818 self
819 }
820
821 #[must_use]
854 pub fn with_extra_router(mut self, router: axum::Router) -> Self {
855 self.extra_router = Some(router);
856 self
857 }
858
859 #[must_use]
865 pub const fn with_trusted_forwarder_max_entries(mut self, max_entries: usize) -> Self {
866 self.trusted_forwarder_max_entries = max_entries;
867 self
868 }
869
870 #[must_use]
873 pub fn with_readiness_check(mut self, check: ReadinessCheck) -> Self {
874 self.readiness_check = Some(check);
875 self
876 }
877
878 #[must_use]
881 pub fn with_max_request_body(mut self, bytes: usize) -> Self {
882 self.max_request_body = bytes;
883 self
884 }
885
886 #[must_use]
888 pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
889 self.request_timeout = timeout;
890 self
891 }
892
893 #[must_use]
895 pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
896 self.shutdown_timeout = timeout;
897 self
898 }
899
900 #[must_use]
902 pub fn with_session_idle_timeout(mut self, timeout: Duration) -> Self {
903 self.session_idle_timeout = timeout;
904 self
905 }
906
907 #[must_use]
909 pub fn with_sse_keep_alive(mut self, interval: Duration) -> Self {
910 self.sse_keep_alive = interval;
911 self
912 }
913
914 #[must_use]
918 pub fn with_max_concurrent_requests(mut self, limit: usize) -> Self {
919 self.max_concurrent_requests = Some(limit);
920 self
921 }
922
923 #[must_use]
931 pub fn with_tls_handshake_timeout(mut self, timeout: Duration) -> Self {
932 self.tls_handshake_timeout = timeout;
933 self
934 }
935
936 #[must_use]
945 pub fn with_max_concurrent_tls_handshakes(mut self, limit: usize) -> Self {
946 self.max_concurrent_tls_handshakes = limit;
947 self
948 }
949
950 #[must_use]
953 pub fn with_tool_rate_limit(mut self, per_minute: u32) -> Self {
954 self.tool_rate_limit = Some(per_minute);
955 self
956 }
957
958 #[must_use]
969 pub fn with_extra_route_rate_limit(mut self, per_minute: u32) -> Self {
970 self.extra_route_rate_limit = Some(per_minute);
971 self
972 }
973
974 #[must_use]
979 pub fn with_tool_rate_limit_burst(mut self, burst: u32) -> Self {
980 self.tool_rate_limit_burst = Some(burst);
981 self
982 }
983
984 #[must_use]
990 pub fn with_extra_route_rate_limit_burst(mut self, burst: u32) -> Self {
991 self.extra_route_rate_limit_burst = Some(burst);
992 self
993 }
994
995 #[must_use]
1015 pub fn with_extra_route_rate_limit_exempt_paths<I, S>(mut self, paths: I) -> Self
1016 where
1017 I: IntoIterator<Item = S>,
1018 S: Into<String>,
1019 {
1020 self.extra_route_rate_limit_exempt_paths = paths.into_iter().map(Into::into).collect();
1021 self
1022 }
1023
1024 #[must_use]
1026 pub const fn with_key_eviction_policy(mut self, policy: KeyEvictionPolicy) -> Self {
1027 self.key_eviction_policy = policy;
1028 self
1029 }
1030
1031 #[must_use]
1043 pub fn with_trusted_proxies<I, S>(mut self, proxies: I) -> Self
1044 where
1045 I: IntoIterator<Item = S>,
1046 S: Into<String>,
1047 {
1048 self.trusted_proxies = proxies.into_iter().map(Into::into).collect();
1049 self
1050 }
1051
1052 #[must_use]
1057 pub fn with_forwarded_header(mut self, mode: ForwardedHeaderMode) -> Self {
1058 self.forwarded_header = Some(mode);
1059 self
1060 }
1061
1062 #[must_use]
1066 pub fn with_reload_callback<F>(mut self, callback: F) -> Self
1067 where
1068 F: FnOnce(ReloadHandle) + Send + 'static,
1069 {
1070 self.on_reload_ready = Some(Box::new(callback));
1071 self
1072 }
1073
1074 #[must_use]
1078 pub fn enable_compression(mut self, min_size: u16) -> Self {
1079 self.compression_enabled = true;
1080 self.compression_min_size = min_size;
1081 self
1082 }
1083
1084 #[must_use]
1089 pub fn enable_admin(mut self, role: impl Into<String>) -> Self {
1090 self.admin_enabled = true;
1091 self.admin_role = role.into();
1092 self
1093 }
1094
1095 #[must_use]
1098 pub fn enable_request_header_logging(mut self) -> Self {
1099 self.log_request_headers = true;
1100 self
1101 }
1102
1103 #[must_use]
1108 pub fn expose_build_metadata(mut self) -> Self {
1109 self.expose_build_metadata = true;
1110 self
1111 }
1112
1113 #[cfg(feature = "metrics")]
1116 #[must_use]
1117 pub fn with_metrics(mut self, bind: impl Into<String>) -> Self {
1118 self.metrics_enabled = true;
1119 self.metrics_bind = bind.into();
1120 self
1121 }
1122
1123 pub fn validate(self) -> Result<Validated<Self>, RmcpServerKitError> {
1156 self.check()?;
1157 Ok(Validated(self))
1158 }
1159
1160 fn check_burst_knobs(&self) -> Result<(), RmcpServerKitError> {
1167 if self.tool_rate_limit_burst == Some(0) {
1168 return Err(RmcpServerKitError::Config(
1169 "tool_rate_limit_burst must be greater than zero".into(),
1170 ));
1171 }
1172 if self.extra_route_rate_limit_burst == Some(0) {
1173 return Err(RmcpServerKitError::Config(
1174 "extra_route_rate_limit_burst must be greater than zero".into(),
1175 ));
1176 }
1177 if self.trusted_forwarder_max_entries == 0
1178 || self.trusted_forwarder_max_entries
1179 > crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES
1180 {
1181 return Err(RmcpServerKitError::Config(format!(
1182 "trusted_forwarder_max_entries must be in 1..={}, got {}",
1183 crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES,
1184 self.trusted_forwarder_max_entries
1185 )));
1186 }
1187 if self.tool_rate_limit_burst.is_some() && self.tool_rate_limit.is_none() {
1188 return Err(RmcpServerKitError::Config(
1189 "tool_rate_limit_burst requires tool_rate_limit to be set".into(),
1190 ));
1191 }
1192 if self.extra_route_rate_limit_burst.is_some() && self.extra_route_rate_limit.is_none() {
1193 return Err(RmcpServerKitError::Config(
1194 "extra_route_rate_limit_burst requires extra_route_rate_limit to be set".into(),
1195 ));
1196 }
1197 if !self.extra_route_rate_limit_exempt_paths.is_empty()
1198 && self.extra_route_rate_limit.is_none()
1199 {
1200 return Err(RmcpServerKitError::Config(
1201 "extra_route_rate_limit_exempt_paths requires extra_route_rate_limit to be set"
1202 .into(),
1203 ));
1204 }
1205 for path in &self.extra_route_rate_limit_exempt_paths {
1206 if path.is_empty() || !path.starts_with('/') {
1207 return Err(RmcpServerKitError::Config(format!(
1208 "extra_route_rate_limit_exempt_paths entries must be non-empty and start with '/': {path:?}"
1209 )));
1210 }
1211 }
1212 if let Some(rl) = self.auth.as_ref().and_then(|a| a.rate_limit.as_ref()) {
1213 if rl.burst == Some(0) {
1214 return Err(RmcpServerKitError::Config(
1215 "auth rate_limit.burst must be greater than zero".into(),
1216 ));
1217 }
1218 if rl.pre_auth_burst == Some(0) {
1219 return Err(RmcpServerKitError::Config(
1220 "auth rate_limit.pre_auth_burst must be greater than zero".into(),
1221 ));
1222 }
1223 }
1224 Ok(())
1225 }
1226
1227 fn check_trusted_forwarder(&self) -> Result<(), RmcpServerKitError> {
1232 for entry in &self.trusted_proxies {
1233 validate_trusted_proxy_entry(entry).map_err(RmcpServerKitError::Config)?;
1234 }
1235 if self.forwarded_header.is_some() && self.trusted_proxies.is_empty() {
1236 return Err(RmcpServerKitError::Config(
1237 "forwarded_header requires trusted_proxies to be nonempty".into(),
1238 ));
1239 }
1240 Ok(())
1241 }
1242
1243 fn check(&self) -> Result<(), RmcpServerKitError> {
1247 if let Err(violation) = crate::config::check_shared_config_invariants(
1262 self.admin_enabled,
1263 self.auth.as_ref().is_some_and(|a| a.enabled),
1264 self.tls_cert_path.is_some(),
1265 self.tls_key_path.is_some(),
1266 self.auth.as_ref().is_some_and(|a| a.mtls.is_some()),
1267 ) {
1268 return Err(RmcpServerKitError::Config(
1269 match violation {
1270 crate::config::SharedConfigViolation::AdminRequiresAuth => {
1271 "admin_enabled=true requires auth to be configured and enabled"
1272 }
1273 crate::config::SharedConfigViolation::TlsCertWithoutKey => {
1274 "tls_cert_path is set but tls_key_path is missing"
1275 }
1276 crate::config::SharedConfigViolation::TlsKeyWithoutCert => {
1277 "tls_key_path is set but tls_cert_path is missing"
1278 }
1279 crate::config::SharedConfigViolation::MtlsRequiresTls => {
1280 "auth.mtls requires TLS: set both tls_cert_path and tls_key_path \
1281 (mTLS client certificates cannot be verified on a plaintext listener)"
1282 }
1283 }
1284 .into(),
1285 ));
1286 }
1287
1288 if self.bind_addr.parse::<SocketAddr>().is_err() {
1290 return Err(RmcpServerKitError::Config(format!(
1291 "bind_addr {:?} is not a valid socket address (expected e.g. 127.0.0.1:8080)",
1292 self.bind_addr
1293 )));
1294 }
1295
1296 if let Some(ref url) = self.public_url
1298 && !(url.starts_with("http://") || url.starts_with("https://"))
1299 {
1300 return Err(RmcpServerKitError::Config(format!(
1301 "public_url {url:?} must start with http:// or https://"
1302 )));
1303 }
1304
1305 for origin in &self.allowed_origins {
1307 if !(origin.starts_with("http://") || origin.starts_with("https://")) {
1308 return Err(RmcpServerKitError::Config(format!(
1309 "allowed_origins entry {origin:?} must start with http:// or https://"
1310 )));
1311 }
1312 }
1313
1314 if self.max_request_body == 0 {
1316 return Err(RmcpServerKitError::Config(
1317 "max_request_body must be greater than zero".into(),
1318 ));
1319 }
1320
1321 if self.extra_route_rate_limit == Some(0) {
1325 return Err(RmcpServerKitError::Config(
1326 "extra_route_rate_limit must be greater than zero".into(),
1327 ));
1328 }
1329
1330 self.check_burst_knobs()?;
1332
1333 self.check_trusted_forwarder()?;
1335
1336 #[cfg(feature = "oauth")]
1338 if let Some(auth_cfg) = &self.auth
1339 && let Some(oauth_cfg) = &auth_cfg.oauth
1340 {
1341 oauth_cfg.validate()?;
1342 }
1343
1344 validate_security_headers(&self.security_headers)?;
1347
1348 if self.max_concurrent_requests == Some(0) {
1352 return Err(RmcpServerKitError::Config(
1353 "max_concurrent_requests must be greater than zero when set".into(),
1354 ));
1355 }
1356
1357 if let Some(auth_cfg) = &self.auth
1361 && let Some(rl) = &auth_cfg.rate_limit
1362 && rl.max_tracked_keys == 0
1363 {
1364 return Err(RmcpServerKitError::Config(
1365 "auth.rate_limit.max_tracked_keys must be greater than zero".into(),
1366 ));
1367 }
1368
1369 check_auth_capacity_knobs(self.auth.as_ref())?;
1370
1371 if self.tls_handshake_timeout == Duration::ZERO {
1376 return Err(RmcpServerKitError::Config(
1377 "tls_handshake_timeout must be greater than zero".into(),
1378 ));
1379 }
1380
1381 if self.max_concurrent_tls_handshakes == 0 {
1386 return Err(RmcpServerKitError::Config(
1387 "max_concurrent_tls_handshakes must be greater than zero".into(),
1388 ));
1389 }
1390
1391 Ok(())
1392 }
1393}
1394
1395#[allow(
1401 missing_debug_implementations,
1402 reason = "contains Arc<AuthState> with non-Debug fields"
1403)]
1404pub struct ReloadHandle {
1405 auth: Option<Arc<AuthState>>,
1406 rbac: Option<Arc<ArcSwap<RbacPolicy>>>,
1407 crl_set: Option<Arc<CrlSet>>,
1408}
1409
1410impl ReloadHandle {
1411 pub fn reload_auth_keys(&self, keys: Vec<crate::auth::ApiKeyEntry>) {
1413 if let Some(ref auth) = self.auth {
1414 auth.reload_keys(keys);
1415 }
1416 }
1417
1418 pub fn reload_rbac(&self, policy: RbacPolicy) {
1420 if let Some(ref rbac) = self.rbac {
1421 rbac.store(Arc::new(policy));
1422 tracing::info!("RBAC policy reloaded");
1423 }
1424 }
1425
1426 pub async fn refresh_crls(&self) -> Result<(), RmcpServerKitError> {
1436 let Some(ref crl_set) = self.crl_set else {
1437 return Err(RmcpServerKitError::Config(
1438 "CRL refresh requested but mTLS CRL support is not configured".into(),
1439 ));
1440 };
1441
1442 crl_set.force_refresh().await
1443 }
1444}
1445
1446#[allow(
1463 clippy::too_many_lines,
1464 clippy::cognitive_complexity,
1465 reason = "middleware layer order is security-critical and must remain visible at one glance; extracting `&mut Router` helpers would obscure the auth/RBAC/origin/rate-limit ordering"
1466)]
1467struct AppRunParams {
1471 tls_paths: Option<(PathBuf, PathBuf)>,
1473 tls_handshake_timeout: Duration,
1475 max_concurrent_tls_handshakes: usize,
1477 mtls_config: Option<MtlsConfig>,
1479 shutdown_timeout: Duration,
1481 auth_state: Option<Arc<AuthState>>,
1483 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
1485 on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
1487 ct: CancellationToken,
1491 session_ct: CancellationToken,
1501 scheme: &'static str,
1503 name: String,
1505}
1506
1507#[allow(
1517 clippy::cognitive_complexity,
1518 reason = "router assembly is intrinsically sequential; splitting harms readability"
1519)]
1520#[allow(
1521 deprecated,
1522 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1523)]
1524fn build_app_router<H, F>(
1525 mut config: McpServerConfig,
1526 handler_factory: F,
1527) -> anyhow::Result<(axum::Router, AppRunParams)>
1528where
1529 H: ServerHandler + 'static,
1530 F: Fn() -> H + Send + Sync + Clone + 'static,
1531{
1532 let ct = CancellationToken::new();
1533 let session_ct = CancellationToken::new();
1534
1535 let allowed_hosts = derive_allowed_hosts(&config.bind_addr, config.public_url.as_deref());
1536 tracing::info!(allowed_hosts = %allowed_hosts.join(", "), "configured Streamable HTTP allowed hosts");
1537
1538 if config.max_concurrent_requests.is_none() {
1539 tracing::warn!(
1540 "max_concurrent_requests is unset: in-flight HTTP requests are unlimited; \
1541 set McpServerConfig::with_max_concurrent_requests or front the server with \
1542 an external concurrency limit"
1543 );
1544 }
1545
1546 let mcp_service = StreamableHttpService::new(
1547 move || Ok(handler_factory()),
1548 {
1549 let mut mgr = LocalSessionManager::default();
1550 mgr.session_config.keep_alive = Some(config.session_idle_timeout);
1551 mgr.into()
1552 },
1553 StreamableHttpServerConfig::default()
1554 .with_allowed_hosts(allowed_hosts)
1555 .with_sse_keep_alive(Some(config.sse_keep_alive))
1556 .with_cancellation_token(session_ct.clone()),
1557 );
1558
1559 let mut mcp_router = axum::Router::new().nest_service("/mcp", mcp_service);
1561
1562 let auth_state: Option<Arc<AuthState>> = match config.auth {
1566 Some(ref auth_config) if auth_config.enabled => {
1567 let rate_limiter = auth_config.rate_limit.as_ref().map(build_rate_limiter);
1568 let pre_auth_limiter = auth_config
1569 .rate_limit
1570 .as_ref()
1571 .map(crate::auth::build_pre_auth_limiter);
1572
1573 #[cfg(feature = "oauth")]
1574 let jwks_cache = auth_config
1575 .oauth
1576 .as_ref()
1577 .map(|c| crate::oauth::JwksCache::new(c).map(Arc::new))
1578 .transpose()
1579 .map_err(|e| std::io::Error::other(format!("JWKS HTTP client: {e}")))?;
1580
1581 Some(Arc::new(AuthState {
1582 api_keys: ArcSwap::new(Arc::new(auth_config.api_keys.clone())),
1583 rate_limiter,
1584 pre_auth_limiter,
1585 #[cfg(feature = "oauth")]
1586 jwks_cache,
1587 seen_identities: crate::auth::SeenIdentitySet::new(),
1588 counters: crate::auth::AuthCounters::default(),
1589 resource_metadata_url: config.public_url.as_ref().map(|url| {
1597 format!(
1598 "{}/.well-known/oauth-protected-resource/mcp",
1599 url.trim_end_matches('/')
1600 )
1601 }),
1602 }))
1603 }
1604 _ => None,
1605 };
1606
1607 let rbac_swap = Arc::new(ArcSwap::new(
1610 config
1611 .rbac
1612 .clone()
1613 .unwrap_or_else(|| Arc::new(RbacPolicy::disabled())),
1614 ));
1615
1616 if config.admin_enabled {
1619 let Some(ref auth_state_ref) = auth_state else {
1620 return Err(anyhow::anyhow!(
1621 "admin_enabled=true requires auth to be configured and enabled"
1622 ));
1623 };
1624 let admin_state = crate::admin::AdminState {
1625 started_at: std::time::Instant::now(),
1626 name: config.name.clone(),
1627 version: config.version.clone(),
1628 auth: Some(Arc::clone(auth_state_ref)),
1629 rbac: Arc::clone(&rbac_swap),
1630 };
1631 let admin_cfg = crate::admin::AdminConfig {
1632 role: config.admin_role.clone(),
1633 };
1634 mcp_router = mcp_router.merge(crate::admin::admin_router(admin_state, &admin_cfg));
1635 tracing::info!(role = %config.admin_role, "/admin/* endpoints enabled");
1636 }
1637
1638 {
1671 let tool_limiter: Option<Arc<ToolRateLimiter>> = config.tool_rate_limit.map(|per_minute| {
1672 build_tool_rate_limiter_with_policy(
1673 per_minute,
1674 config.tool_rate_limit_burst,
1675 config.key_eviction_policy,
1676 )
1677 });
1678
1679 if rbac_swap.load().is_enabled() {
1680 tracing::info!("RBAC enforcement enabled on /mcp");
1681 }
1682 if let Some(limit) = config.tool_rate_limit {
1683 tracing::info!(limit, "tool rate limiting enabled (calls/min per IP)");
1684 }
1685
1686 let rbac_for_mw = Arc::clone(&rbac_swap);
1687 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1688 let p = rbac_for_mw.load_full();
1689 let tl = tool_limiter.clone();
1690 rbac_middleware(p, tl, req, next)
1691 }));
1692 }
1693
1694 if let Some(ref auth_config) = config.auth
1696 && auth_config.enabled
1697 {
1698 let Some(ref state) = auth_state else {
1699 return Err(anyhow::anyhow!("auth state missing despite enabled config"));
1700 };
1701
1702 let methods: Vec<&str> = [
1703 auth_config.mtls.is_some().then_some("mTLS"),
1704 (!auth_config.api_keys.is_empty()).then_some("bearer"),
1705 #[cfg(feature = "oauth")]
1706 auth_config.oauth.is_some().then_some("oauth-jwt"),
1707 ]
1708 .into_iter()
1709 .flatten()
1710 .collect();
1711
1712 tracing::info!(
1713 methods = %methods.join(", "),
1714 api_keys = auth_config.api_keys.len(),
1715 "auth enabled on /mcp"
1716 );
1717
1718 let state_for_mw = Arc::clone(state);
1719 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1720 let s = Arc::clone(&state_for_mw);
1721 auth_middleware(s, req, next)
1722 }));
1723 }
1724
1725 mcp_router = mcp_router.layer(tower_http::timeout::TimeoutLayer::with_status_code(
1728 axum::http::StatusCode::REQUEST_TIMEOUT,
1729 config.request_timeout,
1730 ));
1731
1732 mcp_router = mcp_router.layer(tower_http::limit::RequestBodyLimitLayer::new(
1736 config.max_request_body,
1737 ));
1738
1739 let mut effective_origins = config.allowed_origins.clone();
1746 if effective_origins.is_empty()
1747 && let Some(ref url) = config.public_url
1748 {
1749 if let Some(scheme_end) = url.find("://") {
1754 let scheme_with_sep = url.get(..scheme_end + 3).unwrap_or_default();
1755 let after_scheme = url.get(scheme_end + 3..).unwrap_or_default();
1756 let host_end = after_scheme.find('/').unwrap_or(after_scheme.len());
1757 let host = after_scheme.get(..host_end).unwrap_or_default();
1758 let origin = format!("{scheme_with_sep}{host}");
1759 tracing::info!(
1760 %origin,
1761 "auto-derived allowed origin from public_url"
1762 );
1763 effective_origins.push(origin);
1764 }
1765 }
1766 let allowed_origins: Arc<[String]> = Arc::from(effective_origins);
1767 let cors_origins = Arc::clone(&allowed_origins);
1768 let log_request_headers = config.log_request_headers;
1769
1770 let readyz_route = if let Some(check) = config.readiness_check.take() {
1771 axum::routing::get(move || readyz(Arc::clone(&check)))
1772 } else {
1773 axum::routing::get(healthz)
1774 };
1775
1776 #[allow(
1777 unused_mut,
1778 reason = "the binding is only reassigned when the `oauth` feature adds the \
1779 protected-resource-metadata route below"
1780 )]
1781 let mut router = axum::Router::new()
1782 .route("/healthz", axum::routing::get(healthz))
1783 .route("/readyz", readyz_route)
1784 .route(
1785 "/version",
1786 axum::routing::get({
1787 let payload_bytes: Arc<[u8]> = serialize_version_payload(
1792 &config.name,
1793 &config.version,
1794 config.expose_build_metadata,
1795 );
1796 move || {
1797 let p = Arc::clone(&payload_bytes);
1798 async move {
1799 (
1800 [(axum::http::header::CONTENT_TYPE, "application/json")],
1801 p.to_vec(),
1802 )
1803 }
1804 }
1805 }),
1806 )
1807 .merge(mcp_router);
1808
1809 if let Some(extra) = config.extra_router.take() {
1816 let extra = match config.extra_route_rate_limit {
1817 Some(per_minute) => {
1818 let max_tracked_keys =
1819 NonZeroUsize::new(EXTRA_ROUTE_MAX_TRACKED_KEYS).unwrap_or(NonZeroUsize::MIN);
1820 let limiter = build_extra_route_rate_limiter_with_policy(
1821 per_minute,
1822 config.extra_route_rate_limit_burst,
1823 config.key_eviction_policy,
1824 max_tracked_keys,
1825 );
1826 let exempt: Arc<std::collections::HashSet<String>> = Arc::new(
1827 config
1828 .extra_route_rate_limit_exempt_paths
1829 .iter()
1830 .cloned()
1831 .collect(),
1832 );
1833 tracing::info!(
1834 per_minute,
1835 exempt_paths = exempt.len(),
1836 "extra-route per-IP rate limit enabled"
1837 );
1838 extra.layer(axum::middleware::from_fn(move |req, next| {
1839 let l = Arc::clone(&limiter);
1840 let e = Arc::clone(&exempt);
1841 extra_route_rate_limit_middleware(l, e, req, next)
1842 }))
1843 }
1844 None => extra,
1845 };
1846 router = router.merge(extra);
1847 }
1848
1849 let server_url = derive_server_url(&config);
1856 let resource_url = format!("{server_url}/mcp");
1857
1858 #[cfg(feature = "oauth")]
1859 let prm_metadata = if let Some(ref auth_config) = config.auth
1860 && let Some(ref oauth_config) = auth_config.oauth
1861 {
1862 crate::oauth::protected_resource_metadata(&resource_url, &server_url, oauth_config)
1863 } else {
1864 serde_json::json!({ "resource": resource_url })
1865 };
1866 #[cfg(not(feature = "oauth"))]
1867 let prm_metadata = serde_json::json!({ "resource": resource_url });
1868
1869 let prm_root = prm_metadata.clone();
1875 router = router.route(
1876 "/.well-known/oauth-protected-resource",
1877 axum::routing::get(move || {
1878 let m = prm_root.clone();
1879 async move { axum::Json(m) }
1880 }),
1881 );
1882 router = router.route(
1883 "/.well-known/oauth-protected-resource/mcp",
1884 axum::routing::get(move || {
1885 let m = prm_metadata.clone();
1886 async move { axum::Json(m) }
1887 }),
1888 );
1889
1890 #[cfg(feature = "oauth")]
1895 if let Some(ref auth_config) = config.auth
1896 && let Some(ref oauth_config) = auth_config.oauth
1897 && oauth_config.proxy.is_some()
1898 {
1899 router = install_oauth_proxy_routes(
1900 router,
1901 &server_url,
1902 oauth_config,
1903 auth_state.as_ref(),
1904 config.max_request_body,
1905 &config.admin_role,
1906 )?;
1907 }
1908
1909 if !cors_origins.is_empty() {
1918 let cors = tower_http::cors::CorsLayer::new()
1919 .allow_origin(
1920 cors_origins
1921 .iter()
1922 .filter_map(|o| o.parse::<axum::http::HeaderValue>().ok())
1923 .collect::<Vec<_>>(),
1924 )
1925 .allow_methods([
1926 axum::http::Method::GET,
1927 axum::http::Method::POST,
1928 axum::http::Method::OPTIONS,
1929 ])
1930 .allow_headers([
1931 axum::http::header::CONTENT_TYPE,
1932 axum::http::header::AUTHORIZATION,
1933 ]);
1934 router = router.layer(cors);
1935 }
1936
1937 if config.compression_enabled {
1941 use tower_http::compression::Predicate as _;
1942 let predicate = tower_http::compression::DefaultPredicate::new().and(
1943 tower_http::compression::predicate::SizeAbove::new(u64::from(
1944 config.compression_min_size,
1945 )),
1946 );
1947 router = router.layer(
1948 tower_http::compression::CompressionLayer::new()
1949 .gzip(true)
1950 .br(true)
1951 .compress_when(predicate),
1952 );
1953 tracing::info!(
1954 min_size = config.compression_min_size,
1955 "response compression enabled (gzip, br)"
1956 );
1957 }
1958
1959 if let Some(max) = config.max_concurrent_requests {
1962 let overload_handler = tower::ServiceBuilder::new()
1963 .layer(axum::error_handling::HandleErrorLayer::new(
1964 |_err: tower::BoxError| async {
1965 (
1966 axum::http::StatusCode::SERVICE_UNAVAILABLE,
1967 axum::Json(serde_json::json!({
1968 "error": "overloaded",
1969 "error_description": "server is at capacity, retry later"
1970 })),
1971 )
1972 },
1973 ))
1974 .layer(tower::load_shed::LoadShedLayer::new())
1975 .layer(tower::limit::ConcurrencyLimitLayer::new(max));
1976 router = router.layer(overload_handler);
1977 tracing::info!(max, "global concurrency limit enabled");
1978 }
1979
1980 router = router.fallback(|| async {
1984 (
1985 axum::http::StatusCode::NOT_FOUND,
1986 axum::Json(serde_json::json!({
1987 "error": "not_found",
1988 "error_description": "The requested endpoint does not exist"
1989 })),
1990 )
1991 });
1992
1993 #[cfg(feature = "metrics")]
1995 if config.metrics_enabled {
1996 let metrics = Arc::new(
1997 crate::metrics::McpMetrics::new().map_err(|e| anyhow::anyhow!("metrics init: {e}"))?,
1998 );
1999 let m = Arc::clone(&metrics);
2000 router = router.layer(axum::middleware::from_fn(
2001 move |req: Request<Body>, next: Next| {
2002 let m = Arc::clone(&m);
2003 metrics_middleware(m, req, next)
2004 },
2005 ));
2006 let metrics_bind = config.metrics_bind.clone();
2007 let metrics_shutdown = ct.clone();
2008 tokio::spawn(async move {
2009 if let Err(e) =
2010 crate::metrics::serve_metrics(metrics_bind, metrics, metrics_shutdown).await
2011 {
2012 tracing::error!("metrics listener failed: {e}");
2013 }
2014 });
2015 }
2016
2017 let forward_resolver: Option<Arc<ForwardResolver>> = if config.trusted_proxies.is_empty() {
2025 None
2026 } else {
2027 Some(Arc::new(ForwardResolver {
2030 trusted: config
2031 .trusted_proxies
2032 .iter()
2033 .filter_map(|entry| parse_proxy_net(entry))
2034 .collect(),
2035 mode: config
2036 .forwarded_header
2037 .unwrap_or(ForwardedHeaderMode::XForwardedFor),
2038 max_scanned_entries: config.trusted_forwarder_max_entries,
2039 }))
2040 };
2041 if forward_resolver.is_some() {
2042 tracing::info!(
2043 proxies = config.trusted_proxies.len(),
2044 "trusted-forwarder mode enabled: limiters key by resolved client IP"
2045 );
2046 }
2047 router = router.layer(axum::middleware::from_fn(move |req, next| {
2048 let r = forward_resolver.clone();
2049 normalize_peer_addr_middleware(r, req, next)
2050 }));
2051
2052 router = router.layer(axum::middleware::from_fn(move |req, next| {
2064 let origins = Arc::clone(&allowed_origins);
2065 origin_check_middleware(origins, log_request_headers, req, next)
2066 }));
2067
2068 let is_tls = config.tls_cert_path.is_some();
2077 warn_security_header_overrides(&config.security_headers);
2078 let security_headers_cfg = Arc::new(config.security_headers.clone());
2079 router = router.layer(axum::middleware::from_fn(move |req, next| {
2080 let cfg = Arc::clone(&security_headers_cfg);
2081 security_headers_middleware(is_tls, cfg, req, next)
2082 }));
2083
2084 let scheme = if config.tls_cert_path.is_some() {
2085 "https"
2086 } else {
2087 "http"
2088 };
2089
2090 let tls_paths = match (&config.tls_cert_path, &config.tls_key_path) {
2091 (Some(cert), Some(key)) => Some((cert.clone(), key.clone())),
2092 _ => None,
2093 };
2094 let tls_handshake_timeout = config.tls_handshake_timeout;
2095 let max_concurrent_tls_handshakes = config.max_concurrent_tls_handshakes;
2096 let mtls_config = config.auth.as_ref().and_then(|a| a.mtls.as_ref()).cloned();
2097
2098 Ok((
2099 router,
2100 AppRunParams {
2101 tls_paths,
2102 tls_handshake_timeout,
2103 max_concurrent_tls_handshakes,
2104 mtls_config,
2105 shutdown_timeout: config.shutdown_timeout,
2106 auth_state,
2107 rbac_swap,
2108 on_reload_ready: config.on_reload_ready.take(),
2109 ct,
2110 session_ct,
2111 scheme,
2112 name: config.name.clone(),
2113 },
2114 ))
2115}
2116
2117struct CancelOnDrop(CancellationToken);
2130
2131impl Drop for CancelOnDrop {
2132 fn drop(&mut self) {
2133 self.0.cancel();
2134 }
2135}
2136
2137fn spawn_external_shutdown_bridge(
2141 external: CancellationToken,
2142 internal: CancellationToken,
2143) -> tokio::task::JoinHandle<()> {
2144 tokio::spawn(async move {
2145 tokio::select! {
2149 () = external.cancelled() => internal.cancel(),
2150 () = internal.cancelled() => {}
2151 }
2152 })
2153}
2154
2155pub async fn serve<H, F>(
2175 config: Validated<McpServerConfig>,
2176 handler_factory: F,
2177) -> Result<(), RmcpServerKitError>
2178where
2179 H: ServerHandler + 'static,
2180 F: Fn() -> H + Send + Sync + Clone + 'static,
2181{
2182 let config = config.into_inner();
2183 #[allow(
2184 deprecated,
2185 reason = "internal serve() reads `bind_addr` to construct the listener; field becomes pub(crate) in 1.0"
2186 )]
2187 let bind_addr = config.bind_addr.clone();
2188 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2189 let _cancel_guard = CancelOnDrop(params.ct.clone());
2190
2191 let listener = TcpListener::bind(&bind_addr)
2192 .await
2193 .map_err(|e| io_to_startup(&format!("bind {bind_addr}"), e))?;
2194 log_listening(¶ms.name, params.scheme, &bind_addr);
2195
2196 run_server(
2197 router,
2198 listener,
2199 params.tls_paths,
2200 params.tls_handshake_timeout,
2201 params.max_concurrent_tls_handshakes,
2202 params.mtls_config,
2203 params.shutdown_timeout,
2204 params.auth_state,
2205 params.rbac_swap,
2206 params.on_reload_ready,
2207 params.ct,
2208 params.session_ct,
2209 )
2210 .await
2211 .map_err(anyhow_to_startup)
2212}
2213
2214pub async fn serve_with_listener<H, F>(
2247 listener: TcpListener,
2248 config: Validated<McpServerConfig>,
2249 handler_factory: F,
2250 ready_tx: Option<tokio::sync::oneshot::Sender<SocketAddr>>,
2251 shutdown: Option<CancellationToken>,
2252) -> Result<(), RmcpServerKitError>
2253where
2254 H: ServerHandler + 'static,
2255 F: Fn() -> H + Send + Sync + Clone + 'static,
2256{
2257 let config = config.into_inner();
2258 let local_addr = listener
2259 .local_addr()
2260 .map_err(|e| io_to_startup("listener.local_addr", e))?;
2261 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2262 let _cancel_guard = CancelOnDrop(params.ct.clone());
2263
2264 log_listening(¶ms.name, params.scheme, &local_addr.to_string());
2265
2266 if let Some(external) = shutdown {
2270 let _bridge_task = spawn_external_shutdown_bridge(external, params.ct.clone());
2271 }
2272
2273 if let Some(tx) = ready_tx {
2277 let _ = tx.send(local_addr);
2279 }
2280
2281 run_server(
2282 router,
2283 listener,
2284 params.tls_paths,
2285 params.tls_handshake_timeout,
2286 params.max_concurrent_tls_handshakes,
2287 params.mtls_config,
2288 params.shutdown_timeout,
2289 params.auth_state,
2290 params.rbac_swap,
2291 params.on_reload_ready,
2292 params.ct,
2293 params.session_ct,
2294 )
2295 .await
2296 .map_err(anyhow_to_startup)
2297}
2298
2299#[allow(
2302 clippy::cognitive_complexity,
2303 reason = "tracing::info! macro expansions inflate the score; logic is trivial"
2304)]
2305fn log_listening(name: &str, scheme: &str, addr: &str) {
2306 tracing::info!("{name} listening on {addr}");
2307 tracing::info!(" MCP endpoint: {scheme}://{addr}/mcp");
2308 tracing::info!(" Health check: {scheme}://{addr}/healthz");
2309 tracing::info!(" Readiness: {scheme}://{addr}/readyz");
2310}
2311
2312#[allow(
2335 clippy::too_many_arguments,
2336 clippy::cognitive_complexity,
2337 reason = "server start-up threads TLS, reload state, and graceful shutdown through one flow"
2338)]
2339async fn run_server(
2343 router: axum::Router,
2344 listener: TcpListener,
2345 tls_paths: Option<(PathBuf, PathBuf)>,
2346 tls_handshake_timeout: Duration,
2347 max_concurrent_tls_handshakes: usize,
2348 mtls_config: Option<MtlsConfig>,
2349 shutdown_timeout: Duration,
2350 auth_state: Option<Arc<AuthState>>,
2351 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
2352 mut on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
2353 ct: CancellationToken,
2354 session_ct: CancellationToken,
2355) -> anyhow::Result<()> {
2356 let shutdown_trigger = CancellationToken::new();
2360 {
2361 let trigger = shutdown_trigger.clone();
2362 let parent = ct.clone();
2363 tokio::spawn(async move {
2364 tokio::select! {
2367 () = shutdown_signal() => {}
2368 () = parent.cancelled() => {}
2369 }
2370 trigger.cancel();
2371 });
2372 }
2373
2374 let graceful = {
2375 let trigger = shutdown_trigger.clone();
2376 let ct = ct.clone();
2377 async move {
2378 trigger.cancelled().await;
2379 tracing::info!("shutting down (grace period: {shutdown_timeout:?})");
2380 ct.cancel();
2381 }
2382 };
2383
2384 let force_exit_timer = {
2385 let trigger = shutdown_trigger.clone();
2386 async move {
2387 trigger.cancelled().await;
2388 tokio::time::sleep(shutdown_timeout).await;
2389 }
2390 };
2391
2392 if let Some((cert_path, key_path)) = tls_paths {
2393 let crl_set = if let Some(mtls) = mtls_config.as_ref()
2394 && mtls.crl_enabled
2395 {
2396 let (ca_certs, roots) = load_client_auth_roots(&mtls.ca_cert_path)?;
2397 let (crl_set, discover_rx) =
2398 mtls_revocation::bootstrap_fetch(roots, &ca_certs, mtls.clone())
2399 .await
2400 .map_err(|error| anyhow::anyhow!(error.to_string()))?;
2401 tokio::spawn(mtls_revocation::run_crl_refresher(
2402 Arc::clone(&crl_set),
2403 discover_rx,
2404 ct.clone(),
2405 ));
2406 Some(crl_set)
2407 } else {
2408 None
2409 };
2410
2411 if let Some(cb) = on_reload_ready.take() {
2412 cb(ReloadHandle {
2413 auth: auth_state.clone(),
2414 rbac: Some(Arc::clone(&rbac_swap)),
2415 crl_set: crl_set.clone(),
2416 });
2417 }
2418
2419 let tls_listener = TlsListener::new(
2420 listener,
2421 &cert_path,
2422 &key_path,
2423 mtls_config.as_ref(),
2424 crl_set,
2425 tls_handshake_timeout,
2426 max_concurrent_tls_handshakes,
2427 )?;
2428 let make_svc = router.into_make_service_with_connect_info::<TlsConnInfo>();
2429 tokio::select! {
2432 result = axum::serve(tls_listener, make_svc)
2433 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2434 () = force_exit_timer => {
2435 tracing::warn!("shutdown timeout exceeded, forcing exit");
2436 session_ct.cancel();
2437 }
2438 }
2439 } else {
2440 if let Some(cb) = on_reload_ready.take() {
2441 cb(ReloadHandle {
2442 auth: auth_state,
2443 rbac: Some(rbac_swap),
2444 crl_set: None,
2445 });
2446 }
2447
2448 let make_svc = router.into_make_service_with_connect_info::<SocketAddr>();
2449 tokio::select! {
2452 result = axum::serve(listener, make_svc)
2453 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2454 () = force_exit_timer => {
2455 tracing::warn!("shutdown timeout exceeded, forcing exit");
2456 session_ct.cancel();
2457 }
2458 }
2459 }
2460
2461 Ok(())
2462}
2463
2464#[cfg(feature = "oauth")]
2473fn install_oauth_proxy_routes(
2474 router: axum::Router,
2475 server_url: &str,
2476 oauth_config: &crate::oauth::OAuthConfig,
2477 auth_state: Option<&Arc<AuthState>>,
2478 max_request_body: usize,
2479 admin_role: &str,
2480) -> Result<axum::Router, RmcpServerKitError> {
2481 let Some(ref proxy) = oauth_config.proxy else {
2482 return Ok(router);
2483 };
2484
2485 let http = crate::oauth::OauthHttpClient::with_config(oauth_config)?;
2488
2489 let proxy_router = axum::Router::new();
2495
2496 let asm = crate::oauth::authorization_server_metadata(server_url, oauth_config);
2497 let proxy_router = proxy_router.route(
2498 "/.well-known/oauth-authorization-server",
2499 axum::routing::get(move || {
2500 let m = asm.clone();
2501 async move { axum::Json(m) }
2502 }),
2503 );
2504
2505 let proxy_authorize = proxy.clone();
2506 let proxy_router = proxy_router.route(
2507 "/authorize",
2508 axum::routing::get(
2509 move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
2510 let p = proxy_authorize.clone();
2511 async move { crate::oauth::handle_authorize(&p, &query.unwrap_or_default()) }
2512 },
2513 ),
2514 );
2515
2516 let proxy_token = proxy.clone();
2517 let token_http = http.clone();
2518 let proxy_router = proxy_router.route(
2519 "/token",
2520 axum::routing::post(move |body: String| {
2521 let p = proxy_token.clone();
2522 let h = token_http.clone();
2523 async move { crate::oauth::handle_token(&h, &p, &body).await }
2524 })
2525 .layer(axum::middleware::from_fn(
2526 oauth_token_cache_headers_middleware,
2527 )),
2528 );
2529
2530 let proxy_register = proxy.clone();
2531 let proxy_router = proxy_router.route(
2532 "/register",
2533 axum::routing::post(move |axum::Json(body): axum::Json<serde_json::Value>| {
2534 let p = proxy_register;
2535 async move { axum::Json(crate::oauth::handle_register(&p, &body)) }
2536 })
2537 .layer(axum::middleware::from_fn(
2538 oauth_token_cache_headers_middleware,
2539 )),
2540 );
2541
2542 let admin_routes_enabled = proxy.expose_admin_endpoints
2543 && (proxy.introspection_url.is_some() || proxy.revocation_url.is_some());
2544 if proxy.expose_admin_endpoints
2545 && !proxy.require_auth_on_admin_endpoints
2546 && proxy.allow_unauthenticated_admin_endpoints
2547 {
2548 tracing::warn!(
2552 "OAuth introspect/revoke endpoints are unauthenticated by explicit \
2553 allow_unauthenticated_admin_endpoints opt-out; ensure an \
2554 authenticated reverse proxy fronts these routes"
2555 );
2556 }
2557
2558 let admin_router = if admin_routes_enabled {
2559 build_oauth_admin_router(proxy, http, auth_state, admin_role)?
2560 } else {
2561 axum::Router::new()
2562 };
2563
2564 let proxy_router =
2568 proxy_router
2569 .merge(admin_router)
2570 .layer(tower_http::limit::RequestBodyLimitLayer::new(
2571 max_request_body,
2572 ));
2573
2574 let router = router.merge(proxy_router);
2575
2576 tracing::info!(
2577 introspect = proxy.expose_admin_endpoints && proxy.introspection_url.is_some(),
2578 revoke = proxy.expose_admin_endpoints && proxy.revocation_url.is_some(),
2579 max_request_body,
2580 "OAuth 2.1 proxy endpoints enabled (/authorize, /token, /register)"
2581 );
2582 Ok(router)
2583}
2584
2585#[cfg(feature = "oauth")]
2591fn build_oauth_admin_router(
2592 proxy: &crate::oauth::OAuthProxyConfig,
2593 http: crate::oauth::OauthHttpClient,
2594 auth_state: Option<&Arc<AuthState>>,
2595 admin_role: &str,
2596) -> Result<axum::Router, RmcpServerKitError> {
2597 let mut admin_router = axum::Router::new();
2598 if proxy.introspection_url.is_some() {
2599 let proxy_introspect = proxy.clone();
2600 let introspect_http = http.clone();
2601 admin_router = admin_router.route(
2602 "/introspect",
2603 axum::routing::post(move |body: String| {
2604 let p = proxy_introspect.clone();
2605 let h = introspect_http.clone();
2606 async move { crate::oauth::handle_introspect(&h, &p, &body).await }
2607 }),
2608 );
2609 }
2610 if proxy.revocation_url.is_some() {
2611 let proxy_revoke = proxy.clone();
2612 let revoke_http = http;
2613 admin_router = admin_router.route(
2614 "/revoke",
2615 axum::routing::post(move |body: String| {
2616 let p = proxy_revoke.clone();
2617 let h = revoke_http.clone();
2618 async move { crate::oauth::handle_revoke(&h, &p, &body).await }
2619 }),
2620 );
2621 }
2622
2623 let admin_router = admin_router.layer(axum::middleware::from_fn(
2624 oauth_token_cache_headers_middleware,
2625 ));
2626
2627 if proxy.require_auth_on_admin_endpoints {
2628 let Some(state) = auth_state else {
2629 return Err(RmcpServerKitError::Startup(
2630 "oauth proxy admin endpoints require auth state".into(),
2631 ));
2632 };
2633 let state_for_mw = Arc::clone(state);
2634 let required_role: Arc<str> = Arc::from(admin_role);
2635 Ok(admin_router
2641 .layer(axum::middleware::from_fn(move |req, next| {
2642 let r = Arc::clone(&required_role);
2643 crate::admin::require_admin_role(r, req, next)
2644 }))
2645 .layer(axum::middleware::from_fn(move |req, next| {
2646 let s = Arc::clone(&state_for_mw);
2647 auth_middleware(s, req, next)
2648 })))
2649 } else {
2650 Ok(admin_router)
2651 }
2652}
2653
2654#[allow(
2661 deprecated,
2662 reason = "internal metadata assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
2663)]
2664fn derive_server_url(config: &McpServerConfig) -> String {
2665 config.public_url.as_ref().map_or_else(
2666 || {
2667 let scheme = if config.tls_cert_path.is_some() {
2668 "https"
2669 } else {
2670 "http"
2671 };
2672 format!("{scheme}://{}", config.bind_addr)
2673 },
2674 |url| url.trim_end_matches('/').to_owned(),
2675 )
2676}
2677
2678fn derive_allowed_hosts(bind_addr: &str, public_url: Option<&str>) -> Vec<String> {
2683 let mut hosts = vec![
2684 "localhost".to_owned(),
2685 "127.0.0.1".to_owned(),
2686 "::1".to_owned(),
2687 ];
2688
2689 if let Some(url) = public_url
2690 && let Ok(uri) = url.parse::<axum::http::Uri>()
2691 && let Some(authority) = uri.authority()
2692 {
2693 let host = authority.host().to_owned();
2694 if !hosts.iter().any(|h| h == &host) {
2695 hosts.push(host);
2696 }
2697
2698 let authority = authority.as_str().to_owned();
2699 if !hosts.iter().any(|h| h == &authority) {
2700 hosts.push(authority);
2701 }
2702 }
2703
2704 if let Ok(uri) = format!("http://{bind_addr}").parse::<axum::http::Uri>()
2705 && let Some(authority) = uri.authority()
2706 {
2707 let host = authority.host().to_owned();
2708 if !hosts.iter().any(|h| h == &host) {
2709 hosts.push(host);
2710 }
2711
2712 let authority = authority.as_str().to_owned();
2713 if !hosts.iter().any(|h| h == &authority) {
2714 hosts.push(authority);
2715 }
2716 }
2717
2718 hosts
2719}
2720
2721impl axum::extract::connect_info::Connected<axum::serve::IncomingStream<'_, TlsListener>>
2734 for TlsConnInfo
2735{
2736 fn connect_info(target: axum::serve::IncomingStream<'_, TlsListener>) -> Self {
2737 let addr = *target.remote_addr();
2738 let identity = target.io().identity().cloned();
2739 Self::new(addr, identity)
2740 }
2741}
2742
2743const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
2750
2751const DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256;
2759
2760const TLS_ACCEPT_CHANNEL_CAPACITY: usize = 32;
2765
2766struct TlsListener {
2782 local_addr: SocketAddr,
2785 rx: mpsc::Receiver<(AuthenticatedTlsStream, SocketAddr)>,
2787 acceptor_task: tokio::task::JoinHandle<()>,
2790}
2791
2792impl TlsListener {
2793 fn new(
2794 inner: TcpListener,
2795 cert_path: &Path,
2796 key_path: &Path,
2797 mtls_config: Option<&MtlsConfig>,
2798 crl_set: Option<Arc<CrlSet>>,
2799 handshake_timeout: Duration,
2800 max_concurrent_handshakes: usize,
2801 ) -> anyhow::Result<Self> {
2802 rustls::crypto::ring::default_provider()
2804 .install_default()
2805 .ok();
2806
2807 let certs = load_certs(cert_path)?;
2808 let key = load_key(key_path)?;
2809
2810 let mtls_default_role;
2811
2812 let tls_config = if let Some(mtls) = mtls_config {
2813 mtls_default_role = mtls.default_role.clone();
2814 let verifier: Arc<dyn rustls::server::danger::ClientCertVerifier> = if mtls.crl_enabled
2815 {
2816 let Some(crl_set) = crl_set else {
2817 return Err(anyhow::anyhow!(
2818 "mTLS CRL verifier requested but CRL state was not initialized"
2819 ));
2820 };
2821 Arc::new(DynamicClientCertVerifier::new(crl_set))
2822 } else {
2823 let (_, root_store) = load_client_auth_roots(&mtls.ca_cert_path)?;
2824 if mtls.required {
2825 rustls::server::WebPkiClientVerifier::builder(root_store)
2826 .build()
2827 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2828 } else {
2829 rustls::server::WebPkiClientVerifier::builder(root_store)
2830 .allow_unauthenticated()
2831 .build()
2832 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2833 }
2834 };
2835
2836 tracing::info!(
2837 ca = %mtls.ca_cert_path.display(),
2838 required = mtls.required,
2839 crl_enabled = mtls.crl_enabled,
2840 "mTLS client auth configured"
2841 );
2842
2843 rustls::ServerConfig::builder_with_protocol_versions(&[
2844 &rustls::version::TLS12,
2845 &rustls::version::TLS13,
2846 ])
2847 .with_client_cert_verifier(verifier)
2848 .with_single_cert(certs, key)?
2849 } else {
2850 mtls_default_role = "viewer".to_owned();
2851 rustls::ServerConfig::builder_with_protocol_versions(&[
2852 &rustls::version::TLS12,
2853 &rustls::version::TLS13,
2854 ])
2855 .with_no_client_auth()
2856 .with_single_cert(certs, key)?
2857 };
2858
2859 let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
2860 tracing::info!(
2861 "TLS enabled (cert: {}, key: {})",
2862 cert_path.display(),
2863 key_path.display()
2864 );
2865 let local_addr = inner.local_addr()?;
2866 let (tx, rx) = mpsc::channel(TLS_ACCEPT_CHANNEL_CAPACITY);
2867 let acceptor_task = tokio::spawn(run_tls_acceptor(
2868 inner,
2869 acceptor,
2870 mtls_default_role,
2871 tx,
2872 handshake_timeout,
2873 max_concurrent_handshakes,
2874 ));
2875 Ok(Self {
2876 local_addr,
2877 rx,
2878 acceptor_task,
2879 })
2880 }
2881
2882 fn extract_handshake_identity(
2886 tls_stream: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2887 default_role: &str,
2888 addr: SocketAddr,
2889 ) -> Option<AuthIdentity> {
2890 let (_, server_conn) = tls_stream.get_ref();
2891 let cert_der = server_conn.peer_certificates()?.first()?;
2892 let id = extract_mtls_identity(cert_der.as_ref(), default_role)?;
2893 tracing::debug!(name = %id.name, peer = %addr, "mTLS client cert accepted");
2894 Some(id)
2895 }
2896}
2897
2898async fn run_tls_acceptor(
2909 listener: TcpListener,
2910 acceptor: tokio_rustls::TlsAcceptor,
2911 default_role: String,
2912 tx: mpsc::Sender<(AuthenticatedTlsStream, SocketAddr)>,
2913 handshake_timeout: Duration,
2914 max_concurrent_handshakes: usize,
2915) {
2916 let inflight = Arc::new(Semaphore::new(max_concurrent_handshakes));
2917 loop {
2918 let Ok(permit) = Arc::clone(&inflight).acquire_owned().await else {
2922 return;
2924 };
2925 let (stream, addr) = match listener.accept().await {
2926 Ok(pair) => pair,
2927 Err(e) => {
2928 tracing::debug!("TCP accept error: {e}");
2929 continue;
2930 }
2931 };
2932 if tx.is_closed() {
2933 return;
2935 }
2936 let acceptor = acceptor.clone();
2937 let default_role = default_role.clone();
2938 let tx = tx.clone();
2939 tokio::spawn(async move {
2940 let _permit = permit;
2941 match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
2942 Ok(Ok(tls_stream)) => {
2943 let identity =
2944 TlsListener::extract_handshake_identity(&tls_stream, &default_role, addr);
2945 let wrapped = AuthenticatedTlsStream {
2946 inner: tls_stream,
2947 identity,
2948 };
2949 let _ = tx.send((wrapped, addr)).await;
2952 }
2953 Ok(Err(e)) => {
2954 tracing::debug!("TLS handshake failed from {addr}: {e}");
2955 }
2956 Err(_elapsed) => {
2957 tracing::debug!(
2958 "TLS handshake timed out from {addr} after {handshake_timeout:?}"
2959 );
2960 }
2961 }
2962 });
2963 }
2964}
2965
2966pub(crate) struct AuthenticatedTlsStream {
2978 inner: tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2979 identity: Option<AuthIdentity>,
2980}
2981
2982impl AuthenticatedTlsStream {
2983 #[must_use]
2985 pub(crate) const fn identity(&self) -> Option<&AuthIdentity> {
2986 self.identity.as_ref()
2987 }
2988}
2989
2990impl std::fmt::Debug for AuthenticatedTlsStream {
2991 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2992 f.debug_struct("AuthenticatedTlsStream")
2993 .field("identity", &self.identity.as_ref().map(|id| &id.name))
2994 .finish_non_exhaustive()
2995 }
2996}
2997
2998impl tokio::io::AsyncRead for AuthenticatedTlsStream {
2999 fn poll_read(
3000 mut self: Pin<&mut Self>,
3001 cx: &mut std::task::Context<'_>,
3002 buf: &mut tokio::io::ReadBuf<'_>,
3003 ) -> std::task::Poll<std::io::Result<()>> {
3004 Pin::new(&mut self.inner).poll_read(cx, buf)
3005 }
3006}
3007
3008impl tokio::io::AsyncWrite for AuthenticatedTlsStream {
3009 fn poll_write(
3010 mut self: Pin<&mut Self>,
3011 cx: &mut std::task::Context<'_>,
3012 buf: &[u8],
3013 ) -> std::task::Poll<std::io::Result<usize>> {
3014 Pin::new(&mut self.inner).poll_write(cx, buf)
3015 }
3016
3017 fn poll_flush(
3018 mut self: Pin<&mut Self>,
3019 cx: &mut std::task::Context<'_>,
3020 ) -> std::task::Poll<std::io::Result<()>> {
3021 Pin::new(&mut self.inner).poll_flush(cx)
3022 }
3023
3024 fn poll_shutdown(
3025 mut self: Pin<&mut Self>,
3026 cx: &mut std::task::Context<'_>,
3027 ) -> std::task::Poll<std::io::Result<()>> {
3028 Pin::new(&mut self.inner).poll_shutdown(cx)
3029 }
3030
3031 fn poll_write_vectored(
3032 mut self: Pin<&mut Self>,
3033 cx: &mut std::task::Context<'_>,
3034 bufs: &[std::io::IoSlice<'_>],
3035 ) -> std::task::Poll<std::io::Result<usize>> {
3036 Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
3037 }
3038
3039 fn is_write_vectored(&self) -> bool {
3040 self.inner.is_write_vectored()
3041 }
3042}
3043
3044impl axum::serve::Listener for TlsListener {
3045 type Io = AuthenticatedTlsStream;
3046 type Addr = SocketAddr;
3047
3048 async fn accept(&mut self) -> (Self::Io, Self::Addr) {
3054 if let Some(pair) = self.rx.recv().await {
3055 return pair;
3056 }
3057 tracing::error!("TLS acceptor task terminated; no further connections will be accepted");
3063 std::future::pending().await
3064 }
3065
3066 fn local_addr(&self) -> std::io::Result<Self::Addr> {
3067 Ok(self.local_addr)
3068 }
3069}
3070
3071impl Drop for TlsListener {
3072 fn drop(&mut self) {
3073 self.acceptor_task.abort();
3076 }
3077}
3078
3079fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
3080 use rustls::pki_types::pem::PemObject;
3081 let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_file_iter(path)
3082 .map_err(|e| anyhow::anyhow!("failed to read certs from {}: {e}", path.display()))?
3083 .collect::<Result<_, _>>()
3084 .map_err(|e| anyhow::anyhow!("invalid cert in {}: {e}", path.display()))?;
3085 anyhow::ensure!(
3086 !certs.is_empty(),
3087 "no certificates found in {}",
3088 path.display()
3089 );
3090 Ok(certs)
3091}
3092
3093fn load_client_auth_roots(
3094 path: &Path,
3095) -> anyhow::Result<(
3096 Vec<rustls::pki_types::CertificateDer<'static>>,
3097 Arc<RootCertStore>,
3098)> {
3099 let ca_certs = load_certs(path)?;
3100 let mut root_store = RootCertStore::empty();
3101 for cert in &ca_certs {
3102 root_store
3103 .add(cert.clone())
3104 .map_err(|error| anyhow::anyhow!("invalid CA cert: {error}"))?;
3105 }
3106
3107 Ok((ca_certs, Arc::new(root_store)))
3108}
3109
3110fn load_key(path: &Path) -> anyhow::Result<rustls::pki_types::PrivateKeyDer<'static>> {
3111 use rustls::pki_types::pem::PemObject;
3112 rustls::pki_types::PrivateKeyDer::from_pem_file(path)
3113 .map_err(|e| anyhow::anyhow!("failed to read key from {}: {e}", path.display()))
3114}
3115
3116#[allow(
3118 clippy::unused_async,
3119 reason = "axum route handler signature requires `async fn` even when the body is synchronous"
3120)]
3121async fn healthz() -> impl IntoResponse {
3122 axum::Json(serde_json::json!({
3123 "status": "ok",
3124 }))
3125}
3126
3127fn version_payload(name: &str, version: &str, expose_build_metadata: bool) -> serde_json::Value {
3137 let mut map = serde_json::Map::new();
3138 map.insert("name".into(), name.into());
3139 map.insert("version".into(), version.into());
3140 map.insert(
3141 "rmcp_server_kit_version".into(),
3142 env!("CARGO_PKG_VERSION").into(),
3143 );
3144 if expose_build_metadata {
3145 map.insert(
3146 "build_git_sha".into(),
3147 option_env!("RMCP_SERVER_KIT_BUILD_SHA")
3148 .unwrap_or("unknown")
3149 .into(),
3150 );
3151 map.insert(
3152 "build_timestamp".into(),
3153 option_env!("RMCP_SERVER_KIT_BUILD_TIME")
3154 .unwrap_or("unknown")
3155 .into(),
3156 );
3157 map.insert(
3158 "rust_version".into(),
3159 option_env!("RMCP_SERVER_KIT_RUSTC_VERSION")
3160 .unwrap_or("unknown")
3161 .into(),
3162 );
3163 }
3164 serde_json::Value::Object(map)
3165}
3166
3167fn serialize_version_payload(name: &str, version: &str, expose_build_metadata: bool) -> Arc<[u8]> {
3177 let value = version_payload(name, version, expose_build_metadata);
3178 serde_json::to_vec(&value).map_or_else(|_| Arc::from(&b"{}"[..]), Arc::from)
3179}
3180
3181async fn readyz(check: ReadinessCheck) -> impl IntoResponse {
3186 let status = check().await;
3187 let ready = status
3188 .get("ready")
3189 .and_then(serde_json::Value::as_bool)
3190 .unwrap_or(false);
3191 let code = if ready {
3192 axum::http::StatusCode::OK
3193 } else {
3194 axum::http::StatusCode::SERVICE_UNAVAILABLE
3195 };
3196 (code, axum::Json(status))
3197}
3198
3199async fn shutdown_signal() {
3203 let ctrl_c = tokio::signal::ctrl_c();
3204
3205 #[cfg(unix)]
3206 {
3207 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
3208 Ok(mut term) => {
3209 tokio::select! {
3212 _ = ctrl_c => {}
3213 _ = term.recv() => {}
3214 }
3215 }
3216 Err(e) => {
3217 tracing::warn!(error = %e, "failed to register SIGTERM handler, using SIGINT only");
3218 ctrl_c.await.ok();
3219 }
3220 }
3221 }
3222
3223 #[cfg(not(unix))]
3224 {
3225 ctrl_c.await.ok();
3226 }
3227}
3228
3229#[cfg(feature = "metrics")]
3246fn metrics_labels(req: &Request<Body>) -> (&'static str, String) {
3247 let method = match *req.method() {
3248 axum::http::Method::GET => "GET",
3249 axum::http::Method::POST => "POST",
3250 axum::http::Method::PUT => "PUT",
3251 axum::http::Method::PATCH => "PATCH",
3252 axum::http::Method::DELETE => "DELETE",
3253 axum::http::Method::HEAD => "HEAD",
3254 axum::http::Method::OPTIONS => "OPTIONS",
3255 axum::http::Method::TRACE => "TRACE",
3256 axum::http::Method::CONNECT => "CONNECT",
3257 _ => "OTHER",
3260 };
3261
3262 let path = req
3263 .extensions()
3264 .get::<axum::extract::MatchedPath>()
3265 .map_or_else(
3266 || {
3267 let raw = req.uri().path();
3268 if raw == "/mcp" || raw.starts_with("/mcp/") {
3269 "/mcp".to_owned()
3270 } else {
3271 "<unmatched>".to_owned()
3272 }
3273 },
3274 |matched| matched.as_str().to_owned(),
3275 );
3276
3277 (method, path)
3278}
3279
3280#[cfg(feature = "metrics")]
3290async fn metrics_middleware(
3291 metrics: Arc<crate::metrics::McpMetrics>,
3292 mut req: Request<Body>,
3293 next: Next,
3294) -> axum::response::Response {
3295 let (method, path) = metrics_labels(&req);
3296 let start = std::time::Instant::now();
3297
3298 req.extensions_mut().insert(Arc::clone(&metrics));
3299 let response = next.run(req).await;
3300
3301 let mut status_buf = core::fmt::NumBuffer::<u16>::new();
3302 let status = response.status().as_u16().format_into(&mut status_buf);
3303 let duration = start.elapsed().as_secs_f64();
3304
3305 metrics
3306 .http_requests_total
3307 .with_label_values(&[method, &path, status])
3308 .inc();
3309 metrics
3310 .http_request_duration_seconds
3311 .with_label_values(&[method, &path])
3312 .observe(duration);
3313
3314 response
3315}
3316
3317async fn security_headers_middleware(
3331 is_tls: bool,
3332 cfg: Arc<SecurityHeadersConfig>,
3333 req: Request<Body>,
3334 next: Next,
3335) -> axum::response::Response {
3336 use axum::http::{HeaderName, header};
3337
3338 let mut resp = next.run(req).await;
3339 let headers = resp.headers_mut();
3340
3341 headers.remove(header::SERVER);
3343 headers.remove(HeaderName::from_static("x-powered-by"));
3344
3345 apply_security_header(
3346 headers,
3347 header::X_CONTENT_TYPE_OPTIONS,
3348 cfg.x_content_type_options.as_deref(),
3349 "nosniff",
3350 );
3351 apply_security_header(
3352 headers,
3353 header::X_FRAME_OPTIONS,
3354 cfg.x_frame_options.as_deref(),
3355 "deny",
3356 );
3357 apply_security_header(
3358 headers,
3359 header::CACHE_CONTROL,
3360 cfg.cache_control.as_deref(),
3361 "no-store, max-age=0",
3362 );
3363 apply_security_header(
3364 headers,
3365 header::REFERRER_POLICY,
3366 cfg.referrer_policy.as_deref(),
3367 "no-referrer",
3368 );
3369 apply_security_header(
3370 headers,
3371 HeaderName::from_static("cross-origin-opener-policy"),
3372 cfg.cross_origin_opener_policy.as_deref(),
3373 "same-origin",
3374 );
3375 apply_security_header(
3376 headers,
3377 HeaderName::from_static("cross-origin-resource-policy"),
3378 cfg.cross_origin_resource_policy.as_deref(),
3379 "same-origin",
3380 );
3381 apply_security_header(
3382 headers,
3383 HeaderName::from_static("cross-origin-embedder-policy"),
3384 cfg.cross_origin_embedder_policy.as_deref(),
3385 "require-corp",
3386 );
3387 apply_security_header(
3388 headers,
3389 HeaderName::from_static("permissions-policy"),
3390 cfg.permissions_policy.as_deref(),
3391 "accelerometer=(), camera=(), geolocation=(), microphone=()",
3392 );
3393 apply_security_header(
3394 headers,
3395 HeaderName::from_static("x-permitted-cross-domain-policies"),
3396 cfg.x_permitted_cross_domain_policies.as_deref(),
3397 "none",
3398 );
3399 apply_security_header(
3400 headers,
3401 HeaderName::from_static("content-security-policy"),
3402 cfg.content_security_policy.as_deref(),
3403 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests",
3404 );
3405 apply_security_header(
3406 headers,
3407 HeaderName::from_static("x-dns-prefetch-control"),
3408 cfg.x_dns_prefetch_control.as_deref(),
3409 "off",
3410 );
3411
3412 if is_tls {
3413 apply_security_header(
3414 headers,
3415 header::STRICT_TRANSPORT_SECURITY,
3416 cfg.strict_transport_security.as_deref(),
3417 "max-age=63072000; includeSubDomains",
3418 );
3419 }
3420
3421 resp
3422}
3423
3424fn apply_security_header(
3435 headers: &mut axum::http::HeaderMap,
3436 name: axum::http::HeaderName,
3437 override_value: Option<&str>,
3438 default: &'static str,
3439) {
3440 use axum::http::HeaderValue;
3441
3442 match override_value {
3443 None => {
3444 headers.insert(name, HeaderValue::from_static(default));
3445 }
3446 Some("") => {
3447 }
3449 Some(v) => match HeaderValue::from_str(v) {
3450 Ok(hv) => {
3451 headers.insert(name, hv);
3452 }
3453 Err(err) => {
3454 tracing::error!(
3455 header = %name,
3456 error = %err,
3457 "invalid security header override reached middleware; using default"
3458 );
3459 headers.insert(name, HeaderValue::from_static(default));
3460 }
3461 },
3462 }
3463}
3464
3465fn validate_security_headers(cfg: &SecurityHeadersConfig) -> Result<(), RmcpServerKitError> {
3476 use axum::http::HeaderValue;
3477
3478 let fields: &[(&str, Option<&str>)] = &[
3479 (
3480 "x_content_type_options",
3481 cfg.x_content_type_options.as_deref(),
3482 ),
3483 ("x_frame_options", cfg.x_frame_options.as_deref()),
3484 ("cache_control", cfg.cache_control.as_deref()),
3485 ("referrer_policy", cfg.referrer_policy.as_deref()),
3486 (
3487 "cross_origin_opener_policy",
3488 cfg.cross_origin_opener_policy.as_deref(),
3489 ),
3490 (
3491 "cross_origin_resource_policy",
3492 cfg.cross_origin_resource_policy.as_deref(),
3493 ),
3494 (
3495 "cross_origin_embedder_policy",
3496 cfg.cross_origin_embedder_policy.as_deref(),
3497 ),
3498 ("permissions_policy", cfg.permissions_policy.as_deref()),
3499 (
3500 "x_permitted_cross_domain_policies",
3501 cfg.x_permitted_cross_domain_policies.as_deref(),
3502 ),
3503 (
3504 "content_security_policy",
3505 cfg.content_security_policy.as_deref(),
3506 ),
3507 (
3508 "x_dns_prefetch_control",
3509 cfg.x_dns_prefetch_control.as_deref(),
3510 ),
3511 (
3512 "strict_transport_security",
3513 cfg.strict_transport_security.as_deref(),
3514 ),
3515 ];
3516
3517 for (field, value) in fields {
3518 let Some(v) = value else { continue };
3519 if v.is_empty() {
3520 continue;
3521 }
3522 if let Err(err) = HeaderValue::from_str(v) {
3523 return Err(RmcpServerKitError::Config(format!(
3524 "invalid security_headers.{field}: {err}"
3525 )));
3526 }
3527 }
3528
3529 if let Some(v) = cfg.strict_transport_security.as_deref()
3530 && !v.is_empty()
3531 && v.to_ascii_lowercase().contains("preload")
3532 {
3533 return Err(RmcpServerKitError::Config(format!(
3534 "invalid security_headers.strict_transport_security: {v:?} contains the `preload` directive; \
3535 HSTS preload must be opted into explicitly via a dedicated builder, not via this knob"
3536 )));
3537 }
3538
3539 Ok(())
3540}
3541
3542#[cfg(feature = "oauth")]
3557async fn oauth_token_cache_headers_middleware(
3558 req: Request<Body>,
3559 next: Next,
3560) -> axum::response::Response {
3561 use axum::http::{HeaderValue, header};
3562
3563 let mut resp = next.run(req).await;
3564 let headers = resp.headers_mut();
3565 headers.insert(header::PRAGMA, HeaderValue::from_static("no-cache"));
3566 headers.append(header::VARY, HeaderValue::from_static("Authorization"));
3567 resp
3568}
3569
3570async fn normalize_peer_addr_middleware(
3601 resolver: Option<Arc<ForwardResolver>>,
3602 mut req: Request<Body>,
3603 next: Next,
3604) -> axum::response::Response {
3605 let direct = req
3606 .extensions()
3607 .get::<ConnectInfo<SocketAddr>>()
3608 .map(|ci| ci.0);
3609 let from_tls = req
3610 .extensions()
3611 .get::<ConnectInfo<TlsConnInfo>>()
3612 .map(|ci| ci.0.addr);
3613 if let Some(addr) = direct.or(from_tls) {
3614 if direct.is_none() {
3615 req.extensions_mut().insert(ConnectInfo(addr));
3616 }
3617 req.extensions_mut().insert(PeerAddr::new(addr));
3618 let client_ip = match &resolver {
3619 Some(r) => crate::forwarded::resolve_client_ip(
3620 addr.ip(),
3621 req.headers(),
3622 &r.trusted,
3623 r.mode,
3624 r.max_scanned_entries,
3625 )
3626 .unwrap_or_else(|reason| {
3627 tracing::debug!(
3628 reason = ?reason,
3629 "forwarded-header resolution fell back to direct peer"
3630 );
3631 addr.ip()
3632 }),
3633 None => addr.ip(),
3634 };
3635 req.extensions_mut().insert(ClientIp::new(client_ip));
3636 }
3637 next.run(req).await
3638}
3639
3640fn parse_proxy_net(entry: &str) -> Option<ipnet::IpNet> {
3643 if let Ok(net) = entry.parse::<ipnet::IpNet>() {
3644 return Some(net);
3645 }
3646 entry.parse::<IpAddr>().ok().map(ipnet::IpNet::from)
3647}
3648
3649pub(crate) fn validate_trusted_proxy_entry(entry: &str) -> Result<(), String> {
3659 match parse_proxy_net(entry) {
3660 None => Err(format!(
3661 "trusted_proxies entry {entry:?} is neither a CIDR nor an IP address"
3662 )),
3663 Some(net) if net.prefix_len() == 0 => Err(format!(
3664 "trusted_proxies entry {entry:?}: prefix length 0 is forbidden (marks every peer trusted, enabling client-IP spoofing)"
3665 )),
3666 Some(_) => Ok(()),
3667 }
3668}
3669
3670pub(crate) fn limiter_client_ip(extensions: &axum::http::Extensions) -> Option<IpAddr> {
3674 if let Some(client) = extensions.get::<ClientIp>() {
3675 return Some(client.ip);
3676 }
3677 extensions
3678 .get::<ConnectInfo<SocketAddr>>()
3679 .map(|ci| ci.0.ip())
3680 .or_else(|| {
3681 extensions
3682 .get::<ConnectInfo<TlsConnInfo>>()
3683 .map(|ci| ci.0.addr.ip())
3684 })
3685}
3686
3687#[derive(Clone, PartialEq, Eq, Hash, Debug)]
3700pub(crate) enum RateLimitKey {
3701 Ip(IpAddr),
3703 Unattributed,
3705}
3706
3707impl std::fmt::Display for RateLimitKey {
3708 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3709 match self {
3710 Self::Ip(ip) => write!(f, "{ip}"),
3711 Self::Unattributed => f.write_str("unattributed"),
3712 }
3713 }
3714}
3715
3716static UNATTRIBUTED_WARNED: std::sync::atomic::AtomicBool =
3718 std::sync::atomic::AtomicBool::new(false);
3719
3720pub(crate) fn limiter_client_key(extensions: &axum::http::Extensions) -> RateLimitKey {
3731 if let Some(ip) = limiter_client_ip(extensions) {
3732 return RateLimitKey::Ip(ip);
3733 }
3734 if !UNATTRIBUTED_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
3735 tracing::warn!(
3736 "request carries no resolvable client address; rate limiting is \
3737 falling back to a single shared bucket. This indicates \
3738 rmcp-server-kit middleware composed outside serve()."
3739 );
3740 }
3741 RateLimitKey::Unattributed
3742}
3743
3744pub(crate) type ExtraRouteRateLimiter = BoundedKeyedLimiter<RateLimitKey>;
3748
3749const EXTRA_ROUTE_MAX_TRACKED_KEYS: usize = 10_000;
3755
3756const EXTRA_ROUTE_IDLE_EVICTION: Duration = Duration::from_mins(15);
3759
3760fn build_extra_route_rate_limiter_with_policy(
3767 per_minute: u32,
3768 burst: Option<u32>,
3769 key_eviction_policy: KeyEvictionPolicy,
3770 max_tracked_keys: NonZeroUsize,
3771) -> Arc<ExtraRouteRateLimiter> {
3772 let rate = std::num::NonZeroU32::new(per_minute.max(1)).unwrap_or(std::num::NonZeroU32::MIN);
3773 let mut quota = governor::Quota::per_minute(rate);
3774 if let Some(b) = burst.and_then(std::num::NonZeroU32::new) {
3775 quota = quota.allow_burst(b);
3776 }
3777 Arc::new(BoundedKeyedLimiter::new_with_policy(
3778 quota,
3779 max_tracked_keys,
3780 EXTRA_ROUTE_IDLE_EVICTION,
3781 key_eviction_policy,
3782 ))
3783}
3784
3785async fn extra_route_rate_limit_middleware(
3810 limiter: Arc<ExtraRouteRateLimiter>,
3811 exempt: Arc<std::collections::HashSet<String>>,
3812 req: Request<Body>,
3813 next: Next,
3814) -> axum::response::Response {
3815 if exempt.contains(req.uri().path()) {
3816 return next.run(req).await;
3817 }
3818 let peer_key = limiter_client_key(req.extensions());
3819 match limiter.check_key_detailed(&peer_key) {
3820 Ok(()) => {}
3821 Err(BoundedLimiterDeny::RateLimited(wait)) => {
3822 #[cfg(feature = "metrics")]
3823 crate::metrics::record_rate_limit_deny(req.extensions(), "extra_route");
3824 tracing::warn!(rate_limit_key = %peer_key, "extra route request rate limited");
3825 return RmcpServerKitError::RateLimitedFor {
3826 message: "too many requests to application routes from this source".into(),
3827 retry_after: wait,
3828 }
3829 .into_response();
3830 }
3831 Err(BoundedLimiterDeny::CapacityFull) => {
3832 tracing::warn!(
3833 rate_limit_key = %peer_key,
3834 "extra route limiter rejected unseen key because tracked-key capacity is full"
3835 );
3836 return (
3837 axum::http::StatusCode::SERVICE_UNAVAILABLE,
3838 "rate limiter capacity exhausted",
3839 )
3840 .into_response();
3841 }
3842 }
3843 next.run(req).await
3844}
3845
3846async fn origin_check_middleware(
3852 allowed: Arc<[String]>,
3853 log_request_headers: bool,
3854 req: Request<Body>,
3855 next: Next,
3856) -> axum::response::Response {
3857 let method = req.method().clone();
3858 let path = req.uri().path().to_owned();
3859
3860 log_incoming_request(&method, &path, req.headers(), log_request_headers);
3861
3862 if let Some(origin) = req.headers().get(axum::http::header::ORIGIN) {
3863 let origin_str = origin.to_str().unwrap_or("");
3864 if !allowed.iter().any(|a| a == origin_str) {
3865 tracing::warn!(
3866 origin = origin_str,
3867 %method,
3868 %path,
3869 allowed = ?&*allowed,
3870 "rejected request: Origin not allowed"
3871 );
3872 return (
3873 axum::http::StatusCode::FORBIDDEN,
3874 "Forbidden: Origin not allowed",
3875 )
3876 .into_response();
3877 }
3878 }
3879 next.run(req).await
3880}
3881
3882fn log_incoming_request(
3885 method: &axum::http::Method,
3886 path: &str,
3887 headers: &axum::http::HeaderMap,
3888 log_request_headers: bool,
3889) {
3890 if log_request_headers {
3891 tracing::debug!(
3892 %method,
3893 %path,
3894 headers = %format_request_headers_for_log(headers),
3895 "incoming request"
3896 );
3897 } else {
3898 tracing::debug!(%method, %path, "incoming request");
3899 }
3900}
3901
3902const REDACTED_LOG_HEADERS: [&str; 6] = [
3910 "authorization",
3911 "cookie",
3912 "proxy-authorization",
3913 "forwarded",
3914 "x-forwarded-for",
3915 "x-real-ip",
3916];
3917
3918fn format_request_headers_for_log(headers: &axum::http::HeaderMap) -> String {
3919 headers
3920 .iter()
3921 .map(|(k, v)| {
3922 let name = k.as_str();
3923 if REDACTED_LOG_HEADERS.contains(&name) {
3924 format!("{name}: [REDACTED]")
3925 } else {
3926 format!("{name}: {}", v.to_str().unwrap_or("<non-utf8>"))
3927 }
3928 })
3929 .collect::<Vec<_>>()
3930 .join(", ")
3931}
3932
3933#[allow(
3957 clippy::cognitive_complexity,
3958 reason = "complexity is purely tracing macro expansion (info/warn + match arms); 18 lines of straight-line code, nothing meaningful to extract"
3959)]
3960pub async fn serve_stdio<H>(handler: H) -> Result<(), RmcpServerKitError>
3961where
3962 H: ServerHandler + 'static,
3963{
3964 use rmcp::ServiceExt as _;
3965
3966 tracing::info!("stdio transport: serving on stdin/stdout");
3967 tracing::warn!("stdio mode: auth, RBAC, TLS, and Origin checks are DISABLED");
3968
3969 let transport = rmcp::transport::io::stdio();
3970
3971 let service = handler
3972 .serve(transport)
3973 .await
3974 .map_err(|e| RmcpServerKitError::Startup(format!("stdio initialize failed: {e}")))?;
3975
3976 if let Err(e) = service.waiting().await {
3977 tracing::warn!(error = %e, "stdio session ended with error");
3978 }
3979 tracing::info!("stdio session ended");
3980 Ok(())
3981}
3982
3983#[allow(
3984 deprecated,
3985 reason = "builder methods are the sanctioned transition layer for deprecated public fields"
3986)]
3987impl McpServerConfig {
3988 #[must_use]
3992 pub fn with_tls_paths(mut self, cert_path: Option<PathBuf>, key_path: Option<PathBuf>) -> Self {
3993 self.tls_cert_path = cert_path;
3994 self.tls_key_path = key_path;
3995 self
3996 }
3997
3998 #[must_use]
4002 pub fn with_tls_cert_path(mut self, cert_path: impl Into<PathBuf>) -> Self {
4003 self.tls_cert_path = Some(cert_path.into());
4004 self
4005 }
4006
4007 #[must_use]
4011 pub fn with_tls_key_path(mut self, key_path: impl Into<PathBuf>) -> Self {
4012 self.tls_key_path = Some(key_path.into());
4013 self
4014 }
4015
4016 #[must_use]
4018 pub fn with_optional_auth(mut self, auth: Option<AuthConfig>) -> Self {
4019 self.auth = auth;
4020 self
4021 }
4022
4023 #[must_use]
4025 pub fn with_optional_tool_rate_limit(mut self, per_minute: Option<u32>) -> Self {
4026 self.tool_rate_limit = per_minute;
4027 self
4028 }
4029
4030 #[must_use]
4032 pub fn with_optional_tool_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
4033 self.tool_rate_limit_burst = burst;
4034 self
4035 }
4036
4037 #[must_use]
4039 pub fn with_optional_extra_route_rate_limit(mut self, per_minute: Option<u32>) -> Self {
4040 self.extra_route_rate_limit = per_minute;
4041 self
4042 }
4043
4044 #[must_use]
4046 pub fn with_optional_extra_route_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
4047 self.extra_route_rate_limit_burst = burst;
4048 self
4049 }
4050
4051 #[must_use]
4053 pub fn with_optional_forwarded_header(mut self, mode: Option<ForwardedHeaderMode>) -> Self {
4054 self.forwarded_header = mode;
4055 self
4056 }
4057
4058 #[must_use]
4060 pub fn with_optional_public_url(mut self, url: Option<String>) -> Self {
4061 self.public_url = url;
4062 self
4063 }
4064
4065 #[must_use]
4069 pub fn with_compression_min_size(mut self, min_size: u16) -> Self {
4070 self.compression_min_size = min_size;
4071 self
4072 }
4073
4074 #[must_use]
4076 pub fn with_compression_enabled(mut self, enabled: bool) -> Self {
4077 self.compression_enabled = enabled;
4078 self
4079 }
4080
4081 #[must_use]
4083 pub fn with_optional_max_concurrent_requests(mut self, limit: Option<usize>) -> Self {
4084 self.max_concurrent_requests = limit;
4085 self
4086 }
4087
4088 #[must_use]
4090 pub fn with_admin_enabled(mut self, enabled: bool) -> Self {
4091 self.admin_enabled = enabled;
4092 self
4093 }
4094
4095 #[must_use]
4098 pub fn with_admin_role(mut self, role: impl Into<String>) -> Self {
4099 self.admin_role = role.into();
4100 self
4101 }
4102
4103 #[must_use]
4105 pub fn with_expose_build_metadata(mut self, enabled: bool) -> Self {
4106 self.expose_build_metadata = enabled;
4107 self
4108 }
4109}
4110
4111fn warn_security_header_overrides(cfg: &SecurityHeadersConfig) {
4112 for (field, value) in security_header_overrides(cfg) {
4113 let action = if value.is_empty() {
4114 "omitted"
4115 } else {
4116 "overridden"
4117 };
4118 tracing::warn!(
4119 security_header = field,
4120 action,
4121 "security header configured; inspect server.security_headers.<security_header>"
4122 );
4123 }
4124}
4125
4126fn security_header_overrides(
4127 cfg: &SecurityHeadersConfig,
4128) -> impl Iterator<Item = (&'static str, &str)> {
4129 [
4130 (
4131 "x_content_type_options",
4132 cfg.x_content_type_options.as_deref(),
4133 ),
4134 ("x_frame_options", cfg.x_frame_options.as_deref()),
4135 ("cache_control", cfg.cache_control.as_deref()),
4136 ("referrer_policy", cfg.referrer_policy.as_deref()),
4137 (
4138 "cross_origin_opener_policy",
4139 cfg.cross_origin_opener_policy.as_deref(),
4140 ),
4141 (
4142 "cross_origin_resource_policy",
4143 cfg.cross_origin_resource_policy.as_deref(),
4144 ),
4145 (
4146 "cross_origin_embedder_policy",
4147 cfg.cross_origin_embedder_policy.as_deref(),
4148 ),
4149 ("permissions_policy", cfg.permissions_policy.as_deref()),
4150 (
4151 "x_permitted_cross_domain_policies",
4152 cfg.x_permitted_cross_domain_policies.as_deref(),
4153 ),
4154 (
4155 "content_security_policy",
4156 cfg.content_security_policy.as_deref(),
4157 ),
4158 (
4159 "x_dns_prefetch_control",
4160 cfg.x_dns_prefetch_control.as_deref(),
4161 ),
4162 (
4163 "strict_transport_security",
4164 cfg.strict_transport_security.as_deref(),
4165 ),
4166 ]
4167 .into_iter()
4168 .filter_map(|(field, value)| value.map(|v| (field, v)))
4169}
4170
4171fn check_auth_capacity_knobs(auth: Option<&AuthConfig>) -> Result<(), RmcpServerKitError> {
4172 if let Some(auth_cfg) = auth {
4173 if let Some(rl) = &auth_cfg.rate_limit {
4174 (rl.max_attempts_per_minute != 0).ok_or_else(|| {
4175 RmcpServerKitError::Config(
4176 "auth.rate_limit.max_attempts_per_minute must be nonzero".into(),
4177 )
4178 })?;
4179 (rl.pre_auth_max_per_minute != Some(0)).ok_or_else(|| {
4184 RmcpServerKitError::Config(
4185 "auth.rate_limit.pre_auth_max_per_minute must be nonzero when set".into(),
4186 )
4187 })?;
4188 }
4189 if let Some(mtls) = &auth_cfg.mtls {
4190 check_mtls_capacity_knobs(mtls)?;
4191 }
4192 auth_cfg.check_oauth_feature()?;
4193 }
4194 Ok(())
4195}
4196
4197fn check_mtls_capacity_knobs(mtls: &MtlsConfig) -> Result<(), RmcpServerKitError> {
4198 (mtls.crl_max_concurrent_fetches != 0).ok_or_else(|| {
4199 RmcpServerKitError::Config("auth.mtls.crl_max_concurrent_fetches must be nonzero".into())
4200 })?;
4201 (mtls.crl_discovery_rate_per_min != 0).ok_or_else(|| {
4202 RmcpServerKitError::Config("auth.mtls.crl_discovery_rate_per_min must be nonzero".into())
4203 })?;
4204 (mtls.crl_max_host_semaphores != 0).ok_or_else(|| {
4205 RmcpServerKitError::Config("auth.mtls.crl_max_host_semaphores must be nonzero".into())
4206 })?;
4207 (mtls.crl_max_seen_urls != 0).ok_or_else(|| {
4208 RmcpServerKitError::Config("auth.mtls.crl_max_seen_urls must be nonzero".into())
4209 })?;
4210 (mtls.crl_max_cache_entries != 0).ok_or_else(|| {
4211 RmcpServerKitError::Config("auth.mtls.crl_max_cache_entries must be nonzero".into())
4212 })?;
4213 (mtls.crl_max_response_bytes != 0).ok_or_else(|| {
4218 RmcpServerKitError::Config("auth.mtls.crl_max_response_bytes must be nonzero".into())
4219 })?;
4220 Ok(())
4221}
4222
4223#[cfg(test)]
4224mod tests {
4225 #![allow(
4226 clippy::unwrap_used,
4227 clippy::expect_used,
4228 clippy::panic,
4229 clippy::indexing_slicing,
4230 clippy::unwrap_in_result,
4231 clippy::print_stdout,
4232 clippy::print_stderr,
4233 deprecated,
4234 reason = "internal unit tests legitimately read/write the deprecated `pub` fields they were designed to verify"
4235 )]
4236 use std::{sync::Arc, time::Duration};
4237
4238 use axum::{
4239 body::Body,
4240 http::{Request, StatusCode, header},
4241 response::IntoResponse,
4242 };
4243 use http_body_util::BodyExt;
4244 use tower::ServiceExt as _;
4245
4246 use super::*;
4247
4248 #[tokio::test]
4251 async fn external_shutdown_bridge_exits_when_internal_token_cancels() {
4252 let external = CancellationToken::new();
4253 let internal = CancellationToken::new();
4254 let bridge = spawn_external_shutdown_bridge(external.clone(), internal.clone());
4255
4256 internal.cancel();
4259
4260 let joined = tokio::time::timeout(Duration::from_secs(2), bridge).await;
4261 assert!(
4262 joined.is_ok(),
4263 "bridge task must exit once the internal token is cancelled, \
4264 otherwise it leaks for the lifetime of the process"
4265 );
4266 }
4267
4268 #[tokio::test]
4269 async fn external_shutdown_bridge_still_forwards_external_cancel() {
4270 let external = CancellationToken::new();
4271 let internal = CancellationToken::new();
4272 let bridge = spawn_external_shutdown_bridge(external.clone(), internal.clone());
4273
4274 external.cancel();
4275
4276 let joined = tokio::time::timeout(Duration::from_secs(2), bridge).await;
4277 assert!(joined.is_ok(), "bridge task must exit on external cancel");
4278 assert!(
4279 internal.is_cancelled(),
4280 "external cancellation must still propagate to the internal token"
4281 );
4282 }
4283
4284 #[test]
4285 fn cancel_on_drop_cancels_its_token() {
4286 let ct = CancellationToken::new();
4287 {
4288 let _guard = CancelOnDrop(ct.clone());
4289 assert!(!ct.is_cancelled());
4290 }
4291 assert!(
4292 ct.is_cancelled(),
4293 "dropping the guard must cancel background startup tasks"
4294 );
4295 }
4296
4297 #[test]
4298 fn validate_rejects_mtls_without_tls() {
4299 for (cert, key) in [
4300 (None, None),
4301 (Some("cert.pem"), None),
4302 (None, Some("key.pem")),
4303 ] {
4304 let mut auth = AuthConfig::with_keys(vec![]);
4305 auth.mtls = Some(valid_mtls_config());
4306 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4307 cfg.tls_cert_path = cert.map(Into::into);
4308 cfg.tls_key_path = key.map(Into::into);
4309
4310 let err = cfg
4311 .validate()
4312 .expect_err("mTLS without both TLS paths must be rejected");
4313 let msg = err.to_string();
4314 assert!(
4315 msg.contains("tls_cert_path") && msg.contains("tls_key_path"),
4316 "cert={cert:?} key={key:?}: {msg}"
4317 );
4318 }
4319 }
4320
4321 #[test]
4322 fn validate_accepts_mtls_with_tls() {
4323 let mut auth = AuthConfig::with_keys(vec![]);
4324 auth.mtls = Some(valid_mtls_config());
4325 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4326 cfg.tls_cert_path = Some("cert.pem".into());
4327 cfg.tls_key_path = Some("key.pem".into());
4328
4329 assert!(cfg.validate().is_ok(), "mTLS with both TLS paths is valid");
4330 }
4331
4332 #[test]
4335 fn server_config_new_defaults() {
4336 let cfg = McpServerConfig::new("0.0.0.0:8443", "test-server", "1.0.0");
4337 assert_eq!(cfg.bind_addr, "0.0.0.0:8443");
4338 assert_eq!(cfg.name, "test-server");
4339 assert_eq!(cfg.version, "1.0.0");
4340 assert!(cfg.tls_cert_path.is_none());
4341 assert!(cfg.tls_key_path.is_none());
4342 assert!(cfg.auth.is_none());
4343 assert!(cfg.rbac.is_none());
4344 assert!(cfg.allowed_origins.is_empty());
4345 assert!(cfg.tool_rate_limit.is_none());
4346 assert!(cfg.readiness_check.is_none());
4347 assert_eq!(cfg.max_request_body, 1024 * 1024);
4348 assert_eq!(cfg.request_timeout, Duration::from_mins(2));
4349 assert_eq!(cfg.shutdown_timeout, Duration::from_secs(30));
4350 assert!(!cfg.log_request_headers);
4351 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(10));
4352 assert_eq!(cfg.max_concurrent_tls_handshakes, 256);
4353 }
4354
4355 #[test]
4356 fn tls_handshake_builders_set_fields() {
4357 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4358 .with_tls_handshake_timeout(Duration::from_secs(3))
4359 .with_max_concurrent_tls_handshakes(64);
4360 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(3));
4361 assert_eq!(cfg.max_concurrent_tls_handshakes, 64);
4362 }
4363
4364 #[test]
4365 fn validate_rejects_zero_tls_handshake_timeout() {
4366 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4367 .with_tls_handshake_timeout(Duration::ZERO);
4368 let err = cfg.validate().expect_err("zero handshake timeout");
4369 assert!(err.to_string().contains("tls_handshake_timeout"));
4370 }
4371
4372 #[test]
4373 fn validate_rejects_zero_max_concurrent_tls_handshakes() {
4374 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4375 .with_max_concurrent_tls_handshakes(0);
4376 let err = cfg.validate().expect_err("zero handshake concurrency");
4377 assert!(err.to_string().contains("max_concurrent_tls_handshakes"));
4378 }
4379
4380 #[test]
4381 fn validate_consumes_and_proves() {
4382 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4384 let validated = cfg.validate().expect("valid config");
4385 assert_eq!(validated.as_inner().name, "test-server");
4387 let raw = validated.into_inner();
4389 assert_eq!(raw.name, "test-server");
4390
4391 let mut bad = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4393 bad.max_request_body = 0;
4394 assert!(bad.validate().is_err(), "zero body cap must fail validate");
4395 }
4396
4397 #[test]
4398 fn validate_rejects_zero_max_concurrent_requests() {
4399 let cfg =
4400 McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_max_concurrent_requests(0);
4401 let err = cfg.validate().expect_err("zero concurrency cap must fail");
4402 assert!(
4403 format!("{err}").contains("max_concurrent_requests"),
4404 "error should mention max_concurrent_requests, got: {err}"
4405 );
4406 }
4407
4408 #[test]
4409 fn validate_rejects_zero_max_tracked_keys() {
4410 let rl = crate::auth::RateLimitConfig {
4413 max_attempts_per_minute: 30,
4414 pre_auth_max_per_minute: None,
4415 max_tracked_keys: 0,
4416 idle_eviction: Duration::from_secs(15 * 60),
4417 burst: None,
4418 pre_auth_burst: None,
4419 key_eviction_policy: KeyEvictionPolicy::default(),
4420 };
4421 let auth_cfg = AuthConfig {
4422 enabled: true,
4423 api_keys: Vec::new(),
4424 mtls: None,
4425 rate_limit: Some(rl),
4426 #[cfg(feature = "oauth")]
4427 oauth: None,
4428 #[cfg(not(feature = "oauth"))]
4429 oauth: None,
4430 };
4431 let cfg = McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_auth(auth_cfg);
4432 let err = cfg.validate().expect_err("zero max_tracked_keys must fail");
4433 assert!(
4434 format!("{err}").contains("max_tracked_keys"),
4435 "error should mention max_tracked_keys, got: {err}"
4436 );
4437 }
4438
4439 #[test]
4440 fn derive_allowed_hosts_includes_public_host() {
4441 let hosts = derive_allowed_hosts("0.0.0.0:8080", Some("https://mcp.example.com/mcp"));
4442 assert!(
4443 hosts.iter().any(|h| h == "mcp.example.com"),
4444 "public_url host must be allowed"
4445 );
4446 }
4447
4448 #[test]
4449 fn derive_allowed_hosts_includes_bind_authority() {
4450 let hosts = derive_allowed_hosts("127.0.0.1:8080", None);
4451 assert!(
4452 hosts.iter().any(|h| h == "127.0.0.1"),
4453 "bind host must be allowed"
4454 );
4455 assert!(
4456 hosts.iter().any(|h| h == "127.0.0.1:8080"),
4457 "bind authority must be allowed"
4458 );
4459 }
4460
4461 #[tokio::test]
4464 async fn healthz_returns_ok_json() {
4465 let resp = healthz().await.into_response();
4466 assert_eq!(resp.status(), StatusCode::OK);
4467 let body = resp.into_body().collect().await.unwrap().to_bytes();
4468 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4469 assert_eq!(json["status"], "ok");
4470 assert!(
4471 json.get("name").is_none(),
4472 "healthz must not expose server name"
4473 );
4474 assert!(
4475 json.get("version").is_none(),
4476 "healthz must not expose version"
4477 );
4478 }
4479
4480 #[tokio::test]
4483 async fn readyz_returns_ok_when_ready() {
4484 let check: ReadinessCheck =
4485 Arc::new(|| Box::pin(async { serde_json::json!({"ready": true, "db": "connected"}) }));
4486 let resp = readyz(check).await.into_response();
4487 assert_eq!(resp.status(), StatusCode::OK);
4488 let body = resp.into_body().collect().await.unwrap().to_bytes();
4489 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4490 assert_eq!(json["ready"], true);
4491 assert!(
4492 json.get("name").is_none(),
4493 "readyz must not expose server name"
4494 );
4495 assert!(
4496 json.get("version").is_none(),
4497 "readyz must not expose version"
4498 );
4499 assert_eq!(json["db"], "connected");
4500 }
4501
4502 #[tokio::test]
4503 async fn readyz_returns_503_when_not_ready() {
4504 let check: ReadinessCheck =
4505 Arc::new(|| Box::pin(async { serde_json::json!({"ready": false}) }));
4506 let resp = readyz(check).await.into_response();
4507 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4508 }
4509
4510 #[tokio::test]
4511 async fn readyz_returns_503_when_ready_missing() {
4512 let check: ReadinessCheck =
4513 Arc::new(|| Box::pin(async { serde_json::json!({"status": "starting"}) }));
4514 let resp = readyz(check).await.into_response();
4515 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4517 }
4518
4519 fn peer_probe_router() -> axum::Router {
4524 async fn probe(req: Request<Body>) -> String {
4525 let ci = req
4526 .extensions()
4527 .get::<ConnectInfo<SocketAddr>>()
4528 .map(|c| c.0.to_string())
4529 .unwrap_or_default();
4530 let pa = req
4531 .extensions()
4532 .get::<PeerAddr>()
4533 .map(|p| p.addr.to_string())
4534 .unwrap_or_default();
4535 format!("{ci}|{pa}")
4536 }
4537 axum::Router::new()
4538 .route("/probe", axum::routing::get(probe))
4539 .layer(axum::middleware::from_fn(|req, next| {
4540 normalize_peer_addr_middleware(None, req, next)
4541 }))
4542 }
4543
4544 async fn body_string(resp: axum::response::Response) -> String {
4545 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
4546 String::from_utf8(bytes.to_vec()).unwrap()
4547 }
4548
4549 #[tokio::test]
4550 async fn normalize_preserves_existing_connect_info_and_mirrors_peer_addr() {
4551 let plain: SocketAddr = "10.0.0.1:1111".parse().unwrap();
4554 let tls: SocketAddr = "10.0.0.2:2222".parse().unwrap();
4555 let req = Request::builder()
4556 .uri("/probe")
4557 .extension(ConnectInfo(plain))
4558 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4559 .body(Body::empty())
4560 .unwrap();
4561 let resp = peer_probe_router().oneshot(req).await.unwrap();
4562 assert_eq!(resp.status(), StatusCode::OK);
4563 assert_eq!(body_string(resp).await, format!("{plain}|{plain}"));
4564 }
4565
4566 #[tokio::test]
4567 async fn normalize_inserts_connect_info_and_peer_addr_from_tls() {
4568 let tls: SocketAddr = "192.168.1.7:50443".parse().unwrap();
4569 let req = Request::builder()
4570 .uri("/probe")
4571 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4572 .body(Body::empty())
4573 .unwrap();
4574 let resp = peer_probe_router().oneshot(req).await.unwrap();
4575 assert_eq!(resp.status(), StatusCode::OK);
4576 assert_eq!(body_string(resp).await, format!("{tls}|{tls}"));
4577 }
4578
4579 #[tokio::test]
4580 async fn normalize_no_op_without_any_connect_info() {
4581 let req = Request::builder()
4582 .uri("/probe")
4583 .body(Body::empty())
4584 .unwrap();
4585 let resp = peer_probe_router().oneshot(req).await.unwrap();
4586 assert_eq!(resp.status(), StatusCode::OK);
4587 assert_eq!(body_string(resp).await, "|");
4588 }
4589
4590 #[tokio::test]
4591 async fn peer_addr_extractor_rejects_when_absent() {
4592 async fn h(peer: PeerAddr) -> String {
4593 peer.addr.to_string()
4594 }
4595 let app = axum::Router::new().route("/p", axum::routing::get(h));
4596 let req = Request::builder().uri("/p").body(Body::empty()).unwrap();
4597 let resp = app.oneshot(req).await.unwrap();
4598 assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
4599 }
4600
4601 #[tokio::test]
4602 async fn peer_addr_extractor_returns_value_when_present() {
4603 async fn h(peer: PeerAddr) -> String {
4604 peer.addr.to_string()
4605 }
4606 let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap();
4607 let app = axum::Router::new().route("/p", axum::routing::get(h));
4608 let req = Request::builder()
4609 .uri("/p")
4610 .extension(PeerAddr::new(addr))
4611 .body(Body::empty())
4612 .unwrap();
4613 let resp = app.oneshot(req).await.unwrap();
4614 assert_eq!(resp.status(), StatusCode::OK);
4615 assert_eq!(body_string(resp).await, addr.to_string());
4616 }
4617
4618 #[tokio::test]
4619 async fn peer_addr_via_extension_extractor() {
4620 async fn h(axum::Extension(peer): axum::Extension<PeerAddr>) -> String {
4621 peer.addr.to_string()
4622 }
4623 let addr: SocketAddr = "127.0.0.1:4242".parse().unwrap();
4624 let app = axum::Router::new().route("/p", axum::routing::get(h));
4625 let req = Request::builder()
4626 .uri("/p")
4627 .extension(PeerAddr::new(addr))
4628 .body(Body::empty())
4629 .unwrap();
4630 let resp = app.oneshot(req).await.unwrap();
4631 assert_eq!(resp.status(), StatusCode::OK);
4632 assert_eq!(body_string(resp).await, addr.to_string());
4633 }
4634
4635 fn limited_router(per_minute: u32) -> axum::Router {
4640 limited_router_with_burst(per_minute, None)
4641 }
4642
4643 fn limited_router_with_burst(per_minute: u32, burst: Option<u32>) -> axum::Router {
4645 limited_router_full(per_minute, burst, &[])
4646 }
4647
4648 fn limited_router_full(
4652 per_minute: u32,
4653 burst: Option<u32>,
4654 exempt_paths: &[&str],
4655 ) -> axum::Router {
4656 let limiter = build_extra_route_rate_limiter_with_policy(
4657 per_minute,
4658 burst,
4659 KeyEvictionPolicy::default(),
4660 NonZeroUsize::new(EXTRA_ROUTE_MAX_TRACKED_KEYS).unwrap_or(NonZeroUsize::MIN),
4661 );
4662 let exempt: Arc<std::collections::HashSet<String>> =
4663 Arc::new(exempt_paths.iter().map(|s| (*s).to_owned()).collect());
4664 axum::Router::new()
4665 .route("/limited", axum::routing::get(|| async { "ok" }))
4666 .route("/exempt", axum::routing::get(|| async { "ok" }))
4667 .layer(axum::middleware::from_fn(move |req, next| {
4668 let l = Arc::clone(&limiter);
4669 let e = Arc::clone(&exempt);
4670 extra_route_rate_limit_middleware(l, e, req, next)
4671 }))
4672 }
4673
4674 fn limited_req(ip: &str) -> Request<Body> {
4675 limited_req_to(ip, "/limited")
4676 }
4677
4678 fn limited_req_to(ip: &str, path: &str) -> Request<Body> {
4679 let addr: SocketAddr = format!("{ip}:40000").parse().unwrap();
4680 Request::builder()
4681 .uri(path)
4682 .extension(ConnectInfo(addr))
4683 .body(Body::empty())
4684 .unwrap()
4685 }
4686
4687 #[tokio::test]
4688 async fn extra_route_limiter_denies_over_quota() {
4689 let app = limited_router(2);
4690 for i in 0..2 {
4691 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4692 assert_eq!(resp.status(), StatusCode::OK, "request {i} should pass");
4693 }
4694 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4695 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4696 let body = body_string(resp).await;
4697 assert!(
4698 body.contains("too many requests to application routes"),
4699 "deny body should match the limiter message, got: {body}"
4700 );
4701 }
4702
4703 fn one_tracked_key() -> NonZeroUsize {
4704 NonZeroUsize::new(1).unwrap_or(NonZeroUsize::MIN)
4705 }
4706
4707 #[tokio::test]
4708 async fn extra_route_limiter_capacity_full_returns_503_without_retry_after() {
4709 let limiter = build_extra_route_rate_limiter_with_policy(
4710 10,
4711 None,
4712 KeyEvictionPolicy::RejectNew,
4713 one_tracked_key(),
4714 );
4715 let exempt = Arc::new(std::collections::HashSet::new());
4716 let app = axum::Router::new()
4717 .route("/limited", axum::routing::get(|| async { "ok" }))
4718 .layer(axum::middleware::from_fn(move |req, next| {
4719 let l = Arc::clone(&limiter);
4720 let e = Arc::clone(&exempt);
4721 extra_route_rate_limit_middleware(l, e, req, next)
4722 }));
4723 let established = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4724 assert_eq!(established.status(), StatusCode::OK);
4725
4726 let denied = app.clone().oneshot(limited_req("10.1.1.2")).await.unwrap();
4727
4728 assert_eq!(denied.status(), StatusCode::SERVICE_UNAVAILABLE);
4729 assert!(denied.headers().get(header::RETRY_AFTER).is_none());
4730 }
4731
4732 #[tokio::test]
4733 async fn extra_route_limiter_isolates_keys() {
4734 let app = limited_router(2);
4735 for _ in 0..2 {
4736 let resp = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4737 assert_eq!(resp.status(), StatusCode::OK);
4738 }
4739 let exhausted = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4740 assert_eq!(exhausted.status(), StatusCode::TOO_MANY_REQUESTS);
4741 let other = app.clone().oneshot(limited_req("10.3.3.3")).await.unwrap();
4743 assert_eq!(other.status(), StatusCode::OK);
4744 }
4745
4746 #[tokio::test]
4747 async fn extra_route_limiter_bounds_requests_without_peer() {
4748 let app = limited_router(1);
4752 let mk = || {
4753 Request::builder()
4754 .uri("/limited")
4755 .body(Body::empty())
4756 .unwrap()
4757 };
4758 let first = app.clone().oneshot(mk()).await.unwrap();
4759 assert_eq!(
4760 first.status(),
4761 StatusCode::OK,
4762 "first request consumes quota"
4763 );
4764 let second = app.clone().oneshot(mk()).await.unwrap();
4765 assert_eq!(
4766 second.status(),
4767 StatusCode::TOO_MANY_REQUESTS,
4768 "unattributable requests must share a bounded bucket, not bypass the limiter"
4769 );
4770 }
4771
4772 #[test]
4773 fn limiter_client_key_falls_back_to_unattributed() {
4774 let empty = axum::http::Extensions::new();
4775 assert_eq!(limiter_client_key(&empty), RateLimitKey::Unattributed);
4776 }
4777
4778 #[test]
4779 fn unattributed_key_is_distinct_from_unspecified_ip() {
4780 let unspecified = RateLimitKey::Ip("0.0.0.0".parse::<IpAddr>().unwrap());
4784 assert_ne!(unspecified, RateLimitKey::Unattributed);
4785
4786 let mut set = std::collections::HashSet::new();
4787 set.insert(unspecified);
4788 set.insert(RateLimitKey::Unattributed);
4789 assert_eq!(set.len(), 2, "the two keys must hash to distinct buckets");
4790 }
4791
4792 #[test]
4793 fn rate_limit_key_display_does_not_fabricate_an_ip() {
4794 assert_eq!(
4795 RateLimitKey::Ip("10.1.2.3".parse::<IpAddr>().unwrap()).to_string(),
4796 "10.1.2.3"
4797 );
4798 assert_eq!(RateLimitKey::Unattributed.to_string(), "unattributed");
4799 }
4800
4801 #[tokio::test]
4802 async fn extra_route_limiter_extracts_tls_conn_info() {
4803 let app = limited_router(2);
4804 let mk = || {
4805 let addr: SocketAddr = "192.168.9.9:55555".parse().unwrap();
4806 Request::builder()
4807 .uri("/limited")
4808 .extension(ConnectInfo(TlsConnInfo::new(addr, None)))
4809 .body(Body::empty())
4810 .unwrap()
4811 };
4812 for _ in 0..2 {
4813 assert_eq!(
4814 app.clone().oneshot(mk()).await.unwrap().status(),
4815 StatusCode::OK
4816 );
4817 }
4818 let resp = app.clone().oneshot(mk()).await.unwrap();
4819 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4820 }
4821
4822 #[tokio::test]
4823 async fn extra_route_limiter_exempt_path_bypasses_quota() {
4824 let app = limited_router_full(1, None, &["/exempt"]);
4827 for i in 0..5 {
4828 let resp = app
4829 .clone()
4830 .oneshot(limited_req_to("10.6.6.6", "/exempt"))
4831 .await
4832 .unwrap();
4833 assert_eq!(resp.status(), StatusCode::OK, "exempt request {i}");
4834 }
4835 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
4837 assert_eq!(resp.status(), StatusCode::OK);
4838 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
4840 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4841 }
4842
4843 #[tokio::test]
4844 async fn extra_route_limiter_exemption_is_raw_exact_match() {
4845 let app = limited_router_full(1, None, &["/exempt"]);
4848 let ok = app
4849 .clone()
4850 .oneshot(limited_req_to("10.7.7.7", "/exempt/"))
4851 .await
4852 .unwrap();
4853 assert_eq!(
4854 ok.status(),
4855 StatusCode::NOT_FOUND,
4856 "variant path routes 404"
4857 );
4858 let denied = app
4860 .clone()
4861 .oneshot(limited_req_to("10.7.7.7", "/limited"))
4862 .await
4863 .unwrap();
4864 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4865 }
4866
4867 #[cfg(feature = "metrics")]
4868 #[tokio::test]
4869 async fn extra_route_limiter_deny_increments_counter_exempt_does_not() {
4870 let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
4871 let app = limited_router_full(1, None, &["/exempt"]);
4872 let mk = |path: &str| {
4873 let addr: SocketAddr = "10.8.8.8:40000".parse().unwrap();
4874 Request::builder()
4875 .uri(path)
4876 .extension(ConnectInfo(addr))
4877 .extension(Arc::clone(&metrics))
4878 .body(Body::empty())
4879 .unwrap()
4880 };
4881 let counter = || {
4882 metrics
4883 .rate_limited_total
4884 .with_label_values(&["extra_route"])
4885 .get()
4886 };
4887 for _ in 0..3 {
4889 assert_eq!(
4890 app.clone().oneshot(mk("/exempt")).await.unwrap().status(),
4891 StatusCode::OK
4892 );
4893 }
4894 assert_eq!(counter(), 0, "exempt requests must not count as denies");
4895 assert_eq!(
4897 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4898 StatusCode::OK
4899 );
4900 assert_eq!(counter(), 0);
4901 assert_eq!(
4902 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4903 StatusCode::TOO_MANY_REQUESTS
4904 );
4905 assert_eq!(counter(), 1, "deny must increment the extra_route label");
4906 }
4907
4908 #[test]
4909 fn validate_rejects_exempt_paths_without_base_knob() {
4910 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4911 .with_extra_route_rate_limit_exempt_paths(["/ok"]);
4912 let err = cfg.validate().expect_err("exempt paths without rate limit");
4913 assert!(err.to_string().contains("requires extra_route_rate_limit"));
4914 }
4915
4916 #[test]
4917 fn validate_rejects_malformed_exempt_paths() {
4918 for bad in ["", "no-slash"] {
4919 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4920 .with_extra_route_rate_limit(10)
4921 .with_extra_route_rate_limit_exempt_paths([bad]);
4922 let err = cfg.validate().expect_err("malformed exempt path");
4923 assert!(
4924 err.to_string()
4925 .contains("must be non-empty and start with '/'"),
4926 "entry {bad:?}: {err}"
4927 );
4928 }
4929 }
4930
4931 #[test]
4932 fn validate_accepts_wellformed_exempt_paths() {
4933 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4934 .with_extra_route_rate_limit(10)
4935 .with_extra_route_rate_limit_exempt_paths(["/.well-known/oauth-authorization-server"]);
4936 assert!(cfg.validate().is_ok());
4937 }
4938
4939 #[test]
4940 fn validate_rejects_zero_extra_route_rate_limit() {
4941 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4942 .with_extra_route_rate_limit(0);
4943 let err = cfg.validate().expect_err("zero extra route rate limit");
4944 assert!(err.to_string().contains("extra_route_rate_limit"));
4945 }
4946
4947 #[tokio::test]
4948 async fn extra_route_limiter_burst_allows_initial_spike() {
4949 let app = limited_router_with_burst(1, Some(3));
4950 for i in 0..3 {
4951 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4952 assert_eq!(resp.status(), StatusCode::OK, "burst request {i}");
4953 }
4954 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4955 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4956 }
4957
4958 #[tokio::test]
4959 async fn extra_route_limiter_deny_sets_retry_after() {
4960 let app = limited_router(1);
4961 let ok = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4962 assert_eq!(ok.status(), StatusCode::OK);
4963 let denied = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4964 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4965 let retry_after = denied
4966 .headers()
4967 .get(header::RETRY_AFTER)
4968 .expect("Retry-After present")
4969 .to_str()
4970 .unwrap()
4971 .parse::<u64>()
4972 .unwrap();
4973 assert!(retry_after >= 1, "delta-seconds must be >= 1");
4974 }
4975
4976 #[test]
4977 fn validate_rejects_zero_burst_knobs() {
4978 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4979 .with_tool_rate_limit(10)
4980 .with_tool_rate_limit_burst(0)
4981 .validate()
4982 .expect_err("zero tool burst");
4983 assert!(err.to_string().contains("tool_rate_limit_burst"));
4984
4985 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4986 .with_extra_route_rate_limit(10)
4987 .with_extra_route_rate_limit_burst(0)
4988 .validate()
4989 .expect_err("zero extra route burst");
4990 assert!(err.to_string().contains("extra_route_rate_limit_burst"));
4991 }
4992
4993 #[test]
4994 fn validate_rejects_orphan_burst_knobs() {
4995 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4996 .with_tool_rate_limit_burst(5)
4997 .validate()
4998 .expect_err("orphan tool burst");
4999 assert!(err.to_string().contains("requires tool_rate_limit"));
5000
5001 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5002 .with_extra_route_rate_limit_burst(5)
5003 .validate()
5004 .expect_err("orphan extra route burst");
5005 assert!(err.to_string().contains("requires extra_route_rate_limit"));
5006 }
5007
5008 #[test]
5009 fn validate_rejects_zero_auth_bursts() {
5010 let auth = AuthConfig::with_keys(vec![])
5011 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_burst(0));
5012 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5013 .with_auth(auth)
5014 .validate()
5015 .expect_err("zero auth burst");
5016 assert!(err.to_string().contains("rate_limit.burst"));
5017
5018 let auth = AuthConfig::with_keys(vec![])
5019 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(0));
5020 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5021 .with_auth(auth)
5022 .validate()
5023 .expect_err("zero pre-auth burst");
5024 assert!(err.to_string().contains("pre_auth_burst"));
5025 }
5026
5027 #[test]
5028 fn validate_rejects_zero_pre_auth_max_per_minute() {
5029 let auth = AuthConfig::with_keys(vec![])
5030 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_max_per_minute(0));
5031 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5032 .with_auth(auth)
5033 .validate()
5034 .expect_err("zero pre-auth rate");
5035 assert!(err.to_string().contains("pre_auth_max_per_minute"));
5036 }
5037
5038 fn valid_mtls_config() -> MtlsConfig {
5039 MtlsConfig {
5040 ca_cert_path: "memory://ca.pem".into(),
5041 required: true,
5042 default_role: "viewer".into(),
5043 crl_enabled: true,
5044 crl_refresh_interval: None,
5045 crl_fetch_timeout: Duration::from_secs(30),
5046 crl_stale_grace: Duration::from_secs(24 * 60 * 60),
5047 crl_deny_on_unavailable: false,
5048 crl_end_entity_only: false,
5049 crl_allow_http: true,
5050 crl_enforce_expiration: true,
5051 crl_max_concurrent_fetches: 4,
5052 crl_max_response_bytes: 5 * 1024 * 1024,
5053 crl_discovery_rate_per_min: 60,
5054 crl_max_host_semaphores: 1024,
5055 crl_max_seen_urls: 4096,
5056 crl_max_cache_entries: 1024,
5057 }
5058 }
5059
5060 #[test]
5061 fn validate_rejects_zero_crl_max_response_bytes() {
5062 let mut mtls = valid_mtls_config();
5063 mtls.crl_max_response_bytes = 0;
5064 let mut auth = AuthConfig::with_keys(vec![]);
5065 auth.mtls = Some(mtls);
5066
5067 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
5070 cfg.tls_cert_path = Some("cert.pem".into());
5071 cfg.tls_key_path = Some("key.pem".into());
5072
5073 let err = cfg.validate().expect_err("zero CRL response cap");
5074 assert!(err.to_string().contains("crl_max_response_bytes"));
5075 }
5076
5077 #[test]
5080 fn validate_accepts_pre_auth_burst_without_explicit_pre_auth_rate() {
5081 let auth = AuthConfig::with_keys(vec![])
5082 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(50));
5083 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
5084 assert!(cfg.validate().is_ok(), "pre_auth_burst has no orphan rule");
5085 }
5086
5087 #[test]
5090 fn trusted_forwarder_max_entries_bounds_are_enforced() {
5091 let cfg = |n: usize| {
5092 McpServerConfig::new("127.0.0.1:8080", "t", "0")
5093 .with_trusted_forwarder_max_entries(n)
5094 .validate()
5095 };
5096 assert!(cfg(0).is_err(), "0 would pin every client to the proxy");
5097 assert!(
5098 cfg(crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES + 1).is_err(),
5099 "above the ceiling would re-open the header-bomb vector"
5100 );
5101 assert!(cfg(1).is_ok());
5102 assert!(cfg(crate::forwarded::MAX_SCANNED_ENTRIES).is_ok());
5103 assert!(cfg(crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES).is_ok());
5104 }
5105
5106 #[test]
5107 fn trusted_forwarder_max_entries_defaults_to_the_module_constant() {
5108 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "0");
5109 assert_eq!(
5110 cfg.trusted_forwarder_max_entries,
5111 crate::forwarded::MAX_SCANNED_ENTRIES
5112 );
5113 }
5114
5115 fn forward_resolver(trusted: &[&str], mode: ForwardedHeaderMode) -> Arc<ForwardResolver> {
5116 Arc::new(ForwardResolver {
5117 trusted: trusted.iter().map(|s| s.parse().unwrap()).collect(),
5118 mode,
5119 max_scanned_entries: crate::forwarded::MAX_SCANNED_ENTRIES,
5120 })
5121 }
5122
5123 fn forwarded_probe_router(resolver: Option<Arc<ForwardResolver>>) -> axum::Router {
5125 async fn probe(req: Request<Body>) -> String {
5126 let pa = req
5127 .extensions()
5128 .get::<PeerAddr>()
5129 .map(|p| p.addr.ip().to_string())
5130 .unwrap_or_default();
5131 let ci = req
5132 .extensions()
5133 .get::<ClientIp>()
5134 .map(|c| c.ip.to_string())
5135 .unwrap_or_default();
5136 format!("{pa}|{ci}")
5137 }
5138 axum::Router::new()
5139 .route("/probe", axum::routing::get(probe))
5140 .layer(axum::middleware::from_fn(move |req, next| {
5141 let r = resolver.clone();
5142 normalize_peer_addr_middleware(r, req, next)
5143 }))
5144 }
5145
5146 fn probe_req(peer: &str, header: Option<(&str, &str)>) -> Request<Body> {
5147 let addr: SocketAddr = peer.parse().unwrap();
5148 let mut builder = Request::builder()
5149 .uri("/probe")
5150 .extension(ConnectInfo(addr));
5151 if let Some((name, value)) = header {
5152 builder = builder.header(name, value);
5153 }
5154 builder.body(Body::empty()).unwrap()
5155 }
5156
5157 #[tokio::test]
5158 async fn client_ip_equals_direct_without_resolver() {
5159 let app = forwarded_probe_router(None);
5160 let resp = app
5161 .oneshot(probe_req(
5162 "10.1.2.3:4444",
5163 Some(("x-forwarded-for", "203.0.113.7")),
5164 ))
5165 .await
5166 .unwrap();
5167 assert_eq!(
5168 body_string(resp).await,
5169 "10.1.2.3|10.1.2.3",
5170 "feature off: header ignored, ClientIp == direct"
5171 );
5172 }
5173
5174 #[tokio::test]
5175 async fn client_ip_resolved_for_trusted_peer() {
5176 let app = forwarded_probe_router(Some(forward_resolver(
5177 &["10.0.0.0/8"],
5178 ForwardedHeaderMode::XForwardedFor,
5179 )));
5180 let resp = app
5181 .oneshot(probe_req(
5182 "10.0.0.1:9999",
5183 Some(("x-forwarded-for", "203.0.113.7")),
5184 ))
5185 .await
5186 .unwrap();
5187 assert_eq!(
5188 body_string(resp).await,
5189 "10.0.0.1|203.0.113.7",
5190 "PeerAddr stays direct while ClientIp resolves"
5191 );
5192 }
5193
5194 #[tokio::test]
5195 async fn client_ip_falls_back_to_direct_on_malformed_header() {
5196 let app = forwarded_probe_router(Some(forward_resolver(
5197 &["10.0.0.0/8"],
5198 ForwardedHeaderMode::XForwardedFor,
5199 )));
5200 let resp = app
5201 .oneshot(probe_req(
5202 "10.0.0.1:9999",
5203 Some(("x-forwarded-for", "not-an-ip")),
5204 ))
5205 .await
5206 .unwrap();
5207 assert_eq!(
5208 body_string(resp).await,
5209 "10.0.0.1|10.0.0.1",
5210 "malformed chain falls back to the direct peer"
5211 );
5212 }
5213
5214 #[test]
5215 fn forwarded_header_mode_deserializes_kebab_case() {
5216 #[derive(serde::Deserialize)]
5217 struct Wrapper {
5218 mode: ForwardedHeaderMode,
5219 }
5220 let w: Wrapper = toml::from_str(r#"mode = "x-forwarded-for""#).unwrap();
5221 assert_eq!(w.mode, ForwardedHeaderMode::XForwardedFor);
5222 let w: Wrapper = toml::from_str(r#"mode = "forwarded""#).unwrap();
5223 assert_eq!(w.mode, ForwardedHeaderMode::Forwarded);
5224 assert!(
5225 toml::from_str::<Wrapper>(r#"mode = "XForwardedFor""#).is_err(),
5226 "PascalCase wire value must be rejected"
5227 );
5228 }
5229
5230 #[test]
5231 fn validate_rejects_bad_trusted_proxy_entry() {
5232 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5233 .with_trusted_proxies(["not-a-cidr"]);
5234 let err = cfg.validate().expect_err("bad CIDR");
5235 assert!(err.to_string().contains("trusted_proxies"));
5236 }
5237
5238 #[test]
5239 fn validate_rejects_zero_prefix_trusted_proxy() {
5240 for entry in ["0.0.0.0/0", "::/0"] {
5241 let cfg =
5242 McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([entry]);
5243 let err = cfg.validate().expect_err("zero-prefix CIDR");
5244 assert!(
5245 err.to_string().contains("prefix length 0"),
5246 "entry {entry}: {err}"
5247 );
5248 }
5249 }
5250
5251 #[test]
5252 fn validate_accepts_cidr_and_bare_ip_proxy_entries() {
5253 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([
5254 "10.0.0.0/8",
5255 "192.0.2.1",
5256 "2001:db8::1",
5257 ]);
5258 assert!(cfg.validate().is_ok(), "CIDRs and bare IPs are accepted");
5259 }
5260
5261 #[test]
5262 fn validate_rejects_forwarded_header_without_proxies() {
5263 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5264 .with_forwarded_header(ForwardedHeaderMode::Forwarded);
5265 let err = cfg.validate().expect_err("mode without proxies");
5266 assert!(err.to_string().contains("requires trusted_proxies"));
5267 }
5268
5269 fn origin_router(origins: Vec<String>, log_request_headers: bool) -> axum::Router {
5273 let allowed: Arc<[String]> = Arc::from(origins);
5274 axum::Router::new()
5275 .route("/test", axum::routing::get(|| async { "ok" }))
5276 .layer(axum::middleware::from_fn(move |req, next| {
5277 let a = Arc::clone(&allowed);
5278 origin_check_middleware(a, log_request_headers, req, next)
5279 }))
5280 }
5281
5282 #[tokio::test]
5283 async fn origin_allowed_passes() {
5284 let app = origin_router(vec!["http://localhost:3000".into()], false);
5285 let req = Request::builder()
5286 .uri("/test")
5287 .header(header::ORIGIN, "http://localhost:3000")
5288 .body(Body::empty())
5289 .unwrap();
5290 let resp = app.oneshot(req).await.unwrap();
5291 assert_eq!(resp.status(), StatusCode::OK);
5292 }
5293
5294 #[tokio::test]
5295 async fn origin_rejected_returns_403() {
5296 let app = origin_router(vec!["http://localhost:3000".into()], false);
5297 let req = Request::builder()
5298 .uri("/test")
5299 .header(header::ORIGIN, "http://evil.com")
5300 .body(Body::empty())
5301 .unwrap();
5302 let resp = app.oneshot(req).await.unwrap();
5303 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5304 }
5305
5306 #[tokio::test]
5307 async fn no_origin_header_passes() {
5308 let app = origin_router(vec!["http://localhost:3000".into()], false);
5309 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5310 let resp = app.oneshot(req).await.unwrap();
5311 assert_eq!(resp.status(), StatusCode::OK);
5312 }
5313
5314 #[tokio::test]
5315 async fn empty_allowlist_rejects_any_origin() {
5316 let app = origin_router(vec![], false);
5317 let req = Request::builder()
5318 .uri("/test")
5319 .header(header::ORIGIN, "http://anything.com")
5320 .body(Body::empty())
5321 .unwrap();
5322 let resp = app.oneshot(req).await.unwrap();
5323 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5324 }
5325
5326 #[tokio::test]
5327 async fn empty_allowlist_passes_without_origin() {
5328 let app = origin_router(vec![], false);
5329 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5330 let resp = app.oneshot(req).await.unwrap();
5331 assert_eq!(resp.status(), StatusCode::OK);
5332 }
5333
5334 #[test]
5335 fn format_request_headers_redacts_sensitive_values() {
5336 let mut headers = axum::http::HeaderMap::new();
5337 headers.insert("authorization", "Bearer secret-token".parse().unwrap());
5338 headers.insert("cookie", "sid=abc".parse().unwrap());
5339 headers.insert("x-request-id", "req-123".parse().unwrap());
5340
5341 let out = format_request_headers_for_log(&headers);
5342 assert!(out.contains("authorization: [REDACTED]"));
5343 assert!(out.contains("cookie: [REDACTED]"));
5344 assert!(out.contains("x-request-id: req-123"));
5345 assert!(!out.contains("secret-token"));
5346 }
5347
5348 #[test]
5349 fn format_request_headers_redacts_forwarding_headers() {
5350 let mut headers = axum::http::HeaderMap::new();
5351 headers.insert("forwarded", "for=203.0.113.9;by=10.1.2.3".parse().unwrap());
5352 headers.insert("x-forwarded-for", "203.0.113.9, 10.1.2.3".parse().unwrap());
5353 headers.insert("x-real-ip", "203.0.113.9".parse().unwrap());
5354 headers.insert("x-request-id", "req-123".parse().unwrap());
5355
5356 let out = format_request_headers_for_log(&headers);
5357 for name in ["forwarded", "x-forwarded-for", "x-real-ip"] {
5358 assert!(
5359 out.contains(&format!("{name}: [REDACTED]")),
5360 "{name} carries client IP / proxy topology and must not reach logs; got {out}"
5361 );
5362 }
5363 assert!(
5364 !out.contains("203.0.113.9") && !out.contains("10.1.2.3"),
5365 "no forwarded address may survive redaction; got {out}"
5366 );
5367 assert!(out.contains("x-request-id: req-123"));
5368 }
5369
5370 fn security_router(is_tls: bool) -> axum::Router {
5373 security_router_with(is_tls, SecurityHeadersConfig::default())
5374 }
5375
5376 fn security_router_with(is_tls: bool, cfg: SecurityHeadersConfig) -> axum::Router {
5377 let cfg = Arc::new(cfg);
5378 axum::Router::new()
5379 .route("/test", axum::routing::get(|| async { "ok" }))
5380 .layer(axum::middleware::from_fn(move |req, next| {
5381 let c = Arc::clone(&cfg);
5382 security_headers_middleware(is_tls, c, req, next)
5383 }))
5384 }
5385
5386 #[tokio::test]
5387 async fn security_headers_set_on_response() {
5388 let app = security_router(false);
5389 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5390 let resp = app.oneshot(req).await.unwrap();
5391 assert_eq!(resp.status(), StatusCode::OK);
5392
5393 let h = resp.headers();
5394 assert_eq!(h.get("x-content-type-options").unwrap(), "nosniff");
5395 assert_eq!(h.get("x-frame-options").unwrap(), "deny");
5396 assert_eq!(h.get("cache-control").unwrap(), "no-store, max-age=0");
5397 assert_eq!(h.get("referrer-policy").unwrap(), "no-referrer");
5398 assert_eq!(h.get("cross-origin-opener-policy").unwrap(), "same-origin");
5399 assert_eq!(
5400 h.get("cross-origin-resource-policy").unwrap(),
5401 "same-origin"
5402 );
5403 assert_eq!(
5404 h.get("cross-origin-embedder-policy").unwrap(),
5405 "require-corp"
5406 );
5407 assert_eq!(h.get("x-permitted-cross-domain-policies").unwrap(), "none");
5408 assert!(
5409 h.get("permissions-policy")
5410 .unwrap()
5411 .to_str()
5412 .unwrap()
5413 .contains("camera=()"),
5414 "permissions-policy must restrict browser features"
5415 );
5416 assert_eq!(
5417 h.get("content-security-policy").unwrap(),
5418 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
5419 );
5420 assert_eq!(h.get("x-dns-prefetch-control").unwrap(), "off");
5421 assert!(h.get("strict-transport-security").is_none());
5423 }
5424
5425 #[tokio::test]
5426 async fn hsts_set_when_tls_enabled() {
5427 let app = security_router(true);
5428 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5429 let resp = app.oneshot(req).await.unwrap();
5430
5431 let hsts = resp.headers().get("strict-transport-security").unwrap();
5432 assert!(
5433 hsts.to_str().unwrap().contains("max-age=63072000"),
5434 "HSTS must set 2-year max-age"
5435 );
5436 }
5437
5438 #[tokio::test]
5439 async fn default_csp_matches_guideline() {
5440 let app = security_router(false);
5441 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5442 let resp = app.oneshot(req).await.unwrap();
5443 assert_eq!(
5444 resp.headers().get("content-security-policy").unwrap(),
5445 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
5446 );
5447 }
5448
5449 #[tokio::test]
5450 async fn operator_csp_override_still_wins() {
5451 let cfg = SecurityHeadersConfig {
5452 content_security_policy: Some("default-src 'self'".into()),
5453 ..SecurityHeadersConfig::default()
5454 };
5455 let app = security_router_with(false, cfg);
5456 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5457 let resp = app.oneshot(req).await.unwrap();
5458 assert_eq!(
5459 resp.headers().get("content-security-policy").unwrap(),
5460 "default-src 'self'"
5461 );
5462 }
5463
5464 fn check_with_security_headers(
5470 headers: SecurityHeadersConfig,
5471 ) -> Result<(), RmcpServerKitError> {
5472 let cfg =
5473 McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_security_headers(headers);
5474 cfg.check()
5475 }
5476
5477 #[test]
5478 fn security_headers_config_default_validates() {
5479 check_with_security_headers(SecurityHeadersConfig::default())
5480 .expect("default SecurityHeadersConfig must validate");
5481 }
5482
5483 #[test]
5484 fn security_headers_config_validate_accepts_empty_string() {
5485 let h = SecurityHeadersConfig {
5487 x_content_type_options: Some(String::new()),
5488 x_frame_options: Some(String::new()),
5489 cache_control: Some(String::new()),
5490 referrer_policy: Some(String::new()),
5491 cross_origin_opener_policy: Some(String::new()),
5492 cross_origin_resource_policy: Some(String::new()),
5493 cross_origin_embedder_policy: Some(String::new()),
5494 permissions_policy: Some(String::new()),
5495 x_permitted_cross_domain_policies: Some(String::new()),
5496 content_security_policy: Some(String::new()),
5497 x_dns_prefetch_control: Some(String::new()),
5498 strict_transport_security: Some(String::new()),
5499 };
5500 check_with_security_headers(h).expect("Some(\"\") on every field must validate (omit-all)");
5501 }
5502
5503 #[test]
5504 fn security_headers_config_validate_rejects_bad_value() {
5505 let h = SecurityHeadersConfig {
5507 referrer_policy: Some("\u{0007}".into()),
5508 ..SecurityHeadersConfig::default()
5509 };
5510 let err = check_with_security_headers(h)
5511 .expect_err("control char in referrer_policy must reject");
5512 let msg = err.to_string();
5513 assert!(
5514 msg.contains("referrer_policy"),
5515 "error must name the offending field, got: {msg}"
5516 );
5517 }
5518
5519 #[test]
5520 fn security_headers_config_validate_rejects_hsts_preload() {
5521 let h = SecurityHeadersConfig {
5522 strict_transport_security: Some("max-age=63072000; includeSubDomains; preload".into()),
5523 ..SecurityHeadersConfig::default()
5524 };
5525 let err = check_with_security_headers(h).expect_err("HSTS with preload must reject");
5526 let msg = err.to_string();
5527 assert!(
5528 msg.contains("strict_transport_security"),
5529 "error must name the field, got: {msg}"
5530 );
5531 assert!(
5532 msg.to_lowercase().contains("preload"),
5533 "error must mention `preload`, got: {msg}"
5534 );
5535 }
5536
5537 #[test]
5538 fn security_headers_config_validate_rejects_hsts_preload_uppercase() {
5539 let h = SecurityHeadersConfig {
5541 strict_transport_security: Some("max-age=600; PRELOAD".into()),
5542 ..SecurityHeadersConfig::default()
5543 };
5544 check_with_security_headers(h).expect_err("HSTS preload check must be case-insensitive");
5545 }
5546
5547 #[tokio::test]
5548 async fn security_headers_override_honored() {
5549 let h = SecurityHeadersConfig {
5551 x_frame_options: Some("SAMEORIGIN".into()),
5552 ..SecurityHeadersConfig::default()
5553 };
5554 let app = security_router_with(false, h);
5555 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5556 let resp = app.oneshot(req).await.unwrap();
5557 assert_eq!(resp.status(), StatusCode::OK);
5558
5559 let xfo = resp.headers().get("x-frame-options").unwrap();
5560 assert_eq!(xfo, "SAMEORIGIN");
5561 }
5562
5563 #[tokio::test]
5564 async fn security_headers_empty_string_omits() {
5565 let h = SecurityHeadersConfig {
5567 referrer_policy: Some(String::new()),
5568 ..SecurityHeadersConfig::default()
5569 };
5570 let app = security_router_with(false, h);
5571 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5572 let resp = app.oneshot(req).await.unwrap();
5573 assert_eq!(resp.status(), StatusCode::OK);
5574
5575 assert!(
5576 resp.headers().get("referrer-policy").is_none(),
5577 "Some(\"\") must omit the header"
5578 );
5579 assert_eq!(
5581 resp.headers().get("x-content-type-options").unwrap(),
5582 "nosniff"
5583 );
5584 }
5585
5586 #[tokio::test]
5587 async fn security_headers_hsts_only_when_tls() {
5588 let h = SecurityHeadersConfig {
5590 strict_transport_security: Some("max-age=600".into()),
5591 ..SecurityHeadersConfig::default()
5592 };
5593 let app = security_router_with(false, h);
5594 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5595 let resp = app.oneshot(req).await.unwrap();
5596 assert!(
5597 resp.headers().get("strict-transport-security").is_none(),
5598 "HSTS must remain absent on plaintext deployments even with override"
5599 );
5600 }
5601
5602 #[cfg(feature = "oauth")]
5605 #[tokio::test]
5606 async fn oauth_token_cache_headers_set_pragma_and_vary() {
5607 let app = axum::Router::new()
5608 .route("/token", axum::routing::post(|| async { "{}" }))
5609 .layer(axum::middleware::from_fn(
5610 oauth_token_cache_headers_middleware,
5611 ));
5612 let req = Request::builder()
5613 .method("POST")
5614 .uri("/token")
5615 .body(Body::from("{}"))
5616 .unwrap();
5617 let resp = app.oneshot(req).await.unwrap();
5618 assert_eq!(resp.status(), StatusCode::OK);
5619
5620 let h = resp.headers();
5621 assert_eq!(
5622 h.get("pragma").unwrap(),
5623 "no-cache",
5624 "RFC 6749 §5.1: token responses must set Pragma: no-cache"
5625 );
5626 let vary_values: Vec<String> = h
5627 .get_all("vary")
5628 .iter()
5629 .filter_map(|v| v.to_str().ok().map(str::to_owned))
5630 .collect();
5631 assert!(
5632 vary_values
5633 .iter()
5634 .any(|v| v.eq_ignore_ascii_case("Authorization")),
5635 "RFC 6750 §5.4: Vary must include Authorization, got {vary_values:?}"
5636 );
5637 }
5638
5639 #[cfg(feature = "oauth")]
5640 #[tokio::test]
5641 async fn oauth_token_cache_headers_preserve_existing_vary() {
5642 let app = axum::Router::new()
5645 .route(
5646 "/token",
5647 axum::routing::post(|| async {
5648 axum::response::Response::builder()
5649 .header("vary", "Accept-Encoding")
5650 .body(Body::from("{}"))
5651 .unwrap()
5652 }),
5653 )
5654 .layer(axum::middleware::from_fn(
5655 oauth_token_cache_headers_middleware,
5656 ));
5657 let req = Request::builder()
5658 .method("POST")
5659 .uri("/token")
5660 .body(Body::empty())
5661 .unwrap();
5662 let resp = app.oneshot(req).await.unwrap();
5663
5664 let vary: Vec<String> = resp
5665 .headers()
5666 .get_all("vary")
5667 .iter()
5668 .filter_map(|v| v.to_str().ok().map(str::to_owned))
5669 .collect();
5670 assert!(
5671 vary.iter().any(|v| v.contains("Accept-Encoding")),
5672 "must preserve pre-existing Vary value, got {vary:?}"
5673 );
5674 assert!(
5675 vary.iter().any(|v| v.contains("Authorization")),
5676 "must append Authorization to Vary, got {vary:?}"
5677 );
5678 }
5679
5680 #[test]
5683 fn version_omits_build_fingerprint_by_default() {
5684 let v = version_payload("my-server", "1.2.3", false);
5685 assert_eq!(v["name"], "my-server");
5686 assert_eq!(v["version"], "1.2.3");
5687 assert!(v["rmcp_server_kit_version"].is_string());
5688 assert!(
5689 v.get("build_git_sha").is_none(),
5690 "build sha must be hidden by default"
5691 );
5692 assert!(v.get("build_timestamp").is_none());
5693 assert!(v.get("rust_version").is_none());
5694 }
5695
5696 #[test]
5697 fn version_exposes_all_when_enabled() {
5698 let v = version_payload("my-server", "1.2.3", true);
5699 assert!(v["build_git_sha"].is_string());
5700 assert!(v["build_timestamp"].is_string());
5701 assert!(v["rust_version"].is_string());
5702 assert!(v["rmcp_server_kit_version"].is_string());
5703 }
5704
5705 #[tokio::test]
5708 async fn concurrency_limit_layer_composes_and_serves() {
5709 let app = axum::Router::new()
5713 .route("/ok", axum::routing::get(|| async { "ok" }))
5714 .layer(
5715 tower::ServiceBuilder::new()
5716 .layer(axum::error_handling::HandleErrorLayer::new(
5717 |_err: tower::BoxError| async { StatusCode::SERVICE_UNAVAILABLE },
5718 ))
5719 .layer(tower::load_shed::LoadShedLayer::new())
5720 .layer(tower::limit::ConcurrencyLimitLayer::new(4)),
5721 );
5722 let resp = app
5723 .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap())
5724 .await
5725 .unwrap();
5726 assert_eq!(resp.status(), StatusCode::OK);
5727 }
5728
5729 #[tokio::test]
5732 async fn compression_layer_gzip_encodes_response() {
5733 use tower_http::compression::Predicate as _;
5734
5735 let big_body = "a".repeat(4096);
5736 let app = axum::Router::new()
5737 .route(
5738 "/big",
5739 axum::routing::get(move || {
5740 let body = big_body.clone();
5741 async move { body }
5742 }),
5743 )
5744 .layer(
5745 tower_http::compression::CompressionLayer::new()
5746 .gzip(true)
5747 .br(true)
5748 .compress_when(
5749 tower_http::compression::DefaultPredicate::new()
5750 .and(tower_http::compression::predicate::SizeAbove::new(1024)),
5751 ),
5752 );
5753
5754 let req = Request::builder()
5755 .uri("/big")
5756 .header(header::ACCEPT_ENCODING, "gzip")
5757 .body(Body::empty())
5758 .unwrap();
5759 let resp = app.oneshot(req).await.unwrap();
5760 assert_eq!(resp.status(), StatusCode::OK);
5761 assert_eq!(
5762 resp.headers().get(header::CONTENT_ENCODING).unwrap(),
5763 "gzip"
5764 );
5765 }
5766
5767 #[tokio::test]
5770 async fn tls_handshake_timeout_reaps_idle_connections() {
5771 use tokio::io::AsyncReadExt as _;
5772
5773 let _ = rustls::crypto::ring::default_provider().install_default();
5774
5775 let key = rcgen::KeyPair::generate().expect("generate key");
5777 let cert = rcgen::CertificateParams::new(vec!["localhost".to_owned()])
5778 .expect("cert params")
5779 .self_signed(&key)
5780 .expect("self-signed cert");
5781 let dir = std::env::temp_dir().join(format!(
5782 "rmcp-server-kit-hs-timeout-{}",
5783 std::time::SystemTime::now()
5784 .duration_since(std::time::UNIX_EPOCH)
5785 .expect("clock after epoch")
5786 .as_nanos()
5787 ));
5788 tokio::fs::create_dir_all(&dir).await.expect("temp dir");
5789 let cert_path = dir.join("server.crt");
5790 let key_path = dir.join("server.key");
5791 tokio::fs::write(&cert_path, cert.pem())
5792 .await
5793 .expect("write cert");
5794 tokio::fs::write(&key_path, key.serialize_pem())
5795 .await
5796 .expect("write key");
5797
5798 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
5799 let tls = TlsListener::new(
5800 listener,
5801 &cert_path,
5802 &key_path,
5803 None,
5804 None,
5805 Duration::from_millis(200),
5806 8, )
5808 .expect("tls listener");
5809 let addr = axum::serve::Listener::local_addr(&tls).expect("local addr");
5810
5811 let mut idle = tokio::net::TcpStream::connect(addr).await.expect("connect");
5815 let mut buf = [0_u8; 16];
5816 let read = tokio::time::timeout(Duration::from_secs(2), idle.read(&mut buf))
5817 .await
5818 .expect("server must reap the idle handshake within its timeout");
5819 match read {
5820 Ok(0) | Err(_) => {} Ok(n) => panic!("unexpected {n} bytes from server during reaped handshake"),
5822 }
5823
5824 drop(tls);
5825 }
5826
5827 fn assert_owasp_headers(resp: &axum::response::Response, ctx: &str) {
5830 let h = resp.headers();
5831 assert!(
5832 h.contains_key("x-content-type-options"),
5833 "{ctx}: missing X-Content-Type-Options"
5834 );
5835 assert!(
5836 h.contains_key("x-frame-options"),
5837 "{ctx}: missing X-Frame-Options"
5838 );
5839 assert!(
5840 h.contains_key("strict-transport-security"),
5841 "{ctx}: missing Strict-Transport-Security"
5842 );
5843 assert!(
5844 h.contains_key(header::CONTENT_SECURITY_POLICY),
5845 "{ctx}: missing Content-Security-Policy"
5846 );
5847 }
5848
5849 fn m5_router(configure: impl FnOnce(&mut McpServerConfig)) -> axum::Router {
5850 #[derive(Clone)]
5851 struct H;
5852 impl ServerHandler for H {}
5853 let mut config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0")
5857 .with_allowed_origins(["http://good.example"])
5858 .with_tls("unused.crt", "unused.key");
5859 configure(&mut config);
5860 let (router, _params) = build_app_router(config, || H).expect("build_app_router");
5861 router
5862 }
5863
5864 #[test]
5869 #[should_panic(expected = "Overlapping method route")]
5870 fn extra_router_exact_overlap_with_framework_route_panics() {
5871 #[derive(Clone)]
5872 struct H;
5873 impl ServerHandler for H {}
5874 let config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_extra_router(
5875 axum::Router::new().route("/healthz", axum::routing::get(|| async { "mine" })),
5876 );
5877 let _ = build_app_router(config, || H);
5878 }
5879
5880 #[test]
5884 fn extra_router_non_overlapping_path_under_framework_prefix_is_accepted() {
5885 #[derive(Clone)]
5886 struct H;
5887 impl ServerHandler for H {}
5888 let config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_extra_router(
5889 axum::Router::new().route("/admin/custom", axum::routing::get(|| async { "mine" })),
5890 );
5891 assert!(
5892 build_app_router(config, || H).is_ok(),
5893 "non-overlapping path under a framework prefix must merge cleanly"
5894 );
5895 }
5896
5897 #[tokio::test]
5898 async fn headers_on_rejected_origin_403() {
5899 let app = m5_router(|_| {});
5900 let req = Request::builder()
5901 .uri("/healthz")
5902 .header(header::ORIGIN, "http://evil.example")
5903 .body(Body::empty())
5904 .unwrap();
5905 let resp = app.oneshot(req).await.unwrap();
5906 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5907 assert_owasp_headers(&resp, "origin-403");
5908 }
5909
5910 #[tokio::test]
5911 async fn headers_on_cors_preflight() {
5912 let app = m5_router(|_| {});
5913 let req = Request::builder()
5914 .method(axum::http::Method::OPTIONS)
5915 .uri("/mcp")
5916 .header(header::ORIGIN, "http://good.example")
5917 .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
5918 .body(Body::empty())
5919 .unwrap();
5920 let resp = app.oneshot(req).await.unwrap();
5921 assert_owasp_headers(&resp, "cors-preflight");
5922 }
5923
5924 #[tokio::test]
5925 async fn headers_on_404_fallback() {
5926 let app = m5_router(|_| {});
5927 let req = Request::builder()
5928 .uri("/no-such-route")
5929 .body(Body::empty())
5930 .unwrap();
5931 let resp = app.oneshot(req).await.unwrap();
5932 assert_eq!(resp.status(), StatusCode::NOT_FOUND);
5933 assert_owasp_headers(&resp, "404-fallback");
5934 }
5935
5936 #[tokio::test]
5937 async fn headers_on_overload_503() {
5938 let app = m5_router(|c| c.max_concurrent_requests = Some(0));
5941 let req = Request::builder()
5942 .uri("/healthz")
5943 .body(Body::empty())
5944 .unwrap();
5945 let resp = app.oneshot(req).await.unwrap();
5946 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
5947 assert_owasp_headers(&resp, "overload-503");
5948 }
5949
5950 #[cfg(feature = "oauth")]
5953 fn m6_auth_state() -> (Arc<AuthState>, String, String) {
5954 let (admin_token, admin_hash) = crate::auth::generate_api_key().unwrap();
5955 let (viewer_token, viewer_hash) = crate::auth::generate_api_key().unwrap();
5956 let state = Arc::new(AuthState {
5957 api_keys: ArcSwap::from_pointee(vec![
5958 crate::auth::ApiKeyEntry::new("admin-key", admin_hash, "admin"),
5959 crate::auth::ApiKeyEntry::new("viewer-key", viewer_hash, "viewer"),
5960 ]),
5961 rate_limiter: None,
5962 pre_auth_limiter: None,
5963 jwks_cache: None,
5964 seen_identities: crate::auth::SeenIdentitySet::new(),
5965 counters: crate::auth::AuthCounters::default(),
5966 resource_metadata_url: None,
5967 });
5968 (state, admin_token, viewer_token)
5969 }
5970
5971 #[cfg(feature = "oauth")]
5972 fn m6_admin_router(state: &Arc<AuthState>) -> axum::Router {
5973 let proxy = crate::oauth::OAuthProxyConfig::builder(
5974 "https://idp.example/authorize",
5975 "https://idp.example/token",
5976 "client",
5977 )
5978 .introspection_url("http://127.0.0.1:1/introspect")
5979 .revocation_url("http://127.0.0.1:1/revoke")
5980 .expose_admin_endpoints(true)
5981 .require_auth_on_admin_endpoints(true)
5982 .build();
5983 let http = crate::oauth::OauthHttpClient::new().expect("oauth http client");
5984 build_oauth_admin_router(&proxy, http, Some(state), "admin").expect("admin router")
5985 }
5986
5987 #[cfg(feature = "oauth")]
5988 fn m6_req(path: &str, token: &str) -> Request<Body> {
5989 Request::builder()
5990 .method(axum::http::Method::POST)
5991 .uri(path)
5992 .header(header::AUTHORIZATION, format!("Bearer {token}"))
5993 .body(Body::from("token=abc"))
5994 .unwrap()
5995 }
5996
5997 #[cfg(feature = "oauth")]
5998 #[tokio::test]
5999 async fn oauth_proxy_admin_requires_admin_role() {
6000 let (state, _admin, viewer) = m6_auth_state();
6001 for path in ["/introspect", "/revoke"] {
6002 let app = m6_admin_router(&state);
6003 let resp = app.oneshot(m6_req(path, &viewer)).await.unwrap();
6004 assert_eq!(
6005 resp.status(),
6006 StatusCode::FORBIDDEN,
6007 "an authenticated viewer must be rejected with 403 on {path}"
6008 );
6009 }
6010 }
6011
6012 #[cfg(feature = "oauth")]
6013 #[tokio::test]
6014 async fn oauth_proxy_admin_allows_admin_role() {
6015 let (state, admin, _viewer) = m6_auth_state();
6016 for path in ["/introspect", "/revoke"] {
6017 let app = m6_admin_router(&state);
6018 let resp = app.oneshot(m6_req(path, &admin)).await.unwrap();
6019 assert_ne!(
6023 resp.status(),
6024 StatusCode::FORBIDDEN,
6025 "an authenticated admin must pass the role gate on {path}"
6026 );
6027 assert_ne!(
6028 resp.status(),
6029 StatusCode::UNAUTHORIZED,
6030 "an authenticated admin must pass the auth gate on {path}"
6031 );
6032 }
6033 }
6034
6035 #[cfg(feature = "metrics")]
6043 mod metrics_labels_bounded {
6044 use super::*;
6045
6046 fn labels_for(method: &str, uri: &str) -> (&'static str, String) {
6047 let req = Request::builder()
6048 .method(method)
6049 .uri(uri)
6050 .body(Body::empty())
6051 .unwrap();
6052 metrics_labels(&req)
6053 }
6054
6055 #[test]
6056 fn many_unmatched_paths_collapse_to_one_label() {
6057 let mut seen = std::collections::HashSet::new();
6058 for i in 0..500 {
6059 let (_, path) = labels_for("GET", &format!("/nonexistent-{i}"));
6060 seen.insert(path);
6061 }
6062 assert_eq!(
6063 seen.len(),
6064 1,
6065 "unmatched paths must collapse to a single label, got {seen:?}"
6066 );
6067 assert!(seen.contains("<unmatched>"));
6068 }
6069
6070 #[test]
6071 fn nested_mcp_paths_collapse_to_the_mount_point() {
6072 let mut seen = std::collections::HashSet::new();
6073 for i in 0..200 {
6074 let (_, path) = labels_for("POST", &format!("/mcp/{i}"));
6075 seen.insert(path);
6076 }
6077 let (_, root) = labels_for("POST", "/mcp");
6078 seen.insert(root);
6079 assert_eq!(
6080 seen.len(),
6081 1,
6082 "nested /mcp paths must collapse to one label, got {seen:?}"
6083 );
6084 assert!(seen.contains("/mcp"));
6085 }
6086
6087 #[test]
6088 fn unusual_methods_collapse_to_one_bucket() {
6089 let mut seen = std::collections::HashSet::new();
6090 for verb in ["FROBNICATE", "WIBBLE", "QUUX", "M-SEARCH"] {
6091 let (method, _) = labels_for(verb, "/healthz");
6092 seen.insert(method);
6093 }
6094 assert_eq!(seen, std::collections::HashSet::from(["OTHER"]));
6095 }
6096
6097 #[test]
6098 fn known_methods_keep_their_identity() {
6099 for verb in ["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"] {
6100 let (method, _) = labels_for(verb, "/healthz");
6101 assert_eq!(method, verb);
6102 }
6103 }
6104
6105 #[test]
6106 fn raw_path_never_leaks_into_a_label() {
6107 let (_, path) = labels_for("GET", "/secret-token-abc123");
6108 assert!(
6109 !path.contains("secret-token"),
6110 "raw request path must never become a label value: {path}"
6111 );
6112 }
6113 }
6114}