diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index ab541ff31..b44960581 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -1338,10 +1338,6 @@ pub struct StreamableHttpService { /// than racing to replay the initialize handshake. `None` when no external /// session store is configured (avoids allocating the map). pending_restores: Option, - /// Caches tool input schemas by name for SEP-2243 `Mcp-Param-*` validation. - /// Populated lazily via `get_tool` so the service factory runs at most once - /// per tool name. `None` value means the tool exposes no schema. - tool_schemas: Arc>>>>, } impl Clone for StreamableHttpService { @@ -1351,7 +1347,6 @@ impl Clone for StreamableHttpService { session_manager: self.session_manager.clone(), service_factory: self.service_factory.clone(), pending_restores: self.pending_restores.clone(), - tool_schemas: self.tool_schemas.clone(), } } } @@ -1430,7 +1425,6 @@ where session_manager, service_factory: Arc::new(service_factory), pending_restores, - tool_schemas: Arc::new(std::sync::RwLock::new(HashMap::new())), } } fn get_service(&self) -> Result { @@ -1599,24 +1593,12 @@ where Ok(self.stateless_sse_response(Some(first), receiver, request_ct)) } - /// Returns the cached input schema for `name`, constructing a service once - /// per name to read its `ServerHandler::get_tool` definition. Used to - /// validate SEP-2243 `Mcp-Param-*` headers against the request body. + /// Looks up `name`'s input schema for SEP-2243 `Mcp-Param-*` validation. fn tool_schema(&self, name: &str) -> Option> { - if let Ok(cache) = self.tool_schemas.read() - && let Some(schema) = cache.get(name) - { - return schema.clone(); - } - let schema = self - .get_service() - .ok() - .and_then(|service| service.get_tool(name)) - .map(|tool| tool.input_schema); - if let Ok(mut cache) = self.tool_schemas.write() { - cache.insert(name.to_owned(), schema.clone()); - } - schema + self.get_service() + .ok()? + .get_tool(name) + .map(|tool| tool.input_schema) } /// Spawn a task that runs `serve_server` for the given session, waits for @@ -2269,13 +2251,15 @@ where validate_protocol_version_header(&part.headers, has_per_request_version)?; } } - // Validate SEP-2243 standard headers against the body - validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?; - validate_request_protocol_version_meta(&part.headers, &message)?; - validate_required_protocol_meta(&self.config, &message)?; let service = self .get_service() .map_err(internal_error_response("get service"))?; + // Validate SEP-2243 standard headers against the body + validate_standard_headers(&part.headers, &message, |name| { + service.get_tool(name).map(|tool| tool.input_schema) + })?; + validate_request_protocol_version_meta(&part.headers, &message)?; + validate_required_protocol_meta(&self.config, &message)?; match message { ClientJsonRpcMessage::Request(mut request) => { let negotiates_per_request = has_per_request_version diff --git a/crates/rmcp/tests/test_streamable_http_standard_headers.rs b/crates/rmcp/tests/test_streamable_http_standard_headers.rs index 5bbfedd64..e531ac00f 100644 --- a/crates/rmcp/tests/test_streamable_http_standard_headers.rs +++ b/crates/rmcp/tests/test_streamable_http_standard_headers.rs @@ -1,6 +1,9 @@ #![cfg(not(feature = "local"))] //! SEP-2243 server-side validation of `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers. -use std::sync::Arc; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; use rmcp::{ ServerHandler, @@ -36,14 +39,23 @@ impl ServerHandler for HeaderValidationServer { } async fn spawn_server() -> (reqwest::Client, String, CancellationToken) { + spawn_server_with(|| Ok(HeaderValidationServer)).await +} + +async fn spawn_server_with( + service_factory: impl Fn() -> Result + Send + Sync + 'static, +) -> (reqwest::Client, String, CancellationToken) +where + S: ServerHandler + Send + 'static, +{ let config = StreamableHttpServerConfig::default() .with_legacy_session_mode(false) .with_json_response(true) .with_sse_keep_alive(None) .with_cancellation_token(CancellationToken::new()); let ct = config.cancellation_token.clone(); - let service: StreamableHttpService = - StreamableHttpService::new(|| Ok(HeaderValidationServer), Default::default(), config); + let service: StreamableHttpService = + StreamableHttpService::new(service_factory, Default::default(), config); let router = axum::Router::new().nest_service("/mcp", service); let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -314,6 +326,90 @@ async fn rejects_missing_param_header_with_32020() -> anyhow::Result<()> { Ok(()) } +/// Server whose `deploy` tool gains or loses the `x-mcp-header` annotation at runtime. +#[derive(Clone, Default)] +struct ChangingToolServer { + annotated: Arc, +} + +async fn spawn_changing_tool_server( + server: &ChangingToolServer, +) -> (reqwest::Client, String, CancellationToken) { + let server = server.clone(); + spawn_server_with(move || Ok(server.clone())).await +} + +/// POSTs a `deploy` call whose `Mcp-Param-Region` header contradicts the body. +async fn post_mismatched_region(client: &reqwest::Client, url: &str) -> reqwest::Response { + post_tool_call( + client, + url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("eu-central1"), + ) + .await +} + +impl ServerHandler for ChangingToolServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + fn get_tool(&self, name: &str) -> Option { + if !self.annotated.load(Ordering::SeqCst) { + return None; + } + HeaderValidationServer.get_tool(name) + } +} + +#[tokio::test] +async fn validates_param_header_once_tool_gains_annotation() -> anyhow::Result<()> { + let server = ChangingToolServer::default(); + let (client, url, ct) = spawn_changing_tool_server(&server).await; + + // Header validation passes, so the request fails later on its missing `_meta`. + let response = post_mismatched_region(&client, &url).await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32602); + + server.annotated.store(true, Ordering::SeqCst); + let response = post_mismatched_region(&client, &url).await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn skips_param_header_validation_once_tool_loses_annotation() -> anyhow::Result<()> { + let server = ChangingToolServer { + annotated: Arc::new(AtomicBool::new(true)), + }; + let (client, url, ct) = spawn_changing_tool_server(&server).await; + + let response = post_mismatched_region(&client, &url).await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + server.annotated.store(false, Ordering::SeqCst); + let response = post_mismatched_region(&client, &url).await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32602); + + ct.cancel(); + Ok(()) +} + /// Spawns the server in legacy session mode: the initialize handshake creates /// a real session, exercising the path where the validator used to be bypassed. async fn spawn_legacy_server() -> (reqwest::Client, String, CancellationToken) {