From 526d643403517be0b1a50251a428572a7ce45708 Mon Sep 17 00:00:00 2001 From: Koichi ITO Date: Tue, 17 Mar 2026 02:13:17 +0900 Subject: [PATCH] Use mutex-protected `session_exists?` in `handle_regular_request` ## Motivation and Context `handle_regular_request` was checking `@sessions.key?(session_id)` directly without holding `@mutex`, while concurrent threads could modify `@sessions` via `cleanup_session` or `handle_delete`. This created a TOCTOU race where the check could pass but the session could be deleted before subsequent use. The class already provides a mutex-protected `session_exists?` helper, and `handle_get` already uses it. This change makes `handle_regular_request` consistent with `handle_get`. ## How Has This Been Tested? Added a test that verifies `handle_regular_request` delegates to the mutex-protected `session_exists?` helper instead of accessing `@sessions` directly. All existing tests pass. ## Breaking Change None. --- .../transports/streamable_http_transport.rb | 3 +-- .../streamable_http_transport_test.rb | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/lib/mcp/server/transports/streamable_http_transport.rb b/lib/mcp/server/transports/streamable_http_transport.rb index 2d921bf..2d7b73f 100644 --- a/lib/mcp/server/transports/streamable_http_transport.rb +++ b/lib/mcp/server/transports/streamable_http_transport.rb @@ -259,8 +259,7 @@ def handle_accepted def handle_regular_request(body_string, session_id) unless @stateless - # If session ID is provided, but not in the sessions hash, return an error - if session_id && !@sessions.key?(session_id) + if session_id && !session_exists?(session_id) return [400, { "Content-Type" => "application/json" }, [{ error: "Invalid session ID" }.to_json]] end end diff --git a/test/mcp/server/transports/streamable_http_transport_test.rb b/test/mcp/server/transports/streamable_http_transport_test.rb index 7066e4d..c137bff 100644 --- a/test/mcp/server/transports/streamable_http_transport_test.rb +++ b/test/mcp/server/transports/streamable_http_transport_test.rb @@ -1203,6 +1203,31 @@ class StreamableHTTPTransportTest < ActiveSupport::TestCase assert_equal([], response[2]) end + test "handle_regular_request checks session existence under mutex" do + init_request = create_rack_request( + "POST", + "/", + { "CONTENT_TYPE" => "application/json" }, + { jsonrpc: "2.0", method: "initialize", id: "init" }.to_json, + ) + init_response = @transport.handle_request(init_request) + session_id = init_response[1]["Mcp-Session-Id"] + + @transport.expects(:session_exists?).with(session_id).returns(true) + + request = create_rack_request( + "POST", + "/", + { + "CONTENT_TYPE" => "application/json", + "HTTP_MCP_SESSION_ID" => session_id, + }, + { jsonrpc: "2.0", method: "ping", id: "456" }.to_json, + ) + response = @transport.handle_request(request) + assert_equal(200, response[0]) + end + private def create_rack_request(method, path, headers, body = nil)