1use std::{
2 future::Future,
3 net::{IpAddr, SocketAddr},
4 path::{Path, PathBuf},
5 pin::Pin,
6 sync::Arc,
7 time::Duration,
8};
9
10use arc_swap::ArcSwap;
11use axum::{
12 body::Body,
13 extract::{ConnectInfo, Request},
14 middleware::Next,
15 response::IntoResponse,
16};
17use rmcp::{
18 ServerHandler,
19 transport::streamable_http_server::{
20 StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
21 },
22};
23use rustls::RootCertStore;
24use tokio::{
25 net::TcpListener,
26 sync::{Semaphore, mpsc},
27};
28use tokio_util::sync::CancellationToken;
29
30use crate::{
31 auth::{
32 AuthConfig, AuthIdentity, AuthState, MtlsConfig, TlsConnInfo, auth_middleware,
33 build_rate_limiter, extract_mtls_identity,
34 },
35 bounded_limiter::BoundedKeyedLimiter,
36 error::RmcpServerKitError,
37 mtls_revocation::{self, CrlSet, DynamicClientCertVerifier},
38 rbac::{RbacPolicy, ToolRateLimiter, build_tool_rate_limiter, rbac_middleware},
39};
40
41#[allow(
45 clippy::needless_pass_by_value,
46 reason = "consumed at .map_err(anyhow_to_startup) call sites; by-value matches the closure shape"
47)]
48fn anyhow_to_startup(e: anyhow::Error) -> RmcpServerKitError {
49 RmcpServerKitError::Startup(format!("{e:#}"))
50}
51
52#[allow(
58 clippy::needless_pass_by_value,
59 reason = "consumed at .map_err(|e| io_to_startup(...)) call sites; by-value matches the closure shape"
60)]
61fn io_to_startup(op: &str, e: std::io::Error) -> RmcpServerKitError {
62 RmcpServerKitError::Startup(format!("{op}: {e}"))
63}
64
65pub type ReadinessCheck =
70 Arc<dyn Fn() -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> + Send + Sync>;
71
72#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
117#[non_exhaustive]
118pub struct PeerAddr {
119 pub addr: SocketAddr,
121}
122
123impl PeerAddr {
124 #[must_use]
127 pub(crate) const fn new(addr: SocketAddr) -> Self {
128 Self { addr }
129 }
130}
131
132impl<S: Send + Sync> axum::extract::FromRequestParts<S> for PeerAddr {
141 type Rejection = (axum::http::StatusCode, &'static str);
142
143 #[allow(
144 clippy::unused_async_trait_impl,
145 reason = "async is mandated by the axum FromRequestParts trait signature; this impl only reads a request extension synchronously"
146 )]
147 async fn from_request_parts(
148 parts: &mut axum::http::request::Parts,
149 _state: &S,
150 ) -> Result<Self, Self::Rejection> {
151 parts.extensions.get::<Self>().copied().ok_or((
152 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
153 "peer address unavailable: not running under rmcp-server-kit serve()",
154 ))
155 }
156}
157
158#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
181#[non_exhaustive]
182pub struct ClientIp {
183 pub ip: IpAddr,
185}
186
187impl ClientIp {
188 #[must_use]
191 pub(crate) const fn new(ip: IpAddr) -> Self {
192 Self { ip }
193 }
194}
195
196#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Deserialize)]
201#[serde(rename_all = "kebab-case")]
202#[non_exhaustive]
203pub enum ForwardedHeaderMode {
204 XForwardedFor,
206 Forwarded,
208}
209
210struct ForwardResolver {
213 trusted: Vec<ipnet::IpNet>,
214 mode: ForwardedHeaderMode,
215}
216
217#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
238#[serde(default)]
239#[non_exhaustive]
240pub struct SecurityHeadersConfig {
241 pub x_content_type_options: Option<String>,
243 pub x_frame_options: Option<String>,
245 pub cache_control: Option<String>,
247 pub referrer_policy: Option<String>,
249 pub cross_origin_opener_policy: Option<String>,
251 pub cross_origin_resource_policy: Option<String>,
253 pub cross_origin_embedder_policy: Option<String>,
255 pub permissions_policy: Option<String>,
258 pub x_permitted_cross_domain_policies: Option<String>,
260 pub content_security_policy: Option<String>,
263 pub x_dns_prefetch_control: Option<String>,
265 pub strict_transport_security: Option<String>,
270}
271
272#[allow(
274 missing_debug_implementations,
275 reason = "contains callback/trait objects that don't impl Debug"
276)]
277#[allow(
278 clippy::struct_excessive_bools,
279 reason = "server configuration naturally has many boolean feature flags"
280)]
281#[non_exhaustive]
282pub struct McpServerConfig {
283 #[deprecated(
285 since = "0.13.0",
286 note = "use McpServerConfig::new() / with_bind_addr(); direct field access will become pub(crate) in a future major release"
287 )]
288 pub bind_addr: String,
289 #[deprecated(
291 since = "0.13.0",
292 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
293 )]
294 pub name: String,
295 #[deprecated(
297 since = "0.13.0",
298 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
299 )]
300 pub version: String,
301 #[deprecated(
303 since = "0.13.0",
304 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
305 )]
306 pub tls_cert_path: Option<PathBuf>,
307 #[deprecated(
309 since = "0.13.0",
310 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
311 )]
312 pub tls_key_path: Option<PathBuf>,
313 #[deprecated(
316 since = "0.13.0",
317 note = "use McpServerConfig::with_auth(); direct field access will become pub(crate) in a future major release"
318 )]
319 pub auth: Option<AuthConfig>,
320 #[deprecated(
323 since = "0.13.0",
324 note = "use McpServerConfig::with_rbac(); direct field access will become pub(crate) in a future major release"
325 )]
326 pub rbac: Option<Arc<RbacPolicy>>,
327 #[deprecated(
333 since = "0.13.0",
334 note = "use McpServerConfig::with_allowed_origins(); direct field access will become pub(crate) in a future major release"
335 )]
336 pub allowed_origins: Vec<String>,
337 #[deprecated(
340 since = "0.13.0",
341 note = "use McpServerConfig::with_tool_rate_limit(); direct field access will become pub(crate) in a future major release"
342 )]
343 pub tool_rate_limit: Option<u32>,
344 #[deprecated(
350 since = "1.12.0",
351 note = "use McpServerConfig::with_tool_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
352 )]
353 pub tool_rate_limit_burst: Option<u32>,
354 #[deprecated(
367 since = "1.11.0",
368 note = "use McpServerConfig::with_extra_route_rate_limit(); direct field access will become pub(crate) in a future major release"
369 )]
370 pub extra_route_rate_limit: Option<u32>,
371 #[deprecated(
378 since = "1.12.0",
379 note = "use McpServerConfig::with_extra_route_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
380 )]
381 pub extra_route_rate_limit_burst: Option<u32>,
382 #[deprecated(
395 since = "1.14.0",
396 note = "use McpServerConfig::with_extra_route_rate_limit_exempt_paths(); direct field access will become pub(crate) in a future major release"
397 )]
398 pub extra_route_rate_limit_exempt_paths: Vec<String>,
399 #[deprecated(
407 since = "1.13.0",
408 note = "use McpServerConfig::with_trusted_proxies(); direct field access will become pub(crate) in a future major release"
409 )]
410 pub trusted_proxies: Vec<String>,
411 #[deprecated(
416 since = "1.13.0",
417 note = "use McpServerConfig::with_forwarded_header(); direct field access will become pub(crate) in a future major release"
418 )]
419 pub forwarded_header: Option<ForwardedHeaderMode>,
420 #[deprecated(
423 since = "0.13.0",
424 note = "use McpServerConfig::with_readiness_check(); direct field access will become pub(crate) in a future major release"
425 )]
426 pub readiness_check: Option<ReadinessCheck>,
427 #[deprecated(
430 since = "0.13.0",
431 note = "use McpServerConfig::with_max_request_body(); direct field access will become pub(crate) in a future major release"
432 )]
433 pub max_request_body: usize,
434 #[deprecated(
437 since = "0.13.0",
438 note = "use McpServerConfig::with_request_timeout(); direct field access will become pub(crate) in a future major release"
439 )]
440 pub request_timeout: Duration,
441 #[deprecated(
444 since = "0.13.0",
445 note = "use McpServerConfig::with_shutdown_timeout(); direct field access will become pub(crate) in a future major release"
446 )]
447 pub shutdown_timeout: Duration,
448 #[deprecated(
451 since = "0.13.0",
452 note = "use McpServerConfig::with_session_idle_timeout(); direct field access will become pub(crate) in a future major release"
453 )]
454 pub session_idle_timeout: Duration,
455 #[deprecated(
458 since = "0.13.0",
459 note = "use McpServerConfig::with_sse_keep_alive(); direct field access will become pub(crate) in a future major release"
460 )]
461 pub sse_keep_alive: Duration,
462 #[deprecated(
466 since = "0.13.0",
467 note = "use McpServerConfig::with_reload_callback(); direct field access will become pub(crate) in a future major release"
468 )]
469 pub on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
470 #[deprecated(
477 since = "0.13.0",
478 note = "use McpServerConfig::with_extra_router(); direct field access will become pub(crate) in a future major release"
479 )]
480 pub extra_router: Option<axum::Router>,
481 #[deprecated(
486 since = "0.13.0",
487 note = "use McpServerConfig::with_public_url(); direct field access will become pub(crate) in a future major release"
488 )]
489 pub public_url: Option<String>,
490 #[deprecated(
493 since = "0.13.0",
494 note = "use McpServerConfig::enable_request_header_logging(); direct field access will become pub(crate) in a future major release"
495 )]
496 pub log_request_headers: bool,
497 pub expose_build_metadata: bool,
504 #[deprecated(
507 since = "0.13.0",
508 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
509 )]
510 pub compression_enabled: bool,
511 #[deprecated(
514 since = "0.13.0",
515 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
516 )]
517 pub compression_min_size: u16,
518 #[deprecated(
522 since = "0.13.0",
523 note = "use McpServerConfig::with_max_concurrent_requests(); direct field access will become pub(crate) in a future major release"
524 )]
525 pub max_concurrent_requests: Option<usize>,
526 #[deprecated(
529 since = "0.13.0",
530 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
531 )]
532 pub admin_enabled: bool,
533 #[deprecated(
535 since = "0.13.0",
536 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
537 )]
538 pub admin_role: String,
539 #[cfg(feature = "metrics")]
542 #[deprecated(
543 since = "0.13.0",
544 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
545 )]
546 pub metrics_enabled: bool,
547 #[cfg(feature = "metrics")]
549 #[deprecated(
550 since = "0.13.0",
551 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
552 )]
553 pub metrics_bind: String,
554 #[deprecated(
558 since = "1.5.0",
559 note = "use McpServerConfig::with_security_headers(); direct field access will become pub(crate) in a future major release"
560 )]
561 pub security_headers: SecurityHeadersConfig,
562 #[deprecated(
568 since = "1.9.0",
569 note = "use McpServerConfig::with_tls_handshake_timeout(); direct field access will become pub(crate) in a future major release"
570 )]
571 pub tls_handshake_timeout: Duration,
572 #[deprecated(
579 since = "1.9.0",
580 note = "use McpServerConfig::with_max_concurrent_tls_handshakes(); direct field access will become pub(crate) in a future major release"
581 )]
582 pub max_concurrent_tls_handshakes: usize,
583}
584
585#[allow(
643 missing_debug_implementations,
644 reason = "wraps T which may not implement Debug; manual impl below avoids leaking inner contents into logs"
645)]
646pub struct Validated<T>(T);
647
648impl<T> std::fmt::Debug for Validated<T> {
649 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
650 f.debug_struct("Validated").finish_non_exhaustive()
651 }
652}
653
654impl<T> Validated<T> {
655 #[must_use]
657 pub fn as_inner(&self) -> &T {
658 &self.0
659 }
660
661 #[must_use]
666 pub fn into_inner(self) -> T {
667 self.0
668 }
669}
670
671#[allow(
672 deprecated,
673 reason = "internal builders/validators legitimately read/write the deprecated `pub` fields they were designed to manage"
674)]
675impl McpServerConfig {
676 #[must_use]
684 pub fn new(
685 bind_addr: impl Into<String>,
686 name: impl Into<String>,
687 version: impl Into<String>,
688 ) -> Self {
689 Self {
690 bind_addr: bind_addr.into(),
691 name: name.into(),
692 version: version.into(),
693 tls_cert_path: None,
694 tls_key_path: None,
695 auth: None,
696 rbac: None,
697 allowed_origins: Vec::new(),
698 tool_rate_limit: None,
699 readiness_check: None,
700 max_request_body: 1024 * 1024,
701 request_timeout: Duration::from_mins(2),
702 shutdown_timeout: Duration::from_secs(30),
703 session_idle_timeout: Duration::from_mins(20),
704 sse_keep_alive: Duration::from_secs(15),
705 on_reload_ready: None,
706 extra_router: None,
707 public_url: None,
708 log_request_headers: false,
709 expose_build_metadata: false,
710 compression_enabled: false,
711 compression_min_size: 1024,
712 max_concurrent_requests: None,
713 admin_enabled: false,
714 admin_role: "admin".to_owned(),
715 #[cfg(feature = "metrics")]
716 metrics_enabled: false,
717 #[cfg(feature = "metrics")]
718 metrics_bind: "127.0.0.1:9090".into(),
719 security_headers: SecurityHeadersConfig::default(),
720 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
721 max_concurrent_tls_handshakes: DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES,
722 extra_route_rate_limit: None,
723 tool_rate_limit_burst: None,
724 extra_route_rate_limit_burst: None,
725 extra_route_rate_limit_exempt_paths: Vec::new(),
726 trusted_proxies: Vec::new(),
727 forwarded_header: None,
728 }
729 }
730
731 #[must_use]
741 pub fn with_auth(mut self, auth: AuthConfig) -> Self {
742 self.auth = Some(auth);
743 self
744 }
745
746 #[must_use]
751 pub fn with_security_headers(mut self, headers: SecurityHeadersConfig) -> Self {
752 self.security_headers = headers;
753 self
754 }
755
756 #[must_use]
760 pub fn with_bind_addr(mut self, addr: impl Into<String>) -> Self {
761 self.bind_addr = addr.into();
762 self
763 }
764
765 #[must_use]
768 pub fn with_rbac(mut self, rbac: Arc<RbacPolicy>) -> Self {
769 self.rbac = Some(rbac);
770 self
771 }
772
773 #[must_use]
777 pub fn with_tls(mut self, cert_path: impl Into<PathBuf>, key_path: impl Into<PathBuf>) -> Self {
778 self.tls_cert_path = Some(cert_path.into());
779 self.tls_key_path = Some(key_path.into());
780 self
781 }
782
783 #[must_use]
787 pub fn with_public_url(mut self, url: impl Into<String>) -> Self {
788 self.public_url = Some(url.into());
789 self
790 }
791
792 #[must_use]
796 pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
797 where
798 I: IntoIterator<Item = S>,
799 S: Into<String>,
800 {
801 self.allowed_origins = origins.into_iter().map(Into::into).collect();
802 self
803 }
804
805 #[must_use]
818 pub fn with_extra_router(mut self, router: axum::Router) -> Self {
819 self.extra_router = Some(router);
820 self
821 }
822
823 #[must_use]
826 pub fn with_readiness_check(mut self, check: ReadinessCheck) -> Self {
827 self.readiness_check = Some(check);
828 self
829 }
830
831 #[must_use]
834 pub fn with_max_request_body(mut self, bytes: usize) -> Self {
835 self.max_request_body = bytes;
836 self
837 }
838
839 #[must_use]
841 pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
842 self.request_timeout = timeout;
843 self
844 }
845
846 #[must_use]
848 pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
849 self.shutdown_timeout = timeout;
850 self
851 }
852
853 #[must_use]
855 pub fn with_session_idle_timeout(mut self, timeout: Duration) -> Self {
856 self.session_idle_timeout = timeout;
857 self
858 }
859
860 #[must_use]
862 pub fn with_sse_keep_alive(mut self, interval: Duration) -> Self {
863 self.sse_keep_alive = interval;
864 self
865 }
866
867 #[must_use]
871 pub fn with_max_concurrent_requests(mut self, limit: usize) -> Self {
872 self.max_concurrent_requests = Some(limit);
873 self
874 }
875
876 #[must_use]
884 pub fn with_tls_handshake_timeout(mut self, timeout: Duration) -> Self {
885 self.tls_handshake_timeout = timeout;
886 self
887 }
888
889 #[must_use]
898 pub fn with_max_concurrent_tls_handshakes(mut self, limit: usize) -> Self {
899 self.max_concurrent_tls_handshakes = limit;
900 self
901 }
902
903 #[must_use]
906 pub fn with_tool_rate_limit(mut self, per_minute: u32) -> Self {
907 self.tool_rate_limit = Some(per_minute);
908 self
909 }
910
911 #[must_use]
922 pub fn with_extra_route_rate_limit(mut self, per_minute: u32) -> Self {
923 self.extra_route_rate_limit = Some(per_minute);
924 self
925 }
926
927 #[must_use]
932 pub fn with_tool_rate_limit_burst(mut self, burst: u32) -> Self {
933 self.tool_rate_limit_burst = Some(burst);
934 self
935 }
936
937 #[must_use]
943 pub fn with_extra_route_rate_limit_burst(mut self, burst: u32) -> Self {
944 self.extra_route_rate_limit_burst = Some(burst);
945 self
946 }
947
948 #[must_use]
968 pub fn with_extra_route_rate_limit_exempt_paths<I, S>(mut self, paths: I) -> Self
969 where
970 I: IntoIterator<Item = S>,
971 S: Into<String>,
972 {
973 self.extra_route_rate_limit_exempt_paths = paths.into_iter().map(Into::into).collect();
974 self
975 }
976
977 #[must_use]
989 pub fn with_trusted_proxies<I, S>(mut self, proxies: I) -> Self
990 where
991 I: IntoIterator<Item = S>,
992 S: Into<String>,
993 {
994 self.trusted_proxies = proxies.into_iter().map(Into::into).collect();
995 self
996 }
997
998 #[must_use]
1003 pub fn with_forwarded_header(mut self, mode: ForwardedHeaderMode) -> Self {
1004 self.forwarded_header = Some(mode);
1005 self
1006 }
1007
1008 #[must_use]
1012 pub fn with_reload_callback<F>(mut self, callback: F) -> Self
1013 where
1014 F: FnOnce(ReloadHandle) + Send + 'static,
1015 {
1016 self.on_reload_ready = Some(Box::new(callback));
1017 self
1018 }
1019
1020 #[must_use]
1024 pub fn enable_compression(mut self, min_size: u16) -> Self {
1025 self.compression_enabled = true;
1026 self.compression_min_size = min_size;
1027 self
1028 }
1029
1030 #[must_use]
1035 pub fn enable_admin(mut self, role: impl Into<String>) -> Self {
1036 self.admin_enabled = true;
1037 self.admin_role = role.into();
1038 self
1039 }
1040
1041 #[must_use]
1044 pub fn enable_request_header_logging(mut self) -> Self {
1045 self.log_request_headers = true;
1046 self
1047 }
1048
1049 #[must_use]
1054 pub fn expose_build_metadata(mut self) -> Self {
1055 self.expose_build_metadata = true;
1056 self
1057 }
1058
1059 #[cfg(feature = "metrics")]
1062 #[must_use]
1063 pub fn with_metrics(mut self, bind: impl Into<String>) -> Self {
1064 self.metrics_enabled = true;
1065 self.metrics_bind = bind.into();
1066 self
1067 }
1068
1069 pub fn validate(self) -> Result<Validated<Self>, RmcpServerKitError> {
1102 self.check()?;
1103 Ok(Validated(self))
1104 }
1105
1106 fn check_burst_knobs(&self) -> Result<(), RmcpServerKitError> {
1113 if self.tool_rate_limit_burst == Some(0) {
1114 return Err(RmcpServerKitError::Config(
1115 "tool_rate_limit_burst must be greater than zero".into(),
1116 ));
1117 }
1118 if self.extra_route_rate_limit_burst == Some(0) {
1119 return Err(RmcpServerKitError::Config(
1120 "extra_route_rate_limit_burst must be greater than zero".into(),
1121 ));
1122 }
1123 if self.tool_rate_limit_burst.is_some() && self.tool_rate_limit.is_none() {
1124 return Err(RmcpServerKitError::Config(
1125 "tool_rate_limit_burst requires tool_rate_limit to be set".into(),
1126 ));
1127 }
1128 if self.extra_route_rate_limit_burst.is_some() && self.extra_route_rate_limit.is_none() {
1129 return Err(RmcpServerKitError::Config(
1130 "extra_route_rate_limit_burst requires extra_route_rate_limit to be set".into(),
1131 ));
1132 }
1133 if !self.extra_route_rate_limit_exempt_paths.is_empty()
1134 && self.extra_route_rate_limit.is_none()
1135 {
1136 return Err(RmcpServerKitError::Config(
1137 "extra_route_rate_limit_exempt_paths requires extra_route_rate_limit to be set"
1138 .into(),
1139 ));
1140 }
1141 for path in &self.extra_route_rate_limit_exempt_paths {
1142 if path.is_empty() || !path.starts_with('/') {
1143 return Err(RmcpServerKitError::Config(format!(
1144 "extra_route_rate_limit_exempt_paths entries must be non-empty and start with '/': {path:?}"
1145 )));
1146 }
1147 }
1148 if let Some(rl) = self.auth.as_ref().and_then(|a| a.rate_limit.as_ref()) {
1149 if rl.burst == Some(0) {
1150 return Err(RmcpServerKitError::Config(
1151 "auth rate_limit.burst must be greater than zero".into(),
1152 ));
1153 }
1154 if rl.pre_auth_burst == Some(0) {
1155 return Err(RmcpServerKitError::Config(
1156 "auth rate_limit.pre_auth_burst must be greater than zero".into(),
1157 ));
1158 }
1159 }
1160 Ok(())
1161 }
1162
1163 fn check_trusted_forwarder(&self) -> Result<(), RmcpServerKitError> {
1168 for entry in &self.trusted_proxies {
1169 validate_trusted_proxy_entry(entry).map_err(RmcpServerKitError::Config)?;
1170 }
1171 if self.forwarded_header.is_some() && self.trusted_proxies.is_empty() {
1172 return Err(RmcpServerKitError::Config(
1173 "forwarded_header requires trusted_proxies to be nonempty".into(),
1174 ));
1175 }
1176 Ok(())
1177 }
1178
1179 fn check(&self) -> Result<(), RmcpServerKitError> {
1183 if self.admin_enabled {
1187 let auth_enabled = self.auth.as_ref().is_some_and(|a| a.enabled);
1188 if !auth_enabled {
1189 return Err(RmcpServerKitError::Config(
1190 "admin_enabled=true requires auth to be configured and enabled".into(),
1191 ));
1192 }
1193 }
1194
1195 match (&self.tls_cert_path, &self.tls_key_path) {
1197 (Some(_), None) => {
1198 return Err(RmcpServerKitError::Config(
1199 "tls_cert_path is set but tls_key_path is missing".into(),
1200 ));
1201 }
1202 (None, Some(_)) => {
1203 return Err(RmcpServerKitError::Config(
1204 "tls_key_path is set but tls_cert_path is missing".into(),
1205 ));
1206 }
1207 _ => {}
1208 }
1209
1210 if self.bind_addr.parse::<SocketAddr>().is_err() {
1212 return Err(RmcpServerKitError::Config(format!(
1213 "bind_addr {:?} is not a valid socket address (expected e.g. 127.0.0.1:8080)",
1214 self.bind_addr
1215 )));
1216 }
1217
1218 if let Some(ref url) = self.public_url
1220 && !(url.starts_with("http://") || url.starts_with("https://"))
1221 {
1222 return Err(RmcpServerKitError::Config(format!(
1223 "public_url {url:?} must start with http:// or https://"
1224 )));
1225 }
1226
1227 for origin in &self.allowed_origins {
1229 if !(origin.starts_with("http://") || origin.starts_with("https://")) {
1230 return Err(RmcpServerKitError::Config(format!(
1231 "allowed_origins entry {origin:?} must start with http:// or https://"
1232 )));
1233 }
1234 }
1235
1236 if self.max_request_body == 0 {
1238 return Err(RmcpServerKitError::Config(
1239 "max_request_body must be greater than zero".into(),
1240 ));
1241 }
1242
1243 if self.extra_route_rate_limit == Some(0) {
1247 return Err(RmcpServerKitError::Config(
1248 "extra_route_rate_limit must be greater than zero".into(),
1249 ));
1250 }
1251
1252 self.check_burst_knobs()?;
1254
1255 self.check_trusted_forwarder()?;
1257
1258 #[cfg(feature = "oauth")]
1260 if let Some(auth_cfg) = &self.auth
1261 && let Some(oauth_cfg) = &auth_cfg.oauth
1262 {
1263 oauth_cfg.validate()?;
1264 }
1265
1266 validate_security_headers(&self.security_headers)?;
1269
1270 if self.max_concurrent_requests == Some(0) {
1274 return Err(RmcpServerKitError::Config(
1275 "max_concurrent_requests must be greater than zero when set".into(),
1276 ));
1277 }
1278
1279 if let Some(auth_cfg) = &self.auth
1283 && let Some(rl) = &auth_cfg.rate_limit
1284 && rl.max_tracked_keys == 0
1285 {
1286 return Err(RmcpServerKitError::Config(
1287 "auth.rate_limit.max_tracked_keys must be greater than zero".into(),
1288 ));
1289 }
1290
1291 if self.tls_handshake_timeout == Duration::ZERO {
1296 return Err(RmcpServerKitError::Config(
1297 "tls_handshake_timeout must be greater than zero".into(),
1298 ));
1299 }
1300
1301 if self.max_concurrent_tls_handshakes == 0 {
1306 return Err(RmcpServerKitError::Config(
1307 "max_concurrent_tls_handshakes must be greater than zero".into(),
1308 ));
1309 }
1310
1311 Ok(())
1312 }
1313}
1314
1315#[allow(
1321 missing_debug_implementations,
1322 reason = "contains Arc<AuthState> with non-Debug fields"
1323)]
1324pub struct ReloadHandle {
1325 auth: Option<Arc<AuthState>>,
1326 rbac: Option<Arc<ArcSwap<RbacPolicy>>>,
1327 crl_set: Option<Arc<CrlSet>>,
1328}
1329
1330impl ReloadHandle {
1331 pub fn reload_auth_keys(&self, keys: Vec<crate::auth::ApiKeyEntry>) {
1333 if let Some(ref auth) = self.auth {
1334 auth.reload_keys(keys);
1335 }
1336 }
1337
1338 pub fn reload_rbac(&self, policy: RbacPolicy) {
1340 if let Some(ref rbac) = self.rbac {
1341 rbac.store(Arc::new(policy));
1342 tracing::info!("RBAC policy reloaded");
1343 }
1344 }
1345
1346 pub async fn refresh_crls(&self) -> Result<(), RmcpServerKitError> {
1352 let Some(ref crl_set) = self.crl_set else {
1353 return Err(RmcpServerKitError::Config(
1354 "CRL refresh requested but mTLS CRL support is not configured".into(),
1355 ));
1356 };
1357
1358 crl_set.force_refresh().await
1359 }
1360}
1361
1362#[allow(
1379 clippy::too_many_lines,
1380 clippy::cognitive_complexity,
1381 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"
1382)]
1383struct AppRunParams {
1387 tls_paths: Option<(PathBuf, PathBuf)>,
1389 tls_handshake_timeout: Duration,
1391 max_concurrent_tls_handshakes: usize,
1393 mtls_config: Option<MtlsConfig>,
1395 shutdown_timeout: Duration,
1397 auth_state: Option<Arc<AuthState>>,
1399 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
1401 on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
1403 ct: CancellationToken,
1407 session_ct: CancellationToken,
1417 scheme: &'static str,
1419 name: String,
1421}
1422
1423#[allow(
1433 clippy::cognitive_complexity,
1434 reason = "router assembly is intrinsically sequential; splitting harms readability"
1435)]
1436#[allow(
1437 deprecated,
1438 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1439)]
1440fn build_app_router<H, F>(
1441 mut config: McpServerConfig,
1442 handler_factory: F,
1443) -> anyhow::Result<(axum::Router, AppRunParams)>
1444where
1445 H: ServerHandler + 'static,
1446 F: Fn() -> H + Send + Sync + Clone + 'static,
1447{
1448 let ct = CancellationToken::new();
1449 let session_ct = CancellationToken::new();
1450
1451 let allowed_hosts = derive_allowed_hosts(&config.bind_addr, config.public_url.as_deref());
1452 tracing::info!(allowed_hosts = %allowed_hosts.join(", "), "configured Streamable HTTP allowed hosts");
1453
1454 let mcp_service = StreamableHttpService::new(
1455 move || Ok(handler_factory()),
1456 {
1457 let mut mgr = LocalSessionManager::default();
1458 mgr.session_config.keep_alive = Some(config.session_idle_timeout);
1459 mgr.into()
1460 },
1461 StreamableHttpServerConfig::default()
1462 .with_allowed_hosts(allowed_hosts)
1463 .with_sse_keep_alive(Some(config.sse_keep_alive))
1464 .with_cancellation_token(session_ct.clone()),
1465 );
1466
1467 let mut mcp_router = axum::Router::new().nest_service("/mcp", mcp_service);
1469
1470 let auth_state: Option<Arc<AuthState>> = match config.auth {
1474 Some(ref auth_config) if auth_config.enabled => {
1475 let rate_limiter = auth_config.rate_limit.as_ref().map(build_rate_limiter);
1476 let pre_auth_limiter = auth_config
1477 .rate_limit
1478 .as_ref()
1479 .map(crate::auth::build_pre_auth_limiter);
1480
1481 #[cfg(feature = "oauth")]
1482 let jwks_cache = auth_config
1483 .oauth
1484 .as_ref()
1485 .map(|c| crate::oauth::JwksCache::new(c).map(Arc::new))
1486 .transpose()
1487 .map_err(|e| std::io::Error::other(format!("JWKS HTTP client: {e}")))?;
1488
1489 Some(Arc::new(AuthState {
1490 api_keys: ArcSwap::new(Arc::new(auth_config.api_keys.clone())),
1491 rate_limiter,
1492 pre_auth_limiter,
1493 #[cfg(feature = "oauth")]
1494 jwks_cache,
1495 seen_identities: crate::auth::SeenIdentitySet::new(),
1496 counters: crate::auth::AuthCounters::default(),
1497 }))
1498 }
1499 _ => None,
1500 };
1501
1502 let rbac_swap = Arc::new(ArcSwap::new(
1505 config
1506 .rbac
1507 .clone()
1508 .unwrap_or_else(|| Arc::new(RbacPolicy::disabled())),
1509 ));
1510
1511 if config.admin_enabled {
1514 let Some(ref auth_state_ref) = auth_state else {
1515 return Err(anyhow::anyhow!(
1516 "admin_enabled=true requires auth to be configured and enabled"
1517 ));
1518 };
1519 let admin_state = crate::admin::AdminState {
1520 started_at: std::time::Instant::now(),
1521 name: config.name.clone(),
1522 version: config.version.clone(),
1523 auth: Some(Arc::clone(auth_state_ref)),
1524 rbac: Arc::clone(&rbac_swap),
1525 };
1526 let admin_cfg = crate::admin::AdminConfig {
1527 role: config.admin_role.clone(),
1528 };
1529 mcp_router = mcp_router.merge(crate::admin::admin_router(admin_state, &admin_cfg));
1530 tracing::info!(role = %config.admin_role, "/admin/* endpoints enabled");
1531 }
1532
1533 {
1566 let tool_limiter: Option<Arc<ToolRateLimiter>> = config
1567 .tool_rate_limit
1568 .map(|per_minute| build_tool_rate_limiter(per_minute, config.tool_rate_limit_burst));
1569
1570 if rbac_swap.load().is_enabled() {
1571 tracing::info!("RBAC enforcement enabled on /mcp");
1572 }
1573 if let Some(limit) = config.tool_rate_limit {
1574 tracing::info!(limit, "tool rate limiting enabled (calls/min per IP)");
1575 }
1576
1577 let rbac_for_mw = Arc::clone(&rbac_swap);
1578 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1579 let p = rbac_for_mw.load_full();
1580 let tl = tool_limiter.clone();
1581 rbac_middleware(p, tl, req, next)
1582 }));
1583 }
1584
1585 if let Some(ref auth_config) = config.auth
1587 && auth_config.enabled
1588 {
1589 let Some(ref state) = auth_state else {
1590 return Err(anyhow::anyhow!("auth state missing despite enabled config"));
1591 };
1592
1593 let methods: Vec<&str> = [
1594 auth_config.mtls.is_some().then_some("mTLS"),
1595 (!auth_config.api_keys.is_empty()).then_some("bearer"),
1596 #[cfg(feature = "oauth")]
1597 auth_config.oauth.is_some().then_some("oauth-jwt"),
1598 ]
1599 .into_iter()
1600 .flatten()
1601 .collect();
1602
1603 tracing::info!(
1604 methods = %methods.join(", "),
1605 api_keys = auth_config.api_keys.len(),
1606 "auth enabled on /mcp"
1607 );
1608
1609 let state_for_mw = Arc::clone(state);
1610 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1611 let s = Arc::clone(&state_for_mw);
1612 auth_middleware(s, req, next)
1613 }));
1614 }
1615
1616 mcp_router = mcp_router.layer(tower_http::timeout::TimeoutLayer::with_status_code(
1619 axum::http::StatusCode::REQUEST_TIMEOUT,
1620 config.request_timeout,
1621 ));
1622
1623 mcp_router = mcp_router.layer(tower_http::limit::RequestBodyLimitLayer::new(
1627 config.max_request_body,
1628 ));
1629
1630 let mut effective_origins = config.allowed_origins.clone();
1637 if effective_origins.is_empty()
1638 && let Some(ref url) = config.public_url
1639 {
1640 if let Some(scheme_end) = url.find("://") {
1645 let scheme_with_sep = url.get(..scheme_end + 3).unwrap_or_default();
1646 let after_scheme = url.get(scheme_end + 3..).unwrap_or_default();
1647 let host_end = after_scheme.find('/').unwrap_or(after_scheme.len());
1648 let host = after_scheme.get(..host_end).unwrap_or_default();
1649 let origin = format!("{scheme_with_sep}{host}");
1650 tracing::info!(
1651 %origin,
1652 "auto-derived allowed origin from public_url"
1653 );
1654 effective_origins.push(origin);
1655 }
1656 }
1657 let allowed_origins: Arc<[String]> = Arc::from(effective_origins);
1658 let cors_origins = Arc::clone(&allowed_origins);
1659 let log_request_headers = config.log_request_headers;
1660
1661 let readyz_route = if let Some(check) = config.readiness_check.take() {
1662 axum::routing::get(move || readyz(Arc::clone(&check)))
1663 } else {
1664 axum::routing::get(healthz)
1665 };
1666
1667 #[allow(unused_mut)] let mut router = axum::Router::new()
1669 .route("/healthz", axum::routing::get(healthz))
1670 .route("/readyz", readyz_route)
1671 .route(
1672 "/version",
1673 axum::routing::get({
1674 let payload_bytes: Arc<[u8]> = serialize_version_payload(
1679 &config.name,
1680 &config.version,
1681 config.expose_build_metadata,
1682 );
1683 move || {
1684 let p = Arc::clone(&payload_bytes);
1685 async move {
1686 (
1687 [(axum::http::header::CONTENT_TYPE, "application/json")],
1688 p.to_vec(),
1689 )
1690 }
1691 }
1692 }),
1693 )
1694 .merge(mcp_router);
1695
1696 if let Some(extra) = config.extra_router.take() {
1703 let extra = match config.extra_route_rate_limit {
1704 Some(per_minute) => {
1705 let limiter =
1706 build_extra_route_rate_limiter(per_minute, config.extra_route_rate_limit_burst);
1707 let exempt: Arc<std::collections::HashSet<String>> = Arc::new(
1708 config
1709 .extra_route_rate_limit_exempt_paths
1710 .iter()
1711 .cloned()
1712 .collect(),
1713 );
1714 tracing::info!(
1715 per_minute,
1716 exempt_paths = exempt.len(),
1717 "extra-route per-IP rate limit enabled"
1718 );
1719 extra.layer(axum::middleware::from_fn(move |req, next| {
1720 let l = Arc::clone(&limiter);
1721 let e = Arc::clone(&exempt);
1722 extra_route_rate_limit_middleware(l, e, req, next)
1723 }))
1724 }
1725 None => extra,
1726 };
1727 router = router.merge(extra);
1728 }
1729
1730 let server_url = if let Some(ref url) = config.public_url {
1737 url.trim_end_matches('/').to_owned()
1738 } else {
1739 let prm_scheme = if config.tls_cert_path.is_some() {
1740 "https"
1741 } else {
1742 "http"
1743 };
1744 format!("{prm_scheme}://{}", config.bind_addr)
1745 };
1746 let resource_url = format!("{server_url}/mcp");
1747
1748 #[cfg(feature = "oauth")]
1749 let prm_metadata = if let Some(ref auth_config) = config.auth
1750 && let Some(ref oauth_config) = auth_config.oauth
1751 {
1752 crate::oauth::protected_resource_metadata(&resource_url, &server_url, oauth_config)
1753 } else {
1754 serde_json::json!({ "resource": resource_url })
1755 };
1756 #[cfg(not(feature = "oauth"))]
1757 let prm_metadata = serde_json::json!({ "resource": resource_url });
1758
1759 router = router.route(
1760 "/.well-known/oauth-protected-resource",
1761 axum::routing::get(move || {
1762 let m = prm_metadata.clone();
1763 async move { axum::Json(m) }
1764 }),
1765 );
1766
1767 #[cfg(feature = "oauth")]
1772 if let Some(ref auth_config) = config.auth
1773 && let Some(ref oauth_config) = auth_config.oauth
1774 && oauth_config.proxy.is_some()
1775 {
1776 router = install_oauth_proxy_routes(
1777 router,
1778 &server_url,
1779 oauth_config,
1780 auth_state.as_ref(),
1781 config.max_request_body,
1782 &config.admin_role,
1783 )?;
1784 }
1785
1786 if !cors_origins.is_empty() {
1795 let cors = tower_http::cors::CorsLayer::new()
1796 .allow_origin(
1797 cors_origins
1798 .iter()
1799 .filter_map(|o| o.parse::<axum::http::HeaderValue>().ok())
1800 .collect::<Vec<_>>(),
1801 )
1802 .allow_methods([
1803 axum::http::Method::GET,
1804 axum::http::Method::POST,
1805 axum::http::Method::OPTIONS,
1806 ])
1807 .allow_headers([
1808 axum::http::header::CONTENT_TYPE,
1809 axum::http::header::AUTHORIZATION,
1810 ]);
1811 router = router.layer(cors);
1812 }
1813
1814 if config.compression_enabled {
1818 use tower_http::compression::Predicate as _;
1819 let predicate = tower_http::compression::DefaultPredicate::new().and(
1820 tower_http::compression::predicate::SizeAbove::new(u64::from(
1821 config.compression_min_size,
1822 )),
1823 );
1824 router = router.layer(
1825 tower_http::compression::CompressionLayer::new()
1826 .gzip(true)
1827 .br(true)
1828 .compress_when(predicate),
1829 );
1830 tracing::info!(
1831 min_size = config.compression_min_size,
1832 "response compression enabled (gzip, br)"
1833 );
1834 }
1835
1836 if let Some(max) = config.max_concurrent_requests {
1839 let overload_handler = tower::ServiceBuilder::new()
1840 .layer(axum::error_handling::HandleErrorLayer::new(
1841 |_err: tower::BoxError| async {
1842 (
1843 axum::http::StatusCode::SERVICE_UNAVAILABLE,
1844 axum::Json(serde_json::json!({
1845 "error": "overloaded",
1846 "error_description": "server is at capacity, retry later"
1847 })),
1848 )
1849 },
1850 ))
1851 .layer(tower::load_shed::LoadShedLayer::new())
1852 .layer(tower::limit::ConcurrencyLimitLayer::new(max));
1853 router = router.layer(overload_handler);
1854 tracing::info!(max, "global concurrency limit enabled");
1855 }
1856
1857 router = router.fallback(|| async {
1861 (
1862 axum::http::StatusCode::NOT_FOUND,
1863 axum::Json(serde_json::json!({
1864 "error": "not_found",
1865 "error_description": "The requested endpoint does not exist"
1866 })),
1867 )
1868 });
1869
1870 #[cfg(feature = "metrics")]
1872 if config.metrics_enabled {
1873 let metrics = Arc::new(
1874 crate::metrics::McpMetrics::new().map_err(|e| anyhow::anyhow!("metrics init: {e}"))?,
1875 );
1876 let m = Arc::clone(&metrics);
1877 router = router.layer(axum::middleware::from_fn(
1878 move |req: Request<Body>, next: Next| {
1879 let m = Arc::clone(&m);
1880 metrics_middleware(m, req, next)
1881 },
1882 ));
1883 let metrics_bind = config.metrics_bind.clone();
1884 let metrics_shutdown = ct.clone();
1885 tokio::spawn(async move {
1886 if let Err(e) =
1887 crate::metrics::serve_metrics(metrics_bind, metrics, metrics_shutdown).await
1888 {
1889 tracing::error!("metrics listener failed: {e}");
1890 }
1891 });
1892 }
1893
1894 let forward_resolver: Option<Arc<ForwardResolver>> = if config.trusted_proxies.is_empty() {
1902 None
1903 } else {
1904 Some(Arc::new(ForwardResolver {
1907 trusted: config
1908 .trusted_proxies
1909 .iter()
1910 .filter_map(|entry| parse_proxy_net(entry))
1911 .collect(),
1912 mode: config
1913 .forwarded_header
1914 .unwrap_or(ForwardedHeaderMode::XForwardedFor),
1915 }))
1916 };
1917 if forward_resolver.is_some() {
1918 tracing::info!(
1919 proxies = config.trusted_proxies.len(),
1920 "trusted-forwarder mode enabled: limiters key by resolved client IP"
1921 );
1922 }
1923 router = router.layer(axum::middleware::from_fn(move |req, next| {
1924 let r = forward_resolver.clone();
1925 normalize_peer_addr_middleware(r, req, next)
1926 }));
1927
1928 router = router.layer(axum::middleware::from_fn(move |req, next| {
1940 let origins = Arc::clone(&allowed_origins);
1941 origin_check_middleware(origins, log_request_headers, req, next)
1942 }));
1943
1944 let is_tls = config.tls_cert_path.is_some();
1953 warn_security_header_overrides(&config.security_headers);
1954 let security_headers_cfg = Arc::new(config.security_headers.clone());
1955 router = router.layer(axum::middleware::from_fn(move |req, next| {
1956 let cfg = Arc::clone(&security_headers_cfg);
1957 security_headers_middleware(is_tls, cfg, req, next)
1958 }));
1959
1960 let scheme = if config.tls_cert_path.is_some() {
1961 "https"
1962 } else {
1963 "http"
1964 };
1965
1966 let tls_paths = match (&config.tls_cert_path, &config.tls_key_path) {
1967 (Some(cert), Some(key)) => Some((cert.clone(), key.clone())),
1968 _ => None,
1969 };
1970 let tls_handshake_timeout = config.tls_handshake_timeout;
1971 let max_concurrent_tls_handshakes = config.max_concurrent_tls_handshakes;
1972 let mtls_config = config.auth.as_ref().and_then(|a| a.mtls.as_ref()).cloned();
1973
1974 Ok((
1975 router,
1976 AppRunParams {
1977 tls_paths,
1978 tls_handshake_timeout,
1979 max_concurrent_tls_handshakes,
1980 mtls_config,
1981 shutdown_timeout: config.shutdown_timeout,
1982 auth_state,
1983 rbac_swap,
1984 on_reload_ready: config.on_reload_ready.take(),
1985 ct,
1986 session_ct,
1987 scheme,
1988 name: config.name.clone(),
1989 },
1990 ))
1991}
1992
1993pub async fn serve<H, F>(
2010 config: Validated<McpServerConfig>,
2011 handler_factory: F,
2012) -> Result<(), RmcpServerKitError>
2013where
2014 H: ServerHandler + 'static,
2015 F: Fn() -> H + Send + Sync + Clone + 'static,
2016{
2017 let config = config.into_inner();
2018 #[allow(
2019 deprecated,
2020 reason = "internal serve() reads `bind_addr` to construct the listener; field becomes pub(crate) in 1.0"
2021 )]
2022 let bind_addr = config.bind_addr.clone();
2023 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2024
2025 let listener = TcpListener::bind(&bind_addr)
2026 .await
2027 .map_err(|e| io_to_startup(&format!("bind {bind_addr}"), e))?;
2028 log_listening(¶ms.name, params.scheme, &bind_addr);
2029
2030 run_server(
2031 router,
2032 listener,
2033 params.tls_paths,
2034 params.tls_handshake_timeout,
2035 params.max_concurrent_tls_handshakes,
2036 params.mtls_config,
2037 params.shutdown_timeout,
2038 params.auth_state,
2039 params.rbac_swap,
2040 params.on_reload_ready,
2041 params.ct,
2042 params.session_ct,
2043 )
2044 .await
2045 .map_err(anyhow_to_startup)
2046}
2047
2048pub async fn serve_with_listener<H, F>(
2078 listener: TcpListener,
2079 config: Validated<McpServerConfig>,
2080 handler_factory: F,
2081 ready_tx: Option<tokio::sync::oneshot::Sender<SocketAddr>>,
2082 shutdown: Option<CancellationToken>,
2083) -> Result<(), RmcpServerKitError>
2084where
2085 H: ServerHandler + 'static,
2086 F: Fn() -> H + Send + Sync + Clone + 'static,
2087{
2088 let config = config.into_inner();
2089 let local_addr = listener
2090 .local_addr()
2091 .map_err(|e| io_to_startup("listener.local_addr", e))?;
2092 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2093
2094 log_listening(¶ms.name, params.scheme, &local_addr.to_string());
2095
2096 if let Some(external) = shutdown {
2100 let internal = params.ct.clone();
2101 tokio::spawn(async move {
2102 external.cancelled().await;
2103 internal.cancel();
2104 });
2105 }
2106
2107 if let Some(tx) = ready_tx {
2111 let _ = tx.send(local_addr);
2113 }
2114
2115 run_server(
2116 router,
2117 listener,
2118 params.tls_paths,
2119 params.tls_handshake_timeout,
2120 params.max_concurrent_tls_handshakes,
2121 params.mtls_config,
2122 params.shutdown_timeout,
2123 params.auth_state,
2124 params.rbac_swap,
2125 params.on_reload_ready,
2126 params.ct,
2127 params.session_ct,
2128 )
2129 .await
2130 .map_err(anyhow_to_startup)
2131}
2132
2133#[allow(
2136 clippy::cognitive_complexity,
2137 reason = "tracing::info! macro expansions inflate the score; logic is trivial"
2138)]
2139fn log_listening(name: &str, scheme: &str, addr: &str) {
2140 tracing::info!("{name} listening on {addr}");
2141 tracing::info!(" MCP endpoint: {scheme}://{addr}/mcp");
2142 tracing::info!(" Health check: {scheme}://{addr}/healthz");
2143 tracing::info!(" Readiness: {scheme}://{addr}/readyz");
2144}
2145
2146#[allow(
2169 clippy::too_many_arguments,
2170 clippy::cognitive_complexity,
2171 reason = "server start-up threads TLS, reload state, and graceful shutdown through one flow"
2172)]
2173async fn run_server(
2174 router: axum::Router,
2175 listener: TcpListener,
2176 tls_paths: Option<(PathBuf, PathBuf)>,
2177 tls_handshake_timeout: Duration,
2178 max_concurrent_tls_handshakes: usize,
2179 mtls_config: Option<MtlsConfig>,
2180 shutdown_timeout: Duration,
2181 auth_state: Option<Arc<AuthState>>,
2182 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
2183 mut on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
2184 ct: CancellationToken,
2185 session_ct: CancellationToken,
2186) -> anyhow::Result<()> {
2187 let shutdown_trigger = CancellationToken::new();
2191 {
2192 let trigger = shutdown_trigger.clone();
2193 let parent = ct.clone();
2194 tokio::spawn(async move {
2195 tokio::select! {
2198 () = shutdown_signal() => {}
2199 () = parent.cancelled() => {}
2200 }
2201 trigger.cancel();
2202 });
2203 }
2204
2205 let graceful = {
2206 let trigger = shutdown_trigger.clone();
2207 let ct = ct.clone();
2208 async move {
2209 trigger.cancelled().await;
2210 tracing::info!("shutting down (grace period: {shutdown_timeout:?})");
2211 ct.cancel();
2212 }
2213 };
2214
2215 let force_exit_timer = {
2216 let trigger = shutdown_trigger.clone();
2217 async move {
2218 trigger.cancelled().await;
2219 tokio::time::sleep(shutdown_timeout).await;
2220 }
2221 };
2222
2223 if let Some((cert_path, key_path)) = tls_paths {
2224 let crl_set = if let Some(mtls) = mtls_config.as_ref()
2225 && mtls.crl_enabled
2226 {
2227 let (ca_certs, roots) = load_client_auth_roots(&mtls.ca_cert_path)?;
2228 let (crl_set, discover_rx) =
2229 mtls_revocation::bootstrap_fetch(roots, &ca_certs, mtls.clone())
2230 .await
2231 .map_err(|error| anyhow::anyhow!(error.to_string()))?;
2232 tokio::spawn(mtls_revocation::run_crl_refresher(
2233 Arc::clone(&crl_set),
2234 discover_rx,
2235 ct.clone(),
2236 ));
2237 Some(crl_set)
2238 } else {
2239 None
2240 };
2241
2242 if let Some(cb) = on_reload_ready.take() {
2243 cb(ReloadHandle {
2244 auth: auth_state.clone(),
2245 rbac: Some(Arc::clone(&rbac_swap)),
2246 crl_set: crl_set.clone(),
2247 });
2248 }
2249
2250 let tls_listener = TlsListener::new(
2251 listener,
2252 &cert_path,
2253 &key_path,
2254 mtls_config.as_ref(),
2255 crl_set,
2256 tls_handshake_timeout,
2257 max_concurrent_tls_handshakes,
2258 )?;
2259 let make_svc = router.into_make_service_with_connect_info::<TlsConnInfo>();
2260 tokio::select! {
2263 result = axum::serve(tls_listener, make_svc)
2264 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2265 () = force_exit_timer => {
2266 tracing::warn!("shutdown timeout exceeded, forcing exit");
2267 session_ct.cancel();
2268 }
2269 }
2270 } else {
2271 if let Some(cb) = on_reload_ready.take() {
2272 cb(ReloadHandle {
2273 auth: auth_state,
2274 rbac: Some(rbac_swap),
2275 crl_set: None,
2276 });
2277 }
2278
2279 let make_svc = router.into_make_service_with_connect_info::<SocketAddr>();
2280 tokio::select! {
2283 result = axum::serve(listener, make_svc)
2284 .with_graceful_shutdown(graceful) => { session_ct.cancel(); result?; }
2285 () = force_exit_timer => {
2286 tracing::warn!("shutdown timeout exceeded, forcing exit");
2287 session_ct.cancel();
2288 }
2289 }
2290 }
2291
2292 Ok(())
2293}
2294
2295#[cfg(feature = "oauth")]
2304fn install_oauth_proxy_routes(
2305 router: axum::Router,
2306 server_url: &str,
2307 oauth_config: &crate::oauth::OAuthConfig,
2308 auth_state: Option<&Arc<AuthState>>,
2309 max_request_body: usize,
2310 admin_role: &str,
2311) -> Result<axum::Router, RmcpServerKitError> {
2312 let Some(ref proxy) = oauth_config.proxy else {
2313 return Ok(router);
2314 };
2315
2316 let http = crate::oauth::OauthHttpClient::with_config(oauth_config)?;
2319
2320 let proxy_router = axum::Router::new();
2326
2327 let asm = crate::oauth::authorization_server_metadata(server_url, oauth_config);
2328 let proxy_router = proxy_router.route(
2329 "/.well-known/oauth-authorization-server",
2330 axum::routing::get(move || {
2331 let m = asm.clone();
2332 async move { axum::Json(m) }
2333 }),
2334 );
2335
2336 let proxy_authorize = proxy.clone();
2337 let proxy_router = proxy_router.route(
2338 "/authorize",
2339 axum::routing::get(
2340 move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
2341 let p = proxy_authorize.clone();
2342 async move { crate::oauth::handle_authorize(&p, &query.unwrap_or_default()) }
2343 },
2344 ),
2345 );
2346
2347 let proxy_token = proxy.clone();
2348 let token_http = http.clone();
2349 let proxy_router = proxy_router.route(
2350 "/token",
2351 axum::routing::post(move |body: String| {
2352 let p = proxy_token.clone();
2353 let h = token_http.clone();
2354 async move { crate::oauth::handle_token(&h, &p, &body).await }
2355 })
2356 .layer(axum::middleware::from_fn(
2357 oauth_token_cache_headers_middleware,
2358 )),
2359 );
2360
2361 let proxy_register = proxy.clone();
2362 let proxy_router = proxy_router.route(
2363 "/register",
2364 axum::routing::post(move |axum::Json(body): axum::Json<serde_json::Value>| {
2365 let p = proxy_register;
2366 async move { axum::Json(crate::oauth::handle_register(&p, &body)) }
2367 })
2368 .layer(axum::middleware::from_fn(
2369 oauth_token_cache_headers_middleware,
2370 )),
2371 );
2372
2373 let admin_routes_enabled = proxy.expose_admin_endpoints
2374 && (proxy.introspection_url.is_some() || proxy.revocation_url.is_some());
2375 if proxy.expose_admin_endpoints
2376 && !proxy.require_auth_on_admin_endpoints
2377 && proxy.allow_unauthenticated_admin_endpoints
2378 {
2379 tracing::warn!(
2383 "OAuth introspect/revoke endpoints are unauthenticated by explicit \
2384 allow_unauthenticated_admin_endpoints opt-out; ensure an \
2385 authenticated reverse proxy fronts these routes"
2386 );
2387 }
2388
2389 let admin_router = if admin_routes_enabled {
2390 build_oauth_admin_router(proxy, http, auth_state, admin_role)?
2391 } else {
2392 axum::Router::new()
2393 };
2394
2395 let proxy_router =
2399 proxy_router
2400 .merge(admin_router)
2401 .layer(tower_http::limit::RequestBodyLimitLayer::new(
2402 max_request_body,
2403 ));
2404
2405 let router = router.merge(proxy_router);
2406
2407 tracing::info!(
2408 introspect = proxy.expose_admin_endpoints && proxy.introspection_url.is_some(),
2409 revoke = proxy.expose_admin_endpoints && proxy.revocation_url.is_some(),
2410 max_request_body,
2411 "OAuth 2.1 proxy endpoints enabled (/authorize, /token, /register)"
2412 );
2413 Ok(router)
2414}
2415
2416#[cfg(feature = "oauth")]
2422fn build_oauth_admin_router(
2423 proxy: &crate::oauth::OAuthProxyConfig,
2424 http: crate::oauth::OauthHttpClient,
2425 auth_state: Option<&Arc<AuthState>>,
2426 admin_role: &str,
2427) -> Result<axum::Router, RmcpServerKitError> {
2428 let mut admin_router = axum::Router::new();
2429 if proxy.introspection_url.is_some() {
2430 let proxy_introspect = proxy.clone();
2431 let introspect_http = http.clone();
2432 admin_router = admin_router.route(
2433 "/introspect",
2434 axum::routing::post(move |body: String| {
2435 let p = proxy_introspect.clone();
2436 let h = introspect_http.clone();
2437 async move { crate::oauth::handle_introspect(&h, &p, &body).await }
2438 }),
2439 );
2440 }
2441 if proxy.revocation_url.is_some() {
2442 let proxy_revoke = proxy.clone();
2443 let revoke_http = http;
2444 admin_router = admin_router.route(
2445 "/revoke",
2446 axum::routing::post(move |body: String| {
2447 let p = proxy_revoke.clone();
2448 let h = revoke_http.clone();
2449 async move { crate::oauth::handle_revoke(&h, &p, &body).await }
2450 }),
2451 );
2452 }
2453
2454 let admin_router = admin_router.layer(axum::middleware::from_fn(
2455 oauth_token_cache_headers_middleware,
2456 ));
2457
2458 if proxy.require_auth_on_admin_endpoints {
2459 let Some(state) = auth_state else {
2460 return Err(RmcpServerKitError::Startup(
2461 "oauth proxy admin endpoints require auth state".into(),
2462 ));
2463 };
2464 let state_for_mw = Arc::clone(state);
2465 let required_role: Arc<str> = Arc::from(admin_role);
2466 Ok(admin_router
2472 .layer(axum::middleware::from_fn(move |req, next| {
2473 let r = Arc::clone(&required_role);
2474 crate::admin::require_admin_role(r, req, next)
2475 }))
2476 .layer(axum::middleware::from_fn(move |req, next| {
2477 let s = Arc::clone(&state_for_mw);
2478 auth_middleware(s, req, next)
2479 })))
2480 } else {
2481 Ok(admin_router)
2482 }
2483}
2484
2485fn derive_allowed_hosts(bind_addr: &str, public_url: Option<&str>) -> Vec<String> {
2490 let mut hosts = vec![
2491 "localhost".to_owned(),
2492 "127.0.0.1".to_owned(),
2493 "::1".to_owned(),
2494 ];
2495
2496 if let Some(url) = public_url
2497 && let Ok(uri) = url.parse::<axum::http::Uri>()
2498 && let Some(authority) = uri.authority()
2499 {
2500 let host = authority.host().to_owned();
2501 if !hosts.iter().any(|h| h == &host) {
2502 hosts.push(host);
2503 }
2504
2505 let authority = authority.as_str().to_owned();
2506 if !hosts.iter().any(|h| h == &authority) {
2507 hosts.push(authority);
2508 }
2509 }
2510
2511 if let Ok(uri) = format!("http://{bind_addr}").parse::<axum::http::Uri>()
2512 && let Some(authority) = uri.authority()
2513 {
2514 let host = authority.host().to_owned();
2515 if !hosts.iter().any(|h| h == &host) {
2516 hosts.push(host);
2517 }
2518
2519 let authority = authority.as_str().to_owned();
2520 if !hosts.iter().any(|h| h == &authority) {
2521 hosts.push(authority);
2522 }
2523 }
2524
2525 hosts
2526}
2527
2528impl axum::extract::connect_info::Connected<axum::serve::IncomingStream<'_, TlsListener>>
2541 for TlsConnInfo
2542{
2543 fn connect_info(target: axum::serve::IncomingStream<'_, TlsListener>) -> Self {
2544 let addr = *target.remote_addr();
2545 let identity = target.io().identity().cloned();
2546 Self::new(addr, identity)
2547 }
2548}
2549
2550const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
2557
2558const DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256;
2566
2567const TLS_ACCEPT_CHANNEL_CAPACITY: usize = 32;
2572
2573struct TlsListener {
2589 local_addr: SocketAddr,
2592 rx: mpsc::Receiver<(AuthenticatedTlsStream, SocketAddr)>,
2594 acceptor_task: tokio::task::JoinHandle<()>,
2597}
2598
2599impl TlsListener {
2600 fn new(
2601 inner: TcpListener,
2602 cert_path: &Path,
2603 key_path: &Path,
2604 mtls_config: Option<&MtlsConfig>,
2605 crl_set: Option<Arc<CrlSet>>,
2606 handshake_timeout: Duration,
2607 max_concurrent_handshakes: usize,
2608 ) -> anyhow::Result<Self> {
2609 rustls::crypto::ring::default_provider()
2611 .install_default()
2612 .ok();
2613
2614 let certs = load_certs(cert_path)?;
2615 let key = load_key(key_path)?;
2616
2617 let mtls_default_role;
2618
2619 let tls_config = if let Some(mtls) = mtls_config {
2620 mtls_default_role = mtls.default_role.clone();
2621 let verifier: Arc<dyn rustls::server::danger::ClientCertVerifier> = if mtls.crl_enabled
2622 {
2623 let Some(crl_set) = crl_set else {
2624 return Err(anyhow::anyhow!(
2625 "mTLS CRL verifier requested but CRL state was not initialized"
2626 ));
2627 };
2628 Arc::new(DynamicClientCertVerifier::new(crl_set))
2629 } else {
2630 let (_, root_store) = load_client_auth_roots(&mtls.ca_cert_path)?;
2631 if mtls.required {
2632 rustls::server::WebPkiClientVerifier::builder(root_store)
2633 .build()
2634 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2635 } else {
2636 rustls::server::WebPkiClientVerifier::builder(root_store)
2637 .allow_unauthenticated()
2638 .build()
2639 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2640 }
2641 };
2642
2643 tracing::info!(
2644 ca = %mtls.ca_cert_path.display(),
2645 required = mtls.required,
2646 crl_enabled = mtls.crl_enabled,
2647 "mTLS client auth configured"
2648 );
2649
2650 rustls::ServerConfig::builder_with_protocol_versions(&[
2651 &rustls::version::TLS12,
2652 &rustls::version::TLS13,
2653 ])
2654 .with_client_cert_verifier(verifier)
2655 .with_single_cert(certs, key)?
2656 } else {
2657 mtls_default_role = "viewer".to_owned();
2658 rustls::ServerConfig::builder_with_protocol_versions(&[
2659 &rustls::version::TLS12,
2660 &rustls::version::TLS13,
2661 ])
2662 .with_no_client_auth()
2663 .with_single_cert(certs, key)?
2664 };
2665
2666 let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
2667 tracing::info!(
2668 "TLS enabled (cert: {}, key: {})",
2669 cert_path.display(),
2670 key_path.display()
2671 );
2672 let local_addr = inner.local_addr()?;
2673 let (tx, rx) = mpsc::channel(TLS_ACCEPT_CHANNEL_CAPACITY);
2674 let acceptor_task = tokio::spawn(run_tls_acceptor(
2675 inner,
2676 acceptor,
2677 mtls_default_role,
2678 tx,
2679 handshake_timeout,
2680 max_concurrent_handshakes,
2681 ));
2682 Ok(Self {
2683 local_addr,
2684 rx,
2685 acceptor_task,
2686 })
2687 }
2688
2689 fn extract_handshake_identity(
2693 tls_stream: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2694 default_role: &str,
2695 addr: SocketAddr,
2696 ) -> Option<AuthIdentity> {
2697 let (_, server_conn) = tls_stream.get_ref();
2698 let cert_der = server_conn.peer_certificates()?.first()?;
2699 let id = extract_mtls_identity(cert_der.as_ref(), default_role)?;
2700 tracing::debug!(name = %id.name, peer = %addr, "mTLS client cert accepted");
2701 Some(id)
2702 }
2703}
2704
2705async fn run_tls_acceptor(
2713 listener: TcpListener,
2714 acceptor: tokio_rustls::TlsAcceptor,
2715 default_role: String,
2716 tx: mpsc::Sender<(AuthenticatedTlsStream, SocketAddr)>,
2717 handshake_timeout: Duration,
2718 max_concurrent_handshakes: usize,
2719) {
2720 let inflight = Arc::new(Semaphore::new(max_concurrent_handshakes));
2721 loop {
2722 let Ok(permit) = Arc::clone(&inflight).acquire_owned().await else {
2726 return;
2728 };
2729 let (stream, addr) = match listener.accept().await {
2730 Ok(pair) => pair,
2731 Err(e) => {
2732 tracing::debug!("TCP accept error: {e}");
2733 continue;
2734 }
2735 };
2736 if tx.is_closed() {
2737 return;
2739 }
2740 let acceptor = acceptor.clone();
2741 let default_role = default_role.clone();
2742 let tx = tx.clone();
2743 tokio::spawn(async move {
2744 let _permit = permit;
2745 match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
2746 Ok(Ok(tls_stream)) => {
2747 let identity =
2748 TlsListener::extract_handshake_identity(&tls_stream, &default_role, addr);
2749 let wrapped = AuthenticatedTlsStream {
2750 inner: tls_stream,
2751 identity,
2752 };
2753 let _ = tx.send((wrapped, addr)).await;
2756 }
2757 Ok(Err(e)) => {
2758 tracing::debug!("TLS handshake failed from {addr}: {e}");
2759 }
2760 Err(_elapsed) => {
2761 tracing::debug!(
2762 "TLS handshake timed out from {addr} after {handshake_timeout:?}"
2763 );
2764 }
2765 }
2766 });
2767 }
2768}
2769
2770pub(crate) struct AuthenticatedTlsStream {
2782 inner: tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2783 identity: Option<AuthIdentity>,
2784}
2785
2786impl AuthenticatedTlsStream {
2787 #[must_use]
2789 pub(crate) const fn identity(&self) -> Option<&AuthIdentity> {
2790 self.identity.as_ref()
2791 }
2792}
2793
2794impl std::fmt::Debug for AuthenticatedTlsStream {
2795 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2796 f.debug_struct("AuthenticatedTlsStream")
2797 .field("identity", &self.identity.as_ref().map(|id| &id.name))
2798 .finish_non_exhaustive()
2799 }
2800}
2801
2802impl tokio::io::AsyncRead for AuthenticatedTlsStream {
2803 fn poll_read(
2804 mut self: Pin<&mut Self>,
2805 cx: &mut std::task::Context<'_>,
2806 buf: &mut tokio::io::ReadBuf<'_>,
2807 ) -> std::task::Poll<std::io::Result<()>> {
2808 Pin::new(&mut self.inner).poll_read(cx, buf)
2809 }
2810}
2811
2812impl tokio::io::AsyncWrite for AuthenticatedTlsStream {
2813 fn poll_write(
2814 mut self: Pin<&mut Self>,
2815 cx: &mut std::task::Context<'_>,
2816 buf: &[u8],
2817 ) -> std::task::Poll<std::io::Result<usize>> {
2818 Pin::new(&mut self.inner).poll_write(cx, buf)
2819 }
2820
2821 fn poll_flush(
2822 mut self: Pin<&mut Self>,
2823 cx: &mut std::task::Context<'_>,
2824 ) -> std::task::Poll<std::io::Result<()>> {
2825 Pin::new(&mut self.inner).poll_flush(cx)
2826 }
2827
2828 fn poll_shutdown(
2829 mut self: Pin<&mut Self>,
2830 cx: &mut std::task::Context<'_>,
2831 ) -> std::task::Poll<std::io::Result<()>> {
2832 Pin::new(&mut self.inner).poll_shutdown(cx)
2833 }
2834
2835 fn poll_write_vectored(
2836 mut self: Pin<&mut Self>,
2837 cx: &mut std::task::Context<'_>,
2838 bufs: &[std::io::IoSlice<'_>],
2839 ) -> std::task::Poll<std::io::Result<usize>> {
2840 Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
2841 }
2842
2843 fn is_write_vectored(&self) -> bool {
2844 self.inner.is_write_vectored()
2845 }
2846}
2847
2848impl axum::serve::Listener for TlsListener {
2849 type Io = AuthenticatedTlsStream;
2850 type Addr = SocketAddr;
2851
2852 async fn accept(&mut self) -> (Self::Io, Self::Addr) {
2858 if let Some(pair) = self.rx.recv().await {
2859 return pair;
2860 }
2861 tracing::error!("TLS acceptor task terminated; no further connections will be accepted");
2867 std::future::pending().await
2868 }
2869
2870 fn local_addr(&self) -> std::io::Result<Self::Addr> {
2871 Ok(self.local_addr)
2872 }
2873}
2874
2875impl Drop for TlsListener {
2876 fn drop(&mut self) {
2877 self.acceptor_task.abort();
2880 }
2881}
2882
2883fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
2884 use rustls::pki_types::pem::PemObject;
2885 let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_file_iter(path)
2886 .map_err(|e| anyhow::anyhow!("failed to read certs from {}: {e}", path.display()))?
2887 .collect::<Result<_, _>>()
2888 .map_err(|e| anyhow::anyhow!("invalid cert in {}: {e}", path.display()))?;
2889 anyhow::ensure!(
2890 !certs.is_empty(),
2891 "no certificates found in {}",
2892 path.display()
2893 );
2894 Ok(certs)
2895}
2896
2897fn load_client_auth_roots(
2898 path: &Path,
2899) -> anyhow::Result<(
2900 Vec<rustls::pki_types::CertificateDer<'static>>,
2901 Arc<RootCertStore>,
2902)> {
2903 let ca_certs = load_certs(path)?;
2904 let mut root_store = RootCertStore::empty();
2905 for cert in &ca_certs {
2906 root_store
2907 .add(cert.clone())
2908 .map_err(|error| anyhow::anyhow!("invalid CA cert: {error}"))?;
2909 }
2910
2911 Ok((ca_certs, Arc::new(root_store)))
2912}
2913
2914fn load_key(path: &Path) -> anyhow::Result<rustls::pki_types::PrivateKeyDer<'static>> {
2915 use rustls::pki_types::pem::PemObject;
2916 rustls::pki_types::PrivateKeyDer::from_pem_file(path)
2917 .map_err(|e| anyhow::anyhow!("failed to read key from {}: {e}", path.display()))
2918}
2919
2920#[allow(
2921 clippy::unused_async,
2922 reason = "axum route handler signature requires `async fn` even when the body is synchronous"
2923)]
2924async fn healthz() -> impl IntoResponse {
2925 axum::Json(serde_json::json!({
2926 "status": "ok",
2927 }))
2928}
2929
2930fn version_payload(name: &str, version: &str, expose_build_metadata: bool) -> serde_json::Value {
2940 let mut map = serde_json::Map::new();
2941 map.insert("name".into(), name.into());
2942 map.insert("version".into(), version.into());
2943 map.insert(
2944 "rmcp_server_kit_version".into(),
2945 env!("CARGO_PKG_VERSION").into(),
2946 );
2947 if expose_build_metadata {
2948 map.insert(
2949 "build_git_sha".into(),
2950 option_env!("RMCP_SERVER_KIT_BUILD_SHA")
2951 .unwrap_or("unknown")
2952 .into(),
2953 );
2954 map.insert(
2955 "build_timestamp".into(),
2956 option_env!("RMCP_SERVER_KIT_BUILD_TIME")
2957 .unwrap_or("unknown")
2958 .into(),
2959 );
2960 map.insert(
2961 "rust_version".into(),
2962 option_env!("RMCP_SERVER_KIT_RUSTC_VERSION")
2963 .unwrap_or("unknown")
2964 .into(),
2965 );
2966 }
2967 serde_json::Value::Object(map)
2968}
2969
2970fn serialize_version_payload(name: &str, version: &str, expose_build_metadata: bool) -> Arc<[u8]> {
2980 let value = version_payload(name, version, expose_build_metadata);
2981 serde_json::to_vec(&value).map_or_else(|_| Arc::from(&b"{}"[..]), Arc::from)
2982}
2983
2984async fn readyz(check: ReadinessCheck) -> impl IntoResponse {
2985 let status = check().await;
2986 let ready = status
2987 .get("ready")
2988 .and_then(serde_json::Value::as_bool)
2989 .unwrap_or(false);
2990 let code = if ready {
2991 axum::http::StatusCode::OK
2992 } else {
2993 axum::http::StatusCode::SERVICE_UNAVAILABLE
2994 };
2995 (code, axum::Json(status))
2996}
2997
2998async fn shutdown_signal() {
3002 let ctrl_c = tokio::signal::ctrl_c();
3003
3004 #[cfg(unix)]
3005 {
3006 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
3007 Ok(mut term) => {
3008 tokio::select! {
3011 _ = ctrl_c => {}
3012 _ = term.recv() => {}
3013 }
3014 }
3015 Err(e) => {
3016 tracing::warn!(error = %e, "failed to register SIGTERM handler, using SIGINT only");
3017 ctrl_c.await.ok();
3018 }
3019 }
3020 }
3021
3022 #[cfg(not(unix))]
3023 {
3024 ctrl_c.await.ok();
3025 }
3026}
3027
3028#[cfg(feature = "metrics")]
3045fn metrics_labels(req: &Request<Body>) -> (&'static str, String) {
3046 let method = match *req.method() {
3047 axum::http::Method::GET => "GET",
3048 axum::http::Method::POST => "POST",
3049 axum::http::Method::PUT => "PUT",
3050 axum::http::Method::PATCH => "PATCH",
3051 axum::http::Method::DELETE => "DELETE",
3052 axum::http::Method::HEAD => "HEAD",
3053 axum::http::Method::OPTIONS => "OPTIONS",
3054 axum::http::Method::TRACE => "TRACE",
3055 axum::http::Method::CONNECT => "CONNECT",
3056 _ => "OTHER",
3059 };
3060
3061 let path = req
3062 .extensions()
3063 .get::<axum::extract::MatchedPath>()
3064 .map_or_else(
3065 || {
3066 let raw = req.uri().path();
3067 if raw == "/mcp" || raw.starts_with("/mcp/") {
3068 "/mcp".to_owned()
3069 } else {
3070 "<unmatched>".to_owned()
3071 }
3072 },
3073 |matched| matched.as_str().to_owned(),
3074 );
3075
3076 (method, path)
3077}
3078
3079#[cfg(feature = "metrics")]
3086async fn metrics_middleware(
3087 metrics: Arc<crate::metrics::McpMetrics>,
3088 mut req: Request<Body>,
3089 next: Next,
3090) -> axum::response::Response {
3091 let (method, path) = metrics_labels(&req);
3092 let start = std::time::Instant::now();
3093
3094 req.extensions_mut().insert(Arc::clone(&metrics));
3095 let response = next.run(req).await;
3096
3097 let mut status_buf = core::fmt::NumBuffer::<u16>::new();
3098 let status = response.status().as_u16().format_into(&mut status_buf);
3099 let duration = start.elapsed().as_secs_f64();
3100
3101 metrics
3102 .http_requests_total
3103 .with_label_values(&[method, &path, status])
3104 .inc();
3105 metrics
3106 .http_request_duration_seconds
3107 .with_label_values(&[method, &path])
3108 .observe(duration);
3109
3110 response
3111}
3112
3113async fn security_headers_middleware(
3125 is_tls: bool,
3126 cfg: Arc<SecurityHeadersConfig>,
3127 req: Request<Body>,
3128 next: Next,
3129) -> axum::response::Response {
3130 use axum::http::{HeaderName, header};
3131
3132 let mut resp = next.run(req).await;
3133 let headers = resp.headers_mut();
3134
3135 headers.remove(header::SERVER);
3137 headers.remove(HeaderName::from_static("x-powered-by"));
3138
3139 apply_security_header(
3140 headers,
3141 header::X_CONTENT_TYPE_OPTIONS,
3142 cfg.x_content_type_options.as_deref(),
3143 "nosniff",
3144 );
3145 apply_security_header(
3146 headers,
3147 header::X_FRAME_OPTIONS,
3148 cfg.x_frame_options.as_deref(),
3149 "deny",
3150 );
3151 apply_security_header(
3152 headers,
3153 header::CACHE_CONTROL,
3154 cfg.cache_control.as_deref(),
3155 "no-store, max-age=0",
3156 );
3157 apply_security_header(
3158 headers,
3159 header::REFERRER_POLICY,
3160 cfg.referrer_policy.as_deref(),
3161 "no-referrer",
3162 );
3163 apply_security_header(
3164 headers,
3165 HeaderName::from_static("cross-origin-opener-policy"),
3166 cfg.cross_origin_opener_policy.as_deref(),
3167 "same-origin",
3168 );
3169 apply_security_header(
3170 headers,
3171 HeaderName::from_static("cross-origin-resource-policy"),
3172 cfg.cross_origin_resource_policy.as_deref(),
3173 "same-origin",
3174 );
3175 apply_security_header(
3176 headers,
3177 HeaderName::from_static("cross-origin-embedder-policy"),
3178 cfg.cross_origin_embedder_policy.as_deref(),
3179 "require-corp",
3180 );
3181 apply_security_header(
3182 headers,
3183 HeaderName::from_static("permissions-policy"),
3184 cfg.permissions_policy.as_deref(),
3185 "accelerometer=(), camera=(), geolocation=(), microphone=()",
3186 );
3187 apply_security_header(
3188 headers,
3189 HeaderName::from_static("x-permitted-cross-domain-policies"),
3190 cfg.x_permitted_cross_domain_policies.as_deref(),
3191 "none",
3192 );
3193 apply_security_header(
3194 headers,
3195 HeaderName::from_static("content-security-policy"),
3196 cfg.content_security_policy.as_deref(),
3197 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests",
3198 );
3199 apply_security_header(
3200 headers,
3201 HeaderName::from_static("x-dns-prefetch-control"),
3202 cfg.x_dns_prefetch_control.as_deref(),
3203 "off",
3204 );
3205
3206 if is_tls {
3207 apply_security_header(
3208 headers,
3209 header::STRICT_TRANSPORT_SECURITY,
3210 cfg.strict_transport_security.as_deref(),
3211 "max-age=63072000; includeSubDomains",
3212 );
3213 }
3214
3215 resp
3216}
3217
3218fn apply_security_header(
3229 headers: &mut axum::http::HeaderMap,
3230 name: axum::http::HeaderName,
3231 override_value: Option<&str>,
3232 default: &'static str,
3233) {
3234 use axum::http::HeaderValue;
3235
3236 match override_value {
3237 None => {
3238 headers.insert(name, HeaderValue::from_static(default));
3239 }
3240 Some("") => {
3241 }
3243 Some(v) => match HeaderValue::from_str(v) {
3244 Ok(hv) => {
3245 headers.insert(name, hv);
3246 }
3247 Err(err) => {
3248 tracing::error!(
3249 header = %name,
3250 error = %err,
3251 "invalid security header override reached middleware; using default"
3252 );
3253 headers.insert(name, HeaderValue::from_static(default));
3254 }
3255 },
3256 }
3257}
3258
3259fn validate_security_headers(cfg: &SecurityHeadersConfig) -> Result<(), RmcpServerKitError> {
3270 use axum::http::HeaderValue;
3271
3272 let fields: &[(&str, Option<&str>)] = &[
3273 (
3274 "x_content_type_options",
3275 cfg.x_content_type_options.as_deref(),
3276 ),
3277 ("x_frame_options", cfg.x_frame_options.as_deref()),
3278 ("cache_control", cfg.cache_control.as_deref()),
3279 ("referrer_policy", cfg.referrer_policy.as_deref()),
3280 (
3281 "cross_origin_opener_policy",
3282 cfg.cross_origin_opener_policy.as_deref(),
3283 ),
3284 (
3285 "cross_origin_resource_policy",
3286 cfg.cross_origin_resource_policy.as_deref(),
3287 ),
3288 (
3289 "cross_origin_embedder_policy",
3290 cfg.cross_origin_embedder_policy.as_deref(),
3291 ),
3292 ("permissions_policy", cfg.permissions_policy.as_deref()),
3293 (
3294 "x_permitted_cross_domain_policies",
3295 cfg.x_permitted_cross_domain_policies.as_deref(),
3296 ),
3297 (
3298 "content_security_policy",
3299 cfg.content_security_policy.as_deref(),
3300 ),
3301 (
3302 "x_dns_prefetch_control",
3303 cfg.x_dns_prefetch_control.as_deref(),
3304 ),
3305 (
3306 "strict_transport_security",
3307 cfg.strict_transport_security.as_deref(),
3308 ),
3309 ];
3310
3311 for (field, value) in fields {
3312 let Some(v) = value else { continue };
3313 if v.is_empty() {
3314 continue;
3315 }
3316 if let Err(err) = HeaderValue::from_str(v) {
3317 return Err(RmcpServerKitError::Config(format!(
3318 "invalid security_headers.{field}: {err}"
3319 )));
3320 }
3321 }
3322
3323 if let Some(v) = cfg.strict_transport_security.as_deref()
3324 && !v.is_empty()
3325 && v.to_ascii_lowercase().contains("preload")
3326 {
3327 return Err(RmcpServerKitError::Config(format!(
3328 "invalid security_headers.strict_transport_security: {v:?} contains the `preload` directive; \
3329 HSTS preload must be opted into explicitly via a dedicated builder, not via this knob"
3330 )));
3331 }
3332
3333 Ok(())
3334}
3335
3336#[cfg(feature = "oauth")]
3351async fn oauth_token_cache_headers_middleware(
3352 req: Request<Body>,
3353 next: Next,
3354) -> axum::response::Response {
3355 use axum::http::{HeaderValue, header};
3356
3357 let mut resp = next.run(req).await;
3358 let headers = resp.headers_mut();
3359 headers.insert(header::PRAGMA, HeaderValue::from_static("no-cache"));
3360 headers.append(header::VARY, HeaderValue::from_static("Authorization"));
3361 resp
3362}
3363
3364async fn normalize_peer_addr_middleware(
3393 resolver: Option<Arc<ForwardResolver>>,
3394 mut req: Request<Body>,
3395 next: Next,
3396) -> axum::response::Response {
3397 let direct = req
3398 .extensions()
3399 .get::<ConnectInfo<SocketAddr>>()
3400 .map(|ci| ci.0);
3401 let from_tls = req
3402 .extensions()
3403 .get::<ConnectInfo<TlsConnInfo>>()
3404 .map(|ci| ci.0.addr);
3405 if let Some(addr) = direct.or(from_tls) {
3406 if direct.is_none() {
3407 req.extensions_mut().insert(ConnectInfo(addr));
3408 }
3409 req.extensions_mut().insert(PeerAddr::new(addr));
3410 let client_ip = match &resolver {
3411 Some(r) => {
3412 crate::forwarded::resolve_client_ip(addr.ip(), req.headers(), &r.trusted, r.mode)
3413 .unwrap_or_else(|reason| {
3414 tracing::debug!(
3415 reason = ?reason,
3416 "forwarded-header resolution fell back to direct peer"
3417 );
3418 addr.ip()
3419 })
3420 }
3421 None => addr.ip(),
3422 };
3423 req.extensions_mut().insert(ClientIp::new(client_ip));
3424 }
3425 next.run(req).await
3426}
3427
3428fn parse_proxy_net(entry: &str) -> Option<ipnet::IpNet> {
3431 if let Ok(net) = entry.parse::<ipnet::IpNet>() {
3432 return Some(net);
3433 }
3434 entry.parse::<IpAddr>().ok().map(ipnet::IpNet::from)
3435}
3436
3437pub(crate) fn validate_trusted_proxy_entry(entry: &str) -> Result<(), String> {
3447 match parse_proxy_net(entry) {
3448 None => Err(format!(
3449 "trusted_proxies entry {entry:?} is neither a CIDR nor an IP address"
3450 )),
3451 Some(net) if net.prefix_len() == 0 => Err(format!(
3452 "trusted_proxies entry {entry:?}: prefix length 0 is forbidden (marks every peer trusted, enabling client-IP spoofing)"
3453 )),
3454 Some(_) => Ok(()),
3455 }
3456}
3457
3458pub(crate) fn limiter_client_ip(extensions: &axum::http::Extensions) -> Option<IpAddr> {
3462 if let Some(client) = extensions.get::<ClientIp>() {
3463 return Some(client.ip);
3464 }
3465 extensions
3466 .get::<ConnectInfo<SocketAddr>>()
3467 .map(|ci| ci.0.ip())
3468 .or_else(|| {
3469 extensions
3470 .get::<ConnectInfo<TlsConnInfo>>()
3471 .map(|ci| ci.0.addr.ip())
3472 })
3473}
3474
3475pub(crate) type ExtraRouteRateLimiter = BoundedKeyedLimiter<IpAddr>;
3479
3480const EXTRA_ROUTE_MAX_TRACKED_KEYS: usize = 10_000;
3486
3487const EXTRA_ROUTE_IDLE_EVICTION: Duration = Duration::from_mins(15);
3490
3491fn build_extra_route_rate_limiter(
3498 per_minute: u32,
3499 burst: Option<u32>,
3500) -> Arc<ExtraRouteRateLimiter> {
3501 let rate = std::num::NonZeroU32::new(per_minute.max(1)).unwrap_or(std::num::NonZeroU32::MIN);
3502 let mut quota = governor::Quota::per_minute(rate);
3503 if let Some(b) = burst.and_then(std::num::NonZeroU32::new) {
3504 quota = quota.allow_burst(b);
3505 }
3506 Arc::new(BoundedKeyedLimiter::new(
3507 quota,
3508 EXTRA_ROUTE_MAX_TRACKED_KEYS,
3509 EXTRA_ROUTE_IDLE_EVICTION,
3510 ))
3511}
3512
3513async fn extra_route_rate_limit_middleware(
3535 limiter: Arc<ExtraRouteRateLimiter>,
3536 exempt: Arc<std::collections::HashSet<String>>,
3537 req: Request<Body>,
3538 next: Next,
3539) -> axum::response::Response {
3540 if exempt.contains(req.uri().path()) {
3541 return next.run(req).await;
3542 }
3543 let peer_ip: Option<IpAddr> = limiter_client_ip(req.extensions());
3544 if let Some(ip) = peer_ip
3545 && let Err(wait) = limiter.check_key_wait(&ip)
3546 {
3547 #[cfg(feature = "metrics")]
3548 crate::metrics::record_rate_limit_deny(req.extensions(), "extra_route");
3549 tracing::warn!(%ip, "extra route request rate limited");
3550 return RmcpServerKitError::RateLimitedFor {
3551 message: "too many requests to application routes from this source".into(),
3552 retry_after: wait,
3553 }
3554 .into_response();
3555 }
3556 next.run(req).await
3557}
3558
3559async fn origin_check_middleware(
3563 allowed: Arc<[String]>,
3564 log_request_headers: bool,
3565 req: Request<Body>,
3566 next: Next,
3567) -> axum::response::Response {
3568 let method = req.method().clone();
3569 let path = req.uri().path().to_owned();
3570
3571 log_incoming_request(&method, &path, req.headers(), log_request_headers);
3572
3573 if let Some(origin) = req.headers().get(axum::http::header::ORIGIN) {
3574 let origin_str = origin.to_str().unwrap_or("");
3575 if !allowed.iter().any(|a| a == origin_str) {
3576 tracing::warn!(
3577 origin = origin_str,
3578 %method,
3579 %path,
3580 allowed = ?&*allowed,
3581 "rejected request: Origin not allowed"
3582 );
3583 return (
3584 axum::http::StatusCode::FORBIDDEN,
3585 "Forbidden: Origin not allowed",
3586 )
3587 .into_response();
3588 }
3589 }
3590 next.run(req).await
3591}
3592
3593fn log_incoming_request(
3596 method: &axum::http::Method,
3597 path: &str,
3598 headers: &axum::http::HeaderMap,
3599 log_request_headers: bool,
3600) {
3601 if log_request_headers {
3602 tracing::debug!(
3603 %method,
3604 %path,
3605 headers = %format_request_headers_for_log(headers),
3606 "incoming request"
3607 );
3608 } else {
3609 tracing::debug!(%method, %path, "incoming request");
3610 }
3611}
3612
3613fn format_request_headers_for_log(headers: &axum::http::HeaderMap) -> String {
3614 headers
3615 .iter()
3616 .map(|(k, v)| {
3617 let name = k.as_str();
3618 if name == "authorization" || name == "cookie" || name == "proxy-authorization" {
3619 format!("{name}: [REDACTED]")
3620 } else {
3621 format!("{name}: {}", v.to_str().unwrap_or("<non-utf8>"))
3622 }
3623 })
3624 .collect::<Vec<_>>()
3625 .join(", ")
3626}
3627
3628#[allow(
3652 clippy::cognitive_complexity,
3653 reason = "complexity is purely tracing macro expansion (info/warn + match arms); 18 lines of straight-line code, nothing meaningful to extract"
3654)]
3655pub async fn serve_stdio<H>(handler: H) -> Result<(), RmcpServerKitError>
3656where
3657 H: ServerHandler + 'static,
3658{
3659 use rmcp::ServiceExt as _;
3660
3661 tracing::info!("stdio transport: serving on stdin/stdout");
3662 tracing::warn!("stdio mode: auth, RBAC, TLS, and Origin checks are DISABLED");
3663
3664 let transport = rmcp::transport::io::stdio();
3665
3666 let service = handler
3667 .serve(transport)
3668 .await
3669 .map_err(|e| RmcpServerKitError::Startup(format!("stdio initialize failed: {e}")))?;
3670
3671 if let Err(e) = service.waiting().await {
3672 tracing::warn!(error = %e, "stdio session ended with error");
3673 }
3674 tracing::info!("stdio session ended");
3675 Ok(())
3676}
3677
3678#[allow(
3679 deprecated,
3680 reason = "builder methods are the sanctioned transition layer for deprecated public fields"
3681)]
3682impl McpServerConfig {
3683 #[must_use]
3687 pub fn with_tls_paths(mut self, cert_path: Option<PathBuf>, key_path: Option<PathBuf>) -> Self {
3688 self.tls_cert_path = cert_path;
3689 self.tls_key_path = key_path;
3690 self
3691 }
3692
3693 #[must_use]
3697 pub fn with_tls_cert_path(mut self, cert_path: impl Into<PathBuf>) -> Self {
3698 self.tls_cert_path = Some(cert_path.into());
3699 self
3700 }
3701
3702 #[must_use]
3706 pub fn with_tls_key_path(mut self, key_path: impl Into<PathBuf>) -> Self {
3707 self.tls_key_path = Some(key_path.into());
3708 self
3709 }
3710
3711 #[must_use]
3713 pub fn with_optional_auth(mut self, auth: Option<AuthConfig>) -> Self {
3714 self.auth = auth;
3715 self
3716 }
3717
3718 #[must_use]
3720 pub fn with_optional_tool_rate_limit(mut self, per_minute: Option<u32>) -> Self {
3721 self.tool_rate_limit = per_minute;
3722 self
3723 }
3724
3725 #[must_use]
3727 pub fn with_optional_tool_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
3728 self.tool_rate_limit_burst = burst;
3729 self
3730 }
3731
3732 #[must_use]
3734 pub fn with_optional_extra_route_rate_limit(mut self, per_minute: Option<u32>) -> Self {
3735 self.extra_route_rate_limit = per_minute;
3736 self
3737 }
3738
3739 #[must_use]
3741 pub fn with_optional_extra_route_rate_limit_burst(mut self, burst: Option<u32>) -> Self {
3742 self.extra_route_rate_limit_burst = burst;
3743 self
3744 }
3745
3746 #[must_use]
3748 pub fn with_optional_forwarded_header(mut self, mode: Option<ForwardedHeaderMode>) -> Self {
3749 self.forwarded_header = mode;
3750 self
3751 }
3752
3753 #[must_use]
3755 pub fn with_optional_public_url(mut self, url: Option<String>) -> Self {
3756 self.public_url = url;
3757 self
3758 }
3759
3760 #[must_use]
3764 pub fn with_compression_min_size(mut self, min_size: u16) -> Self {
3765 self.compression_min_size = min_size;
3766 self
3767 }
3768
3769 #[must_use]
3771 pub fn with_compression_enabled(mut self, enabled: bool) -> Self {
3772 self.compression_enabled = enabled;
3773 self
3774 }
3775
3776 #[must_use]
3778 pub fn with_optional_max_concurrent_requests(mut self, limit: Option<usize>) -> Self {
3779 self.max_concurrent_requests = limit;
3780 self
3781 }
3782
3783 #[must_use]
3785 pub fn with_admin_enabled(mut self, enabled: bool) -> Self {
3786 self.admin_enabled = enabled;
3787 self
3788 }
3789
3790 #[must_use]
3793 pub fn with_admin_role(mut self, role: impl Into<String>) -> Self {
3794 self.admin_role = role.into();
3795 self
3796 }
3797
3798 #[must_use]
3800 pub fn with_expose_build_metadata(mut self, enabled: bool) -> Self {
3801 self.expose_build_metadata = enabled;
3802 self
3803 }
3804}
3805
3806fn warn_security_header_overrides(cfg: &SecurityHeadersConfig) {
3807 for (field, value) in security_header_overrides(cfg) {
3808 let action = if value.is_empty() {
3809 "omitted"
3810 } else {
3811 "overridden"
3812 };
3813 tracing::warn!(
3814 security_header = field,
3815 action,
3816 "security header configured; inspect server.security_headers.<security_header>"
3817 );
3818 }
3819}
3820
3821fn security_header_overrides(
3822 cfg: &SecurityHeadersConfig,
3823) -> impl Iterator<Item = (&'static str, &str)> {
3824 [
3825 (
3826 "x_content_type_options",
3827 cfg.x_content_type_options.as_deref(),
3828 ),
3829 ("x_frame_options", cfg.x_frame_options.as_deref()),
3830 ("cache_control", cfg.cache_control.as_deref()),
3831 ("referrer_policy", cfg.referrer_policy.as_deref()),
3832 (
3833 "cross_origin_opener_policy",
3834 cfg.cross_origin_opener_policy.as_deref(),
3835 ),
3836 (
3837 "cross_origin_resource_policy",
3838 cfg.cross_origin_resource_policy.as_deref(),
3839 ),
3840 (
3841 "cross_origin_embedder_policy",
3842 cfg.cross_origin_embedder_policy.as_deref(),
3843 ),
3844 ("permissions_policy", cfg.permissions_policy.as_deref()),
3845 (
3846 "x_permitted_cross_domain_policies",
3847 cfg.x_permitted_cross_domain_policies.as_deref(),
3848 ),
3849 (
3850 "content_security_policy",
3851 cfg.content_security_policy.as_deref(),
3852 ),
3853 (
3854 "x_dns_prefetch_control",
3855 cfg.x_dns_prefetch_control.as_deref(),
3856 ),
3857 (
3858 "strict_transport_security",
3859 cfg.strict_transport_security.as_deref(),
3860 ),
3861 ]
3862 .into_iter()
3863 .filter_map(|(field, value)| value.map(|v| (field, v)))
3864}
3865
3866#[cfg(test)]
3867mod tests {
3868 #![allow(
3869 clippy::unwrap_used,
3870 clippy::expect_used,
3871 clippy::panic,
3872 clippy::indexing_slicing,
3873 clippy::unwrap_in_result,
3874 clippy::print_stdout,
3875 clippy::print_stderr,
3876 deprecated,
3877 reason = "internal unit tests legitimately read/write the deprecated `pub` fields they were designed to verify"
3878 )]
3879 use std::{sync::Arc, time::Duration};
3880
3881 use axum::{
3882 body::Body,
3883 http::{Request, StatusCode, header},
3884 response::IntoResponse,
3885 };
3886 use http_body_util::BodyExt;
3887 use tower::ServiceExt as _;
3888
3889 use super::*;
3890
3891 #[test]
3894 fn server_config_new_defaults() {
3895 let cfg = McpServerConfig::new("0.0.0.0:8443", "test-server", "1.0.0");
3896 assert_eq!(cfg.bind_addr, "0.0.0.0:8443");
3897 assert_eq!(cfg.name, "test-server");
3898 assert_eq!(cfg.version, "1.0.0");
3899 assert!(cfg.tls_cert_path.is_none());
3900 assert!(cfg.tls_key_path.is_none());
3901 assert!(cfg.auth.is_none());
3902 assert!(cfg.rbac.is_none());
3903 assert!(cfg.allowed_origins.is_empty());
3904 assert!(cfg.tool_rate_limit.is_none());
3905 assert!(cfg.readiness_check.is_none());
3906 assert_eq!(cfg.max_request_body, 1024 * 1024);
3907 assert_eq!(cfg.request_timeout, Duration::from_mins(2));
3908 assert_eq!(cfg.shutdown_timeout, Duration::from_secs(30));
3909 assert!(!cfg.log_request_headers);
3910 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(10));
3911 assert_eq!(cfg.max_concurrent_tls_handshakes, 256);
3912 }
3913
3914 #[test]
3915 fn tls_handshake_builders_set_fields() {
3916 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3917 .with_tls_handshake_timeout(Duration::from_secs(3))
3918 .with_max_concurrent_tls_handshakes(64);
3919 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(3));
3920 assert_eq!(cfg.max_concurrent_tls_handshakes, 64);
3921 }
3922
3923 #[test]
3924 fn validate_rejects_zero_tls_handshake_timeout() {
3925 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3926 .with_tls_handshake_timeout(Duration::ZERO);
3927 let err = cfg.validate().expect_err("zero handshake timeout");
3928 assert!(err.to_string().contains("tls_handshake_timeout"));
3929 }
3930
3931 #[test]
3932 fn validate_rejects_zero_max_concurrent_tls_handshakes() {
3933 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3934 .with_max_concurrent_tls_handshakes(0);
3935 let err = cfg.validate().expect_err("zero handshake concurrency");
3936 assert!(err.to_string().contains("max_concurrent_tls_handshakes"));
3937 }
3938
3939 #[test]
3940 fn validate_consumes_and_proves() {
3941 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
3943 let validated = cfg.validate().expect("valid config");
3944 assert_eq!(validated.as_inner().name, "test-server");
3946 let raw = validated.into_inner();
3948 assert_eq!(raw.name, "test-server");
3949
3950 let mut bad = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
3952 bad.max_request_body = 0;
3953 assert!(bad.validate().is_err(), "zero body cap must fail validate");
3954 }
3955
3956 #[test]
3957 fn validate_rejects_zero_max_concurrent_requests() {
3958 let cfg =
3959 McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_max_concurrent_requests(0);
3960 let err = cfg.validate().expect_err("zero concurrency cap must fail");
3961 assert!(
3962 format!("{err}").contains("max_concurrent_requests"),
3963 "error should mention max_concurrent_requests, got: {err}"
3964 );
3965 }
3966
3967 #[test]
3968 fn validate_rejects_zero_max_tracked_keys() {
3969 let rl = crate::auth::RateLimitConfig {
3972 max_attempts_per_minute: 30,
3973 pre_auth_max_per_minute: None,
3974 max_tracked_keys: 0,
3975 idle_eviction: Duration::from_secs(15 * 60),
3976 burst: None,
3977 pre_auth_burst: None,
3978 };
3979 let auth_cfg = AuthConfig {
3980 enabled: true,
3981 api_keys: Vec::new(),
3982 mtls: None,
3983 rate_limit: Some(rl),
3984 #[cfg(feature = "oauth")]
3985 oauth: None,
3986 };
3987 let cfg = McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_auth(auth_cfg);
3988 let err = cfg.validate().expect_err("zero max_tracked_keys must fail");
3989 assert!(
3990 format!("{err}").contains("max_tracked_keys"),
3991 "error should mention max_tracked_keys, got: {err}"
3992 );
3993 }
3994
3995 #[test]
3996 fn derive_allowed_hosts_includes_public_host() {
3997 let hosts = derive_allowed_hosts("0.0.0.0:8080", Some("https://mcp.example.com/mcp"));
3998 assert!(
3999 hosts.iter().any(|h| h == "mcp.example.com"),
4000 "public_url host must be allowed"
4001 );
4002 }
4003
4004 #[test]
4005 fn derive_allowed_hosts_includes_bind_authority() {
4006 let hosts = derive_allowed_hosts("127.0.0.1:8080", None);
4007 assert!(
4008 hosts.iter().any(|h| h == "127.0.0.1"),
4009 "bind host must be allowed"
4010 );
4011 assert!(
4012 hosts.iter().any(|h| h == "127.0.0.1:8080"),
4013 "bind authority must be allowed"
4014 );
4015 }
4016
4017 #[tokio::test]
4020 async fn healthz_returns_ok_json() {
4021 let resp = healthz().await.into_response();
4022 assert_eq!(resp.status(), StatusCode::OK);
4023 let body = resp.into_body().collect().await.unwrap().to_bytes();
4024 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4025 assert_eq!(json["status"], "ok");
4026 assert!(
4027 json.get("name").is_none(),
4028 "healthz must not expose server name"
4029 );
4030 assert!(
4031 json.get("version").is_none(),
4032 "healthz must not expose version"
4033 );
4034 }
4035
4036 #[tokio::test]
4039 async fn readyz_returns_ok_when_ready() {
4040 let check: ReadinessCheck =
4041 Arc::new(|| Box::pin(async { serde_json::json!({"ready": true, "db": "connected"}) }));
4042 let resp = readyz(check).await.into_response();
4043 assert_eq!(resp.status(), StatusCode::OK);
4044 let body = resp.into_body().collect().await.unwrap().to_bytes();
4045 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
4046 assert_eq!(json["ready"], true);
4047 assert!(
4048 json.get("name").is_none(),
4049 "readyz must not expose server name"
4050 );
4051 assert!(
4052 json.get("version").is_none(),
4053 "readyz must not expose version"
4054 );
4055 assert_eq!(json["db"], "connected");
4056 }
4057
4058 #[tokio::test]
4059 async fn readyz_returns_503_when_not_ready() {
4060 let check: ReadinessCheck =
4061 Arc::new(|| Box::pin(async { serde_json::json!({"ready": false}) }));
4062 let resp = readyz(check).await.into_response();
4063 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4064 }
4065
4066 #[tokio::test]
4067 async fn readyz_returns_503_when_ready_missing() {
4068 let check: ReadinessCheck =
4069 Arc::new(|| Box::pin(async { serde_json::json!({"status": "starting"}) }));
4070 let resp = readyz(check).await.into_response();
4071 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
4073 }
4074
4075 fn peer_probe_router() -> axum::Router {
4080 async fn probe(req: Request<Body>) -> String {
4081 let ci = req
4082 .extensions()
4083 .get::<ConnectInfo<SocketAddr>>()
4084 .map(|c| c.0.to_string())
4085 .unwrap_or_default();
4086 let pa = req
4087 .extensions()
4088 .get::<PeerAddr>()
4089 .map(|p| p.addr.to_string())
4090 .unwrap_or_default();
4091 format!("{ci}|{pa}")
4092 }
4093 axum::Router::new()
4094 .route("/probe", axum::routing::get(probe))
4095 .layer(axum::middleware::from_fn(|req, next| {
4096 normalize_peer_addr_middleware(None, req, next)
4097 }))
4098 }
4099
4100 async fn body_string(resp: axum::response::Response) -> String {
4101 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
4102 String::from_utf8(bytes.to_vec()).unwrap()
4103 }
4104
4105 #[tokio::test]
4106 async fn normalize_preserves_existing_connect_info_and_mirrors_peer_addr() {
4107 let plain: SocketAddr = "10.0.0.1:1111".parse().unwrap();
4110 let tls: SocketAddr = "10.0.0.2:2222".parse().unwrap();
4111 let req = Request::builder()
4112 .uri("/probe")
4113 .extension(ConnectInfo(plain))
4114 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4115 .body(Body::empty())
4116 .unwrap();
4117 let resp = peer_probe_router().oneshot(req).await.unwrap();
4118 assert_eq!(resp.status(), StatusCode::OK);
4119 assert_eq!(body_string(resp).await, format!("{plain}|{plain}"));
4120 }
4121
4122 #[tokio::test]
4123 async fn normalize_inserts_connect_info_and_peer_addr_from_tls() {
4124 let tls: SocketAddr = "192.168.1.7:50443".parse().unwrap();
4125 let req = Request::builder()
4126 .uri("/probe")
4127 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
4128 .body(Body::empty())
4129 .unwrap();
4130 let resp = peer_probe_router().oneshot(req).await.unwrap();
4131 assert_eq!(resp.status(), StatusCode::OK);
4132 assert_eq!(body_string(resp).await, format!("{tls}|{tls}"));
4133 }
4134
4135 #[tokio::test]
4136 async fn normalize_no_op_without_any_connect_info() {
4137 let req = Request::builder()
4138 .uri("/probe")
4139 .body(Body::empty())
4140 .unwrap();
4141 let resp = peer_probe_router().oneshot(req).await.unwrap();
4142 assert_eq!(resp.status(), StatusCode::OK);
4143 assert_eq!(body_string(resp).await, "|");
4144 }
4145
4146 #[tokio::test]
4147 async fn peer_addr_extractor_rejects_when_absent() {
4148 async fn h(peer: PeerAddr) -> String {
4149 peer.addr.to_string()
4150 }
4151 let app = axum::Router::new().route("/p", axum::routing::get(h));
4152 let req = Request::builder().uri("/p").body(Body::empty()).unwrap();
4153 let resp = app.oneshot(req).await.unwrap();
4154 assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
4155 }
4156
4157 #[tokio::test]
4158 async fn peer_addr_extractor_returns_value_when_present() {
4159 async fn h(peer: PeerAddr) -> String {
4160 peer.addr.to_string()
4161 }
4162 let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap();
4163 let app = axum::Router::new().route("/p", axum::routing::get(h));
4164 let req = Request::builder()
4165 .uri("/p")
4166 .extension(PeerAddr::new(addr))
4167 .body(Body::empty())
4168 .unwrap();
4169 let resp = app.oneshot(req).await.unwrap();
4170 assert_eq!(resp.status(), StatusCode::OK);
4171 assert_eq!(body_string(resp).await, addr.to_string());
4172 }
4173
4174 #[tokio::test]
4175 async fn peer_addr_via_extension_extractor() {
4176 async fn h(axum::Extension(peer): axum::Extension<PeerAddr>) -> String {
4177 peer.addr.to_string()
4178 }
4179 let addr: SocketAddr = "127.0.0.1:4242".parse().unwrap();
4180 let app = axum::Router::new().route("/p", axum::routing::get(h));
4181 let req = Request::builder()
4182 .uri("/p")
4183 .extension(PeerAddr::new(addr))
4184 .body(Body::empty())
4185 .unwrap();
4186 let resp = app.oneshot(req).await.unwrap();
4187 assert_eq!(resp.status(), StatusCode::OK);
4188 assert_eq!(body_string(resp).await, addr.to_string());
4189 }
4190
4191 fn limited_router(per_minute: u32) -> axum::Router {
4196 limited_router_with_burst(per_minute, None)
4197 }
4198
4199 fn limited_router_with_burst(per_minute: u32, burst: Option<u32>) -> axum::Router {
4201 limited_router_full(per_minute, burst, &[])
4202 }
4203
4204 fn limited_router_full(
4208 per_minute: u32,
4209 burst: Option<u32>,
4210 exempt_paths: &[&str],
4211 ) -> axum::Router {
4212 let limiter = build_extra_route_rate_limiter(per_minute, burst);
4213 let exempt: Arc<std::collections::HashSet<String>> =
4214 Arc::new(exempt_paths.iter().map(|s| (*s).to_owned()).collect());
4215 axum::Router::new()
4216 .route("/limited", axum::routing::get(|| async { "ok" }))
4217 .route("/exempt", axum::routing::get(|| async { "ok" }))
4218 .layer(axum::middleware::from_fn(move |req, next| {
4219 let l = Arc::clone(&limiter);
4220 let e = Arc::clone(&exempt);
4221 extra_route_rate_limit_middleware(l, e, req, next)
4222 }))
4223 }
4224
4225 fn limited_req(ip: &str) -> Request<Body> {
4226 limited_req_to(ip, "/limited")
4227 }
4228
4229 fn limited_req_to(ip: &str, path: &str) -> Request<Body> {
4230 let addr: SocketAddr = format!("{ip}:40000").parse().unwrap();
4231 Request::builder()
4232 .uri(path)
4233 .extension(ConnectInfo(addr))
4234 .body(Body::empty())
4235 .unwrap()
4236 }
4237
4238 #[tokio::test]
4239 async fn extra_route_limiter_denies_over_quota() {
4240 let app = limited_router(2);
4241 for i in 0..2 {
4242 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4243 assert_eq!(resp.status(), StatusCode::OK, "request {i} should pass");
4244 }
4245 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
4246 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4247 let body = body_string(resp).await;
4248 assert!(
4249 body.contains("too many requests to application routes"),
4250 "deny body should match the limiter message, got: {body}"
4251 );
4252 }
4253
4254 #[tokio::test]
4255 async fn extra_route_limiter_isolates_keys() {
4256 let app = limited_router(2);
4257 for _ in 0..2 {
4258 let resp = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4259 assert_eq!(resp.status(), StatusCode::OK);
4260 }
4261 let exhausted = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
4262 assert_eq!(exhausted.status(), StatusCode::TOO_MANY_REQUESTS);
4263 let other = app.clone().oneshot(limited_req("10.3.3.3")).await.unwrap();
4265 assert_eq!(other.status(), StatusCode::OK);
4266 }
4267
4268 #[tokio::test]
4269 async fn extra_route_limiter_fails_open_without_peer() {
4270 let app = limited_router(1);
4271 for i in 0..3 {
4272 let req = Request::builder()
4273 .uri("/limited")
4274 .body(Body::empty())
4275 .unwrap();
4276 let resp = app.clone().oneshot(req).await.unwrap();
4277 assert_eq!(
4278 resp.status(),
4279 StatusCode::OK,
4280 "request {i} should fail open"
4281 );
4282 }
4283 }
4284
4285 #[tokio::test]
4286 async fn extra_route_limiter_extracts_tls_conn_info() {
4287 let app = limited_router(2);
4288 let mk = || {
4289 let addr: SocketAddr = "192.168.9.9:55555".parse().unwrap();
4290 Request::builder()
4291 .uri("/limited")
4292 .extension(ConnectInfo(TlsConnInfo::new(addr, None)))
4293 .body(Body::empty())
4294 .unwrap()
4295 };
4296 for _ in 0..2 {
4297 assert_eq!(
4298 app.clone().oneshot(mk()).await.unwrap().status(),
4299 StatusCode::OK
4300 );
4301 }
4302 let resp = app.clone().oneshot(mk()).await.unwrap();
4303 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4304 }
4305
4306 #[tokio::test]
4307 async fn extra_route_limiter_exempt_path_bypasses_quota() {
4308 let app = limited_router_full(1, None, &["/exempt"]);
4311 for i in 0..5 {
4312 let resp = app
4313 .clone()
4314 .oneshot(limited_req_to("10.6.6.6", "/exempt"))
4315 .await
4316 .unwrap();
4317 assert_eq!(resp.status(), StatusCode::OK, "exempt request {i}");
4318 }
4319 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
4321 assert_eq!(resp.status(), StatusCode::OK);
4322 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
4324 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4325 }
4326
4327 #[tokio::test]
4328 async fn extra_route_limiter_exemption_is_raw_exact_match() {
4329 let app = limited_router_full(1, None, &["/exempt"]);
4332 let ok = app
4333 .clone()
4334 .oneshot(limited_req_to("10.7.7.7", "/exempt/"))
4335 .await
4336 .unwrap();
4337 assert_eq!(
4338 ok.status(),
4339 StatusCode::NOT_FOUND,
4340 "variant path routes 404"
4341 );
4342 let denied = app
4344 .clone()
4345 .oneshot(limited_req_to("10.7.7.7", "/limited"))
4346 .await
4347 .unwrap();
4348 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4349 }
4350
4351 #[cfg(feature = "metrics")]
4352 #[tokio::test]
4353 async fn extra_route_limiter_deny_increments_counter_exempt_does_not() {
4354 let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
4355 let app = limited_router_full(1, None, &["/exempt"]);
4356 let mk = |path: &str| {
4357 let addr: SocketAddr = "10.8.8.8:40000".parse().unwrap();
4358 Request::builder()
4359 .uri(path)
4360 .extension(ConnectInfo(addr))
4361 .extension(Arc::clone(&metrics))
4362 .body(Body::empty())
4363 .unwrap()
4364 };
4365 let counter = || {
4366 metrics
4367 .rate_limited_total
4368 .with_label_values(&["extra_route"])
4369 .get()
4370 };
4371 for _ in 0..3 {
4373 assert_eq!(
4374 app.clone().oneshot(mk("/exempt")).await.unwrap().status(),
4375 StatusCode::OK
4376 );
4377 }
4378 assert_eq!(counter(), 0, "exempt requests must not count as denies");
4379 assert_eq!(
4381 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4382 StatusCode::OK
4383 );
4384 assert_eq!(counter(), 0);
4385 assert_eq!(
4386 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4387 StatusCode::TOO_MANY_REQUESTS
4388 );
4389 assert_eq!(counter(), 1, "deny must increment the extra_route label");
4390 }
4391
4392 #[test]
4393 fn validate_rejects_exempt_paths_without_base_knob() {
4394 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4395 .with_extra_route_rate_limit_exempt_paths(["/ok"]);
4396 let err = cfg.validate().expect_err("exempt paths without rate limit");
4397 assert!(err.to_string().contains("requires extra_route_rate_limit"));
4398 }
4399
4400 #[test]
4401 fn validate_rejects_malformed_exempt_paths() {
4402 for bad in ["", "no-slash"] {
4403 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4404 .with_extra_route_rate_limit(10)
4405 .with_extra_route_rate_limit_exempt_paths([bad]);
4406 let err = cfg.validate().expect_err("malformed exempt path");
4407 assert!(
4408 err.to_string()
4409 .contains("must be non-empty and start with '/'"),
4410 "entry {bad:?}: {err}"
4411 );
4412 }
4413 }
4414
4415 #[test]
4416 fn validate_accepts_wellformed_exempt_paths() {
4417 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4418 .with_extra_route_rate_limit(10)
4419 .with_extra_route_rate_limit_exempt_paths(["/.well-known/oauth-authorization-server"]);
4420 assert!(cfg.validate().is_ok());
4421 }
4422
4423 #[test]
4424 fn validate_rejects_zero_extra_route_rate_limit() {
4425 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4426 .with_extra_route_rate_limit(0);
4427 let err = cfg.validate().expect_err("zero extra route rate limit");
4428 assert!(err.to_string().contains("extra_route_rate_limit"));
4429 }
4430
4431 #[tokio::test]
4432 async fn extra_route_limiter_burst_allows_initial_spike() {
4433 let app = limited_router_with_burst(1, Some(3));
4434 for i in 0..3 {
4435 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4436 assert_eq!(resp.status(), StatusCode::OK, "burst request {i}");
4437 }
4438 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4439 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4440 }
4441
4442 #[tokio::test]
4443 async fn extra_route_limiter_deny_sets_retry_after() {
4444 let app = limited_router(1);
4445 let ok = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4446 assert_eq!(ok.status(), StatusCode::OK);
4447 let denied = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4448 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4449 let retry_after = denied
4450 .headers()
4451 .get(header::RETRY_AFTER)
4452 .expect("Retry-After present")
4453 .to_str()
4454 .unwrap()
4455 .parse::<u64>()
4456 .unwrap();
4457 assert!(retry_after >= 1, "delta-seconds must be >= 1");
4458 }
4459
4460 #[test]
4461 fn validate_rejects_zero_burst_knobs() {
4462 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4463 .with_tool_rate_limit(10)
4464 .with_tool_rate_limit_burst(0)
4465 .validate()
4466 .expect_err("zero tool burst");
4467 assert!(err.to_string().contains("tool_rate_limit_burst"));
4468
4469 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4470 .with_extra_route_rate_limit(10)
4471 .with_extra_route_rate_limit_burst(0)
4472 .validate()
4473 .expect_err("zero extra route burst");
4474 assert!(err.to_string().contains("extra_route_rate_limit_burst"));
4475 }
4476
4477 #[test]
4478 fn validate_rejects_orphan_burst_knobs() {
4479 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4480 .with_tool_rate_limit_burst(5)
4481 .validate()
4482 .expect_err("orphan tool burst");
4483 assert!(err.to_string().contains("requires tool_rate_limit"));
4484
4485 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4486 .with_extra_route_rate_limit_burst(5)
4487 .validate()
4488 .expect_err("orphan extra route burst");
4489 assert!(err.to_string().contains("requires extra_route_rate_limit"));
4490 }
4491
4492 #[test]
4493 fn validate_rejects_zero_auth_bursts() {
4494 let auth = AuthConfig::with_keys(vec![])
4495 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_burst(0));
4496 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4497 .with_auth(auth)
4498 .validate()
4499 .expect_err("zero auth burst");
4500 assert!(err.to_string().contains("rate_limit.burst"));
4501
4502 let auth = AuthConfig::with_keys(vec![])
4503 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(0));
4504 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4505 .with_auth(auth)
4506 .validate()
4507 .expect_err("zero pre-auth burst");
4508 assert!(err.to_string().contains("pre_auth_burst"));
4509 }
4510
4511 #[test]
4514 fn validate_accepts_pre_auth_burst_without_explicit_pre_auth_rate() {
4515 let auth = AuthConfig::with_keys(vec![])
4516 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(50));
4517 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4518 assert!(cfg.validate().is_ok(), "pre_auth_burst has no orphan rule");
4519 }
4520
4521 fn forward_resolver(trusted: &[&str], mode: ForwardedHeaderMode) -> Arc<ForwardResolver> {
4524 Arc::new(ForwardResolver {
4525 trusted: trusted.iter().map(|s| s.parse().unwrap()).collect(),
4526 mode,
4527 })
4528 }
4529
4530 fn forwarded_probe_router(resolver: Option<Arc<ForwardResolver>>) -> axum::Router {
4532 async fn probe(req: Request<Body>) -> String {
4533 let pa = req
4534 .extensions()
4535 .get::<PeerAddr>()
4536 .map(|p| p.addr.ip().to_string())
4537 .unwrap_or_default();
4538 let ci = req
4539 .extensions()
4540 .get::<ClientIp>()
4541 .map(|c| c.ip.to_string())
4542 .unwrap_or_default();
4543 format!("{pa}|{ci}")
4544 }
4545 axum::Router::new()
4546 .route("/probe", axum::routing::get(probe))
4547 .layer(axum::middleware::from_fn(move |req, next| {
4548 let r = resolver.clone();
4549 normalize_peer_addr_middleware(r, req, next)
4550 }))
4551 }
4552
4553 fn probe_req(peer: &str, header: Option<(&str, &str)>) -> Request<Body> {
4554 let addr: SocketAddr = peer.parse().unwrap();
4555 let mut builder = Request::builder()
4556 .uri("/probe")
4557 .extension(ConnectInfo(addr));
4558 if let Some((name, value)) = header {
4559 builder = builder.header(name, value);
4560 }
4561 builder.body(Body::empty()).unwrap()
4562 }
4563
4564 #[tokio::test]
4565 async fn client_ip_equals_direct_without_resolver() {
4566 let app = forwarded_probe_router(None);
4567 let resp = app
4568 .oneshot(probe_req(
4569 "10.1.2.3:4444",
4570 Some(("x-forwarded-for", "203.0.113.7")),
4571 ))
4572 .await
4573 .unwrap();
4574 assert_eq!(
4575 body_string(resp).await,
4576 "10.1.2.3|10.1.2.3",
4577 "feature off: header ignored, ClientIp == direct"
4578 );
4579 }
4580
4581 #[tokio::test]
4582 async fn client_ip_resolved_for_trusted_peer() {
4583 let app = forwarded_probe_router(Some(forward_resolver(
4584 &["10.0.0.0/8"],
4585 ForwardedHeaderMode::XForwardedFor,
4586 )));
4587 let resp = app
4588 .oneshot(probe_req(
4589 "10.0.0.1:9999",
4590 Some(("x-forwarded-for", "203.0.113.7")),
4591 ))
4592 .await
4593 .unwrap();
4594 assert_eq!(
4595 body_string(resp).await,
4596 "10.0.0.1|203.0.113.7",
4597 "PeerAddr stays direct while ClientIp resolves"
4598 );
4599 }
4600
4601 #[tokio::test]
4602 async fn client_ip_falls_back_to_direct_on_malformed_header() {
4603 let app = forwarded_probe_router(Some(forward_resolver(
4604 &["10.0.0.0/8"],
4605 ForwardedHeaderMode::XForwardedFor,
4606 )));
4607 let resp = app
4608 .oneshot(probe_req(
4609 "10.0.0.1:9999",
4610 Some(("x-forwarded-for", "not-an-ip")),
4611 ))
4612 .await
4613 .unwrap();
4614 assert_eq!(
4615 body_string(resp).await,
4616 "10.0.0.1|10.0.0.1",
4617 "malformed chain falls back to the direct peer"
4618 );
4619 }
4620
4621 #[test]
4622 fn forwarded_header_mode_deserializes_kebab_case() {
4623 #[derive(serde::Deserialize)]
4624 struct Wrapper {
4625 mode: ForwardedHeaderMode,
4626 }
4627 let w: Wrapper = toml::from_str(r#"mode = "x-forwarded-for""#).unwrap();
4628 assert_eq!(w.mode, ForwardedHeaderMode::XForwardedFor);
4629 let w: Wrapper = toml::from_str(r#"mode = "forwarded""#).unwrap();
4630 assert_eq!(w.mode, ForwardedHeaderMode::Forwarded);
4631 assert!(
4632 toml::from_str::<Wrapper>(r#"mode = "XForwardedFor""#).is_err(),
4633 "PascalCase wire value must be rejected"
4634 );
4635 }
4636
4637 #[test]
4638 fn validate_rejects_bad_trusted_proxy_entry() {
4639 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4640 .with_trusted_proxies(["not-a-cidr"]);
4641 let err = cfg.validate().expect_err("bad CIDR");
4642 assert!(err.to_string().contains("trusted_proxies"));
4643 }
4644
4645 #[test]
4646 fn validate_rejects_zero_prefix_trusted_proxy() {
4647 for entry in ["0.0.0.0/0", "::/0"] {
4648 let cfg =
4649 McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([entry]);
4650 let err = cfg.validate().expect_err("zero-prefix CIDR");
4651 assert!(
4652 err.to_string().contains("prefix length 0"),
4653 "entry {entry}: {err}"
4654 );
4655 }
4656 }
4657
4658 #[test]
4659 fn validate_accepts_cidr_and_bare_ip_proxy_entries() {
4660 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([
4661 "10.0.0.0/8",
4662 "192.0.2.1",
4663 "2001:db8::1",
4664 ]);
4665 assert!(cfg.validate().is_ok(), "CIDRs and bare IPs are accepted");
4666 }
4667
4668 #[test]
4669 fn validate_rejects_forwarded_header_without_proxies() {
4670 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4671 .with_forwarded_header(ForwardedHeaderMode::Forwarded);
4672 let err = cfg.validate().expect_err("mode without proxies");
4673 assert!(err.to_string().contains("requires trusted_proxies"));
4674 }
4675
4676 fn origin_router(origins: Vec<String>, log_request_headers: bool) -> axum::Router {
4680 let allowed: Arc<[String]> = Arc::from(origins);
4681 axum::Router::new()
4682 .route("/test", axum::routing::get(|| async { "ok" }))
4683 .layer(axum::middleware::from_fn(move |req, next| {
4684 let a = Arc::clone(&allowed);
4685 origin_check_middleware(a, log_request_headers, req, next)
4686 }))
4687 }
4688
4689 #[tokio::test]
4690 async fn origin_allowed_passes() {
4691 let app = origin_router(vec!["http://localhost:3000".into()], false);
4692 let req = Request::builder()
4693 .uri("/test")
4694 .header(header::ORIGIN, "http://localhost:3000")
4695 .body(Body::empty())
4696 .unwrap();
4697 let resp = app.oneshot(req).await.unwrap();
4698 assert_eq!(resp.status(), StatusCode::OK);
4699 }
4700
4701 #[tokio::test]
4702 async fn origin_rejected_returns_403() {
4703 let app = origin_router(vec!["http://localhost:3000".into()], false);
4704 let req = Request::builder()
4705 .uri("/test")
4706 .header(header::ORIGIN, "http://evil.com")
4707 .body(Body::empty())
4708 .unwrap();
4709 let resp = app.oneshot(req).await.unwrap();
4710 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
4711 }
4712
4713 #[tokio::test]
4714 async fn no_origin_header_passes() {
4715 let app = origin_router(vec!["http://localhost:3000".into()], false);
4716 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4717 let resp = app.oneshot(req).await.unwrap();
4718 assert_eq!(resp.status(), StatusCode::OK);
4719 }
4720
4721 #[tokio::test]
4722 async fn empty_allowlist_rejects_any_origin() {
4723 let app = origin_router(vec![], false);
4724 let req = Request::builder()
4725 .uri("/test")
4726 .header(header::ORIGIN, "http://anything.com")
4727 .body(Body::empty())
4728 .unwrap();
4729 let resp = app.oneshot(req).await.unwrap();
4730 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
4731 }
4732
4733 #[tokio::test]
4734 async fn empty_allowlist_passes_without_origin() {
4735 let app = origin_router(vec![], false);
4736 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4737 let resp = app.oneshot(req).await.unwrap();
4738 assert_eq!(resp.status(), StatusCode::OK);
4739 }
4740
4741 #[test]
4742 fn format_request_headers_redacts_sensitive_values() {
4743 let mut headers = axum::http::HeaderMap::new();
4744 headers.insert("authorization", "Bearer secret-token".parse().unwrap());
4745 headers.insert("cookie", "sid=abc".parse().unwrap());
4746 headers.insert("x-request-id", "req-123".parse().unwrap());
4747
4748 let out = format_request_headers_for_log(&headers);
4749 assert!(out.contains("authorization: [REDACTED]"));
4750 assert!(out.contains("cookie: [REDACTED]"));
4751 assert!(out.contains("x-request-id: req-123"));
4752 assert!(!out.contains("secret-token"));
4753 }
4754
4755 fn security_router(is_tls: bool) -> axum::Router {
4758 security_router_with(is_tls, SecurityHeadersConfig::default())
4759 }
4760
4761 fn security_router_with(is_tls: bool, cfg: SecurityHeadersConfig) -> axum::Router {
4762 let cfg = Arc::new(cfg);
4763 axum::Router::new()
4764 .route("/test", axum::routing::get(|| async { "ok" }))
4765 .layer(axum::middleware::from_fn(move |req, next| {
4766 let c = Arc::clone(&cfg);
4767 security_headers_middleware(is_tls, c, req, next)
4768 }))
4769 }
4770
4771 #[tokio::test]
4772 async fn security_headers_set_on_response() {
4773 let app = security_router(false);
4774 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4775 let resp = app.oneshot(req).await.unwrap();
4776 assert_eq!(resp.status(), StatusCode::OK);
4777
4778 let h = resp.headers();
4779 assert_eq!(h.get("x-content-type-options").unwrap(), "nosniff");
4780 assert_eq!(h.get("x-frame-options").unwrap(), "deny");
4781 assert_eq!(h.get("cache-control").unwrap(), "no-store, max-age=0");
4782 assert_eq!(h.get("referrer-policy").unwrap(), "no-referrer");
4783 assert_eq!(h.get("cross-origin-opener-policy").unwrap(), "same-origin");
4784 assert_eq!(
4785 h.get("cross-origin-resource-policy").unwrap(),
4786 "same-origin"
4787 );
4788 assert_eq!(
4789 h.get("cross-origin-embedder-policy").unwrap(),
4790 "require-corp"
4791 );
4792 assert_eq!(h.get("x-permitted-cross-domain-policies").unwrap(), "none");
4793 assert!(
4794 h.get("permissions-policy")
4795 .unwrap()
4796 .to_str()
4797 .unwrap()
4798 .contains("camera=()"),
4799 "permissions-policy must restrict browser features"
4800 );
4801 assert_eq!(
4802 h.get("content-security-policy").unwrap(),
4803 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
4804 );
4805 assert_eq!(h.get("x-dns-prefetch-control").unwrap(), "off");
4806 assert!(h.get("strict-transport-security").is_none());
4808 }
4809
4810 #[tokio::test]
4811 async fn hsts_set_when_tls_enabled() {
4812 let app = security_router(true);
4813 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4814 let resp = app.oneshot(req).await.unwrap();
4815
4816 let hsts = resp.headers().get("strict-transport-security").unwrap();
4817 assert!(
4818 hsts.to_str().unwrap().contains("max-age=63072000"),
4819 "HSTS must set 2-year max-age"
4820 );
4821 }
4822
4823 #[tokio::test]
4824 async fn default_csp_matches_guideline() {
4825 let app = security_router(false);
4826 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4827 let resp = app.oneshot(req).await.unwrap();
4828 assert_eq!(
4829 resp.headers().get("content-security-policy").unwrap(),
4830 "default-src 'none'; form-action 'self'; object-src 'none'; frame-ancestors 'none'; upgrade-insecure-requests"
4831 );
4832 }
4833
4834 #[tokio::test]
4835 async fn operator_csp_override_still_wins() {
4836 let cfg = SecurityHeadersConfig {
4837 content_security_policy: Some("default-src 'self'".into()),
4838 ..SecurityHeadersConfig::default()
4839 };
4840 let app = security_router_with(false, cfg);
4841 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4842 let resp = app.oneshot(req).await.unwrap();
4843 assert_eq!(
4844 resp.headers().get("content-security-policy").unwrap(),
4845 "default-src 'self'"
4846 );
4847 }
4848
4849 fn check_with_security_headers(
4855 headers: SecurityHeadersConfig,
4856 ) -> Result<(), RmcpServerKitError> {
4857 let cfg =
4858 McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_security_headers(headers);
4859 cfg.check()
4860 }
4861
4862 #[test]
4863 fn security_headers_config_default_validates() {
4864 check_with_security_headers(SecurityHeadersConfig::default())
4865 .expect("default SecurityHeadersConfig must validate");
4866 }
4867
4868 #[test]
4869 fn security_headers_config_validate_accepts_empty_string() {
4870 let h = SecurityHeadersConfig {
4872 x_content_type_options: Some(String::new()),
4873 x_frame_options: Some(String::new()),
4874 cache_control: Some(String::new()),
4875 referrer_policy: Some(String::new()),
4876 cross_origin_opener_policy: Some(String::new()),
4877 cross_origin_resource_policy: Some(String::new()),
4878 cross_origin_embedder_policy: Some(String::new()),
4879 permissions_policy: Some(String::new()),
4880 x_permitted_cross_domain_policies: Some(String::new()),
4881 content_security_policy: Some(String::new()),
4882 x_dns_prefetch_control: Some(String::new()),
4883 strict_transport_security: Some(String::new()),
4884 };
4885 check_with_security_headers(h).expect("Some(\"\") on every field must validate (omit-all)");
4886 }
4887
4888 #[test]
4889 fn security_headers_config_validate_rejects_bad_value() {
4890 let h = SecurityHeadersConfig {
4892 referrer_policy: Some("\u{0007}".into()),
4893 ..SecurityHeadersConfig::default()
4894 };
4895 let err = check_with_security_headers(h)
4896 .expect_err("control char in referrer_policy must reject");
4897 let msg = err.to_string();
4898 assert!(
4899 msg.contains("referrer_policy"),
4900 "error must name the offending field, got: {msg}"
4901 );
4902 }
4903
4904 #[test]
4905 fn security_headers_config_validate_rejects_hsts_preload() {
4906 let h = SecurityHeadersConfig {
4907 strict_transport_security: Some("max-age=63072000; includeSubDomains; preload".into()),
4908 ..SecurityHeadersConfig::default()
4909 };
4910 let err = check_with_security_headers(h).expect_err("HSTS with preload must reject");
4911 let msg = err.to_string();
4912 assert!(
4913 msg.contains("strict_transport_security"),
4914 "error must name the field, got: {msg}"
4915 );
4916 assert!(
4917 msg.to_lowercase().contains("preload"),
4918 "error must mention `preload`, got: {msg}"
4919 );
4920 }
4921
4922 #[test]
4923 fn security_headers_config_validate_rejects_hsts_preload_uppercase() {
4924 let h = SecurityHeadersConfig {
4926 strict_transport_security: Some("max-age=600; PRELOAD".into()),
4927 ..SecurityHeadersConfig::default()
4928 };
4929 check_with_security_headers(h).expect_err("HSTS preload check must be case-insensitive");
4930 }
4931
4932 #[tokio::test]
4933 async fn security_headers_override_honored() {
4934 let h = SecurityHeadersConfig {
4936 x_frame_options: Some("SAMEORIGIN".into()),
4937 ..SecurityHeadersConfig::default()
4938 };
4939 let app = security_router_with(false, h);
4940 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4941 let resp = app.oneshot(req).await.unwrap();
4942 assert_eq!(resp.status(), StatusCode::OK);
4943
4944 let xfo = resp.headers().get("x-frame-options").unwrap();
4945 assert_eq!(xfo, "SAMEORIGIN");
4946 }
4947
4948 #[tokio::test]
4949 async fn security_headers_empty_string_omits() {
4950 let h = SecurityHeadersConfig {
4952 referrer_policy: Some(String::new()),
4953 ..SecurityHeadersConfig::default()
4954 };
4955 let app = security_router_with(false, h);
4956 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4957 let resp = app.oneshot(req).await.unwrap();
4958 assert_eq!(resp.status(), StatusCode::OK);
4959
4960 assert!(
4961 resp.headers().get("referrer-policy").is_none(),
4962 "Some(\"\") must omit the header"
4963 );
4964 assert_eq!(
4966 resp.headers().get("x-content-type-options").unwrap(),
4967 "nosniff"
4968 );
4969 }
4970
4971 #[tokio::test]
4972 async fn security_headers_hsts_only_when_tls() {
4973 let h = SecurityHeadersConfig {
4975 strict_transport_security: Some("max-age=600".into()),
4976 ..SecurityHeadersConfig::default()
4977 };
4978 let app = security_router_with(false, h);
4979 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4980 let resp = app.oneshot(req).await.unwrap();
4981 assert!(
4982 resp.headers().get("strict-transport-security").is_none(),
4983 "HSTS must remain absent on plaintext deployments even with override"
4984 );
4985 }
4986
4987 #[cfg(feature = "oauth")]
4990 #[tokio::test]
4991 async fn oauth_token_cache_headers_set_pragma_and_vary() {
4992 let app = axum::Router::new()
4993 .route("/token", axum::routing::post(|| async { "{}" }))
4994 .layer(axum::middleware::from_fn(
4995 oauth_token_cache_headers_middleware,
4996 ));
4997 let req = Request::builder()
4998 .method("POST")
4999 .uri("/token")
5000 .body(Body::from("{}"))
5001 .unwrap();
5002 let resp = app.oneshot(req).await.unwrap();
5003 assert_eq!(resp.status(), StatusCode::OK);
5004
5005 let h = resp.headers();
5006 assert_eq!(
5007 h.get("pragma").unwrap(),
5008 "no-cache",
5009 "RFC 6749 §5.1: token responses must set Pragma: no-cache"
5010 );
5011 let vary_values: Vec<String> = h
5012 .get_all("vary")
5013 .iter()
5014 .filter_map(|v| v.to_str().ok().map(str::to_owned))
5015 .collect();
5016 assert!(
5017 vary_values
5018 .iter()
5019 .any(|v| v.eq_ignore_ascii_case("Authorization")),
5020 "RFC 6750 §5.4: Vary must include Authorization, got {vary_values:?}"
5021 );
5022 }
5023
5024 #[cfg(feature = "oauth")]
5025 #[tokio::test]
5026 async fn oauth_token_cache_headers_preserve_existing_vary() {
5027 let app = axum::Router::new()
5030 .route(
5031 "/token",
5032 axum::routing::post(|| async {
5033 axum::response::Response::builder()
5034 .header("vary", "Accept-Encoding")
5035 .body(Body::from("{}"))
5036 .unwrap()
5037 }),
5038 )
5039 .layer(axum::middleware::from_fn(
5040 oauth_token_cache_headers_middleware,
5041 ));
5042 let req = Request::builder()
5043 .method("POST")
5044 .uri("/token")
5045 .body(Body::empty())
5046 .unwrap();
5047 let resp = app.oneshot(req).await.unwrap();
5048
5049 let vary: Vec<String> = resp
5050 .headers()
5051 .get_all("vary")
5052 .iter()
5053 .filter_map(|v| v.to_str().ok().map(str::to_owned))
5054 .collect();
5055 assert!(
5056 vary.iter().any(|v| v.contains("Accept-Encoding")),
5057 "must preserve pre-existing Vary value, got {vary:?}"
5058 );
5059 assert!(
5060 vary.iter().any(|v| v.contains("Authorization")),
5061 "must append Authorization to Vary, got {vary:?}"
5062 );
5063 }
5064
5065 #[test]
5068 fn version_omits_build_fingerprint_by_default() {
5069 let v = version_payload("my-server", "1.2.3", false);
5070 assert_eq!(v["name"], "my-server");
5071 assert_eq!(v["version"], "1.2.3");
5072 assert!(v["rmcp_server_kit_version"].is_string());
5073 assert!(
5074 v.get("build_git_sha").is_none(),
5075 "build sha must be hidden by default"
5076 );
5077 assert!(v.get("build_timestamp").is_none());
5078 assert!(v.get("rust_version").is_none());
5079 }
5080
5081 #[test]
5082 fn version_exposes_all_when_enabled() {
5083 let v = version_payload("my-server", "1.2.3", true);
5084 assert!(v["build_git_sha"].is_string());
5085 assert!(v["build_timestamp"].is_string());
5086 assert!(v["rust_version"].is_string());
5087 assert!(v["rmcp_server_kit_version"].is_string());
5088 }
5089
5090 #[tokio::test]
5093 async fn concurrency_limit_layer_composes_and_serves() {
5094 let app = axum::Router::new()
5098 .route("/ok", axum::routing::get(|| async { "ok" }))
5099 .layer(
5100 tower::ServiceBuilder::new()
5101 .layer(axum::error_handling::HandleErrorLayer::new(
5102 |_err: tower::BoxError| async { StatusCode::SERVICE_UNAVAILABLE },
5103 ))
5104 .layer(tower::load_shed::LoadShedLayer::new())
5105 .layer(tower::limit::ConcurrencyLimitLayer::new(4)),
5106 );
5107 let resp = app
5108 .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap())
5109 .await
5110 .unwrap();
5111 assert_eq!(resp.status(), StatusCode::OK);
5112 }
5113
5114 #[tokio::test]
5117 async fn compression_layer_gzip_encodes_response() {
5118 use tower_http::compression::Predicate as _;
5119
5120 let big_body = "a".repeat(4096);
5121 let app = axum::Router::new()
5122 .route(
5123 "/big",
5124 axum::routing::get(move || {
5125 let body = big_body.clone();
5126 async move { body }
5127 }),
5128 )
5129 .layer(
5130 tower_http::compression::CompressionLayer::new()
5131 .gzip(true)
5132 .br(true)
5133 .compress_when(
5134 tower_http::compression::DefaultPredicate::new()
5135 .and(tower_http::compression::predicate::SizeAbove::new(1024)),
5136 ),
5137 );
5138
5139 let req = Request::builder()
5140 .uri("/big")
5141 .header(header::ACCEPT_ENCODING, "gzip")
5142 .body(Body::empty())
5143 .unwrap();
5144 let resp = app.oneshot(req).await.unwrap();
5145 assert_eq!(resp.status(), StatusCode::OK);
5146 assert_eq!(
5147 resp.headers().get(header::CONTENT_ENCODING).unwrap(),
5148 "gzip"
5149 );
5150 }
5151
5152 #[tokio::test]
5155 async fn tls_handshake_timeout_reaps_idle_connections() {
5156 use tokio::io::AsyncReadExt as _;
5157
5158 let _ = rustls::crypto::ring::default_provider().install_default();
5159
5160 let key = rcgen::KeyPair::generate().expect("generate key");
5162 let cert = rcgen::CertificateParams::new(vec!["localhost".to_owned()])
5163 .expect("cert params")
5164 .self_signed(&key)
5165 .expect("self-signed cert");
5166 let dir = std::env::temp_dir().join(format!(
5167 "rmcp-server-kit-hs-timeout-{}",
5168 std::time::SystemTime::now()
5169 .duration_since(std::time::UNIX_EPOCH)
5170 .expect("clock after epoch")
5171 .as_nanos()
5172 ));
5173 tokio::fs::create_dir_all(&dir).await.expect("temp dir");
5174 let cert_path = dir.join("server.crt");
5175 let key_path = dir.join("server.key");
5176 tokio::fs::write(&cert_path, cert.pem())
5177 .await
5178 .expect("write cert");
5179 tokio::fs::write(&key_path, key.serialize_pem())
5180 .await
5181 .expect("write key");
5182
5183 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
5184 let tls = TlsListener::new(
5185 listener,
5186 &cert_path,
5187 &key_path,
5188 None,
5189 None,
5190 Duration::from_millis(200),
5191 8, )
5193 .expect("tls listener");
5194 let addr = axum::serve::Listener::local_addr(&tls).expect("local addr");
5195
5196 let mut idle = tokio::net::TcpStream::connect(addr).await.expect("connect");
5200 let mut buf = [0_u8; 16];
5201 let read = tokio::time::timeout(Duration::from_secs(2), idle.read(&mut buf))
5202 .await
5203 .expect("server must reap the idle handshake within its timeout");
5204 match read {
5205 Ok(0) | Err(_) => {} Ok(n) => panic!("unexpected {n} bytes from server during reaped handshake"),
5207 }
5208
5209 drop(tls);
5210 }
5211
5212 fn assert_owasp_headers(resp: &axum::response::Response, ctx: &str) {
5215 let h = resp.headers();
5216 assert!(
5217 h.contains_key("x-content-type-options"),
5218 "{ctx}: missing X-Content-Type-Options"
5219 );
5220 assert!(
5221 h.contains_key("x-frame-options"),
5222 "{ctx}: missing X-Frame-Options"
5223 );
5224 assert!(
5225 h.contains_key("strict-transport-security"),
5226 "{ctx}: missing Strict-Transport-Security"
5227 );
5228 assert!(
5229 h.contains_key(header::CONTENT_SECURITY_POLICY),
5230 "{ctx}: missing Content-Security-Policy"
5231 );
5232 }
5233
5234 fn m5_router(configure: impl FnOnce(&mut McpServerConfig)) -> axum::Router {
5235 #[derive(Clone)]
5236 struct H;
5237 impl ServerHandler for H {}
5238 let mut config = McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0")
5242 .with_allowed_origins(["http://good.example"])
5243 .with_tls("unused.crt", "unused.key");
5244 configure(&mut config);
5245 let (router, _params) = build_app_router(config, || H).expect("build_app_router");
5246 router
5247 }
5248
5249 #[tokio::test]
5250 async fn headers_on_rejected_origin_403() {
5251 let app = m5_router(|_| {});
5252 let req = Request::builder()
5253 .uri("/healthz")
5254 .header(header::ORIGIN, "http://evil.example")
5255 .body(Body::empty())
5256 .unwrap();
5257 let resp = app.oneshot(req).await.unwrap();
5258 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
5259 assert_owasp_headers(&resp, "origin-403");
5260 }
5261
5262 #[tokio::test]
5263 async fn headers_on_cors_preflight() {
5264 let app = m5_router(|_| {});
5265 let req = Request::builder()
5266 .method(axum::http::Method::OPTIONS)
5267 .uri("/mcp")
5268 .header(header::ORIGIN, "http://good.example")
5269 .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
5270 .body(Body::empty())
5271 .unwrap();
5272 let resp = app.oneshot(req).await.unwrap();
5273 assert_owasp_headers(&resp, "cors-preflight");
5274 }
5275
5276 #[tokio::test]
5277 async fn headers_on_404_fallback() {
5278 let app = m5_router(|_| {});
5279 let req = Request::builder()
5280 .uri("/no-such-route")
5281 .body(Body::empty())
5282 .unwrap();
5283 let resp = app.oneshot(req).await.unwrap();
5284 assert_eq!(resp.status(), StatusCode::NOT_FOUND);
5285 assert_owasp_headers(&resp, "404-fallback");
5286 }
5287
5288 #[tokio::test]
5289 async fn headers_on_overload_503() {
5290 let app = m5_router(|c| c.max_concurrent_requests = Some(0));
5293 let req = Request::builder()
5294 .uri("/healthz")
5295 .body(Body::empty())
5296 .unwrap();
5297 let resp = app.oneshot(req).await.unwrap();
5298 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
5299 assert_owasp_headers(&resp, "overload-503");
5300 }
5301
5302 #[cfg(feature = "oauth")]
5305 fn m6_auth_state() -> (Arc<AuthState>, String, String) {
5306 let (admin_token, admin_hash) = crate::auth::generate_api_key().unwrap();
5307 let (viewer_token, viewer_hash) = crate::auth::generate_api_key().unwrap();
5308 let state = Arc::new(AuthState {
5309 api_keys: ArcSwap::from_pointee(vec![
5310 crate::auth::ApiKeyEntry::new("admin-key", admin_hash, "admin"),
5311 crate::auth::ApiKeyEntry::new("viewer-key", viewer_hash, "viewer"),
5312 ]),
5313 rate_limiter: None,
5314 pre_auth_limiter: None,
5315 jwks_cache: None,
5316 seen_identities: crate::auth::SeenIdentitySet::new(),
5317 counters: crate::auth::AuthCounters::default(),
5318 });
5319 (state, admin_token, viewer_token)
5320 }
5321
5322 #[cfg(feature = "oauth")]
5323 fn m6_admin_router(state: &Arc<AuthState>) -> axum::Router {
5324 let proxy = crate::oauth::OAuthProxyConfig::builder(
5325 "https://idp.example/authorize",
5326 "https://idp.example/token",
5327 "client",
5328 )
5329 .introspection_url("http://127.0.0.1:1/introspect")
5330 .revocation_url("http://127.0.0.1:1/revoke")
5331 .expose_admin_endpoints(true)
5332 .require_auth_on_admin_endpoints(true)
5333 .build();
5334 let http = crate::oauth::OauthHttpClient::new().expect("oauth http client");
5335 build_oauth_admin_router(&proxy, http, Some(state), "admin").expect("admin router")
5336 }
5337
5338 #[cfg(feature = "oauth")]
5339 fn m6_req(path: &str, token: &str) -> Request<Body> {
5340 Request::builder()
5341 .method(axum::http::Method::POST)
5342 .uri(path)
5343 .header(header::AUTHORIZATION, format!("Bearer {token}"))
5344 .body(Body::from("token=abc"))
5345 .unwrap()
5346 }
5347
5348 #[cfg(feature = "oauth")]
5349 #[tokio::test]
5350 async fn oauth_proxy_admin_requires_admin_role() {
5351 let (state, _admin, viewer) = m6_auth_state();
5352 for path in ["/introspect", "/revoke"] {
5353 let app = m6_admin_router(&state);
5354 let resp = app.oneshot(m6_req(path, &viewer)).await.unwrap();
5355 assert_eq!(
5356 resp.status(),
5357 StatusCode::FORBIDDEN,
5358 "an authenticated viewer must be rejected with 403 on {path}"
5359 );
5360 }
5361 }
5362
5363 #[cfg(feature = "oauth")]
5364 #[tokio::test]
5365 async fn oauth_proxy_admin_allows_admin_role() {
5366 let (state, admin, _viewer) = m6_auth_state();
5367 for path in ["/introspect", "/revoke"] {
5368 let app = m6_admin_router(&state);
5369 let resp = app.oneshot(m6_req(path, &admin)).await.unwrap();
5370 assert_ne!(
5374 resp.status(),
5375 StatusCode::FORBIDDEN,
5376 "an authenticated admin must pass the role gate on {path}"
5377 );
5378 assert_ne!(
5379 resp.status(),
5380 StatusCode::UNAUTHORIZED,
5381 "an authenticated admin must pass the auth gate on {path}"
5382 );
5383 }
5384 }
5385
5386 #[cfg(feature = "metrics")]
5394 mod metrics_labels_bounded {
5395 use super::*;
5396
5397 fn labels_for(method: &str, uri: &str) -> (&'static str, String) {
5398 let req = Request::builder()
5399 .method(method)
5400 .uri(uri)
5401 .body(Body::empty())
5402 .unwrap();
5403 metrics_labels(&req)
5404 }
5405
5406 #[test]
5407 fn many_unmatched_paths_collapse_to_one_label() {
5408 let mut seen = std::collections::HashSet::new();
5409 for i in 0..500 {
5410 let (_, path) = labels_for("GET", &format!("/nonexistent-{i}"));
5411 seen.insert(path);
5412 }
5413 assert_eq!(
5414 seen.len(),
5415 1,
5416 "unmatched paths must collapse to a single label, got {seen:?}"
5417 );
5418 assert!(seen.contains("<unmatched>"));
5419 }
5420
5421 #[test]
5422 fn nested_mcp_paths_collapse_to_the_mount_point() {
5423 let mut seen = std::collections::HashSet::new();
5424 for i in 0..200 {
5425 let (_, path) = labels_for("POST", &format!("/mcp/{i}"));
5426 seen.insert(path);
5427 }
5428 let (_, root) = labels_for("POST", "/mcp");
5429 seen.insert(root);
5430 assert_eq!(
5431 seen.len(),
5432 1,
5433 "nested /mcp paths must collapse to one label, got {seen:?}"
5434 );
5435 assert!(seen.contains("/mcp"));
5436 }
5437
5438 #[test]
5439 fn unusual_methods_collapse_to_one_bucket() {
5440 let mut seen = std::collections::HashSet::new();
5441 for verb in ["FROBNICATE", "WIBBLE", "QUUX", "M-SEARCH"] {
5442 let (method, _) = labels_for(verb, "/healthz");
5443 seen.insert(method);
5444 }
5445 assert_eq!(seen, std::collections::HashSet::from(["OTHER"]));
5446 }
5447
5448 #[test]
5449 fn known_methods_keep_their_identity() {
5450 for verb in ["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"] {
5451 let (method, _) = labels_for(verb, "/healthz");
5452 assert_eq!(method, verb);
5453 }
5454 }
5455
5456 #[test]
5457 fn raw_path_never_leaks_into_a_label() {
5458 let (_, path) = labels_for("GET", "/secret-token-abc123");
5459 assert!(
5460 !path.contains("secret-token"),
5461 "raw request path must never become a label value: {path}"
5462 );
5463 }
5464 }
5465}