portugalfuturista/rust-sdk
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: prevent streamable HTTP session leak (#934)

Browse source
This commit is contained in:
Dale Seo 2026-06-26 21:09:16 -04:00 committed by GitHub
parent e1af378949
commit dfa7fd6f93
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 68 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
crates/rmcp/src/transport/streamable_http_server/tower.rs
View file

@ -1124,44 +1124,40 @@ where
}
}
} else {
// Capture init params for external store persistence before
// extensions are injected (which would require Clone).
let stored_init_params = match &mut message {
ClientJsonRpcMessage::Request(req) => {
let ClientRequest::InitializeRequest(init_req) = &req.request else {
return Err(unexpected_message_response("initialize request"));
};
// Reject mismatched MCP-Protocol-Version header before binding the session to anything.
validate_header_matches_init_body(
&part.headers,
init_req.params.protocol_version.as_str(),
Some(req.id.clone()),
)?;
let stored_init_params = self
.config
.session_store
.as_ref()
.map(|_| init_req.params.clone());
// inject request part to extensions
req.request.extensions_mut().insert(part);
stored_init_params
}
_ => {
return Err(unexpected_message_response("initialize request"));
}
};
let service = self
.get_service()
.map_err(internal_error_response("get service"))?;
let (session_id, transport) = self
.session_manager
.create_session()
.await
.map_err(internal_error_response("create session"))?;
// Capture init params for external store persistence before
// extensions are injected (which would require Clone).
let stored_init_params = if self.config.session_store.is_some() {
if let ClientJsonRpcMessage::Request(req) = &message {
if let ClientRequest::InitializeRequest(init_req) = &req.request {
Some(init_req.params.clone())
} else {
None
}
} else {
None
}
} else {
None
};
if let ClientJsonRpcMessage::Request(req) = &mut message {
let ClientRequest::InitializeRequest(init_req) = &req.request else {
return Err(unexpected_message_response("initialize request"));
};
// Reject mismatched MCP-Protocol-Version header before binding the session to anything.
validate_header_matches_init_body(
&part.headers,
init_req.params.protocol_version.as_str(),
Some(req.id.clone()),
)?;
// inject request part to extensions
req.request.extensions_mut().insert(part);
} else {
return Err(unexpected_message_response("initialize request"));
}
let service = self
.get_service()
.map_err(internal_error_response("get service"))?;
// spawn a task to serve the session
Self::spawn_session_worker(
self.session_manager.clone(),

40
crates/rmcp/tests/test_streamable_http_protocol_version.rs
View file

@ -1,5 +1,7 @@
#![cfg(not(feature = "local"))]
//! Regression tests for the `MCP-Protocol-Version` header / initialize body consistency check.
use std::sync::Arc;
use rmcp::transport::streamable_http_server::{
StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
};
@ -16,10 +18,17 @@ fn init_body(body_version: &str) -> String {
async fn spawn_server(
config: StreamableHttpServerConfig,
) -> (reqwest::Client, String, CancellationToken) {
spawn_server_with_manager(config, Arc::new(LocalSessionManager::default())).await
}
async fn spawn_server_with_manager(
config: StreamableHttpServerConfig,
session_manager: Arc<LocalSessionManager>,
) -> (reqwest::Client, String, CancellationToken) {
let ct = config.cancellation_token.clone();
let service: StreamableHttpService<Calculator, LocalSessionManager> =
StreamableHttpService::new(|| Ok(Calculator::new()), Default::default(), config);
StreamableHttpService::new(|| Ok(Calculator::new()), session_manager, config);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@ -71,6 +80,17 @@ async fn post_init(
req.send().await.expect("send initialize request")
}
async fn post_non_initialize(client: &reqwest::Client, url: &str) -> reqwest::Response {
client
.post(url)
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#)
.send()
.await
.expect("send non-initialize request")
}
#[tokio::test]
async fn stateless_init_rejects_when_header_older_than_body() -> anyhow::Result<()> {
let (client, url, ct) = spawn_server(stateless_json_config()).await;
@ -147,3 +167,21 @@ async fn stateful_init_rejects_when_header_mismatches_body() -> anyhow::Result<(
ct.cancel();
Ok(())
}
#[tokio::test]
async fn stateful_rejected_initial_posts_do_not_create_sessions() -> anyhow::Result<()> {
let session_manager = Arc::new(LocalSessionManager::default());
let (client, url, ct) =
spawn_server_with_manager(stateful_config(), session_manager.clone()).await;
let response = post_non_initialize(&client, &url).await;
assert_eq!(response.status(), 422);
assert_eq!(session_manager.sessions.read().await.len(), 0);
let response = post_init(&client, &url, Some("2024-11-05"), "2025-11-25").await;
assert_eq!(response.status(), 400);
assert_eq!(session_manager.sessions.read().await.len(), 0);
ct.cancel();
Ok(())
}