Skip to content

Commit 6677eee

Browse files
style: cargo fmt fixes on validate_standard_headers changes (#1275)
Signed-off-by: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com>
1 parent fbed447 commit 6677eee

2 files changed

Lines changed: 220 additions & 4 deletions

File tree

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

Lines changed: 134 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,8 @@ use crate::{
3939
OneshotTransport, TransportAdapterIdentity,
4040
common::{
4141
http_header::{
42-
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_MCP_PROTOCOL_VERSION,
43-
HEADER_SESSION_ID, JSON_MIME_TYPE,
42+
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_MCP_METHOD,
43+
HEADER_MCP_PROTOCOL_VERSION, HEADER_SESSION_ID, JSON_MIME_TYPE,
4444
},
4545
mcp_headers,
4646
server_side_http::{
@@ -690,6 +690,99 @@ mod jsonrpc_http_status_tests {
690690
}
691691
}
692692

693+
#[cfg(test)]
694+
mod standard_header_init_tests {
695+
use super::*;
696+
697+
fn initialize_message() -> ClientJsonRpcMessage {
698+
ClientJsonRpcMessage::request(
699+
ClientRequest::InitializeRequest(InitializeRequest {
700+
params: InitializeRequestParams {
701+
protocol_version: ProtocolVersion::STANDARD_HEADERS,
702+
..Default::default()
703+
},
704+
..Default::default()
705+
}),
706+
RequestId::Number(1),
707+
)
708+
}
709+
710+
fn headers_with(mcp_method: Option<&str>) -> HeaderMap {
711+
let mut headers = HeaderMap::new();
712+
headers.insert(
713+
HEADER_MCP_PROTOCOL_VERSION,
714+
http::HeaderValue::from_static(ProtocolVersion::STANDARD_HEADERS.as_str()),
715+
);
716+
if let Some(method) = mcp_method {
717+
headers.insert(
718+
HEADER_MCP_METHOD,
719+
method.parse::<http::HeaderValue>().unwrap(),
720+
);
721+
}
722+
headers
723+
}
724+
725+
fn no_tool_schema(_: &str) -> Option<Arc<JsonObject>> {
726+
None
727+
}
728+
729+
/// A supplied Mcp-Method header contradicting an initialize body must be
730+
/// rejected at >= STANDARD_HEADERS. Regression test for
731+
/// https://github.com/modelcontextprotocol/rust-sdk/issues/1271
732+
#[test]
733+
fn initialize_rejects_contradicting_mcp_method_header() {
734+
let headers = headers_with(Some("tools/list"));
735+
assert!(
736+
validate_standard_headers(&headers, &initialize_message(), no_tool_schema).is_err()
737+
);
738+
}
739+
740+
/// Absence of the Mcp-Method header on initialize stays accepted: clients
741+
/// emit SEP-2243 headers only after the version has been negotiated.
742+
#[test]
743+
fn initialize_accepts_missing_mcp_method_header() {
744+
let headers = headers_with(None);
745+
assert!(validate_standard_headers(&headers, &initialize_message(), no_tool_schema).is_ok());
746+
}
747+
748+
/// A supplied Mcp-Method header matching the initialize body is accepted.
749+
#[test]
750+
fn initialize_accepts_matching_mcp_method_header() {
751+
let headers = headers_with(Some("initialize"));
752+
assert!(validate_standard_headers(&headers, &initialize_message(), no_tool_schema).is_ok());
753+
}
754+
755+
/// Conflicting duplicate Mcp-Method values are rejected in both orders:
756+
/// only the first value used to be checked, so an appended contradictory
757+
/// value passed silently. Duplicate values are rejected explicitly.
758+
#[test]
759+
fn initialize_rejects_conflicting_duplicates_in_both_orders() {
760+
for (first, second) in [("initialize", "tools/list"), ("tools/list", "initialize")] {
761+
let mut headers = headers_with(Some(first));
762+
headers.append(
763+
HEADER_MCP_METHOD,
764+
http::HeaderValue::from_str(second).unwrap(),
765+
);
766+
assert!(
767+
validate_standard_headers(&headers, &initialize_message(), no_tool_schema).is_err()
768+
);
769+
}
770+
}
771+
772+
/// A present but non-UTF8 Mcp-Method value is rejected: it must not be
773+
/// silently treated as an absent header.
774+
#[test]
775+
fn initialize_rejects_present_non_text_method() {
776+
let mut headers = headers_with(None);
777+
let value = http::HeaderValue::from_bytes(&[0xff]).unwrap();
778+
assert!(value.to_str().is_err());
779+
headers.insert(HEADER_MCP_METHOD, value);
780+
assert!(
781+
validate_standard_headers(&headers, &initialize_message(), no_tool_schema).is_err()
782+
);
783+
}
784+
}
785+
693786
fn jsonrpc_message_response(
694787
message: ServerJsonRpcMessage,
695788
map_protocol_status: bool,
@@ -724,8 +817,11 @@ fn header_mismatch_jsonrpc_response(
724817
/// Validates SEP-2243 `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers against the body.
725818
///
726819
/// Only enforced when the request declares a protocol version `>= STANDARD_HEADERS`.
727-
/// The `initialize` handshake is exempt: clients emit these headers only after the
728-
/// version has been negotiated. `tool_schema` supplies the called tool's input schema
820+
/// The `initialize` handshake is exempt from *requiring* them: clients emit these
821+
/// headers only after the version has been negotiated. But like
822+
/// `validate_header_matches_init_body`, a *supplied* `Mcp-Method` header that
823+
/// contradicts the body is rejected — middleboxes route on the header without
824+
/// parsing the body. `tool_schema` supplies the called tool's input schema
729825
/// so annotated `Mcp-Param-*` headers can be checked (no schema => those are skipped).
730826
fn validate_standard_headers(
731827
headers: &HeaderMap,
@@ -743,6 +839,35 @@ fn validate_standard_headers(
743839
let request_id = match message {
744840
ClientJsonRpcMessage::Request(req) => {
745841
if matches!(&req.request, ClientRequest::InitializeRequest(_)) {
842+
// The handshake may omit SEP-2243 headers, but a supplied
843+
// Mcp-Method header must still agree with the body: middleboxes
844+
// route on the header without parsing the body. `get_all` is
845+
// used so every supplied value is inspected: a contradictory
846+
// duplicate must not hide behind a matching first value, and a
847+
// present-but-non-UTF8 value is rejected rather than silently
848+
// treated as absent. Duplicates (even identical ones) are
849+
// rejected explicitly instead of passing by accident.
850+
let mut method_values = headers.get_all(HEADER_MCP_METHOD).iter();
851+
let Some(first) = method_values.next() else {
852+
return Ok(());
853+
};
854+
if method_values.next().is_some() {
855+
return Err(header_mismatch_jsonrpc_response(
856+
Some(req.id.clone()),
857+
"multiple Mcp-Method header values supplied; initialize requires exactly one",
858+
)
859+
.into());
860+
}
861+
if first != http::HeaderValue::from_static("initialize") {
862+
return Err(header_mismatch_jsonrpc_response(
863+
Some(req.id.clone()),
864+
format!(
865+
"Mcp-Method header `{}` does not match body method `initialize`",
866+
String::from_utf8_lossy(first.as_bytes())
867+
),
868+
)
869+
.into());
870+
}
746871
return Ok(());
747872
}
748873
Some(req.id.clone())
@@ -1860,6 +1985,11 @@ where
18601985
.serve_negotiated_request_directly(service, request, part)
18611986
.await;
18621987
}
1988+
// SEP-2243: run the standard-header validator on the initialize
1989+
// handshake before a session is created for it — a supplied
1990+
// Mcp-Method header contradicting the body must be rejected
1991+
// here, not answered with HTTP 200 and a fresh session.
1992+
validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?;
18631993
// Capture init params for external store persistence before
18641994
// extensions are injected (which would require Clone).
18651995
let stored_init_params = match &mut message {

‎crates/rmcp/tests/test_streamable_http_standard_headers.rs‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,3 +313,89 @@ async fn rejects_missing_param_header_with_32020() -> anyhow::Result<()> {
313313
ct.cancel();
314314
Ok(())
315315
}
316+
317+
/// Spawns the server in legacy session mode: the initialize handshake creates
318+
/// a real session, exercising the path where the validator used to be bypassed.
319+
async fn spawn_legacy_server() -> (reqwest::Client, String, CancellationToken) {
320+
let config = StreamableHttpServerConfig::default()
321+
.with_legacy_session_mode(true)
322+
.with_json_response(true)
323+
.with_sse_keep_alive(None)
324+
.with_cancellation_token(CancellationToken::new());
325+
let ct = config.cancellation_token.clone();
326+
let service: StreamableHttpService<HeaderValidationServer, LocalSessionManager> =
327+
StreamableHttpService::new(|| Ok(HeaderValidationServer), Default::default(), config);
328+
329+
let router = axum::Router::new().nest_service("/mcp", service);
330+
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
331+
let addr = tcp_listener.local_addr().unwrap();
332+
tokio::spawn({
333+
let ct = ct.clone();
334+
async move {
335+
let _ = axum::serve(tcp_listener, router)
336+
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
337+
.await;
338+
}
339+
});
340+
(reqwest::Client::new(), format!("http://{addr}/mcp"), ct)
341+
}
342+
343+
/// POSTs an `initialize` request with the given protocol version and optional
344+
/// `Mcp-Method` header.
345+
async fn post_initialize(
346+
client: &reqwest::Client,
347+
url: &str,
348+
version: &str,
349+
mcp_method: Option<&str>,
350+
) -> reqwest::Response {
351+
let body = serde_json::json!({
352+
"jsonrpc": "2.0",
353+
"id": 1,
354+
"method": "initialize",
355+
"params": {
356+
"protocolVersion": version,
357+
"capabilities": {},
358+
"clientInfo": { "name": "test", "version": "1.0" },
359+
}
360+
});
361+
let mut req = client
362+
.post(url)
363+
.header("Content-Type", "application/json")
364+
.header("Accept", "application/json, text/event-stream")
365+
.header("MCP-Protocol-Version", version)
366+
.body(body.to_string());
367+
if let Some(method) = mcp_method {
368+
req = req.header("Mcp-Method", method);
369+
}
370+
req.send().await.expect("send initialize request")
371+
}
372+
373+
#[tokio::test]
374+
async fn rejects_initialize_with_contradicting_mcp_method_before_session_creation()
375+
-> anyhow::Result<()> {
376+
let (client, url, ct) = spawn_legacy_server().await;
377+
378+
// Regression test: in legacy session mode with no session id, an
379+
// `initialize` body carrying a contradictory `Mcp-Method` header used to
380+
// bypass the validator, returning HTTP 200 and creating a session.
381+
let response = post_initialize(&client, &url, SEP_VERSION, Some("tools/list")).await;
382+
assert_eq!(response.status(), 400);
383+
let body: serde_json::Value = response.json().await?;
384+
assert_eq!(body["error"]["code"], -32020);
385+
386+
ct.cancel();
387+
Ok(())
388+
}
389+
390+
#[tokio::test]
391+
async fn accepts_initialize_with_matching_mcp_method() -> anyhow::Result<()> {
392+
let (client, url, ct) = spawn_legacy_server().await;
393+
394+
// A matching Mcp-Method header on initialize passes validation and the
395+
// handshake completes (HTTP 200, session created).
396+
let response = post_initialize(&client, &url, SEP_VERSION, Some("initialize")).await;
397+
assert_eq!(response.status(), 200);
398+
399+
ct.cancel();
400+
Ok(())
401+
}

0 commit comments

Comments
 (0)