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::{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 session_store: Option<Arc<dyn SessionStore>>,
493 #[deprecated(
496 since = "0.13.0",
497 note = "use McpServerConfig::with_sse_keep_alive(); direct field access will become pub(crate) in a future major release"
498 )]
499 pub sse_keep_alive: Duration,
500 #[deprecated(
504 since = "0.13.0",
505 note = "use McpServerConfig::with_reload_callback(); direct field access will become pub(crate) in a future major release"
506 )]
507 pub on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
508 #[deprecated(
515 since = "0.13.0",
516 note = "use McpServerConfig::with_extra_router(); direct field access will become pub(crate) in a future major release"
517 )]
518 pub extra_router: Option<axum::Router>,
519 #[deprecated(
524 since = "0.13.0",
525 note = "use McpServerConfig::with_public_url(); direct field access will become pub(crate) in a future major release"
526 )]
527 pub public_url: Option<String>,
528 #[deprecated(
531 since = "0.13.0",
532 note = "use McpServerConfig::enable_request_header_logging(); direct field access will become pub(crate) in a future major release"
533 )]
534 pub log_request_headers: bool,
535 pub expose_build_metadata: bool,
542 #[deprecated(
545 since = "0.13.0",
546 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
547 )]
548 pub compression_enabled: bool,
549 #[deprecated(
552 since = "0.13.0",
553 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
554 )]
555 pub compression_min_size: u16,
556 #[deprecated(
560 since = "0.13.0",
561 note = "use McpServerConfig::with_max_concurrent_requests(); direct field access will become pub(crate) in a future major release"
562 )]
563 pub max_concurrent_requests: Option<usize>,
564 #[deprecated(
567 since = "0.13.0",
568 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
569 )]
570 pub admin_enabled: bool,
571 #[deprecated(
573 since = "0.13.0",
574 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
575 )]
576 pub admin_role: String,
577 #[cfg(feature = "metrics")]
580 #[deprecated(
581 since = "0.13.0",
582 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
583 )]
584 pub metrics_enabled: bool,
585 #[cfg(feature = "metrics")]
587 #[deprecated(
588 since = "0.13.0",
589 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
590 )]
591 pub metrics_bind: String,
592 #[deprecated(
596 since = "1.5.0",
597 note = "use McpServerConfig::with_security_headers(); direct field access will become pub(crate) in a future major release"
598 )]
599 pub security_headers: SecurityHeadersConfig,
600 #[deprecated(
606 since = "1.9.0",
607 note = "use McpServerConfig::with_tls_handshake_timeout(); direct field access will become pub(crate) in a future major release"
608 )]
609 pub tls_handshake_timeout: Duration,
610 #[deprecated(
617 since = "1.9.0",
618 note = "use McpServerConfig::with_max_concurrent_tls_handshakes(); direct field access will become pub(crate) in a future major release"
619 )]
620 pub max_concurrent_tls_handshakes: usize,
621}
622
623#[allow(
681 missing_debug_implementations,
682 reason = "wraps T which may not implement Debug; manual impl below avoids leaking inner contents into logs"
683)]
684pub struct Validated<T>(T);
685
686impl<T> std::fmt::Debug for Validated<T> {
687 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
688 f.debug_struct("Validated").finish_non_exhaustive()
689 }
690}
691
692impl<T> Validated<T> {
693 #[must_use]
695 pub fn as_inner(&self) -> &T {
696 &self.0
697 }
698
699 #[must_use]
704 pub fn into_inner(self) -> T {
705 self.0
706 }
707}
708
709#[allow(
710 deprecated,
711 reason = "internal builders/validators legitimately read/write the deprecated `pub` fields they were designed to manage"
712)]
713impl McpServerConfig {
714 #[must_use]
722 pub fn new(
723 bind_addr: impl Into<String>,
724 name: impl Into<String>,
725 version: impl Into<String>,
726 ) -> Self {
727 Self {
728 bind_addr: bind_addr.into(),
729 name: name.into(),
730 version: version.into(),
731 tls_cert_path: None,
732 tls_key_path: None,
733 auth: None,
734 rbac: None,
735 tool_list_filtering: true,
736 allowed_origins: Vec::new(),
737 tool_rate_limit: None,
738 readiness_check: None,
739 max_request_body: 1024 * 1024,
740 request_timeout: Duration::from_mins(2),
741 shutdown_timeout: Duration::from_secs(30),
742 session_idle_timeout: Duration::from_mins(20),
743 session_binding: true,
744 session_binding_secret: None,
745 session_store: None,
746 sse_keep_alive: Duration::from_secs(15),
747 on_reload_ready: None,
748 extra_router: None,
749 public_url: None,
750 log_request_headers: false,
751 expose_build_metadata: false,
752 compression_enabled: false,
753 compression_min_size: 1024,
754 max_concurrent_requests: None,
755 admin_enabled: false,
756 admin_role: "admin".to_owned(),
757 #[cfg(feature = "metrics")]
758 metrics_enabled: false,
759 #[cfg(feature = "metrics")]
760 metrics_bind: "127.0.0.1:9090".into(),
761 security_headers: SecurityHeadersConfig::default(),
762 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
763 max_concurrent_tls_handshakes: DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES,
764 extra_route_rate_limit: None,
765 tool_rate_limit_burst: None,
766 extra_route_rate_limit_burst: None,
767 extra_route_rate_limit_exempt_paths: Vec::new(),
768 key_eviction_policy: KeyEvictionPolicy::default(),
769 trusted_forwarder_max_entries: crate::forwarded::MAX_SCANNED_ENTRIES,
770 trusted_proxies: Vec::new(),
771 forwarded_header: None,
772 }
773 }
774
775 #[must_use]
785 pub fn with_auth(mut self, auth: AuthConfig) -> Self {
786 self.auth = Some(auth);
787 self
788 }
789
790 #[must_use]
795 pub fn with_security_headers(mut self, headers: SecurityHeadersConfig) -> Self {
796 self.security_headers = headers;
797 self
798 }
799
800 #[must_use]
804 pub fn with_bind_addr(mut self, addr: impl Into<String>) -> Self {
805 self.bind_addr = addr.into();
806 self
807 }
808
809 #[must_use]
812 pub fn with_rbac(mut self, rbac: Arc<RbacPolicy>) -> Self {
813 self.rbac = Some(rbac);
814 self
815 }
816
817 #[must_use]
823 pub const fn with_tool_list_filtering(mut self, enabled: bool) -> Self {
824 self.tool_list_filtering = enabled;
825 self
826 }
827
828 #[must_use]
832 pub fn with_tls(mut self, cert_path: impl Into<PathBuf>, key_path: impl Into<PathBuf>) -> Self {
833 self.tls_cert_path = Some(cert_path.into());
834 self.tls_key_path = Some(key_path.into());
835 self
836 }
837
838 #[must_use]
842 pub fn with_public_url(mut self, url: impl Into<String>) -> Self {
843 self.public_url = Some(url.into());
844 self
845 }
846
847 #[must_use]
851 pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
852 where
853 I: IntoIterator<Item = S>,
854 S: Into<String>,
855 {
856 self.allowed_origins = origins.into_iter().map(Into::into).collect();
857 self
858 }
859
860 #[must_use]
893 pub fn with_extra_router(mut self, router: axum::Router) -> Self {
894 self.extra_router = Some(router);
895 self
896 }
897
898 #[must_use]
904 pub const fn with_trusted_forwarder_max_entries(mut self, max_entries: usize) -> Self {
905 self.trusted_forwarder_max_entries = max_entries;
906 self
907 }
908
909 #[must_use]
912 pub fn with_readiness_check(mut self, check: ReadinessCheck) -> Self {
913 self.readiness_check = Some(check);
914 self
915 }
916
917 #[must_use]
920 pub fn with_max_request_body(mut self, bytes: usize) -> Self {
921 self.max_request_body = bytes;
922 self
923 }
924
925 #[must_use]
927 pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
928 self.request_timeout = timeout;
929 self
930 }
931
932 #[must_use]
934 pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
935 self.shutdown_timeout = timeout;
936 self
937 }
938
939 #[must_use]
941 pub fn with_session_idle_timeout(mut self, timeout: Duration) -> Self {
942 self.session_idle_timeout = timeout;
943 self
944 }
945
946 #[must_use]
950 pub const fn with_session_binding(mut self, enabled: bool) -> Self {
951 self.session_binding = enabled;
952 self
953 }
954
955 #[must_use]
957 pub fn with_session_binding_secret(mut self, secret: SecretString) -> Self {
958 self.session_binding_secret = Some(secret);
959 self
960 }
961
962 #[must_use]
964 pub fn with_session_store(mut self, session_store: Arc<dyn SessionStore>) -> Self {
965 self.session_store = Some(session_store);
966 self
967 }
968
969 #[must_use]
971 pub fn with_sse_keep_alive(mut self, interval: Duration) -> Self {
972 self.sse_keep_alive = interval;
973 self
974 }
975
976 #[must_use]
980 pub fn with_max_concurrent_requests(mut self, limit: usize) -> Self {
981 self.max_concurrent_requests = Some(limit);
982 self
983 }
984
985 #[must_use]
993 pub fn with_tls_handshake_timeout(mut self, timeout: Duration) -> Self {
994 self.tls_handshake_timeout = timeout;
995 self
996 }
997
998 #[must_use]
1007 pub fn with_max_concurrent_tls_handshakes(mut self, limit: usize) -> Self {
1008 self.max_concurrent_tls_handshakes = limit;
1009 self
1010 }
1011
1012 #[must_use]
1015 pub fn with_tool_rate_limit(mut self, per_minute: u32) -> Self {
1016 self.tool_rate_limit = Some(per_minute);
1017 self
1018 }
1019
1020 #[must_use]
1031 pub fn with_extra_route_rate_limit(mut self, per_minute: u32) -> Self {
1032 self.extra_route_rate_limit = Some(per_minute);
1033 self
1034 }
1035
1036 #[must_use]
1041 pub fn with_tool_rate_limit_burst(mut self, burst: u32) -> Self {
1042 self.tool_rate_limit_burst = Some(burst);
1043 self
1044 }
1045
1046 #[must_use]
1052 pub fn with_extra_route_rate_limit_burst(mut self, burst: u32) -> Self {
1053 self.extra_route_rate_limit_burst = Some(burst);
1054 self
1055 }
1056
1057 #[must_use]
1077 pub fn with_extra_route_rate_limit_exempt_paths<I, S>(mut self, paths: I) -> Self
1078 where
1079 I: IntoIterator<Item = S>,
1080 S: Into<String>,
1081 {
1082 self.extra_route_rate_limit_exempt_paths = paths.into_iter().map(Into::into).collect();
1083 self
1084 }
1085
1086 #[must_use]
1088 pub const fn with_key_eviction_policy(mut self, policy: KeyEvictionPolicy) -> Self {
1089 self.key_eviction_policy = policy;
1090 self
1091 }
1092
1093 #[must_use]
1105 pub fn with_trusted_proxies<I, S>(mut self, proxies: I) -> Self
1106 where
1107 I: IntoIterator<Item = S>,
1108 S: Into<String>,
1109 {
1110 self.trusted_proxies = proxies.into_iter().map(Into::into).collect();
1111 self
1112 }
1113
1114 #[must_use]
1119 pub fn with_forwarded_header(mut self, mode: ForwardedHeaderMode) -> Self {
1120 self.forwarded_header = Some(mode);
1121 self
1122 }
1123
1124 #[must_use]
1128 pub fn with_reload_callback<F>(mut self, callback: F) -> Self
1129 where
1130 F: FnOnce(ReloadHandle) + Send + 'static,
1131 {
1132 self.on_reload_ready = Some(Box::new(callback));
1133 self
1134 }
1135
1136 #[must_use]
1140 pub fn enable_compression(mut self, min_size: u16) -> Self {
1141 self.compression_enabled = true;
1142 self.compression_min_size = min_size;
1143 self
1144 }
1145
1146 #[must_use]
1151 pub fn enable_admin(mut self, role: impl Into<String>) -> Self {
1152 self.admin_enabled = true;
1153 self.admin_role = role.into();
1154 self
1155 }
1156
1157 #[must_use]
1160 pub fn enable_request_header_logging(mut self) -> Self {
1161 self.log_request_headers = true;
1162 self
1163 }
1164
1165 #[must_use]
1170 pub fn expose_build_metadata(mut self) -> Self {
1171 self.expose_build_metadata = true;
1172 self
1173 }
1174
1175 #[cfg(feature = "metrics")]
1178 #[must_use]
1179 pub fn with_metrics(mut self, bind: impl Into<String>) -> Self {
1180 self.metrics_enabled = true;
1181 self.metrics_bind = bind.into();
1182 self
1183 }
1184
1185 pub fn validate(self) -> Result<Validated<Self>, RmcpServerKitError> {
1218 self.check()?;
1219 Ok(Validated(self))
1220 }
1221
1222 fn check_burst_knobs(&self) -> Result<(), RmcpServerKitError> {
1229 if self.tool_rate_limit_burst == Some(0) {
1230 return Err(RmcpServerKitError::Config(
1231 "tool_rate_limit_burst must be greater than zero".into(),
1232 ));
1233 }
1234 if self.extra_route_rate_limit_burst == Some(0) {
1235 return Err(RmcpServerKitError::Config(
1236 "extra_route_rate_limit_burst must be greater than zero".into(),
1237 ));
1238 }
1239 if self.trusted_forwarder_max_entries == 0
1240 || self.trusted_forwarder_max_entries
1241 > crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES
1242 {
1243 return Err(RmcpServerKitError::Config(format!(
1244 "trusted_forwarder_max_entries must be in 1..={}, got {}",
1245 crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES,
1246 self.trusted_forwarder_max_entries
1247 )));
1248 }
1249 if self.tool_rate_limit_burst.is_some() && self.tool_rate_limit.is_none() {
1250 return Err(RmcpServerKitError::Config(
1251 "tool_rate_limit_burst requires tool_rate_limit to be set".into(),
1252 ));
1253 }
1254 if self.extra_route_rate_limit_burst.is_some() && self.extra_route_rate_limit.is_none() {
1255 return Err(RmcpServerKitError::Config(
1256 "extra_route_rate_limit_burst requires extra_route_rate_limit to be set".into(),
1257 ));
1258 }
1259 if !self.extra_route_rate_limit_exempt_paths.is_empty()
1260 && self.extra_route_rate_limit.is_none()
1261 {
1262 return Err(RmcpServerKitError::Config(
1263 "extra_route_rate_limit_exempt_paths requires extra_route_rate_limit to be set"
1264 .into(),
1265 ));
1266 }
1267 for path in &self.extra_route_rate_limit_exempt_paths {
1268 if path.is_empty() || !path.starts_with('/') {
1269 return Err(RmcpServerKitError::Config(format!(
1270 "extra_route_rate_limit_exempt_paths entries must be non-empty and start with '/': {path:?}"
1271 )));
1272 }
1273 }
1274 if let Some(rl) = self.auth.as_ref().and_then(|a| a.rate_limit.as_ref()) {
1275 if rl.burst == Some(0) {
1276 return Err(RmcpServerKitError::Config(
1277 "auth rate_limit.burst must be greater than zero".into(),
1278 ));
1279 }
1280 if rl.pre_auth_burst == Some(0) {
1281 return Err(RmcpServerKitError::Config(
1282 "auth rate_limit.pre_auth_burst must be greater than zero".into(),
1283 ));
1284 }
1285 }
1286 Ok(())
1287 }
1288
1289 fn check_trusted_forwarder(&self) -> Result<(), RmcpServerKitError> {
1294 for entry in &self.trusted_proxies {
1295 validate_trusted_proxy_entry(entry).map_err(RmcpServerKitError::Config)?;
1296 }
1297 if self.forwarded_header.is_some() && self.trusted_proxies.is_empty() {
1298 return Err(RmcpServerKitError::Config(
1299 "forwarded_header requires trusted_proxies to be nonempty".into(),
1300 ));
1301 }
1302 Ok(())
1303 }
1304
1305 fn check_session_binding_config(&self) -> Result<(), RmcpServerKitError> {
1306 if self.session_store.is_some()
1307 && self.session_binding
1308 && self.auth.as_ref().is_some_and(|auth| auth.enabled)
1309 && self.session_binding_secret.is_none()
1310 {
1311 return Err(RmcpServerKitError::Config(
1312 "session_store with session_binding enabled and auth configured requires \
1313 session_binding_secret: a shared secret is required for cross-instance \
1314 session verification"
1315 .into(),
1316 ));
1317 }
1318
1319 if let Some(secret) = &self.session_binding_secret {
1320 crate::session_binding::validate_configured_secret(
1321 "session_binding_secret",
1322 secret.expose_secret(),
1323 )?;
1324 }
1325 Ok(())
1326 }
1327
1328 fn check(&self) -> Result<(), RmcpServerKitError> {
1332 if let Err(violation) = crate::config::check_shared_config_invariants(
1347 self.admin_enabled,
1348 self.auth.as_ref().is_some_and(|a| a.enabled),
1349 self.tls_cert_path.is_some(),
1350 self.tls_key_path.is_some(),
1351 self.auth.as_ref().is_some_and(|a| a.mtls.is_some()),
1352 ) {
1353 return Err(RmcpServerKitError::Config(
1354 match violation {
1355 crate::config::SharedConfigViolation::AdminRequiresAuth => {
1356 "admin_enabled=true requires auth to be configured and enabled"
1357 }
1358 crate::config::SharedConfigViolation::TlsCertWithoutKey => {
1359 "tls_cert_path is set but tls_key_path is missing"
1360 }
1361 crate::config::SharedConfigViolation::TlsKeyWithoutCert => {
1362 "tls_key_path is set but tls_cert_path is missing"
1363 }
1364 crate::config::SharedConfigViolation::MtlsRequiresTls => {
1365 "auth.mtls requires TLS: set both tls_cert_path and tls_key_path \
1366 (mTLS client certificates cannot be verified on a plaintext listener)"
1367 }
1368 }
1369 .into(),
1370 ));
1371 }
1372
1373 if self.bind_addr.parse::<SocketAddr>().is_err() {
1375 return Err(RmcpServerKitError::Config(format!(
1376 "bind_addr {:?} is not a valid socket address (expected e.g. 127.0.0.1:8080)",
1377 self.bind_addr
1378 )));
1379 }
1380
1381 if let Some(ref url) = self.public_url
1383 && !(url.starts_with("http://") || url.starts_with("https://"))
1384 {
1385 return Err(RmcpServerKitError::Config(format!(
1386 "public_url {url:?} must start with http:// or https://"
1387 )));
1388 }
1389
1390 for origin in &self.allowed_origins {
1392 if !(origin.starts_with("http://") || origin.starts_with("https://")) {
1393 return Err(RmcpServerKitError::Config(format!(
1394 "allowed_origins entry {origin:?} must start with http:// or https://"
1395 )));
1396 }
1397 }
1398
1399 if self.max_request_body == 0 {
1401 return Err(RmcpServerKitError::Config(
1402 "max_request_body must be greater than zero".into(),
1403 ));
1404 }
1405
1406 if self.extra_route_rate_limit == Some(0) {
1410 return Err(RmcpServerKitError::Config(
1411 "extra_route_rate_limit must be greater than zero".into(),
1412 ));
1413 }
1414
1415 self.check_burst_knobs()?;
1417
1418 self.check_trusted_forwarder()?;
1420
1421 #[cfg(feature = "oauth")]
1423 if let Some(auth_cfg) = &self.auth
1424 && let Some(oauth_cfg) = &auth_cfg.oauth
1425 {
1426 oauth_cfg.validate()?;
1427 }
1428
1429 self.check_session_binding_config()?;
1430
1431 validate_security_headers(&self.security_headers)?;
1434
1435 if self.max_concurrent_requests == Some(0) {
1439 return Err(RmcpServerKitError::Config(
1440 "max_concurrent_requests must be greater than zero when set".into(),
1441 ));
1442 }
1443
1444 if let Some(auth_cfg) = &self.auth
1448 && let Some(rl) = &auth_cfg.rate_limit
1449 && rl.max_tracked_keys == 0
1450 {
1451 return Err(RmcpServerKitError::Config(
1452 "auth.rate_limit.max_tracked_keys must be greater than zero".into(),
1453 ));
1454 }
1455
1456 check_auth_capacity_knobs(self.auth.as_ref())?;
1457
1458 if self.tls_handshake_timeout == Duration::ZERO {
1463 return Err(RmcpServerKitError::Config(
1464 "tls_handshake_timeout must be greater than zero".into(),
1465 ));
1466 }
1467
1468 if self.max_concurrent_tls_handshakes == 0 {
1473 return Err(RmcpServerKitError::Config(
1474 "max_concurrent_tls_handshakes must be greater than zero".into(),
1475 ));
1476 }
1477
1478 Ok(())
1479 }
1480}
1481
1482#[allow(
1488 missing_debug_implementations,
1489 reason = "contains Arc<AuthState> with non-Debug fields"
1490)]
1491pub struct ReloadHandle {
1492 auth: Option<Arc<AuthState>>,
1493 rbac: Option<Arc<ArcSwap<RbacPolicy>>>,
1494 crl_set: Option<Arc<CrlSet>>,
1495}
1496
1497impl ReloadHandle {
1498 pub fn reload_auth_keys(&self, keys: Vec<crate::auth::ApiKeyEntry>) {
1500 if let Some(ref auth) = self.auth {
1501 auth.reload_keys(keys);
1502 }
1503 }
1504
1505 pub fn reload_rbac(&self, policy: RbacPolicy) {
1507 if let Some(ref rbac) = self.rbac {
1508 rbac.store(Arc::new(policy));
1509 tracing::info!("RBAC policy reloaded");
1510 }
1511 }
1512
1513 pub async fn refresh_crls(&self) -> Result<(), RmcpServerKitError> {
1523 let Some(ref crl_set) = self.crl_set else {
1524 return Err(RmcpServerKitError::Config(
1525 "CRL refresh requested but mTLS CRL support is not configured".into(),
1526 ));
1527 };
1528
1529 crl_set.force_refresh().await
1530 }
1531}
1532
1533#[allow(
1550 clippy::too_many_lines,
1551 clippy::cognitive_complexity,
1552 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"
1553)]
1554struct AppRunParams {
1558 tls_paths: Option<(PathBuf, PathBuf)>,
1560 tls_handshake_timeout: Duration,
1562 max_concurrent_tls_handshakes: usize,
1564 mtls_config: Option<MtlsConfig>,
1566 shutdown_timeout: Duration,
1568 auth_state: Option<Arc<AuthState>>,
1570 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
1572 on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
1574 ct: CancellationToken,
1578 session_ct: CancellationToken,
1588 scheme: &'static str,
1590 name: String,
1592}
1593
1594#[allow(
1604 clippy::cognitive_complexity,
1605 reason = "router assembly is intrinsically sequential; splitting harms readability"
1606)]
1607#[allow(
1608 deprecated,
1609 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1610)]
1611fn build_app_router<H, F>(
1612 mut config: McpServerConfig,
1613 handler_factory: F,
1614) -> anyhow::Result<(axum::Router, AppRunParams)>
1615where
1616 H: ServerHandler + 'static,
1617 F: Fn() -> H + Send + Sync + Clone + 'static,
1618{
1619 let ct = CancellationToken::new();
1620 let session_ct = CancellationToken::new();
1621
1622 let allowed_hosts = derive_allowed_hosts(&config.bind_addr, config.public_url.as_deref());
1623 tracing::info!(allowed_hosts = %allowed_hosts.join(", "), "configured Streamable HTTP allowed hosts");
1624
1625 if config.max_concurrent_requests.is_none() {
1626 tracing::warn!(
1627 "max_concurrent_requests is unset: in-flight HTTP requests are unlimited; \
1628 set McpServerConfig::with_max_concurrent_requests or front the server with \
1629 an external concurrency limit"
1630 );
1631 }
1632
1633 let rbac_swap = Arc::new(ArcSwap::new(
1636 config
1637 .rbac
1638 .clone()
1639 .unwrap_or_else(|| Arc::new(RbacPolicy::disabled())),
1640 ));
1641
1642 let rbac_for_handler = Arc::clone(&rbac_swap);
1643 let tool_list_filtering = config.tool_list_filtering;
1644 let session_store = config.session_store.take();
1645 let mut rmcp_config = StreamableHttpServerConfig::default()
1646 .with_allowed_hosts(allowed_hosts)
1647 .with_sse_keep_alive(Some(config.sse_keep_alive))
1648 .with_cancellation_token(session_ct.clone());
1649 rmcp_config.session_store = session_store;
1650 let mcp_service = StreamableHttpService::new(
1651 move || {
1652 Ok(RbacContextHandler::new(
1653 handler_factory(),
1654 Arc::clone(&rbac_for_handler),
1655 tool_list_filtering,
1656 ))
1657 },
1658 {
1659 let mut mgr = LocalSessionManager::default();
1660 mgr.session_config.keep_alive = Some(config.session_idle_timeout);
1661 mgr.into()
1662 },
1663 rmcp_config,
1664 );
1665
1666 let mut mcp_router = axum::Router::new().nest_service("/mcp", mcp_service);
1668
1669 let auth_state: Option<Arc<AuthState>> = match config.auth {
1673 Some(ref auth_config) if auth_config.enabled => {
1674 let rate_limiter = auth_config.rate_limit.as_ref().map(build_rate_limiter);
1675 let pre_auth_limiter = auth_config
1676 .rate_limit
1677 .as_ref()
1678 .map(crate::auth::build_pre_auth_limiter);
1679
1680 #[cfg(feature = "oauth")]
1681 let jwks_cache = auth_config
1682 .oauth
1683 .as_ref()
1684 .map(|c| crate::oauth::JwksCache::new(c).map(Arc::new))
1685 .transpose()
1686 .map_err(|e| std::io::Error::other(format!("JWKS HTTP client: {e}")))?;
1687
1688 Some(Arc::new(AuthState {
1689 api_keys: ArcSwap::new(Arc::new(auth_config.api_keys.clone())),
1690 rate_limiter,
1691 pre_auth_limiter,
1692 #[cfg(feature = "oauth")]
1693 jwks_cache,
1694 seen_identities: crate::auth::SeenIdentitySet::new(),
1695 counters: crate::auth::AuthCounters::default(),
1696 resource_metadata_url: config.public_url.as_ref().map(|url| {
1704 format!(
1705 "{}/.well-known/oauth-protected-resource/mcp",
1706 url.trim_end_matches('/')
1707 )
1708 }),
1709 }))
1710 }
1711 _ => None,
1712 };
1713
1714 if config.admin_enabled {
1717 let Some(ref auth_state_ref) = auth_state else {
1718 return Err(anyhow::anyhow!(
1719 "admin_enabled=true requires auth to be configured and enabled"
1720 ));
1721 };
1722 let admin_state = crate::admin::AdminState {
1723 started_at: std::time::Instant::now(),
1724 name: config.name.clone(),
1725 version: config.version.clone(),
1726 auth: Some(Arc::clone(auth_state_ref)),
1727 rbac: Arc::clone(&rbac_swap),
1728 };
1729 let admin_cfg = crate::admin::AdminConfig {
1730 role: config.admin_role.clone(),
1731 };
1732 mcp_router = mcp_router.merge(crate::admin::admin_router(admin_state, &admin_cfg));
1733 tracing::info!(role = %config.admin_role, "/admin/* endpoints enabled");
1734 }
1735
1736 if config.session_binding {
1768 let secret = match config.session_binding_secret.as_ref() {
1769 Some(configured) => configured_session_binding_secret(configured)?,
1770 None => process_session_binding_secret().clone(),
1771 };
1772 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1773 let secret = secret.clone();
1774 session_binding_middleware(secret, req, next)
1775 }));
1776 }
1777
1778 {
1782 let tool_limiter: Option<Arc<ToolRateLimiter>> = config.tool_rate_limit.map(|per_minute| {
1783 build_tool_rate_limiter_with_policy(
1784 per_minute,
1785 config.tool_rate_limit_burst,
1786 config.key_eviction_policy,
1787 )
1788 });
1789
1790 if rbac_swap.load().is_enabled() {
1791 tracing::info!("RBAC enforcement enabled on /mcp");
1792 }
1793 if let Some(limit) = config.tool_rate_limit {
1794 tracing::info!(limit, "tool rate limiting enabled (calls/min per IP)");
1795 }
1796
1797 let rbac_for_mw = Arc::clone(&rbac_swap);
1798 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1799 let p = rbac_for_mw.load_full();
1800 let tl = tool_limiter.clone();
1801 rbac_middleware(p, tl, req, next)
1802 }));
1803 }
1804
1805 if let Some(ref auth_config) = config.auth
1807 && auth_config.enabled
1808 {
1809 let Some(ref state) = auth_state else {
1810 return Err(anyhow::anyhow!("auth state missing despite enabled config"));
1811 };
1812
1813 let methods: Vec<&str> = [
1814 auth_config.mtls.is_some().then_some("mTLS"),
1815 (!auth_config.api_keys.is_empty()).then_some("bearer"),
1816 #[cfg(feature = "oauth")]
1817 auth_config.oauth.is_some().then_some("oauth-jwt"),
1818 ]
1819 .into_iter()
1820 .flatten()
1821 .collect();
1822
1823 tracing::info!(
1824 methods = %methods.join(", "),
1825 api_keys = auth_config.api_keys.len(),
1826 "auth enabled on /mcp"
1827 );
1828
1829 let state_for_mw = Arc::clone(state);
1830 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1831 let s = Arc::clone(&state_for_mw);
1832 auth_middleware(s, req, next)
1833 }));
1834 }
1835
1836 mcp_router = mcp_router.layer(tower_http::timeout::TimeoutLayer::with_status_code(
1839 axum::http::StatusCode::REQUEST_TIMEOUT,
1840 config.request_timeout,
1841 ));
1842
1843 mcp_router = mcp_router.layer(tower_http::limit::RequestBodyLimitLayer::new(
1847 config.max_request_body,
1848 ));
1849
1850 let mut effective_origins = config.allowed_origins.clone();
1857 if effective_origins.is_empty()
1858 && let Some(ref url) = config.public_url
1859 {
1860 if let Some(scheme_end) = url.find("://") {
1865 let scheme_with_sep = url.get(..scheme_end + 3).unwrap_or_default();
1866 let after_scheme = url.get(scheme_end + 3..).unwrap_or_default();
1867 let host_end = after_scheme.find('/').unwrap_or(after_scheme.len());
1868 let host = after_scheme.get(..host_end).unwrap_or_default();
1869 let origin = format!("{scheme_with_sep}{host}");
1870 tracing::info!(
1871 %origin,
1872 "auto-derived allowed origin from public_url"
1873 );
1874 effective_origins.push(origin);
1875 }
1876 }
1877 let allowed_origins: Arc<[String]> = Arc::from(effective_origins);
1878 let cors_origins = Arc::clone(&allowed_origins);
1879 let log_request_headers = config.log_request_headers;
1880
1881 let readyz_route = if let Some(check) = config.readiness_check.take() {
1882 axum::routing::get(move || readyz(Arc::clone(&check)))
1883 } else {
1884 axum::routing::get(healthz)
1885 };
1886
1887 #[allow(
1888 unused_mut,
1889 reason = "the binding is only reassigned when the `oauth` feature adds the \
1890 protected-resource-metadata route below"
1891 )]
1892 let mut router = axum::Router::new()
1893 .route("/healthz", axum::routing::get(healthz))
1894 .route("/readyz", readyz_route)
1895 .route(
1896 "/version",
1897 axum::routing::get({
1898 let payload_bytes: Arc<[u8]> = serialize_version_payload(
1903 &config.name,
1904 &config.version,
1905 config.expose_build_metadata,
1906 );
1907 move || {
1908 let p = Arc::clone(&payload_bytes);
1909 async move {
1910 (
1911 [(axum::http::header::CONTENT_TYPE, "application/json")],
1912 p.to_vec(),
1913 )
1914 }
1915 }
1916 }),
1917 )
1918 .merge(mcp_router);
1919
1920 if let Some(extra) = config.extra_router.take() {
1927 let extra = match config.extra_route_rate_limit {
1928 Some(per_minute) => {
1929 let max_tracked_keys =
1930 NonZeroUsize::new(EXTRA_ROUTE_MAX_TRACKED_KEYS).unwrap_or(NonZeroUsize::MIN);
1931 let limiter = build_extra_route_rate_limiter_with_policy(
1932 per_minute,
1933 config.extra_route_rate_limit_burst,
1934 config.key_eviction_policy,
1935 max_tracked_keys,
1936 );
1937 let exempt: Arc<std::collections::HashSet<String>> = Arc::new(
1938 config
1939 .extra_route_rate_limit_exempt_paths
1940 .iter()
1941 .cloned()
1942 .collect(),
1943 );
1944 tracing::info!(
1945 per_minute,
1946 exempt_paths = exempt.len(),
1947 "extra-route per-IP rate limit enabled"
1948 );
1949 extra.layer(axum::middleware::from_fn(move |req, next| {
1950 let l = Arc::clone(&limiter);
1951 let e = Arc::clone(&exempt);
1952 extra_route_rate_limit_middleware(l, e, req, next)
1953 }))
1954 }
1955 None => extra,
1956 };
1957 router = router.merge(extra);
1958 }
1959
1960 let server_url = derive_server_url(&config);
1967 let resource_url = format!("{server_url}/mcp");
1968
1969 #[cfg(feature = "oauth")]
1970 let prm_metadata = if let Some(ref auth_config) = config.auth
1971 && let Some(ref oauth_config) = auth_config.oauth
1972 {
1973 crate::oauth::protected_resource_metadata(&resource_url, &server_url, oauth_config)
1974 } else {
1975 serde_json::json!({ "resource": resource_url })
1976 };
1977 #[cfg(not(feature = "oauth"))]
1978 let prm_metadata = serde_json::json!({ "resource": resource_url });
1979
1980 let prm_root = prm_metadata.clone();
1986 router = router.route(
1987 "/.well-known/oauth-protected-resource",
1988 axum::routing::get(move || {
1989 let m = prm_root.clone();
1990 async move { axum::Json(m) }
1991 }),
1992 );
1993 router = router.route(
1994 "/.well-known/oauth-protected-resource/mcp",
1995 axum::routing::get(move || {
1996 let m = prm_metadata.clone();
1997 async move { axum::Json(m) }
1998 }),
1999 );
2000
2001 #[cfg(feature = "oauth")]
2006 if let Some(ref auth_config) = config.auth
2007 && let Some(ref oauth_config) = auth_config.oauth
2008 && oauth_config.proxy.is_some()
2009 {
2010 router = install_oauth_proxy_routes(
2011 router,
2012 &server_url,
2013 oauth_config,
2014 auth_state.as_ref(),
2015 config.max_request_body,
2016 &config.admin_role,
2017 )?;
2018 }
2019
2020 if !cors_origins.is_empty() {
2029 let cors = tower_http::cors::CorsLayer::new()
2030 .allow_origin(
2031 cors_origins
2032 .iter()
2033 .filter_map(|o| o.parse::<axum::http::HeaderValue>().ok())
2034 .collect::<Vec<_>>(),
2035 )
2036 .allow_methods([
2037 axum::http::Method::GET,
2038 axum::http::Method::POST,
2039 axum::http::Method::OPTIONS,
2040 ])
2041 .allow_headers([
2042 axum::http::header::CONTENT_TYPE,
2043 axum::http::header::AUTHORIZATION,
2044 ]);
2045 router = router.layer(cors);
2046 }
2047
2048 if config.compression_enabled {
2052 use tower_http::compression::Predicate as _;
2053 let predicate = tower_http::compression::DefaultPredicate::new().and(
2054 tower_http::compression::predicate::SizeAbove::new(u64::from(
2055 config.compression_min_size,
2056 )),
2057 );
2058 router = router.layer(
2059 tower_http::compression::CompressionLayer::new()
2060 .gzip(true)
2061 .br(true)
2062 .compress_when(predicate),
2063 );
2064 tracing::info!(
2065 min_size = config.compression_min_size,
2066 "response compression enabled (gzip, br)"
2067 );
2068 }
2069
2070 if let Some(max) = config.max_concurrent_requests {
2073 let overload_handler = tower::ServiceBuilder::new()
2074 .layer(axum::error_handling::HandleErrorLayer::new(
2075 |_err: tower::BoxError| async {
2076 (
2077 axum::http::StatusCode::SERVICE_UNAVAILABLE,
2078 axum::Json(serde_json::json!({
2079 "error": "overloaded",
2080 "error_description": "server is at capacity, retry later"
2081 })),
2082 )
2083 },
2084 ))
2085 .layer(tower::load_shed::LoadShedLayer::new())
2086 .layer(tower::limit::ConcurrencyLimitLayer::new(max));
2087 router = router.layer(overload_handler);
2088 tracing::info!(max, "global concurrency limit enabled");
2089 }
2090
2091 router = router.fallback(|| async {
2095 (
2096 axum::http::StatusCode::NOT_FOUND,
2097 axum::Json(serde_json::json!({
2098 "error": "not_found",
2099 "error_description": "The requested endpoint does not exist"
2100 })),
2101 )
2102 });
2103
2104 #[cfg(feature = "metrics")]
2106 if config.metrics_enabled {
2107 let metrics = Arc::new(
2108 crate::metrics::McpMetrics::new().map_err(|e| anyhow::anyhow!("metrics init: {e}"))?,
2109 );
2110 let m = Arc::clone(&metrics);
2111 router = router.layer(axum::middleware::from_fn(
2112 move |req: Request<Body>, next: Next| {
2113 let m = Arc::clone(&m);
2114 metrics_middleware(m, req, next)
2115 },
2116 ));
2117 let metrics_bind = config.metrics_bind.clone();
2118 let metrics_shutdown = ct.clone();
2119 tokio::spawn(async move {
2120 if let Err(e) =
2121 crate::metrics::serve_metrics(metrics_bind, metrics, metrics_shutdown).await
2122 {
2123 tracing::error!("metrics listener failed: {e}");
2124 }
2125 });
2126 }
2127
2128 let forward_resolver: Option<Arc<ForwardResolver>> = if config.trusted_proxies.is_empty() {
2136 None
2137 } else {
2138 Some(Arc::new(ForwardResolver {
2141 trusted: config
2142 .trusted_proxies
2143 .iter()
2144 .filter_map(|entry| parse_proxy_net(entry))
2145 .collect(),
2146 mode: config
2147 .forwarded_header
2148 .unwrap_or(ForwardedHeaderMode::XForwardedFor),
2149 max_scanned_entries: config.trusted_forwarder_max_entries,
2150 }))
2151 };
2152 if forward_resolver.is_some() {
2153 tracing::info!(
2154 proxies = config.trusted_proxies.len(),
2155 "trusted-forwarder mode enabled: limiters key by resolved client IP"
2156 );
2157 }
2158 router = router.layer(axum::middleware::from_fn(move |req, next| {
2159 let r = forward_resolver.clone();
2160 normalize_peer_addr_middleware(r, req, next)
2161 }));
2162
2163 router = router.layer(axum::middleware::from_fn(move |req, next| {
2175 let origins = Arc::clone(&allowed_origins);
2176 origin_check_middleware(origins, log_request_headers, req, next)
2177 }));
2178
2179 let is_tls = config.tls_cert_path.is_some();
2188 warn_security_header_overrides(&config.security_headers);
2189 let security_headers_cfg = Arc::new(config.security_headers.clone());
2190 router = router.layer(axum::middleware::from_fn(move |req, next| {
2191 let cfg = Arc::clone(&security_headers_cfg);
2192 security_headers_middleware(is_tls, cfg, req, next)
2193 }));
2194
2195 let scheme = if config.tls_cert_path.is_some() {
2196 "https"
2197 } else {
2198 "http"
2199 };
2200
2201 let tls_paths = match (&config.tls_cert_path, &config.tls_key_path) {
2202 (Some(cert), Some(key)) => Some((cert.clone(), key.clone())),
2203 _ => None,
2204 };
2205 let tls_handshake_timeout = config.tls_handshake_timeout;
2206 let max_concurrent_tls_handshakes = config.max_concurrent_tls_handshakes;
2207 let mtls_config = config.auth.as_ref().and_then(|a| a.mtls.as_ref()).cloned();
2208
2209 Ok((
2210 router,
2211 AppRunParams {
2212 tls_paths,
2213 tls_handshake_timeout,
2214 max_concurrent_tls_handshakes,
2215 mtls_config,
2216 shutdown_timeout: config.shutdown_timeout,
2217 auth_state,
2218 rbac_swap,
2219 on_reload_ready: config.on_reload_ready.take(),
2220 ct,
2221 session_ct,
2222 scheme,
2223 name: config.name.clone(),
2224 },
2225 ))
2226}
2227
2228struct CancelOnDrop(CancellationToken);
2241
2242impl Drop for CancelOnDrop {
2243 fn drop(&mut self) {
2244 self.0.cancel();
2245 }
2246}
2247
2248fn spawn_external_shutdown_bridge(
2252 external: CancellationToken,
2253 internal: CancellationToken,
2254) -> tokio::task::JoinHandle<()> {
2255 tokio::spawn(async move {
2256 tokio::select! {
2260 () = external.cancelled() => internal.cancel(),
2261 () = internal.cancelled() => {}
2262 }
2263 })
2264}
2265
2266pub async fn serve<H, F>(
2286 config: Validated<McpServerConfig>,
2287 handler_factory: F,
2288) -> Result<(), RmcpServerKitError>
2289where
2290 H: ServerHandler + 'static,
2291 F: Fn() -> H + Send + Sync + Clone + 'static,
2292{
2293 let config = config.into_inner();
2294 #[allow(
2295 deprecated,
2296 reason = "internal serve() reads `bind_addr` to construct the listener; field becomes pub(crate) in 1.0"
2297 )]
2298 let bind_addr = config.bind_addr.clone();
2299 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2300 let _cancel_guard = CancelOnDrop(params.ct.clone());
2301
2302 let listener = TcpListener::bind(&bind_addr)
2303 .await
2304 .map_err(|e| io_to_startup(&format!("bind {bind_addr}"), e))?;
2305 log_listening(¶ms.name, params.scheme, &bind_addr);
2306
2307 run_server(
2308 router,
2309 listener,
2310 params.tls_paths,
2311 params.tls_handshake_timeout,
2312 params.max_concurrent_tls_handshakes,
2313 params.mtls_config,
2314 params.shutdown_timeout,
2315 params.auth_state,
2316 params.rbac_swap,
2317 params.on_reload_ready,
2318 params.ct,
2319 params.session_ct,
2320 )
2321 .await
2322 .map_err(anyhow_to_startup)
2323}
2324
2325pub async fn serve_with_listener<H, F>(
2358 listener: TcpListener,
2359 config: Validated<McpServerConfig>,
2360 handler_factory: F,
2361 ready_tx: Option<tokio::sync::oneshot::Sender<SocketAddr>>,
2362 shutdown: Option<CancellationToken>,
2363) -> Result<(), RmcpServerKitError>
2364where
2365 H: ServerHandler + 'static,
2366 F: Fn() -> H + Send + Sync + Clone + 'static,
2367{
2368 let config = config.into_inner();
2369 let local_addr = listener
2370 .local_addr()
2371 .map_err(|e| io_to_startup("listener.local_addr", e))?;
2372 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2373 let _cancel_guard = CancelOnDrop(params.ct.clone());
2374
2375 log_listening(¶ms.name, params.scheme, &local_addr.to_string());
2376
2377 if let Some(external) = shutdown {
2381 let _bridge_task = spawn_external_shutdown_bridge(external, params.ct.clone());
2382 }
2383
2384 if let Some(tx) = ready_tx {
2388 let _ = tx.send(local_addr);
2390 }
2391
2392 run_server(
2393 router,
2394 listener,
2395 params.tls_paths,
2396 params.tls_handshake_timeout,
2397 params.max_concurrent_tls_handshakes,
2398 params.mtls_config,
2399 params.shutdown_timeout,
2400 params.auth_state,
2401 params.rbac_swap,
2402 params.on_reload_ready,
2403 params.ct,
2404 params.session_ct,
2405 )
2406 .await
2407 .map_err(anyhow_to_startup)
2408}
2409
2410#[allow(
2413 clippy::cognitive_complexity,
2414 reason = "tracing::info! macro expansions inflate the score; logic is trivial"
2415)]
2416fn log_listening(name: &str, scheme: &str, addr: &str) {
2417 tracing::info!("{name} listening on {addr}");
2418 tracing::info!(" MCP endpoint: {scheme}://{addr}/mcp");
2419 tracing::info!(" Health check: {scheme}://{addr}/healthz");
2420 tracing::info!(" Readiness: {scheme}://{addr}/readyz");
2421}
2422
2423#[allow(
2446 clippy::too_many_arguments,
2447 clippy::cognitive_complexity,
2448 reason = "server start-up threads TLS, reload state, and graceful shutdown through one flow"
2449)]
2450async fn run_server(
2454 router: axum::Router,
2455 listener: TcpListener,
2456 tls_paths: Option<(PathBuf, PathBuf)>,
2457 tls_handshake_timeout: Duration,
2458 max_concurrent_tls_handshakes: usize,
2459 mtls_config: Option<MtlsConfig>,
2460 shutdown_timeout: Duration,
2461 auth_state: Option<Arc<AuthState>>,
2462 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
2463 mut on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
2464 ct: CancellationToken,
2465 session_ct: CancellationToken,
2466) -> anyhow::Result<()> {
2467 let shutdown_trigger = CancellationToken::new();
2471 {
2472 let trigger = shutdown_trigger.clone();
2473 let parent = ct.clone();
2474 tokio::spawn(async move {
2475 tokio::select! {
2478 () = shutdown_signal() => {}
2479 () = parent.cancelled() => {}
2480 }
2481 trigger.cancel();
2482 });
2483 }
2484
2485 let graceful = {
2486 let trigger = shutdown_trigger.clone();
2487 let ct = ct.clone();
2488 async move {
2489 trigger.cancelled().await;
2490 tracing::info!("shutting down (grace period: {shutdown_timeout:?})");
2491 ct.cancel();
2492 }
2493 };
2494
2495 let force_exit_timer = {
2496 let trigger = shutdown_trigger.clone();
2497 async move {
2498 trigger.cancelled().await;
2499 tokio::time::sleep(shutdown_timeout).await;
2500 }
2501 };
2502
2503 if let Some((cert_path, key_path)) = tls_paths {
2504 let crl_set = if let Some(mtls) = mtls_config.as_ref()
2505 && mtls.crl_enabled
2506 {
2507 let (ca_certs, roots) = load_client_auth_roots(&mtls.ca_cert_path)?;
2508 let (crl_set, discover_rx) =
2509 mtls_revocation::bootstrap_fetch(roots, &ca_certs, mtls.clone())
2510 .await
2511 .map_err(|error| anyhow::anyhow!(error.to_string()))?;
2512 tokio::spawn(mtls_revocation::run_crl_refresher(
2513 Arc::clone(&crl_set),
2514 discover_rx,
2515 ct.clone(),
2516 ));
2517 Some(crl_set)
2518 } else {
2519 None
2520 };
2521
2522 if let Some(cb) = on_reload_ready.take() {
2523 cb(ReloadHandle {
2524 auth: auth_state.clone(),
2525 rbac: Some(Arc::clone(&rbac_swap)),
2526 crl_set: crl_set.clone(),
2527 });
2528 }
2529
2530 let tls_listener = TlsListener::new(
2531 listener,
2532 &cert_path,
2533 &key_path,
2534 mtls_config.as_ref(),
2535 crl_set,
2536 tls_handshake_timeout,
2537 max_concurrent_tls_handshakes,
2538 )?;
2539 let make_svc = router.into_make_service_with_connect_info::<TlsConnInfo>();
2540 tokio::select! {
2543 result = axum::serve(tls_listener, make_svc)
2544 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2545 () = force_exit_timer => {
2546 tracing::warn!("shutdown timeout exceeded, forcing exit");
2547 session_ct.cancel();
2548 }
2549 }
2550 } else {
2551 if let Some(cb) = on_reload_ready.take() {
2552 cb(ReloadHandle {
2553 auth: auth_state,
2554 rbac: Some(rbac_swap),
2555 crl_set: None,
2556 });
2557 }
2558
2559 let make_svc = router.into_make_service_with_connect_info::<SocketAddr>();
2560 tokio::select! {
2563 result = axum::serve(listener, make_svc)
2564 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2565 () = force_exit_timer => {
2566 tracing::warn!("shutdown timeout exceeded, forcing exit");
2567 session_ct.cancel();
2568 }
2569 }
2570 }
2571
2572 Ok(())
2573}
2574
2575#[cfg(feature = "oauth")]
2584fn install_oauth_proxy_routes(
2585 router: axum::Router,
2586 server_url: &str,
2587 oauth_config: &crate::oauth::OAuthConfig,
2588 auth_state: Option<&Arc<AuthState>>,
2589 max_request_body: usize,
2590 admin_role: &str,
2591) -> Result<axum::Router, RmcpServerKitError> {
2592 let Some(ref proxy) = oauth_config.proxy else {
2593 return Ok(router);
2594 };
2595
2596 let http = crate::oauth::OauthHttpClient::with_config(oauth_config)?;
2599
2600 let proxy_router = axum::Router::new();
2606
2607 let asm = crate::oauth::authorization_server_metadata(server_url, oauth_config);
2608 let proxy_router = proxy_router.route(
2609 "/.well-known/oauth-authorization-server",
2610 axum::routing::get(move || {
2611 let m = asm.clone();
2612 async move { axum::Json(m) }
2613 }),
2614 );
2615
2616 let proxy_authorize = proxy.clone();
2617 let proxy_router = proxy_router.route(
2618 "/authorize",
2619 axum::routing::get(
2620 move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
2621 let p = proxy_authorize.clone();
2622 async move { crate::oauth::handle_authorize(&p, &query.unwrap_or_default()) }
2623 },
2624 ),
2625 );
2626
2627 let proxy_token = proxy.clone();
2628 let token_http = http.clone();
2629 let proxy_router = proxy_router.route(
2630 "/token",
2631 axum::routing::post(move |body: String| {
2632 let p = proxy_token.clone();
2633 let h = token_http.clone();
2634 async move { crate::oauth::handle_token(&h, &p, &body).await }
2635 })
2636 .layer(axum::middleware::from_fn(
2637 oauth_token_cache_headers_middleware,
2638 )),
2639 );
2640
2641 let proxy_register = proxy.clone();
2642 let proxy_router = proxy_router.route(
2643 "/register",
2644 axum::routing::post(move |axum::Json(body): axum::Json<serde_json::Value>| {
2645 let p = proxy_register;
2646 async move { axum::Json(crate::oauth::handle_register(&p, &body)) }
2647 })
2648 .layer(axum::middleware::from_fn(
2649 oauth_token_cache_headers_middleware,
2650 )),
2651 );
2652
2653 let admin_routes_enabled = proxy.expose_admin_endpoints
2654 && (proxy.introspection_url.is_some() || proxy.revocation_url.is_some());
2655 if proxy.expose_admin_endpoints
2656 && !proxy.require_auth_on_admin_endpoints
2657 && proxy.allow_unauthenticated_admin_endpoints
2658 {
2659 tracing::warn!(
2663 "OAuth introspect/revoke endpoints are unauthenticated by explicit \
2664 allow_unauthenticated_admin_endpoints opt-out; ensure an \
2665 authenticated reverse proxy fronts these routes"
2666 );
2667 }
2668
2669 let admin_router = if admin_routes_enabled {
2670 build_oauth_admin_router(proxy, http, auth_state, admin_role)?
2671 } else {
2672 axum::Router::new()
2673 };
2674
2675 let proxy_router =
2679 proxy_router
2680 .merge(admin_router)
2681 .layer(tower_http::limit::RequestBodyLimitLayer::new(
2682 max_request_body,
2683 ));
2684
2685 let router = router.merge(proxy_router);
2686
2687 tracing::info!(
2688 introspect = proxy.expose_admin_endpoints && proxy.introspection_url.is_some(),
2689 revoke = proxy.expose_admin_endpoints && proxy.revocation_url.is_some(),
2690 max_request_body,
2691 "OAuth 2.1 proxy endpoints enabled (/authorize, /token, /register)"
2692 );
2693 Ok(router)
2694}
2695
2696#[cfg(feature = "oauth")]
2702fn build_oauth_admin_router(
2703 proxy: &crate::oauth::OAuthProxyConfig,
2704 http: crate::oauth::OauthHttpClient,
2705 auth_state: Option<&Arc<AuthState>>,
2706 admin_role: &str,
2707) -> Result<axum::Router, RmcpServerKitError> {
2708 let mut admin_router = axum::Router::new();
2709 if proxy.introspection_url.is_some() {
2710 let proxy_introspect = proxy.clone();
2711 let introspect_http = http.clone();
2712 admin_router = admin_router.route(
2713 "/introspect",
2714 axum::routing::post(move |body: String| {
2715 let p = proxy_introspect.clone();
2716 let h = introspect_http.clone();
2717 async move { crate::oauth::handle_introspect(&h, &p, &body).await }
2718 }),
2719 );
2720 }
2721 if proxy.revocation_url.is_some() {
2722 let proxy_revoke = proxy.clone();
2723 let revoke_http = http;
2724 admin_router = admin_router.route(
2725 "/revoke",
2726 axum::routing::post(move |body: String| {
2727 let p = proxy_revoke.clone();
2728 let h = revoke_http.clone();
2729 async move { crate::oauth::handle_revoke(&h, &p, &body).await }
2730 }),
2731 );
2732 }
2733
2734 let admin_router = admin_router.layer(axum::middleware::from_fn(
2735 oauth_token_cache_headers_middleware,
2736 ));
2737
2738 if proxy.require_auth_on_admin_endpoints {
2739 let Some(state) = auth_state else {
2740 return Err(RmcpServerKitError::Startup(
2741 "oauth proxy admin endpoints require auth state".into(),
2742 ));
2743 };
2744 let state_for_mw = Arc::clone(state);
2745 let required_role: Arc<str> = Arc::from(admin_role);
2746 Ok(admin_router
2752 .layer(axum::middleware::from_fn(move |req, next| {
2753 let r = Arc::clone(&required_role);
2754 crate::admin::require_admin_role(r, req, next)
2755 }))
2756 .layer(axum::middleware::from_fn(move |req, next| {
2757 let s = Arc::clone(&state_for_mw);
2758 auth_middleware(s, req, next)
2759 })))
2760 } else {
2761 Ok(admin_router)
2762 }
2763}
2764
2765#[allow(
2772 deprecated,
2773 reason = "internal metadata assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
2774)]
2775fn derive_server_url(config: &McpServerConfig) -> String {
2776 config.public_url.as_ref().map_or_else(
2777 || {
2778 let scheme = if config.tls_cert_path.is_some() {
2779 "https"
2780 } else {
2781 "http"
2782 };
2783 format!("{scheme}://{}", config.bind_addr)
2784 },
2785 |url| url.trim_end_matches('/').to_owned(),
2786 )
2787}
2788
2789fn derive_allowed_hosts(bind_addr: &str, public_url: Option<&str>) -> Vec<String> {
2794 let mut hosts = vec![
2795 "localhost".to_owned(),
2796 "127.0.0.1".to_owned(),
2797 "::1".to_owned(),
2798 ];
2799
2800 if let Some(url) = public_url
2801 && let Ok(uri) = url.parse::<axum::http::Uri>()
2802 && let Some(authority) = uri.authority()
2803 {
2804 let host = authority.host().to_owned();
2805 if !hosts.iter().any(|h| h == &host) {
2806 hosts.push(host);
2807 }
2808
2809 let authority = authority.as_str().to_owned();
2810 if !hosts.iter().any(|h| h == &authority) {
2811 hosts.push(authority);
2812 }
2813 }
2814
2815 if let Ok(uri) = format!("http://{bind_addr}").parse::<axum::http::Uri>()
2816 && let Some(authority) = uri.authority()
2817 {
2818 let host = authority.host().to_owned();
2819 if !hosts.iter().any(|h| h == &host) {
2820 hosts.push(host);
2821 }
2822
2823 let authority = authority.as_str().to_owned();
2824 if !hosts.iter().any(|h| h == &authority) {
2825 hosts.push(authority);
2826 }
2827 }
2828
2829 hosts
2830}
2831
2832impl axum::extract::connect_info::Connected<axum::serve::IncomingStream<'_, TlsListener>>
2845 for TlsConnInfo
2846{
2847 fn connect_info(target: axum::serve::IncomingStream<'_, TlsListener>) -> Self {
2848 let addr = *target.remote_addr();
2849 let identity = target.io().identity().cloned();
2850 Self::new(addr, identity)
2851 }
2852}
2853
2854const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
2861
2862const DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256;
2870
2871const TLS_ACCEPT_CHANNEL_CAPACITY: usize = 32;
2876
2877struct TlsListener {
2893 local_addr: SocketAddr,
2896 rx: mpsc::Receiver<(AuthenticatedTlsStream, SocketAddr)>,
2898 acceptor_task: tokio::task::JoinHandle<()>,
2901}
2902
2903impl TlsListener {
2904 fn new(
2905 inner: TcpListener,
2906 cert_path: &Path,
2907 key_path: &Path,
2908 mtls_config: Option<&MtlsConfig>,
2909 crl_set: Option<Arc<CrlSet>>,
2910 handshake_timeout: Duration,
2911 max_concurrent_handshakes: usize,
2912 ) -> anyhow::Result<Self> {
2913 rustls::crypto::ring::default_provider()
2915 .install_default()
2916 .ok();
2917
2918 let certs = load_certs(cert_path)?;
2919 let key = load_key(key_path)?;
2920
2921 let mtls_default_role;
2922
2923 let tls_config = if let Some(mtls) = mtls_config {
2924 mtls_default_role = mtls.default_role.clone();
2925 let verifier: Arc<dyn rustls::server::danger::ClientCertVerifier> = if mtls.crl_enabled
2926 {
2927 let Some(crl_set) = crl_set else {
2928 return Err(anyhow::anyhow!(
2929 "mTLS CRL verifier requested but CRL state was not initialized"
2930 ));
2931 };
2932 Arc::new(DynamicClientCertVerifier::new(crl_set))
2933 } else {
2934 let (_, root_store) = load_client_auth_roots(&mtls.ca_cert_path)?;
2935 if mtls.required {
2936 rustls::server::WebPkiClientVerifier::builder(root_store)
2937 .build()
2938 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2939 } else {
2940 rustls::server::WebPkiClientVerifier::builder(root_store)
2941 .allow_unauthenticated()
2942 .build()
2943 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2944 }
2945 };
2946
2947 tracing::info!(
2948 ca = %mtls.ca_cert_path.display(),
2949 required = mtls.required,
2950 crl_enabled = mtls.crl_enabled,
2951 "mTLS client auth configured"
2952 );
2953
2954 rustls::ServerConfig::builder_with_protocol_versions(&[
2955 &rustls::version::TLS12,
2956 &rustls::version::TLS13,
2957 ])
2958 .with_client_cert_verifier(verifier)
2959 .with_single_cert(certs, key)?
2960 } else {
2961 mtls_default_role = "viewer".to_owned();
2962 rustls::ServerConfig::builder_with_protocol_versions(&[
2963 &rustls::version::TLS12,
2964 &rustls::version::TLS13,
2965 ])
2966 .with_no_client_auth()
2967 .with_single_cert(certs, key)?
2968 };
2969
2970 let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
2971 tracing::info!(
2972 "TLS enabled (cert: {}, key: {})",
2973 cert_path.display(),
2974 key_path.display()
2975 );
2976 let local_addr = inner.local_addr()?;
2977 let (tx, rx) = mpsc::channel(TLS_ACCEPT_CHANNEL_CAPACITY);
2978 let acceptor_task = tokio::spawn(run_tls_acceptor(
2979 inner,
2980 acceptor,
2981 mtls_default_role,
2982 tx,
2983 handshake_timeout,
2984 max_concurrent_handshakes,
2985 ));
2986 Ok(Self {
2987 local_addr,
2988 rx,
2989 acceptor_task,
2990 })
2991 }
2992
2993 fn extract_handshake_identity(
2997 tls_stream: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2998 default_role: &str,
2999 addr: SocketAddr,
3000 ) -> Option<AuthIdentity> {
3001 let (_, server_conn) = tls_stream.get_ref();
3002 let cert_der = server_conn.peer_certificates()?.first()?;
3003 let id = extract_mtls_identity(cert_der.as_ref(), default_role)?;
3004 tracing::debug!(name = %id.name, peer = %addr, "mTLS client cert accepted");
3005 Some(id)
3006 }
3007}
3008
3009async fn run_tls_acceptor(
3020 listener: TcpListener,
3021 acceptor: tokio_rustls::TlsAcceptor,
3022 default_role: String,
3023 tx: mpsc::Sender<(AuthenticatedTlsStream, SocketAddr)>,
3024 handshake_timeout: Duration,
3025 max_concurrent_handshakes: usize,
3026) {
3027 let inflight = Arc::new(Semaphore::new(max_concurrent_handshakes));
3028 loop {
3029 let Ok(permit) = Arc::clone(&inflight).acquire_owned().await else {
3033 return;
3035 };
3036 let (stream, addr) = match listener.accept().await {
3037 Ok(pair) => pair,
3038 Err(e) => {
3039 tracing::debug!("TCP accept error: {e}");
3040 continue;
3041 }
3042 };
3043 if tx.is_closed() {
3044 return;
3046 }
3047 let acceptor = acceptor.clone();
3048 let default_role = default_role.clone();
3049 let tx = tx.clone();
3050 tokio::spawn(async move {
3051 let _permit = permit;
3052 match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
3053 Ok(Ok(tls_stream)) => {
3054 let identity =
3055 TlsListener::extract_handshake_identity(&tls_stream, &default_role, addr);
3056 let wrapped = AuthenticatedTlsStream {
3057 inner: tls_stream,
3058 identity,
3059 };
3060 let _ = tx.send((wrapped, addr)).await;
3063 }
3064 Ok(Err(e)) => {
3065 tracing::debug!("TLS handshake failed from {addr}: {e}");
3066 }
3067 Err(_elapsed) => {
3068 tracing::debug!(
3069 "TLS handshake timed out from {addr} after {handshake_timeout:?}"
3070 );
3071 }
3072 }
3073 });
3074 }
3075}
3076
3077pub(crate) struct AuthenticatedTlsStream {
3089 inner: tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
3090 identity: Option<AuthIdentity>,
3091}
3092
3093impl AuthenticatedTlsStream {
3094 #[must_use]
3096 pub(crate) const fn identity(&self) -> Option<&AuthIdentity> {
3097 self.identity.as_ref()
3098 }
3099}
3100
3101impl std::fmt::Debug for AuthenticatedTlsStream {
3102 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3103 f.debug_struct("AuthenticatedTlsStream")
3104 .field("identity", &self.identity.as_ref().map(|id| &id.name))
3105 .finish_non_exhaustive()
3106 }
3107}
3108
3109impl tokio::io::AsyncRead for AuthenticatedTlsStream {
3110 fn poll_read(
3111 mut self: Pin<&mut Self>,
3112 cx: &mut std::task::Context<'_>,
3113 buf: &mut tokio::io::ReadBuf<'_>,
3114 ) -> std::task::Poll<std::io::Result<()>> {
3115 Pin::new(&mut self.inner).poll_read(cx, buf)
3116 }
3117}
3118
3119impl tokio::io::AsyncWrite for AuthenticatedTlsStream {
3120 fn poll_write(
3121 mut self: Pin<&mut Self>,
3122 cx: &mut std::task::Context<'_>,
3123 buf: &[u8],
3124 ) -> std::task::Poll<std::io::Result<usize>> {
3125 Pin::new(&mut self.inner).poll_write(cx, buf)
3126 }
3127
3128 fn poll_flush(
3129 mut self: Pin<&mut Self>,
3130 cx: &mut std::task::Context<'_>,
3131 ) -> std::task::Poll<std::io::Result<()>> {
3132 Pin::new(&mut self.inner).poll_flush(cx)
3133 }
3134
3135 fn poll_shutdown(
3136 mut self: Pin<&mut Self>,
3137 cx: &mut std::task::Context<'_>,
3138 ) -> std::task::Poll<std::io::Result<()>> {
3139 Pin::new(&mut self.inner).poll_shutdown(cx)
3140 }
3141
3142 fn poll_write_vectored(
3143 mut self: Pin<&mut Self>,
3144 cx: &mut std::task::Context<'_>,
3145 bufs: &[std::io::IoSlice<'_>],
3146 ) -> std::task::Poll<std::io::Result<usize>> {
3147 Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
3148 }
3149
3150 fn is_write_vectored(&self) -> bool {
3151 self.inner.is_write_vectored()
3152 }
3153}
3154
3155impl axum::serve::Listener for TlsListener {
3156 type Io = AuthenticatedTlsStream;
3157 type Addr = SocketAddr;
3158
3159 async fn accept(&mut self) -> (Self::Io, Self::Addr) {
3165 if let Some(pair) = self.rx.recv().await {
3166 return pair;
3167 }
3168 tracing::error!("TLS acceptor task terminated; no further connections will be accepted");
3174 std::future::pending().await
3175 }
3176
3177 fn local_addr(&self) -> std::io::Result<Self::Addr> {
3178 Ok(self.local_addr)
3179 }
3180}
3181
3182impl Drop for TlsListener {
3183 fn drop(&mut self) {
3184 self.acceptor_task.abort();
3187 }
3188}
3189
3190fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
3191 use rustls::pki_types::pem::PemObject;
3192 let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_file_iter(path)
3193 .map_err(|e| anyhow::anyhow!("failed to read certs from {}: {e}", path.display()))?
3194 .collect::<Result<_, _>>()
3195 .map_err(|e| anyhow::anyhow!("invalid cert in {}: {e}", path.display()))?;
3196 anyhow::ensure!(
3197 !certs.is_empty(),
3198 "no certificates found in {}",
3199 path.display()
3200 );
3201 Ok(certs)
3202}
3203
3204fn load_client_auth_roots(
3205 path: &Path,
3206) -> anyhow::Result<(
3207 Vec<rustls::pki_types::CertificateDer<'static>>,
3208 Arc<RootCertStore>,
3209)> {
3210 let ca_certs = load_certs(path)?;
3211 let mut root_store = RootCertStore::empty();
3212 for cert in &ca_certs {
3213 root_store
3214 .add(cert.clone())
3215 .map_err(|error| anyhow::anyhow!("invalid CA cert: {error}"))?;
3216 }
3217
3218 Ok((ca_certs, Arc::new(root_store)))
3219}
3220
3221fn load_key(path: &Path) -> anyhow::Result<rustls::pki_types::PrivateKeyDer<'static>> {
3222 use rustls::pki_types::pem::PemObject;
3223 rustls::pki_types::PrivateKeyDer::from_pem_file(path)
3224 .map_err(|e| anyhow::anyhow!("failed to read key from {}: {e}", path.display()))
3225}
3226
3227#[allow(
3229 clippy::unused_async,
3230 reason = "axum route handler signature requires `async fn` even when the body is synchronous"
3231)]
3232async fn healthz() -> impl IntoResponse {
3233 axum::Json(serde_json::json!({
3234 "status": "ok",
3235 }))
3236}
3237
3238fn version_payload(name: &str, version: &str, expose_build_metadata: bool) -> serde_json::Value {
3248 let mut map = serde_json::Map::new();
3249 map.insert("name".into(), name.into());
3250 map.insert("version".into(), version.into());
3251 map.insert(
3252 "rmcp_server_kit_version".into(),
3253 env!("CARGO_PKG_VERSION").into(),
3254 );
3255 if expose_build_metadata {
3256 map.insert(
3257 "build_git_sha".into(),
3258 option_env!("RMCP_SERVER_KIT_BUILD_SHA")
3259 .unwrap_or("unknown")
3260 .into(),
3261 );
3262 map.insert(
3263 "build_timestamp".into(),
3264 option_env!("RMCP_SERVER_KIT_BUILD_TIME")
3265 .unwrap_or("unknown")
3266 .into(),
3267 );
3268 map.insert(
3269 "rust_version".into(),
3270 option_env!("RMCP_SERVER_KIT_RUSTC_VERSION")
3271 .unwrap_or("unknown")
3272 .into(),
3273 );
3274 }
3275 serde_json::Value::Object(map)
3276}
3277
3278fn serialize_version_payload(name: &str, version: &str, expose_build_metadata: bool) -> Arc<[u8]> {
3288 let value = version_payload(name, version, expose_build_metadata);
3289 serde_json::to_vec(&value).map_or_else(|_| Arc::from(&b"{}"[..]), Arc::from)
3290}
3291
3292async fn readyz(check: ReadinessCheck) -> impl IntoResponse {
3297 let status = check().await;
3298 let ready = status
3299 .get("ready")
3300 .and_then(serde_json::Value::as_bool)
3301 .unwrap_or(false);
3302 let code = if ready {
3303 axum::http::StatusCode::OK
3304 } else {
3305 axum::http::StatusCode::SERVICE_UNAVAILABLE
3306 };
3307 (code, axum::Json(status))
3308}
3309
3310async fn shutdown_signal() {
3314 let ctrl_c = tokio::signal::ctrl_c();
3315
3316 #[cfg(unix)]
3317 {
3318 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
3319 Ok(mut term) => {
3320 tokio::select! {
3323 _ = ctrl_c => {}
3324 _ = term.recv() => {}
3325 }
3326 }
3327 Err(e) => {
3328 tracing::warn!(error = %e, "failed to register SIGTERM handler, using SIGINT only");
3329 ctrl_c.await.ok();
3330 }
3331 }
3332 }
3333
3334 #[cfg(not(unix))]
3335 {
3336 ctrl_c.await.ok();
3337 }
3338}
3339
3340#[cfg(feature = "metrics")]
3357fn metrics_labels(req: &Request<Body>) -> (&'static str, String) {
3358 let method = match *req.method() {
3359 axum::http::Method::GET => "GET",
3360 axum::http::Method::POST => "POST",
3361 axum::http::Method::PUT => "PUT",
3362 axum::http::Method::PATCH => "PATCH",
3363 axum::http::Method::DELETE => "DELETE",
3364 axum::http::Method::HEAD => "HEAD",
3365 axum::http::Method::OPTIONS => "OPTIONS",
3366 axum::http::Method::TRACE => "TRACE",
3367 axum::http::Method::CONNECT => "CONNECT",
3368 _ => "OTHER",
3371 };
3372
3373 let path = req
3374 .extensions()
3375 .get::<axum::extract::MatchedPath>()
3376 .map_or_else(
3377 || {
3378 let raw = req.uri().path();
3379 if raw == "/mcp" || raw.starts_with("/mcp/") {
3380 "/mcp".to_owned()
3381 } else {
3382 "<unmatched>".to_owned()
3383 }
3384 },
3385 |matched| matched.as_str().to_owned(),
3386 );
3387
3388 (method, path)
3389}
3390
3391#[cfg(feature = "metrics")]
3401async fn metrics_middleware(
3402 metrics: Arc<crate::metrics::McpMetrics>,
3403 mut req: Request<Body>,
3404 next: Next,
3405) -> axum::response::Response {
3406 let (method, path) = metrics_labels(&req);
3407 let start = std::time::Instant::now();
3408
3409 req.extensions_mut().insert(Arc::clone(&metrics));
3410 let response = next.run(req).await;
3411
3412 let mut status_buf = core::fmt::NumBuffer::<u16>::new();
3413 let status = response.status().as_u16().format_into(&mut status_buf);
3414 let duration = start.elapsed().as_secs_f64();
3415
3416 metrics
3417 .http_requests_total
3418 .with_label_values(&[method, &path, status])
3419 .inc();
3420 metrics
3421 .http_request_duration_seconds
3422 .with_label_values(&[method, &path])
3423 .observe(duration);
3424
3425 response
3426}
3427
3428async fn security_headers_middleware(
3442 is_tls: bool,
3443 cfg: Arc<SecurityHeadersConfig>,
3444 req: Request<Body>,
3445 next: Next,
3446) -> axum::response::Response {
3447 use axum::http::{HeaderName, header};
3448
3449 let mut resp = next.run(req).await;
3450 let headers = resp.headers_mut();
3451
3452 headers.remove(header::SERVER);
3454 headers.remove(HeaderName::from_static("x-powered-by"));
3455
3456 apply_security_header(
3457 headers,
3458 header::X_CONTENT_TYPE_OPTIONS,
3459 cfg.x_content_type_options.as_deref(),
3460 "nosniff",
3461 );
3462 apply_security_header(
3463 headers,
3464 header::X_FRAME_OPTIONS,
3465 cfg.x_frame_options.as_deref(),
3466 "deny",
3467 );
3468 apply_security_header(
3469 headers,
3470 header::CACHE_CONTROL,
3471 cfg.cache_control.as_deref(),
3472 "no-store, max-age=0",
3473 );
3474 apply_security_header(
3475 headers,
3476 header::REFERRER_POLICY,
3477 cfg.referrer_policy.as_deref(),
3478 "no-referrer",
3479 );
3480 apply_security_header(
3481 headers,
3482 HeaderName::from_static("cross-origin-opener-policy"),
3483 cfg.cross_origin_opener_policy.as_deref(),
3484 "same-origin",
3485 );
3486 apply_security_header(
3487 headers,
3488 HeaderName::from_static("cross-origin-resource-policy"),
3489 cfg.cross_origin_resource_policy.as_deref(),
3490 "same-origin",
3491 );
3492 apply_security_header(
3493 headers,
3494 HeaderName::from_static("cross-origin-embedder-policy"),
3495 cfg.cross_origin_embedder_policy.as_deref(),
3496 "require-corp",
3497 );
3498 apply_security_header(
3499 headers,
3500 HeaderName::from_static("permissions-policy"),
3501 cfg.permissions_policy.as_deref(),
3502 "accelerometer=(), camera=(), geolocation=(), microphone=()",
3503 );
3504 apply_security_header(
3505 headers,
3506 HeaderName::from_static("x-permitted-cross-domain-policies"),
3507 cfg.x_permitted_cross_domain_policies.as_deref(),
3508 "none",
3509 );
3510 apply_security_header(
3511 headers,
3512 HeaderName::from_static("content-security-policy"),
3513 cfg.content_security_policy.as_deref(),
3514 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests",
3515 );
3516 apply_security_header(
3517 headers,
3518 HeaderName::from_static("x-dns-prefetch-control"),
3519 cfg.x_dns_prefetch_control.as_deref(),
3520 "off",
3521 );
3522
3523 if is_tls {
3524 apply_security_header(
3525 headers,
3526 header::STRICT_TRANSPORT_SECURITY,
3527 cfg.strict_transport_security.as_deref(),
3528 "max-age=63072000; includeSubDomains",
3529 );
3530 }
3531
3532 resp
3533}
3534
3535fn apply_security_header(
3546 headers: &mut axum::http::HeaderMap,
3547 name: axum::http::HeaderName,
3548 override_value: Option<&str>,
3549 default: &'static str,
3550) {
3551 use axum::http::HeaderValue;
3552
3553 match override_value {
3554 None => {
3555 headers.insert(name, HeaderValue::from_static(default));
3556 }
3557 Some("") => {
3558 }
3560 Some(v) => match HeaderValue::from_str(v) {
3561 Ok(hv) => {
3562 headers.insert(name, hv);
3563 }
3564 Err(err) => {
3565 tracing::error!(
3566 header = %name,
3567 error = %err,
3568 "invalid security header override reached middleware; using default"
3569 );
3570 headers.insert(name, HeaderValue::from_static(default));
3571 }
3572 },
3573 }
3574}
3575
3576fn validate_security_headers(cfg: &SecurityHeadersConfig) -> Result<(), RmcpServerKitError> {
3587 use axum::http::HeaderValue;
3588
3589 let fields: &[(&str, Option<&str>)] = &[
3590 (
3591 "x_content_type_options",
3592 cfg.x_content_type_options.as_deref(),
3593 ),
3594 ("x_frame_options", cfg.x_frame_options.as_deref()),
3595 ("cache_control", cfg.cache_control.as_deref()),
3596 ("referrer_policy", cfg.referrer_policy.as_deref()),
3597 (
3598 "cross_origin_opener_policy",
3599 cfg.cross_origin_opener_policy.as_deref(),
3600 ),
3601 (
3602 "cross_origin_resource_policy",
3603 cfg.cross_origin_resource_policy.as_deref(),
3604 ),
3605 (
3606 "cross_origin_embedder_policy",
3607 cfg.cross_origin_embedder_policy.as_deref(),
3608 ),
3609 ("permissions_policy", cfg.permissions_policy.as_deref()),
3610 (
3611 "x_permitted_cross_domain_policies",
3612 cfg.x_permitted_cross_domain_policies.as_deref(),
3613 ),
3614 (
3615 "content_security_policy",
3616 cfg.content_security_policy.as_deref(),
3617 ),
3618 (
3619 "x_dns_prefetch_control",
3620 cfg.x_dns_prefetch_control.as_deref(),
3621 ),
3622 (
3623 "strict_transport_security",
3624 cfg.strict_transport_security.as_deref(),
3625 ),
3626 ];
3627
3628 for (field, value) in fields {
3629 let Some(v) = value else { continue };
3630 if v.is_empty() {
3631 continue;
3632 }
3633 if let Err(err) = HeaderValue::from_str(v) {
3634 return Err(RmcpServerKitError::Config(format!(
3635 "invalid security_headers.{field}: {err}"
3636 )));
3637 }
3638 }
3639
3640 if let Some(v) = cfg.strict_transport_security.as_deref()
3641 && !v.is_empty()
3642 && v.to_ascii_lowercase().contains("preload")
3643 {
3644 return Err(RmcpServerKitError::Config(format!(
3645 "invalid security_headers.strict_transport_security: {v:?} contains the `preload` directive; \
3646 HSTS preload must be opted into explicitly via a dedicated builder, not via this knob"
3647 )));
3648 }
3649
3650 Ok(())
3651}
3652
3653#[cfg(feature = "oauth")]
3668async fn oauth_token_cache_headers_middleware(
3669 req: Request<Body>,
3670 next: Next,
3671) -> axum::response::Response {
3672 use axum::http::{HeaderValue, header};
3673
3674 let mut resp = next.run(req).await;
3675 let headers = resp.headers_mut();
3676 headers.insert(header::PRAGMA, HeaderValue::from_static("no-cache"));
3677 headers.append(header::VARY, HeaderValue::from_static("Authorization"));
3678 resp
3679}
3680
3681async fn normalize_peer_addr_middleware(
3712 resolver: Option<Arc<ForwardResolver>>,
3713 mut req: Request<Body>,
3714 next: Next,
3715) -> axum::response::Response {
3716 let direct = req
3717 .extensions()
3718 .get::<ConnectInfo<SocketAddr>>()
3719 .map(|ci| ci.0);
3720 let from_tls = req
3721 .extensions()
3722 .get::<ConnectInfo<TlsConnInfo>>()
3723 .map(|ci| ci.0.addr);
3724 if let Some(addr) = direct.or(from_tls) {
3725 if direct.is_none() {
3726 req.extensions_mut().insert(ConnectInfo(addr));
3727 }
3728 req.extensions_mut().insert(PeerAddr::new(addr));
3729 let client_ip = match &resolver {
3730 Some(r) => crate::forwarded::resolve_client_ip(
3731 addr.ip(),
3732 req.headers(),
3733 &r.trusted,
3734 r.mode,
3735 r.max_scanned_entries,
3736 )
3737 .unwrap_or_else(|reason| {
3738 tracing::debug!(
3739 reason = ?reason,
3740 "forwarded-header resolution fell back to direct peer"
3741 );
3742 addr.ip()
3743 }),
3744 None => addr.ip(),
3745 };
3746 req.extensions_mut().insert(ClientIp::new(client_ip));
3747 }
3748 next.run(req).await
3749}
3750
3751fn parse_proxy_net(entry: &str) -> Option<ipnet::IpNet> {
3754 if let Ok(net) = entry.parse::<ipnet::IpNet>() {
3755 return Some(net);
3756 }
3757 entry.parse::<IpAddr>().ok().map(ipnet::IpNet::from)
3758}
3759
3760pub(crate) fn validate_trusted_proxy_entry(entry: &str) -> Result<(), String> {
3770 match parse_proxy_net(entry) {
3771 None => Err(format!(
3772 "trusted_proxies entry {entry:?} is neither a CIDR nor an IP address"
3773 )),
3774 Some(net) if net.prefix_len() == 0 => Err(format!(
3775 "trusted_proxies entry {entry:?}: prefix length 0 is forbidden (marks every peer trusted, enabling client-IP spoofing)"
3776 )),
3777 Some(_) => Ok(()),
3778 }
3779}
3780
3781pub(crate) fn limiter_client_ip(extensions: &axum::http::Extensions) -> Option<IpAddr> {
3785 if let Some(client) = extensions.get::<ClientIp>() {
3786 return Some(client.ip);
3787 }
3788 extensions
3789 .get::<ConnectInfo<SocketAddr>>()
3790 .map(|ci| ci.0.ip())
3791 .or_else(|| {
3792 extensions
3793 .get::<ConnectInfo<TlsConnInfo>>()
3794 .map(|ci| ci.0.addr.ip())
3795 })
3796}
3797
3798#[derive(Clone, PartialEq, Eq, Hash, Debug)]
3811pub(crate) enum RateLimitKey {
3812 Ip(IpAddr),
3814 Unattributed,
3816}
3817
3818impl std::fmt::Display for RateLimitKey {
3819 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3820 match self {
3821 Self::Ip(ip) => write!(f, "{ip}"),
3822 Self::Unattributed => f.write_str("unattributed"),
3823 }
3824 }
3825}
3826
3827static UNATTRIBUTED_WARNED: std::sync::atomic::AtomicBool =
3829 std::sync::atomic::AtomicBool::new(false);
3830
3831pub(crate) fn limiter_client_key(extensions: &axum::http::Extensions) -> RateLimitKey {
3842 if let Some(ip) = limiter_client_ip(extensions) {
3843 return RateLimitKey::Ip(ip);
3844 }
3845 if !UNATTRIBUTED_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
3846 tracing::warn!(
3847 "request carries no resolvable client address; rate limiting is \
3848 falling back to a single shared bucket. This indicates \
3849 rmcp-server-kit middleware composed outside serve()."
3850 );
3851 }
3852 RateLimitKey::Unattributed
3853}
3854
3855pub(crate) type ExtraRouteRateLimiter = BoundedKeyedLimiter<RateLimitKey>;
3859
3860const EXTRA_ROUTE_MAX_TRACKED_KEYS: usize = 10_000;
3866
3867const EXTRA_ROUTE_IDLE_EVICTION: Duration = Duration::from_mins(15);
3870
3871fn build_extra_route_rate_limiter_with_policy(
3878 per_minute: u32,
3879 burst: Option<u32>,
3880 key_eviction_policy: KeyEvictionPolicy,
3881 max_tracked_keys: NonZeroUsize,
3882) -> Arc<ExtraRouteRateLimiter> {
3883 let rate = std::num::NonZeroU32::new(per_minute.max(1)).unwrap_or(std::num::NonZeroU32::MIN);
3884 let mut quota = governor::Quota::per_minute(rate);
3885 if let Some(b) = burst.and_then(std::num::NonZeroU32::new) {
3886 quota = quota.allow_burst(b);
3887 }
3888 Arc::new(BoundedKeyedLimiter::new_with_policy(
3889 quota,
3890 max_tracked_keys,
3891 EXTRA_ROUTE_IDLE_EVICTION,
3892 key_eviction_policy,
3893 ))
3894}
3895
3896async fn extra_route_rate_limit_middleware(
3921 limiter: Arc<ExtraRouteRateLimiter>,
3922 exempt: Arc<std::collections::HashSet<String>>,
3923 req: Request<Body>,
3924 next: Next,
3925) -> axum::response::Response {
3926 if exempt.contains(req.uri().path()) {
3927 return next.run(req).await;
3928 }
3929 let peer_key = limiter_client_key(req.extensions());
3930 match limiter.check_key_detailed(&peer_key) {
3931 Ok(()) => {}
3932 Err(BoundedLimiterDeny::RateLimited(wait)) => {
3933 #[cfg(feature = "metrics")]
3934 crate::metrics::record_rate_limit_deny(req.extensions(), "extra_route");
3935 tracing::warn!(rate_limit_key = %peer_key, "extra route request rate limited");
3936 return RmcpServerKitError::RateLimitedFor {
3937 message: "too many requests to application routes from this source".into(),
3938 retry_after: wait,
3939 }
3940 .into_response();
3941 }
3942 Err(BoundedLimiterDeny::CapacityFull) => {
3943 tracing::warn!(
3944 rate_limit_key = %peer_key,
3945 "extra route limiter rejected unseen key because tracked-key capacity is full"
3946 );
3947 return (
3948 axum::http::StatusCode::SERVICE_UNAVAILABLE,
3949 "rate limiter capacity exhausted",
3950 )
3951 .into_response();
3952 }
3953 }
3954 next.run(req).await
3955}
3956
3957async fn origin_check_middleware(
3963 allowed: Arc<[String]>,
3964 log_request_headers: bool,
3965 req: Request<Body>,
3966 next: Next,
3967) -> axum::response::Response {
3968 let method = req.method().clone();
3969 let path = req.uri().path().to_owned();
3970
3971 log_incoming_request(&method, &path, req.headers(), log_request_headers);
3972
3973 if let Some(origin) = req.headers().get(axum::http::header::ORIGIN) {
3974 let origin_str = origin.to_str().unwrap_or("");
3975 if !allowed.iter().any(|a| a == origin_str) {
3976 tracing::warn!(
3977 origin = origin_str,
3978 %method,
3979 %path,
3980 allowed = ?&*allowed,
3981 "rejected request: Origin not allowed"
3982 );
3983 return (
3984 axum::http::StatusCode::FORBIDDEN,
3985 "Forbidden: Origin not allowed",
3986 )
3987 .into_response();
3988 }
3989 }
3990 next.run(req).await
3991}
3992
3993fn log_incoming_request(
3996 method: &axum::http::Method,
3997 path: &str,
3998 headers: &axum::http::HeaderMap,
3999 log_request_headers: bool,
4000) {
4001 if log_request_headers {
4002 tracing::debug!(
4003 %method,
4004 %path,
4005 headers = %format_request_headers_for_log(headers),
4006 "incoming request"
4007 );
4008 } else {
4009 tracing::debug!(%method, %path, "incoming request");
4010 }
4011}
4012
4013const REDACTED_LOG_HEADERS: [&str; 6] = [
4021 "authorization",
4022 "cookie",
4023 "proxy-authorization",
4024 "forwarded",
4025 "x-forwarded-for",
4026 "x-real-ip",
4027];
4028
4029fn format_request_headers_for_log(headers: &axum::http::HeaderMap) -> String {
4030 headers
4031 .iter()
4032 .map(|(k, v)| {
4033 let name = k.as_str();
4034 if REDACTED_LOG_HEADERS.contains(&name) {
4035 format!("{name}: [REDACTED]")
4036 } else {
4037 format!("{name}: {}", v.to_str().unwrap_or("<non-utf8>"))
4038 }
4039 })
4040 .collect::<Vec<_>>()
4041 .join(", ")
4042}
4043
4044#[allow(
4068 clippy::cognitive_complexity,
4069 reason = "complexity is purely tracing macro expansion (info/warn + match arms); 18 lines of straight-line code, nothing meaningful to extract"
4070)]
4071pub async fn serve_stdio<H>(handler: H) -> Result<(), RmcpServerKitError>
4072where
4073 H: ServerHandler + 'static,
4074{
4075 use rmcp::ServiceExt as _;
4076
4077 tracing::info!("stdio transport: serving on stdin/stdout");
4078 tracing::warn!("stdio mode: auth, RBAC, TLS, and Origin checks are DISABLED");
4079
4080 let transport = rmcp::transport::io::stdio();
4081
4082 let service = handler
4083 .serve(transport)
4084 .await
4085 .map_err(|e| RmcpServerKitError::Startup(format!("stdio initialize failed: {e}")))?;
4086
4087 if let Err(e) = service.waiting().await {
4088 tracing::warn!(error = %e, "stdio session ended with error");
4089 }
4090 tracing::info!("stdio session ended");
4091 Ok(())
4092}
4093
4094#[allow(
4095 deprecated,
4096 reason = "builder methods are the sanctioned transition layer for deprecated public fields"
4097)]
4098impl McpServerConfig {
4099 #[must_use]
4103 pub fn with_tls_paths(mut self, cert_path: Option<PathBuf>, key_path: Option<PathBuf>) -> Self {
4104 self.tls_cert_path = cert_path;
4105 self.tls_key_path = key_path;
4106 self
4107 }
4108
4109 #[must_use]
4113 pub fn with_tls_cert_path(mut self, cert_path: impl Into<PathBuf>) -> Self {
4114 self.tls_cert_path = Some(cert_path.into());
4115 self
4116 }
4117
4118 #[must_use]
4122 pub fn with_tls_key_path(mut self, key_path: impl Into<PathBuf>) -> Self {
4123 self.tls_key_path = Some(key_path.into());
4124 self
4125 }
4126
4127 #[must_use]
4129 pub fn with_optional_auth(mut self, auth: Option<AuthConfig>) -> Self {
4130 self.auth = auth;
4131 self
4132 }
4133
4134 #[must_use]
4136 pub fn with_optional_session_binding_secret(mut self, secret: Option<SecretString>) -> Self {
4137 self.session_binding_secret = secret;
4138 self
4139 }
4140
4141 #[must_use]
4143 pub fn with_optional_tool_rate_limit(mut self, per_minute: Option<u32>) -> Self {
4144 self.tool_rate_limit = per_minute;
4145 self
4146 }
4147
4148 #[must_use]
4150 pub fn with_optional_tool_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
4151 self.tool_rate_limit_burst = burst;
4152 self
4153 }
4154
4155 #[must_use]
4157 pub fn with_optional_extra_route_rate_limit(mut self, per_minute: Option<u32>) -> Self {
4158 self.extra_route_rate_limit = per_minute;
4159 self
4160 }
4161
4162 #[must_use]
4164 pub fn with_optional_extra_route_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
4165 self.extra_route_rate_limit_burst = burst;
4166 self
4167 }
4168
4169 #[must_use]
4171 pub fn with_optional_forwarded_header(mut self, mode: Option<ForwardedHeaderMode>) -> Self {
4172 self.forwarded_header = mode;
4173 self
4174 }
4175
4176 #[must_use]
4178 pub fn with_optional_public_url(mut self, url: Option<String>) -> Self {
4179 self.public_url = url;
4180 self
4181 }
4182
4183 #[must_use]
4187 pub fn with_compression_min_size(mut self, min_size: u16) -> Self {
4188 self.compression_min_size = min_size;
4189 self
4190 }
4191
4192 #[must_use]
4194 pub fn with_compression_enabled(mut self, enabled: bool) -> Self {
4195 self.compression_enabled = enabled;
4196 self
4197 }
4198
4199 #[must_use]
4201 pub fn with_optional_max_concurrent_requests(mut self, limit: Option<usize>) -> Self {
4202 self.max_concurrent_requests = limit;
4203 self
4204 }
4205
4206 #[must_use]
4208 pub fn with_admin_enabled(mut self, enabled: bool) -> Self {
4209 self.admin_enabled = enabled;
4210 self
4211 }
4212
4213 #[must_use]
4216 pub fn with_admin_role(mut self, role: impl Into<String>) -> Self {
4217 self.admin_role = role.into();
4218 self
4219 }
4220
4221 #[must_use]
4223 pub fn with_expose_build_metadata(mut self, enabled: bool) -> Self {
4224 self.expose_build_metadata = enabled;
4225 self
4226 }
4227}
4228
4229fn warn_security_header_overrides(cfg: &SecurityHeadersConfig) {
4230 for (field, value) in security_header_overrides(cfg) {
4231 let action = if value.is_empty() {
4232 "omitted"
4233 } else {
4234 "overridden"
4235 };
4236 tracing::warn!(
4237 security_header = field,
4238 action,
4239 "security header configured; inspect server.security_headers.<security_header>"
4240 );
4241 }
4242}
4243
4244fn security_header_overrides(
4245 cfg: &SecurityHeadersConfig,
4246) -> impl Iterator<Item = (&'static str, &str)> {
4247 [
4248 (
4249 "x_content_type_options",
4250 cfg.x_content_type_options.as_deref(),
4251 ),
4252 ("x_frame_options", cfg.x_frame_options.as_deref()),
4253 ("cache_control", cfg.cache_control.as_deref()),
4254 ("referrer_policy", cfg.referrer_policy.as_deref()),
4255 (
4256 "cross_origin_opener_policy",
4257 cfg.cross_origin_opener_policy.as_deref(),
4258 ),
4259 (
4260 "cross_origin_resource_policy",
4261 cfg.cross_origin_resource_policy.as_deref(),
4262 ),
4263 (
4264 "cross_origin_embedder_policy",
4265 cfg.cross_origin_embedder_policy.as_deref(),
4266 ),
4267 ("permissions_policy", cfg.permissions_policy.as_deref()),
4268 (
4269 "x_permitted_cross_domain_policies",
4270 cfg.x_permitted_cross_domain_policies.as_deref(),
4271 ),
4272 (
4273 "content_security_policy",
4274 cfg.content_security_policy.as_deref(),
4275 ),
4276 (
4277 "x_dns_prefetch_control",
4278 cfg.x_dns_prefetch_control.as_deref(),
4279 ),
4280 (
4281 "strict_transport_security",
4282 cfg.strict_transport_security.as_deref(),
4283 ),
4284 ]
4285 .into_iter()
4286 .filter_map(|(field, value)| value.map(|v| (field, v)))
4287}
4288
4289fn check_auth_capacity_knobs(auth: Option<&AuthConfig>) -> Result<(), RmcpServerKitError> {
4290 if let Some(auth_cfg) = auth {
4291 if let Some(rl) = &auth_cfg.rate_limit {
4292 (rl.max_attempts_per_minute != 0).ok_or_else(|| {
4293 RmcpServerKitError::Config(
4294 "auth.rate_limit.max_attempts_per_minute must be nonzero".into(),
4295 )
4296 })?;
4297 (rl.pre_auth_max_per_minute != Some(0)).ok_or_else(|| {
4302 RmcpServerKitError::Config(
4303 "auth.rate_limit.pre_auth_max_per_minute must be nonzero when set".into(),
4304 )
4305 })?;
4306 }
4307 if let Some(mtls) = &auth_cfg.mtls {
4308 check_mtls_capacity_knobs(mtls)?;
4309 }
4310 auth_cfg.check_oauth_feature()?;
4311 }
4312 Ok(())
4313}
4314
4315fn check_mtls_capacity_knobs(mtls: &MtlsConfig) -> Result<(), RmcpServerKitError> {
4316 (mtls.crl_max_concurrent_fetches != 0).ok_or_else(|| {
4317 RmcpServerKitError::Config("auth.mtls.crl_max_concurrent_fetches must be nonzero".into())
4318 })?;
4319 (mtls.crl_discovery_rate_per_min != 0).ok_or_else(|| {
4320 RmcpServerKitError::Config("auth.mtls.crl_discovery_rate_per_min must be nonzero".into())
4321 })?;
4322 (mtls.crl_max_host_semaphores != 0).ok_or_else(|| {
4323 RmcpServerKitError::Config("auth.mtls.crl_max_host_semaphores must be nonzero".into())
4324 })?;
4325 (mtls.crl_max_seen_urls != 0).ok_or_else(|| {
4326 RmcpServerKitError::Config("auth.mtls.crl_max_seen_urls must be nonzero".into())
4327 })?;
4328 (mtls.crl_max_cache_entries != 0).ok_or_else(|| {
4329 RmcpServerKitError::Config("auth.mtls.crl_max_cache_entries must be nonzero".into())
4330 })?;
4331 (mtls.crl_max_response_bytes != 0).ok_or_else(|| {
4336 RmcpServerKitError::Config("auth.mtls.crl_max_response_bytes must be nonzero".into())
4337 })?;
4338 Ok(())
4339}
4340
4341#[cfg(test)]
4342mod tests {
4343 #![allow(
4344 clippy::unwrap_used,
4345 clippy::expect_used,
4346 clippy::panic,
4347 clippy::indexing_slicing,
4348 clippy::unwrap_in_result,
4349 clippy::print_stdout,
4350 clippy::print_stderr,
4351 deprecated,
4352 reason = "internal unit tests legitimately read/write the deprecated `pub` fields they were designed to verify"
4353 )]
4354 use std::{sync::Arc, time::Duration};
4355
4356 use axum::{
4357 body::Body,
4358 http::{Request, StatusCode, header},
4359 response::IntoResponse,
4360 };
4361 use http_body_util::BodyExt;
4362 use tower::ServiceExt as _;
4363
4364 use super::*;
4365
4366 #[tokio::test]
4369 async fn external_shutdown_bridge_exits_when_internal_token_cancels() {
4370 let external = CancellationToken::new();
4371 let internal = CancellationToken::new();
4372 let bridge = spawn_external_shutdown_bridge(external.clone(), internal.clone());
4373
4374 internal.cancel();
4377
4378 let joined = tokio::time::timeout(Duration::from_secs(2), bridge).await;
4379 assert!(
4380 joined.is_ok(),
4381 "bridge task must exit once the internal token is cancelled, \
4382 otherwise it leaks for the lifetime of the process"
4383 );
4384 }
4385
4386 #[tokio::test]
4387 async fn external_shutdown_bridge_still_forwards_external_cancel() {
4388 let external = CancellationToken::new();
4389 let internal = CancellationToken::new();
4390 let bridge = spawn_external_shutdown_bridge(external.clone(), internal.clone());
4391
4392 external.cancel();
4393
4394 let joined = tokio::time::timeout(Duration::from_secs(2), bridge).await;
4395 assert!(joined.is_ok(), "bridge task must exit on external cancel");
4396 assert!(
4397 internal.is_cancelled(),
4398 "external cancellation must still propagate to the internal token"
4399 );
4400 }
4401
4402 #[test]
4403 fn cancel_on_drop_cancels_its_token() {
4404 let ct = CancellationToken::new();
4405 {
4406 let _guard = CancelOnDrop(ct.clone());
4407 assert!(!ct.is_cancelled());
4408 }
4409 assert!(
4410 ct.is_cancelled(),
4411 "dropping the guard must cancel background startup tasks"
4412 );
4413 }
4414
4415 #[test]
4416 fn validate_rejects_mtls_without_tls() {
4417 for (cert, key) in [
4418 (None, None),
4419 (Some("cert.pem"), None),
4420 (None, Some("key.pem")),
4421 ] {
4422 let mut auth = AuthConfig::with_keys(vec![]);
4423 auth.mtls = Some(valid_mtls_config());
4424 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4425 cfg.tls_cert_path = cert.map(Into::into);
4426 cfg.tls_key_path = key.map(Into::into);
4427
4428 let err = cfg
4429 .validate()
4430 .expect_err("mTLS without both TLS paths must be rejected");
4431 let msg = err.to_string();
4432 assert!(
4433 msg.contains("tls_cert_path") && msg.contains("tls_key_path"),
4434 "cert={cert:?} key={key:?}: {msg}"
4435 );
4436 }
4437 }
4438
4439 #[test]
4440 fn validate_accepts_mtls_with_tls() {
4441 let mut auth = AuthConfig::with_keys(vec![]);
4442 auth.mtls = Some(valid_mtls_config());
4443 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4444 cfg.tls_cert_path = Some("cert.pem".into());
4445 cfg.tls_key_path = Some("key.pem".into());
4446
4447 assert!(cfg.validate().is_ok(), "mTLS with both TLS paths is valid");
4448 }
4449
4450 #[test]
4453 fn server_config_new_defaults() {
4454 let cfg = McpServerConfig::new("0.0.0.0:8443", "test-server", "1.0.0");
4455 assert_eq!(cfg.bind_addr, "0.0.0.0:8443");
4456 assert_eq!(cfg.name, "test-server");
4457 assert_eq!(cfg.version, "1.0.0");
4458 assert!(cfg.tls_cert_path.is_none());
4459 assert!(cfg.tls_key_path.is_none());
4460 assert!(cfg.auth.is_none());
4461 assert!(cfg.rbac.is_none());
4462 assert!(cfg.allowed_origins.is_empty());
4463 assert!(cfg.tool_rate_limit.is_none());
4464 assert!(cfg.readiness_check.is_none());
4465 assert_eq!(cfg.max_request_body, 1024 * 1024);
4466 assert_eq!(cfg.request_timeout, Duration::from_mins(2));
4467 assert_eq!(cfg.shutdown_timeout, Duration::from_secs(30));
4468 assert!(!cfg.log_request_headers);
4469 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(10));
4470 assert_eq!(cfg.max_concurrent_tls_handshakes, 256);
4471 assert!(cfg.session_store.is_none());
4472 assert!(cfg.session_binding_secret.is_none());
4473 }
4474
4475 #[derive(Default)]
4476 struct TestSessionStore;
4477
4478 #[async_trait::async_trait]
4479 impl SessionStore for TestSessionStore {
4480 async fn load(
4481 &self,
4482 _session_id: &str,
4483 ) -> Result<
4484 Option<rmcp::transport::streamable_http_server::session::SessionState>,
4485 rmcp::transport::streamable_http_server::session::SessionStoreError,
4486 > {
4487 Ok(None)
4488 }
4489
4490 async fn store(
4491 &self,
4492 _session_id: &str,
4493 _state: &rmcp::transport::streamable_http_server::session::SessionState,
4494 ) -> Result<(), rmcp::transport::streamable_http_server::session::SessionStoreError>
4495 {
4496 Ok(())
4497 }
4498
4499 async fn delete(
4500 &self,
4501 _session_id: &str,
4502 ) -> Result<(), rmcp::transport::streamable_http_server::session::SessionStoreError>
4503 {
4504 Ok(())
4505 }
4506 }
4507
4508 fn test_session_store() -> Arc<dyn SessionStore> {
4509 Arc::new(TestSessionStore)
4510 }
4511
4512 fn shared_session_binding_secret() -> SecretString {
4513 SecretString::from("0123456789abcdef0123456789abcdef")
4514 }
4515
4516 #[test]
4517 fn session_store_defaults_to_none() {
4518 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4519
4520 assert!(cfg.session_store.is_none());
4521 }
4522
4523 #[test]
4524 fn validate_rejects_session_store_without_binding_secret() {
4525 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4526 .with_auth(AuthConfig::with_keys(vec![]))
4527 .with_session_store(test_session_store());
4528
4529 let err = cfg
4530 .validate()
4531 .expect_err("authenticated shared-store binding needs a shared secret");
4532 let msg = err.to_string();
4533 assert!(msg.contains("session_store"), "{msg}");
4534 assert!(msg.contains("session_binding"), "{msg}");
4535 assert!(msg.contains("shared secret"), "{msg}");
4536 }
4537
4538 #[test]
4539 fn validate_allows_session_store_with_binding_secret() {
4540 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4541 .with_auth(AuthConfig::with_keys(vec![]))
4542 .with_session_store(test_session_store())
4543 .with_session_binding_secret(shared_session_binding_secret());
4544
4545 assert!(cfg.validate().is_ok());
4546 }
4547
4548 #[test]
4549 fn validate_allows_session_store_when_binding_disabled() {
4550 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4551 .with_auth(AuthConfig::with_keys(vec![]))
4552 .with_session_binding(false)
4553 .with_session_store(test_session_store());
4554
4555 assert!(cfg.validate().is_ok());
4556 }
4557
4558 #[test]
4559 fn validate_allows_binding_secret_without_session_store() {
4560 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4561 .with_auth(AuthConfig::with_keys(vec![]))
4562 .with_session_binding_secret(shared_session_binding_secret());
4563
4564 assert!(cfg.validate().is_ok());
4565 }
4566
4567 #[test]
4568 fn tls_handshake_builders_set_fields() {
4569 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4570 .with_tls_handshake_timeout(Duration::from_secs(3))
4571 .with_max_concurrent_tls_handshakes(64);
4572 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(3));
4573 assert_eq!(cfg.max_concurrent_tls_handshakes, 64);
4574 }
4575
4576 #[test]
4577 fn validate_rejects_zero_tls_handshake_timeout() {
4578 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4579 .with_tls_handshake_timeout(Duration::ZERO);
4580 let err = cfg.validate().expect_err("zero handshake timeout");
4581 assert!(err.to_string().contains("tls_handshake_timeout"));
4582 }
4583
4584 #[test]
4585 fn validate_rejects_zero_max_concurrent_tls_handshakes() {
4586 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4587 .with_max_concurrent_tls_handshakes(0);
4588 let err = cfg.validate().expect_err("zero handshake concurrency");
4589 assert!(err.to_string().contains("max_concurrent_tls_handshakes"));
4590 }
4591
4592 #[test]
4593 fn validate_consumes_and_proves() {
4594 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4596 let validated = cfg.validate().expect("valid config");
4597 assert_eq!(validated.as_inner().name, "test-server");
4599 let raw = validated.into_inner();
4601 assert_eq!(raw.name, "test-server");
4602
4603 let mut bad = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
4605 bad.max_request_body = 0;
4606 assert!(bad.validate().is_err(), "zero body cap must fail validate");
4607 }
4608
4609 #[test]
4610 fn validate_rejects_zero_max_concurrent_requests() {
4611 let cfg =
4612 McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_max_concurrent_requests(0);
4613 let err = cfg.validate().expect_err("zero concurrency cap must fail");
4614 assert!(
4615 format!("{err}").contains("max_concurrent_requests"),
4616 "error should mention max_concurrent_requests, got: {err}"
4617 );
4618 }
4619
4620 #[test]
4621 fn validate_rejects_zero_max_tracked_keys() {
4622 let rl = crate::auth::RateLimitConfig {
4625 max_attempts_per_minute: 30,
4626 pre_auth_max_per_minute: None,
4627 max_tracked_keys: 0,
4628 idle_eviction: Duration::from_secs(15 * 60),
4629 burst: None,
4630 pre_auth_burst: None,
4631 key_eviction_policy: KeyEvictionPolicy::default(),
4632 };
4633 let auth_cfg = AuthConfig {
4634 enabled: true,
4635 api_keys: Vec::new(),
4636 mtls: None,
4637 rate_limit: Some(rl),
4638 #[cfg(feature = "oauth")]
4639 oauth: None,
4640 #[cfg(not(feature = "oauth"))]
4641 oauth: None,
4642 };
4643 let cfg = McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_auth(auth_cfg);
4644 let err = cfg.validate().expect_err("zero max_tracked_keys must fail");
4645 assert!(
4646 format!("{err}").contains("max_tracked_keys"),
4647 "error should mention max_tracked_keys, got: {err}"
4648 );
4649 }
4650
4651 #[test]
4652 fn derive_allowed_hosts_includes_public_host() {
4653 let hosts = derive_allowed_hosts("0.0.0.0:8080", Some("https://mcp.example.com/mcp"));
4654 assert!(
4655 hosts.iter().any(|h| h == "mcp.example.com"),
4656 "public_url host must be allowed"
4657 );
4658 }
4659
4660 #[test]
4661 fn derive_allowed_hosts_includes_bind_authority() {
4662 let hosts = derive_allowed_hosts("127.0.0.1:8080", None);
4663 assert!(
4664 hosts.iter().any(|h| h == "127.0.0.1"),
4665 "bind host must be allowed"
4666 );
4667 assert!(
4668 hosts.iter().any(|h| h == "127.0.0.1:8080"),
4669 "bind authority must be allowed"
4670 );
4671 }
4672
4673 #[tokio::test]
4676 async fn healthz_returns_ok_json() {
4677 let resp = healthz().await.into_response();
4678 assert_eq!(resp.status(), StatusCode::OK);
4679 let body = resp.into_body().collect().await.unwrap().to_bytes();
4680 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4681 assert_eq!(json["status"], "ok");
4682 assert!(
4683 json.get("name").is_none(),
4684 "healthz must not expose server name"
4685 );
4686 assert!(
4687 json.get("version").is_none(),
4688 "healthz must not expose version"
4689 );
4690 }
4691
4692 #[tokio::test]
4695 async fn readyz_returns_ok_when_ready() {
4696 let check: ReadinessCheck =
4697 Arc::new(|| Box::pin(async { serde_json::json!({"ready": true, "db": "connected"}) }));
4698 let resp = readyz(check).await.into_response();
4699 assert_eq!(resp.status(), StatusCode::OK);
4700 let body = resp.into_body().collect().await.unwrap().to_bytes();
4701 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4702 assert_eq!(json["ready"], true);
4703 assert!(
4704 json.get("name").is_none(),
4705 "readyz must not expose server name"
4706 );
4707 assert!(
4708 json.get("version").is_none(),
4709 "readyz must not expose version"
4710 );
4711 assert_eq!(json["db"], "connected");
4712 }
4713
4714 #[tokio::test]
4715 async fn readyz_returns_503_when_not_ready() {
4716 let check: ReadinessCheck =
4717 Arc::new(|| Box::pin(async { serde_json::json!({"ready": false}) }));
4718 let resp = readyz(check).await.into_response();
4719 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4720 }
4721
4722 #[tokio::test]
4723 async fn readyz_returns_503_when_ready_missing() {
4724 let check: ReadinessCheck =
4725 Arc::new(|| Box::pin(async { serde_json::json!({"status": "starting"}) }));
4726 let resp = readyz(check).await.into_response();
4727 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4729 }
4730
4731 fn peer_probe_router() -> axum::Router {
4736 async fn probe(req: Request<Body>) -> String {
4737 let ci = req
4738 .extensions()
4739 .get::<ConnectInfo<SocketAddr>>()
4740 .map(|c| c.0.to_string())
4741 .unwrap_or_default();
4742 let pa = req
4743 .extensions()
4744 .get::<PeerAddr>()
4745 .map(|p| p.addr.to_string())
4746 .unwrap_or_default();
4747 format!("{ci}|{pa}")
4748 }
4749 axum::Router::new()
4750 .route("/probe", axum::routing::get(probe))
4751 .layer(axum::middleware::from_fn(|req, next| {
4752 normalize_peer_addr_middleware(None, req, next)
4753 }))
4754 }
4755
4756 async fn body_string(resp: axum::response::Response) -> String {
4757 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
4758 String::from_utf8(bytes.to_vec()).unwrap()
4759 }
4760
4761 #[tokio::test]
4762 async fn normalize_preserves_existing_connect_info_and_mirrors_peer_addr() {
4763 let plain: SocketAddr = "10.0.0.1:1111".parse().unwrap();
4766 let tls: SocketAddr = "10.0.0.2:2222".parse().unwrap();
4767 let req = Request::builder()
4768 .uri("/probe")
4769 .extension(ConnectInfo(plain))
4770 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4771 .body(Body::empty())
4772 .unwrap();
4773 let resp = peer_probe_router().oneshot(req).await.unwrap();
4774 assert_eq!(resp.status(), StatusCode::OK);
4775 assert_eq!(body_string(resp).await, format!("{plain}|{plain}"));
4776 }
4777
4778 #[tokio::test]
4779 async fn normalize_inserts_connect_info_and_peer_addr_from_tls() {
4780 let tls: SocketAddr = "192.168.1.7:50443".parse().unwrap();
4781 let req = Request::builder()
4782 .uri("/probe")
4783 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4784 .body(Body::empty())
4785 .unwrap();
4786 let resp = peer_probe_router().oneshot(req).await.unwrap();
4787 assert_eq!(resp.status(), StatusCode::OK);
4788 assert_eq!(body_string(resp).await, format!("{tls}|{tls}"));
4789 }
4790
4791 #[tokio::test]
4792 async fn normalize_no_op_without_any_connect_info() {
4793 let req = Request::builder()
4794 .uri("/probe")
4795 .body(Body::empty())
4796 .unwrap();
4797 let resp = peer_probe_router().oneshot(req).await.unwrap();
4798 assert_eq!(resp.status(), StatusCode::OK);
4799 assert_eq!(body_string(resp).await, "|");
4800 }
4801
4802 #[tokio::test]
4803 async fn peer_addr_extractor_rejects_when_absent() {
4804 async fn h(peer: PeerAddr) -> String {
4805 peer.addr.to_string()
4806 }
4807 let app = axum::Router::new().route("/p", axum::routing::get(h));
4808 let req = Request::builder().uri("/p").body(Body::empty()).unwrap();
4809 let resp = app.oneshot(req).await.unwrap();
4810 assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
4811 }
4812
4813 #[tokio::test]
4814 async fn peer_addr_extractor_returns_value_when_present() {
4815 async fn h(peer: PeerAddr) -> String {
4816 peer.addr.to_string()
4817 }
4818 let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap();
4819 let app = axum::Router::new().route("/p", axum::routing::get(h));
4820 let req = Request::builder()
4821 .uri("/p")
4822 .extension(PeerAddr::new(addr))
4823 .body(Body::empty())
4824 .unwrap();
4825 let resp = app.oneshot(req).await.unwrap();
4826 assert_eq!(resp.status(), StatusCode::OK);
4827 assert_eq!(body_string(resp).await, addr.to_string());
4828 }
4829
4830 #[tokio::test]
4831 async fn peer_addr_via_extension_extractor() {
4832 async fn h(axum::Extension(peer): axum::Extension<PeerAddr>) -> String {
4833 peer.addr.to_string()
4834 }
4835 let addr: SocketAddr = "127.0.0.1:4242".parse().unwrap();
4836 let app = axum::Router::new().route("/p", axum::routing::get(h));
4837 let req = Request::builder()
4838 .uri("/p")
4839 .extension(PeerAddr::new(addr))
4840 .body(Body::empty())
4841 .unwrap();
4842 let resp = app.oneshot(req).await.unwrap();
4843 assert_eq!(resp.status(), StatusCode::OK);
4844 assert_eq!(body_string(resp).await, addr.to_string());
4845 }
4846
4847 fn limited_router(per_minute: u32) -> axum::Router {
4852 limited_router_with_burst(per_minute, None)
4853 }
4854
4855 fn limited_router_with_burst(per_minute: u32, burst: Option<u32>) -> axum::Router {
4857 limited_router_full(per_minute, burst, &[])
4858 }
4859
4860 fn limited_router_full(
4864 per_minute: u32,
4865 burst: Option<u32>,
4866 exempt_paths: &[&str],
4867 ) -> axum::Router {
4868 let limiter = build_extra_route_rate_limiter_with_policy(
4869 per_minute,
4870 burst,
4871 KeyEvictionPolicy::default(),
4872 NonZeroUsize::new(EXTRA_ROUTE_MAX_TRACKED_KEYS).unwrap_or(NonZeroUsize::MIN),
4873 );
4874 let exempt: Arc<std::collections::HashSet<String>> =
4875 Arc::new(exempt_paths.iter().map(|s| (*s).to_owned()).collect());
4876 axum::Router::new()
4877 .route("/limited", axum::routing::get(|| async { "ok" }))
4878 .route("/exempt", axum::routing::get(|| async { "ok" }))
4879 .layer(axum::middleware::from_fn(move |req, next| {
4880 let l = Arc::clone(&limiter);
4881 let e = Arc::clone(&exempt);
4882 extra_route_rate_limit_middleware(l, e, req, next)
4883 }))
4884 }
4885
4886 fn limited_req(ip: &str) -> Request<Body> {
4887 limited_req_to(ip, "/limited")
4888 }
4889
4890 fn limited_req_to(ip: &str, path: &str) -> Request<Body> {
4891 let addr: SocketAddr = format!("{ip}:40000").parse().unwrap();
4892 Request::builder()
4893 .uri(path)
4894 .extension(ConnectInfo(addr))
4895 .body(Body::empty())
4896 .unwrap()
4897 }
4898
4899 #[tokio::test]
4900 async fn extra_route_limiter_denies_over_quota() {
4901 let app = limited_router(2);
4902 for i in 0..2 {
4903 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4904 assert_eq!(resp.status(), StatusCode::OK, "request {i} should pass");
4905 }
4906 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4907 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4908 let body = body_string(resp).await;
4909 assert!(
4910 body.contains("too many requests to application routes"),
4911 "deny body should match the limiter message, got: {body}"
4912 );
4913 }
4914
4915 fn one_tracked_key() -> NonZeroUsize {
4916 NonZeroUsize::new(1).unwrap_or(NonZeroUsize::MIN)
4917 }
4918
4919 #[tokio::test]
4920 async fn extra_route_limiter_capacity_full_returns_503_without_retry_after() {
4921 let limiter = build_extra_route_rate_limiter_with_policy(
4922 10,
4923 None,
4924 KeyEvictionPolicy::RejectNew,
4925 one_tracked_key(),
4926 );
4927 let exempt = Arc::new(std::collections::HashSet::new());
4928 let app = axum::Router::new()
4929 .route("/limited", axum::routing::get(|| async { "ok" }))
4930 .layer(axum::middleware::from_fn(move |req, next| {
4931 let l = Arc::clone(&limiter);
4932 let e = Arc::clone(&exempt);
4933 extra_route_rate_limit_middleware(l, e, req, next)
4934 }));
4935 let established = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4936 assert_eq!(established.status(), StatusCode::OK);
4937
4938 let denied = app.clone().oneshot(limited_req("10.1.1.2")).await.unwrap();
4939
4940 assert_eq!(denied.status(), StatusCode::SERVICE_UNAVAILABLE);
4941 assert!(denied.headers().get(header::RETRY_AFTER).is_none());
4942 }
4943
4944 #[tokio::test]
4945 async fn extra_route_limiter_isolates_keys() {
4946 let app = limited_router(2);
4947 for _ in 0..2 {
4948 let resp = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4949 assert_eq!(resp.status(), StatusCode::OK);
4950 }
4951 let exhausted = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4952 assert_eq!(exhausted.status(), StatusCode::TOO_MANY_REQUESTS);
4953 let other = app.clone().oneshot(limited_req("10.3.3.3")).await.unwrap();
4955 assert_eq!(other.status(), StatusCode::OK);
4956 }
4957
4958 #[tokio::test]
4959 async fn extra_route_limiter_bounds_requests_without_peer() {
4960 let app = limited_router(1);
4964 let mk = || {
4965 Request::builder()
4966 .uri("/limited")
4967 .body(Body::empty())
4968 .unwrap()
4969 };
4970 let first = app.clone().oneshot(mk()).await.unwrap();
4971 assert_eq!(
4972 first.status(),
4973 StatusCode::OK,
4974 "first request consumes quota"
4975 );
4976 let second = app.clone().oneshot(mk()).await.unwrap();
4977 assert_eq!(
4978 second.status(),
4979 StatusCode::TOO_MANY_REQUESTS,
4980 "unattributable requests must share a bounded bucket, not bypass the limiter"
4981 );
4982 }
4983
4984 #[test]
4985 fn limiter_client_key_falls_back_to_unattributed() {
4986 let empty = axum::http::Extensions::new();
4987 assert_eq!(limiter_client_key(&empty), RateLimitKey::Unattributed);
4988 }
4989
4990 #[test]
4991 fn unattributed_key_is_distinct_from_unspecified_ip() {
4992 let unspecified = RateLimitKey::Ip("0.0.0.0".parse::<IpAddr>().unwrap());
4996 assert_ne!(unspecified, RateLimitKey::Unattributed);
4997
4998 let mut set = std::collections::HashSet::new();
4999 set.insert(unspecified);
5000 set.insert(RateLimitKey::Unattributed);
5001 assert_eq!(set.len(), 2, "the two keys must hash to distinct buckets");
5002 }
5003
5004 #[test]
5005 fn rate_limit_key_display_does_not_fabricate_an_ip() {
5006 assert_eq!(
5007 RateLimitKey::Ip("10.1.2.3".parse::<IpAddr>().unwrap()).to_string(),
5008 "10.1.2.3"
5009 );
5010 assert_eq!(RateLimitKey::Unattributed.to_string(), "unattributed");
5011 }
5012
5013 #[tokio::test]
5014 async fn extra_route_limiter_extracts_tls_conn_info() {
5015 let app = limited_router(2);
5016 let mk = || {
5017 let addr: SocketAddr = "192.168.9.9:55555".parse().unwrap();
5018 Request::builder()
5019 .uri("/limited")
5020 .extension(ConnectInfo(TlsConnInfo::new(addr, None)))
5021 .body(Body::empty())
5022 .unwrap()
5023 };
5024 for _ in 0..2 {
5025 assert_eq!(
5026 app.clone().oneshot(mk()).await.unwrap().status(),
5027 StatusCode::OK
5028 );
5029 }
5030 let resp = app.clone().oneshot(mk()).await.unwrap();
5031 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
5032 }
5033
5034 #[tokio::test]
5035 async fn extra_route_limiter_exempt_path_bypasses_quota() {
5036 let app = limited_router_full(1, None, &["/exempt"]);
5039 for i in 0..5 {
5040 let resp = app
5041 .clone()
5042 .oneshot(limited_req_to("10.6.6.6", "/exempt"))
5043 .await
5044 .unwrap();
5045 assert_eq!(resp.status(), StatusCode::OK, "exempt request {i}");
5046 }
5047 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
5049 assert_eq!(resp.status(), StatusCode::OK);
5050 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
5052 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
5053 }
5054
5055 #[tokio::test]
5056 async fn extra_route_limiter_exemption_is_raw_exact_match() {
5057 let app = limited_router_full(1, None, &["/exempt"]);
5060 let ok = app
5061 .clone()
5062 .oneshot(limited_req_to("10.7.7.7", "/exempt/"))
5063 .await
5064 .unwrap();
5065 assert_eq!(
5066 ok.status(),
5067 StatusCode::NOT_FOUND,
5068 "variant path routes 404"
5069 );
5070 let denied = app
5072 .clone()
5073 .oneshot(limited_req_to("10.7.7.7", "/limited"))
5074 .await
5075 .unwrap();
5076 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
5077 }
5078
5079 #[cfg(feature = "metrics")]
5080 #[tokio::test]
5081 async fn extra_route_limiter_deny_increments_counter_exempt_does_not() {
5082 let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
5083 let app = limited_router_full(1, None, &["/exempt"]);
5084 let mk = |path: &str| {
5085 let addr: SocketAddr = "10.8.8.8:40000".parse().unwrap();
5086 Request::builder()
5087 .uri(path)
5088 .extension(ConnectInfo(addr))
5089 .extension(Arc::clone(&metrics))
5090 .body(Body::empty())
5091 .unwrap()
5092 };
5093 let counter = || {
5094 metrics
5095 .rate_limited_total
5096 .with_label_values(&["extra_route"])
5097 .get()
5098 };
5099 for _ in 0..3 {
5101 assert_eq!(
5102 app.clone().oneshot(mk("/exempt")).await.unwrap().status(),
5103 StatusCode::OK
5104 );
5105 }
5106 assert_eq!(counter(), 0, "exempt requests must not count as denies");
5107 assert_eq!(
5109 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
5110 StatusCode::OK
5111 );
5112 assert_eq!(counter(), 0);
5113 assert_eq!(
5114 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
5115 StatusCode::TOO_MANY_REQUESTS
5116 );
5117 assert_eq!(counter(), 1, "deny must increment the extra_route label");
5118 }
5119
5120 #[test]
5121 fn validate_rejects_exempt_paths_without_base_knob() {
5122 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5123 .with_extra_route_rate_limit_exempt_paths(["/ok"]);
5124 let err = cfg.validate().expect_err("exempt paths without rate limit");
5125 assert!(err.to_string().contains("requires extra_route_rate_limit"));
5126 }
5127
5128 #[test]
5129 fn validate_rejects_malformed_exempt_paths() {
5130 for bad in ["", "no-slash"] {
5131 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5132 .with_extra_route_rate_limit(10)
5133 .with_extra_route_rate_limit_exempt_paths([bad]);
5134 let err = cfg.validate().expect_err("malformed exempt path");
5135 assert!(
5136 err.to_string()
5137 .contains("must be non-empty and start with '/'"),
5138 "entry {bad:?}: {err}"
5139 );
5140 }
5141 }
5142
5143 #[test]
5144 fn validate_accepts_wellformed_exempt_paths() {
5145 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5146 .with_extra_route_rate_limit(10)
5147 .with_extra_route_rate_limit_exempt_paths(["/.well-known/oauth-authorization-server"]);
5148 assert!(cfg.validate().is_ok());
5149 }
5150
5151 #[test]
5152 fn validate_rejects_zero_extra_route_rate_limit() {
5153 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
5154 .with_extra_route_rate_limit(0);
5155 let err = cfg.validate().expect_err("zero extra route rate limit");
5156 assert!(err.to_string().contains("extra_route_rate_limit"));
5157 }
5158
5159 #[tokio::test]
5160 async fn extra_route_limiter_burst_allows_initial_spike() {
5161 let app = limited_router_with_burst(1, Some(3));
5162 for i in 0..3 {
5163 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
5164 assert_eq!(resp.status(), StatusCode::OK, "burst request {i}");
5165 }
5166 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
5167 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
5168 }
5169
5170 #[tokio::test]
5171 async fn extra_route_limiter_deny_sets_retry_after() {
5172 let app = limited_router(1);
5173 let ok = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
5174 assert_eq!(ok.status(), StatusCode::OK);
5175 let denied = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
5176 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
5177 let retry_after = denied
5178 .headers()
5179 .get(header::RETRY_AFTER)
5180 .expect("Retry-After present")
5181 .to_str()
5182 .unwrap()
5183 .parse::<u64>()
5184 .unwrap();
5185 assert!(retry_after >= 1, "delta-seconds must be >= 1");
5186 }
5187
5188 #[test]
5189 fn validate_rejects_zero_burst_knobs() {
5190 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5191 .with_tool_rate_limit(10)
5192 .with_tool_rate_limit_burst(0)
5193 .validate()
5194 .expect_err("zero tool burst");
5195 assert!(err.to_string().contains("tool_rate_limit_burst"));
5196
5197 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5198 .with_extra_route_rate_limit(10)
5199 .with_extra_route_rate_limit_burst(0)
5200 .validate()
5201 .expect_err("zero extra route burst");
5202 assert!(err.to_string().contains("extra_route_rate_limit_burst"));
5203 }
5204
5205 #[test]
5206 fn validate_rejects_orphan_burst_knobs() {
5207 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5208 .with_tool_rate_limit_burst(5)
5209 .validate()
5210 .expect_err("orphan tool burst");
5211 assert!(err.to_string().contains("requires tool_rate_limit"));
5212
5213 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5214 .with_extra_route_rate_limit_burst(5)
5215 .validate()
5216 .expect_err("orphan extra route burst");
5217 assert!(err.to_string().contains("requires extra_route_rate_limit"));
5218 }
5219
5220 #[test]
5221 fn validate_rejects_zero_auth_bursts() {
5222 let auth = AuthConfig::with_keys(vec![])
5223 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_burst(0));
5224 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5225 .with_auth(auth)
5226 .validate()
5227 .expect_err("zero auth burst");
5228 assert!(err.to_string().contains("rate_limit.burst"));
5229
5230 let auth = AuthConfig::with_keys(vec![])
5231 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(0));
5232 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5233 .with_auth(auth)
5234 .validate()
5235 .expect_err("zero pre-auth burst");
5236 assert!(err.to_string().contains("pre_auth_burst"));
5237 }
5238
5239 #[test]
5240 fn validate_rejects_zero_pre_auth_max_per_minute() {
5241 let auth = AuthConfig::with_keys(vec![])
5242 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_max_per_minute(0));
5243 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5244 .with_auth(auth)
5245 .validate()
5246 .expect_err("zero pre-auth rate");
5247 assert!(err.to_string().contains("pre_auth_max_per_minute"));
5248 }
5249
5250 fn valid_mtls_config() -> MtlsConfig {
5251 MtlsConfig {
5252 ca_cert_path: "memory://ca.pem".into(),
5253 required: true,
5254 default_role: "viewer".into(),
5255 crl_enabled: true,
5256 crl_refresh_interval: None,
5257 crl_fetch_timeout: Duration::from_secs(30),
5258 crl_stale_grace: Duration::from_secs(24 * 60 * 60),
5259 crl_deny_on_unavailable: false,
5260 crl_end_entity_only: false,
5261 crl_allow_http: true,
5262 crl_enforce_expiration: true,
5263 crl_max_concurrent_fetches: 4,
5264 crl_max_response_bytes: 5 * 1024 * 1024,
5265 crl_discovery_rate_per_min: 60,
5266 crl_max_host_semaphores: 1024,
5267 crl_max_seen_urls: 4096,
5268 crl_max_cache_entries: 1024,
5269 }
5270 }
5271
5272 #[test]
5273 fn validate_rejects_zero_crl_max_response_bytes() {
5274 let mut mtls = valid_mtls_config();
5275 mtls.crl_max_response_bytes = 0;
5276 let mut auth = AuthConfig::with_keys(vec![]);
5277 auth.mtls = Some(mtls);
5278
5279 let mut cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
5282 cfg.tls_cert_path = Some("cert.pem".into());
5283 cfg.tls_key_path = Some("key.pem".into());
5284
5285 let err = cfg.validate().expect_err("zero CRL response cap");
5286 assert!(err.to_string().contains("crl_max_response_bytes"));
5287 }
5288
5289 #[test]
5292 fn validate_accepts_pre_auth_burst_without_explicit_pre_auth_rate() {
5293 let auth = AuthConfig::with_keys(vec![])
5294 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(50));
5295 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
5296 assert!(cfg.validate().is_ok(), "pre_auth_burst has no orphan rule");
5297 }
5298
5299 #[test]
5302 fn trusted_forwarder_max_entries_bounds_are_enforced() {
5303 let cfg = |n: usize| {
5304 McpServerConfig::new("127.0.0.1:8080", "t", "0")
5305 .with_trusted_forwarder_max_entries(n)
5306 .validate()
5307 };
5308 assert!(cfg(0).is_err(), "0 would pin every client to the proxy");
5309 assert!(
5310 cfg(crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES + 1).is_err(),
5311 "above the ceiling would re-open the header-bomb vector"
5312 );
5313 assert!(cfg(1).is_ok());
5314 assert!(cfg(crate::forwarded::MAX_SCANNED_ENTRIES).is_ok());
5315 assert!(cfg(crate::forwarded::MAX_CONFIGURABLE_SCANNED_ENTRIES).is_ok());
5316 }
5317
5318 #[test]
5319 fn trusted_forwarder_max_entries_defaults_to_the_module_constant() {
5320 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "0");
5321 assert_eq!(
5322 cfg.trusted_forwarder_max_entries,
5323 crate::forwarded::MAX_SCANNED_ENTRIES
5324 );
5325 }
5326
5327 fn forward_resolver(trusted: &[&str], mode: ForwardedHeaderMode) -> Arc<ForwardResolver> {
5328 Arc::new(ForwardResolver {
5329 trusted: trusted.iter().map(|s| s.parse().unwrap()).collect(),
5330 mode,
5331 max_scanned_entries: crate::forwarded::MAX_SCANNED_ENTRIES,
5332 })
5333 }
5334
5335 fn forwarded_probe_router(resolver: Option<Arc<ForwardResolver>>) -> axum::Router {
5337 async fn probe(req: Request<Body>) -> String {
5338 let pa = req
5339 .extensions()
5340 .get::<PeerAddr>()
5341 .map(|p| p.addr.ip().to_string())
5342 .unwrap_or_default();
5343 let ci = req
5344 .extensions()
5345 .get::<ClientIp>()
5346 .map(|c| c.ip.to_string())
5347 .unwrap_or_default();
5348 format!("{pa}|{ci}")
5349 }
5350 axum::Router::new()
5351 .route("/probe", axum::routing::get(probe))
5352 .layer(axum::middleware::from_fn(move |req, next| {
5353 let r = resolver.clone();
5354 normalize_peer_addr_middleware(r, req, next)
5355 }))
5356 }
5357
5358 fn probe_req(peer: &str, header: Option<(&str, &str)>) -> Request<Body> {
5359 let addr: SocketAddr = peer.parse().unwrap();
5360 let mut builder = Request::builder()
5361 .uri("/probe")
5362 .extension(ConnectInfo(addr));
5363 if let Some((name, value)) = header {
5364 builder = builder.header(name, value);
5365 }
5366 builder.body(Body::empty()).unwrap()
5367 }
5368
5369 #[tokio::test]
5370 async fn client_ip_equals_direct_without_resolver() {
5371 let app = forwarded_probe_router(None);
5372 let resp = app
5373 .oneshot(probe_req(
5374 "10.1.2.3:4444",
5375 Some(("x-forwarded-for", "203.0.113.7")),
5376 ))
5377 .await
5378 .unwrap();
5379 assert_eq!(
5380 body_string(resp).await,
5381 "10.1.2.3|10.1.2.3",
5382 "feature off: header ignored, ClientIp == direct"
5383 );
5384 }
5385
5386 #[tokio::test]
5387 async fn client_ip_resolved_for_trusted_peer() {
5388 let app = forwarded_probe_router(Some(forward_resolver(
5389 &["10.0.0.0/8"],
5390 ForwardedHeaderMode::XForwardedFor,
5391 )));
5392 let resp = app
5393 .oneshot(probe_req(
5394 "10.0.0.1:9999",
5395 Some(("x-forwarded-for", "203.0.113.7")),
5396 ))
5397 .await
5398 .unwrap();
5399 assert_eq!(
5400 body_string(resp).await,
5401 "10.0.0.1|203.0.113.7",
5402 "PeerAddr stays direct while ClientIp resolves"
5403 );
5404 }
5405
5406 #[tokio::test]
5407 async fn client_ip_falls_back_to_direct_on_malformed_header() {
5408 let app = forwarded_probe_router(Some(forward_resolver(
5409 &["10.0.0.0/8"],
5410 ForwardedHeaderMode::XForwardedFor,
5411 )));
5412 let resp = app
5413 .oneshot(probe_req(
5414 "10.0.0.1:9999",
5415 Some(("x-forwarded-for", "not-an-ip")),
5416 ))
5417 .await
5418 .unwrap();
5419 assert_eq!(
5420 body_string(resp).await,
5421 "10.0.0.1|10.0.0.1",
5422 "malformed chain falls back to the direct peer"
5423 );
5424 }
5425
5426 #[test]
5427 fn forwarded_header_mode_deserializes_kebab_case() {
5428 #[derive(serde::Deserialize)]
5429 struct Wrapper {
5430 mode: ForwardedHeaderMode,
5431 }
5432 let w: Wrapper = toml::from_str(r#"mode = "x-forwarded-for""#).unwrap();
5433 assert_eq!(w.mode, ForwardedHeaderMode::XForwardedFor);
5434 let w: Wrapper = toml::from_str(r#"mode = "forwarded""#).unwrap();
5435 assert_eq!(w.mode, ForwardedHeaderMode::Forwarded);
5436 assert!(
5437 toml::from_str::<Wrapper>(r#"mode = "XForwardedFor""#).is_err(),
5438 "PascalCase wire value must be rejected"
5439 );
5440 }
5441
5442 #[test]
5443 fn validate_rejects_bad_trusted_proxy_entry() {
5444 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5445 .with_trusted_proxies(["not-a-cidr"]);
5446 let err = cfg.validate().expect_err("bad CIDR");
5447 assert!(err.to_string().contains("trusted_proxies"));
5448 }
5449
5450 #[test]
5451 fn validate_rejects_zero_prefix_trusted_proxy() {
5452 for entry in ["0.0.0.0/0", "::/0"] {
5453 let cfg =
5454 McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([entry]);
5455 let err = cfg.validate().expect_err("zero-prefix CIDR");
5456 assert!(
5457 err.to_string().contains("prefix length 0"),
5458 "entry {entry}: {err}"
5459 );
5460 }
5461 }
5462
5463 #[test]
5464 fn validate_accepts_cidr_and_bare_ip_proxy_entries() {
5465 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([
5466 "10.0.0.0/8",
5467 "192.0.2.1",
5468 "2001:db8::1",
5469 ]);
5470 assert!(cfg.validate().is_ok(), "CIDRs and bare IPs are accepted");
5471 }
5472
5473 #[test]
5474 fn validate_rejects_forwarded_header_without_proxies() {
5475 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
5476 .with_forwarded_header(ForwardedHeaderMode::Forwarded);
5477 let err = cfg.validate().expect_err("mode without proxies");
5478 assert!(err.to_string().contains("requires trusted_proxies"));
5479 }
5480
5481 fn origin_router(origins: Vec<String>, log_request_headers: bool) -> axum::Router {
5485 let allowed: Arc<[String]> = Arc::from(origins);
5486 axum::Router::new()
5487 .route("/test", axum::routing::get(|| async { "ok" }))
5488 .layer(axum::middleware::from_fn(move |req, next| {
5489 let a = Arc::clone(&allowed);
5490 origin_check_middleware(a, log_request_headers, req, next)
5491 }))
5492 }
5493
5494 #[tokio::test]
5495 async fn origin_allowed_passes() {
5496 let app = origin_router(vec!["http://localhost:3000".into()], false);
5497 let req = Request::builder()
5498 .uri("/test")
5499 .header(header::ORIGIN, "http://localhost:3000")
5500 .body(Body::empty())
5501 .unwrap();
5502 let resp = app.oneshot(req).await.unwrap();
5503 assert_eq!(resp.status(), StatusCode::OK);
5504 }
5505
5506 #[tokio::test]
5507 async fn origin_rejected_returns_403() {
5508 let app = origin_router(vec!["http://localhost:3000".into()], false);
5509 let req = Request::builder()
5510 .uri("/test")
5511 .header(header::ORIGIN, "http://evil.com")
5512 .body(Body::empty())
5513 .unwrap();
5514 let resp = app.oneshot(req).await.unwrap();
5515 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5516 }
5517
5518 #[tokio::test]
5519 async fn no_origin_header_passes() {
5520 let app = origin_router(vec!["http://localhost:3000".into()], false);
5521 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5522 let resp = app.oneshot(req).await.unwrap();
5523 assert_eq!(resp.status(), StatusCode::OK);
5524 }
5525
5526 #[tokio::test]
5527 async fn empty_allowlist_rejects_any_origin() {
5528 let app = origin_router(vec![], false);
5529 let req = Request::builder()
5530 .uri("/test")
5531 .header(header::ORIGIN, "http://anything.com")
5532 .body(Body::empty())
5533 .unwrap();
5534 let resp = app.oneshot(req).await.unwrap();
5535 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5536 }
5537
5538 #[tokio::test]
5539 async fn empty_allowlist_passes_without_origin() {
5540 let app = origin_router(vec![], false);
5541 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5542 let resp = app.oneshot(req).await.unwrap();
5543 assert_eq!(resp.status(), StatusCode::OK);
5544 }
5545
5546 #[test]
5547 fn format_request_headers_redacts_sensitive_values() {
5548 let mut headers = axum::http::HeaderMap::new();
5549 headers.insert("authorization", "Bearer secret-token".parse().unwrap());
5550 headers.insert("cookie", "sid=abc".parse().unwrap());
5551 headers.insert("x-request-id", "req-123".parse().unwrap());
5552
5553 let out = format_request_headers_for_log(&headers);
5554 assert!(out.contains("authorization: [REDACTED]"));
5555 assert!(out.contains("cookie: [REDACTED]"));
5556 assert!(out.contains("x-request-id: req-123"));
5557 assert!(!out.contains("secret-token"));
5558 }
5559
5560 #[test]
5561 fn format_request_headers_redacts_forwarding_headers() {
5562 let mut headers = axum::http::HeaderMap::new();
5563 headers.insert("forwarded", "for=203.0.113.9;by=10.1.2.3".parse().unwrap());
5564 headers.insert("x-forwarded-for", "203.0.113.9, 10.1.2.3".parse().unwrap());
5565 headers.insert("x-real-ip", "203.0.113.9".parse().unwrap());
5566 headers.insert("x-request-id", "req-123".parse().unwrap());
5567
5568 let out = format_request_headers_for_log(&headers);
5569 for name in ["forwarded", "x-forwarded-for", "x-real-ip"] {
5570 assert!(
5571 out.contains(&format!("{name}: [REDACTED]")),
5572 "{name} carries client IP / proxy topology and must not reach logs; got {out}"
5573 );
5574 }
5575 assert!(
5576 !out.contains("203.0.113.9") && !out.contains("10.1.2.3"),
5577 "no forwarded address may survive redaction; got {out}"
5578 );
5579 assert!(out.contains("x-request-id: req-123"));
5580 }
5581
5582 fn security_router(is_tls: bool) -> axum::Router {
5585 security_router_with(is_tls, SecurityHeadersConfig::default())
5586 }
5587
5588 fn security_router_with(is_tls: bool, cfg: SecurityHeadersConfig) -> axum::Router {
5589 let cfg = Arc::new(cfg);
5590 axum::Router::new()
5591 .route("/test", axum::routing::get(|| async { "ok" }))
5592 .layer(axum::middleware::from_fn(move |req, next| {
5593 let c = Arc::clone(&cfg);
5594 security_headers_middleware(is_tls, c, req, next)
5595 }))
5596 }
5597
5598 #[tokio::test]
5599 async fn security_headers_set_on_response() {
5600 let app = security_router(false);
5601 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5602 let resp = app.oneshot(req).await.unwrap();
5603 assert_eq!(resp.status(), StatusCode::OK);
5604
5605 let h = resp.headers();
5606 assert_eq!(h.get("x-content-type-options").unwrap(), "nosniff");
5607 assert_eq!(h.get("x-frame-options").unwrap(), "deny");
5608 assert_eq!(h.get("cache-control").unwrap(), "no-store, max-age=0");
5609 assert_eq!(h.get("referrer-policy").unwrap(), "no-referrer");
5610 assert_eq!(h.get("cross-origin-opener-policy").unwrap(), "same-origin");
5611 assert_eq!(
5612 h.get("cross-origin-resource-policy").unwrap(),
5613 "same-origin"
5614 );
5615 assert_eq!(
5616 h.get("cross-origin-embedder-policy").unwrap(),
5617 "require-corp"
5618 );
5619 assert_eq!(h.get("x-permitted-cross-domain-policies").unwrap(), "none");
5620 assert!(
5621 h.get("permissions-policy")
5622 .unwrap()
5623 .to_str()
5624 .unwrap()
5625 .contains("camera=()"),
5626 "permissions-policy must restrict browser features"
5627 );
5628 assert_eq!(
5629 h.get("content-security-policy").unwrap(),
5630 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
5631 );
5632 assert_eq!(h.get("x-dns-prefetch-control").unwrap(), "off");
5633 assert!(h.get("strict-transport-security").is_none());
5635 }
5636
5637 #[tokio::test]
5638 async fn hsts_set_when_tls_enabled() {
5639 let app = security_router(true);
5640 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5641 let resp = app.oneshot(req).await.unwrap();
5642
5643 let hsts = resp.headers().get("strict-transport-security").unwrap();
5644 assert!(
5645 hsts.to_str().unwrap().contains("max-age=63072000"),
5646 "HSTS must set 2-year max-age"
5647 );
5648 }
5649
5650 #[tokio::test]
5651 async fn default_csp_matches_guideline() {
5652 let app = security_router(false);
5653 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5654 let resp = app.oneshot(req).await.unwrap();
5655 assert_eq!(
5656 resp.headers().get("content-security-policy").unwrap(),
5657 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
5658 );
5659 }
5660
5661 #[tokio::test]
5662 async fn operator_csp_override_still_wins() {
5663 let cfg = SecurityHeadersConfig {
5664 content_security_policy: Some("default-src 'self'".into()),
5665 ..SecurityHeadersConfig::default()
5666 };
5667 let app = security_router_with(false, cfg);
5668 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5669 let resp = app.oneshot(req).await.unwrap();
5670 assert_eq!(
5671 resp.headers().get("content-security-policy").unwrap(),
5672 "default-src 'self'"
5673 );
5674 }
5675
5676 fn check_with_security_headers(
5682 headers: SecurityHeadersConfig,
5683 ) -> Result<(), RmcpServerKitError> {
5684 let cfg =
5685 McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_security_headers(headers);
5686 cfg.check()
5687 }
5688
5689 #[test]
5690 fn security_headers_config_default_validates() {
5691 check_with_security_headers(SecurityHeadersConfig::default())
5692 .expect("default SecurityHeadersConfig must validate");
5693 }
5694
5695 #[test]
5696 fn security_headers_config_validate_accepts_empty_string() {
5697 let h = SecurityHeadersConfig {
5699 x_content_type_options: Some(String::new()),
5700 x_frame_options: Some(String::new()),
5701 cache_control: Some(String::new()),
5702 referrer_policy: Some(String::new()),
5703 cross_origin_opener_policy: Some(String::new()),
5704 cross_origin_resource_policy: Some(String::new()),
5705 cross_origin_embedder_policy: Some(String::new()),
5706 permissions_policy: Some(String::new()),
5707 x_permitted_cross_domain_policies: Some(String::new()),
5708 content_security_policy: Some(String::new()),
5709 x_dns_prefetch_control: Some(String::new()),
5710 strict_transport_security: Some(String::new()),
5711 };
5712 check_with_security_headers(h).expect("Some(\"\") on every field must validate (omit-all)");
5713 }
5714
5715 #[test]
5716 fn security_headers_config_validate_rejects_bad_value() {
5717 let h = SecurityHeadersConfig {
5719 referrer_policy: Some("\u{0007}".into()),
5720 ..SecurityHeadersConfig::default()
5721 };
5722 let err = check_with_security_headers(h)
5723 .expect_err("control char in referrer_policy must reject");
5724 let msg = err.to_string();
5725 assert!(
5726 msg.contains("referrer_policy"),
5727 "error must name the offending field, got: {msg}"
5728 );
5729 }
5730
5731 #[test]
5732 fn security_headers_config_validate_rejects_hsts_preload() {
5733 let h = SecurityHeadersConfig {
5734 strict_transport_security: Some("max-age=63072000; includeSubDomains; preload".into()),
5735 ..SecurityHeadersConfig::default()
5736 };
5737 let err = check_with_security_headers(h).expect_err("HSTS with preload must reject");
5738 let msg = err.to_string();
5739 assert!(
5740 msg.contains("strict_transport_security"),
5741 "error must name the field, got: {msg}"
5742 );
5743 assert!(
5744 msg.to_lowercase().contains("preload"),
5745 "error must mention `preload`, got: {msg}"
5746 );
5747 }
5748
5749 #[test]
5750 fn security_headers_config_validate_rejects_hsts_preload_uppercase() {
5751 let h = SecurityHeadersConfig {
5753 strict_transport_security: Some("max-age=600; PRELOAD".into()),
5754 ..SecurityHeadersConfig::default()
5755 };
5756 check_with_security_headers(h).expect_err("HSTS preload check must be case-insensitive");
5757 }
5758
5759 #[tokio::test]
5760 async fn security_headers_override_honored() {
5761 let h = SecurityHeadersConfig {
5763 x_frame_options: Some("SAMEORIGIN".into()),
5764 ..SecurityHeadersConfig::default()
5765 };
5766 let app = security_router_with(false, h);
5767 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5768 let resp = app.oneshot(req).await.unwrap();
5769 assert_eq!(resp.status(), StatusCode::OK);
5770
5771 let xfo = resp.headers().get("x-frame-options").unwrap();
5772 assert_eq!(xfo, "SAMEORIGIN");
5773 }
5774
5775 #[tokio::test]
5776 async fn security_headers_empty_string_omits() {
5777 let h = SecurityHeadersConfig {
5779 referrer_policy: Some(String::new()),
5780 ..SecurityHeadersConfig::default()
5781 };
5782 let app = security_router_with(false, h);
5783 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5784 let resp = app.oneshot(req).await.unwrap();
5785 assert_eq!(resp.status(), StatusCode::OK);
5786
5787 assert!(
5788 resp.headers().get("referrer-policy").is_none(),
5789 "Some(\"\") must omit the header"
5790 );
5791 assert_eq!(
5793 resp.headers().get("x-content-type-options").unwrap(),
5794 "nosniff"
5795 );
5796 }
5797
5798 #[tokio::test]
5799 async fn security_headers_hsts_only_when_tls() {
5800 let h = SecurityHeadersConfig {
5802 strict_transport_security: Some("max-age=600".into()),
5803 ..SecurityHeadersConfig::default()
5804 };
5805 let app = security_router_with(false, h);
5806 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
5807 let resp = app.oneshot(req).await.unwrap();
5808 assert!(
5809 resp.headers().get("strict-transport-security").is_none(),
5810 "HSTS must remain absent on plaintext deployments even with override"
5811 );
5812 }
5813
5814 #[cfg(feature = "oauth")]
5817 #[tokio::test]
5818 async fn oauth_token_cache_headers_set_pragma_and_vary() {
5819 let app = axum::Router::new()
5820 .route("/token", axum::routing::post(|| async { "{}" }))
5821 .layer(axum::middleware::from_fn(
5822 oauth_token_cache_headers_middleware,
5823 ));
5824 let req = Request::builder()
5825 .method("POST")
5826 .uri("/token")
5827 .body(Body::from("{}"))
5828 .unwrap();
5829 let resp = app.oneshot(req).await.unwrap();
5830 assert_eq!(resp.status(), StatusCode::OK);
5831
5832 let h = resp.headers();
5833 assert_eq!(
5834 h.get("pragma").unwrap(),
5835 "no-cache",
5836 "RFC 6749 §5.1: token responses must set Pragma: no-cache"
5837 );
5838 let vary_values: Vec<String> = h
5839 .get_all("vary")
5840 .iter()
5841 .filter_map(|v| v.to_str().ok().map(str::to_owned))
5842 .collect();
5843 assert!(
5844 vary_values
5845 .iter()
5846 .any(|v| v.eq_ignore_ascii_case("Authorization")),
5847 "RFC 6750 §5.4: Vary must include Authorization, got {vary_values:?}"
5848 );
5849 }
5850
5851 #[cfg(feature = "oauth")]
5852 #[tokio::test]
5853 async fn oauth_token_cache_headers_preserve_existing_vary() {
5854 let app = axum::Router::new()
5857 .route(
5858 "/token",
5859 axum::routing::post(|| async {
5860 axum::response::Response::builder()
5861 .header("vary", "Accept-Encoding")
5862 .body(Body::from("{}"))
5863 .unwrap()
5864 }),
5865 )
5866 .layer(axum::middleware::from_fn(
5867 oauth_token_cache_headers_middleware,
5868 ));
5869 let req = Request::builder()
5870 .method("POST")
5871 .uri("/token")
5872 .body(Body::empty())
5873 .unwrap();
5874 let resp = app.oneshot(req).await.unwrap();
5875
5876 let vary: Vec<String> = resp
5877 .headers()
5878 .get_all("vary")
5879 .iter()
5880 .filter_map(|v| v.to_str().ok().map(str::to_owned))
5881 .collect();
5882 assert!(
5883 vary.iter().any(|v| v.contains("Accept-Encoding")),
5884 "must preserve pre-existing Vary value, got {vary:?}"
5885 );
5886 assert!(
5887 vary.iter().any(|v| v.contains("Authorization")),
5888 "must append Authorization to Vary, got {vary:?}"
5889 );
5890 }
5891
5892 #[test]
5895 fn version_omits_build_fingerprint_by_default() {
5896 let v = version_payload("my-server", "1.2.3", false);
5897 assert_eq!(v["name"], "my-server");
5898 assert_eq!(v["version"], "1.2.3");
5899 assert!(v["rmcp_server_kit_version"].is_string());
5900 assert!(
5901 v.get("build_git_sha").is_none(),
5902 "build sha must be hidden by default"
5903 );
5904 assert!(v.get("build_timestamp").is_none());
5905 assert!(v.get("rust_version").is_none());
5906 }
5907
5908 #[test]
5909 fn version_exposes_all_when_enabled() {
5910 let v = version_payload("my-server", "1.2.3", true);
5911 assert!(v["build_git_sha"].is_string());
5912 assert!(v["build_timestamp"].is_string());
5913 assert!(v["rust_version"].is_string());
5914 assert!(v["rmcp_server_kit_version"].is_string());
5915 }
5916
5917 #[tokio::test]
5920 async fn concurrency_limit_layer_composes_and_serves() {
5921 let app = axum::Router::new()
5925 .route("/ok", axum::routing::get(|| async { "ok" }))
5926 .layer(
5927 tower::ServiceBuilder::new()
5928 .layer(axum::error_handling::HandleErrorLayer::new(
5929 |_err: tower::BoxError| async { StatusCode::SERVICE_UNAVAILABLE },
5930 ))
5931 .layer(tower::load_shed::LoadShedLayer::new())
5932 .layer(tower::limit::ConcurrencyLimitLayer::new(4)),
5933 );
5934 let resp = app
5935 .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap())
5936 .await
5937 .unwrap();
5938 assert_eq!(resp.status(), StatusCode::OK);
5939 }
5940
5941 #[tokio::test]
5944 async fn compression_layer_gzip_encodes_response() {
5945 use tower_http::compression::Predicate as _;
5946
5947 let big_body = "a".repeat(4096);
5948 let app = axum::Router::new()
5949 .route(
5950 "/big",
5951 axum::routing::get(move || {
5952 let body = big_body.clone();
5953 async move { body }
5954 }),
5955 )
5956 .layer(
5957 tower_http::compression::CompressionLayer::new()
5958 .gzip(true)
5959 .br(true)
5960 .compress_when(
5961 tower_http::compression::DefaultPredicate::new()
5962 .and(tower_http::compression::predicate::SizeAbove::new(1024)),
5963 ),
5964 );
5965
5966 let req = Request::builder()
5967 .uri("/big")
5968 .header(header::ACCEPT_ENCODING, "gzip")
5969 .body(Body::empty())
5970 .unwrap();
5971 let resp = app.oneshot(req).await.unwrap();
5972 assert_eq!(resp.status(), StatusCode::OK);
5973 assert_eq!(
5974 resp.headers().get(header::CONTENT_ENCODING).unwrap(),
5975 "gzip"
5976 );
5977 }
5978
5979 #[tokio::test]
5982 async fn tls_handshake_timeout_reaps_idle_connections() {
5983 use tokio::io::AsyncReadExt as _;
5984
5985 let _ = rustls::crypto::ring::default_provider().install_default();
5986
5987 let key = rcgen::KeyPair::generate().expect("generate key");
5989 let cert = rcgen::CertificateParams::new(vec!["localhost".to_owned()])
5990 .expect("cert params")
5991 .self_signed(&key)
5992 .expect("self-signed cert");
5993 let dir = std::env::temp_dir().join(format!(
5994 "rmcp-server-kit-hs-timeout-{}",
5995 std::time::SystemTime::now()
5996 .duration_since(std::time::UNIX_EPOCH)
5997 .expect("clock after epoch")
5998 .as_nanos()
5999 ));
6000 tokio::fs::create_dir_all(&dir).await.expect("temp dir");
6001 let cert_path = dir.join("server.crt");
6002 let key_path = dir.join("server.key");
6003 tokio::fs::write(&cert_path, cert.pem())
6004 .await
6005 .expect("write cert");
6006 tokio::fs::write(&key_path, key.serialize_pem())
6007 .await
6008 .expect("write key");
6009
6010 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
6011 let tls = TlsListener::new(
6012 listener,
6013 &cert_path,
6014 &key_path,
6015 None,
6016 None,
6017 Duration::from_millis(200),
6018 8, )
6020 .expect("tls listener");
6021 let addr = axum::serve::Listener::local_addr(&tls).expect("local addr");
6022
6023 let mut idle = tokio::net::TcpStream::connect(addr).await.expect("connect");
6027 let mut buf = [0_u8; 16];
6028 let read = tokio::time::timeout(Duration::from_secs(2), idle.read(&mut buf))
6029 .await
6030 .expect("server must reap the idle handshake within its timeout");
6031 match read {
6032 Ok(0) | Err(_) => {} Ok(n) => panic!("unexpected {n} bytes from server during reaped handshake"),
6034 }
6035
6036 drop(tls);
6037 }
6038
6039 fn assert_owasp_headers(resp: &axum::response::Response, ctx: &str) {
6042 let h = resp.headers();
6043 assert!(
6044 h.contains_key("x-content-type-options"),
6045 "{ctx}: missing X-Content-Type-Options"
6046 );
6047 assert!(
6048 h.contains_key("x-frame-options"),
6049 "{ctx}: missing X-Frame-Options"
6050 );
6051 assert!(
6052 h.contains_key("strict-transport-security"),
6053 "{ctx}: missing Strict-Transport-Security"
6054 );
6055 assert!(
6056 h.contains_key(header::CONTENT_SECURITY_POLICY),
6057 "{ctx}: missing Content-Security-Policy"
6058 );
6059 }
6060
6061 fn m5_router(configure: impl FnOnce(&mut McpServerConfig)) -> axum::Router {
6062 #[derive(Clone)]
6063 struct H;
6064 impl ServerHandler for H {}
6065 let mut config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0")
6069 .with_allowed_origins(["http://good.example"])
6070 .with_tls("unused.crt", "unused.key");
6071 configure(&mut config);
6072 let (router, _params) = build_app_router(config, || H).expect("build_app_router");
6073 router
6074 }
6075
6076 #[test]
6081 #[should_panic(expected = "Overlapping method route")]
6082 fn extra_router_exact_overlap_with_framework_route_panics() {
6083 #[derive(Clone)]
6084 struct H;
6085 impl ServerHandler for H {}
6086 let config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_extra_router(
6087 axum::Router::new().route("/healthz", axum::routing::get(|| async { "mine" })),
6088 );
6089 let _ = build_app_router(config, || H);
6090 }
6091
6092 #[test]
6096 fn extra_router_non_overlapping_path_under_framework_prefix_is_accepted() {
6097 #[derive(Clone)]
6098 struct H;
6099 impl ServerHandler for H {}
6100 let config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_extra_router(
6101 axum::Router::new().route("/admin/custom", axum::routing::get(|| async { "mine" })),
6102 );
6103 assert!(
6104 build_app_router(config, || H).is_ok(),
6105 "non-overlapping path under a framework prefix must merge cleanly"
6106 );
6107 }
6108
6109 #[tokio::test]
6110 async fn headers_on_rejected_origin_403() {
6111 let app = m5_router(|_| {});
6112 let req = Request::builder()
6113 .uri("/healthz")
6114 .header(header::ORIGIN, "http://evil.example")
6115 .body(Body::empty())
6116 .unwrap();
6117 let resp = app.oneshot(req).await.unwrap();
6118 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
6119 assert_owasp_headers(&resp, "origin-403");
6120 }
6121
6122 #[tokio::test]
6123 async fn headers_on_cors_preflight() {
6124 let app = m5_router(|_| {});
6125 let req = Request::builder()
6126 .method(axum::http::Method::OPTIONS)
6127 .uri("/mcp")
6128 .header(header::ORIGIN, "http://good.example")
6129 .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
6130 .body(Body::empty())
6131 .unwrap();
6132 let resp = app.oneshot(req).await.unwrap();
6133 assert_owasp_headers(&resp, "cors-preflight");
6134 }
6135
6136 #[tokio::test]
6137 async fn headers_on_404_fallback() {
6138 let app = m5_router(|_| {});
6139 let req = Request::builder()
6140 .uri("/no-such-route")
6141 .body(Body::empty())
6142 .unwrap();
6143 let resp = app.oneshot(req).await.unwrap();
6144 assert_eq!(resp.status(), StatusCode::NOT_FOUND);
6145 assert_owasp_headers(&resp, "404-fallback");
6146 }
6147
6148 #[tokio::test]
6149 async fn headers_on_overload_503() {
6150 let app = m5_router(|c| c.max_concurrent_requests = Some(0));
6153 let req = Request::builder()
6154 .uri("/healthz")
6155 .body(Body::empty())
6156 .unwrap();
6157 let resp = app.oneshot(req).await.unwrap();
6158 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
6159 assert_owasp_headers(&resp, "overload-503");
6160 }
6161
6162 #[cfg(feature = "oauth")]
6165 fn m6_auth_state() -> (Arc<AuthState>, String, String) {
6166 let (admin_token, admin_hash) = crate::auth::generate_api_key().unwrap();
6167 let (viewer_token, viewer_hash) = crate::auth::generate_api_key().unwrap();
6168 let state = Arc::new(AuthState {
6169 api_keys: ArcSwap::from_pointee(vec![
6170 crate::auth::ApiKeyEntry::new("admin-key", admin_hash, "admin"),
6171 crate::auth::ApiKeyEntry::new("viewer-key", viewer_hash, "viewer"),
6172 ]),
6173 rate_limiter: None,
6174 pre_auth_limiter: None,
6175 jwks_cache: None,
6176 seen_identities: crate::auth::SeenIdentitySet::new(),
6177 counters: crate::auth::AuthCounters::default(),
6178 resource_metadata_url: None,
6179 });
6180 (state, admin_token, viewer_token)
6181 }
6182
6183 #[cfg(feature = "oauth")]
6184 fn m6_admin_router(state: &Arc<AuthState>) -> axum::Router {
6185 let proxy = crate::oauth::OAuthProxyConfig::builder(
6186 "https://idp.example/authorize",
6187 "https://idp.example/token",
6188 "client",
6189 )
6190 .introspection_url("http://127.0.0.1:1/introspect")
6191 .revocation_url("http://127.0.0.1:1/revoke")
6192 .expose_admin_endpoints(true)
6193 .require_auth_on_admin_endpoints(true)
6194 .build();
6195 let http = crate::oauth::OauthHttpClient::new().expect("oauth http client");
6196 build_oauth_admin_router(&proxy, http, Some(state), "admin").expect("admin router")
6197 }
6198
6199 #[cfg(feature = "oauth")]
6200 fn m6_req(path: &str, token: &str) -> Request<Body> {
6201 Request::builder()
6202 .method(axum::http::Method::POST)
6203 .uri(path)
6204 .header(header::AUTHORIZATION, format!("Bearer {token}"))
6205 .body(Body::from("token=abc"))
6206 .unwrap()
6207 }
6208
6209 #[cfg(feature = "oauth")]
6210 #[tokio::test]
6211 async fn oauth_proxy_admin_requires_admin_role() {
6212 let (state, _admin, viewer) = m6_auth_state();
6213 for path in ["/introspect", "/revoke"] {
6214 let app = m6_admin_router(&state);
6215 let resp = app.oneshot(m6_req(path, &viewer)).await.unwrap();
6216 assert_eq!(
6217 resp.status(),
6218 StatusCode::FORBIDDEN,
6219 "an authenticated viewer must be rejected with 403 on {path}"
6220 );
6221 }
6222 }
6223
6224 #[cfg(feature = "oauth")]
6225 #[tokio::test]
6226 async fn oauth_proxy_admin_allows_admin_role() {
6227 let (state, admin, _viewer) = m6_auth_state();
6228 for path in ["/introspect", "/revoke"] {
6229 let app = m6_admin_router(&state);
6230 let resp = app.oneshot(m6_req(path, &admin)).await.unwrap();
6231 assert_ne!(
6235 resp.status(),
6236 StatusCode::FORBIDDEN,
6237 "an authenticated admin must pass the role gate on {path}"
6238 );
6239 assert_ne!(
6240 resp.status(),
6241 StatusCode::UNAUTHORIZED,
6242 "an authenticated admin must pass the auth gate on {path}"
6243 );
6244 }
6245 }
6246
6247 #[cfg(feature = "metrics")]
6255 mod metrics_labels_bounded {
6256 use super::*;
6257
6258 fn labels_for(method: &str, uri: &str) -> (&'static str, String) {
6259 let req = Request::builder()
6260 .method(method)
6261 .uri(uri)
6262 .body(Body::empty())
6263 .unwrap();
6264 metrics_labels(&req)
6265 }
6266
6267 #[test]
6268 fn many_unmatched_paths_collapse_to_one_label() {
6269 let mut seen = std::collections::HashSet::new();
6270 for i in 0..500 {
6271 let (_, path) = labels_for("GET", &format!("/nonexistent-{i}"));
6272 seen.insert(path);
6273 }
6274 assert_eq!(
6275 seen.len(),
6276 1,
6277 "unmatched paths must collapse to a single label, got {seen:?}"
6278 );
6279 assert!(seen.contains("<unmatched>"));
6280 }
6281
6282 #[test]
6283 fn nested_mcp_paths_collapse_to_the_mount_point() {
6284 let mut seen = std::collections::HashSet::new();
6285 for i in 0..200 {
6286 let (_, path) = labels_for("POST", &format!("/mcp/{i}"));
6287 seen.insert(path);
6288 }
6289 let (_, root) = labels_for("POST", "/mcp");
6290 seen.insert(root);
6291 assert_eq!(
6292 seen.len(),
6293 1,
6294 "nested /mcp paths must collapse to one label, got {seen:?}"
6295 );
6296 assert!(seen.contains("/mcp"));
6297 }
6298
6299 #[test]
6300 fn unusual_methods_collapse_to_one_bucket() {
6301 let mut seen = std::collections::HashSet::new();
6302 for verb in ["FROBNICATE", "WIBBLE", "QUUX", "M-SEARCH"] {
6303 let (method, _) = labels_for(verb, "/healthz");
6304 seen.insert(method);
6305 }
6306 assert_eq!(seen, std::collections::HashSet::from(["OTHER"]));
6307 }
6308
6309 #[test]
6310 fn known_methods_keep_their_identity() {
6311 for verb in ["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"] {
6312 let (method, _) = labels_for(verb, "/healthz");
6313 assert_eq!(method, verb);
6314 }
6315 }
6316
6317 #[test]
6318 fn raw_path_never_leaks_into_a_label() {
6319 let (_, path) = labels_for("GET", "/secret-token-abc123");
6320 assert!(
6321 !path.contains("secret-token"),
6322 "raw request path must never become a label value: {path}"
6323 );
6324 }
6325 }
6326}