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