1use std::collections::HashMap;
4use std::convert::Infallible;
5use std::time::Duration;
6
7use anyhow::{Context, Result};
8use axum::Router;
9use tokio::process::Command;
10use tower::timeout::TimeoutLayer;
11use tower::util::BoxCloneService;
12use tower::{Layer, ServiceExt};
13use tower_mcp::SessionHandle;
14use tower_mcp::auth::{AuthLayer, StaticBearerValidator};
15use tower_mcp::client::StdioClientTransport;
16use tower_mcp::proxy::McpProxy;
17use tower_mcp::{RouterRequest, RouterResponse};
18use tower_resilience::retry::RetryLayer;
19
20use crate::admin::BackendMeta;
21use crate::alias;
22use crate::cache;
23use crate::coalesce;
24use crate::config::{AuthConfig, ProxyConfig, TransportType};
25use crate::filter::CapabilityFilterService;
26#[cfg(feature = "oauth")]
27use crate::rbac::{RbacConfig, RbacService};
28use crate::validation::{ValidationConfig, ValidationService};
29
30pub struct Proxy {
32 router: Router,
33 session_handle: SessionHandle,
34 inner: McpProxy,
35 config: ProxyConfig,
36 #[cfg(feature = "discovery")]
37 discovery_index: Option<crate::discovery::SharedDiscoveryIndex>,
38}
39
40impl Proxy {
41 pub async fn from_config(config: ProxyConfig) -> Result<Self> {
47 let (mcp_proxy, cb_handles) = build_mcp_proxy(&config).await?;
48 let proxy_for_admin = mcp_proxy.clone();
49 #[cfg_attr(not(feature = "discovery"), allow(unused_mut))]
51 let mut proxy_for_caller = mcp_proxy.clone();
52 let proxy_for_management = mcp_proxy.clone();
53
54 #[cfg(feature = "metrics")]
56 let metrics_handle = if config.observability.metrics.enabled {
57 tracing::info!("Prometheus metrics enabled at /admin/metrics");
58 let builder = metrics_exporter_prometheus::PrometheusBuilder::new();
59 let handle = builder
60 .install_recorder()
61 .context("installing Prometheus metrics recorder")?;
62 Some(handle)
63 } else {
64 None
65 };
66 #[cfg(not(feature = "metrics"))]
67 let metrics_handle = None;
68
69 let (service, cache_handle) = build_middleware_stack(&config, mcp_proxy)?;
70
71 let (router, session_handle) =
72 tower_mcp::transport::http::HttpTransport::from_service(service)
73 .into_router_with_handle();
74
75 let router = apply_auth(&config, router).await?;
77
78 let backend_meta: std::collections::HashMap<String, BackendMeta> = config
80 .backends
81 .iter()
82 .map(|b| {
83 (
84 b.name.clone(),
85 BackendMeta {
86 transport: format!("{:?}", b.transport).to_lowercase(),
87 },
88 )
89 })
90 .collect();
91
92 let admin_state = crate::admin::spawn_health_checker(
94 proxy_for_admin,
95 config.proxy.name.clone(),
96 config.proxy.version.clone(),
97 config.backends.len(),
98 backend_meta,
99 );
100 let router = router.nest(
101 "/admin",
102 crate::admin::admin_router(
103 admin_state.clone(),
104 metrics_handle,
105 session_handle.clone(),
106 cache_handle,
107 proxy_for_management,
108 &config,
109 config.source_path.clone(),
110 cb_handles,
111 ),
112 );
113 tracing::info!("Admin API enabled at /admin/backends");
114
115 #[cfg(feature = "discovery")]
117 let discovery_enabled = config.proxy.tool_discovery
118 || config.proxy.tool_exposure == crate::config::ToolExposure::Search;
119 #[cfg(feature = "discovery")]
120 let (discovery_index, discovery_tools) = if discovery_enabled {
121 let index =
122 crate::discovery::build_index(&mut proxy_for_caller, &config.proxy.separator).await;
123 let tools = crate::discovery::build_discovery_tools(index.clone());
124 (Some(index), Some(tools))
125 } else {
126 (None, None)
127 };
128 #[cfg(not(feature = "discovery"))]
129 let discovery_tools: Option<Vec<tower_mcp::Tool>> = None;
130
131 if let Err(e) = crate::admin_tools::register_admin_tools(
133 &proxy_for_caller,
134 admin_state,
135 session_handle.clone(),
136 &config,
137 discovery_tools,
138 )
139 .await
140 {
141 tracing::warn!("Failed to register admin tools: {e}");
142 } else {
143 tracing::info!("MCP admin tools registered under proxy/ namespace");
144 }
145
146 Ok(Self {
147 router,
148 session_handle,
149 inner: proxy_for_caller,
150 config,
151 #[cfg(feature = "discovery")]
152 discovery_index,
153 })
154 }
155
156 pub fn session_handle(&self) -> &SessionHandle {
158 &self.session_handle
159 }
160
161 pub fn mcp_proxy(&self) -> &McpProxy {
165 &self.inner
166 }
167
168 pub fn enable_hot_reload(&self, config_path: std::path::PathBuf) {
173 tracing::info!("Hot reload enabled, watching config file for changes");
174 crate::reload::spawn_config_watcher(
175 config_path,
176 self.inner.clone(),
177 #[cfg(feature = "discovery")]
178 self.discovery_index
179 .as_ref()
180 .map(|idx| (idx.clone(), self.config.proxy.separator.clone())),
181 );
182 }
183
184 pub fn into_router(self) -> (Router, SessionHandle) {
196 (self.router, self.session_handle)
197 }
198
199 pub async fn serve(self) -> Result<()> {
204 let addr = format!(
205 "{}:{}",
206 self.config.proxy.listen.host, self.config.proxy.listen.port
207 );
208
209 tracing::info!(listen = %addr, "Proxy ready");
210
211 let listener = tokio::net::TcpListener::bind(&addr)
212 .await
213 .with_context(|| format!("binding to {}", addr))?;
214
215 let shutdown_timeout = Duration::from_secs(self.config.proxy.shutdown_timeout_seconds);
216 axum::serve(listener, self.router)
217 .with_graceful_shutdown(shutdown_signal(shutdown_timeout))
218 .await
219 .context("server error")?;
220
221 tracing::info!("Proxy shut down");
222 Ok(())
223 }
224}
225
226pub type CbHandle = tower_resilience::circuitbreaker::CircuitBreakerHandle;
228
229fn build_breaker_layer(
237 cb: &crate::config::CircuitBreakerConfig,
238 backend_name: &str,
239) -> (
240 tower_resilience::circuitbreaker::CircuitBreakerLayer,
241 CbHandle,
242) {
243 tower_resilience::circuitbreaker::CircuitBreakerLayer::builder()
244 .failure_rate_threshold(cb.failure_rate_threshold)
245 .minimum_number_of_calls(cb.minimum_calls)
246 .sliding_window_size(cb.minimum_calls)
247 .wait_duration_in_open(Duration::from_secs(cb.wait_duration_seconds))
248 .permitted_calls_in_half_open(cb.permitted_calls_in_half_open)
249 .name(format!("{backend_name}-cb"))
250 .build_with_handle()
251}
252
253type InfallibleService = BoxCloneService<RouterRequest, RouterResponse, Infallible>;
254type MwService = BoxCloneService<RouterRequest, RouterResponse, tower::BoxError>;
255type MwStage = Box<dyn Fn(MwService) -> MwService + Send + Sync>;
256
257#[derive(Default)]
272struct BackendMiddlewareLayer {
273 retry: Option<RetryLayer<RouterRequest, RouterResponse, Infallible>>,
275 hedge: Option<tower_resilience::hedge::HedgeLayer>,
277 stages: Vec<MwStage>,
280 outlier: Option<crate::outlier::OutlierDetectionLayer>,
282}
283
284impl BackendMiddlewareLayer {
285 fn is_empty(&self) -> bool {
286 self.retry.is_none()
287 && self.hedge.is_none()
288 && self.stages.is_empty()
289 && self.outlier.is_none()
290 }
291
292 fn stage<L>(&mut self, layer: L)
293 where
294 L: Layer<MwService> + Send + Sync + 'static,
295 L::Service:
296 tower::Service<RouterRequest, Response = RouterResponse> + Clone + Send + 'static,
297 <L::Service as tower::Service<RouterRequest>>::Error: std::fmt::Display,
298 <L::Service as tower::Service<RouterRequest>>::Future: Send + 'static,
299 {
300 self.stages.push(Box::new(move |inner| {
301 BoxCloneService::new(
302 layer
303 .layer(inner)
304 .map_err(|e| -> tower::BoxError { e.to_string().into() }),
305 )
306 }));
307 }
308}
309
310impl<S> Layer<S> for BackendMiddlewareLayer
311where
312 S: tower::Service<RouterRequest, Response = RouterResponse, Error = Infallible>
313 + Clone
314 + Send
315 + 'static,
316 S::Future: Send + 'static,
317{
318 type Service = InfallibleService;
319
320 fn layer(&self, base: S) -> InfallibleService {
321 fn widen(e: Infallible) -> tower::BoxError {
322 match e {}
323 }
324
325 let after_retry: InfallibleService = match &self.retry {
329 Some(retry) => BoxCloneService::new(retry.layer(base)),
330 None => BoxCloneService::new(base),
331 };
332
333 let mut svc: MwService = match &self.hedge {
337 Some(hedge) => BoxCloneService::new(
338 hedge
339 .layer(after_retry)
340 .map_err(|e| -> tower::BoxError { e.to_string().into() }),
341 ),
342 None => BoxCloneService::new(after_retry.map_err(widen)),
343 };
344 for stage in &self.stages {
345 svc = stage(svc);
346 }
347
348 let folded: InfallibleService =
350 BoxCloneService::new(tower_mcp::transport::CatchError::new(svc));
351
352 match &self.outlier {
354 Some(outlier) => BoxCloneService::new(outlier.layer(folded)),
355 None => folded,
356 }
357 }
358}
359
360async fn build_mcp_proxy(config: &ProxyConfig) -> Result<(McpProxy, HashMap<String, CbHandle>)> {
363 let mut builder = McpProxy::builder(&config.proxy.name, &config.proxy.version)
364 .separator(&config.proxy.separator);
365 let mut cb_handles: HashMap<String, CbHandle> = HashMap::new();
366
367 if let Some(instructions) = &config.proxy.instructions {
368 builder = builder.instructions(instructions);
369 }
370
371 let outlier_detector = {
374 let max_pct = config
375 .backends
376 .iter()
377 .filter_map(|b| b.outlier_detection.as_ref())
378 .map(|od| od.max_ejection_percent)
379 .max();
380 max_pct.map(crate::outlier::OutlierDetector::new)
381 };
382
383 for backend in &config.backends {
384 tracing::info!(name = %backend.name, transport = ?backend.transport, "Adding backend");
385
386 match backend.transport {
387 TransportType::Stdio => {
388 let command = backend.command.as_deref().unwrap();
389 let args: Vec<&str> = backend.args.iter().map(|s| s.as_str()).collect();
390
391 let mut cmd = Command::new(command);
392 cmd.args(&args);
393
394 for (key, value) in &backend.env {
395 cmd.env(key, value);
396 }
397
398 let transport = StdioClientTransport::spawn_command(&mut cmd)
399 .await
400 .with_context(|| format!("spawning backend '{}'", backend.name))?;
401
402 builder = builder.backend(&backend.name, transport).await;
403 }
404 TransportType::Http => {
405 let url = backend.url.as_deref().unwrap();
406 let mut transport = tower_mcp::client::HttpClientTransport::new(url);
407 if let Some(token) = &backend.bearer_token {
408 transport = transport.bearer_token(token);
409 }
410
411 builder = builder.backend(&backend.name, transport).await;
412 }
413 #[cfg(feature = "websocket")]
414 TransportType::Websocket => {
415 let url = backend.url.as_deref().unwrap();
416 tracing::info!(url = %url, "Connecting to WebSocket backend");
417 let transport = if let Some(token) = &backend.bearer_token {
418 crate::ws_transport::WebSocketClientTransport::connect_with_bearer_token(
419 url, token,
420 )
421 .await
422 .with_context(|| {
423 format!("connecting to WebSocket backend '{}'", backend.name)
424 })?
425 } else {
426 crate::ws_transport::WebSocketClientTransport::connect(url)
427 .await
428 .with_context(|| {
429 format!("connecting to WebSocket backend '{}'", backend.name)
430 })?
431 };
432
433 builder = builder.backend(&backend.name, transport).await;
434 }
435 #[cfg(not(feature = "websocket"))]
436 TransportType::Websocket => {
437 anyhow::bail!(
438 "WebSocket transport requires the 'websocket' feature. \
439 Rebuild with: cargo install mcp-proxy --features websocket"
440 );
441 }
442 }
443
444 let mut mw = BackendMiddlewareLayer::default();
448
449 if let Some(retry_cfg) = &backend.retry {
451 tracing::info!(
452 backend = %backend.name,
453 max_retries = retry_cfg.max_retries,
454 initial_backoff_ms = retry_cfg.initial_backoff_ms,
455 max_backoff_ms = retry_cfg.max_backoff_ms,
456 "Applying retry policy"
457 );
458 mw.retry = Some(crate::retry::build_retry_layer(retry_cfg, &backend.name));
459 }
460
461 if let Some(hedge_cfg) = &backend.hedging {
463 let delay = Duration::from_millis(hedge_cfg.delay_ms);
464 let max_attempts = hedge_cfg.max_hedges + 1; tracing::info!(
466 backend = %backend.name,
467 delay_ms = hedge_cfg.delay_ms,
468 max_hedges = hedge_cfg.max_hedges,
469 "Applying request hedging"
470 );
471 let layer = if delay.is_zero() {
472 tower_resilience::hedge::HedgeLayer::builder()
473 .no_delay()
474 .max_hedged_attempts(max_attempts)
475 .name(format!("{}-hedge", backend.name))
476 .build()
477 } else {
478 tower_resilience::hedge::HedgeLayer::builder()
479 .delay(delay)
480 .max_hedged_attempts(max_attempts)
481 .name(format!("{}-hedge", backend.name))
482 .build()
483 };
484 mw.hedge = Some(layer);
485 }
486
487 if let Some(cc) = &backend.concurrency {
489 tracing::info!(
490 backend = %backend.name,
491 max = cc.max_concurrent,
492 "Applying concurrency limit"
493 );
494 mw.stage(tower::limit::ConcurrencyLimitLayer::new(cc.max_concurrent));
495 }
496
497 if let Some(rl) = &backend.rate_limit {
499 tracing::info!(
500 backend = %backend.name,
501 requests = rl.requests,
502 period_seconds = rl.period_seconds,
503 "Applying rate limit"
504 );
505 let layer = tower_resilience::ratelimiter::RateLimiterLayer::builder()
506 .limit_for_period(rl.requests)
507 .refresh_period(Duration::from_secs(rl.period_seconds))
508 .name(format!("{}-ratelimit", backend.name))
509 .build();
510 mw.stage(layer);
511 }
512
513 if let Some(timeout) = &backend.timeout {
515 tracing::info!(
516 backend = %backend.name,
517 seconds = timeout.seconds,
518 "Applying timeout"
519 );
520 mw.stage(TimeoutLayer::new(Duration::from_secs(timeout.seconds)));
521 }
522
523 if let Some(cb) = &backend.circuit_breaker {
525 tracing::info!(
526 backend = %backend.name,
527 failure_rate = cb.failure_rate_threshold,
528 wait_seconds = cb.wait_duration_seconds,
529 "Applying circuit breaker"
530 );
531 let (layer, handle) = build_breaker_layer(cb, &backend.name);
532 cb_handles.insert(backend.name.clone(), handle);
533 mw.stage(layer);
534 }
535
536 if let Some(od) = &backend.outlier_detection
538 && let Some(ref detector) = outlier_detector
539 {
540 tracing::info!(
541 backend = %backend.name,
542 consecutive_errors = od.consecutive_errors,
543 base_ejection_seconds = od.base_ejection_seconds,
544 max_ejection_percent = od.max_ejection_percent,
545 "Applying outlier detection"
546 );
547 let layer = crate::outlier::OutlierDetectionLayer::new(
548 backend.name.clone(),
549 od.clone(),
550 detector.clone(),
551 );
552 mw.outlier = Some(layer);
553 }
554
555 if !mw.is_empty() {
556 builder = builder.backend_layer(mw);
557 }
558 }
559
560 let result = builder.build().await?;
561
562 if !result.skipped.is_empty() {
563 for s in &result.skipped {
564 tracing::warn!("Skipped backend: {s}");
565 }
566 }
567
568 Ok((result.proxy, cb_handles))
569}
570
571#[cfg(feature = "oauth")]
578fn oauth_scope_layer(
579 required_scopes: &[String],
580) -> Option<tower_mcp::oauth::ScopeEnforcementLayer> {
581 if required_scopes.is_empty() {
582 return None;
583 }
584 let policy = tower_mcp::oauth::ScopePolicy::new().default_scopes(
585 tower_mcp::oauth::ScopeRequirement::all(required_scopes.iter().cloned()),
586 );
587 Some(tower_mcp::oauth::ScopeEnforcementLayer::new(policy))
588}
589
590fn build_middleware_stack(
592 config: &ProxyConfig,
593 proxy: McpProxy,
594) -> Result<(
595 BoxCloneService<RouterRequest, RouterResponse, Infallible>,
596 Option<cache::CacheHandle>,
597)> {
598 let mut service: BoxCloneService<RouterRequest, RouterResponse, Infallible> =
599 BoxCloneService::new(proxy);
600 let mut cache_handle: Option<cache::CacheHandle> = None;
601
602 let injection_rules: Vec<_> = config
604 .backends
605 .iter()
606 .filter(|b| !b.default_args.is_empty() || !b.inject_args.is_empty())
607 .map(|b| {
608 let namespace = format!("{}{}", b.name, config.proxy.separator);
609 tracing::info!(
610 backend = %b.name,
611 default_args = b.default_args.len(),
612 tool_rules = b.inject_args.len(),
613 "Applying argument injection"
614 );
615 crate::inject::InjectionRules::new(
616 namespace,
617 b.default_args.clone(),
618 b.inject_args.clone(),
619 )
620 })
621 .collect();
622
623 if !injection_rules.is_empty() {
624 service = BoxCloneService::new(crate::inject::InjectArgsService::new(
625 service,
626 injection_rules,
627 ));
628 }
629
630 let param_overrides: Vec<_> = config
632 .backends
633 .iter()
634 .filter(|b| !b.param_overrides.is_empty())
635 .flat_map(|b| {
636 let namespace = format!("{}{}", b.name, config.proxy.separator);
637 tracing::info!(
638 backend = %b.name,
639 overrides = b.param_overrides.len(),
640 "Applying parameter overrides"
641 );
642 b.param_overrides
643 .iter()
644 .map(move |c| crate::param_override::ToolOverride::new(&namespace, c))
645 })
646 .collect();
647
648 if !param_overrides.is_empty() {
649 service = BoxCloneService::new(crate::param_override::ParamOverrideService::new(
650 service,
651 param_overrides,
652 ));
653 }
654
655 let canary_mappings: std::collections::HashMap<String, (String, u32, u32)> = config
657 .backends
658 .iter()
659 .filter_map(|b| {
660 b.canary_of.as_ref().map(|primary_name| {
661 let primary_weight = config
663 .backends
664 .iter()
665 .find(|p| p.name == *primary_name)
666 .map(|p| p.weight)
667 .unwrap_or(100);
668 (
669 primary_name.clone(),
670 (b.name.clone(), primary_weight, b.weight),
671 )
672 })
673 })
674 .collect();
675
676 if !canary_mappings.is_empty() {
677 for (primary, (canary, pw, cw)) in &canary_mappings {
678 tracing::info!(
679 primary = %primary,
680 canary = %canary,
681 primary_weight = pw,
682 canary_weight = cw,
683 "Enabling canary routing"
684 );
685 }
686 service = BoxCloneService::new(crate::canary::CanaryService::new(
687 service,
688 canary_mappings,
689 &config.proxy.separator,
690 ));
691 }
692
693 let mut failover_groups: std::collections::HashMap<String, Vec<(u32, String)>> =
696 std::collections::HashMap::new();
697 for b in &config.backends {
698 if let Some(ref primary) = b.failover_for {
699 failover_groups
700 .entry(primary.clone())
701 .or_default()
702 .push((b.priority, b.name.clone()));
703 }
704 }
705 let failover_mappings: std::collections::HashMap<String, Vec<String>> = failover_groups
707 .into_iter()
708 .map(|(primary, mut backends)| {
709 backends.sort_by_key(|(priority, _)| *priority);
710 let names: Vec<String> = backends.into_iter().map(|(_, name)| name).collect();
711 (primary, names)
712 })
713 .collect();
714
715 if !failover_mappings.is_empty() {
716 for (primary, failovers) in &failover_mappings {
717 tracing::info!(
718 primary = %primary,
719 failovers = ?failovers,
720 "Enabling failover routing"
721 );
722 }
723 service = BoxCloneService::new(crate::failover::FailoverService::new(
724 service,
725 failover_mappings,
726 &config.proxy.separator,
727 ));
728 }
729
730 let mirror_mappings: std::collections::HashMap<String, (String, u32)> = config
732 .backends
733 .iter()
734 .filter_map(|b| {
735 b.mirror_of
736 .as_ref()
737 .map(|source| (source.clone(), (b.name.clone(), b.mirror_percent)))
738 })
739 .collect();
740
741 if !mirror_mappings.is_empty() {
742 for (source, (mirror, pct)) in &mirror_mappings {
743 tracing::info!(
744 source = %source,
745 mirror = %mirror,
746 percent = pct,
747 "Enabling traffic mirroring"
748 );
749 }
750 service = BoxCloneService::new(crate::mirror::MirrorService::new(
751 service,
752 mirror_mappings,
753 &config.proxy.separator,
754 ));
755 }
756
757 let cache_configs: Vec<_> = config
759 .backends
760 .iter()
761 .filter_map(|b| {
762 b.cache
763 .as_ref()
764 .map(|c| (format!("{}{}", b.name, config.proxy.separator), c))
765 })
766 .collect();
767
768 if !cache_configs.is_empty() {
769 for (ns, cfg) in &cache_configs {
770 tracing::info!(
771 backend = %ns.trim_end_matches(&config.proxy.separator),
772 resource_ttl = cfg.resource_ttl_seconds,
773 tool_ttl = cfg.tool_ttl_seconds,
774 max_entries = cfg.max_entries,
775 "Applying response cache"
776 );
777 }
778 let (cache_svc, handle) = cache::CacheService::new(service, cache_configs, &config.cache);
779 service = BoxCloneService::new(cache_svc);
780 cache_handle = Some(handle);
781 }
782
783 if config.performance.coalesce_requests {
785 tracing::info!("Request coalescing enabled");
786 service = BoxCloneService::new(coalesce::CoalesceService::new(service));
787 }
788
789 if config.security.max_argument_size.is_some() {
791 let validation = ValidationConfig {
792 max_argument_size: config.security.max_argument_size,
793 };
794 if let Some(max) = validation.max_argument_size {
795 tracing::info!(max_argument_size = max, "Applying request validation");
796 }
797 service = BoxCloneService::new(ValidationService::new(service, validation));
798 }
799
800 let filters: Vec<_> = config
802 .backends
803 .iter()
804 .filter_map(|b| b.build_filter(&config.proxy.separator).transpose())
805 .collect::<anyhow::Result<Vec<_>>>()?;
806
807 if !filters.is_empty() {
808 for f in &filters {
809 tracing::info!(
810 backend = %f.namespace.trim_end_matches(&config.proxy.separator),
811 tool_filter = ?f.tool_filter,
812 resource_filter = ?f.resource_filter,
813 prompt_filter = ?f.prompt_filter,
814 "Applying capability filter"
815 );
816 }
817 service = BoxCloneService::new(CapabilityFilterService::new(service, filters));
818 }
819
820 if config.proxy.tool_exposure == crate::config::ToolExposure::Search {
822 let prefix = format!("proxy{}", config.proxy.separator);
823 tracing::info!(
824 prefix = %prefix,
825 "Search mode: ListTools will only show proxy/ namespace tools"
826 );
827 service =
828 BoxCloneService::new(crate::filter::SearchModeFilterService::new(service, prefix));
829 }
830
831 let alias_mappings: Vec<_> = config
833 .backends
834 .iter()
835 .flat_map(|b| {
836 let ns = format!("{}{}", b.name, config.proxy.separator);
837 b.aliases
838 .iter()
839 .map(move |a| (ns.clone(), a.from.clone(), a.to.clone()))
840 })
841 .collect();
842
843 if let Some(alias_map) = alias::AliasMap::new(alias_mappings) {
844 let count = alias_map.forward.len();
845 tracing::info!(aliases = count, "Applying tool aliases");
846 service = BoxCloneService::new(alias::AliasService::new(service, alias_map));
847 }
848
849 if !config.composite_tools.is_empty() {
851 let count = config.composite_tools.len();
852 tracing::info!(composite_tools = count, "Applying composite tool fan-out");
853 service = BoxCloneService::new(crate::composite::CompositeService::new(
854 service,
855 config.composite_tools.clone(),
856 ));
857 }
858
859 #[cfg(feature = "oauth")]
861 if matches!(
862 &config.auth,
863 Some(AuthConfig::Bearer {
864 scoped_tokens,
865 ..
866 }) if !scoped_tokens.is_empty()
867 ) {
868 tracing::info!("Enabling bearer token scoping middleware");
869 service = BoxCloneService::new(crate::bearer_scope::BearerScopingService::new(service));
870 }
871
872 #[cfg(feature = "oauth")]
874 {
875 let rbac_config = match &config.auth {
876 Some(
877 AuthConfig::Jwt {
878 roles,
879 role_mapping: Some(mapping),
880 ..
881 }
882 | AuthConfig::OAuth {
883 roles,
884 role_mapping: Some(mapping),
885 ..
886 },
887 ) if !roles.is_empty() => {
888 tracing::info!(
889 roles = roles.len(),
890 claim = %mapping.claim,
891 "Enabling RBAC"
892 );
893 Some(RbacConfig::new(roles, mapping))
894 }
895 _ => None,
896 };
897
898 if let Some(rbac) = rbac_config {
899 service = BoxCloneService::new(RbacService::new(service, rbac));
900 }
901
902 let required_scopes: &[String] = match &config.auth {
909 Some(AuthConfig::OAuth {
910 required_scopes, ..
911 }) => required_scopes,
912 _ => &[],
913 };
914 if let Some(layer) = oauth_scope_layer(required_scopes) {
915 tracing::info!(
916 scopes = ?required_scopes,
917 "Enabling OAuth required_scopes enforcement"
918 );
919 service = BoxCloneService::new(tower::Layer::layer(&layer, service));
920 }
921
922 let forward_namespaces: std::collections::HashSet<String> = config
924 .backends
925 .iter()
926 .filter(|b| b.forward_auth)
927 .map(|b| format!("{}{}", b.name, config.proxy.separator))
928 .collect();
929
930 if !forward_namespaces.is_empty() {
931 tracing::info!(
932 backends = ?forward_namespaces,
933 "Enabling token passthrough for forward_auth backends"
934 );
935 service = BoxCloneService::new(crate::token::TokenPassthroughService::new(
936 service,
937 forward_namespaces,
938 ));
939 }
940 }
941
942 #[cfg(feature = "metrics")]
944 if config.observability.metrics.enabled {
945 service = BoxCloneService::new(crate::metrics::MetricsService::new(service));
946 }
947
948 if config.observability.access_log.enabled {
950 tracing::info!("Access logging enabled (target: mcp::access)");
951 service = BoxCloneService::new(crate::access_log::AccessLogService::new(
952 service,
953 &config.proxy.separator,
954 ));
955 }
956
957 if config.observability.audit {
959 tracing::info!("Audit logging enabled (target: mcp::audit)");
960 let audited = tower::Layer::layer(&tower_mcp::AuditLayer::new(), service);
961 service = BoxCloneService::new(tower_mcp::CatchError::new(audited));
962 }
963
964 if let Some(ref rl) = config.proxy.rate_limit {
966 tracing::info!(
967 requests = rl.requests,
968 period_seconds = rl.period_seconds,
969 "Applying global rate limit"
970 );
971 let layer = tower_resilience::ratelimiter::RateLimiterLayer::builder()
972 .limit_for_period(rl.requests)
973 .refresh_period(Duration::from_secs(rl.period_seconds))
974 .name("global-ratelimit")
975 .build();
976 let limited = tower::Layer::layer(&layer, service);
977 service = BoxCloneService::new(tower_mcp::CatchError::new(limited));
978 }
979
980 Ok((service, cache_handle))
981}
982
983async fn apply_auth(config: &ProxyConfig, router: Router) -> Result<Router> {
985 let router = if let Some(auth) = &config.auth {
986 match auth {
987 AuthConfig::Bearer {
988 tokens,
989 scoped_tokens,
990 } => {
991 let total = tokens.len() + scoped_tokens.len();
992 if scoped_tokens.is_empty() {
993 tracing::info!(token_count = total, "Enabling bearer token auth");
995 let validator = StaticBearerValidator::new(tokens.iter().cloned());
996 let layer = AuthLayer::new(validator);
997 router.layer(layer)
998 } else {
999 #[cfg(feature = "oauth")]
1001 {
1002 tracing::info!(
1003 token_count = total,
1004 scoped = scoped_tokens.len(),
1005 "Enabling bearer token auth with per-token scoping"
1006 );
1007 let layer =
1008 crate::bearer_scope::ScopedBearerAuthLayer::new(tokens, scoped_tokens);
1009 router.layer(layer)
1010 }
1011 #[cfg(not(feature = "oauth"))]
1012 {
1013 anyhow::bail!(
1014 "Per-token tool scoping requires the 'oauth' feature. \
1015 Rebuild with: cargo install mcp-proxy --features oauth"
1016 );
1017 }
1018 }
1019 }
1020 #[cfg(feature = "oauth")]
1021 AuthConfig::Jwt {
1022 issuer,
1023 audience,
1024 jwks_uri,
1025 ..
1026 } => {
1027 tracing::info!(
1028 issuer = %issuer,
1029 audience = %audience,
1030 jwks_uri = %jwks_uri,
1031 "Enabling JWT auth (JWKS)"
1032 );
1033 let validator = tower_mcp::oauth::JwksValidator::builder(jwks_uri)
1034 .expected_audience(audience)
1035 .expected_issuer(issuer)
1036 .build()
1037 .await
1038 .context("building JWKS validator")?;
1039
1040 let addr = format!(
1041 "http://{}:{}",
1042 config.proxy.listen.host, config.proxy.listen.port
1043 );
1044 let metadata = tower_mcp::oauth::ProtectedResourceMetadata::new(&addr)
1045 .authorization_server(issuer);
1046
1047 let layer = tower_mcp::oauth::OAuthLayer::new(validator, metadata);
1048 router.layer(layer)
1049 }
1050 #[cfg(not(feature = "oauth"))]
1051 AuthConfig::Jwt { .. } => {
1052 anyhow::bail!(
1053 "JWT auth requires the 'oauth' feature. Rebuild with: cargo install mcp-proxy --features oauth"
1054 );
1055 }
1056 #[cfg(feature = "oauth")]
1057 AuthConfig::OAuth {
1058 issuer,
1059 audience,
1060 token_validation,
1061 jwks_uri,
1062 introspection_endpoint,
1063 client_id,
1064 client_secret,
1065 ..
1066 } => {
1067 use crate::config::TokenValidationStrategy;
1068
1069 tracing::info!(
1070 issuer = %issuer,
1071 audience = %audience,
1072 strategy = ?token_validation,
1073 "Enabling OAuth 2.1 auth"
1074 );
1075
1076 let discovered = crate::introspection::discover_auth_server(issuer)
1078 .await
1079 .context("discovering OAuth authorization server")?;
1080
1081 let effective_jwks_uri = jwks_uri
1082 .as_deref()
1083 .or(discovered.jwks_uri.as_deref())
1084 .ok_or_else(|| {
1085 anyhow::anyhow!(
1086 "JWKS URI not found via discovery and not configured manually"
1087 )
1088 })?;
1089
1090 let effective_introspection = introspection_endpoint
1091 .as_deref()
1092 .or(discovered.introspection_endpoint.as_deref());
1093
1094 let addr = format!(
1095 "http://{}:{}",
1096 config.proxy.listen.host, config.proxy.listen.port
1097 );
1098 let metadata = tower_mcp::oauth::ProtectedResourceMetadata::new(&addr)
1099 .authorization_server(issuer);
1100
1101 match token_validation {
1102 TokenValidationStrategy::Jwt => {
1103 let validator =
1104 tower_mcp::oauth::JwksValidator::builder(effective_jwks_uri)
1105 .expected_audience(audience)
1106 .expected_issuer(issuer)
1107 .build()
1108 .await
1109 .context("building JWKS validator")?;
1110 let layer = tower_mcp::oauth::OAuthLayer::new(validator, metadata);
1111 router.layer(layer)
1112 }
1113 TokenValidationStrategy::Introspection => {
1114 let endpoint = effective_introspection.ok_or_else(|| {
1115 anyhow::anyhow!(
1116 "introspection endpoint not found via discovery and not configured"
1117 )
1118 })?;
1119 let validator = crate::introspection::IntrospectionValidator::new(
1120 endpoint,
1121 client_id.as_deref().unwrap(),
1122 client_secret.as_deref().unwrap(),
1123 )
1124 .expected_audience(audience);
1125 let layer = tower_mcp::oauth::OAuthLayer::new(validator, metadata);
1126 router.layer(layer)
1127 }
1128 TokenValidationStrategy::Both => {
1129 let endpoint = effective_introspection.ok_or_else(|| {
1130 anyhow::anyhow!(
1131 "introspection endpoint not found via discovery and not configured"
1132 )
1133 })?;
1134 let jwt_validator =
1135 tower_mcp::oauth::JwksValidator::builder(effective_jwks_uri)
1136 .expected_audience(audience)
1137 .expected_issuer(issuer)
1138 .build()
1139 .await
1140 .context("building JWKS validator")?;
1141 let introspection_validator =
1142 crate::introspection::IntrospectionValidator::new(
1143 endpoint,
1144 client_id.as_deref().unwrap(),
1145 client_secret.as_deref().unwrap(),
1146 )
1147 .expected_audience(audience);
1148 let fallback = crate::introspection::FallbackValidator::new(
1149 jwt_validator,
1150 introspection_validator,
1151 );
1152 let layer = tower_mcp::oauth::OAuthLayer::new(fallback, metadata);
1153 router.layer(layer)
1154 }
1155 }
1156 }
1157 #[cfg(not(feature = "oauth"))]
1158 AuthConfig::OAuth { .. } => {
1159 anyhow::bail!(
1160 "OAuth auth requires the 'oauth' feature. Rebuild with: cargo install mcp-proxy --features oauth"
1161 );
1162 }
1163 }
1164 } else {
1165 router
1166 };
1167 Ok(router)
1168}
1169
1170pub async fn shutdown_signal(timeout: Duration) {
1172 let ctrl_c = tokio::signal::ctrl_c();
1173 #[cfg(unix)]
1174 {
1175 let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
1176 .expect("SIGTERM handler");
1177 tokio::select! {
1178 _ = ctrl_c => {},
1179 _ = sigterm.recv() => {},
1180 }
1181 }
1182 #[cfg(not(unix))]
1183 {
1184 ctrl_c.await.ok();
1185 }
1186 tracing::info!(
1187 timeout_seconds = timeout.as_secs(),
1188 "Shutdown signal received, draining connections"
1189 );
1190}
1191
1192#[cfg(all(test, feature = "oauth"))]
1193mod scope_enforcement_tests {
1194 use std::collections::HashMap;
1195
1196 use tower::{Layer, Service};
1197 use tower_mcp::oauth::token::TokenClaims;
1198 use tower_mcp::protocol::{CallToolParams, McpRequest, RequestId};
1199 use tower_mcp::router::Extensions;
1200
1201 use super::oauth_scope_layer;
1202 use crate::test_util::MockService;
1203
1204 fn call_with_scope(scope: Option<&str>) -> tower_mcp::RouterRequest {
1206 let mut extensions = Extensions::new();
1207 if let Some(scope) = scope {
1208 extensions.insert(TokenClaims {
1209 sub: Some("user".into()),
1210 iss: None,
1211 aud: None,
1212 exp: None,
1213 scope: Some(scope.to_string()),
1214 client_id: None,
1215 extra: HashMap::new(),
1216 });
1217 }
1218 tower_mcp::RouterRequest {
1219 id: RequestId::Number(1),
1220 inner: McpRequest::CallTool(CallToolParams {
1221 name: "fs/read".into(),
1222 arguments: serde_json::json!({}),
1223 input_responses: None,
1224 request_state: None,
1225 meta: None,
1226 task: None,
1227 }),
1228 extensions,
1229 }
1230 }
1231
1232 #[test]
1233 fn no_required_scopes_yields_no_layer() {
1234 assert!(oauth_scope_layer(&[]).is_none());
1235 }
1236
1237 #[tokio::test]
1238 async fn token_missing_required_scope_is_rejected() {
1239 let required = vec!["mcp:access".to_string()];
1240 let layer = oauth_scope_layer(&required).expect("layer for non-empty scopes");
1241 let mut svc = layer.layer(MockService::with_tools(&["fs/read"]));
1242
1243 let resp = svc
1245 .call(call_with_scope(Some("other:scope")))
1246 .await
1247 .unwrap();
1248 let err = resp.inner.unwrap_err();
1249 assert!(
1250 err.message.to_lowercase().contains("scope"),
1251 "expected insufficient-scope error, got: {}",
1252 err.message
1253 );
1254 }
1255
1256 #[tokio::test]
1257 async fn token_with_all_required_scopes_is_allowed() {
1258 let required = vec!["mcp:access".to_string(), "mcp:read".to_string()];
1259 let layer = oauth_scope_layer(&required).expect("layer for non-empty scopes");
1260 let mut svc = layer.layer(MockService::with_tools(&["fs/read"]));
1261
1262 let resp = svc
1263 .call(call_with_scope(Some("mcp:access mcp:read mcp:extra")))
1264 .await
1265 .unwrap();
1266 assert!(
1267 resp.inner.is_ok(),
1268 "token carrying all required scopes should be allowed"
1269 );
1270 }
1271}
1272
1273#[cfg(test)]
1274mod middleware_stack_tests {
1275 use std::pin::Pin;
1281 use std::sync::Arc;
1282 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
1283 use std::task::{Context, Poll};
1284
1285 use tower::Service;
1286 use tower_mcp::CallToolResult;
1287 use tower_mcp::protocol::{CallToolParams, McpRequest, McpResponse, RequestId};
1288 use tower_mcp::router::Extensions;
1289 use tower_mcp_types::JsonRpcError;
1290
1291 use super::*;
1292
1293 #[derive(Clone)]
1296 struct FlakyBackend {
1297 hits: Arc<AtomicUsize>,
1298 slow: Arc<AtomicBool>,
1299 fail_responses: Arc<AtomicUsize>,
1300 }
1301
1302 impl FlakyBackend {
1303 fn new() -> (Self, Arc<AtomicUsize>, Arc<AtomicBool>, Arc<AtomicUsize>) {
1304 let hits = Arc::new(AtomicUsize::new(0));
1305 let slow = Arc::new(AtomicBool::new(false));
1306 let fail_responses = Arc::new(AtomicUsize::new(0));
1307 (
1308 Self {
1309 hits: hits.clone(),
1310 slow: slow.clone(),
1311 fail_responses: fail_responses.clone(),
1312 },
1313 hits,
1314 slow,
1315 fail_responses,
1316 )
1317 }
1318 }
1319
1320 impl tower::Service<RouterRequest> for FlakyBackend {
1321 type Response = RouterResponse;
1322 type Error = Infallible;
1323 type Future = Pin<Box<dyn Future<Output = Result<RouterResponse, Infallible>> + Send>>;
1324
1325 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
1326 Poll::Ready(Ok(()))
1327 }
1328
1329 fn call(&mut self, req: RouterRequest) -> Self::Future {
1330 let hits = self.hits.clone();
1331 let slow = self.slow.clone();
1332 let fail_responses = self.fail_responses.clone();
1333 Box::pin(async move {
1334 hits.fetch_add(1, Ordering::SeqCst);
1335 if slow.load(Ordering::SeqCst) {
1336 tokio::time::sleep(Duration::from_millis(100)).await;
1337 }
1338 let inner = if fail_responses
1339 .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
1340 .is_ok()
1341 {
1342 Err(JsonRpcError::internal_error("transient backend failure"))
1343 } else {
1344 Ok(McpResponse::CallTool(CallToolResult::text("pong")))
1345 };
1346 Ok(RouterResponse { id: req.id, inner })
1347 })
1348 }
1349 }
1350
1351 fn req() -> RouterRequest {
1352 RouterRequest {
1353 id: RequestId::Number(1),
1354 inner: McpRequest::CallTool(CallToolParams {
1355 name: "ping".to_string(),
1356 arguments: serde_json::json!({}),
1357 input_responses: None,
1358 request_state: None,
1359 meta: None,
1360 task: None,
1361 }),
1362 extensions: Extensions::new(),
1363 }
1364 }
1365
1366 async fn drive(svc: &mut InfallibleService) -> RouterResponse {
1367 svc.ready()
1368 .await
1369 .expect("infallible")
1370 .call(req())
1371 .await
1372 .expect("infallible")
1373 }
1374
1375 fn test_breaker(
1376 minimum_calls: usize,
1377 wait_in_open: Duration,
1378 permitted_in_half_open: usize,
1379 ) -> tower_resilience::circuitbreaker::CircuitBreakerLayer {
1380 let (layer, _handle) = tower_resilience::circuitbreaker::CircuitBreakerLayer::builder()
1381 .failure_rate_threshold(0.5)
1382 .minimum_number_of_calls(minimum_calls)
1383 .sliding_window_size(minimum_calls)
1387 .wait_duration_in_open(wait_in_open)
1388 .permitted_calls_in_half_open(permitted_in_half_open)
1389 .name("test-cb")
1390 .build_with_handle();
1391 layer
1392 }
1393
1394 #[tokio::test]
1398 async fn timeout_and_breaker_compose_and_breaker_trips() {
1399 let (backend, hits, slow, _fail) = FlakyBackend::new();
1400 slow.store(true, Ordering::SeqCst);
1401
1402 let mut mw = BackendMiddlewareLayer::default();
1403 mw.stage(TimeoutLayer::new(Duration::from_millis(10)));
1404 mw.stage(test_breaker(4, Duration::from_millis(200), 2));
1405
1406 let mut svc = mw.layer(backend);
1407
1408 for i in 1..=4 {
1409 let resp = drive(&mut svc).await;
1410 assert!(resp.inner.is_err(), "call {i} should time out");
1411 }
1412 assert_eq!(
1413 hits.load(Ordering::SeqCst),
1414 4,
1415 "all four calls reach the backend"
1416 );
1417
1418 let resp = drive(&mut svc).await;
1419 assert!(resp.inner.is_err(), "open breaker rejects");
1420 assert_eq!(
1421 hits.load(Ordering::SeqCst),
1422 4,
1423 "rejected call must not reach the backend"
1424 );
1425
1426 tokio::time::sleep(Duration::from_millis(300)).await;
1427 slow.store(false, Ordering::SeqCst);
1428 for i in 1..=3 {
1429 let resp = drive(&mut svc).await;
1430 assert!(resp.inner.is_ok(), "recovered call {i} succeeds");
1431 }
1432 assert_eq!(hits.load(Ordering::SeqCst), 7);
1433 }
1434
1435 #[tokio::test]
1439 async fn retry_and_timeout_compose() {
1440 let (backend, hits, _slow, fail) = FlakyBackend::new();
1441 fail.store(2, Ordering::SeqCst);
1442
1443 let cfg = crate::config::RetryConfig {
1444 max_retries: 3,
1445 initial_backoff_ms: 1,
1446 max_backoff_ms: 5,
1447 budget_percent: None,
1448 min_retries_per_sec: 10,
1449 };
1450 let mut mw = BackendMiddlewareLayer {
1451 retry: Some(crate::retry::build_retry_layer(&cfg, "test")),
1452 ..Default::default()
1453 };
1454 mw.stage(TimeoutLayer::new(Duration::from_millis(500)));
1455
1456 let mut svc = mw.layer(backend);
1457
1458 let resp = drive(&mut svc).await;
1459 assert!(resp.inner.is_ok(), "retries absorb the transient failures");
1460 assert_eq!(
1461 hits.load(Ordering::SeqCst),
1462 3,
1463 "two failures plus the successful attempt"
1464 );
1465 }
1466
1467 #[tokio::test]
1469 async fn empty_middleware_layer_is_identity() {
1470 let (backend, hits, _slow, _fail) = FlakyBackend::new();
1471 let mw = BackendMiddlewareLayer::default();
1472 assert!(mw.is_empty());
1473
1474 let mut svc = mw.layer(backend);
1475 let resp = drive(&mut svc).await;
1476 assert!(resp.inner.is_ok());
1477 assert_eq!(hits.load(Ordering::SeqCst), 1);
1478 }
1479
1480 #[tokio::test]
1484 async fn config_mapped_breaker_opens_at_minimum_calls() {
1485 let (backend, hits, slow, _fail) = FlakyBackend::new();
1486 slow.store(true, Ordering::SeqCst);
1487
1488 let cfg = crate::config::CircuitBreakerConfig {
1489 failure_rate_threshold: 0.5,
1490 minimum_calls: 4,
1491 wait_duration_seconds: 60,
1492 permitted_calls_in_half_open: 1,
1493 };
1494 let (breaker, _handle) = build_breaker_layer(&cfg, "test");
1495
1496 let mut mw = BackendMiddlewareLayer::default();
1497 mw.stage(TimeoutLayer::new(Duration::from_millis(10)));
1498 mw.stage(breaker);
1499
1500 let mut svc = mw.layer(backend);
1501
1502 for i in 1..=4 {
1503 let resp = drive(&mut svc).await;
1504 assert!(resp.inner.is_err(), "call {i} should time out");
1505 }
1506 assert_eq!(hits.load(Ordering::SeqCst), 4);
1507
1508 let resp = drive(&mut svc).await;
1509 assert!(resp.inner.is_err(), "open breaker rejects");
1510 assert_eq!(
1511 hits.load(Ordering::SeqCst),
1512 4,
1513 "rejected call must not reach the backend"
1514 );
1515 }
1516}