1use std::{
2 future::Future,
3 net::{IpAddr, SocketAddr},
4 path::{Path, PathBuf},
5 pin::Pin,
6 sync::Arc,
7 time::Duration,
8};
9
10use arc_swap::ArcSwap;
11use axum::{
12 body::Body,
13 extract::{ConnectInfo, Request},
14 middleware::Next,
15 response::IntoResponse,
16};
17use rmcp::{
18 ServerHandler,
19 transport::streamable_http_server::{
20 StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
21 },
22};
23use rustls::RootCertStore;
24use tokio::{
25 net::TcpListener,
26 sync::{Semaphore, mpsc},
27};
28use tokio_util::sync::CancellationToken;
29
30use crate::{
31 auth::{
32 AuthConfig, AuthIdentity, AuthState, MtlsConfig, TlsConnInfo, auth_middleware,
33 build_rate_limiter, extract_mtls_identity,
34 },
35 bounded_limiter::BoundedKeyedLimiter,
36 error::McpxError,
37 mtls_revocation::{self, CrlSet, DynamicClientCertVerifier},
38 rbac::{RbacPolicy, ToolRateLimiter, build_tool_rate_limiter, rbac_middleware},
39};
40
41#[allow(
45 clippy::needless_pass_by_value,
46 reason = "consumed at .map_err(anyhow_to_startup) call sites; by-value matches the closure shape"
47)]
48fn anyhow_to_startup(e: anyhow::Error) -> McpxError {
49 McpxError::Startup(format!("{e:#}"))
50}
51
52#[allow(
58 clippy::needless_pass_by_value,
59 reason = "consumed at .map_err(|e| io_to_startup(...)) call sites; by-value matches the closure shape"
60)]
61fn io_to_startup(op: &str, e: std::io::Error) -> McpxError {
62 McpxError::Startup(format!("{op}: {e}"))
63}
64
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 async fn from_request_parts(
144 parts: &mut axum::http::request::Parts,
145 _state: &S,
146 ) -> Result<Self, Self::Rejection> {
147 parts.extensions.get::<Self>().copied().ok_or((
148 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
149 "peer address unavailable: not running under rmcp-server-kit serve()",
150 ))
151 }
152}
153
154#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
177#[non_exhaustive]
178pub struct ClientIp {
179 pub ip: IpAddr,
181}
182
183impl ClientIp {
184 #[must_use]
187 pub(crate) const fn new(ip: IpAddr) -> Self {
188 Self { ip }
189 }
190}
191
192#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Deserialize)]
197#[serde(rename_all = "kebab-case")]
198#[non_exhaustive]
199pub enum ForwardedHeaderMode {
200 XForwardedFor,
202 Forwarded,
204}
205
206struct ForwardResolver {
209 trusted: Vec<ipnet::IpNet>,
210 mode: ForwardedHeaderMode,
211}
212
213#[derive(Debug, Clone, Default)]
234#[non_exhaustive]
235pub struct SecurityHeadersConfig {
236 pub x_content_type_options: Option<String>,
238 pub x_frame_options: Option<String>,
240 pub cache_control: Option<String>,
242 pub referrer_policy: Option<String>,
244 pub cross_origin_opener_policy: Option<String>,
246 pub cross_origin_resource_policy: Option<String>,
248 pub cross_origin_embedder_policy: Option<String>,
250 pub permissions_policy: Option<String>,
253 pub x_permitted_cross_domain_policies: Option<String>,
255 pub content_security_policy: Option<String>,
258 pub x_dns_prefetch_control: Option<String>,
260 pub strict_transport_security: Option<String>,
265}
266
267#[allow(
269 missing_debug_implementations,
270 reason = "contains callback/trait objects that don't impl Debug"
271)]
272#[allow(
273 clippy::struct_excessive_bools,
274 reason = "server configuration naturally has many boolean feature flags"
275)]
276#[non_exhaustive]
277pub struct McpServerConfig {
278 #[deprecated(
280 since = "0.13.0",
281 note = "use McpServerConfig::new() / with_bind_addr(); direct field access will become pub(crate) in a future major release"
282 )]
283 pub bind_addr: String,
284 #[deprecated(
286 since = "0.13.0",
287 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
288 )]
289 pub name: String,
290 #[deprecated(
292 since = "0.13.0",
293 note = "set via McpServerConfig::new(); direct field access will become pub(crate) in a future major release"
294 )]
295 pub version: String,
296 #[deprecated(
298 since = "0.13.0",
299 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
300 )]
301 pub tls_cert_path: Option<PathBuf>,
302 #[deprecated(
304 since = "0.13.0",
305 note = "use McpServerConfig::with_tls(); direct field access will become pub(crate) in a future major release"
306 )]
307 pub tls_key_path: Option<PathBuf>,
308 #[deprecated(
311 since = "0.13.0",
312 note = "use McpServerConfig::with_auth(); direct field access will become pub(crate) in a future major release"
313 )]
314 pub auth: Option<AuthConfig>,
315 #[deprecated(
318 since = "0.13.0",
319 note = "use McpServerConfig::with_rbac(); direct field access will become pub(crate) in a future major release"
320 )]
321 pub rbac: Option<Arc<RbacPolicy>>,
322 #[deprecated(
328 since = "0.13.0",
329 note = "use McpServerConfig::with_allowed_origins(); direct field access will become pub(crate) in a future major release"
330 )]
331 pub allowed_origins: Vec<String>,
332 #[deprecated(
335 since = "0.13.0",
336 note = "use McpServerConfig::with_tool_rate_limit(); direct field access will become pub(crate) in a future major release"
337 )]
338 pub tool_rate_limit: Option<u32>,
339 #[deprecated(
345 since = "1.12.0",
346 note = "use McpServerConfig::with_tool_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
347 )]
348 pub tool_rate_limit_burst: Option<u32>,
349 #[deprecated(
362 since = "1.11.0",
363 note = "use McpServerConfig::with_extra_route_rate_limit(); direct field access will become pub(crate) in a future major release"
364 )]
365 pub extra_route_rate_limit: Option<u32>,
366 #[deprecated(
373 since = "1.12.0",
374 note = "use McpServerConfig::with_extra_route_rate_limit_burst(); direct field access will become pub(crate) in a future major release"
375 )]
376 pub extra_route_rate_limit_burst: Option<u32>,
377 #[deprecated(
390 since = "1.14.0",
391 note = "use McpServerConfig::with_extra_route_rate_limit_exempt_paths(); direct field access will become pub(crate) in a future major release"
392 )]
393 pub extra_route_rate_limit_exempt_paths: Vec<String>,
394 #[deprecated(
402 since = "1.13.0",
403 note = "use McpServerConfig::with_trusted_proxies(); direct field access will become pub(crate) in a future major release"
404 )]
405 pub trusted_proxies: Vec<String>,
406 #[deprecated(
411 since = "1.13.0",
412 note = "use McpServerConfig::with_forwarded_header(); direct field access will become pub(crate) in a future major release"
413 )]
414 pub forwarded_header: Option<ForwardedHeaderMode>,
415 #[deprecated(
418 since = "0.13.0",
419 note = "use McpServerConfig::with_readiness_check(); direct field access will become pub(crate) in a future major release"
420 )]
421 pub readiness_check: Option<ReadinessCheck>,
422 #[deprecated(
425 since = "0.13.0",
426 note = "use McpServerConfig::with_max_request_body(); direct field access will become pub(crate) in a future major release"
427 )]
428 pub max_request_body: usize,
429 #[deprecated(
432 since = "0.13.0",
433 note = "use McpServerConfig::with_request_timeout(); direct field access will become pub(crate) in a future major release"
434 )]
435 pub request_timeout: Duration,
436 #[deprecated(
439 since = "0.13.0",
440 note = "use McpServerConfig::with_shutdown_timeout(); direct field access will become pub(crate) in a future major release"
441 )]
442 pub shutdown_timeout: Duration,
443 #[deprecated(
446 since = "0.13.0",
447 note = "use McpServerConfig::with_session_idle_timeout(); direct field access will become pub(crate) in a future major release"
448 )]
449 pub session_idle_timeout: Duration,
450 #[deprecated(
453 since = "0.13.0",
454 note = "use McpServerConfig::with_sse_keep_alive(); direct field access will become pub(crate) in a future major release"
455 )]
456 pub sse_keep_alive: Duration,
457 #[deprecated(
461 since = "0.13.0",
462 note = "use McpServerConfig::with_reload_callback(); direct field access will become pub(crate) in a future major release"
463 )]
464 pub on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
465 #[deprecated(
472 since = "0.13.0",
473 note = "use McpServerConfig::with_extra_router(); direct field access will become pub(crate) in a future major release"
474 )]
475 pub extra_router: Option<axum::Router>,
476 #[deprecated(
481 since = "0.13.0",
482 note = "use McpServerConfig::with_public_url(); direct field access will become pub(crate) in a future major release"
483 )]
484 pub public_url: Option<String>,
485 #[deprecated(
488 since = "0.13.0",
489 note = "use McpServerConfig::enable_request_header_logging(); direct field access will become pub(crate) in a future major release"
490 )]
491 pub log_request_headers: bool,
492 #[deprecated(
495 since = "0.13.0",
496 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
497 )]
498 pub compression_enabled: bool,
499 #[deprecated(
502 since = "0.13.0",
503 note = "use McpServerConfig::enable_compression(); direct field access will become pub(crate) in a future major release"
504 )]
505 pub compression_min_size: u16,
506 #[deprecated(
510 since = "0.13.0",
511 note = "use McpServerConfig::with_max_concurrent_requests(); direct field access will become pub(crate) in a future major release"
512 )]
513 pub max_concurrent_requests: Option<usize>,
514 #[deprecated(
517 since = "0.13.0",
518 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
519 )]
520 pub admin_enabled: bool,
521 #[deprecated(
523 since = "0.13.0",
524 note = "use McpServerConfig::enable_admin(); direct field access will become pub(crate) in a future major release"
525 )]
526 pub admin_role: String,
527 #[cfg(feature = "metrics")]
530 #[deprecated(
531 since = "0.13.0",
532 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
533 )]
534 pub metrics_enabled: bool,
535 #[cfg(feature = "metrics")]
537 #[deprecated(
538 since = "0.13.0",
539 note = "use McpServerConfig::with_metrics(); direct field access will become pub(crate) in a future major release"
540 )]
541 pub metrics_bind: String,
542 #[deprecated(
546 since = "1.5.0",
547 note = "use McpServerConfig::with_security_headers(); direct field access will become pub(crate) in a future major release"
548 )]
549 pub security_headers: SecurityHeadersConfig,
550 #[deprecated(
556 since = "1.9.0",
557 note = "use McpServerConfig::with_tls_handshake_timeout(); direct field access will become pub(crate) in a future major release"
558 )]
559 pub tls_handshake_timeout: Duration,
560 #[deprecated(
567 since = "1.9.0",
568 note = "use McpServerConfig::with_max_concurrent_tls_handshakes(); direct field access will become pub(crate) in a future major release"
569 )]
570 pub max_concurrent_tls_handshakes: usize,
571}
572
573#[allow(
631 missing_debug_implementations,
632 reason = "wraps T which may not implement Debug; manual impl below avoids leaking inner contents into logs"
633)]
634pub struct Validated<T>(T);
635
636impl<T> std::fmt::Debug for Validated<T> {
637 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
638 f.debug_struct("Validated").finish_non_exhaustive()
639 }
640}
641
642impl<T> Validated<T> {
643 #[must_use]
645 pub fn as_inner(&self) -> &T {
646 &self.0
647 }
648
649 #[must_use]
654 pub fn into_inner(self) -> T {
655 self.0
656 }
657}
658
659#[allow(
660 deprecated,
661 reason = "internal builders/validators legitimately read/write the deprecated `pub` fields they were designed to manage"
662)]
663impl McpServerConfig {
664 #[must_use]
672 pub fn new(
673 bind_addr: impl Into<String>,
674 name: impl Into<String>,
675 version: impl Into<String>,
676 ) -> Self {
677 Self {
678 bind_addr: bind_addr.into(),
679 name: name.into(),
680 version: version.into(),
681 tls_cert_path: None,
682 tls_key_path: None,
683 auth: None,
684 rbac: None,
685 allowed_origins: Vec::new(),
686 tool_rate_limit: None,
687 readiness_check: None,
688 max_request_body: 1024 * 1024,
689 request_timeout: Duration::from_mins(2),
690 shutdown_timeout: Duration::from_secs(30),
691 session_idle_timeout: Duration::from_mins(20),
692 sse_keep_alive: Duration::from_secs(15),
693 on_reload_ready: None,
694 extra_router: None,
695 public_url: None,
696 log_request_headers: false,
697 compression_enabled: false,
698 compression_min_size: 1024,
699 max_concurrent_requests: None,
700 admin_enabled: false,
701 admin_role: "admin".to_owned(),
702 #[cfg(feature = "metrics")]
703 metrics_enabled: false,
704 #[cfg(feature = "metrics")]
705 metrics_bind: "127.0.0.1:9090".into(),
706 security_headers: SecurityHeadersConfig::default(),
707 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
708 max_concurrent_tls_handshakes: DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES,
709 extra_route_rate_limit: None,
710 tool_rate_limit_burst: None,
711 extra_route_rate_limit_burst: None,
712 extra_route_rate_limit_exempt_paths: Vec::new(),
713 trusted_proxies: Vec::new(),
714 forwarded_header: None,
715 }
716 }
717
718 #[must_use]
728 pub fn with_auth(mut self, auth: AuthConfig) -> Self {
729 self.auth = Some(auth);
730 self
731 }
732
733 #[must_use]
738 pub fn with_security_headers(mut self, headers: SecurityHeadersConfig) -> Self {
739 self.security_headers = headers;
740 self
741 }
742
743 #[must_use]
747 pub fn with_bind_addr(mut self, addr: impl Into<String>) -> Self {
748 self.bind_addr = addr.into();
749 self
750 }
751
752 #[must_use]
755 pub fn with_rbac(mut self, rbac: Arc<RbacPolicy>) -> Self {
756 self.rbac = Some(rbac);
757 self
758 }
759
760 #[must_use]
764 pub fn with_tls(mut self, cert_path: impl Into<PathBuf>, key_path: impl Into<PathBuf>) -> Self {
765 self.tls_cert_path = Some(cert_path.into());
766 self.tls_key_path = Some(key_path.into());
767 self
768 }
769
770 #[must_use]
774 pub fn with_public_url(mut self, url: impl Into<String>) -> Self {
775 self.public_url = Some(url.into());
776 self
777 }
778
779 #[must_use]
783 pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
784 where
785 I: IntoIterator<Item = S>,
786 S: Into<String>,
787 {
788 self.allowed_origins = origins.into_iter().map(Into::into).collect();
789 self
790 }
791
792 #[must_use]
805 pub fn with_extra_router(mut self, router: axum::Router) -> Self {
806 self.extra_router = Some(router);
807 self
808 }
809
810 #[must_use]
813 pub fn with_readiness_check(mut self, check: ReadinessCheck) -> Self {
814 self.readiness_check = Some(check);
815 self
816 }
817
818 #[must_use]
821 pub fn with_max_request_body(mut self, bytes: usize) -> Self {
822 self.max_request_body = bytes;
823 self
824 }
825
826 #[must_use]
828 pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
829 self.request_timeout = timeout;
830 self
831 }
832
833 #[must_use]
835 pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
836 self.shutdown_timeout = timeout;
837 self
838 }
839
840 #[must_use]
842 pub fn with_session_idle_timeout(mut self, timeout: Duration) -> Self {
843 self.session_idle_timeout = timeout;
844 self
845 }
846
847 #[must_use]
849 pub fn with_sse_keep_alive(mut self, interval: Duration) -> Self {
850 self.sse_keep_alive = interval;
851 self
852 }
853
854 #[must_use]
858 pub fn with_max_concurrent_requests(mut self, limit: usize) -> Self {
859 self.max_concurrent_requests = Some(limit);
860 self
861 }
862
863 #[must_use]
871 pub fn with_tls_handshake_timeout(mut self, timeout: Duration) -> Self {
872 self.tls_handshake_timeout = timeout;
873 self
874 }
875
876 #[must_use]
885 pub fn with_max_concurrent_tls_handshakes(mut self, limit: usize) -> Self {
886 self.max_concurrent_tls_handshakes = limit;
887 self
888 }
889
890 #[must_use]
893 pub fn with_tool_rate_limit(mut self, per_minute: u32) -> Self {
894 self.tool_rate_limit = Some(per_minute);
895 self
896 }
897
898 #[must_use]
909 pub fn with_extra_route_rate_limit(mut self, per_minute: u32) -> Self {
910 self.extra_route_rate_limit = Some(per_minute);
911 self
912 }
913
914 #[must_use]
919 pub fn with_tool_rate_limit_burst(mut self, burst: u32) -> Self {
920 self.tool_rate_limit_burst = Some(burst);
921 self
922 }
923
924 #[must_use]
930 pub fn with_extra_route_rate_limit_burst(mut self, burst: u32) -> Self {
931 self.extra_route_rate_limit_burst = Some(burst);
932 self
933 }
934
935 #[must_use]
955 pub fn with_extra_route_rate_limit_exempt_paths<I, S>(mut self, paths: I) -> Self
956 where
957 I: IntoIterator<Item = S>,
958 S: Into<String>,
959 {
960 self.extra_route_rate_limit_exempt_paths = paths.into_iter().map(Into::into).collect();
961 self
962 }
963
964 #[must_use]
976 pub fn with_trusted_proxies<I, S>(mut self, proxies: I) -> Self
977 where
978 I: IntoIterator<Item = S>,
979 S: Into<String>,
980 {
981 self.trusted_proxies = proxies.into_iter().map(Into::into).collect();
982 self
983 }
984
985 #[must_use]
990 pub fn with_forwarded_header(mut self, mode: ForwardedHeaderMode) -> Self {
991 self.forwarded_header = Some(mode);
992 self
993 }
994
995 #[must_use]
999 pub fn with_reload_callback<F>(mut self, callback: F) -> Self
1000 where
1001 F: FnOnce(ReloadHandle) + Send + 'static,
1002 {
1003 self.on_reload_ready = Some(Box::new(callback));
1004 self
1005 }
1006
1007 #[must_use]
1011 pub fn enable_compression(mut self, min_size: u16) -> Self {
1012 self.compression_enabled = true;
1013 self.compression_min_size = min_size;
1014 self
1015 }
1016
1017 #[must_use]
1022 pub fn enable_admin(mut self, role: impl Into<String>) -> Self {
1023 self.admin_enabled = true;
1024 self.admin_role = role.into();
1025 self
1026 }
1027
1028 #[must_use]
1031 pub fn enable_request_header_logging(mut self) -> Self {
1032 self.log_request_headers = true;
1033 self
1034 }
1035
1036 #[cfg(feature = "metrics")]
1039 #[must_use]
1040 pub fn with_metrics(mut self, bind: impl Into<String>) -> Self {
1041 self.metrics_enabled = true;
1042 self.metrics_bind = bind.into();
1043 self
1044 }
1045
1046 pub fn validate(self) -> Result<Validated<Self>, McpxError> {
1079 self.check()?;
1080 Ok(Validated(self))
1081 }
1082
1083 fn check_burst_knobs(&self) -> Result<(), McpxError> {
1090 if self.tool_rate_limit_burst == Some(0) {
1091 return Err(McpxError::Config(
1092 "tool_rate_limit_burst must be greater than zero".into(),
1093 ));
1094 }
1095 if self.extra_route_rate_limit_burst == Some(0) {
1096 return Err(McpxError::Config(
1097 "extra_route_rate_limit_burst must be greater than zero".into(),
1098 ));
1099 }
1100 if self.tool_rate_limit_burst.is_some() && self.tool_rate_limit.is_none() {
1101 return Err(McpxError::Config(
1102 "tool_rate_limit_burst requires tool_rate_limit to be set".into(),
1103 ));
1104 }
1105 if self.extra_route_rate_limit_burst.is_some() && self.extra_route_rate_limit.is_none() {
1106 return Err(McpxError::Config(
1107 "extra_route_rate_limit_burst requires extra_route_rate_limit to be set".into(),
1108 ));
1109 }
1110 if !self.extra_route_rate_limit_exempt_paths.is_empty()
1111 && self.extra_route_rate_limit.is_none()
1112 {
1113 return Err(McpxError::Config(
1114 "extra_route_rate_limit_exempt_paths requires extra_route_rate_limit to be set"
1115 .into(),
1116 ));
1117 }
1118 for path in &self.extra_route_rate_limit_exempt_paths {
1119 if path.is_empty() || !path.starts_with('/') {
1120 return Err(McpxError::Config(format!(
1121 "extra_route_rate_limit_exempt_paths entries must be non-empty and start with '/': {path:?}"
1122 )));
1123 }
1124 }
1125 if let Some(rl) = self.auth.as_ref().and_then(|a| a.rate_limit.as_ref()) {
1126 if rl.burst == Some(0) {
1127 return Err(McpxError::Config(
1128 "auth rate_limit.burst must be greater than zero".into(),
1129 ));
1130 }
1131 if rl.pre_auth_burst == Some(0) {
1132 return Err(McpxError::Config(
1133 "auth rate_limit.pre_auth_burst must be greater than zero".into(),
1134 ));
1135 }
1136 }
1137 Ok(())
1138 }
1139
1140 fn check_trusted_forwarder(&self) -> Result<(), McpxError> {
1145 for entry in &self.trusted_proxies {
1146 if parse_proxy_net(entry).is_none() {
1147 return Err(McpxError::Config(format!(
1148 "trusted_proxies entry {entry:?} is neither a CIDR nor an IP address"
1149 )));
1150 }
1151 }
1152 if self.forwarded_header.is_some() && self.trusted_proxies.is_empty() {
1153 return Err(McpxError::Config(
1154 "forwarded_header requires trusted_proxies to be nonempty".into(),
1155 ));
1156 }
1157 Ok(())
1158 }
1159
1160 fn check(&self) -> Result<(), McpxError> {
1164 if self.admin_enabled {
1168 let auth_enabled = self.auth.as_ref().is_some_and(|a| a.enabled);
1169 if !auth_enabled {
1170 return Err(McpxError::Config(
1171 "admin_enabled=true requires auth to be configured and enabled".into(),
1172 ));
1173 }
1174 }
1175
1176 match (&self.tls_cert_path, &self.tls_key_path) {
1178 (Some(_), None) => {
1179 return Err(McpxError::Config(
1180 "tls_cert_path is set but tls_key_path is missing".into(),
1181 ));
1182 }
1183 (None, Some(_)) => {
1184 return Err(McpxError::Config(
1185 "tls_key_path is set but tls_cert_path is missing".into(),
1186 ));
1187 }
1188 _ => {}
1189 }
1190
1191 if self.bind_addr.parse::<SocketAddr>().is_err() {
1193 return Err(McpxError::Config(format!(
1194 "bind_addr {:?} is not a valid socket address (expected e.g. 127.0.0.1:8080)",
1195 self.bind_addr
1196 )));
1197 }
1198
1199 if let Some(ref url) = self.public_url
1201 && !(url.starts_with("http://") || url.starts_with("https://"))
1202 {
1203 return Err(McpxError::Config(format!(
1204 "public_url {url:?} must start with http:// or https://"
1205 )));
1206 }
1207
1208 for origin in &self.allowed_origins {
1210 if !(origin.starts_with("http://") || origin.starts_with("https://")) {
1211 return Err(McpxError::Config(format!(
1212 "allowed_origins entry {origin:?} must start with http:// or https://"
1213 )));
1214 }
1215 }
1216
1217 if self.max_request_body == 0 {
1219 return Err(McpxError::Config(
1220 "max_request_body must be greater than zero".into(),
1221 ));
1222 }
1223
1224 if self.extra_route_rate_limit == Some(0) {
1228 return Err(McpxError::Config(
1229 "extra_route_rate_limit must be greater than zero".into(),
1230 ));
1231 }
1232
1233 self.check_burst_knobs()?;
1235
1236 self.check_trusted_forwarder()?;
1238
1239 #[cfg(feature = "oauth")]
1241 if let Some(auth_cfg) = &self.auth
1242 && let Some(oauth_cfg) = &auth_cfg.oauth
1243 {
1244 oauth_cfg.validate()?;
1245 }
1246
1247 validate_security_headers(&self.security_headers)?;
1250
1251 if self.max_concurrent_requests == Some(0) {
1255 return Err(McpxError::Config(
1256 "max_concurrent_requests must be greater than zero when set".into(),
1257 ));
1258 }
1259
1260 if let Some(auth_cfg) = &self.auth
1264 && let Some(rl) = &auth_cfg.rate_limit
1265 && rl.max_tracked_keys == 0
1266 {
1267 return Err(McpxError::Config(
1268 "auth.rate_limit.max_tracked_keys must be greater than zero".into(),
1269 ));
1270 }
1271
1272 if self.tls_handshake_timeout == Duration::ZERO {
1277 return Err(McpxError::Config(
1278 "tls_handshake_timeout must be greater than zero".into(),
1279 ));
1280 }
1281
1282 if self.max_concurrent_tls_handshakes == 0 {
1287 return Err(McpxError::Config(
1288 "max_concurrent_tls_handshakes must be greater than zero".into(),
1289 ));
1290 }
1291
1292 Ok(())
1293 }
1294}
1295
1296#[allow(
1302 missing_debug_implementations,
1303 reason = "contains Arc<AuthState> with non-Debug fields"
1304)]
1305pub struct ReloadHandle {
1306 auth: Option<Arc<AuthState>>,
1307 rbac: Option<Arc<ArcSwap<RbacPolicy>>>,
1308 crl_set: Option<Arc<CrlSet>>,
1309}
1310
1311impl ReloadHandle {
1312 pub fn reload_auth_keys(&self, keys: Vec<crate::auth::ApiKeyEntry>) {
1314 if let Some(ref auth) = self.auth {
1315 auth.reload_keys(keys);
1316 }
1317 }
1318
1319 pub fn reload_rbac(&self, policy: RbacPolicy) {
1321 if let Some(ref rbac) = self.rbac {
1322 rbac.store(Arc::new(policy));
1323 tracing::info!("RBAC policy reloaded");
1324 }
1325 }
1326
1327 pub async fn refresh_crls(&self) -> Result<(), McpxError> {
1333 let Some(ref crl_set) = self.crl_set else {
1334 return Err(McpxError::Config(
1335 "CRL refresh requested but mTLS CRL support is not configured".into(),
1336 ));
1337 };
1338
1339 crl_set.force_refresh().await
1340 }
1341}
1342
1343#[allow(
1360 clippy::too_many_lines,
1361 clippy::cognitive_complexity,
1362 reason = "middleware layer order is security-critical and must remain visible at one glance; extracting `&mut Router` helpers would obscure the auth/RBAC/origin/rate-limit ordering"
1363)]
1364struct AppRunParams {
1368 tls_paths: Option<(PathBuf, PathBuf)>,
1370 tls_handshake_timeout: Duration,
1372 max_concurrent_tls_handshakes: usize,
1374 mtls_config: Option<MtlsConfig>,
1376 shutdown_timeout: Duration,
1378 auth_state: Option<Arc<AuthState>>,
1380 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
1382 on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
1384 ct: CancellationToken,
1388 scheme: &'static str,
1390 name: String,
1392}
1393
1394#[allow(
1404 clippy::cognitive_complexity,
1405 reason = "router assembly is intrinsically sequential; splitting harms readability"
1406)]
1407#[allow(
1408 deprecated,
1409 reason = "internal router assembly reads deprecated `pub` config fields by design until 1.0 makes them pub(crate)"
1410)]
1411fn build_app_router<H, F>(
1412 mut config: McpServerConfig,
1413 handler_factory: F,
1414) -> anyhow::Result<(axum::Router, AppRunParams)>
1415where
1416 H: ServerHandler + 'static,
1417 F: Fn() -> H + Send + Sync + Clone + 'static,
1418{
1419 let ct = CancellationToken::new();
1420
1421 let allowed_hosts = derive_allowed_hosts(&config.bind_addr, config.public_url.as_deref());
1422 tracing::info!(allowed_hosts = %allowed_hosts.join(", "), "configured Streamable HTTP allowed hosts");
1423
1424 let mcp_service = StreamableHttpService::new(
1425 move || Ok(handler_factory()),
1426 {
1427 let mut mgr = LocalSessionManager::default();
1428 mgr.session_config.keep_alive = Some(config.session_idle_timeout);
1429 mgr.into()
1430 },
1431 StreamableHttpServerConfig::default()
1432 .with_allowed_hosts(allowed_hosts)
1433 .with_sse_keep_alive(Some(config.sse_keep_alive))
1434 .with_cancellation_token(ct.child_token()),
1435 );
1436
1437 let mut mcp_router = axum::Router::new().nest_service("/mcp", mcp_service);
1439
1440 let auth_state: Option<Arc<AuthState>> = match config.auth {
1444 Some(ref auth_config) if auth_config.enabled => {
1445 let rate_limiter = auth_config.rate_limit.as_ref().map(build_rate_limiter);
1446 let pre_auth_limiter = auth_config
1447 .rate_limit
1448 .as_ref()
1449 .map(crate::auth::build_pre_auth_limiter);
1450
1451 #[cfg(feature = "oauth")]
1452 let jwks_cache = auth_config
1453 .oauth
1454 .as_ref()
1455 .map(|c| crate::oauth::JwksCache::new(c).map(Arc::new))
1456 .transpose()
1457 .map_err(|e| std::io::Error::other(format!("JWKS HTTP client: {e}")))?;
1458
1459 Some(Arc::new(AuthState {
1460 api_keys: ArcSwap::new(Arc::new(auth_config.api_keys.clone())),
1461 rate_limiter,
1462 pre_auth_limiter,
1463 #[cfg(feature = "oauth")]
1464 jwks_cache,
1465 seen_identities: crate::auth::SeenIdentitySet::new(),
1466 counters: crate::auth::AuthCounters::default(),
1467 }))
1468 }
1469 _ => None,
1470 };
1471
1472 let rbac_swap = Arc::new(ArcSwap::new(
1475 config
1476 .rbac
1477 .clone()
1478 .unwrap_or_else(|| Arc::new(RbacPolicy::disabled())),
1479 ));
1480
1481 if config.admin_enabled {
1484 let Some(ref auth_state_ref) = auth_state else {
1485 return Err(anyhow::anyhow!(
1486 "admin_enabled=true requires auth to be configured and enabled"
1487 ));
1488 };
1489 let admin_state = crate::admin::AdminState {
1490 started_at: std::time::Instant::now(),
1491 name: config.name.clone(),
1492 version: config.version.clone(),
1493 auth: Some(Arc::clone(auth_state_ref)),
1494 rbac: Arc::clone(&rbac_swap),
1495 };
1496 let admin_cfg = crate::admin::AdminConfig {
1497 role: config.admin_role.clone(),
1498 };
1499 mcp_router = mcp_router.merge(crate::admin::admin_router(admin_state, &admin_cfg));
1500 tracing::info!(role = %config.admin_role, "/admin/* endpoints enabled");
1501 }
1502
1503 {
1536 let tool_limiter: Option<Arc<ToolRateLimiter>> = config
1537 .tool_rate_limit
1538 .map(|per_minute| build_tool_rate_limiter(per_minute, config.tool_rate_limit_burst));
1539
1540 if rbac_swap.load().is_enabled() {
1541 tracing::info!("RBAC enforcement enabled on /mcp");
1542 }
1543 if let Some(limit) = config.tool_rate_limit {
1544 tracing::info!(limit, "tool rate limiting enabled (calls/min per IP)");
1545 }
1546
1547 let rbac_for_mw = Arc::clone(&rbac_swap);
1548 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1549 let p = rbac_for_mw.load_full();
1550 let tl = tool_limiter.clone();
1551 rbac_middleware(p, tl, req, next)
1552 }));
1553 }
1554
1555 if let Some(ref auth_config) = config.auth
1557 && auth_config.enabled
1558 {
1559 let Some(ref state) = auth_state else {
1560 return Err(anyhow::anyhow!("auth state missing despite enabled config"));
1561 };
1562
1563 let methods: Vec<&str> = [
1564 auth_config.mtls.is_some().then_some("mTLS"),
1565 (!auth_config.api_keys.is_empty()).then_some("bearer"),
1566 #[cfg(feature = "oauth")]
1567 auth_config.oauth.is_some().then_some("oauth-jwt"),
1568 ]
1569 .into_iter()
1570 .flatten()
1571 .collect();
1572
1573 tracing::info!(
1574 methods = %methods.join(", "),
1575 api_keys = auth_config.api_keys.len(),
1576 "auth enabled on /mcp"
1577 );
1578
1579 let state_for_mw = Arc::clone(state);
1580 mcp_router = mcp_router.layer(axum::middleware::from_fn(move |req, next| {
1581 let s = Arc::clone(&state_for_mw);
1582 auth_middleware(s, req, next)
1583 }));
1584 }
1585
1586 mcp_router = mcp_router.layer(tower_http::timeout::TimeoutLayer::with_status_code(
1589 axum::http::StatusCode::REQUEST_TIMEOUT,
1590 config.request_timeout,
1591 ));
1592
1593 mcp_router = mcp_router.layer(tower_http::limit::RequestBodyLimitLayer::new(
1597 config.max_request_body,
1598 ));
1599
1600 let mut effective_origins = config.allowed_origins.clone();
1607 if effective_origins.is_empty()
1608 && let Some(ref url) = config.public_url
1609 {
1610 if let Some(scheme_end) = url.find("://") {
1615 let scheme_with_sep = url.get(..scheme_end + 3).unwrap_or_default();
1616 let after_scheme = url.get(scheme_end + 3..).unwrap_or_default();
1617 let host_end = after_scheme.find('/').unwrap_or(after_scheme.len());
1618 let host = after_scheme.get(..host_end).unwrap_or_default();
1619 let origin = format!("{scheme_with_sep}{host}");
1620 tracing::info!(
1621 %origin,
1622 "auto-derived allowed origin from public_url"
1623 );
1624 effective_origins.push(origin);
1625 }
1626 }
1627 let allowed_origins: Arc<[String]> = Arc::from(effective_origins);
1628 let cors_origins = Arc::clone(&allowed_origins);
1629 let log_request_headers = config.log_request_headers;
1630
1631 let readyz_route = if let Some(check) = config.readiness_check.take() {
1632 axum::routing::get(move || readyz(Arc::clone(&check)))
1633 } else {
1634 axum::routing::get(healthz)
1635 };
1636
1637 #[allow(unused_mut)] let mut router = axum::Router::new()
1639 .route("/healthz", axum::routing::get(healthz))
1640 .route("/readyz", readyz_route)
1641 .route(
1642 "/version",
1643 axum::routing::get({
1644 let payload_bytes: Arc<[u8]> =
1649 serialize_version_payload(&config.name, &config.version);
1650 move || {
1651 let p = Arc::clone(&payload_bytes);
1652 async move {
1653 (
1654 [(axum::http::header::CONTENT_TYPE, "application/json")],
1655 p.to_vec(),
1656 )
1657 }
1658 }
1659 }),
1660 )
1661 .merge(mcp_router);
1662
1663 if let Some(extra) = config.extra_router.take() {
1670 let extra = match config.extra_route_rate_limit {
1671 Some(per_minute) => {
1672 let limiter =
1673 build_extra_route_rate_limiter(per_minute, config.extra_route_rate_limit_burst);
1674 let exempt: Arc<std::collections::HashSet<String>> = Arc::new(
1675 config
1676 .extra_route_rate_limit_exempt_paths
1677 .iter()
1678 .cloned()
1679 .collect(),
1680 );
1681 tracing::info!(
1682 per_minute,
1683 exempt_paths = exempt.len(),
1684 "extra-route per-IP rate limit enabled"
1685 );
1686 extra.layer(axum::middleware::from_fn(move |req, next| {
1687 let l = Arc::clone(&limiter);
1688 let e = Arc::clone(&exempt);
1689 extra_route_rate_limit_middleware(l, e, req, next)
1690 }))
1691 }
1692 None => extra,
1693 };
1694 router = router.merge(extra);
1695 }
1696
1697 let server_url = if let Some(ref url) = config.public_url {
1704 url.trim_end_matches('/').to_owned()
1705 } else {
1706 let prm_scheme = if config.tls_cert_path.is_some() {
1707 "https"
1708 } else {
1709 "http"
1710 };
1711 format!("{prm_scheme}://{}", config.bind_addr)
1712 };
1713 let resource_url = format!("{server_url}/mcp");
1714
1715 #[cfg(feature = "oauth")]
1716 let prm_metadata = if let Some(ref auth_config) = config.auth
1717 && let Some(ref oauth_config) = auth_config.oauth
1718 {
1719 crate::oauth::protected_resource_metadata(&resource_url, &server_url, oauth_config)
1720 } else {
1721 serde_json::json!({ "resource": resource_url })
1722 };
1723 #[cfg(not(feature = "oauth"))]
1724 let prm_metadata = serde_json::json!({ "resource": resource_url });
1725
1726 router = router.route(
1727 "/.well-known/oauth-protected-resource",
1728 axum::routing::get(move || {
1729 let m = prm_metadata.clone();
1730 async move { axum::Json(m) }
1731 }),
1732 );
1733
1734 #[cfg(feature = "oauth")]
1739 if let Some(ref auth_config) = config.auth
1740 && let Some(ref oauth_config) = auth_config.oauth
1741 && oauth_config.proxy.is_some()
1742 {
1743 router = install_oauth_proxy_routes(
1744 router,
1745 &server_url,
1746 oauth_config,
1747 auth_state.as_ref(),
1748 config.max_request_body,
1749 )?;
1750 }
1751
1752 let is_tls = config.tls_cert_path.is_some();
1755 let security_headers_cfg = Arc::new(config.security_headers.clone());
1756 router = router.layer(axum::middleware::from_fn(move |req, next| {
1757 let cfg = Arc::clone(&security_headers_cfg);
1758 security_headers_middleware(is_tls, cfg, req, next)
1759 }));
1760
1761 if !cors_origins.is_empty() {
1765 let cors = tower_http::cors::CorsLayer::new()
1766 .allow_origin(
1767 cors_origins
1768 .iter()
1769 .filter_map(|o| o.parse::<axum::http::HeaderValue>().ok())
1770 .collect::<Vec<_>>(),
1771 )
1772 .allow_methods([
1773 axum::http::Method::GET,
1774 axum::http::Method::POST,
1775 axum::http::Method::OPTIONS,
1776 ])
1777 .allow_headers([
1778 axum::http::header::CONTENT_TYPE,
1779 axum::http::header::AUTHORIZATION,
1780 ]);
1781 router = router.layer(cors);
1782 }
1783
1784 if config.compression_enabled {
1788 use tower_http::compression::Predicate as _;
1789 let predicate = tower_http::compression::DefaultPredicate::new().and(
1790 tower_http::compression::predicate::SizeAbove::new(u64::from(
1791 config.compression_min_size,
1792 )),
1793 );
1794 router = router.layer(
1795 tower_http::compression::CompressionLayer::new()
1796 .gzip(true)
1797 .br(true)
1798 .compress_when(predicate),
1799 );
1800 tracing::info!(
1801 min_size = config.compression_min_size,
1802 "response compression enabled (gzip, br)"
1803 );
1804 }
1805
1806 if let Some(max) = config.max_concurrent_requests {
1809 let overload_handler = tower::ServiceBuilder::new()
1810 .layer(axum::error_handling::HandleErrorLayer::new(
1811 |_err: tower::BoxError| async {
1812 (
1813 axum::http::StatusCode::SERVICE_UNAVAILABLE,
1814 axum::Json(serde_json::json!({
1815 "error": "overloaded",
1816 "error_description": "server is at capacity, retry later"
1817 })),
1818 )
1819 },
1820 ))
1821 .layer(tower::load_shed::LoadShedLayer::new())
1822 .layer(tower::limit::ConcurrencyLimitLayer::new(max));
1823 router = router.layer(overload_handler);
1824 tracing::info!(max, "global concurrency limit enabled");
1825 }
1826
1827 router = router.fallback(|| async {
1831 (
1832 axum::http::StatusCode::NOT_FOUND,
1833 axum::Json(serde_json::json!({
1834 "error": "not_found",
1835 "error_description": "The requested endpoint does not exist"
1836 })),
1837 )
1838 });
1839
1840 #[cfg(feature = "metrics")]
1842 if config.metrics_enabled {
1843 let metrics = Arc::new(
1844 crate::metrics::McpMetrics::new().map_err(|e| anyhow::anyhow!("metrics init: {e}"))?,
1845 );
1846 let m = Arc::clone(&metrics);
1847 router = router.layer(axum::middleware::from_fn(
1848 move |req: Request<Body>, next: Next| {
1849 let m = Arc::clone(&m);
1850 metrics_middleware(m, req, next)
1851 },
1852 ));
1853 let metrics_bind = config.metrics_bind.clone();
1854 let metrics_shutdown = ct.clone();
1855 tokio::spawn(async move {
1856 if let Err(e) =
1857 crate::metrics::serve_metrics(metrics_bind, metrics, metrics_shutdown).await
1858 {
1859 tracing::error!("metrics listener failed: {e}");
1860 }
1861 });
1862 }
1863
1864 let forward_resolver: Option<Arc<ForwardResolver>> = if config.trusted_proxies.is_empty() {
1872 None
1873 } else {
1874 Some(Arc::new(ForwardResolver {
1877 trusted: config
1878 .trusted_proxies
1879 .iter()
1880 .filter_map(|entry| parse_proxy_net(entry))
1881 .collect(),
1882 mode: config
1883 .forwarded_header
1884 .unwrap_or(ForwardedHeaderMode::XForwardedFor),
1885 }))
1886 };
1887 if forward_resolver.is_some() {
1888 tracing::info!(
1889 proxies = config.trusted_proxies.len(),
1890 "trusted-forwarder mode enabled: limiters key by resolved client IP"
1891 );
1892 }
1893 router = router.layer(axum::middleware::from_fn(move |req, next| {
1894 let r = forward_resolver.clone();
1895 normalize_peer_addr_middleware(r, req, next)
1896 }));
1897
1898 router = router.layer(axum::middleware::from_fn(move |req, next| {
1909 let origins = Arc::clone(&allowed_origins);
1910 origin_check_middleware(origins, log_request_headers, req, next)
1911 }));
1912
1913 let scheme = if config.tls_cert_path.is_some() {
1914 "https"
1915 } else {
1916 "http"
1917 };
1918
1919 let tls_paths = match (&config.tls_cert_path, &config.tls_key_path) {
1920 (Some(cert), Some(key)) => Some((cert.clone(), key.clone())),
1921 _ => None,
1922 };
1923 let tls_handshake_timeout = config.tls_handshake_timeout;
1924 let max_concurrent_tls_handshakes = config.max_concurrent_tls_handshakes;
1925 let mtls_config = config.auth.as_ref().and_then(|a| a.mtls.as_ref()).cloned();
1926
1927 Ok((
1928 router,
1929 AppRunParams {
1930 tls_paths,
1931 tls_handshake_timeout,
1932 max_concurrent_tls_handshakes,
1933 mtls_config,
1934 shutdown_timeout: config.shutdown_timeout,
1935 auth_state,
1936 rbac_swap,
1937 on_reload_ready: config.on_reload_ready.take(),
1938 ct,
1939 scheme,
1940 name: config.name.clone(),
1941 },
1942 ))
1943}
1944
1945pub async fn serve<H, F>(
1962 config: Validated<McpServerConfig>,
1963 handler_factory: F,
1964) -> Result<(), McpxError>
1965where
1966 H: ServerHandler + 'static,
1967 F: Fn() -> H + Send + Sync + Clone + 'static,
1968{
1969 let config = config.into_inner();
1970 #[allow(
1971 deprecated,
1972 reason = "internal serve() reads `bind_addr` to construct the listener; field becomes pub(crate) in 1.0"
1973 )]
1974 let bind_addr = config.bind_addr.clone();
1975 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
1976
1977 let listener = TcpListener::bind(&bind_addr)
1978 .await
1979 .map_err(|e| io_to_startup(&format!("bind {bind_addr}"), e))?;
1980 log_listening(¶ms.name, params.scheme, &bind_addr);
1981
1982 run_server(
1983 router,
1984 listener,
1985 params.tls_paths,
1986 params.tls_handshake_timeout,
1987 params.max_concurrent_tls_handshakes,
1988 params.mtls_config,
1989 params.shutdown_timeout,
1990 params.auth_state,
1991 params.rbac_swap,
1992 params.on_reload_ready,
1993 params.ct,
1994 )
1995 .await
1996 .map_err(anyhow_to_startup)
1997}
1998
1999pub async fn serve_with_listener<H, F>(
2029 listener: TcpListener,
2030 config: Validated<McpServerConfig>,
2031 handler_factory: F,
2032 ready_tx: Option<tokio::sync::oneshot::Sender<SocketAddr>>,
2033 shutdown: Option<CancellationToken>,
2034) -> Result<(), McpxError>
2035where
2036 H: ServerHandler + 'static,
2037 F: Fn() -> H + Send + Sync + Clone + 'static,
2038{
2039 let config = config.into_inner();
2040 let local_addr = listener
2041 .local_addr()
2042 .map_err(|e| io_to_startup("listener.local_addr", e))?;
2043 let (router, params) = build_app_router(config, handler_factory).map_err(anyhow_to_startup)?;
2044
2045 log_listening(¶ms.name, params.scheme, &local_addr.to_string());
2046
2047 if let Some(external) = shutdown {
2051 let internal = params.ct.clone();
2052 tokio::spawn(async move {
2053 external.cancelled().await;
2054 internal.cancel();
2055 });
2056 }
2057
2058 if let Some(tx) = ready_tx {
2062 let _ = tx.send(local_addr);
2064 }
2065
2066 run_server(
2067 router,
2068 listener,
2069 params.tls_paths,
2070 params.tls_handshake_timeout,
2071 params.max_concurrent_tls_handshakes,
2072 params.mtls_config,
2073 params.shutdown_timeout,
2074 params.auth_state,
2075 params.rbac_swap,
2076 params.on_reload_ready,
2077 params.ct,
2078 )
2079 .await
2080 .map_err(anyhow_to_startup)
2081}
2082
2083#[allow(
2086 clippy::cognitive_complexity,
2087 reason = "tracing::info! macro expansions inflate the score; logic is trivial"
2088)]
2089fn log_listening(name: &str, scheme: &str, addr: &str) {
2090 tracing::info!("{name} listening on {addr}");
2091 tracing::info!(" MCP endpoint: {scheme}://{addr}/mcp");
2092 tracing::info!(" Health check: {scheme}://{addr}/healthz");
2093 tracing::info!(" Readiness: {scheme}://{addr}/readyz");
2094}
2095
2096#[allow(
2119 clippy::too_many_arguments,
2120 clippy::cognitive_complexity,
2121 reason = "server start-up threads TLS, reload state, and graceful shutdown through one flow"
2122)]
2123async fn run_server(
2124 router: axum::Router,
2125 listener: TcpListener,
2126 tls_paths: Option<(PathBuf, PathBuf)>,
2127 tls_handshake_timeout: Duration,
2128 max_concurrent_tls_handshakes: usize,
2129 mtls_config: Option<MtlsConfig>,
2130 shutdown_timeout: Duration,
2131 auth_state: Option<Arc<AuthState>>,
2132 rbac_swap: Arc<ArcSwap<RbacPolicy>>,
2133 mut on_reload_ready: Option<Box<dyn FnOnce(ReloadHandle) + Send>>,
2134 ct: CancellationToken,
2135) -> anyhow::Result<()> {
2136 let shutdown_trigger = CancellationToken::new();
2140 {
2141 let trigger = shutdown_trigger.clone();
2142 let parent = ct.clone();
2143 tokio::spawn(async move {
2144 tokio::select! {
2147 () = shutdown_signal() => {}
2148 () = parent.cancelled() => {}
2149 }
2150 trigger.cancel();
2151 });
2152 }
2153
2154 let graceful = {
2155 let trigger = shutdown_trigger.clone();
2156 let ct = ct.clone();
2157 async move {
2158 trigger.cancelled().await;
2159 tracing::info!("shutting down (grace period: {shutdown_timeout:?})");
2160 ct.cancel();
2161 }
2162 };
2163
2164 let force_exit_timer = {
2165 let trigger = shutdown_trigger.clone();
2166 async move {
2167 trigger.cancelled().await;
2168 tokio::time::sleep(shutdown_timeout).await;
2169 }
2170 };
2171
2172 if let Some((cert_path, key_path)) = tls_paths {
2173 let crl_set = if let Some(mtls) = mtls_config.as_ref()
2174 && mtls.crl_enabled
2175 {
2176 let (ca_certs, roots) = load_client_auth_roots(&mtls.ca_cert_path)?;
2177 let (crl_set, discover_rx) =
2178 mtls_revocation::bootstrap_fetch(roots, &ca_certs, mtls.clone())
2179 .await
2180 .map_err(|error| anyhow::anyhow!(error.to_string()))?;
2181 tokio::spawn(mtls_revocation::run_crl_refresher(
2182 Arc::clone(&crl_set),
2183 discover_rx,
2184 ct.clone(),
2185 ));
2186 Some(crl_set)
2187 } else {
2188 None
2189 };
2190
2191 if let Some(cb) = on_reload_ready.take() {
2192 cb(ReloadHandle {
2193 auth: auth_state.clone(),
2194 rbac: Some(Arc::clone(&rbac_swap)),
2195 crl_set: crl_set.clone(),
2196 });
2197 }
2198
2199 let tls_listener = TlsListener::new(
2200 listener,
2201 &cert_path,
2202 &key_path,
2203 mtls_config.as_ref(),
2204 crl_set,
2205 tls_handshake_timeout,
2206 max_concurrent_tls_handshakes,
2207 )?;
2208 let make_svc = router.into_make_service_with_connect_info::<TlsConnInfo>();
2209 tokio::select! {
2212 result = axum::serve(tls_listener, make_svc)
2213 .with_graceful_shutdown(graceful) => { result?; }
2214 () = force_exit_timer => {
2215 tracing::warn!("shutdown timeout exceeded, forcing exit");
2216 }
2217 }
2218 } else {
2219 if let Some(cb) = on_reload_ready.take() {
2220 cb(ReloadHandle {
2221 auth: auth_state,
2222 rbac: Some(rbac_swap),
2223 crl_set: None,
2224 });
2225 }
2226
2227 let make_svc = router.into_make_service_with_connect_info::<SocketAddr>();
2228 tokio::select! {
2231 result = axum::serve(listener, make_svc)
2232 .with_graceful_shutdown(graceful) => { result?; }
2233 () = force_exit_timer => {
2234 tracing::warn!("shutdown timeout exceeded, forcing exit");
2235 }
2236 }
2237 }
2238
2239 Ok(())
2240}
2241
2242#[cfg(feature = "oauth")]
2251fn install_oauth_proxy_routes(
2252 router: axum::Router,
2253 server_url: &str,
2254 oauth_config: &crate::oauth::OAuthConfig,
2255 auth_state: Option<&Arc<AuthState>>,
2256 max_request_body: usize,
2257) -> Result<axum::Router, McpxError> {
2258 let Some(ref proxy) = oauth_config.proxy else {
2259 return Ok(router);
2260 };
2261
2262 let http = crate::oauth::OauthHttpClient::with_config(oauth_config)?;
2265
2266 let proxy_router = axum::Router::new();
2272
2273 let asm = crate::oauth::authorization_server_metadata(server_url, oauth_config);
2274 let proxy_router = proxy_router.route(
2275 "/.well-known/oauth-authorization-server",
2276 axum::routing::get(move || {
2277 let m = asm.clone();
2278 async move { axum::Json(m) }
2279 }),
2280 );
2281
2282 let proxy_authorize = proxy.clone();
2283 let proxy_router = proxy_router.route(
2284 "/authorize",
2285 axum::routing::get(
2286 move |axum::extract::RawQuery(query): axum::extract::RawQuery| {
2287 let p = proxy_authorize.clone();
2288 async move { crate::oauth::handle_authorize(&p, &query.unwrap_or_default()) }
2289 },
2290 ),
2291 );
2292
2293 let proxy_token = proxy.clone();
2294 let token_http = http.clone();
2295 let proxy_router = proxy_router.route(
2296 "/token",
2297 axum::routing::post(move |body: String| {
2298 let p = proxy_token.clone();
2299 let h = token_http.clone();
2300 async move { crate::oauth::handle_token(&h, &p, &body).await }
2301 })
2302 .layer(axum::middleware::from_fn(
2303 oauth_token_cache_headers_middleware,
2304 )),
2305 );
2306
2307 let proxy_register = proxy.clone();
2308 let proxy_router = proxy_router.route(
2309 "/register",
2310 axum::routing::post(move |axum::Json(body): axum::Json<serde_json::Value>| {
2311 let p = proxy_register;
2312 async move { axum::Json(crate::oauth::handle_register(&p, &body)) }
2313 })
2314 .layer(axum::middleware::from_fn(
2315 oauth_token_cache_headers_middleware,
2316 )),
2317 );
2318
2319 let admin_routes_enabled = proxy.expose_admin_endpoints
2320 && (proxy.introspection_url.is_some() || proxy.revocation_url.is_some());
2321 if proxy.expose_admin_endpoints
2322 && !proxy.require_auth_on_admin_endpoints
2323 && proxy.allow_unauthenticated_admin_endpoints
2324 {
2325 tracing::warn!(
2329 "OAuth introspect/revoke endpoints are unauthenticated by explicit \
2330 allow_unauthenticated_admin_endpoints opt-out; ensure an \
2331 authenticated reverse proxy fronts these routes"
2332 );
2333 }
2334
2335 let admin_router = if admin_routes_enabled {
2336 build_oauth_admin_router(proxy, http, auth_state)?
2337 } else {
2338 axum::Router::new()
2339 };
2340
2341 let proxy_router =
2345 proxy_router
2346 .merge(admin_router)
2347 .layer(tower_http::limit::RequestBodyLimitLayer::new(
2348 max_request_body,
2349 ));
2350
2351 let router = router.merge(proxy_router);
2352
2353 tracing::info!(
2354 introspect = proxy.expose_admin_endpoints && proxy.introspection_url.is_some(),
2355 revoke = proxy.expose_admin_endpoints && proxy.revocation_url.is_some(),
2356 max_request_body,
2357 "OAuth 2.1 proxy endpoints enabled (/authorize, /token, /register)"
2358 );
2359 Ok(router)
2360}
2361
2362#[cfg(feature = "oauth")]
2368fn build_oauth_admin_router(
2369 proxy: &crate::oauth::OAuthProxyConfig,
2370 http: crate::oauth::OauthHttpClient,
2371 auth_state: Option<&Arc<AuthState>>,
2372) -> Result<axum::Router, McpxError> {
2373 let mut admin_router = axum::Router::new();
2374 if proxy.introspection_url.is_some() {
2375 let proxy_introspect = proxy.clone();
2376 let introspect_http = http.clone();
2377 admin_router = admin_router.route(
2378 "/introspect",
2379 axum::routing::post(move |body: String| {
2380 let p = proxy_introspect.clone();
2381 let h = introspect_http.clone();
2382 async move { crate::oauth::handle_introspect(&h, &p, &body).await }
2383 }),
2384 );
2385 }
2386 if proxy.revocation_url.is_some() {
2387 let proxy_revoke = proxy.clone();
2388 let revoke_http = http;
2389 admin_router = admin_router.route(
2390 "/revoke",
2391 axum::routing::post(move |body: String| {
2392 let p = proxy_revoke.clone();
2393 let h = revoke_http.clone();
2394 async move { crate::oauth::handle_revoke(&h, &p, &body).await }
2395 }),
2396 );
2397 }
2398
2399 let admin_router = admin_router.layer(axum::middleware::from_fn(
2400 oauth_token_cache_headers_middleware,
2401 ));
2402
2403 if proxy.require_auth_on_admin_endpoints {
2404 let Some(state) = auth_state else {
2405 return Err(McpxError::Startup(
2406 "oauth proxy admin endpoints require auth state".into(),
2407 ));
2408 };
2409 let state_for_mw = Arc::clone(state);
2410 Ok(
2411 admin_router.layer(axum::middleware::from_fn(move |req, next| {
2412 let s = Arc::clone(&state_for_mw);
2413 auth_middleware(s, req, next)
2414 })),
2415 )
2416 } else {
2417 Ok(admin_router)
2418 }
2419}
2420
2421fn derive_allowed_hosts(bind_addr: &str, public_url: Option<&str>) -> Vec<String> {
2426 let mut hosts = vec![
2427 "localhost".to_owned(),
2428 "127.0.0.1".to_owned(),
2429 "::1".to_owned(),
2430 ];
2431
2432 if let Some(url) = public_url
2433 && let Ok(uri) = url.parse::<axum::http::Uri>()
2434 && let Some(authority) = uri.authority()
2435 {
2436 let host = authority.host().to_owned();
2437 if !hosts.iter().any(|h| h == &host) {
2438 hosts.push(host);
2439 }
2440
2441 let authority = authority.as_str().to_owned();
2442 if !hosts.iter().any(|h| h == &authority) {
2443 hosts.push(authority);
2444 }
2445 }
2446
2447 if let Ok(uri) = format!("http://{bind_addr}").parse::<axum::http::Uri>()
2448 && let Some(authority) = uri.authority()
2449 {
2450 let host = authority.host().to_owned();
2451 if !hosts.iter().any(|h| h == &host) {
2452 hosts.push(host);
2453 }
2454
2455 let authority = authority.as_str().to_owned();
2456 if !hosts.iter().any(|h| h == &authority) {
2457 hosts.push(authority);
2458 }
2459 }
2460
2461 hosts
2462}
2463
2464impl axum::extract::connect_info::Connected<axum::serve::IncomingStream<'_, TlsListener>>
2477 for TlsConnInfo
2478{
2479 fn connect_info(target: axum::serve::IncomingStream<'_, TlsListener>) -> Self {
2480 let addr = *target.remote_addr();
2481 let identity = target.io().identity().cloned();
2482 Self::new(addr, identity)
2483 }
2484}
2485
2486const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
2493
2494const DEFAULT_MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256;
2502
2503const TLS_ACCEPT_CHANNEL_CAPACITY: usize = 32;
2508
2509struct TlsListener {
2525 local_addr: SocketAddr,
2528 rx: mpsc::Receiver<(AuthenticatedTlsStream, SocketAddr)>,
2530 acceptor_task: tokio::task::JoinHandle<()>,
2533}
2534
2535impl TlsListener {
2536 fn new(
2537 inner: TcpListener,
2538 cert_path: &Path,
2539 key_path: &Path,
2540 mtls_config: Option<&MtlsConfig>,
2541 crl_set: Option<Arc<CrlSet>>,
2542 handshake_timeout: Duration,
2543 max_concurrent_handshakes: usize,
2544 ) -> anyhow::Result<Self> {
2545 rustls::crypto::ring::default_provider()
2547 .install_default()
2548 .ok();
2549
2550 let certs = load_certs(cert_path)?;
2551 let key = load_key(key_path)?;
2552
2553 let mtls_default_role;
2554
2555 let tls_config = if let Some(mtls) = mtls_config {
2556 mtls_default_role = mtls.default_role.clone();
2557 let verifier: Arc<dyn rustls::server::danger::ClientCertVerifier> = if mtls.crl_enabled
2558 {
2559 let Some(crl_set) = crl_set else {
2560 return Err(anyhow::anyhow!(
2561 "mTLS CRL verifier requested but CRL state was not initialized"
2562 ));
2563 };
2564 Arc::new(DynamicClientCertVerifier::new(crl_set))
2565 } else {
2566 let (_, root_store) = load_client_auth_roots(&mtls.ca_cert_path)?;
2567 if mtls.required {
2568 rustls::server::WebPkiClientVerifier::builder(root_store)
2569 .build()
2570 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2571 } else {
2572 rustls::server::WebPkiClientVerifier::builder(root_store)
2573 .allow_unauthenticated()
2574 .build()
2575 .map_err(|e| anyhow::anyhow!("mTLS verifier error: {e}"))?
2576 }
2577 };
2578
2579 tracing::info!(
2580 ca = %mtls.ca_cert_path.display(),
2581 required = mtls.required,
2582 crl_enabled = mtls.crl_enabled,
2583 "mTLS client auth configured"
2584 );
2585
2586 rustls::ServerConfig::builder_with_protocol_versions(&[
2587 &rustls::version::TLS12,
2588 &rustls::version::TLS13,
2589 ])
2590 .with_client_cert_verifier(verifier)
2591 .with_single_cert(certs, key)?
2592 } else {
2593 mtls_default_role = "viewer".to_owned();
2594 rustls::ServerConfig::builder_with_protocol_versions(&[
2595 &rustls::version::TLS12,
2596 &rustls::version::TLS13,
2597 ])
2598 .with_no_client_auth()
2599 .with_single_cert(certs, key)?
2600 };
2601
2602 let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
2603 tracing::info!(
2604 "TLS enabled (cert: {}, key: {})",
2605 cert_path.display(),
2606 key_path.display()
2607 );
2608 let local_addr = inner.local_addr()?;
2609 let (tx, rx) = mpsc::channel(TLS_ACCEPT_CHANNEL_CAPACITY);
2610 let acceptor_task = tokio::spawn(run_tls_acceptor(
2611 inner,
2612 acceptor,
2613 mtls_default_role,
2614 tx,
2615 handshake_timeout,
2616 max_concurrent_handshakes,
2617 ));
2618 Ok(Self {
2619 local_addr,
2620 rx,
2621 acceptor_task,
2622 })
2623 }
2624
2625 fn extract_handshake_identity(
2629 tls_stream: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2630 default_role: &str,
2631 addr: SocketAddr,
2632 ) -> Option<AuthIdentity> {
2633 let (_, server_conn) = tls_stream.get_ref();
2634 let cert_der = server_conn.peer_certificates()?.first()?;
2635 let id = extract_mtls_identity(cert_der.as_ref(), default_role)?;
2636 tracing::debug!(name = %id.name, peer = %addr, "mTLS client cert accepted");
2637 Some(id)
2638 }
2639}
2640
2641async fn run_tls_acceptor(
2649 listener: TcpListener,
2650 acceptor: tokio_rustls::TlsAcceptor,
2651 default_role: String,
2652 tx: mpsc::Sender<(AuthenticatedTlsStream, SocketAddr)>,
2653 handshake_timeout: Duration,
2654 max_concurrent_handshakes: usize,
2655) {
2656 let inflight = Arc::new(Semaphore::new(max_concurrent_handshakes));
2657 loop {
2658 let Ok(permit) = Arc::clone(&inflight).acquire_owned().await else {
2662 return;
2664 };
2665 let (stream, addr) = match listener.accept().await {
2666 Ok(pair) => pair,
2667 Err(e) => {
2668 tracing::debug!("TCP accept error: {e}");
2669 continue;
2670 }
2671 };
2672 if tx.is_closed() {
2673 return;
2675 }
2676 let acceptor = acceptor.clone();
2677 let default_role = default_role.clone();
2678 let tx = tx.clone();
2679 tokio::spawn(async move {
2680 let _permit = permit;
2681 match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
2682 Ok(Ok(tls_stream)) => {
2683 let identity =
2684 TlsListener::extract_handshake_identity(&tls_stream, &default_role, addr);
2685 let wrapped = AuthenticatedTlsStream {
2686 inner: tls_stream,
2687 identity,
2688 };
2689 let _ = tx.send((wrapped, addr)).await;
2692 }
2693 Ok(Err(e)) => {
2694 tracing::debug!("TLS handshake failed from {addr}: {e}");
2695 }
2696 Err(_elapsed) => {
2697 tracing::debug!(
2698 "TLS handshake timed out from {addr} after {handshake_timeout:?}"
2699 );
2700 }
2701 }
2702 });
2703 }
2704}
2705
2706pub(crate) struct AuthenticatedTlsStream {
2718 inner: tokio_rustls::server::TlsStream<tokio::net::TcpStream>,
2719 identity: Option<AuthIdentity>,
2720}
2721
2722impl AuthenticatedTlsStream {
2723 #[must_use]
2725 pub(crate) const fn identity(&self) -> Option<&AuthIdentity> {
2726 self.identity.as_ref()
2727 }
2728}
2729
2730impl std::fmt::Debug for AuthenticatedTlsStream {
2731 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2732 f.debug_struct("AuthenticatedTlsStream")
2733 .field("identity", &self.identity.as_ref().map(|id| &id.name))
2734 .finish_non_exhaustive()
2735 }
2736}
2737
2738impl tokio::io::AsyncRead for AuthenticatedTlsStream {
2739 fn poll_read(
2740 mut self: Pin<&mut Self>,
2741 cx: &mut std::task::Context<'_>,
2742 buf: &mut tokio::io::ReadBuf<'_>,
2743 ) -> std::task::Poll<std::io::Result<()>> {
2744 Pin::new(&mut self.inner).poll_read(cx, buf)
2745 }
2746}
2747
2748impl tokio::io::AsyncWrite for AuthenticatedTlsStream {
2749 fn poll_write(
2750 mut self: Pin<&mut Self>,
2751 cx: &mut std::task::Context<'_>,
2752 buf: &[u8],
2753 ) -> std::task::Poll<std::io::Result<usize>> {
2754 Pin::new(&mut self.inner).poll_write(cx, buf)
2755 }
2756
2757 fn poll_flush(
2758 mut self: Pin<&mut Self>,
2759 cx: &mut std::task::Context<'_>,
2760 ) -> std::task::Poll<std::io::Result<()>> {
2761 Pin::new(&mut self.inner).poll_flush(cx)
2762 }
2763
2764 fn poll_shutdown(
2765 mut self: Pin<&mut Self>,
2766 cx: &mut std::task::Context<'_>,
2767 ) -> std::task::Poll<std::io::Result<()>> {
2768 Pin::new(&mut self.inner).poll_shutdown(cx)
2769 }
2770
2771 fn poll_write_vectored(
2772 mut self: Pin<&mut Self>,
2773 cx: &mut std::task::Context<'_>,
2774 bufs: &[std::io::IoSlice<'_>],
2775 ) -> std::task::Poll<std::io::Result<usize>> {
2776 Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
2777 }
2778
2779 fn is_write_vectored(&self) -> bool {
2780 self.inner.is_write_vectored()
2781 }
2782}
2783
2784impl axum::serve::Listener for TlsListener {
2785 type Io = AuthenticatedTlsStream;
2786 type Addr = SocketAddr;
2787
2788 async fn accept(&mut self) -> (Self::Io, Self::Addr) {
2794 if let Some(pair) = self.rx.recv().await {
2795 return pair;
2796 }
2797 tracing::error!("TLS acceptor task terminated; no further connections will be accepted");
2803 std::future::pending().await
2804 }
2805
2806 fn local_addr(&self) -> std::io::Result<Self::Addr> {
2807 Ok(self.local_addr)
2808 }
2809}
2810
2811impl Drop for TlsListener {
2812 fn drop(&mut self) {
2813 self.acceptor_task.abort();
2816 }
2817}
2818
2819fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
2820 use rustls::pki_types::pem::PemObject;
2821 let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_file_iter(path)
2822 .map_err(|e| anyhow::anyhow!("failed to read certs from {}: {e}", path.display()))?
2823 .collect::<Result<_, _>>()
2824 .map_err(|e| anyhow::anyhow!("invalid cert in {}: {e}", path.display()))?;
2825 anyhow::ensure!(
2826 !certs.is_empty(),
2827 "no certificates found in {}",
2828 path.display()
2829 );
2830 Ok(certs)
2831}
2832
2833fn load_client_auth_roots(
2834 path: &Path,
2835) -> anyhow::Result<(
2836 Vec<rustls::pki_types::CertificateDer<'static>>,
2837 Arc<RootCertStore>,
2838)> {
2839 let ca_certs = load_certs(path)?;
2840 let mut root_store = RootCertStore::empty();
2841 for cert in &ca_certs {
2842 root_store
2843 .add(cert.clone())
2844 .map_err(|error| anyhow::anyhow!("invalid CA cert: {error}"))?;
2845 }
2846
2847 Ok((ca_certs, Arc::new(root_store)))
2848}
2849
2850fn load_key(path: &Path) -> anyhow::Result<rustls::pki_types::PrivateKeyDer<'static>> {
2851 use rustls::pki_types::pem::PemObject;
2852 rustls::pki_types::PrivateKeyDer::from_pem_file(path)
2853 .map_err(|e| anyhow::anyhow!("failed to read key from {}: {e}", path.display()))
2854}
2855
2856#[allow(
2857 clippy::unused_async,
2858 reason = "axum route handler signature requires `async fn` even when the body is synchronous"
2859)]
2860async fn healthz() -> impl IntoResponse {
2861 axum::Json(serde_json::json!({
2862 "status": "ok",
2863 }))
2864}
2865
2866fn version_payload(name: &str, version: &str) -> serde_json::Value {
2873 serde_json::json!({
2874 "name": name,
2875 "version": version,
2876 "build_git_sha": option_env!("RMCP_SERVER_KIT_BUILD_SHA").unwrap_or("unknown"),
2877 "build_timestamp": option_env!("RMCP_SERVER_KIT_BUILD_TIME").unwrap_or("unknown"),
2878 "rust_version": option_env!("RMCP_SERVER_KIT_RUSTC_VERSION").unwrap_or("unknown"),
2879 "mcpx_version": env!("CARGO_PKG_VERSION"),
2880 })
2881}
2882
2883fn serialize_version_payload(name: &str, version: &str) -> Arc<[u8]> {
2893 let value = version_payload(name, version);
2894 serde_json::to_vec(&value).map_or_else(|_| Arc::from(&b"{}"[..]), Arc::from)
2895}
2896
2897async fn readyz(check: ReadinessCheck) -> impl IntoResponse {
2898 let status = check().await;
2899 let ready = status
2900 .get("ready")
2901 .and_then(serde_json::Value::as_bool)
2902 .unwrap_or(false);
2903 let code = if ready {
2904 axum::http::StatusCode::OK
2905 } else {
2906 axum::http::StatusCode::SERVICE_UNAVAILABLE
2907 };
2908 (code, axum::Json(status))
2909}
2910
2911async fn shutdown_signal() {
2915 let ctrl_c = tokio::signal::ctrl_c();
2916
2917 #[cfg(unix)]
2918 {
2919 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
2920 Ok(mut term) => {
2921 tokio::select! {
2924 _ = ctrl_c => {}
2925 _ = term.recv() => {}
2926 }
2927 }
2928 Err(e) => {
2929 tracing::warn!(error = %e, "failed to register SIGTERM handler, using SIGINT only");
2930 ctrl_c.await.ok();
2931 }
2932 }
2933 }
2934
2935 #[cfg(not(unix))]
2936 {
2937 ctrl_c.await.ok();
2938 }
2939}
2940
2941#[cfg(feature = "metrics")]
2952async fn metrics_middleware(
2953 metrics: Arc<crate::metrics::McpMetrics>,
2954 mut req: Request<Body>,
2955 next: Next,
2956) -> axum::response::Response {
2957 let method = req.method().to_string();
2958 let path = req.uri().path().to_owned();
2959 let start = std::time::Instant::now();
2960
2961 req.extensions_mut().insert(Arc::clone(&metrics));
2962 let response = next.run(req).await;
2963
2964 let status = response.status().as_u16().to_string();
2965 let duration = start.elapsed().as_secs_f64();
2966
2967 metrics
2968 .http_requests_total
2969 .with_label_values(&[&method, &path, &status])
2970 .inc();
2971 metrics
2972 .http_request_duration_seconds
2973 .with_label_values(&[&method, &path])
2974 .observe(duration);
2975
2976 response
2977}
2978
2979async fn security_headers_middleware(
2991 is_tls: bool,
2992 cfg: Arc<SecurityHeadersConfig>,
2993 req: Request<Body>,
2994 next: Next,
2995) -> axum::response::Response {
2996 use axum::http::{HeaderName, header};
2997
2998 let mut resp = next.run(req).await;
2999 let headers = resp.headers_mut();
3000
3001 headers.remove(header::SERVER);
3003 headers.remove(HeaderName::from_static("x-powered-by"));
3004
3005 apply_security_header(
3006 headers,
3007 header::X_CONTENT_TYPE_OPTIONS,
3008 cfg.x_content_type_options.as_deref(),
3009 "nosniff",
3010 );
3011 apply_security_header(
3012 headers,
3013 header::X_FRAME_OPTIONS,
3014 cfg.x_frame_options.as_deref(),
3015 "deny",
3016 );
3017 apply_security_header(
3018 headers,
3019 header::CACHE_CONTROL,
3020 cfg.cache_control.as_deref(),
3021 "no-store, max-age=0",
3022 );
3023 apply_security_header(
3024 headers,
3025 header::REFERRER_POLICY,
3026 cfg.referrer_policy.as_deref(),
3027 "no-referrer",
3028 );
3029 apply_security_header(
3030 headers,
3031 HeaderName::from_static("cross-origin-opener-policy"),
3032 cfg.cross_origin_opener_policy.as_deref(),
3033 "same-origin",
3034 );
3035 apply_security_header(
3036 headers,
3037 HeaderName::from_static("cross-origin-resource-policy"),
3038 cfg.cross_origin_resource_policy.as_deref(),
3039 "same-origin",
3040 );
3041 apply_security_header(
3042 headers,
3043 HeaderName::from_static("cross-origin-embedder-policy"),
3044 cfg.cross_origin_embedder_policy.as_deref(),
3045 "require-corp",
3046 );
3047 apply_security_header(
3048 headers,
3049 HeaderName::from_static("permissions-policy"),
3050 cfg.permissions_policy.as_deref(),
3051 "accelerometer=(), camera=(), geolocation=(), microphone=()",
3052 );
3053 apply_security_header(
3054 headers,
3055 HeaderName::from_static("x-permitted-cross-domain-policies"),
3056 cfg.x_permitted_cross_domain_policies.as_deref(),
3057 "none",
3058 );
3059 apply_security_header(
3060 headers,
3061 HeaderName::from_static("content-security-policy"),
3062 cfg.content_security_policy.as_deref(),
3063 "default-src 'none'; frame-ancestors 'none'",
3064 );
3065 apply_security_header(
3066 headers,
3067 HeaderName::from_static("x-dns-prefetch-control"),
3068 cfg.x_dns_prefetch_control.as_deref(),
3069 "off",
3070 );
3071
3072 if is_tls {
3073 apply_security_header(
3074 headers,
3075 header::STRICT_TRANSPORT_SECURITY,
3076 cfg.strict_transport_security.as_deref(),
3077 "max-age=63072000; includeSubDomains",
3078 );
3079 }
3080
3081 resp
3082}
3083
3084fn apply_security_header(
3095 headers: &mut axum::http::HeaderMap,
3096 name: axum::http::HeaderName,
3097 override_value: Option<&str>,
3098 default: &'static str,
3099) {
3100 use axum::http::HeaderValue;
3101
3102 match override_value {
3103 None => {
3104 headers.insert(name, HeaderValue::from_static(default));
3105 }
3106 Some("") => {
3107 }
3109 Some(v) => match HeaderValue::from_str(v) {
3110 Ok(hv) => {
3111 headers.insert(name, hv);
3112 }
3113 Err(err) => {
3114 tracing::error!(
3115 header = %name,
3116 error = %err,
3117 "invalid security header override reached middleware; using default"
3118 );
3119 headers.insert(name, HeaderValue::from_static(default));
3120 }
3121 },
3122 }
3123}
3124
3125fn validate_security_headers(cfg: &SecurityHeadersConfig) -> Result<(), McpxError> {
3136 use axum::http::HeaderValue;
3137
3138 let fields: &[(&str, Option<&str>)] = &[
3139 (
3140 "x_content_type_options",
3141 cfg.x_content_type_options.as_deref(),
3142 ),
3143 ("x_frame_options", cfg.x_frame_options.as_deref()),
3144 ("cache_control", cfg.cache_control.as_deref()),
3145 ("referrer_policy", cfg.referrer_policy.as_deref()),
3146 (
3147 "cross_origin_opener_policy",
3148 cfg.cross_origin_opener_policy.as_deref(),
3149 ),
3150 (
3151 "cross_origin_resource_policy",
3152 cfg.cross_origin_resource_policy.as_deref(),
3153 ),
3154 (
3155 "cross_origin_embedder_policy",
3156 cfg.cross_origin_embedder_policy.as_deref(),
3157 ),
3158 ("permissions_policy", cfg.permissions_policy.as_deref()),
3159 (
3160 "x_permitted_cross_domain_policies",
3161 cfg.x_permitted_cross_domain_policies.as_deref(),
3162 ),
3163 (
3164 "content_security_policy",
3165 cfg.content_security_policy.as_deref(),
3166 ),
3167 (
3168 "x_dns_prefetch_control",
3169 cfg.x_dns_prefetch_control.as_deref(),
3170 ),
3171 (
3172 "strict_transport_security",
3173 cfg.strict_transport_security.as_deref(),
3174 ),
3175 ];
3176
3177 for (field, value) in fields {
3178 let Some(v) = value else { continue };
3179 if v.is_empty() {
3180 continue;
3181 }
3182 if let Err(err) = HeaderValue::from_str(v) {
3183 return Err(McpxError::Config(format!(
3184 "invalid security_headers.{field}: {err}"
3185 )));
3186 }
3187 }
3188
3189 if let Some(v) = cfg.strict_transport_security.as_deref()
3190 && !v.is_empty()
3191 && v.to_ascii_lowercase().contains("preload")
3192 {
3193 return Err(McpxError::Config(format!(
3194 "invalid security_headers.strict_transport_security: {v:?} contains the `preload` directive; \
3195 HSTS preload must be opted into explicitly via a dedicated builder, not via this knob"
3196 )));
3197 }
3198
3199 Ok(())
3200}
3201
3202#[cfg(feature = "oauth")]
3217async fn oauth_token_cache_headers_middleware(
3218 req: Request<Body>,
3219 next: Next,
3220) -> axum::response::Response {
3221 use axum::http::{HeaderValue, header};
3222
3223 let mut resp = next.run(req).await;
3224 let headers = resp.headers_mut();
3225 headers.insert(header::PRAGMA, HeaderValue::from_static("no-cache"));
3226 headers.append(header::VARY, HeaderValue::from_static("Authorization"));
3227 resp
3228}
3229
3230async fn normalize_peer_addr_middleware(
3259 resolver: Option<Arc<ForwardResolver>>,
3260 mut req: Request<Body>,
3261 next: Next,
3262) -> axum::response::Response {
3263 let direct = req
3264 .extensions()
3265 .get::<ConnectInfo<SocketAddr>>()
3266 .map(|ci| ci.0);
3267 let from_tls = req
3268 .extensions()
3269 .get::<ConnectInfo<TlsConnInfo>>()
3270 .map(|ci| ci.0.addr);
3271 if let Some(addr) = direct.or(from_tls) {
3272 if direct.is_none() {
3273 req.extensions_mut().insert(ConnectInfo(addr));
3274 }
3275 req.extensions_mut().insert(PeerAddr::new(addr));
3276 let client_ip = match &resolver {
3277 Some(r) => {
3278 crate::forwarded::resolve_client_ip(addr.ip(), req.headers(), &r.trusted, r.mode)
3279 .unwrap_or_else(|reason| {
3280 tracing::debug!(
3281 reason = ?reason,
3282 "forwarded-header resolution fell back to direct peer"
3283 );
3284 addr.ip()
3285 })
3286 }
3287 None => addr.ip(),
3288 };
3289 req.extensions_mut().insert(ClientIp::new(client_ip));
3290 }
3291 next.run(req).await
3292}
3293
3294fn parse_proxy_net(entry: &str) -> Option<ipnet::IpNet> {
3297 if let Ok(net) = entry.parse::<ipnet::IpNet>() {
3298 return Some(net);
3299 }
3300 entry.parse::<IpAddr>().ok().map(ipnet::IpNet::from)
3301}
3302
3303pub(crate) fn limiter_client_ip(extensions: &axum::http::Extensions) -> Option<IpAddr> {
3307 if let Some(client) = extensions.get::<ClientIp>() {
3308 return Some(client.ip);
3309 }
3310 extensions
3311 .get::<ConnectInfo<SocketAddr>>()
3312 .map(|ci| ci.0.ip())
3313 .or_else(|| {
3314 extensions
3315 .get::<ConnectInfo<TlsConnInfo>>()
3316 .map(|ci| ci.0.addr.ip())
3317 })
3318}
3319
3320pub(crate) type ExtraRouteRateLimiter = BoundedKeyedLimiter<IpAddr>;
3324
3325const EXTRA_ROUTE_MAX_TRACKED_KEYS: usize = 10_000;
3331
3332const EXTRA_ROUTE_IDLE_EVICTION: Duration = Duration::from_mins(15);
3335
3336fn build_extra_route_rate_limiter(
3343 per_minute: u32,
3344 burst: Option<u32>,
3345) -> Arc<ExtraRouteRateLimiter> {
3346 let rate = std::num::NonZeroU32::new(per_minute.max(1)).unwrap_or(std::num::NonZeroU32::MIN);
3347 let mut quota = governor::Quota::per_minute(rate);
3348 if let Some(b) = burst.and_then(std::num::NonZeroU32::new) {
3349 quota = quota.allow_burst(b);
3350 }
3351 Arc::new(BoundedKeyedLimiter::new(
3352 quota,
3353 EXTRA_ROUTE_MAX_TRACKED_KEYS,
3354 EXTRA_ROUTE_IDLE_EVICTION,
3355 ))
3356}
3357
3358async fn extra_route_rate_limit_middleware(
3380 limiter: Arc<ExtraRouteRateLimiter>,
3381 exempt: Arc<std::collections::HashSet<String>>,
3382 req: Request<Body>,
3383 next: Next,
3384) -> axum::response::Response {
3385 if exempt.contains(req.uri().path()) {
3386 return next.run(req).await;
3387 }
3388 let peer_ip: Option<IpAddr> = limiter_client_ip(req.extensions());
3389 if let Some(ip) = peer_ip
3390 && let Err(wait) = limiter.check_key_wait(&ip)
3391 {
3392 #[cfg(feature = "metrics")]
3393 crate::metrics::record_rate_limit_deny(req.extensions(), "extra_route");
3394 tracing::warn!(%ip, "extra route request rate limited");
3395 return McpxError::RateLimitedFor {
3396 message: "too many requests to application routes from this source".into(),
3397 retry_after: wait,
3398 }
3399 .into_response();
3400 }
3401 next.run(req).await
3402}
3403
3404async fn origin_check_middleware(
3408 allowed: Arc<[String]>,
3409 log_request_headers: bool,
3410 req: Request<Body>,
3411 next: Next,
3412) -> axum::response::Response {
3413 let method = req.method().clone();
3414 let path = req.uri().path().to_owned();
3415
3416 log_incoming_request(&method, &path, req.headers(), log_request_headers);
3417
3418 if let Some(origin) = req.headers().get(axum::http::header::ORIGIN) {
3419 let origin_str = origin.to_str().unwrap_or("");
3420 if !allowed.iter().any(|a| a == origin_str) {
3421 tracing::warn!(
3422 origin = origin_str,
3423 %method,
3424 %path,
3425 allowed = ?&*allowed,
3426 "rejected request: Origin not allowed"
3427 );
3428 return (
3429 axum::http::StatusCode::FORBIDDEN,
3430 "Forbidden: Origin not allowed",
3431 )
3432 .into_response();
3433 }
3434 }
3435 next.run(req).await
3436}
3437
3438fn log_incoming_request(
3441 method: &axum::http::Method,
3442 path: &str,
3443 headers: &axum::http::HeaderMap,
3444 log_request_headers: bool,
3445) {
3446 if log_request_headers {
3447 tracing::debug!(
3448 %method,
3449 %path,
3450 headers = %format_request_headers_for_log(headers),
3451 "incoming request"
3452 );
3453 } else {
3454 tracing::debug!(%method, %path, "incoming request");
3455 }
3456}
3457
3458fn format_request_headers_for_log(headers: &axum::http::HeaderMap) -> String {
3459 headers
3460 .iter()
3461 .map(|(k, v)| {
3462 let name = k.as_str();
3463 if name == "authorization" || name == "cookie" || name == "proxy-authorization" {
3464 format!("{name}: [REDACTED]")
3465 } else {
3466 format!("{name}: {}", v.to_str().unwrap_or("<non-utf8>"))
3467 }
3468 })
3469 .collect::<Vec<_>>()
3470 .join(", ")
3471}
3472
3473#[allow(
3497 clippy::cognitive_complexity,
3498 reason = "complexity is purely tracing macro expansion (info/warn + match arms); 18 lines of straight-line code, nothing meaningful to extract"
3499)]
3500pub async fn serve_stdio<H>(handler: H) -> Result<(), McpxError>
3501where
3502 H: ServerHandler + 'static,
3503{
3504 use rmcp::ServiceExt as _;
3505
3506 tracing::info!("stdio transport: serving on stdin/stdout");
3507 tracing::warn!("stdio mode: auth, RBAC, TLS, and Origin checks are DISABLED");
3508
3509 let transport = rmcp::transport::io::stdio();
3510
3511 let service = handler
3512 .serve(transport)
3513 .await
3514 .map_err(|e| McpxError::Startup(format!("stdio initialize failed: {e}")))?;
3515
3516 if let Err(e) = service.waiting().await {
3517 tracing::warn!(error = %e, "stdio session ended with error");
3518 }
3519 tracing::info!("stdio session ended");
3520 Ok(())
3521}
3522
3523#[cfg(test)]
3524mod tests {
3525 #![allow(
3526 clippy::unwrap_used,
3527 clippy::expect_used,
3528 clippy::panic,
3529 clippy::indexing_slicing,
3530 clippy::unwrap_in_result,
3531 clippy::print_stdout,
3532 clippy::print_stderr,
3533 deprecated,
3534 reason = "internal unit tests legitimately read/write the deprecated `pub` fields they were designed to verify"
3535 )]
3536 use std::{sync::Arc, time::Duration};
3537
3538 use axum::{
3539 body::Body,
3540 http::{Request, StatusCode, header},
3541 response::IntoResponse,
3542 };
3543 use http_body_util::BodyExt;
3544 use tower::ServiceExt as _;
3545
3546 use super::*;
3547
3548 #[test]
3551 fn server_config_new_defaults() {
3552 let cfg = McpServerConfig::new("0.0.0.0:8443", "test-server", "1.0.0");
3553 assert_eq!(cfg.bind_addr, "0.0.0.0:8443");
3554 assert_eq!(cfg.name, "test-server");
3555 assert_eq!(cfg.version, "1.0.0");
3556 assert!(cfg.tls_cert_path.is_none());
3557 assert!(cfg.tls_key_path.is_none());
3558 assert!(cfg.auth.is_none());
3559 assert!(cfg.rbac.is_none());
3560 assert!(cfg.allowed_origins.is_empty());
3561 assert!(cfg.tool_rate_limit.is_none());
3562 assert!(cfg.readiness_check.is_none());
3563 assert_eq!(cfg.max_request_body, 1024 * 1024);
3564 assert_eq!(cfg.request_timeout, Duration::from_mins(2));
3565 assert_eq!(cfg.shutdown_timeout, Duration::from_secs(30));
3566 assert!(!cfg.log_request_headers);
3567 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(10));
3568 assert_eq!(cfg.max_concurrent_tls_handshakes, 256);
3569 }
3570
3571 #[test]
3572 fn tls_handshake_builders_set_fields() {
3573 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3574 .with_tls_handshake_timeout(Duration::from_secs(3))
3575 .with_max_concurrent_tls_handshakes(64);
3576 assert_eq!(cfg.tls_handshake_timeout, Duration::from_secs(3));
3577 assert_eq!(cfg.max_concurrent_tls_handshakes, 64);
3578 }
3579
3580 #[test]
3581 fn validate_rejects_zero_tls_handshake_timeout() {
3582 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3583 .with_tls_handshake_timeout(Duration::ZERO);
3584 let err = cfg.validate().expect_err("zero handshake timeout");
3585 assert!(err.to_string().contains("tls_handshake_timeout"));
3586 }
3587
3588 #[test]
3589 fn validate_rejects_zero_max_concurrent_tls_handshakes() {
3590 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
3591 .with_max_concurrent_tls_handshakes(0);
3592 let err = cfg.validate().expect_err("zero handshake concurrency");
3593 assert!(err.to_string().contains("max_concurrent_tls_handshakes"));
3594 }
3595
3596 #[test]
3597 fn validate_consumes_and_proves() {
3598 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
3600 let validated = cfg.validate().expect("valid config");
3601 assert_eq!(validated.as_inner().name, "test-server");
3603 let raw = validated.into_inner();
3605 assert_eq!(raw.name, "test-server");
3606
3607 let mut bad = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0");
3609 bad.max_request_body = 0;
3610 assert!(bad.validate().is_err(), "zero body cap must fail validate");
3611 }
3612
3613 #[test]
3614 fn validate_rejects_zero_max_concurrent_requests() {
3615 let cfg =
3616 McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_max_concurrent_requests(0);
3617 let err = cfg.validate().expect_err("zero concurrency cap must fail");
3618 assert!(
3619 format!("{err}").contains("max_concurrent_requests"),
3620 "error should mention max_concurrent_requests, got: {err}"
3621 );
3622 }
3623
3624 #[test]
3625 fn validate_rejects_zero_max_tracked_keys() {
3626 let rl = crate::auth::RateLimitConfig {
3629 max_attempts_per_minute: 30,
3630 pre_auth_max_per_minute: None,
3631 max_tracked_keys: 0,
3632 idle_eviction: Duration::from_secs(15 * 60),
3633 burst: None,
3634 pre_auth_burst: None,
3635 };
3636 let auth_cfg = AuthConfig {
3637 enabled: true,
3638 api_keys: Vec::new(),
3639 mtls: None,
3640 rate_limit: Some(rl),
3641 #[cfg(feature = "oauth")]
3642 oauth: None,
3643 };
3644 let cfg = McpServerConfig::new("127.0.0.1:8080", "test", "1.0.0").with_auth(auth_cfg);
3645 let err = cfg.validate().expect_err("zero max_tracked_keys must fail");
3646 assert!(
3647 format!("{err}").contains("max_tracked_keys"),
3648 "error should mention max_tracked_keys, got: {err}"
3649 );
3650 }
3651
3652 #[test]
3653 fn derive_allowed_hosts_includes_public_host() {
3654 let hosts = derive_allowed_hosts("0.0.0.0:8080", Some("https://mcp.example.com/mcp"));
3655 assert!(
3656 hosts.iter().any(|h| h == "mcp.example.com"),
3657 "public_url host must be allowed"
3658 );
3659 }
3660
3661 #[test]
3662 fn derive_allowed_hosts_includes_bind_authority() {
3663 let hosts = derive_allowed_hosts("127.0.0.1:8080", None);
3664 assert!(
3665 hosts.iter().any(|h| h == "127.0.0.1"),
3666 "bind host must be allowed"
3667 );
3668 assert!(
3669 hosts.iter().any(|h| h == "127.0.0.1:8080"),
3670 "bind authority must be allowed"
3671 );
3672 }
3673
3674 #[tokio::test]
3677 async fn healthz_returns_ok_json() {
3678 let resp = healthz().await.into_response();
3679 assert_eq!(resp.status(), StatusCode::OK);
3680 let body = resp.into_body().collect().await.unwrap().to_bytes();
3681 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
3682 assert_eq!(json["status"], "ok");
3683 assert!(
3684 json.get("name").is_none(),
3685 "healthz must not expose server name"
3686 );
3687 assert!(
3688 json.get("version").is_none(),
3689 "healthz must not expose version"
3690 );
3691 }
3692
3693 #[tokio::test]
3696 async fn readyz_returns_ok_when_ready() {
3697 let check: ReadinessCheck =
3698 Arc::new(|| Box::pin(async { serde_json::json!({"ready": true, "db": "connected"}) }));
3699 let resp = readyz(check).await.into_response();
3700 assert_eq!(resp.status(), StatusCode::OK);
3701 let body = resp.into_body().collect().await.unwrap().to_bytes();
3702 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
3703 assert_eq!(json["ready"], true);
3704 assert!(
3705 json.get("name").is_none(),
3706 "readyz must not expose server name"
3707 );
3708 assert!(
3709 json.get("version").is_none(),
3710 "readyz must not expose version"
3711 );
3712 assert_eq!(json["db"], "connected");
3713 }
3714
3715 #[tokio::test]
3716 async fn readyz_returns_503_when_not_ready() {
3717 let check: ReadinessCheck =
3718 Arc::new(|| Box::pin(async { serde_json::json!({"ready": false}) }));
3719 let resp = readyz(check).await.into_response();
3720 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
3721 }
3722
3723 #[tokio::test]
3724 async fn readyz_returns_503_when_ready_missing() {
3725 let check: ReadinessCheck =
3726 Arc::new(|| Box::pin(async { serde_json::json!({"status": "starting"}) }));
3727 let resp = readyz(check).await.into_response();
3728 assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
3730 }
3731
3732 fn peer_probe_router() -> axum::Router {
3737 async fn probe(req: Request<Body>) -> String {
3738 let ci = req
3739 .extensions()
3740 .get::<ConnectInfo<SocketAddr>>()
3741 .map(|c| c.0.to_string())
3742 .unwrap_or_default();
3743 let pa = req
3744 .extensions()
3745 .get::<PeerAddr>()
3746 .map(|p| p.addr.to_string())
3747 .unwrap_or_default();
3748 format!("{ci}|{pa}")
3749 }
3750 axum::Router::new()
3751 .route("/probe", axum::routing::get(probe))
3752 .layer(axum::middleware::from_fn(|req, next| {
3753 normalize_peer_addr_middleware(None, req, next)
3754 }))
3755 }
3756
3757 async fn body_string(resp: axum::response::Response) -> String {
3758 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
3759 String::from_utf8(bytes.to_vec()).unwrap()
3760 }
3761
3762 #[tokio::test]
3763 async fn normalize_preserves_existing_connect_info_and_mirrors_peer_addr() {
3764 let plain: SocketAddr = "10.0.0.1:1111".parse().unwrap();
3767 let tls: SocketAddr = "10.0.0.2:2222".parse().unwrap();
3768 let req = Request::builder()
3769 .uri("/probe")
3770 .extension(ConnectInfo(plain))
3771 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
3772 .body(Body::empty())
3773 .unwrap();
3774 let resp = peer_probe_router().oneshot(req).await.unwrap();
3775 assert_eq!(resp.status(), StatusCode::OK);
3776 assert_eq!(body_string(resp).await, format!("{plain}|{plain}"));
3777 }
3778
3779 #[tokio::test]
3780 async fn normalize_inserts_connect_info_and_peer_addr_from_tls() {
3781 let tls: SocketAddr = "192.168.1.7:50443".parse().unwrap();
3782 let req = Request::builder()
3783 .uri("/probe")
3784 .extension(ConnectInfo(TlsConnInfo::new(tls, None)))
3785 .body(Body::empty())
3786 .unwrap();
3787 let resp = peer_probe_router().oneshot(req).await.unwrap();
3788 assert_eq!(resp.status(), StatusCode::OK);
3789 assert_eq!(body_string(resp).await, format!("{tls}|{tls}"));
3790 }
3791
3792 #[tokio::test]
3793 async fn normalize_no_op_without_any_connect_info() {
3794 let req = Request::builder()
3795 .uri("/probe")
3796 .body(Body::empty())
3797 .unwrap();
3798 let resp = peer_probe_router().oneshot(req).await.unwrap();
3799 assert_eq!(resp.status(), StatusCode::OK);
3800 assert_eq!(body_string(resp).await, "|");
3801 }
3802
3803 #[tokio::test]
3804 async fn peer_addr_extractor_rejects_when_absent() {
3805 async fn h(peer: PeerAddr) -> String {
3806 peer.addr.to_string()
3807 }
3808 let app = axum::Router::new().route("/p", axum::routing::get(h));
3809 let req = Request::builder().uri("/p").body(Body::empty()).unwrap();
3810 let resp = app.oneshot(req).await.unwrap();
3811 assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
3812 }
3813
3814 #[tokio::test]
3815 async fn peer_addr_extractor_returns_value_when_present() {
3816 async fn h(peer: PeerAddr) -> String {
3817 peer.addr.to_string()
3818 }
3819 let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap();
3820 let app = axum::Router::new().route("/p", axum::routing::get(h));
3821 let req = Request::builder()
3822 .uri("/p")
3823 .extension(PeerAddr::new(addr))
3824 .body(Body::empty())
3825 .unwrap();
3826 let resp = app.oneshot(req).await.unwrap();
3827 assert_eq!(resp.status(), StatusCode::OK);
3828 assert_eq!(body_string(resp).await, addr.to_string());
3829 }
3830
3831 #[tokio::test]
3832 async fn peer_addr_via_extension_extractor() {
3833 async fn h(axum::Extension(peer): axum::Extension<PeerAddr>) -> String {
3834 peer.addr.to_string()
3835 }
3836 let addr: SocketAddr = "127.0.0.1:4242".parse().unwrap();
3837 let app = axum::Router::new().route("/p", axum::routing::get(h));
3838 let req = Request::builder()
3839 .uri("/p")
3840 .extension(PeerAddr::new(addr))
3841 .body(Body::empty())
3842 .unwrap();
3843 let resp = app.oneshot(req).await.unwrap();
3844 assert_eq!(resp.status(), StatusCode::OK);
3845 assert_eq!(body_string(resp).await, addr.to_string());
3846 }
3847
3848 fn limited_router(per_minute: u32) -> axum::Router {
3853 limited_router_with_burst(per_minute, None)
3854 }
3855
3856 fn limited_router_with_burst(per_minute: u32, burst: Option<u32>) -> axum::Router {
3858 limited_router_full(per_minute, burst, &[])
3859 }
3860
3861 fn limited_router_full(
3865 per_minute: u32,
3866 burst: Option<u32>,
3867 exempt_paths: &[&str],
3868 ) -> axum::Router {
3869 let limiter = build_extra_route_rate_limiter(per_minute, burst);
3870 let exempt: Arc<std::collections::HashSet<String>> =
3871 Arc::new(exempt_paths.iter().map(|s| (*s).to_owned()).collect());
3872 axum::Router::new()
3873 .route("/limited", axum::routing::get(|| async { "ok" }))
3874 .route("/exempt", axum::routing::get(|| async { "ok" }))
3875 .layer(axum::middleware::from_fn(move |req, next| {
3876 let l = Arc::clone(&limiter);
3877 let e = Arc::clone(&exempt);
3878 extra_route_rate_limit_middleware(l, e, req, next)
3879 }))
3880 }
3881
3882 fn limited_req(ip: &str) -> Request<Body> {
3883 limited_req_to(ip, "/limited")
3884 }
3885
3886 fn limited_req_to(ip: &str, path: &str) -> Request<Body> {
3887 let addr: SocketAddr = format!("{ip}:40000").parse().unwrap();
3888 Request::builder()
3889 .uri(path)
3890 .extension(ConnectInfo(addr))
3891 .body(Body::empty())
3892 .unwrap()
3893 }
3894
3895 #[tokio::test]
3896 async fn extra_route_limiter_denies_over_quota() {
3897 let app = limited_router(2);
3898 for i in 0..2 {
3899 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
3900 assert_eq!(resp.status(), StatusCode::OK, "request {i} should pass");
3901 }
3902 let resp = app.clone().oneshot(limited_req("10.1.1.1")).await.unwrap();
3903 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
3904 let body = body_string(resp).await;
3905 assert!(
3906 body.contains("too many requests to application routes"),
3907 "deny body should match the limiter message, got: {body}"
3908 );
3909 }
3910
3911 #[tokio::test]
3912 async fn extra_route_limiter_isolates_keys() {
3913 let app = limited_router(2);
3914 for _ in 0..2 {
3915 let resp = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
3916 assert_eq!(resp.status(), StatusCode::OK);
3917 }
3918 let exhausted = app.clone().oneshot(limited_req("10.2.2.2")).await.unwrap();
3919 assert_eq!(exhausted.status(), StatusCode::TOO_MANY_REQUESTS);
3920 let other = app.clone().oneshot(limited_req("10.3.3.3")).await.unwrap();
3922 assert_eq!(other.status(), StatusCode::OK);
3923 }
3924
3925 #[tokio::test]
3926 async fn extra_route_limiter_fails_open_without_peer() {
3927 let app = limited_router(1);
3928 for i in 0..3 {
3929 let req = Request::builder()
3930 .uri("/limited")
3931 .body(Body::empty())
3932 .unwrap();
3933 let resp = app.clone().oneshot(req).await.unwrap();
3934 assert_eq!(
3935 resp.status(),
3936 StatusCode::OK,
3937 "request {i} should fail open"
3938 );
3939 }
3940 }
3941
3942 #[tokio::test]
3943 async fn extra_route_limiter_extracts_tls_conn_info() {
3944 let app = limited_router(2);
3945 let mk = || {
3946 let addr: SocketAddr = "192.168.9.9:55555".parse().unwrap();
3947 Request::builder()
3948 .uri("/limited")
3949 .extension(ConnectInfo(TlsConnInfo::new(addr, None)))
3950 .body(Body::empty())
3951 .unwrap()
3952 };
3953 for _ in 0..2 {
3954 assert_eq!(
3955 app.clone().oneshot(mk()).await.unwrap().status(),
3956 StatusCode::OK
3957 );
3958 }
3959 let resp = app.clone().oneshot(mk()).await.unwrap();
3960 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
3961 }
3962
3963 #[tokio::test]
3964 async fn extra_route_limiter_exempt_path_bypasses_quota() {
3965 let app = limited_router_full(1, None, &["/exempt"]);
3968 for i in 0..5 {
3969 let resp = app
3970 .clone()
3971 .oneshot(limited_req_to("10.6.6.6", "/exempt"))
3972 .await
3973 .unwrap();
3974 assert_eq!(resp.status(), StatusCode::OK, "exempt request {i}");
3975 }
3976 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
3978 assert_eq!(resp.status(), StatusCode::OK);
3979 let resp = app.clone().oneshot(limited_req("10.6.6.6")).await.unwrap();
3981 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
3982 }
3983
3984 #[tokio::test]
3985 async fn extra_route_limiter_exemption_is_raw_exact_match() {
3986 let app = limited_router_full(1, None, &["/exempt"]);
3989 let ok = app
3990 .clone()
3991 .oneshot(limited_req_to("10.7.7.7", "/exempt/"))
3992 .await
3993 .unwrap();
3994 assert_eq!(
3995 ok.status(),
3996 StatusCode::NOT_FOUND,
3997 "variant path routes 404"
3998 );
3999 let denied = app
4001 .clone()
4002 .oneshot(limited_req_to("10.7.7.7", "/limited"))
4003 .await
4004 .unwrap();
4005 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4006 }
4007
4008 #[cfg(feature = "metrics")]
4009 #[tokio::test]
4010 async fn extra_route_limiter_deny_increments_counter_exempt_does_not() {
4011 let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
4012 let app = limited_router_full(1, None, &["/exempt"]);
4013 let mk = |path: &str| {
4014 let addr: SocketAddr = "10.8.8.8:40000".parse().unwrap();
4015 Request::builder()
4016 .uri(path)
4017 .extension(ConnectInfo(addr))
4018 .extension(Arc::clone(&metrics))
4019 .body(Body::empty())
4020 .unwrap()
4021 };
4022 let counter = || {
4023 metrics
4024 .rate_limited_total
4025 .with_label_values(&["extra_route"])
4026 .get()
4027 };
4028 for _ in 0..3 {
4030 assert_eq!(
4031 app.clone().oneshot(mk("/exempt")).await.unwrap().status(),
4032 StatusCode::OK
4033 );
4034 }
4035 assert_eq!(counter(), 0, "exempt requests must not count as denies");
4036 assert_eq!(
4038 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4039 StatusCode::OK
4040 );
4041 assert_eq!(counter(), 0);
4042 assert_eq!(
4043 app.clone().oneshot(mk("/limited")).await.unwrap().status(),
4044 StatusCode::TOO_MANY_REQUESTS
4045 );
4046 assert_eq!(counter(), 1, "deny must increment the extra_route label");
4047 }
4048
4049 #[test]
4050 fn validate_rejects_exempt_paths_without_base_knob() {
4051 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4052 .with_extra_route_rate_limit_exempt_paths(["/ok"]);
4053 let err = cfg.validate().expect_err("exempt paths without rate limit");
4054 assert!(err.to_string().contains("requires extra_route_rate_limit"));
4055 }
4056
4057 #[test]
4058 fn validate_rejects_malformed_exempt_paths() {
4059 for bad in ["", "no-slash"] {
4060 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4061 .with_extra_route_rate_limit(10)
4062 .with_extra_route_rate_limit_exempt_paths([bad]);
4063 let err = cfg.validate().expect_err("malformed exempt path");
4064 assert!(
4065 err.to_string()
4066 .contains("must be non-empty and start with '/'"),
4067 "entry {bad:?}: {err}"
4068 );
4069 }
4070 }
4071
4072 #[test]
4073 fn validate_accepts_wellformed_exempt_paths() {
4074 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4075 .with_extra_route_rate_limit(10)
4076 .with_extra_route_rate_limit_exempt_paths(["/.well-known/oauth-authorization-server"]);
4077 assert!(cfg.validate().is_ok());
4078 }
4079
4080 #[test]
4081 fn validate_rejects_zero_extra_route_rate_limit() {
4082 let cfg = McpServerConfig::new("127.0.0.1:8080", "test-server", "1.0.0")
4083 .with_extra_route_rate_limit(0);
4084 let err = cfg.validate().expect_err("zero extra route rate limit");
4085 assert!(err.to_string().contains("extra_route_rate_limit"));
4086 }
4087
4088 #[tokio::test]
4089 async fn extra_route_limiter_burst_allows_initial_spike() {
4090 let app = limited_router_with_burst(1, Some(3));
4091 for i in 0..3 {
4092 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4093 assert_eq!(resp.status(), StatusCode::OK, "burst request {i}");
4094 }
4095 let resp = app.clone().oneshot(limited_req("10.4.4.4")).await.unwrap();
4096 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
4097 }
4098
4099 #[tokio::test]
4100 async fn extra_route_limiter_deny_sets_retry_after() {
4101 let app = limited_router(1);
4102 let ok = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4103 assert_eq!(ok.status(), StatusCode::OK);
4104 let denied = app.clone().oneshot(limited_req("10.5.5.5")).await.unwrap();
4105 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
4106 let retry_after = denied
4107 .headers()
4108 .get(header::RETRY_AFTER)
4109 .expect("Retry-After present")
4110 .to_str()
4111 .unwrap()
4112 .parse::<u64>()
4113 .unwrap();
4114 assert!(retry_after >= 1, "delta-seconds must be >= 1");
4115 }
4116
4117 #[test]
4118 fn validate_rejects_zero_burst_knobs() {
4119 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4120 .with_tool_rate_limit(10)
4121 .with_tool_rate_limit_burst(0)
4122 .validate()
4123 .expect_err("zero tool burst");
4124 assert!(err.to_string().contains("tool_rate_limit_burst"));
4125
4126 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4127 .with_extra_route_rate_limit(10)
4128 .with_extra_route_rate_limit_burst(0)
4129 .validate()
4130 .expect_err("zero extra route burst");
4131 assert!(err.to_string().contains("extra_route_rate_limit_burst"));
4132 }
4133
4134 #[test]
4135 fn validate_rejects_orphan_burst_knobs() {
4136 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4137 .with_tool_rate_limit_burst(5)
4138 .validate()
4139 .expect_err("orphan tool burst");
4140 assert!(err.to_string().contains("requires tool_rate_limit"));
4141
4142 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4143 .with_extra_route_rate_limit_burst(5)
4144 .validate()
4145 .expect_err("orphan extra route burst");
4146 assert!(err.to_string().contains("requires extra_route_rate_limit"));
4147 }
4148
4149 #[test]
4150 fn validate_rejects_zero_auth_bursts() {
4151 let auth = AuthConfig::with_keys(vec![])
4152 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_burst(0));
4153 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4154 .with_auth(auth)
4155 .validate()
4156 .expect_err("zero auth burst");
4157 assert!(err.to_string().contains("rate_limit.burst"));
4158
4159 let auth = AuthConfig::with_keys(vec![])
4160 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(0));
4161 let err = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4162 .with_auth(auth)
4163 .validate()
4164 .expect_err("zero pre-auth burst");
4165 assert!(err.to_string().contains("pre_auth_burst"));
4166 }
4167
4168 #[test]
4171 fn validate_accepts_pre_auth_burst_without_explicit_pre_auth_rate() {
4172 let auth = AuthConfig::with_keys(vec![])
4173 .with_rate_limit(crate::auth::RateLimitConfig::new(10).with_pre_auth_burst(50));
4174 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_auth(auth);
4175 assert!(cfg.validate().is_ok(), "pre_auth_burst has no orphan rule");
4176 }
4177
4178 fn forward_resolver(trusted: &[&str], mode: ForwardedHeaderMode) -> Arc<ForwardResolver> {
4181 Arc::new(ForwardResolver {
4182 trusted: trusted.iter().map(|s| s.parse().unwrap()).collect(),
4183 mode,
4184 })
4185 }
4186
4187 fn forwarded_probe_router(resolver: Option<Arc<ForwardResolver>>) -> axum::Router {
4189 async fn probe(req: Request<Body>) -> String {
4190 let pa = req
4191 .extensions()
4192 .get::<PeerAddr>()
4193 .map(|p| p.addr.ip().to_string())
4194 .unwrap_or_default();
4195 let ci = req
4196 .extensions()
4197 .get::<ClientIp>()
4198 .map(|c| c.ip.to_string())
4199 .unwrap_or_default();
4200 format!("{pa}|{ci}")
4201 }
4202 axum::Router::new()
4203 .route("/probe", axum::routing::get(probe))
4204 .layer(axum::middleware::from_fn(move |req, next| {
4205 let r = resolver.clone();
4206 normalize_peer_addr_middleware(r, req, next)
4207 }))
4208 }
4209
4210 fn probe_req(peer: &str, header: Option<(&str, &str)>) -> Request<Body> {
4211 let addr: SocketAddr = peer.parse().unwrap();
4212 let mut builder = Request::builder()
4213 .uri("/probe")
4214 .extension(ConnectInfo(addr));
4215 if let Some((name, value)) = header {
4216 builder = builder.header(name, value);
4217 }
4218 builder.body(Body::empty()).unwrap()
4219 }
4220
4221 #[tokio::test]
4222 async fn client_ip_equals_direct_without_resolver() {
4223 let app = forwarded_probe_router(None);
4224 let resp = app
4225 .oneshot(probe_req(
4226 "10.1.2.3:4444",
4227 Some(("x-forwarded-for", "203.0.113.7")),
4228 ))
4229 .await
4230 .unwrap();
4231 assert_eq!(
4232 body_string(resp).await,
4233 "10.1.2.3|10.1.2.3",
4234 "feature off: header ignored, ClientIp == direct"
4235 );
4236 }
4237
4238 #[tokio::test]
4239 async fn client_ip_resolved_for_trusted_peer() {
4240 let app = forwarded_probe_router(Some(forward_resolver(
4241 &["10.0.0.0/8"],
4242 ForwardedHeaderMode::XForwardedFor,
4243 )));
4244 let resp = app
4245 .oneshot(probe_req(
4246 "10.0.0.1:9999",
4247 Some(("x-forwarded-for", "203.0.113.7")),
4248 ))
4249 .await
4250 .unwrap();
4251 assert_eq!(
4252 body_string(resp).await,
4253 "10.0.0.1|203.0.113.7",
4254 "PeerAddr stays direct while ClientIp resolves"
4255 );
4256 }
4257
4258 #[tokio::test]
4259 async fn client_ip_falls_back_to_direct_on_malformed_header() {
4260 let app = forwarded_probe_router(Some(forward_resolver(
4261 &["10.0.0.0/8"],
4262 ForwardedHeaderMode::XForwardedFor,
4263 )));
4264 let resp = app
4265 .oneshot(probe_req(
4266 "10.0.0.1:9999",
4267 Some(("x-forwarded-for", "not-an-ip")),
4268 ))
4269 .await
4270 .unwrap();
4271 assert_eq!(
4272 body_string(resp).await,
4273 "10.0.0.1|10.0.0.1",
4274 "malformed chain falls back to the direct peer"
4275 );
4276 }
4277
4278 #[test]
4279 fn forwarded_header_mode_deserializes_kebab_case() {
4280 #[derive(serde::Deserialize)]
4281 struct Wrapper {
4282 mode: ForwardedHeaderMode,
4283 }
4284 let w: Wrapper = toml::from_str(r#"mode = "x-forwarded-for""#).unwrap();
4285 assert_eq!(w.mode, ForwardedHeaderMode::XForwardedFor);
4286 let w: Wrapper = toml::from_str(r#"mode = "forwarded""#).unwrap();
4287 assert_eq!(w.mode, ForwardedHeaderMode::Forwarded);
4288 assert!(
4289 toml::from_str::<Wrapper>(r#"mode = "XForwardedFor""#).is_err(),
4290 "PascalCase wire value must be rejected"
4291 );
4292 }
4293
4294 #[test]
4295 fn validate_rejects_bad_trusted_proxy_entry() {
4296 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4297 .with_trusted_proxies(["not-a-cidr"]);
4298 let err = cfg.validate().expect_err("bad CIDR");
4299 assert!(err.to_string().contains("trusted_proxies"));
4300 }
4301
4302 #[test]
4303 fn validate_accepts_cidr_and_bare_ip_proxy_entries() {
4304 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0").with_trusted_proxies([
4305 "10.0.0.0/8",
4306 "192.0.2.1",
4307 "2001:db8::1",
4308 ]);
4309 assert!(cfg.validate().is_ok(), "CIDRs and bare IPs are accepted");
4310 }
4311
4312 #[test]
4313 fn validate_rejects_forwarded_header_without_proxies() {
4314 let cfg = McpServerConfig::new("127.0.0.1:8080", "t", "1.0.0")
4315 .with_forwarded_header(ForwardedHeaderMode::Forwarded);
4316 let err = cfg.validate().expect_err("mode without proxies");
4317 assert!(err.to_string().contains("requires trusted_proxies"));
4318 }
4319
4320 fn origin_router(origins: Vec<String>, log_request_headers: bool) -> axum::Router {
4324 let allowed: Arc<[String]> = Arc::from(origins);
4325 axum::Router::new()
4326 .route("/test", axum::routing::get(|| async { "ok" }))
4327 .layer(axum::middleware::from_fn(move |req, next| {
4328 let a = Arc::clone(&allowed);
4329 origin_check_middleware(a, log_request_headers, req, next)
4330 }))
4331 }
4332
4333 #[tokio::test]
4334 async fn origin_allowed_passes() {
4335 let app = origin_router(vec!["http://localhost:3000".into()], false);
4336 let req = Request::builder()
4337 .uri("/test")
4338 .header(header::ORIGIN, "http://localhost:3000")
4339 .body(Body::empty())
4340 .unwrap();
4341 let resp = app.oneshot(req).await.unwrap();
4342 assert_eq!(resp.status(), StatusCode::OK);
4343 }
4344
4345 #[tokio::test]
4346 async fn origin_rejected_returns_403() {
4347 let app = origin_router(vec!["http://localhost:3000".into()], false);
4348 let req = Request::builder()
4349 .uri("/test")
4350 .header(header::ORIGIN, "http://evil.com")
4351 .body(Body::empty())
4352 .unwrap();
4353 let resp = app.oneshot(req).await.unwrap();
4354 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
4355 }
4356
4357 #[tokio::test]
4358 async fn no_origin_header_passes() {
4359 let app = origin_router(vec!["http://localhost:3000".into()], false);
4360 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4361 let resp = app.oneshot(req).await.unwrap();
4362 assert_eq!(resp.status(), StatusCode::OK);
4363 }
4364
4365 #[tokio::test]
4366 async fn empty_allowlist_rejects_any_origin() {
4367 let app = origin_router(vec![], false);
4368 let req = Request::builder()
4369 .uri("/test")
4370 .header(header::ORIGIN, "http://anything.com")
4371 .body(Body::empty())
4372 .unwrap();
4373 let resp = app.oneshot(req).await.unwrap();
4374 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
4375 }
4376
4377 #[tokio::test]
4378 async fn empty_allowlist_passes_without_origin() {
4379 let app = origin_router(vec![], false);
4380 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4381 let resp = app.oneshot(req).await.unwrap();
4382 assert_eq!(resp.status(), StatusCode::OK);
4383 }
4384
4385 #[test]
4386 fn format_request_headers_redacts_sensitive_values() {
4387 let mut headers = axum::http::HeaderMap::new();
4388 headers.insert("authorization", "Bearer secret-token".parse().unwrap());
4389 headers.insert("cookie", "sid=abc".parse().unwrap());
4390 headers.insert("x-request-id", "req-123".parse().unwrap());
4391
4392 let out = format_request_headers_for_log(&headers);
4393 assert!(out.contains("authorization: [REDACTED]"));
4394 assert!(out.contains("cookie: [REDACTED]"));
4395 assert!(out.contains("x-request-id: req-123"));
4396 assert!(!out.contains("secret-token"));
4397 }
4398
4399 fn security_router(is_tls: bool) -> axum::Router {
4402 security_router_with(is_tls, SecurityHeadersConfig::default())
4403 }
4404
4405 fn security_router_with(is_tls: bool, cfg: SecurityHeadersConfig) -> axum::Router {
4406 let cfg = Arc::new(cfg);
4407 axum::Router::new()
4408 .route("/test", axum::routing::get(|| async { "ok" }))
4409 .layer(axum::middleware::from_fn(move |req, next| {
4410 let c = Arc::clone(&cfg);
4411 security_headers_middleware(is_tls, c, req, next)
4412 }))
4413 }
4414
4415 #[tokio::test]
4416 async fn security_headers_set_on_response() {
4417 let app = security_router(false);
4418 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4419 let resp = app.oneshot(req).await.unwrap();
4420 assert_eq!(resp.status(), StatusCode::OK);
4421
4422 let h = resp.headers();
4423 assert_eq!(h.get("x-content-type-options").unwrap(), "nosniff");
4424 assert_eq!(h.get("x-frame-options").unwrap(), "deny");
4425 assert_eq!(h.get("cache-control").unwrap(), "no-store, max-age=0");
4426 assert_eq!(h.get("referrer-policy").unwrap(), "no-referrer");
4427 assert_eq!(h.get("cross-origin-opener-policy").unwrap(), "same-origin");
4428 assert_eq!(
4429 h.get("cross-origin-resource-policy").unwrap(),
4430 "same-origin"
4431 );
4432 assert_eq!(
4433 h.get("cross-origin-embedder-policy").unwrap(),
4434 "require-corp"
4435 );
4436 assert_eq!(h.get("x-permitted-cross-domain-policies").unwrap(), "none");
4437 assert!(
4438 h.get("permissions-policy")
4439 .unwrap()
4440 .to_str()
4441 .unwrap()
4442 .contains("camera=()"),
4443 "permissions-policy must restrict browser features"
4444 );
4445 assert_eq!(
4446 h.get("content-security-policy").unwrap(),
4447 "default-src 'none'; frame-ancestors 'none'"
4448 );
4449 assert_eq!(h.get("x-dns-prefetch-control").unwrap(), "off");
4450 assert!(h.get("strict-transport-security").is_none());
4452 }
4453
4454 #[tokio::test]
4455 async fn hsts_set_when_tls_enabled() {
4456 let app = security_router(true);
4457 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4458 let resp = app.oneshot(req).await.unwrap();
4459
4460 let hsts = resp.headers().get("strict-transport-security").unwrap();
4461 assert!(
4462 hsts.to_str().unwrap().contains("max-age=63072000"),
4463 "HSTS must set 2-year max-age"
4464 );
4465 }
4466
4467 fn check_with_security_headers(headers: SecurityHeadersConfig) -> Result<(), McpxError> {
4473 let cfg =
4474 McpServerConfig::new("127.0.0.1:8080", "test", "0.0.0").with_security_headers(headers);
4475 cfg.check()
4476 }
4477
4478 #[test]
4479 fn security_headers_config_default_validates() {
4480 check_with_security_headers(SecurityHeadersConfig::default())
4481 .expect("default SecurityHeadersConfig must validate");
4482 }
4483
4484 #[test]
4485 fn security_headers_config_validate_accepts_empty_string() {
4486 let h = SecurityHeadersConfig {
4488 x_content_type_options: Some(String::new()),
4489 x_frame_options: Some(String::new()),
4490 cache_control: Some(String::new()),
4491 referrer_policy: Some(String::new()),
4492 cross_origin_opener_policy: Some(String::new()),
4493 cross_origin_resource_policy: Some(String::new()),
4494 cross_origin_embedder_policy: Some(String::new()),
4495 permissions_policy: Some(String::new()),
4496 x_permitted_cross_domain_policies: Some(String::new()),
4497 content_security_policy: Some(String::new()),
4498 x_dns_prefetch_control: Some(String::new()),
4499 strict_transport_security: Some(String::new()),
4500 };
4501 check_with_security_headers(h).expect("Some(\"\") on every field must validate (omit-all)");
4502 }
4503
4504 #[test]
4505 fn security_headers_config_validate_rejects_bad_value() {
4506 let h = SecurityHeadersConfig {
4508 referrer_policy: Some("\u{0007}".into()),
4509 ..SecurityHeadersConfig::default()
4510 };
4511 let err = check_with_security_headers(h)
4512 .expect_err("control char in referrer_policy must reject");
4513 let msg = err.to_string();
4514 assert!(
4515 msg.contains("referrer_policy"),
4516 "error must name the offending field, got: {msg}"
4517 );
4518 }
4519
4520 #[test]
4521 fn security_headers_config_validate_rejects_hsts_preload() {
4522 let h = SecurityHeadersConfig {
4523 strict_transport_security: Some("max-age=63072000; includeSubDomains; preload".into()),
4524 ..SecurityHeadersConfig::default()
4525 };
4526 let err = check_with_security_headers(h).expect_err("HSTS with preload must reject");
4527 let msg = err.to_string();
4528 assert!(
4529 msg.contains("strict_transport_security"),
4530 "error must name the field, got: {msg}"
4531 );
4532 assert!(
4533 msg.to_lowercase().contains("preload"),
4534 "error must mention `preload`, got: {msg}"
4535 );
4536 }
4537
4538 #[test]
4539 fn security_headers_config_validate_rejects_hsts_preload_uppercase() {
4540 let h = SecurityHeadersConfig {
4542 strict_transport_security: Some("max-age=600; PRELOAD".into()),
4543 ..SecurityHeadersConfig::default()
4544 };
4545 check_with_security_headers(h).expect_err("HSTS preload check must be case-insensitive");
4546 }
4547
4548 #[tokio::test]
4549 async fn security_headers_override_honored() {
4550 let h = SecurityHeadersConfig {
4552 x_frame_options: Some("SAMEORIGIN".into()),
4553 ..SecurityHeadersConfig::default()
4554 };
4555 let app = security_router_with(false, h);
4556 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4557 let resp = app.oneshot(req).await.unwrap();
4558 assert_eq!(resp.status(), StatusCode::OK);
4559
4560 let xfo = resp.headers().get("x-frame-options").unwrap();
4561 assert_eq!(xfo, "SAMEORIGIN");
4562 }
4563
4564 #[tokio::test]
4565 async fn security_headers_empty_string_omits() {
4566 let h = SecurityHeadersConfig {
4568 referrer_policy: Some(String::new()),
4569 ..SecurityHeadersConfig::default()
4570 };
4571 let app = security_router_with(false, h);
4572 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4573 let resp = app.oneshot(req).await.unwrap();
4574 assert_eq!(resp.status(), StatusCode::OK);
4575
4576 assert!(
4577 resp.headers().get("referrer-policy").is_none(),
4578 "Some(\"\") must omit the header"
4579 );
4580 assert_eq!(
4582 resp.headers().get("x-content-type-options").unwrap(),
4583 "nosniff"
4584 );
4585 }
4586
4587 #[tokio::test]
4588 async fn security_headers_hsts_only_when_tls() {
4589 let h = SecurityHeadersConfig {
4591 strict_transport_security: Some("max-age=600".into()),
4592 ..SecurityHeadersConfig::default()
4593 };
4594 let app = security_router_with(false, h);
4595 let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
4596 let resp = app.oneshot(req).await.unwrap();
4597 assert!(
4598 resp.headers().get("strict-transport-security").is_none(),
4599 "HSTS must remain absent on plaintext deployments even with override"
4600 );
4601 }
4602
4603 #[cfg(feature = "oauth")]
4606 #[tokio::test]
4607 async fn oauth_token_cache_headers_set_pragma_and_vary() {
4608 let app = axum::Router::new()
4609 .route("/token", axum::routing::post(|| async { "{}" }))
4610 .layer(axum::middleware::from_fn(
4611 oauth_token_cache_headers_middleware,
4612 ));
4613 let req = Request::builder()
4614 .method("POST")
4615 .uri("/token")
4616 .body(Body::from("{}"))
4617 .unwrap();
4618 let resp = app.oneshot(req).await.unwrap();
4619 assert_eq!(resp.status(), StatusCode::OK);
4620
4621 let h = resp.headers();
4622 assert_eq!(
4623 h.get("pragma").unwrap(),
4624 "no-cache",
4625 "RFC 6749 §5.1: token responses must set Pragma: no-cache"
4626 );
4627 let vary_values: Vec<String> = h
4628 .get_all("vary")
4629 .iter()
4630 .filter_map(|v| v.to_str().ok().map(str::to_owned))
4631 .collect();
4632 assert!(
4633 vary_values
4634 .iter()
4635 .any(|v| v.eq_ignore_ascii_case("Authorization")),
4636 "RFC 6750 §5.4: Vary must include Authorization, got {vary_values:?}"
4637 );
4638 }
4639
4640 #[cfg(feature = "oauth")]
4641 #[tokio::test]
4642 async fn oauth_token_cache_headers_preserve_existing_vary() {
4643 let app = axum::Router::new()
4646 .route(
4647 "/token",
4648 axum::routing::post(|| async {
4649 axum::response::Response::builder()
4650 .header("vary", "Accept-Encoding")
4651 .body(axum::body::Body::from("{}"))
4652 .unwrap()
4653 }),
4654 )
4655 .layer(axum::middleware::from_fn(
4656 oauth_token_cache_headers_middleware,
4657 ));
4658 let req = Request::builder()
4659 .method("POST")
4660 .uri("/token")
4661 .body(Body::empty())
4662 .unwrap();
4663 let resp = app.oneshot(req).await.unwrap();
4664
4665 let vary: Vec<String> = resp
4666 .headers()
4667 .get_all("vary")
4668 .iter()
4669 .filter_map(|v| v.to_str().ok().map(str::to_owned))
4670 .collect();
4671 assert!(
4672 vary.iter().any(|v| v.contains("Accept-Encoding")),
4673 "must preserve pre-existing Vary value, got {vary:?}"
4674 );
4675 assert!(
4676 vary.iter().any(|v| v.contains("Authorization")),
4677 "must append Authorization to Vary, got {vary:?}"
4678 );
4679 }
4680
4681 #[test]
4684 fn version_payload_contains_expected_fields() {
4685 let v = version_payload("my-server", "1.2.3");
4686 assert_eq!(v["name"], "my-server");
4687 assert_eq!(v["version"], "1.2.3");
4688 assert!(v["build_git_sha"].is_string());
4689 assert!(v["build_timestamp"].is_string());
4690 assert!(v["rust_version"].is_string());
4691 assert!(v["mcpx_version"].is_string());
4692 }
4693
4694 #[tokio::test]
4697 async fn concurrency_limit_layer_composes_and_serves() {
4698 let app = axum::Router::new()
4702 .route("/ok", axum::routing::get(|| async { "ok" }))
4703 .layer(
4704 tower::ServiceBuilder::new()
4705 .layer(axum::error_handling::HandleErrorLayer::new(
4706 |_err: tower::BoxError| async { StatusCode::SERVICE_UNAVAILABLE },
4707 ))
4708 .layer(tower::load_shed::LoadShedLayer::new())
4709 .layer(tower::limit::ConcurrencyLimitLayer::new(4)),
4710 );
4711 let resp = app
4712 .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap())
4713 .await
4714 .unwrap();
4715 assert_eq!(resp.status(), StatusCode::OK);
4716 }
4717
4718 #[tokio::test]
4721 async fn compression_layer_gzip_encodes_response() {
4722 use tower_http::compression::Predicate as _;
4723
4724 let big_body = "a".repeat(4096);
4725 let app = axum::Router::new()
4726 .route(
4727 "/big",
4728 axum::routing::get(move || {
4729 let body = big_body.clone();
4730 async move { body }
4731 }),
4732 )
4733 .layer(
4734 tower_http::compression::CompressionLayer::new()
4735 .gzip(true)
4736 .br(true)
4737 .compress_when(
4738 tower_http::compression::DefaultPredicate::new()
4739 .and(tower_http::compression::predicate::SizeAbove::new(1024)),
4740 ),
4741 );
4742
4743 let req = Request::builder()
4744 .uri("/big")
4745 .header(header::ACCEPT_ENCODING, "gzip")
4746 .body(Body::empty())
4747 .unwrap();
4748 let resp = app.oneshot(req).await.unwrap();
4749 assert_eq!(resp.status(), StatusCode::OK);
4750 assert_eq!(
4751 resp.headers().get(header::CONTENT_ENCODING).unwrap(),
4752 "gzip"
4753 );
4754 }
4755
4756 #[tokio::test]
4759 async fn tls_handshake_timeout_reaps_idle_connections() {
4760 use tokio::io::AsyncReadExt as _;
4761
4762 let _ = rustls::crypto::ring::default_provider().install_default();
4763
4764 let key = rcgen::KeyPair::generate().expect("generate key");
4766 let cert = rcgen::CertificateParams::new(vec!["localhost".to_owned()])
4767 .expect("cert params")
4768 .self_signed(&key)
4769 .expect("self-signed cert");
4770 let dir = std::env::temp_dir().join(format!(
4771 "rmcp-server-kit-hs-timeout-{}",
4772 std::time::SystemTime::now()
4773 .duration_since(std::time::UNIX_EPOCH)
4774 .expect("clock after epoch")
4775 .as_nanos()
4776 ));
4777 tokio::fs::create_dir_all(&dir).await.expect("temp dir");
4778 let cert_path = dir.join("server.crt");
4779 let key_path = dir.join("server.key");
4780 tokio::fs::write(&cert_path, cert.pem())
4781 .await
4782 .expect("write cert");
4783 tokio::fs::write(&key_path, key.serialize_pem())
4784 .await
4785 .expect("write key");
4786
4787 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
4788 let tls = TlsListener::new(
4789 listener,
4790 &cert_path,
4791 &key_path,
4792 None,
4793 None,
4794 Duration::from_millis(200),
4795 8, )
4797 .expect("tls listener");
4798 let addr = axum::serve::Listener::local_addr(&tls).expect("local addr");
4799
4800 let mut idle = tokio::net::TcpStream::connect(addr).await.expect("connect");
4804 let mut buf = [0_u8; 16];
4805 let read = tokio::time::timeout(Duration::from_secs(2), idle.read(&mut buf))
4806 .await
4807 .expect("server must reap the idle handshake within its timeout");
4808 match read {
4809 Ok(0) | Err(_) => {} Ok(n) => panic!("unexpected {n} bytes from server during reaped handshake"),
4811 }
4812
4813 drop(tls);
4814 }
4815}