Skip to content

Commit 44bb310

Browse files
ddworkenddworken
committed
Fix uncaught exception in MCP server
1 parent 9dad266 commit 44bb310

File tree

2 files changed

+177
-17
lines changed

2 files changed

+177
-17
lines changed

‎src/mcp/shared/session.py‎

Lines changed: 39 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from mcp.shared.exceptions import McpError
1515
from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage
1616
from mcp.types import (
17+
INVALID_PARAMS,
1718
CancelledNotification,
1819
ClientNotification,
1920
ClientRequest,
@@ -351,26 +352,47 @@ async def _receive_loop(self) -> None:
351352
if isinstance(message, Exception):
352353
await self._handle_incoming(message)
353354
elif isinstance(message.message.root, JSONRPCRequest):
354-
validated_request = self._receive_request_type.model_validate(
355-
message.message.root.model_dump(
356-
by_alias=True, mode="json", exclude_none=True
355+
try:
356+
validated_request = self._receive_request_type.model_validate(
357+
message.message.root.model_dump(
358+
by_alias=True, mode="json", exclude_none=True
359+
)
360+
)
361+
responder = RequestResponder(
362+
request_id=message.message.root.id,
363+
request_meta=validated_request.root.params.meta
364+
if validated_request.root.params
365+
else None,
366+
request=validated_request,
367+
session=self,
368+
on_complete=lambda r: self._in_flight.pop(
369+
r.request_id, None),
357370
)
358-
)
359-
responder = RequestResponder(
360-
request_id=message.message.root.id,
361-
request_meta=validated_request.root.params.meta
362-
if validated_request.root.params
363-
else None,
364-
request=validated_request,
365-
session=self,
366-
on_complete=lambda r: self._in_flight.pop(r.request_id, None),
367-
)
368371

369-
self._in_flight[responder.request_id] = responder
370-
await self._received_request(responder)
372+
self._in_flight[responder.request_id] = responder
373+
await self._received_request(responder)
371374

372-
if not responder._completed: # type: ignore[reportPrivateUsage]
373-
await self._handle_incoming(responder)
375+
if not responder._completed: # type: ignore[reportPrivateUsage]
376+
await self._handle_incoming(responder)
377+
except Exception as e:
378+
# For request validation errors, send a proper JSON-RPC error
379+
# response instead of crashing the server
380+
logging.warning(
381+
f"Failed to validate request: {e}. "
382+
f"Message was: {message.message.root}"
383+
)
384+
error_response = JSONRPCError(
385+
jsonrpc="2.0",
386+
id=message.message.root.id,
387+
error=ErrorData(
388+
code=INVALID_PARAMS,
389+
message="Invalid request parameters",
390+
data="",
391+
),
392+
)
393+
session_message = SessionMessage(
394+
message=JSONRPCMessage(error_response))
395+
await self._write_stream.send(session_message)
374396

375397
elif isinstance(message.message.root, JSONRPCNotification):
376398
try:

‎tests/issues/test_malformed_input.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
# Claude Debug
2+
"""Test for HackerOne vulnerability report #3156202 - malformed input DOS."""
3+
4+
import anyio
5+
import pytest
6+
7+
from mcp.server.session import ServerSession
8+
from mcp.shared.message import SessionMessage
9+
from mcp.types import (
10+
INVALID_PARAMS,
11+
JSONRPCError,
12+
JSONRPCMessage,
13+
JSONRPCRequest,
14+
)
15+
16+
17+
@pytest.mark.anyio
18+
async def test_malformed_initialize_request_does_not_crash_server():
19+
"""
20+
Test that malformed initialize requests return proper error responses
21+
instead of crashing the server (HackerOne #3156202).
22+
"""
23+
# Create in-memory streams for testing
24+
read_send_stream, read_receive_stream = anyio.create_memory_object_stream(10)
25+
write_send_stream, write_receive_stream = anyio.create_memory_object_stream(10)
26+
27+
# Create a malformed initialize request (missing required params field)
28+
malformed_request = JSONRPCRequest(
29+
jsonrpc="2.0",
30+
id="f20fe86132ed4cd197f89a7134de5685",
31+
method="initialize",
32+
# params=None # Missing required params field
33+
)
34+
35+
# Wrap in session message
36+
request_message = SessionMessage(message=JSONRPCMessage(malformed_request))
37+
38+
# Start a server session
39+
async with ServerSession(
40+
read_stream=read_receive_stream,
41+
write_stream=write_send_stream,
42+
initialization_options={},
43+
):
44+
# Send the malformed request
45+
await read_send_stream.send(request_message)
46+
47+
# Give the session time to process the request
48+
await anyio.sleep(0.1)
49+
50+
# Check that we received an error response instead of a crash
51+
try:
52+
response_message = write_receive_stream.receive_nowait()
53+
response = response_message.message.root
54+
55+
# Verify it's a proper JSON-RPC error response
56+
assert isinstance(response, JSONRPCError)
57+
assert response.jsonrpc == "2.0"
58+
assert response.id == "f20fe86132ed4cd197f89a7134de5685"
59+
assert response.error.code == INVALID_PARAMS
60+
assert "Invalid request parameters" in response.error.message
61+
62+
# Verify the session is still alive and can handle more requests
63+
# Send another malformed request to confirm server stability
64+
another_malformed_request = JSONRPCRequest(
65+
jsonrpc="2.0",
66+
id="test_id_2",
67+
method="tools/call",
68+
# params=None # Missing required params
69+
)
70+
another_request_message = SessionMessage(
71+
message=JSONRPCMessage(another_malformed_request)
72+
)
73+
74+
await read_send_stream.send(another_request_message)
75+
await anyio.sleep(0.1)
76+
77+
# Should get another error response, not a crash
78+
second_response_message = write_receive_stream.receive_nowait()
79+
second_response = second_response_message.message.root
80+
81+
assert isinstance(second_response, JSONRPCError)
82+
assert second_response.id == "test_id_2"
83+
assert second_response.error.code == INVALID_PARAMS
84+
85+
except anyio.WouldBlock:
86+
pytest.fail("No response received - server likely crashed")
87+
88+
89+
@pytest.mark.anyio
90+
async def test_multiple_concurrent_malformed_requests():
91+
"""
92+
Test that multiple concurrent malformed requests don't crash the server.
93+
"""
94+
# Create in-memory streams for testing
95+
read_send_stream, read_receive_stream = anyio.create_memory_object_stream(100)
96+
write_send_stream, write_receive_stream = anyio.create_memory_object_stream(100)
97+
98+
# Start a server session
99+
async with ServerSession(
100+
read_stream=read_receive_stream,
101+
write_stream=write_send_stream,
102+
initialization_options={},
103+
):
104+
# Send multiple malformed requests concurrently
105+
malformed_requests = []
106+
for i in range(10):
107+
malformed_request = JSONRPCRequest(
108+
jsonrpc="2.0",
109+
id=f"malformed_{i}",
110+
method="initialize",
111+
# params=None # Missing required params
112+
)
113+
request_message = SessionMessage(message=JSONRPCMessage(malformed_request))
114+
malformed_requests.append(request_message)
115+
116+
# Send all requests
117+
for request in malformed_requests:
118+
await read_send_stream.send(request)
119+
120+
# Give time to process
121+
await anyio.sleep(0.2)
122+
123+
# Verify we get error responses for all requests
124+
error_responses = []
125+
try:
126+
while True:
127+
response_message = write_receive_stream.receive_nowait()
128+
error_responses.append(response_message.message.root)
129+
except anyio.WouldBlock:
130+
pass # No more messages
131+
132+
# Should have received 10 error responses
133+
assert len(error_responses) == 10
134+
135+
for i, response in enumerate(error_responses):
136+
assert isinstance(response, JSONRPCError)
137+
assert response.id == f"malformed_{i}"
138+
assert response.error.code == INVALID_PARAMS

0 commit comments

Comments
 (0)