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,
22 session::{EventStore, SessionStore, local::LocalSessionManager},
23 },
24};
25use rustls::RootCertStore;
26use secrecy::{ExposeSecret as _, SecretString};
27use tokio::{
28 net::TcpListener,
29 sync::{Semaphore, mpsc},
30};
31use tokio_util::sync::CancellationToken;
32
33use crate::{
34 auth::{
35 AuthConfig, AuthIdentity, AuthState, MtlsConfig, TlsConnInfo, auth_middleware,
36 build_rate_limiter, extract_mtls_identity,
37 },
38 bounded_limiter::{BoundedKeyedLimiter, BoundedLimiterDeny, KeyEvictionPolicy},
39 error::RmcpServerKitError,
40 mtls_revocation::{self, CrlSet, DynamicClientCertVerifier},
41 rbac::{RbacPolicy, ToolRateLimiter, build_tool_rate_limiter_with_policy, rbac_middleware},
42 rbac_context::RbacContextHandler,
43 session_binding::{
44 configured_session_binding_secret, process_session_binding_secret,
45 session_binding_middleware,
46 },
47};
48
49#[allow(
53 clippy::needless_pass_by_value,
54 reason = "consumed at .map_err(anyhow_to_startup) call sites; by-value matches the closure shape"
55)]
56fn anyhow_to_startup(e: anyhow::Error) -> RmcpServerKitError {
57 RmcpServerKitError::Startup(format!("{e:#}"))
58}
59
60#[allow(
66 clippy::needless_pass_by_value,
67 reason = "consumed at .map_err(|e| io_to_startup(...)) call sites; by-value matches the closure shape"
68)]
69fn io_to_startup(op: &str, e: std::io::Error) -> RmcpServerKitError {
70 RmcpServerKitError::Startup(format!("{op}: {e}"))
71}
72
73pub type ReadinessCheck =
78 Arc<dyn Fn() -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> + Send + Sync>;
79
80#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
125#[non_exhaustive]
126pub struct PeerAddr {
127 pub addr: SocketAddr,
129}
130
131impl PeerAddr {
132 #[must_use]
135 pub(crate) const fn new(addr: SocketAddr) -> Self {
136 Self { addr }
137 }
138}
139
140impl<S: Send + Sync> axum::extract::FromRequestParts<S> for PeerAddr {
149 type Rejection = (axum::http::StatusCode, &'static str);
150
151 #[allow(
152 clippy::unused_async_trait_impl,
153 reason = "async is mandated by the axum FromRequestParts trait signature; this impl only reads a request extension synchronously"
154 )]
155 async fn from_request_parts(
156 parts: &mut axum::http::request::Parts,
157 _state: &S,
158 ) -> Result<Self, Self::Rejection> {
159 parts.extensions.get::<Self>().copied().ok_or((
160 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
161 "peer address unavailable: not running under rmcp-server-kit serve()",
162 ))
163 }
164}
165
166#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
189#[non_exhaustive]
190pub struct ClientIp {
191 pub ip: IpAddr,
193}
194
195impl ClientIp {
196 #[must_use]
199 pub(crate) const fn new(ip: IpAddr) -> Self {
200 Self { ip }
201 }
202}
203
204#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Deserialize)]
209#[serde(rename_all = "kebab-case")]
210#[non_exhaustive]
211pub enum ForwardedHeaderMode {
212 XForwardedFor,
214 Forwarded,
216}
217
218struct ForwardResolver {
221 trusted: Vec<ipnet::IpNet>,
222 mode: ForwardedHeaderMode,
223 max_scanned_entries: usize,
224}
225
226#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
247#[serde(default)]
248#[serde(deny_unknown_fields)]
249#[non_exhaustive]
250pub struct SecurityHeadersConfig {
251 pub x_content_type_options: Option<String>,
253 pub x_frame_options: Option<String>,
255 pub cache_control: Option<String>,
257 pub referrer_policy: Option<String>,
259 pub cross_origin_opener_policy: Option<String>,
261 pub cross_origin_resource_policy: Option<String>,
263 pub cross_origin_embedder_policy: Option<String>,
265 pub permissions_policy: Option<String>,
268 pub x_permitted_cross_domain_policies: Option<String>,
270 pub content_security_policy: Option<String>,
273 pub x_dns_prefetch_control: Option<String>,
275 pub strict_transport_security: Option<String>,
280}
281
282#[allow(
284 missing_debug_implementations,
285 reason = "contains callback/trait objects that don't impl Debug"
286)]
287#[allow(
288 clippy::struct_excessive_bools,
289 reason = "server configuration naturally has many boolean feature flags"
290)]
291#[non_exhaustive]
292pub struct McpServerConfig {
293 #[deprecated(
295 since = "0.13.0",
296 note = "use McpServerConfig::new() / with_bind_addr(); direct field access will become pub(crate) in a future major release"
297 )]
298 pub bind_addr: String,
299 #[deprecated(
301 since = "0.13.0",
302 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
303 )]
304 pub name: String,
305 #[deprecated(
307 since = "0.13.0",
308 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
309 )]
310 pub version: String,
311 #[deprecated(
313 since = "0.13.0",
314 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
315 )]
316 pub tls_cert_path: Option<PathBuf>,
317 #[deprecated(
319 since = "0.13.0",
320 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
321 )]
322 pub tls_key_path: Option<PathBuf>,
323 #[deprecated(
326 since = "0.13.0",
327 note = "use McpServerConfig::with_auth(); direct field access will become pub(crate) in a future major release"
328 )]
329 pub auth: Option<AuthConfig>,
330 #[deprecated(
333 since = "0.13.0",
334 note = "use McpServerConfig::with_rbac(); direct field access will become pub(crate) in a future major release"
335 )]
336 pub rbac: Option<Arc<RbacPolicy>>,
337 pub tool_list_filtering: bool,
340 #[deprecated(
346 since = "0.13.0",
347 note = "use McpServerConfig::with_allowed_origins(); direct field access will become pub(crate) in a future major release"
348 )]
349 pub allowed_origins: Vec<String>,
350 #[deprecated(
353 since = "0.13.0",
354 note = "use McpServerConfig::with_tool_rate_limit(); direct field access will become pub(crate) in a future major release"
355 )]
356 pub tool_rate_limit: Option<u32>,
357 #[deprecated(
363 since = "1.12.0",
364 note = "use McpServerConfig::with_tool_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
365 )]
366 pub tool_rate_limit_burst: Option<u32>,
367 #[deprecated(
380 since = "1.11.0",
381 note = "use McpServerConfig::with_extra_route_rate_limit(); direct field access will become pub(crate) in a future major release"
382 )]
383 pub extra_route_rate_limit: Option<u32>,
384 #[deprecated(
391 since = "1.12.0",
392 note = "use McpServerConfig::with_extra_route_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
393 )]
394 pub extra_route_rate_limit_burst: Option<u32>,
395 #[deprecated(
408 since = "1.14.0",
409 note = "use McpServerConfig::with_extra_route_rate_limit_exempt_paths(); direct field access will become pub(crate) in a future major release"
410 )]
411 pub extra_route_rate_limit_exempt_paths: Vec<String>,
412
413 pub key_eviction_policy: KeyEvictionPolicy,
415
416 pub trusted_forwarder_max_entries: usize,
423 #[deprecated(
431 since = "1.13.0",
432 note = "use McpServerConfig::with_trusted_proxies(); direct field access will become pub(crate) in a future major release"
433 )]
434 pub trusted_proxies: Vec<String>,
435 #[deprecated(
440 since = "1.13.0",
441 note = "use McpServerConfig::with_forwarded_header(); direct field access will become pub(crate) in a future major release"
442 )]
443 pub forwarded_header: Option<ForwardedHeaderMode>,
444 #[deprecated(
447 since = "0.13.0",
448 note = "use McpServerConfig::with_readiness_check(); direct field access will become pub(crate) in a future major release"
449 )]
450 pub readiness_check: Option<ReadinessCheck>,
451 #[deprecated(
454 since = "0.13.0",
455 note = "use McpServerConfig::with_max_request_body(); direct field access will become pub(crate) in a future major release"
456 )]
457 pub max_request_body: usize,
458 #[deprecated(
461 since = "0.13.0",
462 note = "use McpServerConfig::with_request_timeout(); direct field access will become pub(crate) in a future major release"
463 )]
464 pub request_timeout: Duration,
465 #[deprecated(
468 since = "0.13.0",
469 note = "use McpServerConfig::with_shutdown_timeout(); direct field access will become pub(crate) in a future major release"
470 )]
471 pub shutdown_timeout: Duration,
472 #[deprecated(
475 since = "0.13.0",
476 note = "use McpServerConfig::with_session_idle_timeout(); direct field access will become pub(crate) in a future major release"
477 )]
478 pub session_idle_timeout: Duration,
479 pub session_binding: bool,
487 pub session_binding_secret: Option<SecretString>,
491 pub task_binding: bool,
504 pub session_store: Option<Arc<dyn SessionStore>>,
506 pub event_store: Option<Arc<dyn EventStore>>,
519 #[deprecated(
522 since = "0.13.0",
523 note = "use McpServerConfig::with_sse_keep_alive(); direct field access will become pub(crate) in a future major release"
524 )]
525 pub sse_keep_alive: Duration,
526 #[deprecated(
530 since = "0.13.0",
531 note = "use McpServerConfig::with_reload_callback(); direct field access will become pub(crate) in a future major release"
532 )]
533 pub on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
534 #[deprecated(
541 since = "0.13.0",
542 note = "use McpServerConfig::with_extra_router(); direct field access will become pub(crate) in a future major release"
543 )]
544 pub extra_router: Option<axum::Router>,
545 #[deprecated(
550 since = "0.13.0",
551 note = "use McpServerConfig::with_public_url(); direct field access will become pub(crate) in a future major release"
552 )]
553 pub public_url: Option<String>,
554 #[deprecated(
557 since = "0.13.0",
558 note = "use McpServerConfig::enable_request_header_logging(); direct field access will become pub(crate) in a future major release"
559 )]
560 pub log_request_headers: bool,
561 pub expose_build_metadata: bool,
568 #[deprecated(
571 since = "0.13.0",
572 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
573 )]
574 pub compression_enabled: bool,
575 #[deprecated(
578 since = "0.13.0",
579 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
580 )]
581 pub compression_min_size: u16,
582 #[deprecated(
586 since = "0.13.0",
587 note = "use McpServerConfig::with_max_concurrent_requests(); direct field access will become pub(crate) in a future major release"
588 )]
589 pub max_concurrent_requests: Option<usize>,
590 #[deprecated(
593 since = "0.13.0",
594 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
595 )]
596 pub admin_enabled: bool,
597 #[deprecated(
599 since = "0.13.0",
600 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
601 )]
602 pub admin_role: String,
603 #[cfg(feature = "metrics")]
606 #[deprecated(
607 since = "0.13.0",
608 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
609 )]
610 pub metrics_enabled: bool,
611 #[cfg(feature = "metrics")]
613 #[deprecated(
614 since = "0.13.0",
615 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
616 )]
617 pub metrics_bind: String,
618 #[deprecated(
622 since = "1.5.0",
623 note = "use McpServerConfig::with_security_headers(); direct field access will become pub(crate) in a future major release"
624 )]
625 pub security_headers: SecurityHeadersConfig,
626 #[deprecated(
632 since = "1.9.0",
633 note = "use McpServerConfig::with_tls_handshake_timeout(); direct field access will become pub(crate) in a future major release"
634 )]
635 pub tls_handshake_timeout: Duration,
636 #[deprecated(
643 since = "1.9.0",
644 note = "use McpServerConfig::with_max_concurrent_tls_handshakes(); direct field access will become pub(crate) in a future major release"
645 )]
646 pub max_concurrent_tls_handshakes: usize,
647}
648
649#[allow(
707 missing_debug_implementations,
708 reason = "wraps T which may not implement Debug; manual impl below avoids leaking inner contents into logs"
709)]
710pub struct Validated<T>(T);
711
712impl<T> std::fmt::Debug for Validated<T> {
713 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
714 f.debug_struct("Validated").finish_non_exhaustive()
715 }
716}
717
718impl<T> Validated<T> {
719 #[must_use]
721 pub fn as_inner(&self) -> &T {
722 &self.0
723 }
724
725 #[must_use]
730 pub fn into_inner(self) -> T {
731 self.0
732 }
733}
734
735#[allow(
736 deprecated,
737 reason = "internal builders/validators legitimately read/write the deprecated `pub` fields they were designed to manage"
738)]
739impl McpServerConfig {
740 #[must_use]
748 pub fn new(
749 bind_addr: impl Into<String>,
750 name: impl Into<String>,
751 version: impl Into<String>,
752 ) -> Self {
753 Self {
754 bind_addr: bind_addr.into(),
755 name: name.into(),
756 version: version.into(),
757 tls_cert_path: None,
758 tls_key_path: None,
759 auth: None,
760 rbac: None,
761 tool_list_filtering: true,
762 allowed_origins: Vec::new(),
763 tool_rate_limit: None,
764 readiness_check: None,
765 max_request_body: 1024 * 1024,
766 request_timeout: Duration::from_mins(2),
767 shutdown_timeout: Duration::from_secs(30),
768 session_idle_timeout: Duration::from_mins(20),
769 session_binding: true,
770 session_binding_secret: None,
771 task_binding: false,
772 session_store: None,
773 event_store: None,
774 sse_keep_alive: Duration::from_secs(15),
775 on_reload_ready: None,
776 extra_router: None,
777 public_url: None,
778 log_request_headers: false,
779 expose_build_metadata: false,
780 compression_enabled: false,
781 compression_min_size: 1024,
782 max_concurrent_requests: None,
783 admin_enabled: false,
784 admin_role: "admin".to_owned(),
785 #[cfg(feature = "metrics")]
786 metrics_enabled: false,
787 #[cfg(feature = "metrics")]
788 metrics_bind: "127.0.0.1:9090".into(),
789 security_headers: SecurityHeadersConfig::default(),
790 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
791 max_concurrent_tls_handshakes: DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES,
792 extra_route_rate_limit: None,
793 tool_rate_limit_burst: None,
794 extra_route_rate_limit_burst: None,
795 extra_route_rate_limit_exempt_paths: Vec::new(),
796 key_eviction_policy: KeyEvictionPolicy::default(),
797 trusted_forwarder_max_entries: crate::forwarded::MAX_SCANNED_ENTRIES,
798 trusted_proxies: Vec::new(),
799 forwarded_header: None,
800 }
801 }
802
803 #[must_use]
813 pub fn with_auth(mut self, auth: AuthConfig) -> Self {
814 self.auth = Some(auth);
815 self
816 }
817
818 #[must_use]
823 pub fn with_security_headers(mut self, headers: SecurityHeadersConfig) -> Self {
824 self.security_headers = headers;
825 self
826 }
827
828 #[must_use]
832 pub fn with_bind_addr(mut self, addr: impl Into<String>) -> Self {
833 self.bind_addr = addr.into();
834 self
835 }
836
837 #[must_use]
840 pub fn with_rbac(mut self, rbac: Arc<RbacPolicy>) -> Self {
841 self.rbac = Some(rbac);
842 self
843 }
844
845 #[must_use]
851 pub const fn with_tool_list_filtering(mut self, enabled: bool) -> Self {
852 self.tool_list_filtering = enabled;
853 self
854 }
855
856 #[must_use]
860 pub fn with_tls(mut self, cert_path: impl Into<PathBuf>, key_path: impl Into<PathBuf>) -> Self {
861 self.tls_cert_path = Some(cert_path.into());
862 self.tls_key_path = Some(key_path.into());
863 self
864 }
865
866 #[must_use]
870 pub fn with_public_url(mut self, url: impl Into<String>) -> Self {
871 self.public_url = Some(url.into());
872 self
873 }
874
875 #[must_use]
879 pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
880 where
881 I: IntoIterator<Item = S>,
882 S: Into<String>,
883 {
884 self.allowed_origins = origins.into_iter().map(Into::into).collect();
885 self
886 }
887
888 #[must_use]
921 pub fn with_extra_router(mut self, router: axum::Router) -> Self {
922 self.extra_router = Some(router);
923 self
924 }
925
926 #[must_use]
932 pub const fn with_trusted_forwarder_max_entries(mut self, max_entries: usize) -> Self {
933 self.trusted_forwarder_max_entries = max_entries;
934 self
935 }
936
937 #[must_use]
940 pub fn with_readiness_check(mut self, check: ReadinessCheck) -> Self {
941 self.readiness_check = Some(check);
942 self
943 }
944
945 #[must_use]
948 pub fn with_max_request_body(mut self, bytes: usize) -> Self {
949 self.max_request_body = bytes;
950 self
951 }
952
953 #[must_use]
955 pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
956 self.request_timeout = timeout;
957 self
958 }
959
960 #[must_use]
962 pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
963 self.shutdown_timeout = timeout;
964 self
965 }
966
967 #[must_use]
969 pub fn with_session_idle_timeout(mut self, timeout: Duration) -> Self {
970 self.session_idle_timeout = timeout;
971 self
972 }
973
974 #[must_use]
978 pub const fn with_session_binding(mut self, enabled: bool) -> Self {
979 self.session_binding = enabled;
980 self
981 }
982
983 #[must_use]
985 pub fn with_session_binding_secret(mut self, secret: SecretString) -> Self {
986 self.session_binding_secret = Some(secret);
987 self
988 }
989
990 #[must_use]
997 pub const fn with_task_binding(mut self, enabled: bool) -> Self {
998 self.task_binding = enabled;
999 self
1000 }
1001
1002 #[must_use]
1004 pub fn with_session_store(mut self, session_store: Arc<dyn SessionStore>) -> Self {
1005 self.session_store = Some(session_store);
1006 self
1007 }
1008
1009 #[must_use]
1014 pub fn with_event_store(mut self, event_store: Arc<dyn EventStore>) -> Self {
1015 self.event_store = Some(event_store);
1016 self
1017 }
1018
1019 #[must_use]
1021 pub fn with_sse_keep_alive(mut self, interval: Duration) -> Self {
1022 self.sse_keep_alive = interval;
1023 self
1024 }
1025
1026 #[must_use]
1030 pub fn with_max_concurrent_requests(mut self, limit: usize) -> Self {
1031 self.max_concurrent_requests = Some(limit);
1032 self
1033 }
1034
1035 #[must_use]
1043 pub fn with_tls_handshake_timeout(mut self, timeout: Duration) -> Self {
1044 self.tls_handshake_timeout = timeout;
1045 self
1046 }
1047
1048 #[must_use]
1057 pub fn with_max_concurrent_tls_handshakes(mut self, limit: usize) -> Self {
1058 self.max_concurrent_tls_handshakes = limit;
1059 self
1060 }
1061
1062 #[must_use]
1065 pub fn with_tool_rate_limit(mut self, per_minute: u32) -> Self {
1066 self.tool_rate_limit = Some(per_minute);
1067 self
1068 }
1069
1070 #[must_use]
1081 pub fn with_extra_route_rate_limit(mut self, per_minute: u32) -> Self {
1082 self.extra_route_rate_limit = Some(per_minute);
1083 self
1084 }
1085
1086 #[must_use]
1091 pub fn with_tool_rate_limit_burst(mut self, burst: u32) -> Self {
1092 self.tool_rate_limit_burst = Some(burst);
1093 self
1094 }
1095
1096 #[must_use]
1102 pub fn with_extra_route_rate_limit_burst(mut self, burst: u32) -> Self {
1103 self.extra_route_rate_limit_burst = Some(burst);
1104 self
1105 }
1106
1107 #[must_use]
1127 pub fn with_extra_route_rate_limit_exempt_paths<I, S>(mut self, paths: I) -> Self
1128 where
1129 I: IntoIterator<Item = S>,
1130 S: Into<String>,
1131 {
1132 self.extra_route_rate_limit_exempt_paths = paths.into_iter().map(Into::into).collect();
1133 self
1134 }
1135
1136 #[must_use]
1138 pub const fn with_key_eviction_policy(mut self, policy: KeyEvictionPolicy) -> Self {
1139 self.key_eviction_policy = policy;
1140 self
1141 }
1142
1143 #[must_use]
1155 pub fn with_trusted_proxies<I, S>(mut self, proxies: I) -> Self
1156 where
1157 I: IntoIterator<Item = S>,
1158 S: Into<String>,
1159 {
1160 self.trusted_proxies = proxies.into_iter().map(Into::into).collect();
1161 self
1162 }
1163
1164 #[must_use]
1169 pub fn with_forwarded_header(mut self, mode: ForwardedHeaderMode) -> Self {
1170 self.forwarded_header = Some(mode);
1171 self
1172 }
1173
1174 #[must_use]
1178 pub fn with_reload_callback<F>(mut self, callback: F) -> Self
1179 where
1180 F: FnOnce(ReloadHandle) + Send + 'static,
1181 {
1182 self.on_reload_ready = Some(Box::new(callback));
1183 self
1184 }
1185
1186 #[must_use]
1190 pub fn enable_compression(mut self, min_size: u16) -> Self {
1191 self.compression_enabled = true;
1192 self.compression_min_size = min_size;
1193 self
1194 }
1195
1196 #[must_use]
1201 pub fn enable_admin(mut self, role: impl Into<String>) -> Self {
1202 self.admin_enabled = true;
1203 self.admin_role = role.into();
1204 self
1205 }
1206
1207 #[must_use]
1210 pub fn enable_request_header_logging(mut self) -> Self {
1211 self.log_request_headers = true;
1212 self
1213 }
1214
1215 #[must_use]
1220 pub fn expose_build_metadata(mut self) -> Self {
1221 self.expose_build_metadata = true;
1222 self
1223 }
1224
1225 #[cfg(feature = "metrics")]
1228 #[must_use]
1229 pub fn with_metrics(mut self, bind: impl Into<String>) -> Self {
1230 self.metrics_enabled = true;
1231 self.metrics_bind = bind.into();
1232 self
1233 }
1234
1235 pub fn validate(self) -> Result<Validated<Self>, RmcpServerKitError> {
1268 self.check()?;
1269 Ok(Validated(self))
1270 }
1271
1272 fn check_burst_knobs(&self) -> Result<(), RmcpServerKitError> {
1279 if self.tool_rate_limit_burst == Some(0) {
1280 return Err(RmcpServerKitError::Config(
1281 "tool_rate_limit_burst must be greater than zero".into(),
1282 ));
1283 }
1284 if self.extra_route_rate_limit_burst == Some(0) {
1285 return Err(RmcpServerKitError::Config(
1286 "extra_route_rate_limit_burst must be greater than zero".into(),
1287 ));
1288 }
1289 if self.trusted_forwarder_max_entries == 0
1290 || self.trusted_forwarder_max_entries
1291 > crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES
1292 {
1293 return Err(RmcpServerKitError::Config(format!(
1294 "trusted_forwarder_max_entries must be in 1..={}, got {}",
1295 crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES,
1296 self.trusted_forwarder_max_entries
1297 )));
1298 }
1299 if self.tool_rate_limit_burst.is_some() && self.tool_rate_limit.is_none() {
1300 return Err(RmcpServerKitError::Config(
1301 "tool_rate_limit_burst requires tool_rate_limit to be set".into(),
1302 ));
1303 }
1304 if self.extra_route_rate_limit_burst.is_some() && self.extra_route_rate_limit.is_none() {
1305 return Err(RmcpServerKitError::Config(
1306 "extra_route_rate_limit_burst requires extra_route_rate_limit to be set".into(),
1307 ));
1308 }
1309 if !self.extra_route_rate_limit_exempt_paths.is_empty()
1310 && self.extra_route_rate_limit.is_none()
1311 {
1312 return Err(RmcpServerKitError::Config(
1313 "extra_route_rate_limit_exempt_paths requires extra_route_rate_limit to be set"
1314 .into(),
1315 ));
1316 }
1317 for path in &self.extra_route_rate_limit_exempt_paths {
1318 if path.is_empty() || !path.starts_with('/') {
1319 return Err(RmcpServerKitError::Config(format!(
1320 "extra_route_rate_limit_exempt_paths entries must be non-empty and start with '/': {path:?}"
1321 )));
1322 }
1323 }
1324 if let Some(rl) = self.auth.as_ref().and_then(|a| a.rate_limit.as_ref()) {
1325 if rl.burst == Some(0) {
1326 return Err(RmcpServerKitError::Config(
1327 "auth rate_limit.burst must be greater than zero".into(),
1328 ));
1329 }
1330 if rl.pre_auth_burst == Some(0) {
1331 return Err(RmcpServerKitError::Config(
1332 "auth rate_limit.pre_auth_burst must be greater than zero".into(),
1333 ));
1334 }
1335 }
1336 Ok(())
1337 }
1338
1339 fn check_trusted_forwarder(&self) -> Result<(), RmcpServerKitError> {
1344 for entry in &self.trusted_proxies {
1345 validate_trusted_proxy_entry(entry).map_err(RmcpServerKitError::Config)?;
1346 }
1347 if self.forwarded_header.is_some() && self.trusted_proxies.is_empty() {
1348 return Err(RmcpServerKitError::Config(
1349 "forwarded_header requires trusted_proxies to be nonempty".into(),
1350 ));
1351 }
1352 Ok(())
1353 }
1354
1355 fn check_session_binding_config(&self) -> Result<(), RmcpServerKitError> {
1356 if self.session_store.is_some()
1357 && self.session_binding
1358 && self.auth.as_ref().is_some_and(|auth| auth.enabled)
1359 && self.session_binding_secret.is_none()
1360 {
1361 return Err(RmcpServerKitError::Config(
1362 "session_store with session_binding enabled and auth configured requires \
1363 session_binding_secret: a shared secret is required for cross-instance \
1364 session verification"
1365 .into(),
1366 ));
1367 }
1368
1369 if let Some(secret) = &self.session_binding_secret {
1370 crate::session_binding::validate_configured_secret(
1371 "session_binding_secret",
1372 secret.expose_secret(),
1373 )?;
1374 }
1375 Ok(())
1376 }
1377
1378 fn check(&self) -> Result<(), RmcpServerKitError> {
1382 if let Err(violation) = crate::config::check_shared_config_invariants(
1397 self.admin_enabled,
1398 self.auth.as_ref().is_some_and(|a| a.enabled),
1399 self.tls_cert_path.is_some(),
1400 self.tls_key_path.is_some(),
1401 self.auth.as_ref().is_some_and(|a| a.mtls.is_some()),
1402 ) {
1403 return Err(RmcpServerKitError::Config(
1404 match violation {
1405 crate::config::SharedConfigViolation::AdminRequiresAuth => {
1406 "admin_enabled=true requires auth to be configured and enabled"
1407 }
1408 crate::config::SharedConfigViolation::TlsCertWithoutKey => {
1409 "tls_cert_path is set but tls_key_path is missing"
1410 }
1411 crate::config::SharedConfigViolation::TlsKeyWithoutCert => {
1412 "tls_key_path is set but tls_cert_path is missing"
1413 }
1414 crate::config::SharedConfigViolation::MtlsRequiresTls => {
1415 "auth.mtls requires TLS: set both tls_cert_path and tls_key_path \
1416 (mTLS client certificates cannot be verified on a plaintext listener)"
1417 }
1418 }
1419 .into(),
1420 ));
1421 }
1422
1423 if let Some(auth) = &self.auth {
1424 auth.validate_api_key_names()?;
1425 }
1426
1427 if self.bind_addr.parse::<SocketAddr>().is_err() {
1429 return Err(RmcpServerKitError::Config(format!(
1430 "bind_addr {:?} is not a valid socket address (expected e.g. 127.0.0.1:8080)",
1431 self.bind_addr
1432 )));
1433 }
1434
1435 if let Some(ref url) = self.public_url
1437 && !(url.starts_with("http://") || url.starts_with("https://"))
1438 {
1439 return Err(RmcpServerKitError::Config(format!(
1440 "public_url {url:?} must start with http:// or https://"
1441 )));
1442 }
1443
1444 for origin in &self.allowed_origins {
1446 if !(origin.starts_with("http://") || origin.starts_with("https://")) {
1447 return Err(RmcpServerKitError::Config(format!(
1448 "allowed_origins entry {origin:?} must start with http:// or https://"
1449 )));
1450 }
1451 }
1452
1453 if self.max_request_body == 0 {
1455 return Err(RmcpServerKitError::Config(
1456 "max_request_body must be greater than zero".into(),
1457 ));
1458 }
1459
1460 if self.extra_route_rate_limit == Some(0) {
1464 return Err(RmcpServerKitError::Config(
1465 "extra_route_rate_limit must be greater than zero".into(),
1466 ));
1467 }
1468
1469 self.check_burst_knobs()?;
1471
1472 self.check_trusted_forwarder()?;
1474
1475 #[cfg(feature = "oauth")]
1477 if let Some(auth_cfg) = &self.auth
1478 && let Some(oauth_cfg) = &auth_cfg.oauth
1479 {
1480 oauth_cfg.validate()?;
1481 }
1482
1483 self.check_session_binding_config()?;
1484
1485 validate_security_headers(&self.security_headers)?;
1488
1489 if self.max_concurrent_requests == Some(0) {
1493 return Err(RmcpServerKitError::Config(
1494 "max_concurrent_requests must be greater than zero when set".into(),
1495 ));
1496 }
1497
1498 if let Some(auth_cfg) = &self.auth
1502 && let Some(rl) = &auth_cfg.rate_limit
1503 && rl.max_tracked_keys == 0
1504 {
1505 return Err(RmcpServerKitError::Config(
1506 "auth.rate_limit.max_tracked_keys must be greater than zero".into(),
1507 ));
1508 }
1509
1510 check_auth_capacity_knobs(self.auth.as_ref())?;
1511
1512 if self.tls_handshake_timeout == Duration::ZERO {
1517 return Err(RmcpServerKitError::Config(
1518 "tls_handshake_timeout must be greater than zero".into(),
1519 ));
1520 }
1521
1522 if self.max_concurrent_tls_handshakes == 0 {
1527 return Err(RmcpServerKitError::Config(
1528 "max_concurrent_tls_handshakes must be greater than zero".into(),
1529 ));
1530 }
1531
1532 Ok(())
1533 }
1534}
1535
1536#[allow(
1542 missing_debug_implementations,
1543 reason = "contains Arc<AuthState> with non-Debug fields"
1544)]
1545pub struct ReloadHandle {
1546 auth: Option<Arc<AuthState>>,
1547 rbac: Option<Arc<ArcSwap<RbacPolicy>>>,
1548 crl_set: Option<Arc<CrlSet>>,
1549}
1550
1551impl ReloadHandle {
1552 pub fn try_reload_auth_keys(
1569 &self,
1570 keys: Vec<crate::auth::ApiKeyEntry>,
1571 ) -> Result<(), RmcpServerKitError> {
1572 if let Some(ref auth) = self.auth {
1573 auth.try_reload_keys(keys)
1574 } else {
1575 Ok(())
1576 }
1577 }
1578
1579 pub fn reload_auth_keys(&self, keys: Vec<crate::auth::ApiKeyEntry>) {
1588 if let Err(error) = self.try_reload_auth_keys(keys) {
1589 tracing::error!(%error, "API key hot reload rejected: keys left unchanged");
1590 }
1591 }
1592
1593 pub fn reload_rbac(&self, policy: RbacPolicy) {
1595 if let Some(ref rbac) = self.rbac {
1596 rbac.store(Arc::new(policy));
1597 tracing::info!("RBAC policy reloaded");
1598 }
1599 }
1600
1601 pub async fn refresh_crls(&self) -> Result<(), RmcpServerKitError> {
1611 let Some(ref crl_set) = self.crl_set else {
1612 return Err(RmcpServerKitError::Config(
1613 "CRL refresh requested but mTLS CRL support is not configured".into(),
1614 ));
1615 };
1616
1617 crl_set.force_refresh().await
1618 }
1619}
1620
1621#[allow(
1638 clippy::too_many_lines,
1639 clippy::cognitive_complexity,
1640 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"
1641)]
1642struct AppRunParams {
1646 tls_paths: Option<(PathBuf, PathBuf)>,
1648 tls_handshake_timeout: Duration,
1650 max_concurrent_tls_handshakes: usize,
1652 mtls_config: Option<MtlsConfig>,
1654 shutdown_timeout: Duration,
1656 auth_state: Option<Arc<AuthState>>,
1658 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
1660 on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
1662 ct: CancellationToken,
1666 session_ct: CancellationToken,
1676 scheme: &'static str,
1678 name: String,
1680}
1681
1682type BindingSecrets = (
1699 Option<crate::session_binding::SessionBindingSecret>,
1700 Option<crate::session_binding::SessionBindingSecret>,
1701);
1702
1703fn resolve_binding_secret(config: &McpServerConfig) -> anyhow::Result<BindingSecrets> {
1704 #[allow(
1705 deprecated,
1706 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1707 )]
1708 if !config.session_binding && !config.task_binding {
1709 return Ok((None, None));
1710 }
1711 #[allow(
1712 deprecated,
1713 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1714 )]
1715 let secret = match config.session_binding_secret.as_ref() {
1716 Some(configured) => configured_session_binding_secret(configured)?,
1717 None => process_session_binding_secret().clone(),
1718 };
1719 #[allow(
1720 deprecated,
1721 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1722 )]
1723 let pair = (
1724 config.session_binding.then(|| secret.clone()),
1725 config.task_binding.then_some(secret),
1726 );
1727 Ok(pair)
1728}
1729
1730#[allow(
1731 clippy::cognitive_complexity,
1732 reason = "router assembly is intrinsically sequential; splitting harms readability"
1733)]
1734#[allow(
1735 deprecated,
1736 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1737)]
1738fn build_app_router<H, F>(
1739 mut config: McpServerConfig,
1740 handler_factory: F,
1741) -> anyhow::Result<(axum::Router, AppRunParams)>
1742where
1743 H: ServerHandler + 'static,
1744 F: Fn() -> H + Send + Sync + Clone + 'static,
1745{
1746 let ct = CancellationToken::new();
1747 let session_ct = CancellationToken::new();
1748
1749 let allowed_hosts = derive_allowed_hosts(&config.bind_addr, config.public_url.as_deref());
1750 tracing::info!(allowed_hosts = %allowed_hosts.join(", "), "configured Streamable HTTP allowed hosts");
1751
1752 if config.max_concurrent_requests.is_none() {
1753 tracing::warn!(
1754 "max_concurrent_requests is unset: in-flight HTTP requests are unlimited; \
1755 set McpServerConfig::with_max_concurrent_requests or front the server with \
1756 an external concurrency limit"
1757 );
1758 }
1759
1760 let rbac_swap = Arc::new(ArcSwap::new(
1763 config
1764 .rbac
1765 .clone()
1766 .unwrap_or_else(|| Arc::new(RbacPolicy::disabled())),
1767 ));
1768
1769 let rbac_for_handler = Arc::clone(&rbac_swap);
1770 let tool_list_filtering = config.tool_list_filtering;
1771 let session_store = config.session_store.take();
1772 let mut rmcp_config = StreamableHttpServerConfig::default()
1773 .with_allowed_hosts(allowed_hosts)
1774 .with_sse_keep_alive(Some(config.sse_keep_alive))
1775 .with_cancellation_token(session_ct.clone());
1776 rmcp_config.session_store = session_store;
1777 let event_store = config.event_store.take();
1778 let (binding_secret, task_binding_secret) = resolve_binding_secret(&config)?;
1779 let mcp_service = StreamableHttpService::new(
1780 move || {
1781 Ok(RbacContextHandler::new(
1782 handler_factory(),
1783 Arc::clone(&rbac_for_handler),
1784 tool_list_filtering,
1785 )
1786 .with_task_binding(task_binding_secret.clone()))
1787 },
1788 {
1789 let mut mgr = LocalSessionManager::default();
1790 mgr.session_config.keep_alive = Some(config.session_idle_timeout);
1791 if let Some(event_store) = event_store {
1792 mgr = mgr.with_event_store(event_store);
1793 }
1794 mgr.into()
1795 },
1796 rmcp_config,
1797 );
1798
1799 let mut mcp_router = axum::Router::new().nest_service("/mcp", mcp_service);
1801
1802 let auth_state: Option<Arc<AuthState>> = match config.auth {
1806 Some(ref auth_config) if auth_config.enabled => {
1807 let rate_limiter = auth_config.rate_limit.as_ref().map(build_rate_limiter);
1808 let pre_auth_limiter = auth_config
1809 .rate_limit
1810 .as_ref()
1811 .map(crate::auth::build_pre_auth_limiter);
1812
1813 #[cfg(feature = "oauth")]
1814 let jwks_cache = auth_config
1815 .oauth
1816 .as_ref()
1817 .map(|c| crate::oauth::JwksCache::new(c).map(Arc::new))
1818 .transpose()
1819 .map_err(|e| std::io::Error::other(format!("JWKS HTTP client: {e}")))?;
1820
1821 Some(Arc::new(AuthState {
1822 api_keys: ArcSwap::new(Arc::new(auth_config.api_keys.clone())),
1823 rate_limiter,
1824 pre_auth_limiter,
1825 #[cfg(feature = "oauth")]
1826 jwks_cache,
1827 seen_identities: crate::auth::SeenIdentitySet::new(),
1828 counters: crate::auth::AuthCounters::default(),
1829 resource_metadata_url: config.public_url.as_ref().map(|url| {
1837 format!(
1838 "{}/.well-known/oauth-protected-resource/mcp",
1839 url.trim_end_matches('/')
1840 )
1841 }),
1842 }))
1843 }
1844 _ => None,
1845 };
1846
1847 if config.admin_enabled {
1850 let Some(ref auth_state_ref) = auth_state else {
1851 return Err(anyhow::anyhow!(
1852 "admin_enabled=true requires auth to be configured and enabled"
1853 ));
1854 };
1855 let admin_state = crate::admin::AdminState {
1856 started_at: std::time::Instant::now(),
1857 name: config.name.clone(),
1858 version: config.version.clone(),
1859 auth: Some(Arc::clone(auth_state_ref)),
1860 rbac: Arc::clone(&rbac_swap),
1861 };
1862 let admin_cfg = crate::admin::AdminConfig {
1863 role: config.admin_role.clone(),
1864 };
1865 mcp_router = mcp_router.merge(crate::admin::admin_router(admin_state, &admin_cfg));
1866 tracing::info!(role = %config.admin_role, "/admin/* endpoints enabled");
1867 }
1868
1869 if let Some(secret) = binding_secret {
1901 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1902 let secret = secret.clone();
1903 session_binding_middleware(secret, req, next)
1904 }));
1905 }
1906
1907 {
1911 let tool_limiter: Option<Arc<ToolRateLimiter>> = config.tool_rate_limit.map(|per_minute| {
1912 build_tool_rate_limiter_with_policy(
1913 per_minute,
1914 config.tool_rate_limit_burst,
1915 config.key_eviction_policy,
1916 )
1917 });
1918
1919 if rbac_swap.load().is_enabled() {
1920 tracing::info!("RBAC enforcement enabled on /mcp");
1921 }
1922 if let Some(limit) = config.tool_rate_limit {
1923 tracing::info!(limit, "tool rate limiting enabled (calls/min per IP)");
1924 }
1925
1926 let rbac_for_mw = Arc::clone(&rbac_swap);
1927 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1928 let p = rbac_for_mw.load_full();
1929 let tl = tool_limiter.clone();
1930 rbac_middleware(p, tl, req, next)
1931 }));
1932 }
1933
1934 if let Some(ref auth_config) = config.auth
1936 && auth_config.enabled
1937 {
1938 let Some(ref state) = auth_state else {
1939 return Err(anyhow::anyhow!("auth state missing despite enabled config"));
1940 };
1941
1942 let methods: Vec<&str> = [
1943 auth_config.mtls.is_some().then_some("mTLS"),
1944 (!auth_config.api_keys.is_empty()).then_some("bearer"),
1945 #[cfg(feature = "oauth")]
1946 auth_config.oauth.is_some().then_some("oauth-jwt"),
1947 ]
1948 .into_iter()
1949 .flatten()
1950 .collect();
1951
1952 tracing::info!(
1953 methods = %methods.join(", "),
1954 api_keys = auth_config.api_keys.len(),
1955 "auth enabled on /mcp"
1956 );
1957
1958 let state_for_mw = Arc::clone(state);
1959 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1960 let s = Arc::clone(&state_for_mw);
1961 auth_middleware(s, req, next)
1962 }));
1963 }
1964
1965 mcp_router = mcp_router.layer(tower_http::timeout::TimeoutLayer::with_status_code(
1968 axum::http::StatusCode::REQUEST_TIMEOUT,
1969 config.request_timeout,
1970 ));
1971
1972 mcp_router = mcp_router.layer(tower_http::limit::RequestBodyLimitLayer::new(
1976 config.max_request_body,
1977 ));
1978
1979 let mut effective_origins = config.allowed_origins.clone();
1986 if effective_origins.is_empty()
1987 && let Some(ref url) = config.public_url
1988 {
1989 if let Some(scheme_end) = url.find("://") {
1994 let scheme_with_sep = url.get(..scheme_end + 3).unwrap_or_default();
1995 let after_scheme = url.get(scheme_end + 3..).unwrap_or_default();
1996 let host_end = after_scheme.find('/').unwrap_or(after_scheme.len());
1997 let host = after_scheme.get(..host_end).unwrap_or_default();
1998 let origin = format!("{scheme_with_sep}{host}");
1999 tracing::info!(
2000 %origin,
2001 "auto-derived allowed origin from public_url"
2002 );
2003 effective_origins.push(origin);
2004 }
2005 }
2006 let allowed_origins: Arc<[String]> = Arc::from(effective_origins);
2007 let cors_origins = Arc::clone(&allowed_origins);
2008 let log_request_headers = config.log_request_headers;
2009
2010 let readyz_route = if let Some(check) = config.readiness_check.take() {
2011 axum::routing::get(move || readyz(Arc::clone(&check)))
2012 } else {
2013 axum::routing::get(healthz)
2014 };
2015
2016 #[allow(
2017 unused_mut,
2018 reason = "the binding is only reassigned when the `oauth` feature adds the \
2019 protected-resource-metadata route below"
2020 )]
2021 let mut router = axum::Router::new()
2022 .route("/healthz", axum::routing::get(healthz))
2023 .route("/readyz", readyz_route)
2024 .route(
2025 "/version",
2026 axum::routing::get({
2027 let payload_bytes: Arc<[u8]> = serialize_version_payload(
2032 &config.name,
2033 &config.version,
2034 config.expose_build_metadata,
2035 );
2036 move || {
2037 let p = Arc::clone(&payload_bytes);
2038 async move {
2039 (
2040 [(axum::http::header::CONTENT_TYPE, "application/json")],
2041 p.to_vec(),
2042 )
2043 }
2044 }
2045 }),
2046 )
2047 .merge(mcp_router);
2048
2049 if let Some(extra) = config.extra_router.take() {
2056 let extra = match config.extra_route_rate_limit {
2057 Some(per_minute) => {
2058 let max_tracked_keys =
2059 NonZeroUsize::new(EXTRA_ROUTE_MAX_TRACKED_KEYS).unwrap_or(NonZeroUsize::MIN);
2060 let limiter = build_extra_route_rate_limiter_with_policy(
2061 per_minute,
2062 config.extra_route_rate_limit_burst,
2063 config.key_eviction_policy,
2064 max_tracked_keys,
2065 );
2066 let exempt: Arc<std::collections::HashSet<String>> = Arc::new(
2067 config
2068 .extra_route_rate_limit_exempt_paths
2069 .iter()
2070 .cloned()
2071 .collect(),
2072 );
2073 tracing::info!(
2074 per_minute,
2075 exempt_paths = exempt.len(),
2076 "extra-route per-IP rate limit enabled"
2077 );
2078 extra.layer(axum::middleware::from_fn(move |req, next| {
2079 let l = Arc::clone(&limiter);
2080 let e = Arc::clone(&exempt);
2081 extra_route_rate_limit_middleware(l, e, req, next)
2082 }))
2083 }
2084 None => extra,
2085 };
2086 router = router.merge(extra);
2087 }
2088
2089 let server_url = derive_server_url(&config);
2096 let resource_url = format!("{server_url}/mcp");
2097
2098 #[cfg(feature = "oauth")]
2099 let prm_metadata = if let Some(ref auth_config) = config.auth
2100 && let Some(ref oauth_config) = auth_config.oauth
2101 {
2102 crate::oauth::protected_resource_metadata(&resource_url, &server_url, oauth_config)
2103 } else {
2104 serde_json::json!({ "resource": resource_url })
2105 };
2106 #[cfg(not(feature = "oauth"))]
2107 let prm_metadata = serde_json::json!({ "resource": resource_url });
2108
2109 let prm_root = prm_metadata.clone();
2115 router = router.route(
2116 "/.well-known/oauth-protected-resource",
2117 axum::routing::get(move || {
2118 let m = prm_root.clone();
2119 async move { axum::Json(m) }
2120 }),
2121 );
2122 router = router.route(
2123 "/.well-known/oauth-protected-resource/mcp",
2124 axum::routing::get(move || {
2125 let m = prm_metadata.clone();
2126 async move { axum::Json(m) }
2127 }),
2128 );
2129
2130 #[cfg(feature = "oauth")]
2135 if let Some(ref auth_config) = config.auth
2136 && let Some(ref oauth_config) = auth_config.oauth
2137 && oauth_config.proxy.is_some()
2138 {
2139 router = install_oauth_proxy_routes(
2140 router,
2141 &server_url,
2142 oauth_config,
2143 auth_state.as_ref(),
2144 config.max_request_body,
2145 &config.admin_role,
2146 )?;
2147 }
2148
2149 if !cors_origins.is_empty() {
2158 let cors = tower_http::cors::CorsLayer::new()
2159 .allow_origin(
2160 cors_origins
2161 .iter()
2162 .filter_map(|o| o.parse::<axum::http::HeaderValue>().ok())
2163 .collect::<Vec<_>>(),
2164 )
2165 .allow_methods([
2166 axum::http::Method::GET,
2167 axum::http::Method::POST,
2168 axum::http::Method::OPTIONS,
2169 ])
2170 .allow_headers([
2171 axum::http::header::CONTENT_TYPE,
2172 axum::http::header::AUTHORIZATION,
2173 ]);
2174 router = router.layer(cors);
2175 }
2176
2177 if config.compression_enabled {
2181 use tower_http::compression::Predicate as _;
2182 let predicate = tower_http::compression::DefaultPredicate::new().and(
2183 tower_http::compression::predicate::SizeAbove::new(u64::from(
2184 config.compression_min_size,
2185 )),
2186 );
2187 router = router.layer(
2188 tower_http::compression::CompressionLayer::new()
2189 .gzip(true)
2190 .br(true)
2191 .compress_when(predicate),
2192 );
2193 tracing::info!(
2194 min_size = config.compression_min_size,
2195 "response compression enabled (gzip, br)"
2196 );
2197 }
2198
2199 if let Some(max) = config.max_concurrent_requests {
2202 let overload_handler = tower::ServiceBuilder::new()
2203 .layer(axum::error_handling::HandleErrorLayer::new(
2204 |_err: tower::BoxError| async {
2205 (
2206 axum::http::StatusCode::SERVICE_UNAVAILABLE,
2207 axum::Json(serde_json::json!({
2208 "error": "overloaded",
2209 "error_description": "server is at capacity, retry later"
2210 })),
2211 )
2212 },
2213 ))
2214 .layer(tower::load_shed::LoadShedLayer::new())
2215 .layer(tower::limit::ConcurrencyLimitLayer::new(max));
2216 router = router.layer(overload_handler);
2217 tracing::info!(max, "global concurrency limit enabled");
2218 }
2219
2220 router = router.fallback(|| async {
2224 (
2225 axum::http::StatusCode::NOT_FOUND,
2226 axum::Json(serde_json::json!({
2227 "error": "not_found",
2228 "error_description": "The requested endpoint does not exist"
2229 })),
2230 )
2231 });
2232
2233 #[cfg(feature = "metrics")]
2235 if config.metrics_enabled {
2236 let metrics = Arc::new(
2237 crate::metrics::McpMetrics::new().map_err(|e| anyhow::anyhow!("metrics init: {e}"))?,
2238 );
2239 let m = Arc::clone(&metrics);
2240 router = router.layer(axum::middleware::from_fn(
2241 move |req: Request<Body>, next: Next| {
2242 let m = Arc::clone(&m);
2243 metrics_middleware(m, req, next)
2244 },
2245 ));
2246 let metrics_bind = config.metrics_bind.clone();
2247 let metrics_shutdown = ct.clone();
2248 tokio::spawn(async move {
2249 if let Err(e) =
2250 crate::metrics::serve_metrics(metrics_bind, metrics, metrics_shutdown).await
2251 {
2252 tracing::error!("metrics listener failed: {e}");
2253 }
2254 });
2255 }
2256
2257 let forward_resolver: Option<Arc<ForwardResolver>> = if config.trusted_proxies.is_empty() {
2265 None
2266 } else {
2267 Some(Arc::new(ForwardResolver {
2270 trusted: config
2271 .trusted_proxies
2272 .iter()
2273 .filter_map(|entry| parse_proxy_net(entry))
2274 .collect(),
2275 mode: config
2276 .forwarded_header
2277 .unwrap_or(ForwardedHeaderMode::XForwardedFor),
2278 max_scanned_entries: config.trusted_forwarder_max_entries,
2279 }))
2280 };
2281 if forward_resolver.is_some() {
2282 tracing::info!(
2283 proxies = config.trusted_proxies.len(),
2284 "trusted-forwarder mode enabled: limiters key by resolved client IP"
2285 );
2286 }
2287 router = router.layer(axum::middleware::from_fn(move |req, next| {
2288 let r = forward_resolver.clone();
2289 normalize_peer_addr_middleware(r, req, next)
2290 }));
2291
2292 router = router.layer(axum::middleware::from_fn(move |req, next| {
2304 let origins = Arc::clone(&allowed_origins);
2305 origin_check_middleware(origins, log_request_headers, req, next)
2306 }));
2307
2308 let is_tls = config.tls_cert_path.is_some();
2317 warn_security_header_overrides(&config.security_headers);
2318 let security_headers_cfg = Arc::new(config.security_headers.clone());
2319 router = router.layer(axum::middleware::from_fn(move |req, next| {
2320 let cfg = Arc::clone(&security_headers_cfg);
2321 security_headers_middleware(is_tls, cfg, req, next)
2322 }));
2323
2324 let scheme = if config.tls_cert_path.is_some() {
2325 "https"
2326 } else {
2327 "http"
2328 };
2329
2330 let tls_paths = match (&config.tls_cert_path, &config.tls_key_path) {
2331 (Some(cert), Some(key)) => Some((cert.clone(), key.clone())),
2332 _ => None,
2333 };
2334 let tls_handshake_timeout = config.tls_handshake_timeout;
2335 let max_concurrent_tls_handshakes = config.max_concurrent_tls_handshakes;
2336 let mtls_config = config.auth.as_ref().and_then(|a| a.mtls.as_ref()).cloned();
2337
2338 Ok((
2339 router,
2340 AppRunParams {
2341 tls_paths,
2342 tls_handshake_timeout,
2343 max_concurrent_tls_handshakes,
2344 mtls_config,
2345 shutdown_timeout: config.shutdown_timeout,
2346 auth_state,
2347 rbac_swap,
2348 on_reload_ready: config.on_reload_ready.take(),
2349 ct,
2350 session_ct,
2351 scheme,
2352 name: config.name.clone(),
2353 },
2354 ))
2355}
2356
2357struct CancelOnDrop(CancellationToken);
2370
2371impl Drop for CancelOnDrop {
2372 fn drop(&mut self) {
2373 self.0.cancel();
2374 }
2375}
2376
2377fn spawn_external_shutdown_bridge(
2381 external: CancellationToken,
2382 internal: CancellationToken,
2383) -> tokio::task::JoinHandle<()> {
2384 tokio::spawn(async move {
2385 tokio::select! {
2389 () = external.cancelled() => internal.cancel(),
2390 () = internal.cancelled() => {}
2391 }
2392 })
2393}
2394
2395pub async fn serve<H, F>(
2415 config: Validated<McpServerConfig>,
2416 handler_factory: F,
2417) -> Result<(), RmcpServerKitError>
2418where
2419 H: ServerHandler + 'static,
2420 F: Fn() -> H + Send + Sync + Clone + 'static,
2421{
2422 let config = config.into_inner();
2423 #[allow(
2424 deprecated,
2425 reason = "internal serve() reads `bind_addr` to construct the listener; field becomes pub(crate) in 1.0"
2426 )]
2427 let bind_addr = config.bind_addr.clone();
2428 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2429 let _cancel_guard = CancelOnDrop(params.ct.clone());
2430
2431 let listener = TcpListener::bind(&bind_addr)
2432 .await
2433 .map_err(|e| io_to_startup(&format!("bind {bind_addr}"), e))?;
2434 log_listening(¶ms.name, params.scheme, &bind_addr);
2435
2436 run_server(
2437 router,
2438 listener,
2439 params.tls_paths,
2440 params.tls_handshake_timeout,
2441 params.max_concurrent_tls_handshakes,
2442 params.mtls_config,
2443 params.shutdown_timeout,
2444 params.auth_state,
2445 params.rbac_swap,
2446 params.on_reload_ready,
2447 params.ct,
2448 params.session_ct,
2449 )
2450 .await
2451 .map_err(anyhow_to_startup)
2452}
2453
2454pub async fn serve_with_listener<H, F>(
2487 listener: TcpListener,
2488 config: Validated<McpServerConfig>,
2489 handler_factory: F,
2490 ready_tx: Option<tokio::sync::oneshot::Sender<SocketAddr>>,
2491 shutdown: Option<CancellationToken>,
2492) -> Result<(), RmcpServerKitError>
2493where
2494 H: ServerHandler + 'static,
2495 F: Fn() -> H + Send + Sync + Clone + 'static,
2496{
2497 let config = config.into_inner();
2498 let local_addr = listener
2499 .local_addr()
2500 .map_err(|e| io_to_startup("listener.local_addr", e))?;
2501 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2502 let _cancel_guard = CancelOnDrop(params.ct.clone());
2503
2504 log_listening(¶ms.name, params.scheme, &local_addr.to_string());
2505
2506 if let Some(external) = shutdown {
2510 let _bridge_task = spawn_external_shutdown_bridge(external, params.ct.clone());
2511 }
2512
2513 if let Some(tx) = ready_tx {
2517 let _ = tx.send(local_addr);
2519 }
2520
2521 run_server(
2522 router,
2523 listener,
2524 params.tls_paths,
2525 params.tls_handshake_timeout,
2526 params.max_concurrent_tls_handshakes,
2527 params.mtls_config,
2528 params.shutdown_timeout,
2529 params.auth_state,
2530 params.rbac_swap,
2531 params.on_reload_ready,
2532 params.ct,
2533 params.session_ct,
2534 )
2535 .await
2536 .map_err(anyhow_to_startup)
2537}
2538
2539#[allow(
2542 clippy::cognitive_complexity,
2543 reason = "tracing::info! macro expansions inflate the score; logic is trivial"
2544)]
2545fn log_listening(name: &str, scheme: &str, addr: &str) {
2546 tracing::info!("{name} listening on {addr}");
2547 tracing::info!(" MCP endpoint: {scheme}://{addr}/mcp");
2548 tracing::info!(" Health check: {scheme}://{addr}/healthz");
2549 tracing::info!(" Readiness: {scheme}://{addr}/readyz");
2550}
2551
2552#[allow(
2575 clippy::too_many_arguments,
2576 clippy::cognitive_complexity,
2577 reason = "server start-up threads TLS, reload state, and graceful shutdown through one flow"
2578)]
2579async fn run_server(
2583 router: axum::Router,
2584 listener: TcpListener,
2585 tls_paths: Option<(PathBuf, PathBuf)>,
2586 tls_handshake_timeout: Duration,
2587 max_concurrent_tls_handshakes: usize,
2588 mtls_config: Option<MtlsConfig>,
2589 shutdown_timeout: Duration,
2590 auth_state: Option<Arc<AuthState>>,
2591 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
2592 mut on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
2593 ct: CancellationToken,
2594 session_ct: CancellationToken,
2595) -> anyhow::Result<()> {
2596 let shutdown_trigger = CancellationToken::new();
2600 {
2601 let trigger = shutdown_trigger.clone();
2602 let parent = ct.clone();
2603 tokio::spawn(async move {
2604 tokio::select! {
2607 () = shutdown_signal() => {}
2608 () = parent.cancelled() => {}
2609 }
2610 trigger.cancel();
2611 });
2612 }
2613
2614 let graceful = {
2615 let trigger = shutdown_trigger.clone();
2616 let ct = ct.clone();
2617 async move {
2618 trigger.cancelled().await;
2619 tracing::info!("shutting down (grace period: {shutdown_timeout:?})");
2620 ct.cancel();
2621 }
2622 };
2623
2624 let force_exit_timer = {
2625 let trigger = shutdown_trigger.clone();
2626 async move {
2627 trigger.cancelled().await;
2628 tokio::time::sleep(shutdown_timeout).await;
2629 }
2630 };
2631
2632 if let Some((cert_path, key_path)) = tls_paths {
2633 let crl_set = if let Some(mtls) = mtls_config.as_ref()
2634 && mtls.crl_enabled
2635 {
2636 let (ca_certs, roots) = load_client_auth_roots(&mtls.ca_cert_path)?;
2637 let (crl_set, discover_rx) =
2638 mtls_revocation::bootstrap_fetch(roots, &ca_certs, mtls.clone())
2639 .await
2640 .map_err(|error| anyhow::anyhow!(error.to_string()))?;
2641 tokio::spawn(mtls_revocation::run_crl_refresher(
2642 Arc::clone(&crl_set),
2643 discover_rx,
2644 ct.clone(),
2645 ));
2646 Some(crl_set)
2647 } else {
2648 None
2649 };
2650
2651 if let Some(cb) = on_reload_ready.take() {
2652 cb(ReloadHandle {
2653 auth: auth_state.clone(),
2654 rbac: Some(Arc::clone(&rbac_swap)),
2655 crl_set: crl_set.clone(),
2656 });
2657 }
2658
2659 let tls_listener = TlsListener::new(
2660 listener,
2661 &cert_path,
2662 &key_path,
2663 mtls_config.as_ref(),
2664 crl_set,
2665 tls_handshake_timeout,
2666 max_concurrent_tls_handshakes,
2667 )?;
2668 let make_svc = router.into_make_service_with_connect_info::<TlsConnInfo>();
2669 tokio::select! {
2672 result = axum::serve(tls_listener, make_svc)
2673 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2674 () = force_exit_timer => {
2675 tracing::warn!("shutdown timeout exceeded, forcing exit");
2676 session_ct.cancel();
2677 }
2678 }
2679 } else {
2680 if let Some(cb) = on_reload_ready.take() {
2681 cb(ReloadHandle {
2682 auth: auth_state,
2683 rbac: Some(rbac_swap),
2684 crl_set: None,
2685 });
2686 }
2687
2688 let make_svc = router.into_make_service_with_connect_info::<SocketAddr>();
2689 tokio::select! {
2692 result = axum::serve(listener, make_svc)
2693 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2694 () = force_exit_timer => {
2695 tracing::warn!("shutdown timeout exceeded, forcing exit");
2696 session_ct.cancel();
2697 }
2698 }
2699 }
2700
2701 Ok(())
2702}
2703
2704#[cfg(feature = "oauth")]
2713fn install_oauth_proxy_routes(
2714 router: axum::Router,
2715 server_url: &str,
2716 oauth_config: &crate::oauth::OAuthConfig,
2717 auth_state: Option<&Arc<AuthState>>,
2718 max_request_body: usize,
2719 admin_role: &str,
2720) -> Result<axum::Router, RmcpServerKitError> {
2721 let Some(ref proxy) = oauth_config.proxy else {
2722 return Ok(router);
2723 };
2724
2725 let http = crate::oauth::OauthHttpClient::with_config(oauth_config)?;
2728
2729 let proxy_router = axum::Router::new();
2735
2736 let asm = crate::oauth::authorization_server_metadata(server_url, oauth_config);
2737 let proxy_router = proxy_router.route(
2738 "/.well-known/oauth-authorization-server",
2739 axum::routing::get(move || {
2740 let m = asm.clone();
2741 async move { axum::Json(m) }
2742 }),
2743 );
2744
2745 let proxy_authorize = proxy.clone();
2746 let proxy_router = proxy_router.route(
2747 "/authorize",
2748 axum::routing::get(
2749 move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
2750 let p = proxy_authorize.clone();
2751 async move { crate::oauth::handle_authorize(&p, &query.unwrap_or_default()) }
2752 },
2753 ),
2754 );
2755
2756 let proxy_token = proxy.clone();
2757 let token_http = http.clone();
2758 let proxy_router = proxy_router.route(
2759 "/token",
2760 axum::routing::post(move |body: String| {
2761 let p = proxy_token.clone();
2762 let h = token_http.clone();
2763 async move { crate::oauth::handle_token(&h, &p, &body).await }
2764 })
2765 .layer(axum::middleware::from_fn(
2766 oauth_token_cache_headers_middleware,
2767 )),
2768 );
2769
2770 let proxy_register = proxy.clone();
2771 let proxy_router = proxy_router.route(
2772 "/register",
2773 axum::routing::post(move |axum::Json(body): axum::Json<serde_json::Value>| {
2774 let p = proxy_register;
2775 async move { axum::Json(crate::oauth::handle_register(&p, &body)) }
2776 })
2777 .layer(axum::middleware::from_fn(
2778 oauth_token_cache_headers_middleware,
2779 )),
2780 );
2781
2782 let admin_routes_enabled = proxy.expose_admin_endpoints
2783 && (proxy.introspection_url.is_some() || proxy.revocation_url.is_some());
2784 if proxy.expose_admin_endpoints
2785 && !proxy.require_auth_on_admin_endpoints
2786 && proxy.allow_unauthenticated_admin_endpoints
2787 {
2788 tracing::warn!(
2792 "OAuth introspect/revoke endpoints are unauthenticated by explicit \
2793 allow_unauthenticated_admin_endpoints opt-out; ensure an \
2794 authenticated reverse proxy fronts these routes"
2795 );
2796 }
2797
2798 let admin_router = if admin_routes_enabled {
2799 build_oauth_admin_router(proxy, http, auth_state, admin_role)?
2800 } else {
2801 axum::Router::new()
2802 };
2803
2804 let proxy_router =
2808 proxy_router
2809 .merge(admin_router)
2810 .layer(tower_http::limit::RequestBodyLimitLayer::new(
2811 max_request_body,
2812 ));
2813
2814 let router = router.merge(proxy_router);
2815
2816 tracing::info!(
2817 introspect = proxy.expose_admin_endpoints && proxy.introspection_url.is_some(),
2818 revoke = proxy.expose_admin_endpoints && proxy.revocation_url.is_some(),
2819 max_request_body,
2820 "OAuth 2.1 proxy endpoints enabled (/authorize, /token, /register)"
2821 );
2822 Ok(router)
2823}
2824
2825#[cfg(feature = "oauth")]
2831fn build_oauth_admin_router(
2832 proxy: &crate::oauth::OAuthProxyConfig,
2833 http: crate::oauth::OauthHttpClient,
2834 auth_state: Option<&Arc<AuthState>>,
2835 admin_role: &str,
2836) -> Result<axum::Router, RmcpServerKitError> {
2837 let mut admin_router = axum::Router::new();
2838 if proxy.introspection_url.is_some() {
2839 let proxy_introspect = proxy.clone();
2840 let introspect_http = http.clone();
2841 admin_router = admin_router.route(
2842 "/introspect",
2843 axum::routing::post(move |body: String| {
2844 let p = proxy_introspect.clone();
2845 let h = introspect_http.clone();
2846 async move { crate::oauth::handle_introspect(&h, &p, &body).await }
2847 }),
2848 );
2849 }
2850 if proxy.revocation_url.is_some() {
2851 let proxy_revoke = proxy.clone();
2852 let revoke_http = http;
2853 admin_router = admin_router.route(
2854 "/revoke",
2855 axum::routing::post(move |body: String| {
2856 let p = proxy_revoke.clone();
2857 let h = revoke_http.clone();
2858 async move { crate::oauth::handle_revoke(&h, &p, &body).await }
2859 }),
2860 );
2861 }
2862
2863 let admin_router = admin_router.layer(axum::middleware::from_fn(
2864 oauth_token_cache_headers_middleware,
2865 ));
2866
2867 if proxy.require_auth_on_admin_endpoints {
2868 let Some(state) = auth_state else {
2869 return Err(RmcpServerKitError::Startup(
2870 "oauth proxy admin endpoints require auth state".into(),
2871 ));
2872 };
2873 let state_for_mw = Arc::clone(state);
2874 let required_role: Arc<str> = Arc::from(admin_role);
2875 Ok(admin_router
2881 .layer(axum::middleware::from_fn(move |req, next| {
2882 let r = Arc::clone(&required_role);
2883 crate::admin::require_admin_role(r, req, next)
2884 }))
2885 .layer(axum::middleware::from_fn(move |req, next| {
2886 let s = Arc::clone(&state_for_mw);
2887 auth_middleware(s, req, next)
2888 })))
2889 } else {
2890 Ok(admin_router)
2891 }
2892}
2893
2894#[allow(
2901 deprecated,
2902 reason = "internal metadata assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
2903)]
2904fn derive_server_url(config: &McpServerConfig) -> String {
2905 config.public_url.as_ref().map_or_else(
2906 || {
2907 let scheme = if config.tls_cert_path.is_some() {
2908 "https"
2909 } else {
2910 "http"
2911 };
2912 format!("{scheme}://{}", config.bind_addr)
2913 },
2914 |url| url.trim_end_matches('/').to_owned(),
2915 )
2916}
2917
2918fn derive_allowed_hosts(bind_addr: &str, public_url: Option<&str>) -> Vec<String> {
2923 let mut hosts = vec![
2924 "localhost".to_owned(),
2925 "127.0.0.1".to_owned(),
2926 "::1".to_owned(),
2927 ];
2928
2929 if let Some(url) = public_url
2930 && let Ok(uri) = url.parse::<axum::http::Uri>()
2931 && let Some(authority) = uri.authority()
2932 {
2933 let host = authority.host().to_owned();
2934 if !hosts.iter().any(|h| h == &host) {
2935 hosts.push(host);
2936 }
2937
2938 let authority = authority.as_str().to_owned();
2939 if !hosts.iter().any(|h| h == &authority) {
2940 hosts.push(authority);
2941 }
2942 }
2943
2944 if let Ok(uri) = format!("http://{bind_addr}").parse::<axum::http::Uri>()
2945 && let Some(authority) = uri.authority()
2946 {
2947 let host = authority.host().to_owned();
2948 if !hosts.iter().any(|h| h == &host) {
2949 hosts.push(host);
2950 }
2951
2952 let authority = authority.as_str().to_owned();
2953 if !hosts.iter().any(|h| h == &authority) {
2954 hosts.push(authority);
2955 }
2956 }
2957
2958 hosts
2959}
2960
2961impl axum::extract::connect_info::Connected<axum::serve::IncomingStream<'_, TlsListener>>
2974 for TlsConnInfo
2975{
2976 fn connect_info(target: axum::serve::IncomingStream<'_, TlsListener>) -> Self {
2977 let addr = *target.remote_addr();
2978 let identity = target.io().identity().cloned();
2979 Self::new(addr, identity)
2980 }
2981}
2982
2983const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
2990
2991const DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256;
2999
3000const TLS_ACCEPT_CHANNEL_CAPACITY: usize = 32;
3005
3006struct TlsListener {
3022 local_addr: SocketAddr,
3025 rx: mpsc::Receiver<(AuthenticatedTlsStream, SocketAddr)>,
3027 acceptor_task: tokio::task::JoinHandle<()>,
3030}
3031
3032impl TlsListener {
3033 fn new(
3034 inner: TcpListener,
3035 cert_path: &Path,
3036 key_path: &Path,
3037 mtls_config: Option<&MtlsConfig>,
3038 crl_set: Option<Arc<CrlSet>>,
3039 handshake_timeout: Duration,
3040 max_concurrent_handshakes: usize,
3041 ) -> anyhow::Result<Self> {
3042 rustls::crypto::ring::default_provider()
3044 .install_default()
3045 .ok();
3046
3047 let certs = load_certs(cert_path)?;
3048 let key = load_key(key_path)?;
3049
3050 let mtls_default_role;
3051
3052 let tls_config = if let Some(mtls) = mtls_config {
3053 mtls_default_role = mtls.default_role.clone();
3054 let verifier: Arc<dyn rustls::server::danger::ClientCertVerifier> = if mtls.crl_enabled
3055 {
3056 let Some(crl_set) = crl_set else {
3057 return Err(anyhow::anyhow!(
3058 "mTLS CRL verifier requested but CRL state was not initialized"
3059 ));
3060 };
3061 Arc::new(DynamicClientCertVerifier::new(crl_set))
3062 } else {
3063 let (_, root_store) = load_client_auth_roots(&mtls.ca_cert_path)?;
3064 if mtls.required {
3065 rustls::server::WebPkiClientVerifier::builder(root_store)
3066 .build()
3067 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
3068 } else {
3069 rustls::server::WebPkiClientVerifier::builder(root_store)
3070 .allow_unauthenticated()
3071 .build()
3072 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
3073 }
3074 };
3075
3076 tracing::info!(
3077 ca = %mtls.ca_cert_path.display(),
3078 required = mtls.required,
3079 crl_enabled = mtls.crl_enabled,
3080 "mTLS client auth configured"
3081 );
3082
3083 rustls::ServerConfig::builder_with_protocol_versions(&[
3084 &rustls::version::TLS12,
3085 &rustls::version::TLS13,
3086 ])
3087 .with_client_cert_verifier(verifier)
3088 .with_single_cert(certs, key)?
3089 } else {
3090 mtls_default_role = "viewer".to_owned();
3091 rustls::ServerConfig::builder_with_protocol_versions(&[
3092 &rustls::version::TLS12,
3093 &rustls::version::TLS13,
3094 ])
3095 .with_no_client_auth()
3096 .with_single_cert(certs, key)?
3097 };
3098
3099 let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
3100 tracing::info!(
3101 "TLS enabled (cert: {}, key: {})",
3102 cert_path.display(),
3103 key_path.display()
3104 );
3105 let local_addr = inner.local_addr()?;
3106 let (tx, rx) = mpsc::channel(TLS_ACCEPT_CHANNEL_CAPACITY);
3107 let acceptor_task = tokio::spawn(run_tls_acceptor(
3108 inner,
3109 acceptor,
3110 mtls_default_role,
3111 tx,
3112 handshake_timeout,
3113 max_concurrent_handshakes,
3114 ));
3115 Ok(Self {
3116 local_addr,
3117 rx,
3118 acceptor_task,
3119 })
3120 }
3121
3122 fn extract_handshake_identity(
3126 tls_stream: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
3127 default_role: &str,
3128 addr: SocketAddr,
3129 ) -> Option<AuthIdentity> {
3130 let (_, server_conn) = tls_stream.get_ref();
3131 let cert_der = server_conn.peer_certificates()?.first()?;
3132 let id = extract_mtls_identity(cert_der.as_ref(), default_role)?;
3133 tracing::debug!(name = %id.name, peer = %addr, "mTLS client cert accepted");
3134 Some(id)
3135 }
3136}
3137
3138async fn run_tls_acceptor(
3149 listener: TcpListener,
3150 acceptor: tokio_rustls::TlsAcceptor,
3151 default_role: String,
3152 tx: mpsc::Sender<(AuthenticatedTlsStream, SocketAddr)>,
3153 handshake_timeout: Duration,
3154 max_concurrent_handshakes: usize,
3155) {
3156 let inflight = Arc::new(Semaphore::new(max_concurrent_handshakes));
3157 loop {
3158 let Ok(permit) = Arc::clone(&inflight).acquire_owned().await else {
3162 return;
3164 };
3165 let (stream, addr) = match listener.accept().await {
3166 Ok(pair) => pair,
3167 Err(e) => {
3168 tracing::debug!("TCP accept error: {e}");
3169 continue;
3170 }
3171 };
3172 if tx.is_closed() {
3173 return;
3175 }
3176 let acceptor = acceptor.clone();
3177 let default_role = default_role.clone();
3178 let tx = tx.clone();
3179 tokio::spawn(async move {
3180 let _permit = permit;
3181 match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
3182 Ok(Ok(tls_stream)) => {
3183 let identity =
3184 TlsListener::extract_handshake_identity(&tls_stream, &default_role, addr);
3185 let wrapped = AuthenticatedTlsStream {
3186 inner: tls_stream,
3187 identity,
3188 };
3189 let _ = tx.send((wrapped, addr)).await;
3192 }
3193 Ok(Err(e)) => {
3194 tracing::debug!("TLS handshake failed from {addr}: {e}");
3195 }
3196 Err(_elapsed) => {
3197 tracing::debug!(
3198 "TLS handshake timed out from {addr} after {handshake_timeout:?}"
3199 );
3200 }
3201 }
3202 });
3203 }
3204}
3205
3206pub(crate) struct AuthenticatedTlsStream {
3218 inner: tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
3219 identity: Option<AuthIdentity>,
3220}
3221
3222impl AuthenticatedTlsStream {
3223 #[must_use]
3225 pub(crate) const fn identity(&self) -> Option<&AuthIdentity> {
3226 self.identity.as_ref()
3227 }
3228}
3229
3230impl std::fmt::Debug for AuthenticatedTlsStream {
3231 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3232 f.debug_struct("AuthenticatedTlsStream")
3233 .field("identity", &self.identity.as_ref().map(|id| &id.name))
3234 .finish_non_exhaustive()
3235 }
3236}
3237
3238impl tokio::io::AsyncRead for AuthenticatedTlsStream {
3239 fn poll_read(
3240 mut self: Pin<&mut Self>,
3241 cx: &mut std::task::Context<'_>,
3242 buf: &mut tokio::io::ReadBuf<'_>,
3243 ) -> std::task::Poll<std::io::Result<()>> {
3244 Pin::new(&mut self.inner).poll_read(cx, buf)
3245 }
3246}
3247
3248impl tokio::io::AsyncWrite for AuthenticatedTlsStream {
3249 fn poll_write(
3250 mut self: Pin<&mut Self>,
3251 cx: &mut std::task::Context<'_>,
3252 buf: &[u8],
3253 ) -> std::task::Poll<std::io::Result<usize>> {
3254 Pin::new(&mut self.inner).poll_write(cx, buf)
3255 }
3256
3257 fn poll_flush(
3258 mut self: Pin<&mut Self>,
3259 cx: &mut std::task::Context<'_>,
3260 ) -> std::task::Poll<std::io::Result<()>> {
3261 Pin::new(&mut self.inner).poll_flush(cx)
3262 }
3263
3264 fn poll_shutdown(
3265 mut self: Pin<&mut Self>,
3266 cx: &mut std::task::Context<'_>,
3267 ) -> std::task::Poll<std::io::Result<()>> {
3268 Pin::new(&mut self.inner).poll_shutdown(cx)
3269 }
3270
3271 fn poll_write_vectored(
3272 mut self: Pin<&mut Self>,
3273 cx: &mut std::task::Context<'_>,
3274 bufs: &[std::io::IoSlice<'_>],
3275 ) -> std::task::Poll<std::io::Result<usize>> {
3276 Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
3277 }
3278
3279 fn is_write_vectored(&self) -> bool {
3280 self.inner.is_write_vectored()
3281 }
3282}
3283
3284impl axum::serve::Listener for TlsListener {
3285 type Io = AuthenticatedTlsStream;
3286 type Addr = SocketAddr;
3287
3288 async fn accept(&mut self) -> (Self::Io, Self::Addr) {
3294 if let Some(pair) = self.rx.recv().await {
3295 return pair;
3296 }
3297 tracing::error!("TLS acceptor task terminated; no further connections will be accepted");
3303 std::future::pending().await
3304 }
3305
3306 fn local_addr(&self) -> std::io::Result<Self::Addr> {
3307 Ok(self.local_addr)
3308 }
3309}
3310
3311impl Drop for TlsListener {
3312 fn drop(&mut self) {
3313 self.acceptor_task.abort();
3316 }
3317}
3318
3319fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
3320 use rustls::pki_types::pem::PemObject;
3321 let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_file_iter(path)
3322 .map_err(|e| anyhow::anyhow!("failed to read certs from {}: {e}", path.display()))?
3323 .collect::<Result<_, _>>()
3324 .map_err(|e| anyhow::anyhow!("invalid cert in {}: {e}", path.display()))?;
3325 anyhow::ensure!(
3326 !certs.is_empty(),
3327 "no certificates found in {}",
3328 path.display()
3329 );
3330 Ok(certs)
3331}
3332
3333fn load_client_auth_roots(
3334 path: &Path,
3335) -> anyhow::Result<(
3336 Vec<rustls::pki_types::CertificateDer<'static>>,
3337 Arc<RootCertStore>,
3338)> {
3339 let ca_certs = load_certs(path)?;
3340 let mut root_store = RootCertStore::empty();
3341 for cert in &ca_certs {
3342 root_store
3343 .add(cert.clone())
3344 .map_err(|error| anyhow::anyhow!("invalid CA cert: {error}"))?;
3345 }
3346
3347 Ok((ca_certs, Arc::new(root_store)))
3348}
3349
3350fn load_key(path: &Path) -> anyhow::Result<rustls::pki_types::PrivateKeyDer<'static>> {
3351 use rustls::pki_types::pem::PemObject;
3352 rustls::pki_types::PrivateKeyDer::from_pem_file(path)
3353 .map_err(|e| anyhow::anyhow!("failed to read key from {}: {e}", path.display()))
3354}
3355
3356#[allow(
3358 clippy::unused_async,
3359 reason = "axum route handler signature requires `async fn` even when the body is synchronous"
3360)]
3361async fn healthz() -> impl IntoResponse {
3362 axum::Json(serde_json::json!({
3363 "status": "ok",
3364 }))
3365}
3366
3367fn version_payload(name: &str, version: &str, expose_build_metadata: bool) -> serde_json::Value {
3377 let mut map = serde_json::Map::new();
3378 map.insert("name".into(), name.into());
3379 map.insert("version".into(), version.into());
3380 map.insert(
3381 "rmcp_server_kit_version".into(),
3382 env!("CARGO_PKG_VERSION").into(),
3383 );
3384 if expose_build_metadata {
3385 map.insert(
3386 "build_git_sha".into(),
3387 option_env!("RMCP_SERVER_KIT_BUILD_SHA")
3388 .unwrap_or("unknown")
3389 .into(),
3390 );
3391 map.insert(
3392 "build_timestamp".into(),
3393 option_env!("RMCP_SERVER_KIT_BUILD_TIME")
3394 .unwrap_or("unknown")
3395 .into(),
3396 );
3397 map.insert(
3398 "rust_version".into(),
3399 option_env!("RMCP_SERVER_KIT_RUSTC_VERSION")
3400 .unwrap_or("unknown")
3401 .into(),
3402 );
3403 }
3404 serde_json::Value::Object(map)
3405}
3406
3407fn serialize_version_payload(name: &str, version: &str, expose_build_metadata: bool) -> Arc<[u8]> {
3417 let value = version_payload(name, version, expose_build_metadata);
3418 serde_json::to_vec(&value).map_or_else(|_| Arc::from(&b"{}"[..]), Arc::from)
3419}
3420
3421async fn readyz(check: ReadinessCheck) -> impl IntoResponse {
3426 let status = check().await;
3427 let ready = status
3428 .get("ready")
3429 .and_then(serde_json::Value::as_bool)
3430 .unwrap_or(false);
3431 let code = if ready {
3432 axum::http::StatusCode::OK
3433 } else {
3434 axum::http::StatusCode::SERVICE_UNAVAILABLE
3435 };
3436 (code, axum::Json(status))
3437}
3438
3439async fn shutdown_signal() {
3443 let ctrl_c = tokio::signal::ctrl_c();
3444
3445 #[cfg(unix)]
3446 {
3447 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
3448 Ok(mut term) => {
3449 tokio::select! {
3452 _ = ctrl_c => {}
3453 _ = term.recv() => {}
3454 }
3455 }
3456 Err(e) => {
3457 tracing::warn!(error = %e, "failed to register SIGTERM handler, using SIGINT only");
3458 ctrl_c.await.ok();
3459 }
3460 }
3461 }
3462
3463 #[cfg(not(unix))]
3464 {
3465 ctrl_c.await.ok();
3466 }
3467}
3468
3469#[cfg(feature = "metrics")]
3486fn metrics_labels(req: &Request<Body>) -> (&'static str, String) {
3487 let method = match *req.method() {
3488 axum::http::Method::GET => "GET",
3489 axum::http::Method::POST => "POST",
3490 axum::http::Method::PUT => "PUT",
3491 axum::http::Method::PATCH => "PATCH",
3492 axum::http::Method::DELETE => "DELETE",
3493 axum::http::Method::HEAD => "HEAD",
3494 axum::http::Method::OPTIONS => "OPTIONS",
3495 axum::http::Method::TRACE => "TRACE",
3496 axum::http::Method::CONNECT => "CONNECT",
3497 _ => "OTHER",
3500 };
3501
3502 let path = req
3503 .extensions()
3504 .get::<axum::extract::MatchedPath>()
3505 .map_or_else(
3506 || {
3507 let raw = req.uri().path();
3508 if raw == "/mcp" || raw.starts_with("/mcp/") {
3509 "/mcp".to_owned()
3510 } else {
3511 "<unmatched>".to_owned()
3512 }
3513 },
3514 |matched| matched.as_str().to_owned(),
3515 );
3516
3517 (method, path)
3518}
3519
3520#[cfg(feature = "metrics")]
3530async fn metrics_middleware(
3531 metrics: Arc<crate::metrics::McpMetrics>,
3532 mut req: Request<Body>,
3533 next: Next,
3534) -> axum::response::Response {
3535 let (method, path) = metrics_labels(&req);
3536 let start = std::time::Instant::now();
3537
3538 req.extensions_mut().insert(Arc::clone(&metrics));
3539 let response = next.run(req).await;
3540
3541 let mut status_buf = core::fmt::NumBuffer::<u16>::new();
3542 let status = response.status().as_u16().format_into(&mut status_buf);
3543 let duration = start.elapsed().as_secs_f64();
3544
3545 metrics
3546 .http_requests_total
3547 .with_label_values(&[method, &path, status])
3548 .inc();
3549 metrics
3550 .http_request_duration_seconds
3551 .with_label_values(&[method, &path])
3552 .observe(duration);
3553
3554 response
3555}
3556
3557async fn security_headers_middleware(
3571 is_tls: bool,
3572 cfg: Arc<SecurityHeadersConfig>,
3573 req: Request<Body>,
3574 next: Next,
3575) -> axum::response::Response {
3576 use axum::http::{HeaderName, header};
3577
3578 let mut resp = next.run(req).await;
3579 let headers = resp.headers_mut();
3580
3581 headers.remove(header::SERVER);
3583 headers.remove(HeaderName::from_static("x-powered-by"));
3584
3585 apply_security_header(
3586 headers,
3587 header::X_CONTENT_TYPE_OPTIONS,
3588 cfg.x_content_type_options.as_deref(),
3589 "nosniff",
3590 );
3591 apply_security_header(
3592 headers,
3593 header::X_FRAME_OPTIONS,
3594 cfg.x_frame_options.as_deref(),
3595 "deny",
3596 );
3597 apply_security_header(
3598 headers,
3599 header::CACHE_CONTROL,
3600 cfg.cache_control.as_deref(),
3601 "no-store, max-age=0",
3602 );
3603 apply_security_header(
3604 headers,
3605 header::REFERRER_POLICY,
3606 cfg.referrer_policy.as_deref(),
3607 "no-referrer",
3608 );
3609 apply_security_header(
3610 headers,
3611 HeaderName::from_static("cross-origin-opener-policy"),
3612 cfg.cross_origin_opener_policy.as_deref(),
3613 "same-origin",
3614 );
3615 apply_security_header(
3616 headers,
3617 HeaderName::from_static("cross-origin-resource-policy"),
3618 cfg.cross_origin_resource_policy.as_deref(),
3619 "same-origin",
3620 );
3621 apply_security_header(
3622 headers,
3623 HeaderName::from_static("cross-origin-embedder-policy"),
3624 cfg.cross_origin_embedder_policy.as_deref(),
3625 "require-corp",
3626 );
3627 apply_security_header(
3628 headers,
3629 HeaderName::from_static("permissions-policy"),
3630 cfg.permissions_policy.as_deref(),
3631 "accelerometer=(), camera=(), geolocation=(), microphone=()",
3632 );
3633 apply_security_header(
3634 headers,
3635 HeaderName::from_static("x-permitted-cross-domain-policies"),
3636 cfg.x_permitted_cross_domain_policies.as_deref(),
3637 "none",
3638 );
3639 apply_security_header(
3640 headers,
3641 HeaderName::from_static("content-security-policy"),
3642 cfg.content_security_policy.as_deref(),
3643 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests",
3644 );
3645 apply_security_header(
3646 headers,
3647 HeaderName::from_static("x-dns-prefetch-control"),
3648 cfg.x_dns_prefetch_control.as_deref(),
3649 "off",
3650 );
3651
3652 if is_tls {
3653 apply_security_header(
3654 headers,
3655 header::STRICT_TRANSPORT_SECURITY,
3656 cfg.strict_transport_security.as_deref(),
3657 "max-age=63072000; includeSubDomains",
3658 );
3659 }
3660
3661 resp
3662}
3663
3664fn apply_security_header(
3675 headers: &mut axum::http::HeaderMap,
3676 name: axum::http::HeaderName,
3677 override_value: Option<&str>,
3678 default: &'static str,
3679) {
3680 use axum::http::HeaderValue;
3681
3682 match override_value {
3683 None => {
3684 headers.insert(name, HeaderValue::from_static(default));
3685 }
3686 Some("") => {
3687 }
3689 Some(v) => match HeaderValue::from_str(v) {
3690 Ok(hv) => {
3691 headers.insert(name, hv);
3692 }
3693 Err(err) => {
3694 tracing::error!(
3695 header = %name,
3696 error = %err,
3697 "invalid security header override reached middleware; using default"
3698 );
3699 headers.insert(name, HeaderValue::from_static(default));
3700 }
3701 },
3702 }
3703}
3704
3705fn validate_security_headers(cfg: &SecurityHeadersConfig) -> Result<(), RmcpServerKitError> {
3716 use axum::http::HeaderValue;
3717
3718 let fields: &[(&str, Option<&str>)] = &[
3719 (
3720 "x_content_type_options",
3721 cfg.x_content_type_options.as_deref(),
3722 ),
3723 ("x_frame_options", cfg.x_frame_options.as_deref()),
3724 ("cache_control", cfg.cache_control.as_deref()),
3725 ("referrer_policy", cfg.referrer_policy.as_deref()),
3726 (
3727 "cross_origin_opener_policy",
3728 cfg.cross_origin_opener_policy.as_deref(),
3729 ),
3730 (
3731 "cross_origin_resource_policy",
3732 cfg.cross_origin_resource_policy.as_deref(),
3733 ),
3734 (
3735 "cross_origin_embedder_policy",
3736 cfg.cross_origin_embedder_policy.as_deref(),
3737 ),
3738 ("permissions_policy", cfg.permissions_policy.as_deref()),
3739 (
3740 "x_permitted_cross_domain_policies",
3741 cfg.x_permitted_cross_domain_policies.as_deref(),
3742 ),
3743 (
3744 "content_security_policy",
3745 cfg.content_security_policy.as_deref(),
3746 ),
3747 (
3748 "x_dns_prefetch_control",
3749 cfg.x_dns_prefetch_control.as_deref(),
3750 ),
3751 (
3752 "strict_transport_security",
3753 cfg.strict_transport_security.as_deref(),
3754 ),
3755 ];
3756
3757 for (field, value) in fields {
3758 let Some(v) = value else { continue };
3759 if v.is_empty() {
3760 continue;
3761 }
3762 if let Err(err) = HeaderValue::from_str(v) {
3763 return Err(RmcpServerKitError::Config(format!(
3764 "invalid security_headers.{field}: {err}"
3765 )));
3766 }
3767 }
3768
3769 if let Some(v) = cfg.strict_transport_security.as_deref()
3770 && !v.is_empty()
3771 && v.to_ascii_lowercase().contains("preload")
3772 {
3773 return Err(RmcpServerKitError::Config(format!(
3774 "invalid security_headers.strict_transport_security: {v:?} contains the `preload` directive; \
3775 HSTS preload must be opted into explicitly via a dedicated builder, not via this knob"
3776 )));
3777 }
3778
3779 Ok(())
3780}
3781
3782#[cfg(feature = "oauth")]
3797async fn oauth_token_cache_headers_middleware(
3798 req: Request<Body>,
3799 next: Next,
3800) -> axum::response::Response {
3801 use axum::http::{HeaderValue, header};
3802
3803 let mut resp = next.run(req).await;
3804 let headers = resp.headers_mut();
3805 headers.insert(header::PRAGMA, HeaderValue::from_static("no-cache"));
3806 headers.append(header::VARY, HeaderValue::from_static("Authorization"));
3807 resp
3808}
3809
3810async fn normalize_peer_addr_middleware(
3841 resolver: Option<Arc<ForwardResolver>>,
3842 mut req: Request<Body>,
3843 next: Next,
3844) -> axum::response::Response {
3845 let direct = req
3846 .extensions()
3847 .get::<ConnectInfo<SocketAddr>>()
3848 .map(|ci| ci.0);
3849 let from_tls = req
3850 .extensions()
3851 .get::<ConnectInfo<TlsConnInfo>>()
3852 .map(|ci| ci.0.addr);
3853 if let Some(addr) = direct.or(from_tls) {
3854 if direct.is_none() {
3855 req.extensions_mut().insert(ConnectInfo(addr));
3856 }
3857 req.extensions_mut().insert(PeerAddr::new(addr));
3858 let client_ip = match &resolver {
3859 Some(r) => crate::forwarded::resolve_client_ip(
3860 addr.ip(),
3861 req.headers(),
3862 &r.trusted,
3863 r.mode,
3864 r.max_scanned_entries,
3865 )
3866 .unwrap_or_else(|reason| {
3867 tracing::debug!(
3868 reason = ?reason,
3869 "forwarded-header resolution fell back to direct peer"
3870 );
3871 addr.ip()
3872 }),
3873 None => addr.ip(),
3874 };
3875 req.extensions_mut().insert(ClientIp::new(client_ip));
3876 }
3877 next.run(req).await
3878}
3879
3880fn parse_proxy_net(entry: &str) -> Option<ipnet::IpNet> {
3883 if let Ok(net) = entry.parse::<ipnet::IpNet>() {
3884 return Some(net);
3885 }
3886 entry.parse::<IpAddr>().ok().map(ipnet::IpNet::from)
3887}
3888
3889pub(crate) fn validate_trusted_proxy_entry(entry: &str) -> Result<(), String> {
3899 match parse_proxy_net(entry) {
3900 None => Err(format!(
3901 "trusted_proxies entry {entry:?} is neither a CIDR nor an IP address"
3902 )),
3903 Some(net) if net.prefix_len() == 0 => Err(format!(
3904 "trusted_proxies entry {entry:?}: prefix length 0 is forbidden (marks every peer trusted, enabling client-IP spoofing)"
3905 )),
3906 Some(_) => Ok(()),
3907 }
3908}
3909
3910pub(crate) fn limiter_client_ip(extensions: &axum::http::Extensions) -> Option<IpAddr> {
3914 if let Some(client) = extensions.get::<ClientIp>() {
3915 return Some(client.ip);
3916 }
3917 extensions
3918 .get::<ConnectInfo<SocketAddr>>()
3919 .map(|ci| ci.0.ip())
3920 .or_else(|| {
3921 extensions
3922 .get::<ConnectInfo<TlsConnInfo>>()
3923 .map(|ci| ci.0.addr.ip())
3924 })
3925}
3926
3927#[derive(Clone, PartialEq, Eq, Hash, Debug)]
3940pub(crate) enum RateLimitKey {
3941 Ip(IpAddr),
3943 Unattributed,
3945}
3946
3947impl std::fmt::Display for RateLimitKey {
3948 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3949 match self {
3950 Self::Ip(ip) => write!(f, "{ip}"),
3951 Self::Unattributed => f.write_str("unattributed"),
3952 }
3953 }
3954}
3955
3956static UNATTRIBUTED_WARNED: std::sync::atomic::AtomicBool =
3958 std::sync::atomic::AtomicBool::new(false);
3959
3960pub(crate) fn limiter_client_key(extensions: &axum::http::Extensions) -> RateLimitKey {
3971 if let Some(ip) = limiter_client_ip(extensions) {
3972 return RateLimitKey::Ip(ip);
3973 }
3974 if !UNATTRIBUTED_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
3975 tracing::warn!(
3976 "request carries no resolvable client address; rate limiting is \
3977 falling back to a single shared bucket. This indicates \
3978 rmcp-server-kit middleware composed outside serve()."
3979 );
3980 }
3981 RateLimitKey::Unattributed
3982}
3983
3984pub(crate) type ExtraRouteRateLimiter = BoundedKeyedLimiter<RateLimitKey>;
3988
3989const EXTRA_ROUTE_MAX_TRACKED_KEYS: usize = 10_000;
3995
3996const EXTRA_ROUTE_IDLE_EVICTION: Duration = Duration::from_mins(15);
3999
4000fn build_extra_route_rate_limiter_with_policy(
4007 per_minute: u32,
4008 burst: Option<u32>,
4009 key_eviction_policy: KeyEvictionPolicy,
4010 max_tracked_keys: NonZeroUsize,
4011) -> Arc<ExtraRouteRateLimiter> {
4012 let rate = std::num::NonZeroU32::new(per_minute.max(1)).unwrap_or(std::num::NonZeroU32::MIN);
4013 let mut quota = governor::Quota::per_minute(rate);
4014 if let Some(b) = burst.and_then(std::num::NonZeroU32::new) {
4015 quota = quota.allow_burst(b);
4016 }
4017 Arc::new(BoundedKeyedLimiter::new_with_policy(
4018 quota,
4019 max_tracked_keys,
4020 EXTRA_ROUTE_IDLE_EVICTION,
4021 key_eviction_policy,
4022 ))
4023}
4024
4025async fn extra_route_rate_limit_middleware(
4050 limiter: Arc<ExtraRouteRateLimiter>,
4051 exempt: Arc<std::collections::HashSet<String>>,
4052 req: Request<Body>,
4053 next: Next,
4054) -> axum::response::Response {
4055 if exempt.contains(req.uri().path()) {
4056 return next.run(req).await;
4057 }
4058 let peer_key = limiter_client_key(req.extensions());
4059 match limiter.check_key_detailed(&peer_key) {
4060 Ok(()) => {}
4061 Err(BoundedLimiterDeny::RateLimited(wait)) => {
4062 #[cfg(feature = "metrics")]
4063 crate::metrics::record_rate_limit_deny(req.extensions(), "extra_route");
4064 tracing::warn!(rate_limit_key = %peer_key, "extra route request rate limited");
4065 return RmcpServerKitError::RateLimitedFor {
4066 message: "too many requests to application routes from this source".into(),
4067 retry_after: wait,
4068 }
4069 .into_response();
4070 }
4071 Err(BoundedLimiterDeny::CapacityFull) => {
4072 tracing::warn!(
4073 rate_limit_key = %peer_key,
4074 "extra route limiter rejected unseen key because tracked-key capacity is full"
4075 );
4076 return (
4077 axum::http::StatusCode::SERVICE_UNAVAILABLE,
4078 "rate limiter capacity exhausted",
4079 )
4080 .into_response();
4081 }
4082 }
4083 next.run(req).await
4084}
4085
4086async fn origin_check_middleware(
4092 allowed: Arc<[String]>,
4093 log_request_headers: bool,
4094 req: Request<Body>,
4095 next: Next,
4096) -> axum::response::Response {
4097 let method = req.method().clone();
4098 let path = req.uri().path().to_owned();
4099
4100 log_incoming_request(&method, &path, req.headers(), log_request_headers);
4101
4102 if let Some(origin) = req.headers().get(axum::http::header::ORIGIN) {
4103 let origin_str = origin.to_str().unwrap_or("");
4104 if !allowed.iter().any(|a| a == origin_str) {
4105 tracing::warn!(
4106 origin = origin_str,
4107 %method,
4108 %path,
4109 allowed = ?&*allowed,
4110 "rejected request: Origin not allowed"
4111 );
4112 return (
4113 axum::http::StatusCode::FORBIDDEN,
4114 "Forbidden: Origin not allowed",
4115 )
4116 .into_response();
4117 }
4118 }
4119 next.run(req).await
4120}
4121
4122fn log_incoming_request(
4125 method: &axum::http::Method,
4126 path: &str,
4127 headers: &axum::http::HeaderMap,
4128 log_request_headers: bool,
4129) {
4130 if log_request_headers {
4131 tracing::debug!(
4132 %method,
4133 %path,
4134 headers = %format_request_headers_for_log(headers),
4135 "incoming request"
4136 );
4137 } else {
4138 tracing::debug!(%method, %path, "incoming request");
4139 }
4140}
4141
4142const REDACTED_LOG_HEADERS: [&str; 6] = [
4150 "authorization",
4151 "cookie",
4152 "proxy-authorization",
4153 "forwarded",
4154 "x-forwarded-for",
4155 "x-real-ip",
4156];
4157
4158fn format_request_headers_for_log(headers: &axum::http::HeaderMap) -> String {
4159 headers
4160 .iter()
4161 .map(|(k, v)| {
4162 let name = k.as_str();
4163 if REDACTED_LOG_HEADERS.contains(&name) {
4164 format!("{name}: [REDACTED]")
4165 } else {
4166 format!("{name}: {}", v.to_str().unwrap_or("<non-utf8>"))
4167 }
4168 })
4169 .collect::<Vec<_>>()
4170 .join(", ")
4171}
4172
4173#[allow(
4197 clippy::cognitive_complexity,
4198 reason = "complexity is purely tracing macro expansion (info/warn + match arms); 18 lines of straight-line code, nothing meaningful to extract"
4199)]
4200pub async fn serve_stdio<H>(handler: H) -> Result<(), RmcpServerKitError>
4201where
4202 H: ServerHandler + 'static,
4203{
4204 use rmcp::ServiceExt as _;
4205
4206 tracing::info!("stdio transport: serving on stdin/stdout");
4207 tracing::warn!("stdio mode: auth, RBAC, TLS, and Origin checks are DISABLED");
4208
4209 let transport = rmcp::transport::io::stdio();
4210
4211 let service = handler
4212 .serve(transport)
4213 .await
4214 .map_err(|e| RmcpServerKitError::Startup(format!("stdio initialize failed: {e}")))?;
4215
4216 if let Err(e) = service.waiting().await {
4217 tracing::warn!(error = %e, "stdio session ended with error");
4218 }
4219 tracing::info!("stdio session ended");
4220 Ok(())
4221}
4222
4223#[allow(
4224 deprecated,
4225 reason = "builder methods are the sanctioned transition layer for deprecated public fields"
4226)]
4227impl McpServerConfig {
4228 #[must_use]
4232 pub fn with_tls_paths(mut self, cert_path: Option<PathBuf>, key_path: Option<PathBuf>) -> Self {
4233 self.tls_cert_path = cert_path;
4234 self.tls_key_path = key_path;
4235 self
4236 }
4237
4238 #[must_use]
4242 pub fn with_tls_cert_path(mut self, cert_path: impl Into<PathBuf>) -> Self {
4243 self.tls_cert_path = Some(cert_path.into());
4244 self
4245 }
4246
4247 #[must_use]
4251 pub fn with_tls_key_path(mut self, key_path: impl Into<PathBuf>) -> Self {
4252 self.tls_key_path = Some(key_path.into());
4253 self
4254 }
4255
4256 #[must_use]
4258 pub fn with_optional_auth(mut self, auth: Option<AuthConfig>) -> Self {
4259 self.auth = auth;
4260 self
4261 }
4262
4263 #[must_use]
4265 pub fn with_optional_session_binding_secret(mut self, secret: Option<SecretString>) -> Self {
4266 self.session_binding_secret = secret;
4267 self
4268 }
4269
4270 #[must_use]
4272 pub fn with_optional_tool_rate_limit(mut self, per_minute: Option<u32>) -> Self {
4273 self.tool_rate_limit = per_minute;
4274 self
4275 }
4276
4277 #[must_use]
4279 pub fn with_optional_tool_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
4280 self.tool_rate_limit_burst = burst;
4281 self
4282 }
4283
4284 #[must_use]
4286 pub fn with_optional_extra_route_rate_limit(mut self, per_minute: Option<u32>) -> Self {
4287 self.extra_route_rate_limit = per_minute;
4288 self
4289 }
4290
4291 #[must_use]
4293 pub fn with_optional_extra_route_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
4294 self.extra_route_rate_limit_burst = burst;
4295 self
4296 }
4297
4298 #[must_use]
4300 pub fn with_optional_forwarded_header(mut self, mode: Option<ForwardedHeaderMode>) -> Self {
4301 self.forwarded_header = mode;
4302 self
4303 }
4304
4305 #[must_use]
4307 pub fn with_optional_public_url(mut self, url: Option<String>) -> Self {
4308 self.public_url = url;
4309 self
4310 }
4311
4312 #[must_use]
4316 pub fn with_compression_min_size(mut self, min_size: u16) -> Self {
4317 self.compression_min_size = min_size;
4318 self
4319 }
4320
4321 #[must_use]
4323 pub fn with_compression_enabled(mut self, enabled: bool) -> Self {
4324 self.compression_enabled = enabled;
4325 self
4326 }
4327
4328 #[must_use]
4330 pub fn with_optional_max_concurrent_requests(mut self, limit: Option<usize>) -> Self {
4331 self.max_concurrent_requests = limit;
4332 self
4333 }
4334
4335 #[must_use]
4337 pub fn with_admin_enabled(mut self, enabled: bool) -> Self {
4338 self.admin_enabled = enabled;
4339 self
4340 }
4341
4342 #[must_use]
4345 pub fn with_admin_role(mut self, role: impl Into<String>) -> Self {
4346 self.admin_role = role.into();
4347 self
4348 }
4349
4350 #[must_use]
4352 pub fn with_expose_build_metadata(mut self, enabled: bool) -> Self {
4353 self.expose_build_metadata = enabled;
4354 self
4355 }
4356}
4357
4358fn warn_security_header_overrides(cfg: &SecurityHeadersConfig) {
4359 for (field, value) in security_header_overrides(cfg) {
4360 let action = if value.is_empty() {
4361 "omitted"
4362 } else {
4363 "overridden"
4364 };
4365 tracing::warn!(
4366 security_header = field,
4367 action,
4368 "security header configured; inspect server.security_headers.<security_header>"
4369 );
4370 }
4371}
4372
4373fn security_header_overrides(
4374 cfg: &SecurityHeadersConfig,
4375) -> impl Iterator<Item = (&'static str, &str)> {
4376 [
4377 (
4378 "x_content_type_options",
4379 cfg.x_content_type_options.as_deref(),
4380 ),
4381 ("x_frame_options", cfg.x_frame_options.as_deref()),
4382 ("cache_control", cfg.cache_control.as_deref()),
4383 ("referrer_policy", cfg.referrer_policy.as_deref()),
4384 (
4385 "cross_origin_opener_policy",
4386 cfg.cross_origin_opener_policy.as_deref(),
4387 ),
4388 (
4389 "cross_origin_resource_policy",
4390 cfg.cross_origin_resource_policy.as_deref(),
4391 ),
4392 (
4393 "cross_origin_embedder_policy",
4394 cfg.cross_origin_embedder_policy.as_deref(),
4395 ),
4396 ("permissions_policy", cfg.permissions_policy.as_deref()),
4397 (
4398 "x_permitted_cross_domain_policies",
4399 cfg.x_permitted_cross_domain_policies.as_deref(),
4400 ),
4401 (
4402 "content_security_policy",
4403 cfg.content_security_policy.as_deref(),
4404 ),
4405 (
4406 "x_dns_prefetch_control",
4407 cfg.x_dns_prefetch_control.as_deref(),
4408 ),
4409 (
4410 "strict_transport_security",
4411 cfg.strict_transport_security.as_deref(),
4412 ),
4413 ]
4414 .into_iter()
4415 .filter_map(|(field, value)| value.map(|v| (field, v)))
4416}
4417
4418fn check_auth_capacity_knobs(auth: Option<&AuthConfig>) -> Result<(), RmcpServerKitError> {
4419 if let Some(auth_cfg) = auth {
4420 if let Some(rl) = &auth_cfg.rate_limit {
4421 (rl.max_attempts_per_minute != 0).ok_or_else(|| {
4422 RmcpServerKitError::Config(
4423 "auth.rate_limit.max_attempts_per_minute must be nonzero".into(),
4424 )
4425 })?;
4426 (rl.pre_auth_max_per_minute != Some(0)).ok_or_else(|| {
4431 RmcpServerKitError::Config(
4432 "auth.rate_limit.pre_auth_max_per_minute must be nonzero when set".into(),
4433 )
4434 })?;
4435 }
4436 if let Some(mtls) = &auth_cfg.mtls {
4437 check_mtls_capacity_knobs(mtls)?;
4438 }
4439 auth_cfg.check_oauth_feature()?;
4440 }
4441 Ok(())
4442}
4443
4444fn check_mtls_capacity_knobs(mtls: &MtlsConfig) -> Result<(), RmcpServerKitError> {
4445 (mtls.crl_max_concurrent_fetches != 0).ok_or_else(|| {
4446 RmcpServerKitError::Config("auth.mtls.crl_max_concurrent_fetches must be nonzero".into())
4447 })?;
4448 (mtls.crl_discovery_rate_per_min != 0).ok_or_else(|| {
4449 RmcpServerKitError::Config("auth.mtls.crl_discovery_rate_per_min must be nonzero".into())
4450 })?;
4451 (mtls.crl_max_host_semaphores != 0).ok_or_else(|| {
4452 RmcpServerKitError::Config("auth.mtls.crl_max_host_semaphores must be nonzero".into())
4453 })?;
4454 (mtls.crl_max_seen_urls != 0).ok_or_else(|| {
4455 RmcpServerKitError::Config("auth.mtls.crl_max_seen_urls must be nonzero".into())
4456 })?;
4457 (mtls.crl_max_cache_entries != 0).ok_or_else(|| {
4458 RmcpServerKitError::Config("auth.mtls.crl_max_cache_entries must be nonzero".into())
4459 })?;
4460 (mtls.crl_max_response_bytes != 0).ok_or_else(|| {
4465 RmcpServerKitError::Config("auth.mtls.crl_max_response_bytes must be nonzero".into())
4466 })?;
4467 Ok(())
4468}
4469
4470#[cfg(test)]
4471mod tests {
4472 #![allow(
4473 clippy::unwrap_used,
4474 clippy::expect_used,
4475 clippy::panic,
4476 clippy::indexing_slicing,
4477 clippy::unwrap_in_result,
4478 clippy::print_stdout,
4479 clippy::print_stderr,
4480 deprecated,
4481 reason = "internal unit tests legitimately read/write the deprecated `pub` fields they were designed to verify"
4482 )]
4483 use std::{sync::Arc, time::Duration};
4484
4485 use axum::{
4486 body::Body,
4487 http::{Request, StatusCode, header},
4488 response::IntoResponse,
4489 };
4490 use http_body_util::BodyExt;
4491 use tower::ServiceExt as _;
4492
4493 use super::*;
4494
4495 #[tokio::test]
4498 async fn external_shutdown_bridge_exits_when_internal_token_cancels() {
4499 let external = CancellationToken::new();
4500 let internal = CancellationToken::new();
4501 let bridge = spawn_external_shutdown_bridge(external.clone(), internal.clone());
4502
4503 internal.cancel();
4506
4507 let joined = tokio::time::timeout(Duration::from_secs(2), bridge).await;
4508 assert!(
4509 joined.is_ok(),
4510 "bridge task must exit once the internal token is cancelled, \
4511 otherwise it leaks for the lifetime of the process"
4512 );
4513 }
4514
4515 #[tokio::test]
4516 async fn external_shutdown_bridge_still_forwards_external_cancel() {
4517 let external = CancellationToken::new();
4518 let internal = CancellationToken::new();
4519 let bridge = spawn_external_shutdown_bridge(external.clone(), internal.clone());
4520
4521 external.cancel();
4522
4523 let joined = tokio::time::timeout(Duration::from_secs(2), bridge).await;
4524 assert!(joined.is_ok(), "bridge task must exit on external cancel");
4525 assert!(
4526 internal.is_cancelled(),
4527 "external cancellation must still propagate to the internal token"
4528 );
4529 }
4530
4531 #[test]
4532 fn cancel_on_drop_cancels_its_token() {
4533 let ct = CancellationToken::new();
4534 {
4535 let _guard = CancelOnDrop(ct.clone());
4536 assert!(!ct.is_cancelled());
4537 }
4538 assert!(
4539 ct.is_cancelled(),
4540 "dropping the guard must cancel background startup tasks"
4541 );
4542 }
4543
4544 #[test]
4545 fn validate_rejects_mtls_without_tls() {
4546 for (cert, key) in [
4547 (None, None),
4548 (Some("cert.pem"), None),
4549 (None, Some("key.pem")),
4550 ] {
4551 let mut auth = AuthConfig::with_keys(vec![]);
4552 auth.mtls = Some(valid_mtls_config());
4553 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4554 cfg.tls_cert_path = cert.map(Into::into);
4555 cfg.tls_key_path = key.map(Into::into);
4556
4557 let err = cfg
4558 .validate()
4559 .expect_err("mTLS without both TLS paths must be rejected");
4560 let msg = err.to_string();
4561 assert!(
4562 msg.contains("tls_cert_path") && msg.contains("tls_key_path"),
4563 "cert={cert:?} key={key:?}: {msg}"
4564 );
4565 }
4566 }
4567
4568 #[test]
4569 fn validate_accepts_mtls_with_tls() {
4570 let mut auth = AuthConfig::with_keys(vec![]);
4571 auth.mtls = Some(valid_mtls_config());
4572 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4573 cfg.tls_cert_path = Some("cert.pem".into());
4574 cfg.tls_key_path = Some("key.pem".into());
4575
4576 assert!(cfg.validate().is_ok(), "mTLS with both TLS paths is valid");
4577 }
4578
4579 #[test]
4580 fn validate_rejects_blank_api_key_name() {
4581 let blank = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(
4582 AuthConfig::with_keys(vec![crate::auth::ApiKeyEntry::new("", "hash", "viewer")]),
4583 );
4584 let err = blank
4585 .validate()
4586 .expect_err("blank API-key name must be rejected");
4587 assert!(err.to_string().contains("api_keys[0]"), "{err}");
4588
4589 let whitespace = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(
4590 AuthConfig::with_keys(vec![crate::auth::ApiKeyEntry::new(" ", "hash", "viewer")]),
4591 );
4592 assert!(whitespace.validate().is_err());
4593
4594 let ok = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(
4595 AuthConfig::with_keys(vec![crate::auth::ApiKeyEntry::new(
4596 "viewer-key",
4597 "hash",
4598 "viewer",
4599 )]),
4600 );
4601 assert!(ok.validate().is_ok(), "a normal name must still validate");
4602 }
4603
4604 fn reload_test_state(name: &str) -> (Arc<AuthState>, String) {
4605 let (token, hash) = crate::auth::generate_api_key().unwrap();
4606 let state = Arc::new(AuthState {
4607 api_keys: ArcSwap::from_pointee(vec![crate::auth::ApiKeyEntry::new(name, hash, "ops")]),
4608 rate_limiter: None,
4609 pre_auth_limiter: None,
4610 #[cfg(feature = "oauth")]
4611 jwks_cache: None,
4612 seen_identities: crate::auth::SeenIdentitySet::new(),
4613 counters: crate::auth::AuthCounters::default(),
4614 resource_metadata_url: None,
4615 });
4616 (state, token)
4617 }
4618
4619 #[test]
4620 fn try_reload_auth_keys_rejects_blank_name() {
4621 let (state, _token) = reload_test_state("prev-key");
4622 let handle = ReloadHandle {
4623 auth: Some(state),
4624 rbac: None,
4625 crl_set: None,
4626 };
4627 let err = handle
4628 .try_reload_auth_keys(vec![crate::auth::ApiKeyEntry::new("", "h", "ops")])
4629 .expect_err("blank API-key name must be rejected on reload");
4630 assert!(err.to_string().contains("api_keys[0]"), "{err}");
4631 }
4632
4633 #[test]
4634 fn reload_auth_keys_blank_name_leaves_previous_keys() {
4635 let (state, token) = reload_test_state("prev-key");
4636 let handle = ReloadHandle {
4637 auth: Some(Arc::clone(&state)),
4638 rbac: None,
4639 crl_set: None,
4640 };
4641 handle.reload_auth_keys(vec![crate::auth::ApiKeyEntry::new(" ", "h", "ops")]);
4642
4643 let installed = state.api_keys.load();
4644 assert!(
4645 crate::auth::verify_bearer_token(&token, &installed).is_some(),
4646 "the previous key must still authenticate after a rejected reload"
4647 );
4648 }
4649
4650 #[test]
4653 fn server_config_new_defaults() {
4654 let cfg = McpServerConfig::new("0.0.0.0:8443", "test-server", "1.0.0");
4655 assert_eq!(cfg.bind_addr, "0.0.0.0:8443");
4656 assert_eq!(cfg.name, "test-server");
4657 assert_eq!(cfg.version, "1.0.0");
4658 assert!(cfg.tls_cert_path.is_none());
4659 assert!(cfg.tls_key_path.is_none());
4660 assert!(cfg.auth.is_none());
4661 assert!(cfg.rbac.is_none());
4662 assert!(cfg.allowed_origins.is_empty());
4663 assert!(cfg.tool_rate_limit.is_none());
4664 assert!(cfg.readiness_check.is_none());
4665 assert_eq!(cfg.max_request_body, 1024 * 1024);
4666 assert_eq!(cfg.request_timeout, Duration::from_mins(2));
4667 assert_eq!(cfg.shutdown_timeout, Duration::from_secs(30));
4668 assert!(!cfg.log_request_headers);
4669 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(10));
4670 assert_eq!(cfg.max_concurrent_tls_handshakes, 256);
4671 assert!(cfg.session_store.is_none());
4672 assert!(cfg.session_binding_secret.is_none());
4673 }
4674
4675 #[derive(Default)]
4676 struct TestSessionStore;
4677
4678 #[async_trait::async_trait]
4679 impl SessionStore for TestSessionStore {
4680 async fn load(
4681 &self,
4682 _session_id: &str,
4683 ) -> Result<
4684 Option<rmcp::transport::streamable_http_server::session::SessionState>,
4685 rmcp::transport::streamable_http_server::session::SessionStoreError,
4686 > {
4687 Ok(None)
4688 }
4689
4690 async fn store(
4691 &self,
4692 _session_id: &str,
4693 _state: &rmcp::transport::streamable_http_server::session::SessionState,
4694 ) -> Result<(), rmcp::transport::streamable_http_server::session::SessionStoreError>
4695 {
4696 Ok(())
4697 }
4698
4699 async fn delete(
4700 &self,
4701 _session_id: &str,
4702 ) -> Result<(), rmcp::transport::streamable_http_server::session::SessionStoreError>
4703 {
4704 Ok(())
4705 }
4706 }
4707
4708 fn test_session_store() -> Arc<dyn SessionStore> {
4709 Arc::new(TestSessionStore)
4710 }
4711
4712 fn shared_session_binding_secret() -> SecretString {
4713 SecretString::from("0123456789abcdef0123456789abcdef")
4714 }
4715
4716 #[test]
4717 fn session_store_defaults_to_none() {
4718 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4719
4720 assert!(cfg.session_store.is_none());
4721 }
4722
4723 #[test]
4724 fn event_store_defaults_to_none() {
4725 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4726
4727 assert!(cfg.event_store.is_none());
4728 }
4729
4730 #[test]
4731 fn validate_rejects_session_store_without_binding_secret() {
4732 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4733 .with_auth(AuthConfig::with_keys(vec![]))
4734 .with_session_store(test_session_store());
4735
4736 let err = cfg
4737 .validate()
4738 .expect_err("authenticated shared-store binding needs a shared secret");
4739 let msg = err.to_string();
4740 assert!(msg.contains("session_store"), "{msg}");
4741 assert!(msg.contains("session_binding"), "{msg}");
4742 assert!(msg.contains("shared secret"), "{msg}");
4743 }
4744
4745 #[test]
4746 fn validate_allows_session_store_with_binding_secret() {
4747 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4748 .with_auth(AuthConfig::with_keys(vec![]))
4749 .with_session_store(test_session_store())
4750 .with_session_binding_secret(shared_session_binding_secret());
4751
4752 assert!(cfg.validate().is_ok());
4753 }
4754
4755 #[test]
4756 fn validate_allows_session_store_when_binding_disabled() {
4757 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4758 .with_auth(AuthConfig::with_keys(vec![]))
4759 .with_session_binding(false)
4760 .with_session_store(test_session_store());
4761
4762 assert!(cfg.validate().is_ok());
4763 }
4764
4765 #[test]
4766 fn validate_allows_binding_secret_without_session_store() {
4767 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4768 .with_auth(AuthConfig::with_keys(vec![]))
4769 .with_session_binding_secret(shared_session_binding_secret());
4770
4771 assert!(cfg.validate().is_ok());
4772 }
4773
4774 #[test]
4775 fn tls_handshake_builders_set_fields() {
4776 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4777 .with_tls_handshake_timeout(Duration::from_secs(3))
4778 .with_max_concurrent_tls_handshakes(64);
4779 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(3));
4780 assert_eq!(cfg.max_concurrent_tls_handshakes, 64);
4781 }
4782
4783 #[test]
4784 fn validate_rejects_zero_tls_handshake_timeout() {
4785 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4786 .with_tls_handshake_timeout(Duration::ZERO);
4787 let err = cfg.validate().expect_err("zero handshake timeout");
4788 assert!(err.to_string().contains("tls_handshake_timeout"));
4789 }
4790
4791 #[test]
4792 fn validate_rejects_zero_max_concurrent_tls_handshakes() {
4793 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4794 .with_max_concurrent_tls_handshakes(0);
4795 let err = cfg.validate().expect_err("zero handshake concurrency");
4796 assert!(err.to_string().contains("max_concurrent_tls_handshakes"));
4797 }
4798
4799 #[test]
4800 fn validate_consumes_and_proves() {
4801 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4803 let validated = cfg.validate().expect("valid config");
4804 assert_eq!(validated.as_inner().name, "test-server");
4806 let raw = validated.into_inner();
4808 assert_eq!(raw.name, "test-server");
4809
4810 let mut bad = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4812 bad.max_request_body = 0;
4813 assert!(bad.validate().is_err(), "zero body cap must fail validate");
4814 }
4815
4816 #[test]
4817 fn validate_rejects_zero_max_concurrent_requests() {
4818 let cfg =
4819 McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_max_concurrent_requests(0);
4820 let err = cfg.validate().expect_err("zero concurrency cap must fail");
4821 assert!(
4822 format!("{err}").contains("max_concurrent_requests"),
4823 "error should mention max_concurrent_requests, got: {err}"
4824 );
4825 }
4826
4827 #[test]
4828 fn validate_rejects_zero_max_tracked_keys() {
4829 let rl = crate::auth::RateLimitConfig {
4832 max_attempts_per_minute: 30,
4833 pre_auth_max_per_minute: None,
4834 max_tracked_keys: 0,
4835 idle_eviction: Duration::from_secs(15 * 60),
4836 burst: None,
4837 pre_auth_burst: None,
4838 key_eviction_policy: KeyEvictionPolicy::default(),
4839 };
4840 let auth_cfg = AuthConfig {
4841 enabled: true,
4842 api_keys: Vec::new(),
4843 mtls: None,
4844 rate_limit: Some(rl),
4845 #[cfg(feature = "oauth")]
4846 oauth: None,
4847 #[cfg(not(feature = "oauth"))]
4848 oauth: None,
4849 };
4850 let cfg = McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_auth(auth_cfg);
4851 let err = cfg.validate().expect_err("zero max_tracked_keys must fail");
4852 assert!(
4853 format!("{err}").contains("max_tracked_keys"),
4854 "error should mention max_tracked_keys, got: {err}"
4855 );
4856 }
4857
4858 #[test]
4859 fn derive_allowed_hosts_includes_public_host() {
4860 let hosts = derive_allowed_hosts("0.0.0.0:8080", Some("https://mcp.example.com/mcp"));
4861 assert!(
4862 hosts.iter().any(|h| h == "mcp.example.com"),
4863 "public_url host must be allowed"
4864 );
4865 }
4866
4867 #[test]
4868 fn derive_allowed_hosts_includes_bind_authority() {
4869 let hosts = derive_allowed_hosts("127.0.0.1:8080", None);
4870 assert!(
4871 hosts.iter().any(|h| h == "127.0.0.1"),
4872 "bind host must be allowed"
4873 );
4874 assert!(
4875 hosts.iter().any(|h| h == "127.0.0.1:8080"),
4876 "bind authority must be allowed"
4877 );
4878 }
4879
4880 #[tokio::test]
4883 async fn healthz_returns_ok_json() {
4884 let resp = healthz().await.into_response();
4885 assert_eq!(resp.status(), StatusCode::OK);
4886 let body = resp.into_body().collect().await.unwrap().to_bytes();
4887 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4888 assert_eq!(json["status"], "ok");
4889 assert!(
4890 json.get("name").is_none(),
4891 "healthz must not expose server name"
4892 );
4893 assert!(
4894 json.get("version").is_none(),
4895 "healthz must not expose version"
4896 );
4897 }
4898
4899 #[tokio::test]
4902 async fn readyz_returns_ok_when_ready() {
4903 let check: ReadinessCheck =
4904 Arc::new(|| Box::pin(async { serde_json::json!({"ready": true, "db": "connected"}) }));
4905 let resp = readyz(check).await.into_response();
4906 assert_eq!(resp.status(), StatusCode::OK);
4907 let body = resp.into_body().collect().await.unwrap().to_bytes();
4908 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4909 assert_eq!(json["ready"], true);
4910 assert!(
4911 json.get("name").is_none(),
4912 "readyz must not expose server name"
4913 );
4914 assert!(
4915 json.get("version").is_none(),
4916 "readyz must not expose version"
4917 );
4918 assert_eq!(json["db"], "connected");
4919 }
4920
4921 #[tokio::test]
4922 async fn readyz_returns_503_when_not_ready() {
4923 let check: ReadinessCheck =
4924 Arc::new(|| Box::pin(async { serde_json::json!({"ready": false}) }));
4925 let resp = readyz(check).await.into_response();
4926 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4927 }
4928
4929 #[tokio::test]
4930 async fn readyz_returns_503_when_ready_missing() {
4931 let check: ReadinessCheck =
4932 Arc::new(|| Box::pin(async { serde_json::json!({"status": "starting"}) }));
4933 let resp = readyz(check).await.into_response();
4934 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4936 }
4937
4938 fn peer_probe_router() -> axum::Router {
4943 async fn probe(req: Request<Body>) -> String {
4944 let ci = req
4945 .extensions()
4946 .get::<ConnectInfo<SocketAddr>>()
4947 .map(|c| c.0.to_string())
4948 .unwrap_or_default();
4949 let pa = req
4950 .extensions()
4951 .get::<PeerAddr>()
4952 .map(|p| p.addr.to_string())
4953 .unwrap_or_default();
4954 format!("{ci}|{pa}")
4955 }
4956 axum::Router::new()
4957 .route("/probe", axum::routing::get(probe))
4958 .layer(axum::middleware::from_fn(|req, next| {
4959 normalize_peer_addr_middleware(None, req, next)
4960 }))
4961 }
4962
4963 async fn body_string(resp: axum::response::Response) -> String {
4964 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
4965 String::from_utf8(bytes.to_vec()).unwrap()
4966 }
4967
4968 #[tokio::test]
4969 async fn normalize_preserves_existing_connect_info_and_mirrors_peer_addr() {
4970 let plain: SocketAddr = "10.0.0.1:1111".parse().unwrap();
4973 let tls: SocketAddr = "10.0.0.2:2222".parse().unwrap();
4974 let req = Request::builder()
4975 .uri("/probe")
4976 .extension(ConnectInfo(plain))
4977 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4978 .body(Body::empty())
4979 .unwrap();
4980 let resp = peer_probe_router().oneshot(req).await.unwrap();
4981 assert_eq!(resp.status(), StatusCode::OK);
4982 assert_eq!(body_string(resp).await, format!("{plain}|{plain}"));
4983 }
4984
4985 #[tokio::test]
4986 async fn normalize_inserts_connect_info_and_peer_addr_from_tls() {
4987 let tls: SocketAddr = "192.168.1.7:50443".parse().unwrap();
4988 let req = Request::builder()
4989 .uri("/probe")
4990 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4991 .body(Body::empty())
4992 .unwrap();
4993 let resp = peer_probe_router().oneshot(req).await.unwrap();
4994 assert_eq!(resp.status(), StatusCode::OK);
4995 assert_eq!(body_string(resp).await, format!("{tls}|{tls}"));
4996 }
4997
4998 #[tokio::test]
4999 async fn normalize_no_op_without_any_connect_info() {
5000 let req = Request::builder()
5001 .uri("/probe")
5002 .body(Body::empty())
5003 .unwrap();
5004 let resp = peer_probe_router().oneshot(req).await.unwrap();
5005 assert_eq!(resp.status(), StatusCode::OK);
5006 assert_eq!(body_string(resp).await, "|");
5007 }
5008
5009 #[tokio::test]
5010 async fn peer_addr_extractor_rejects_when_absent() {
5011 async fn h(peer: PeerAddr) -> String {
5012 peer.addr.to_string()
5013 }
5014 let app = axum::Router::new().route("/p", axum::routing::get(h));
5015 let req = Request::builder().uri("/p").body(Body::empty()).unwrap();
5016 let resp = app.oneshot(req).await.unwrap();
5017 assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
5018 }
5019
5020 #[tokio::test]
5021 async fn peer_addr_extractor_returns_value_when_present() {
5022 async fn h(peer: PeerAddr) -> String {
5023 peer.addr.to_string()
5024 }
5025 let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap();
5026 let app = axum::Router::new().route("/p", axum::routing::get(h));
5027 let req = Request::builder()
5028 .uri("/p")
5029 .extension(PeerAddr::new(addr))
5030 .body(Body::empty())
5031 .unwrap();
5032 let resp = app.oneshot(req).await.unwrap();
5033 assert_eq!(resp.status(), StatusCode::OK);
5034 assert_eq!(body_string(resp).await, addr.to_string());
5035 }
5036
5037 #[tokio::test]
5038 async fn peer_addr_via_extension_extractor() {
5039 async fn h(axum::Extension(peer): axum::Extension<PeerAddr>) -> String {
5040 peer.addr.to_string()
5041 }
5042 let addr: SocketAddr = "127.0.0.1:4242".parse().unwrap();
5043 let app = axum::Router::new().route("/p", axum::routing::get(h));
5044 let req = Request::builder()
5045 .uri("/p")
5046 .extension(PeerAddr::new(addr))
5047 .body(Body::empty())
5048 .unwrap();
5049 let resp = app.oneshot(req).await.unwrap();
5050 assert_eq!(resp.status(), StatusCode::OK);
5051 assert_eq!(body_string(resp).await, addr.to_string());
5052 }
5053
5054 fn limited_router(per_minute: u32) -> axum::Router {
5059 limited_router_with_burst(per_minute, None)
5060 }
5061
5062 fn limited_router_with_burst(per_minute: u32, burst: Option<u32>) -> axum::Router {
5064 limited_router_full(per_minute, burst, &[])
5065 }
5066
5067 fn limited_router_full(
5071 per_minute: u32,
5072 burst: Option<u32>,
5073 exempt_paths: &[&str],
5074 ) -> axum::Router {
5075 let limiter = build_extra_route_rate_limiter_with_policy(
5076 per_minute,
5077 burst,
5078 KeyEvictionPolicy::default(),
5079 NonZeroUsize::new(EXTRA_ROUTE_MAX_TRACKED_KEYS).unwrap_or(NonZeroUsize::MIN),
5080 );
5081 let exempt: Arc<std::collections::HashSet<String>> =
5082 Arc::new(exempt_paths.iter().map(|s| (*s).to_owned()).collect());
5083 axum::Router::new()
5084 .route("/limited", axum::routing::get(|| async { "ok" }))
5085 .route("/exempt", axum::routing::get(|| async { "ok" }))
5086 .layer(axum::middleware::from_fn(move |req, next| {
5087 let l = Arc::clone(&limiter);
5088 let e = Arc::clone(&exempt);
5089 extra_route_rate_limit_middleware(l, e, req, next)
5090 }))
5091 }
5092
5093 fn limited_req(ip: &str) -> Request<Body> {
5094 limited_req_to(ip, "/limited")
5095 }
5096
5097 fn limited_req_to(ip: &str, path: &str) -> Request<Body> {
5098 let addr: SocketAddr = format!("{ip}:40000").parse().unwrap();
5099 Request::builder()
5100 .uri(path)
5101 .extension(ConnectInfo(addr))
5102 .body(Body::empty())
5103 .unwrap()
5104 }
5105
5106 #[tokio::test]
5107 async fn extra_route_limiter_denies_over_quota() {
5108 let app = limited_router(2);
5109 for i in 0..2 {
5110 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
5111 assert_eq!(resp.status(), StatusCode::OK, "request {i} should pass");
5112 }
5113 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
5114 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
5115 let body = body_string(resp).await;
5116 assert!(
5117 body.contains("too many requests to application routes"),
5118 "deny body should match the limiter message, got: {body}"
5119 );
5120 }
5121
5122 fn one_tracked_key() -> NonZeroUsize {
5123 NonZeroUsize::new(1).unwrap_or(NonZeroUsize::MIN)
5124 }
5125
5126 #[tokio::test]
5127 async fn extra_route_limiter_capacity_full_returns_503_without_retry_after() {
5128 let limiter = build_extra_route_rate_limiter_with_policy(
5129 10,
5130 None,
5131 KeyEvictionPolicy::RejectNew,
5132 one_tracked_key(),
5133 );
5134 let exempt = Arc::new(std::collections::HashSet::new());
5135 let app = axum::Router::new()
5136 .route("/limited", axum::routing::get(|| async { "ok" }))
5137 .layer(axum::middleware::from_fn(move |req, next| {
5138 let l = Arc::clone(&limiter);
5139 let e = Arc::clone(&exempt);
5140 extra_route_rate_limit_middleware(l, e, req, next)
5141 }));
5142 let established = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
5143 assert_eq!(established.status(), StatusCode::OK);
5144
5145 let denied = app.clone().oneshot(limited_req("10.1.1.2")).await.unwrap();
5146
5147 assert_eq!(denied.status(), StatusCode::SERVICE_UNAVAILABLE);
5148 assert!(denied.headers().get(header::RETRY_AFTER).is_none());
5149 }
5150
5151 #[tokio::test]
5152 async fn extra_route_limiter_isolates_keys() {
5153 let app = limited_router(2);
5154 for _ in 0..2 {
5155 let resp = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
5156 assert_eq!(resp.status(), StatusCode::OK);
5157 }
5158 let exhausted = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
5159 assert_eq!(exhausted.status(), StatusCode::TOO_MANY_REQUESTS);
5160 let other = app.clone().oneshot(limited_req("10.3.3.3")).await.unwrap();
5162 assert_eq!(other.status(), StatusCode::OK);
5163 }
5164
5165 #[tokio::test]
5166 async fn extra_route_limiter_bounds_requests_without_peer() {
5167 let app = limited_router(1);
5171 let mk = || {
5172 Request::builder()
5173 .uri("/limited")
5174 .body(Body::empty())
5175 .unwrap()
5176 };
5177 let first = app.clone().oneshot(mk()).await.unwrap();
5178 assert_eq!(
5179 first.status(),
5180 StatusCode::OK,
5181 "first request consumes quota"
5182 );
5183 let second = app.clone().oneshot(mk()).await.unwrap();
5184 assert_eq!(
5185 second.status(),
5186 StatusCode::TOO_MANY_REQUESTS,
5187 "unattributable requests must share a bounded bucket, not bypass the limiter"
5188 );
5189 }
5190
5191 #[test]
5192 fn limiter_client_key_falls_back_to_unattributed() {
5193 let empty = axum::http::Extensions::new();
5194 assert_eq!(limiter_client_key(&empty), RateLimitKey::Unattributed);
5195 }
5196
5197 #[test]
5198 fn unattributed_key_is_distinct_from_unspecified_ip() {
5199 let unspecified = RateLimitKey::Ip("0.0.0.0".parse::<IpAddr>().unwrap());
5203 assert_ne!(unspecified, RateLimitKey::Unattributed);
5204
5205 let mut set = std::collections::HashSet::new();
5206 set.insert(unspecified);
5207 set.insert(RateLimitKey::Unattributed);
5208 assert_eq!(set.len(), 2, "the two keys must hash to distinct buckets");
5209 }
5210
5211 #[test]
5212 fn rate_limit_key_display_does_not_fabricate_an_ip() {
5213 assert_eq!(
5214 RateLimitKey::Ip("10.1.2.3".parse::<IpAddr>().unwrap()).to_string(),
5215 "10.1.2.3"
5216 );
5217 assert_eq!(RateLimitKey::Unattributed.to_string(), "unattributed");
5218 }
5219
5220 #[tokio::test]
5221 async fn extra_route_limiter_extracts_tls_conn_info() {
5222 let app = limited_router(2);
5223 let mk = || {
5224 let addr: SocketAddr = "192.168.9.9:55555".parse().unwrap();
5225 Request::builder()
5226 .uri("/limited")
5227 .extension(ConnectInfo(TlsConnInfo::new(addr, None)))
5228 .body(Body::empty())
5229 .unwrap()
5230 };
5231 for _ in 0..2 {
5232 assert_eq!(
5233 app.clone().oneshot(mk()).await.unwrap().status(),
5234 StatusCode::OK
5235 );
5236 }
5237 let resp = app.clone().oneshot(mk()).await.unwrap();
5238 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
5239 }
5240
5241 #[tokio::test]
5242 async fn extra_route_limiter_exempt_path_bypasses_quota() {
5243 let app = limited_router_full(1, None, &["/exempt"]);
5246 for i in 0..5 {
5247 let resp = app
5248 .clone()
5249 .oneshot(limited_req_to("10.6.6.6", "/exempt"))
5250 .await
5251 .unwrap();
5252 assert_eq!(resp.status(), StatusCode::OK, "exempt request {i}");
5253 }
5254 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
5256 assert_eq!(resp.status(), StatusCode::OK);
5257 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
5259 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
5260 }
5261
5262 #[tokio::test]
5263 async fn extra_route_limiter_exemption_is_raw_exact_match() {
5264 let app = limited_router_full(1, None, &["/exempt"]);
5267 let ok = app
5268 .clone()
5269 .oneshot(limited_req_to("10.7.7.7", "/exempt/"))
5270 .await
5271 .unwrap();
5272 assert_eq!(
5273 ok.status(),
5274 StatusCode::NOT_FOUND,
5275 "variant path routes 404"
5276 );
5277 let denied = app
5279 .clone()
5280 .oneshot(limited_req_to("10.7.7.7", "/limited"))
5281 .await
5282 .unwrap();
5283 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
5284 }
5285
5286 #[cfg(feature = "metrics")]
5287 #[tokio::test]
5288 async fn extra_route_limiter_deny_increments_counter_exempt_does_not() {
5289 let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
5290 let app = limited_router_full(1, None, &["/exempt"]);
5291 let mk = |path: &str| {
5292 let addr: SocketAddr = "10.8.8.8:40000".parse().unwrap();
5293 Request::builder()
5294 .uri(path)
5295 .extension(ConnectInfo(addr))
5296 .extension(Arc::clone(&metrics))
5297 .body(Body::empty())
5298 .unwrap()
5299 };
5300 let counter = || {
5301 metrics
5302 .rate_limited_total
5303 .with_label_values(&["extra_route"])
5304 .get()
5305 };
5306 for _ in 0..3 {
5308 assert_eq!(
5309 app.clone().oneshot(mk("/exempt")).await.unwrap().status(),
5310 StatusCode::OK
5311 );
5312 }
5313 assert_eq!(counter(), 0, "exempt requests must not count as denies");
5314 assert_eq!(
5316 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
5317 StatusCode::OK
5318 );
5319 assert_eq!(counter(), 0);
5320 assert_eq!(
5321 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
5322 StatusCode::TOO_MANY_REQUESTS
5323 );
5324 assert_eq!(counter(), 1, "deny must increment the extra_route label");
5325 }
5326
5327 #[test]
5328 fn validate_rejects_exempt_paths_without_base_knob() {
5329 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5330 .with_extra_route_rate_limit_exempt_paths(["/ok"]);
5331 let err = cfg.validate().expect_err("exempt paths without rate limit");
5332 assert!(err.to_string().contains("requires extra_route_rate_limit"));
5333 }
5334
5335 #[test]
5336 fn validate_rejects_malformed_exempt_paths() {
5337 for bad in ["", "no-slash"] {
5338 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5339 .with_extra_route_rate_limit(10)
5340 .with_extra_route_rate_limit_exempt_paths([bad]);
5341 let err = cfg.validate().expect_err("malformed exempt path");
5342 assert!(
5343 err.to_string()
5344 .contains("must be non-empty and start with '/'"),
5345 "entry {bad:?}: {err}"
5346 );
5347 }
5348 }
5349
5350 #[test]
5351 fn validate_accepts_wellformed_exempt_paths() {
5352 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5353 .with_extra_route_rate_limit(10)
5354 .with_extra_route_rate_limit_exempt_paths(["/.well-known/oauth-authorization-server"]);
5355 assert!(cfg.validate().is_ok());
5356 }
5357
5358 #[test]
5359 fn validate_rejects_zero_extra_route_rate_limit() {
5360 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5361 .with_extra_route_rate_limit(0);
5362 let err = cfg.validate().expect_err("zero extra route rate limit");
5363 assert!(err.to_string().contains("extra_route_rate_limit"));
5364 }
5365
5366 #[tokio::test]
5367 async fn extra_route_limiter_burst_allows_initial_spike() {
5368 let app = limited_router_with_burst(1, Some(3));
5369 for i in 0..3 {
5370 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
5371 assert_eq!(resp.status(), StatusCode::OK, "burst request {i}");
5372 }
5373 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
5374 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
5375 }
5376
5377 #[tokio::test]
5378 async fn extra_route_limiter_deny_sets_retry_after() {
5379 let app = limited_router(1);
5380 let ok = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
5381 assert_eq!(ok.status(), StatusCode::OK);
5382 let denied = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
5383 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
5384 let retry_after = denied
5385 .headers()
5386 .get(header::RETRY_AFTER)
5387 .expect("Retry-After present")
5388 .to_str()
5389 .unwrap()
5390 .parse::<u64>()
5391 .unwrap();
5392 assert!(retry_after >= 1, "delta-seconds must be >= 1");
5393 }
5394
5395 #[test]
5396 fn validate_rejects_zero_burst_knobs() {
5397 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5398 .with_tool_rate_limit(10)
5399 .with_tool_rate_limit_burst(0)
5400 .validate()
5401 .expect_err("zero tool burst");
5402 assert!(err.to_string().contains("tool_rate_limit_burst"));
5403
5404 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5405 .with_extra_route_rate_limit(10)
5406 .with_extra_route_rate_limit_burst(0)
5407 .validate()
5408 .expect_err("zero extra route burst");
5409 assert!(err.to_string().contains("extra_route_rate_limit_burst"));
5410 }
5411
5412 #[test]
5413 fn validate_rejects_orphan_burst_knobs() {
5414 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5415 .with_tool_rate_limit_burst(5)
5416 .validate()
5417 .expect_err("orphan tool burst");
5418 assert!(err.to_string().contains("requires tool_rate_limit"));
5419
5420 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5421 .with_extra_route_rate_limit_burst(5)
5422 .validate()
5423 .expect_err("orphan extra route burst");
5424 assert!(err.to_string().contains("requires extra_route_rate_limit"));
5425 }
5426
5427 #[test]
5428 fn validate_rejects_zero_auth_bursts() {
5429 let auth = AuthConfig::with_keys(vec![])
5430 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_burst(0));
5431 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5432 .with_auth(auth)
5433 .validate()
5434 .expect_err("zero auth burst");
5435 assert!(err.to_string().contains("rate_limit.burst"));
5436
5437 let auth = AuthConfig::with_keys(vec![])
5438 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(0));
5439 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5440 .with_auth(auth)
5441 .validate()
5442 .expect_err("zero pre-auth burst");
5443 assert!(err.to_string().contains("pre_auth_burst"));
5444 }
5445
5446 #[test]
5447 fn validate_rejects_zero_pre_auth_max_per_minute() {
5448 let auth = AuthConfig::with_keys(vec![])
5449 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_max_per_minute(0));
5450 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5451 .with_auth(auth)
5452 .validate()
5453 .expect_err("zero pre-auth rate");
5454 assert!(err.to_string().contains("pre_auth_max_per_minute"));
5455 }
5456
5457 fn valid_mtls_config() -> MtlsConfig {
5458 MtlsConfig {
5459 ca_cert_path: "memory://ca.pem".into(),
5460 required: true,
5461 default_role: "viewer".into(),
5462 crl_enabled: true,
5463 crl_refresh_interval: None,
5464 crl_fetch_timeout: Duration::from_secs(30),
5465 crl_stale_grace: Duration::from_secs(24 * 60 * 60),
5466 crl_deny_on_unavailable: false,
5467 crl_end_entity_only: false,
5468 crl_allow_http: true,
5469 crl_enforce_expiration: true,
5470 crl_max_concurrent_fetches: 4,
5471 crl_max_response_bytes: 5 * 1024 * 1024,
5472 crl_discovery_rate_per_min: 60,
5473 crl_max_host_semaphores: 1024,
5474 crl_max_seen_urls: 4096,
5475 crl_max_cache_entries: 1024,
5476 }
5477 }
5478
5479 #[test]
5480 fn validate_rejects_zero_crl_max_response_bytes() {
5481 let mut mtls = valid_mtls_config();
5482 mtls.crl_max_response_bytes = 0;
5483 let mut auth = AuthConfig::with_keys(vec![]);
5484 auth.mtls = Some(mtls);
5485
5486 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
5489 cfg.tls_cert_path = Some("cert.pem".into());
5490 cfg.tls_key_path = Some("key.pem".into());
5491
5492 let err = cfg.validate().expect_err("zero CRL response cap");
5493 assert!(err.to_string().contains("crl_max_response_bytes"));
5494 }
5495
5496 #[test]
5499 fn validate_accepts_pre_auth_burst_without_explicit_pre_auth_rate() {
5500 let auth = AuthConfig::with_keys(vec![])
5501 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(50));
5502 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
5503 assert!(cfg.validate().is_ok(), "pre_auth_burst has no orphan rule");
5504 }
5505
5506 #[test]
5509 fn trusted_forwarder_max_entries_bounds_are_enforced() {
5510 let cfg = |n: usize| {
5511 McpServerConfig::new("127.0.0.1:8080", "t", "0")
5512 .with_trusted_forwarder_max_entries(n)
5513 .validate()
5514 };
5515 assert!(cfg(0).is_err(), "0 would pin every client to the proxy");
5516 assert!(
5517 cfg(crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES + 1).is_err(),
5518 "above the ceiling would re-open the header-bomb vector"
5519 );
5520 assert!(cfg(1).is_ok());
5521 assert!(cfg(crate::forwarded::MAX_SCANNED_ENTRIES).is_ok());
5522 assert!(cfg(crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES).is_ok());
5523 }
5524
5525 #[test]
5526 fn trusted_forwarder_max_entries_defaults_to_the_module_constant() {
5527 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "0");
5528 assert_eq!(
5529 cfg.trusted_forwarder_max_entries,
5530 crate::forwarded::MAX_SCANNED_ENTRIES
5531 );
5532 }
5533
5534 fn forward_resolver(trusted: &[&str], mode: ForwardedHeaderMode) -> Arc<ForwardResolver> {
5535 Arc::new(ForwardResolver {
5536 trusted: trusted.iter().map(|s| s.parse().unwrap()).collect(),
5537 mode,
5538 max_scanned_entries: crate::forwarded::MAX_SCANNED_ENTRIES,
5539 })
5540 }
5541
5542 fn forwarded_probe_router(resolver: Option<Arc<ForwardResolver>>) -> axum::Router {
5544 async fn probe(req: Request<Body>) -> String {
5545 let pa = req
5546 .extensions()
5547 .get::<PeerAddr>()
5548 .map(|p| p.addr.ip().to_string())
5549 .unwrap_or_default();
5550 let ci = req
5551 .extensions()
5552 .get::<ClientIp>()
5553 .map(|c| c.ip.to_string())
5554 .unwrap_or_default();
5555 format!("{pa}|{ci}")
5556 }
5557 axum::Router::new()
5558 .route("/probe", axum::routing::get(probe))
5559 .layer(axum::middleware::from_fn(move |req, next| {
5560 let r = resolver.clone();
5561 normalize_peer_addr_middleware(r, req, next)
5562 }))
5563 }
5564
5565 fn probe_req(peer: &str, header: Option<(&str, &str)>) -> Request<Body> {
5566 let addr: SocketAddr = peer.parse().unwrap();
5567 let mut builder = Request::builder()
5568 .uri("/probe")
5569 .extension(ConnectInfo(addr));
5570 if let Some((name, value)) = header {
5571 builder = builder.header(name, value);
5572 }
5573 builder.body(Body::empty()).unwrap()
5574 }
5575
5576 #[tokio::test]
5577 async fn client_ip_equals_direct_without_resolver() {
5578 let app = forwarded_probe_router(None);
5579 let resp = app
5580 .oneshot(probe_req(
5581 "10.1.2.3:4444",
5582 Some(("x-forwarded-for", "203.0.113.7")),
5583 ))
5584 .await
5585 .unwrap();
5586 assert_eq!(
5587 body_string(resp).await,
5588 "10.1.2.3|10.1.2.3",
5589 "feature off: header ignored, ClientIp == direct"
5590 );
5591 }
5592
5593 #[tokio::test]
5594 async fn client_ip_resolved_for_trusted_peer() {
5595 let app = forwarded_probe_router(Some(forward_resolver(
5596 &["10.0.0.0/8"],
5597 ForwardedHeaderMode::XForwardedFor,
5598 )));
5599 let resp = app
5600 .oneshot(probe_req(
5601 "10.0.0.1:9999",
5602 Some(("x-forwarded-for", "203.0.113.7")),
5603 ))
5604 .await
5605 .unwrap();
5606 assert_eq!(
5607 body_string(resp).await,
5608 "10.0.0.1|203.0.113.7",
5609 "PeerAddr stays direct while ClientIp resolves"
5610 );
5611 }
5612
5613 #[tokio::test]
5614 async fn client_ip_falls_back_to_direct_on_malformed_header() {
5615 let app = forwarded_probe_router(Some(forward_resolver(
5616 &["10.0.0.0/8"],
5617 ForwardedHeaderMode::XForwardedFor,
5618 )));
5619 let resp = app
5620 .oneshot(probe_req(
5621 "10.0.0.1:9999",
5622 Some(("x-forwarded-for", "not-an-ip")),
5623 ))
5624 .await
5625 .unwrap();
5626 assert_eq!(
5627 body_string(resp).await,
5628 "10.0.0.1|10.0.0.1",
5629 "malformed chain falls back to the direct peer"
5630 );
5631 }
5632
5633 #[test]
5634 fn forwarded_header_mode_deserializes_kebab_case() {
5635 #[derive(serde::Deserialize)]
5636 struct Wrapper {
5637 mode: ForwardedHeaderMode,
5638 }
5639 let w: Wrapper = toml::from_str(r#"mode = "x-forwarded-for""#).unwrap();
5640 assert_eq!(w.mode, ForwardedHeaderMode::XForwardedFor);
5641 let w: Wrapper = toml::from_str(r#"mode = "forwarded""#).unwrap();
5642 assert_eq!(w.mode, ForwardedHeaderMode::Forwarded);
5643 assert!(
5644 toml::from_str::<Wrapper>(r#"mode = "XForwardedFor""#).is_err(),
5645 "PascalCase wire value must be rejected"
5646 );
5647 }
5648
5649 #[test]
5650 fn validate_rejects_bad_trusted_proxy_entry() {
5651 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5652 .with_trusted_proxies(["not-a-cidr"]);
5653 let err = cfg.validate().expect_err("bad CIDR");
5654 assert!(err.to_string().contains("trusted_proxies"));
5655 }
5656
5657 #[test]
5658 fn validate_rejects_zero_prefix_trusted_proxy() {
5659 for entry in ["0.0.0.0/0", "::/0"] {
5660 let cfg =
5661 McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([entry]);
5662 let err = cfg.validate().expect_err("zero-prefix CIDR");
5663 assert!(
5664 err.to_string().contains("prefix length 0"),
5665 "entry {entry}: {err}"
5666 );
5667 }
5668 }
5669
5670 #[test]
5671 fn validate_accepts_cidr_and_bare_ip_proxy_entries() {
5672 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([
5673 "10.0.0.0/8",
5674 "192.0.2.1",
5675 "2001:db8::1",
5676 ]);
5677 assert!(cfg.validate().is_ok(), "CIDRs and bare IPs are accepted");
5678 }
5679
5680 #[test]
5681 fn validate_rejects_forwarded_header_without_proxies() {
5682 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5683 .with_forwarded_header(ForwardedHeaderMode::Forwarded);
5684 let err = cfg.validate().expect_err("mode without proxies");
5685 assert!(err.to_string().contains("requires trusted_proxies"));
5686 }
5687
5688 fn origin_router(origins: Vec<String>, log_request_headers: bool) -> axum::Router {
5692 let allowed: Arc<[String]> = Arc::from(origins);
5693 axum::Router::new()
5694 .route("/test", axum::routing::get(|| async { "ok" }))
5695 .layer(axum::middleware::from_fn(move |req, next| {
5696 let a = Arc::clone(&allowed);
5697 origin_check_middleware(a, log_request_headers, req, next)
5698 }))
5699 }
5700
5701 #[tokio::test]
5702 async fn origin_allowed_passes() {
5703 let app = origin_router(vec!["http://localhost:3000".into()], false);
5704 let req = Request::builder()
5705 .uri("/test")
5706 .header(header::ORIGIN, "http://localhost:3000")
5707 .body(Body::empty())
5708 .unwrap();
5709 let resp = app.oneshot(req).await.unwrap();
5710 assert_eq!(resp.status(), StatusCode::OK);
5711 }
5712
5713 #[tokio::test]
5714 async fn origin_rejected_returns_403() {
5715 let app = origin_router(vec!["http://localhost:3000".into()], false);
5716 let req = Request::builder()
5717 .uri("/test")
5718 .header(header::ORIGIN, "http://evil.com")
5719 .body(Body::empty())
5720 .unwrap();
5721 let resp = app.oneshot(req).await.unwrap();
5722 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5723 }
5724
5725 #[tokio::test]
5726 async fn no_origin_header_passes() {
5727 let app = origin_router(vec!["http://localhost:3000".into()], false);
5728 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5729 let resp = app.oneshot(req).await.unwrap();
5730 assert_eq!(resp.status(), StatusCode::OK);
5731 }
5732
5733 #[tokio::test]
5734 async fn empty_allowlist_rejects_any_origin() {
5735 let app = origin_router(vec![], false);
5736 let req = Request::builder()
5737 .uri("/test")
5738 .header(header::ORIGIN, "http://anything.com")
5739 .body(Body::empty())
5740 .unwrap();
5741 let resp = app.oneshot(req).await.unwrap();
5742 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5743 }
5744
5745 #[tokio::test]
5746 async fn empty_allowlist_passes_without_origin() {
5747 let app = origin_router(vec![], false);
5748 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5749 let resp = app.oneshot(req).await.unwrap();
5750 assert_eq!(resp.status(), StatusCode::OK);
5751 }
5752
5753 #[test]
5754 fn format_request_headers_redacts_sensitive_values() {
5755 let mut headers = axum::http::HeaderMap::new();
5756 headers.insert("authorization", "Bearer secret-token".parse().unwrap());
5757 headers.insert("cookie", "sid=abc".parse().unwrap());
5758 headers.insert("x-request-id", "req-123".parse().unwrap());
5759
5760 let out = format_request_headers_for_log(&headers);
5761 assert!(out.contains("authorization: [REDACTED]"));
5762 assert!(out.contains("cookie: [REDACTED]"));
5763 assert!(out.contains("x-request-id: req-123"));
5764 assert!(!out.contains("secret-token"));
5765 }
5766
5767 #[test]
5768 fn format_request_headers_redacts_forwarding_headers() {
5769 let mut headers = axum::http::HeaderMap::new();
5770 headers.insert("forwarded", "for=203.0.113.9;by=10.1.2.3".parse().unwrap());
5771 headers.insert("x-forwarded-for", "203.0.113.9, 10.1.2.3".parse().unwrap());
5772 headers.insert("x-real-ip", "203.0.113.9".parse().unwrap());
5773 headers.insert("x-request-id", "req-123".parse().unwrap());
5774
5775 let out = format_request_headers_for_log(&headers);
5776 for name in ["forwarded", "x-forwarded-for", "x-real-ip"] {
5777 assert!(
5778 out.contains(&format!("{name}: [REDACTED]")),
5779 "{name} carries client IP / proxy topology and must not reach logs; got {out}"
5780 );
5781 }
5782 assert!(
5783 !out.contains("203.0.113.9") && !out.contains("10.1.2.3"),
5784 "no forwarded address may survive redaction; got {out}"
5785 );
5786 assert!(out.contains("x-request-id: req-123"));
5787 }
5788
5789 fn security_router(is_tls: bool) -> axum::Router {
5792 security_router_with(is_tls, SecurityHeadersConfig::default())
5793 }
5794
5795 fn security_router_with(is_tls: bool, cfg: SecurityHeadersConfig) -> axum::Router {
5796 let cfg = Arc::new(cfg);
5797 axum::Router::new()
5798 .route("/test", axum::routing::get(|| async { "ok" }))
5799 .layer(axum::middleware::from_fn(move |req, next| {
5800 let c = Arc::clone(&cfg);
5801 security_headers_middleware(is_tls, c, req, next)
5802 }))
5803 }
5804
5805 #[tokio::test]
5806 async fn security_headers_set_on_response() {
5807 let app = security_router(false);
5808 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5809 let resp = app.oneshot(req).await.unwrap();
5810 assert_eq!(resp.status(), StatusCode::OK);
5811
5812 let h = resp.headers();
5813 assert_eq!(h.get("x-content-type-options").unwrap(), "nosniff");
5814 assert_eq!(h.get("x-frame-options").unwrap(), "deny");
5815 assert_eq!(h.get("cache-control").unwrap(), "no-store, max-age=0");
5816 assert_eq!(h.get("referrer-policy").unwrap(), "no-referrer");
5817 assert_eq!(h.get("cross-origin-opener-policy").unwrap(), "same-origin");
5818 assert_eq!(
5819 h.get("cross-origin-resource-policy").unwrap(),
5820 "same-origin"
5821 );
5822 assert_eq!(
5823 h.get("cross-origin-embedder-policy").unwrap(),
5824 "require-corp"
5825 );
5826 assert_eq!(h.get("x-permitted-cross-domain-policies").unwrap(), "none");
5827 assert!(
5828 h.get("permissions-policy")
5829 .unwrap()
5830 .to_str()
5831 .unwrap()
5832 .contains("camera=()"),
5833 "permissions-policy must restrict browser features"
5834 );
5835 assert_eq!(
5836 h.get("content-security-policy").unwrap(),
5837 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
5838 );
5839 assert_eq!(h.get("x-dns-prefetch-control").unwrap(), "off");
5840 assert!(h.get("strict-transport-security").is_none());
5842 }
5843
5844 #[tokio::test]
5845 async fn hsts_set_when_tls_enabled() {
5846 let app = security_router(true);
5847 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5848 let resp = app.oneshot(req).await.unwrap();
5849
5850 let hsts = resp.headers().get("strict-transport-security").unwrap();
5851 assert!(
5852 hsts.to_str().unwrap().contains("max-age=63072000"),
5853 "HSTS must set 2-year max-age"
5854 );
5855 }
5856
5857 #[tokio::test]
5858 async fn default_csp_matches_guideline() {
5859 let app = security_router(false);
5860 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5861 let resp = app.oneshot(req).await.unwrap();
5862 assert_eq!(
5863 resp.headers().get("content-security-policy").unwrap(),
5864 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
5865 );
5866 }
5867
5868 #[tokio::test]
5869 async fn operator_csp_override_still_wins() {
5870 let cfg = SecurityHeadersConfig {
5871 content_security_policy: Some("default-src 'self'".into()),
5872 ..SecurityHeadersConfig::default()
5873 };
5874 let app = security_router_with(false, cfg);
5875 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5876 let resp = app.oneshot(req).await.unwrap();
5877 assert_eq!(
5878 resp.headers().get("content-security-policy").unwrap(),
5879 "default-src 'self'"
5880 );
5881 }
5882
5883 fn check_with_security_headers(
5889 headers: SecurityHeadersConfig,
5890 ) -> Result<(), RmcpServerKitError> {
5891 let cfg =
5892 McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_security_headers(headers);
5893 cfg.check()
5894 }
5895
5896 #[test]
5897 fn security_headers_config_default_validates() {
5898 check_with_security_headers(SecurityHeadersConfig::default())
5899 .expect("default SecurityHeadersConfig must validate");
5900 }
5901
5902 #[test]
5903 fn security_headers_config_validate_accepts_empty_string() {
5904 let h = SecurityHeadersConfig {
5906 x_content_type_options: Some(String::new()),
5907 x_frame_options: Some(String::new()),
5908 cache_control: Some(String::new()),
5909 referrer_policy: Some(String::new()),
5910 cross_origin_opener_policy: Some(String::new()),
5911 cross_origin_resource_policy: Some(String::new()),
5912 cross_origin_embedder_policy: Some(String::new()),
5913 permissions_policy: Some(String::new()),
5914 x_permitted_cross_domain_policies: Some(String::new()),
5915 content_security_policy: Some(String::new()),
5916 x_dns_prefetch_control: Some(String::new()),
5917 strict_transport_security: Some(String::new()),
5918 };
5919 check_with_security_headers(h).expect("Some(\"\") on every field must validate (omit-all)");
5920 }
5921
5922 #[test]
5923 fn security_headers_config_validate_rejects_bad_value() {
5924 let h = SecurityHeadersConfig {
5926 referrer_policy: Some("\u{0007}".into()),
5927 ..SecurityHeadersConfig::default()
5928 };
5929 let err = check_with_security_headers(h)
5930 .expect_err("control char in referrer_policy must reject");
5931 let msg = err.to_string();
5932 assert!(
5933 msg.contains("referrer_policy"),
5934 "error must name the offending field, got: {msg}"
5935 );
5936 }
5937
5938 #[test]
5939 fn security_headers_config_validate_rejects_hsts_preload() {
5940 let h = SecurityHeadersConfig {
5941 strict_transport_security: Some("max-age=63072000; includeSubDomains; preload".into()),
5942 ..SecurityHeadersConfig::default()
5943 };
5944 let err = check_with_security_headers(h).expect_err("HSTS with preload must reject");
5945 let msg = err.to_string();
5946 assert!(
5947 msg.contains("strict_transport_security"),
5948 "error must name the field, got: {msg}"
5949 );
5950 assert!(
5951 msg.to_lowercase().contains("preload"),
5952 "error must mention `preload`, got: {msg}"
5953 );
5954 }
5955
5956 #[test]
5957 fn security_headers_config_validate_rejects_hsts_preload_uppercase() {
5958 let h = SecurityHeadersConfig {
5960 strict_transport_security: Some("max-age=600; PRELOAD".into()),
5961 ..SecurityHeadersConfig::default()
5962 };
5963 check_with_security_headers(h).expect_err("HSTS preload check must be case-insensitive");
5964 }
5965
5966 #[tokio::test]
5967 async fn security_headers_override_honored() {
5968 let h = SecurityHeadersConfig {
5970 x_frame_options: Some("SAMEORIGIN".into()),
5971 ..SecurityHeadersConfig::default()
5972 };
5973 let app = security_router_with(false, h);
5974 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5975 let resp = app.oneshot(req).await.unwrap();
5976 assert_eq!(resp.status(), StatusCode::OK);
5977
5978 let xfo = resp.headers().get("x-frame-options").unwrap();
5979 assert_eq!(xfo, "SAMEORIGIN");
5980 }
5981
5982 #[tokio::test]
5983 async fn security_headers_empty_string_omits() {
5984 let h = SecurityHeadersConfig {
5986 referrer_policy: Some(String::new()),
5987 ..SecurityHeadersConfig::default()
5988 };
5989 let app = security_router_with(false, h);
5990 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5991 let resp = app.oneshot(req).await.unwrap();
5992 assert_eq!(resp.status(), StatusCode::OK);
5993
5994 assert!(
5995 resp.headers().get("referrer-policy").is_none(),
5996 "Some(\"\") must omit the header"
5997 );
5998 assert_eq!(
6000 resp.headers().get("x-content-type-options").unwrap(),
6001 "nosniff"
6002 );
6003 }
6004
6005 #[tokio::test]
6006 async fn security_headers_hsts_only_when_tls() {
6007 let h = SecurityHeadersConfig {
6009 strict_transport_security: Some("max-age=600".into()),
6010 ..SecurityHeadersConfig::default()
6011 };
6012 let app = security_router_with(false, h);
6013 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
6014 let resp = app.oneshot(req).await.unwrap();
6015 assert!(
6016 resp.headers().get("strict-transport-security").is_none(),
6017 "HSTS must remain absent on plaintext deployments even with override"
6018 );
6019 }
6020
6021 #[cfg(feature = "oauth")]
6024 #[tokio::test]
6025 async fn oauth_token_cache_headers_set_pragma_and_vary() {
6026 let app = axum::Router::new()
6027 .route("/token", axum::routing::post(|| async { "{}" }))
6028 .layer(axum::middleware::from_fn(
6029 oauth_token_cache_headers_middleware,
6030 ));
6031 let req = Request::builder()
6032 .method("POST")
6033 .uri("/token")
6034 .body(Body::from("{}"))
6035 .unwrap();
6036 let resp = app.oneshot(req).await.unwrap();
6037 assert_eq!(resp.status(), StatusCode::OK);
6038
6039 let h = resp.headers();
6040 assert_eq!(
6041 h.get("pragma").unwrap(),
6042 "no-cache",
6043 "RFC 6749 §5.1: token responses must set Pragma: no-cache"
6044 );
6045 let vary_values: Vec<String> = h
6046 .get_all("vary")
6047 .iter()
6048 .filter_map(|v| v.to_str().ok().map(str::to_owned))
6049 .collect();
6050 assert!(
6051 vary_values
6052 .iter()
6053 .any(|v| v.eq_ignore_ascii_case("Authorization")),
6054 "RFC 6750 §5.4: Vary must include Authorization, got {vary_values:?}"
6055 );
6056 }
6057
6058 #[cfg(feature = "oauth")]
6059 #[tokio::test]
6060 async fn oauth_token_cache_headers_preserve_existing_vary() {
6061 let app = axum::Router::new()
6064 .route(
6065 "/token",
6066 axum::routing::post(|| async {
6067 axum::response::Response::builder()
6068 .header("vary", "Accept-Encoding")
6069 .body(Body::from("{}"))
6070 .unwrap()
6071 }),
6072 )
6073 .layer(axum::middleware::from_fn(
6074 oauth_token_cache_headers_middleware,
6075 ));
6076 let req = Request::builder()
6077 .method("POST")
6078 .uri("/token")
6079 .body(Body::empty())
6080 .unwrap();
6081 let resp = app.oneshot(req).await.unwrap();
6082
6083 let vary: Vec<String> = resp
6084 .headers()
6085 .get_all("vary")
6086 .iter()
6087 .filter_map(|v| v.to_str().ok().map(str::to_owned))
6088 .collect();
6089 assert!(
6090 vary.iter().any(|v| v.contains("Accept-Encoding")),
6091 "must preserve pre-existing Vary value, got {vary:?}"
6092 );
6093 assert!(
6094 vary.iter().any(|v| v.contains("Authorization")),
6095 "must append Authorization to Vary, got {vary:?}"
6096 );
6097 }
6098
6099 #[test]
6102 fn version_omits_build_fingerprint_by_default() {
6103 let v = version_payload("my-server", "1.2.3", false);
6104 assert_eq!(v["name"], "my-server");
6105 assert_eq!(v["version"], "1.2.3");
6106 assert!(v["rmcp_server_kit_version"].is_string());
6107 assert!(
6108 v.get("build_git_sha").is_none(),
6109 "build sha must be hidden by default"
6110 );
6111 assert!(v.get("build_timestamp").is_none());
6112 assert!(v.get("rust_version").is_none());
6113 }
6114
6115 #[test]
6116 fn version_exposes_all_when_enabled() {
6117 let v = version_payload("my-server", "1.2.3", true);
6118 assert!(v["build_git_sha"].is_string());
6119 assert!(v["build_timestamp"].is_string());
6120 assert!(v["rust_version"].is_string());
6121 assert!(v["rmcp_server_kit_version"].is_string());
6122 }
6123
6124 #[tokio::test]
6127 async fn concurrency_limit_layer_composes_and_serves() {
6128 let app = axum::Router::new()
6132 .route("/ok", axum::routing::get(|| async { "ok" }))
6133 .layer(
6134 tower::ServiceBuilder::new()
6135 .layer(axum::error_handling::HandleErrorLayer::new(
6136 |_err: tower::BoxError| async { StatusCode::SERVICE_UNAVAILABLE },
6137 ))
6138 .layer(tower::load_shed::LoadShedLayer::new())
6139 .layer(tower::limit::ConcurrencyLimitLayer::new(4)),
6140 );
6141 let resp = app
6142 .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap())
6143 .await
6144 .unwrap();
6145 assert_eq!(resp.status(), StatusCode::OK);
6146 }
6147
6148 #[tokio::test]
6151 async fn compression_layer_gzip_encodes_response() {
6152 use tower_http::compression::Predicate as _;
6153
6154 let big_body = "a".repeat(4096);
6155 let app = axum::Router::new()
6156 .route(
6157 "/big",
6158 axum::routing::get(move || {
6159 let body = big_body.clone();
6160 async move { body }
6161 }),
6162 )
6163 .layer(
6164 tower_http::compression::CompressionLayer::new()
6165 .gzip(true)
6166 .br(true)
6167 .compress_when(
6168 tower_http::compression::DefaultPredicate::new()
6169 .and(tower_http::compression::predicate::SizeAbove::new(1024)),
6170 ),
6171 );
6172
6173 let req = Request::builder()
6174 .uri("/big")
6175 .header(header::ACCEPT_ENCODING, "gzip")
6176 .body(Body::empty())
6177 .unwrap();
6178 let resp = app.oneshot(req).await.unwrap();
6179 assert_eq!(resp.status(), StatusCode::OK);
6180 assert_eq!(
6181 resp.headers().get(header::CONTENT_ENCODING).unwrap(),
6182 "gzip"
6183 );
6184 }
6185
6186 #[tokio::test]
6189 async fn tls_handshake_timeout_reaps_idle_connections() {
6190 use tokio::io::AsyncReadExt as _;
6191
6192 let _ = rustls::crypto::ring::default_provider().install_default();
6193
6194 let key = rcgen::KeyPair::generate().expect("generate key");
6196 let cert = rcgen::CertificateParams::new(vec!["localhost".to_owned()])
6197 .expect("cert params")
6198 .self_signed(&key)
6199 .expect("self-signed cert");
6200 let dir = std::env::temp_dir().join(format!(
6201 "rmcp-server-kit-hs-timeout-{}",
6202 std::time::SystemTime::now()
6203 .duration_since(std::time::UNIX_EPOCH)
6204 .expect("clock after epoch")
6205 .as_nanos()
6206 ));
6207 tokio::fs::create_dir_all(&dir).await.expect("temp dir");
6208 let cert_path = dir.join("server.crt");
6209 let key_path = dir.join("server.key");
6210 tokio::fs::write(&cert_path, cert.pem())
6211 .await
6212 .expect("write cert");
6213 tokio::fs::write(&key_path, key.serialize_pem())
6214 .await
6215 .expect("write key");
6216
6217 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
6218 let tls = TlsListener::new(
6219 listener,
6220 &cert_path,
6221 &key_path,
6222 None,
6223 None,
6224 Duration::from_millis(200),
6225 8, )
6227 .expect("tls listener");
6228 let addr = axum::serve::Listener::local_addr(&tls).expect("local addr");
6229
6230 let mut idle = tokio::net::TcpStream::connect(addr).await.expect("connect");
6234 let mut buf = [0_u8; 16];
6235 let read = tokio::time::timeout(Duration::from_secs(2), idle.read(&mut buf))
6236 .await
6237 .expect("server must reap the idle handshake within its timeout");
6238 match read {
6239 Ok(0) | Err(_) => {} Ok(n) => panic!("unexpected {n} bytes from server during reaped handshake"),
6241 }
6242
6243 drop(tls);
6244 }
6245
6246 fn assert_owasp_headers(resp: &axum::response::Response, ctx: &str) {
6249 let h = resp.headers();
6250 assert!(
6251 h.contains_key("x-content-type-options"),
6252 "{ctx}: missing X-Content-Type-Options"
6253 );
6254 assert!(
6255 h.contains_key("x-frame-options"),
6256 "{ctx}: missing X-Frame-Options"
6257 );
6258 assert!(
6259 h.contains_key("strict-transport-security"),
6260 "{ctx}: missing Strict-Transport-Security"
6261 );
6262 assert!(
6263 h.contains_key(header::CONTENT_SECURITY_POLICY),
6264 "{ctx}: missing Content-Security-Policy"
6265 );
6266 }
6267
6268 fn m5_router(configure: impl FnOnce(&mut McpServerConfig)) -> axum::Router {
6269 #[derive(Clone)]
6270 struct H;
6271 impl ServerHandler for H {}
6272 let mut config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0")
6276 .with_allowed_origins(["http://good.example"])
6277 .with_tls("unused.crt", "unused.key");
6278 configure(&mut config);
6279 let (router, _params) = build_app_router(config, || H).expect("build_app_router");
6280 router
6281 }
6282
6283 #[test]
6288 #[should_panic(expected = "Overlapping method route")]
6289 fn extra_router_exact_overlap_with_framework_route_panics() {
6290 #[derive(Clone)]
6291 struct H;
6292 impl ServerHandler for H {}
6293 let config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_extra_router(
6294 axum::Router::new().route("/healthz", axum::routing::get(|| async { "mine" })),
6295 );
6296 let _ = build_app_router(config, || H);
6297 }
6298
6299 #[test]
6303 fn extra_router_non_overlapping_path_under_framework_prefix_is_accepted() {
6304 #[derive(Clone)]
6305 struct H;
6306 impl ServerHandler for H {}
6307 let config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_extra_router(
6308 axum::Router::new().route("/admin/custom", axum::routing::get(|| async { "mine" })),
6309 );
6310 assert!(
6311 build_app_router(config, || H).is_ok(),
6312 "non-overlapping path under a framework prefix must merge cleanly"
6313 );
6314 }
6315
6316 #[tokio::test]
6317 async fn headers_on_rejected_origin_403() {
6318 let app = m5_router(|_| {});
6319 let req = Request::builder()
6320 .uri("/healthz")
6321 .header(header::ORIGIN, "http://evil.example")
6322 .body(Body::empty())
6323 .unwrap();
6324 let resp = app.oneshot(req).await.unwrap();
6325 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
6326 assert_owasp_headers(&resp, "origin-403");
6327 }
6328
6329 #[tokio::test]
6330 async fn headers_on_cors_preflight() {
6331 let app = m5_router(|_| {});
6332 let req = Request::builder()
6333 .method(axum::http::Method::OPTIONS)
6334 .uri("/mcp")
6335 .header(header::ORIGIN, "http://good.example")
6336 .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
6337 .body(Body::empty())
6338 .unwrap();
6339 let resp = app.oneshot(req).await.unwrap();
6340 assert_owasp_headers(&resp, "cors-preflight");
6341 }
6342
6343 #[tokio::test]
6344 async fn headers_on_404_fallback() {
6345 let app = m5_router(|_| {});
6346 let req = Request::builder()
6347 .uri("/no-such-route")
6348 .body(Body::empty())
6349 .unwrap();
6350 let resp = app.oneshot(req).await.unwrap();
6351 assert_eq!(resp.status(), StatusCode::NOT_FOUND);
6352 assert_owasp_headers(&resp, "404-fallback");
6353 }
6354
6355 #[tokio::test]
6356 async fn headers_on_overload_503() {
6357 let app = m5_router(|c| c.max_concurrent_requests = Some(0));
6360 let req = Request::builder()
6361 .uri("/healthz")
6362 .body(Body::empty())
6363 .unwrap();
6364 let resp = app.oneshot(req).await.unwrap();
6365 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
6366 assert_owasp_headers(&resp, "overload-503");
6367 }
6368
6369 #[cfg(feature = "oauth")]
6372 fn m6_auth_state() -> (Arc<AuthState>, String, String) {
6373 let (admin_token, admin_hash) = crate::auth::generate_api_key().unwrap();
6374 let (viewer_token, viewer_hash) = crate::auth::generate_api_key().unwrap();
6375 let state = Arc::new(AuthState {
6376 api_keys: ArcSwap::from_pointee(vec![
6377 crate::auth::ApiKeyEntry::new("admin-key", admin_hash, "admin"),
6378 crate::auth::ApiKeyEntry::new("viewer-key", viewer_hash, "viewer"),
6379 ]),
6380 rate_limiter: None,
6381 pre_auth_limiter: None,
6382 jwks_cache: None,
6383 seen_identities: crate::auth::SeenIdentitySet::new(),
6384 counters: crate::auth::AuthCounters::default(),
6385 resource_metadata_url: None,
6386 });
6387 (state, admin_token, viewer_token)
6388 }
6389
6390 #[cfg(feature = "oauth")]
6391 fn m6_admin_router(state: &Arc<AuthState>) -> axum::Router {
6392 let proxy = crate::oauth::OAuthProxyConfig::builder(
6393 "https://idp.example/authorize",
6394 "https://idp.example/token",
6395 "client",
6396 )
6397 .introspection_url("http://127.0.0.1:1/introspect")
6398 .revocation_url("http://127.0.0.1:1/revoke")
6399 .expose_admin_endpoints(true)
6400 .require_auth_on_admin_endpoints(true)
6401 .build();
6402 let http = crate::oauth::OauthHttpClient::new().expect("oauth http client");
6403 build_oauth_admin_router(&proxy, http, Some(state), "admin").expect("admin router")
6404 }
6405
6406 #[cfg(feature = "oauth")]
6407 fn m6_req(path: &str, token: &str) -> Request<Body> {
6408 Request::builder()
6409 .method(axum::http::Method::POST)
6410 .uri(path)
6411 .header(header::AUTHORIZATION, format!("Bearer {token}"))
6412 .body(Body::from("token=abc"))
6413 .unwrap()
6414 }
6415
6416 #[cfg(feature = "oauth")]
6417 #[tokio::test]
6418 async fn oauth_proxy_admin_requires_admin_role() {
6419 let (state, _admin, viewer) = m6_auth_state();
6420 for path in ["/introspect", "/revoke"] {
6421 let app = m6_admin_router(&state);
6422 let resp = app.oneshot(m6_req(path, &viewer)).await.unwrap();
6423 assert_eq!(
6424 resp.status(),
6425 StatusCode::FORBIDDEN,
6426 "an authenticated viewer must be rejected with 403 on {path}"
6427 );
6428 }
6429 }
6430
6431 #[cfg(feature = "oauth")]
6432 #[tokio::test]
6433 async fn oauth_proxy_admin_allows_admin_role() {
6434 let (state, admin, _viewer) = m6_auth_state();
6435 for path in ["/introspect", "/revoke"] {
6436 let app = m6_admin_router(&state);
6437 let resp = app.oneshot(m6_req(path, &admin)).await.unwrap();
6438 assert_ne!(
6442 resp.status(),
6443 StatusCode::FORBIDDEN,
6444 "an authenticated admin must pass the role gate on {path}"
6445 );
6446 assert_ne!(
6447 resp.status(),
6448 StatusCode::UNAUTHORIZED,
6449 "an authenticated admin must pass the auth gate on {path}"
6450 );
6451 }
6452 }
6453
6454 #[cfg(feature = "metrics")]
6462 mod metrics_labels_bounded {
6463 use super::*;
6464
6465 fn labels_for(method: &str, uri: &str) -> (&'static str, String) {
6466 let req = Request::builder()
6467 .method(method)
6468 .uri(uri)
6469 .body(Body::empty())
6470 .unwrap();
6471 metrics_labels(&req)
6472 }
6473
6474 #[test]
6475 fn many_unmatched_paths_collapse_to_one_label() {
6476 let mut seen = std::collections::HashSet::new();
6477 for i in 0..500 {
6478 let (_, path) = labels_for("GET", &format!("/nonexistent-{i}"));
6479 seen.insert(path);
6480 }
6481 assert_eq!(
6482 seen.len(),
6483 1,
6484 "unmatched paths must collapse to a single label, got {seen:?}"
6485 );
6486 assert!(seen.contains("<unmatched>"));
6487 }
6488
6489 #[test]
6490 fn nested_mcp_paths_collapse_to_the_mount_point() {
6491 let mut seen = std::collections::HashSet::new();
6492 for i in 0..200 {
6493 let (_, path) = labels_for("POST", &format!("/mcp/{i}"));
6494 seen.insert(path);
6495 }
6496 let (_, root) = labels_for("POST", "/mcp");
6497 seen.insert(root);
6498 assert_eq!(
6499 seen.len(),
6500 1,
6501 "nested /mcp paths must collapse to one label, got {seen:?}"
6502 );
6503 assert!(seen.contains("/mcp"));
6504 }
6505
6506 #[test]
6507 fn unusual_methods_collapse_to_one_bucket() {
6508 let mut seen = std::collections::HashSet::new();
6509 for verb in ["FROBNICATE", "WIBBLE", "QUUX", "M-SEARCH"] {
6510 let (method, _) = labels_for(verb, "/healthz");
6511 seen.insert(method);
6512 }
6513 assert_eq!(seen, std::collections::HashSet::from(["OTHER"]));
6514 }
6515
6516 #[test]
6517 fn known_methods_keep_their_identity() {
6518 for verb in ["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"] {
6519 let (method, _) = labels_for(verb, "/healthz");
6520 assert_eq!(method, verb);
6521 }
6522 }
6523
6524 #[test]
6525 fn raw_path_never_leaks_into_a_label() {
6526 let (_, path) = labels_for("GET", "/secret-token-abc123");
6527 assert!(
6528 !path.contains("secret-token"),
6529 "raw request path must never become a label value: {path}"
6530 );
6531 }
6532 }
6533}