feat!: validate Mcp-Param-* headers on server

Author: DaleSeo

crates/rmcp/src/transport/streamable_http_server/tower.rs

@@ -18,7 +18,7 @@ use crate::{
     model::{
         ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData,
         GetExtensions, Implementation, InitializeRequest, InitializeRequestParams,
-        InitializedNotification, JsonRpcError, ProtocolVersion, RequestId,
+        InitializedNotification, JsonObject, JsonRpcError, ProtocolVersion, RequestId,
     },
     serve_server,
     service::serve_directly,
@@ -274,19 +274,20 @@ fn header_mismatch_jsonrpc_response(
         .expect("valid response")
 }
 
-/// Validates SEP-2243 `Mcp-Method` / `Mcp-Name` headers against the request body.
+/// Validates SEP-2243 `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers against the body.
 ///
 /// Only enforced when the request declares a protocol version `>= STANDARD_HEADERS`.
 /// The `initialize` handshake is exempt: clients emit these headers only after the
-/// version has been negotiated. `Mcp-Param-*` validation is not enforced here because
-/// the transport layer has no synchronous access to tool input schemas.
+/// version has been negotiated. `tool_schema` supplies the called tool's input schema
+/// so annotated `Mcp-Param-*` headers can be checked (no schema => those are skipped).
 #[expect(
     clippy::result_large_err,
     reason = "BoxResponse is intentionally large; matches other handlers in this file"
 )]
 fn validate_standard_headers(
     headers: &HeaderMap,
     message: &ClientJsonRpcMessage,
+    tool_schema: impl Fn(&str) -> Option<Arc<JsonObject>>,
 ) -> Result<(), BoxResponse> {
     let version_requires_headers = headers
         .get(HEADER_MCP_PROTOCOL_VERSION)
@@ -310,7 +311,16 @@ fn validate_standard_headers(
     let Ok(value) = serde_json::to_value(message) else {
         return Ok(());
     };
-    if let Err(reason) = mcp_headers::validate_request_headers(headers, &value, None) {
+    // For tools/call, look up the tool schema so Mcp-Param-* headers are validated.
+    let schema = value
+        .get("method")
+        .and_then(|method| method.as_str())
+        .filter(|method| *method == "tools/call")
+        .and_then(|_| value.get("params"))
+        .and_then(|params| params.get("name"))
+        .and_then(|name| name.as_str())
+        .and_then(tool_schema);
+    if let Err(reason) = mcp_headers::validate_request_headers(headers, &value, schema.as_deref()) {
         return Err(header_mismatch_jsonrpc_response(request_id, reason));
     }
     Ok(())
@@ -611,6 +621,10 @@ pub struct StreamableHttpService<S, M> {
     pending_restores: Option<
         Arc<tokio::sync::RwLock<HashMap<SessionId, tokio::sync::watch::Sender<Option<bool>>>>>,
     >,
+    /// 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<std::sync::RwLock<HashMap<String, Option<Arc<JsonObject>>>>>,
 }
 
 impl<S, M> Clone for StreamableHttpService<S, M> {
@@ -620,14 +634,15 @@ impl<S, M> Clone for StreamableHttpService<S, M> {
             session_manager: self.session_manager.clone(),
             service_factory: self.service_factory.clone(),
             pending_restores: self.pending_restores.clone(),
+            tool_schemas: self.tool_schemas.clone(),
         }
     }
 }
 
 impl<RequestBody, S, M> tower_service::Service<Request<RequestBody>> for StreamableHttpService<S, M>
 where
     RequestBody: Body + Send + 'static,
-    S: crate::Service<RoleServer> + Send + 'static,
+    S: crate::ServerHandler + Send + 'static,
     M: SessionManager,
     RequestBody::Error: Display,
     RequestBody::Data: Send + 'static,
@@ -681,7 +696,7 @@ impl Drop for PendingRestoreGuard {
 
 impl<S, M> StreamableHttpService<S, M>
 where
-    S: crate::Service<RoleServer> + Send + 'static,
+    S: crate::ServerHandler + Send + 'static,
     M: SessionManager,
 {
     pub fn new(
@@ -700,12 +715,33 @@ 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<S, std::io::Error> {
         (self.service_factory)()
     }
 
+    /// 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.
+    fn tool_schema(&self, name: &str) -> Option<Arc<JsonObject>> {
+        if let Ok(cache) = self.tool_schemas.read() {
+            if 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
+    }
+
     /// Spawn a task that runs `serve_server` for the given session, waits for
     /// it to finish, and then calls `close_session`.
     ///
@@ -720,7 +756,7 @@ where
         transport: M::Transport,
         init_done_tx: Option<tokio::sync::oneshot::Sender<()>>,
     ) where
-        S: crate::Service<RoleServer> + Send + 'static,
+        S: crate::ServerHandler + Send + 'static,
         M: SessionManager,
     {
         tokio::spawn(async move {
@@ -763,7 +799,7 @@ where
         parts: &http::request::Parts,
     ) -> Result<bool, std::io::Error>
     where
-        S: crate::Service<RoleServer> + Send + 'static,
+        S: crate::ServerHandler + Send + 'static,
         M: SessionManager,
     {
         // Both fields are Some iff a session store is configured.
@@ -1140,7 +1176,7 @@ where
                 // Validate MCP-Protocol-Version header (per 2025-06-18 spec)
                 validate_protocol_version_header(&part.headers)?;
                 // Validate SEP-2243 standard headers against the body
-                validate_standard_headers(&part.headers, &message)?;
+                validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?;
 
                 // inject request part to extensions
                 match &mut message {
@@ -1294,7 +1330,7 @@ where
                 }
             }
             // Validate SEP-2243 standard headers against the body
-            validate_standard_headers(&part.headers, &message)?;
+            validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?;
             let service = self
                 .get_service()
                 .map_err(internal_error_response("get service"))?;

crates/rmcp/tests/test_streamable_http_standard_headers.rs

@@ -1,7 +1,13 @@
 #![cfg(not(feature = "local"))]
-//! SEP-2243 server-side validation of `Mcp-Method` / `Mcp-Name` headers.
-use rmcp::transport::streamable_http_server::{
-    StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
+//! SEP-2243 server-side validation of `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers.
+use std::sync::Arc;
+
+use rmcp::{
+    ServerHandler,
+    model::{ServerCapabilities, ServerInfo, Tool},
+    transport::streamable_http_server::{
+        StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
+    },
 };
 use tokio_util::sync::CancellationToken;
 
@@ -145,3 +151,115 @@ async fn skips_validation_for_pre_sep_version() -> anyhow::Result<()> {
     ct.cancel();
     Ok(())
 }
+
+/// Server exposing one tool whose `region` argument is promoted to `Mcp-Param-Region`.
+#[derive(Clone, Default)]
+struct ParamHeaderServer;
+
+impl ServerHandler for ParamHeaderServer {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
+    }
+
+    fn get_tool(&self, name: &str) -> Option<Tool> {
+        if name != "deploy" {
+            return None;
+        }
+        let schema = serde_json::json!({
+            "type": "object",
+            "properties": { "region": { "type": "string", "x-mcp-header": "Region" } }
+        });
+        let schema = schema.as_object().expect("object schema").clone();
+        Some(Tool::new("deploy", "deploy a thing", Arc::new(schema)))
+    }
+}
+
+async fn spawn_param_server() -> (reqwest::Client, String, CancellationToken) {
+    let config = StreamableHttpServerConfig::default()
+        .with_stateful_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<ParamHeaderServer, LocalSessionManager> =
+        StreamableHttpService::new(|| Ok(ParamHeaderServer), 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();
+    let addr = tcp_listener.local_addr().unwrap();
+    tokio::spawn({
+        let ct = ct.clone();
+        async move {
+            let _ = axum::serve(tcp_listener, router)
+                .with_graceful_shutdown(async move { ct.cancelled_owned().await })
+                .await;
+        }
+    });
+    (reqwest::Client::new(), format!("http://{addr}/mcp"), ct)
+}
+
+/// POSTs a `deploy` tools/call with a `region` argument and optional `Mcp-Param-Region` header.
+async fn post_deploy(
+    client: &reqwest::Client,
+    url: &str,
+    region_arg: &str,
+    param_region: Option<&str>,
+) -> reqwest::Response {
+    let body = format!(
+        r#"{{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{{"name":"deploy","arguments":{{"region":"{region_arg}"}}}}}}"#
+    );
+    let mut req = client
+        .post(url)
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("MCP-Protocol-Version", SEP_VERSION)
+        .header("Mcp-Method", "tools/call")
+        .header("Mcp-Name", "deploy")
+        .body(body);
+    if let Some(region) = param_region {
+        req = req.header("Mcp-Param-Region", region);
+    }
+    req.send().await.expect("send deploy request")
+}
+
+#[tokio::test]
+async fn accepts_matching_param_header() -> anyhow::Result<()> {
+    let (client, url, ct) = spawn_param_server().await;
+
+    let response = post_deploy(&client, &url, "us-west1", Some("us-west1")).await;
+    let body: serde_json::Value = response.json().await?;
+    assert_ne!(
+        body["error"]["code"], -32001,
+        "matching Mcp-Param-* must not be rejected, got: {body}"
+    );
+
+    ct.cancel();
+    Ok(())
+}
+
+#[tokio::test]
+async fn rejects_param_mismatch_with_32001() -> anyhow::Result<()> {
+    let (client, url, ct) = spawn_param_server().await;
+
+    let response = post_deploy(&client, &url, "us-west1", Some("eu-central1")).await;
+    assert_eq!(response.status(), 400);
+    let body: serde_json::Value = response.json().await?;
+    assert_eq!(body["error"]["code"], -32001);
+
+    ct.cancel();
+    Ok(())
+}
+
+#[tokio::test]
+async fn rejects_missing_param_header_with_32001() -> anyhow::Result<()> {
+    let (client, url, ct) = spawn_param_server().await;
+
+    // `region` argument is present but the annotated `Mcp-Param-Region` header is absent.
+    let response = post_deploy(&client, &url, "us-west1", None).await;
+    assert_eq!(response.status(), 400);
+    let body: serde_json::Value = response.json().await?;
+    assert_eq!(body["error"]["code"], -32001);
+
+    ct.cancel();
+    Ok(())
+}