Merge commit from fork
Reject duplicate SSE connections with 409 to prevent stream hijacking
This commit is contained in:
commit
e189d7852a
2 changed files with 75 additions and 4 deletions
|
|
@ -142,6 +142,7 @@ module MCP
|
||||||
|
|
||||||
return missing_session_id_response unless session_id
|
return missing_session_id_response unless session_id
|
||||||
return session_not_found_response unless session_exists?(session_id)
|
return session_not_found_response unless session_exists?(session_id)
|
||||||
|
return session_already_connected_response if get_session_stream(session_id)
|
||||||
|
|
||||||
setup_sse_stream(session_id)
|
setup_sse_stream(session_id)
|
||||||
end
|
end
|
||||||
|
|
@ -305,6 +306,14 @@ module MCP
|
||||||
[404, { "Content-Type" => "application/json" }, [{ error: "Session not found" }.to_json]]
|
[404, { "Content-Type" => "application/json" }, [{ error: "Session not found" }.to_json]]
|
||||||
end
|
end
|
||||||
|
|
||||||
|
def session_already_connected_response
|
||||||
|
[
|
||||||
|
409,
|
||||||
|
{ "Content-Type" => "application/json" },
|
||||||
|
[{ error: "Conflict: Only one SSE stream is allowed per session" }.to_json],
|
||||||
|
]
|
||||||
|
end
|
||||||
|
|
||||||
def setup_sse_stream(session_id)
|
def setup_sse_stream(session_id)
|
||||||
body = create_sse_body(session_id)
|
body = create_sse_body(session_id)
|
||||||
|
|
||||||
|
|
@ -319,17 +328,22 @@ module MCP
|
||||||
|
|
||||||
def create_sse_body(session_id)
|
def create_sse_body(session_id)
|
||||||
proc do |stream|
|
proc do |stream|
|
||||||
store_stream_for_session(session_id, stream)
|
stored = store_stream_for_session(session_id, stream)
|
||||||
start_keepalive_thread(session_id)
|
start_keepalive_thread(session_id) if stored
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
def store_stream_for_session(session_id, stream)
|
def store_stream_for_session(session_id, stream)
|
||||||
@mutex.synchronize do
|
@mutex.synchronize do
|
||||||
if @sessions[session_id]
|
session = @sessions[session_id]
|
||||||
@sessions[session_id][:stream] = stream
|
if session && !session[:stream]
|
||||||
|
session[:stream] = stream
|
||||||
else
|
else
|
||||||
|
# Either session was removed, or another request already established a stream.
|
||||||
stream.close
|
stream.close
|
||||||
|
# `stream.close` may return a truthy value depending on the stream class.
|
||||||
|
# Explicitly return nil to guarantee a falsy return for callers.
|
||||||
|
nil
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
|
||||||
|
|
@ -272,6 +272,63 @@ module MCP
|
||||||
assert_equal "Missing session ID", body["error"]
|
assert_equal "Missing session ID", body["error"]
|
||||||
end
|
end
|
||||||
|
|
||||||
|
test "rejects duplicate SSE connection with 409" do
|
||||||
|
# Create a session
|
||||||
|
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"]
|
||||||
|
|
||||||
|
# Simulate an active SSE stream by storing a stream object in the session
|
||||||
|
mock_stream = StringIO.new
|
||||||
|
@transport.instance_variable_get(:@sessions)[session_id][:stream] = mock_stream
|
||||||
|
|
||||||
|
# Attempt a second GET request for the same session
|
||||||
|
get_request = create_rack_request(
|
||||||
|
"GET",
|
||||||
|
"/",
|
||||||
|
{ "HTTP_MCP_SESSION_ID" => session_id },
|
||||||
|
)
|
||||||
|
|
||||||
|
response = @transport.handle_request(get_request)
|
||||||
|
assert_equal 409, response[0]
|
||||||
|
assert_equal({ "Content-Type" => "application/json" }, response[1])
|
||||||
|
|
||||||
|
body = JSON.parse(response[2][0])
|
||||||
|
assert_equal "Conflict: Only one SSE stream is allowed per session", body["error"]
|
||||||
|
end
|
||||||
|
|
||||||
|
test "store_stream_for_session does not overwrite existing stream (TOCTOU guard)" do
|
||||||
|
# Create a session
|
||||||
|
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"]
|
||||||
|
|
||||||
|
# Establish stream A
|
||||||
|
stream_a = StringIO.new
|
||||||
|
@transport.send(:store_stream_for_session, session_id, stream_a)
|
||||||
|
assert_equal stream_a, @transport.instance_variable_get(:@sessions)[session_id][:stream]
|
||||||
|
|
||||||
|
# Attempt to store stream B (simulating a racing request)
|
||||||
|
stream_b = StringIO.new
|
||||||
|
@transport.send(:store_stream_for_session, session_id, stream_b)
|
||||||
|
|
||||||
|
# Stream A should still be the active stream
|
||||||
|
assert_equal stream_a, @transport.instance_variable_get(:@sessions)[session_id][:stream]
|
||||||
|
|
||||||
|
# Stream B should have been closed
|
||||||
|
assert stream_b.closed?
|
||||||
|
end
|
||||||
|
|
||||||
test "handles GET request with invalid session ID" do
|
test "handles GET request with invalid session ID" do
|
||||||
request = create_rack_request(
|
request = create_rack_request(
|
||||||
"GET",
|
"GET",
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue