-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Fix uncaught exception in MCP server #822
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
+212
−19
Merged
Changes from all commits
Commits
Show all changes
13 commits
Select commit
Hold shift + click to select a range
44bb310
Fix uncaught exception in MCP server
ddworken 900548c
Fix tests
ddworken 6933486
Fix tests
ddworken 9f83b50
Fix tests
ddworken 92865e8
Fix tests
ddworken aba774d
Fix tests
ddworken 3ecc02a
Move extra debug information to logging.debug
ddworken 23d4af4
Merge branch 'main' into main
ddworken 3757bad
Fix lint
ddworken 4e3937f
Merge branch 'main' into main
ddworken 605a6d4
Rerun tests
ddworken a2f889b
Merge branch 'main' into main
ddworken acb45c6
Merge branch 'main' into main
ddworken File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -15,6 +15,7 @@ | |
from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage | ||
from mcp.types import ( | ||
CONNECTION_CLOSED, | ||
INVALID_PARAMS, | ||
CancelledNotification, | ||
ClientNotification, | ||
ClientRequest, | ||
|
@@ -354,27 +355,47 @@ async def _receive_loop(self) -> None: | |
if isinstance(message, Exception): | ||
await self._handle_incoming(message) | ||
elif isinstance(message.message.root, JSONRPCRequest): | ||
validated_request = self._receive_request_type.model_validate( | ||
message.message.root.model_dump( | ||
by_alias=True, mode="json", exclude_none=True | ||
try: | ||
validated_request = self._receive_request_type.model_validate( | ||
message.message.root.model_dump( | ||
by_alias=True, mode="json", exclude_none=True | ||
) | ||
) | ||
) | ||
responder = RequestResponder( | ||
request_id=message.message.root.id, | ||
request_meta=validated_request.root.params.meta | ||
if validated_request.root.params | ||
else None, | ||
request=validated_request, | ||
session=self, | ||
on_complete=lambda r: self._in_flight.pop(r.request_id, None), | ||
message_metadata=message.metadata, | ||
) | ||
|
||
self._in_flight[responder.request_id] = responder | ||
await self._received_request(responder) | ||
responder = RequestResponder( | ||
request_id=message.message.root.id, | ||
request_meta=validated_request.root.params.meta | ||
if validated_request.root.params | ||
else None, | ||
request=validated_request, | ||
session=self, | ||
on_complete=lambda r: self._in_flight.pop( | ||
r.request_id, None), | ||
message_metadata=message.metadata, | ||
) | ||
self._in_flight[responder.request_id] = responder | ||
await self._received_request(responder) | ||
|
||
if not responder._completed: # type: ignore[reportPrivateUsage] | ||
await self._handle_incoming(responder) | ||
if not responder._completed: # type: ignore[reportPrivateUsage] | ||
await self._handle_incoming(responder) | ||
except Exception as e: | ||
# For request validation errors, send a proper JSON-RPC error | ||
# response instead of crashing the server | ||
logging.warning(f"Failed to validate request: {e}") | ||
logging.debug( | ||
f"Message that failed validation: {message.message.root}" | ||
) | ||
error_response = JSONRPCError( | ||
jsonrpc="2.0", | ||
ddworken marked this conversation as resolved.
Show resolved
Hide resolved
|
||
id=message.message.root.id, | ||
error=ErrorData( | ||
code=INVALID_PARAMS, | ||
message="Invalid request parameters", | ||
data="", | ||
), | ||
) | ||
session_message = SessionMessage( | ||
message=JSONRPCMessage(error_response)) | ||
Comment on lines
+396
to
+397
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. How this got merged? |
||
await self._write_stream.send(session_message) | ||
|
||
elif isinstance(message.message.root, JSONRPCNotification): | ||
try: | ||
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,172 @@ | ||
# Claude Debug | ||
"""Test for HackerOne vulnerability report #3156202 - malformed input DOS.""" | ||
|
||
import anyio | ||
import pytest | ||
|
||
from mcp.server.models import InitializationOptions | ||
from mcp.server.session import ServerSession | ||
from mcp.shared.message import SessionMessage | ||
from mcp.types import ( | ||
INVALID_PARAMS, | ||
JSONRPCError, | ||
JSONRPCMessage, | ||
JSONRPCRequest, | ||
ServerCapabilities, | ||
) | ||
|
||
|
||
@pytest.mark.anyio | ||
async def test_malformed_initialize_request_does_not_crash_server(): | ||
""" | ||
Test that malformed initialize requests return proper error responses | ||
instead of crashing the server (HackerOne #3156202). | ||
""" | ||
# Create in-memory streams for testing | ||
read_send_stream, read_receive_stream = anyio.create_memory_object_stream[ | ||
SessionMessage | Exception | ||
](10) | ||
write_send_stream, write_receive_stream = anyio.create_memory_object_stream[ | ||
SessionMessage | ||
](10) | ||
|
||
try: | ||
# Create a malformed initialize request (missing required params field) | ||
malformed_request = JSONRPCRequest( | ||
jsonrpc="2.0", | ||
id="f20fe86132ed4cd197f89a7134de5685", | ||
method="initialize", | ||
# params=None # Missing required params field | ||
) | ||
|
||
# Wrap in session message | ||
request_message = SessionMessage(message=JSONRPCMessage(malformed_request)) | ||
|
||
# Start a server session | ||
async with ServerSession( | ||
read_stream=read_receive_stream, | ||
write_stream=write_send_stream, | ||
init_options=InitializationOptions( | ||
server_name="test_server", | ||
server_version="1.0.0", | ||
capabilities=ServerCapabilities(), | ||
), | ||
): | ||
# Send the malformed request | ||
await read_send_stream.send(request_message) | ||
|
||
# Give the session time to process the request | ||
await anyio.sleep(0.1) | ||
|
||
# Check that we received an error response instead of a crash | ||
try: | ||
response_message = write_receive_stream.receive_nowait() | ||
response = response_message.message.root | ||
|
||
# Verify it's a proper JSON-RPC error response | ||
assert isinstance(response, JSONRPCError) | ||
assert response.jsonrpc == "2.0" | ||
assert response.id == "f20fe86132ed4cd197f89a7134de5685" | ||
assert response.error.code == INVALID_PARAMS | ||
assert "Invalid request parameters" in response.error.message | ||
|
||
# Verify the session is still alive and can handle more requests | ||
# Send another malformed request to confirm server stability | ||
another_malformed_request = JSONRPCRequest( | ||
jsonrpc="2.0", | ||
id="test_id_2", | ||
method="tools/call", | ||
# params=None # Missing required params | ||
) | ||
another_request_message = SessionMessage( | ||
message=JSONRPCMessage(another_malformed_request) | ||
) | ||
|
||
await read_send_stream.send(another_request_message) | ||
await anyio.sleep(0.1) | ||
|
||
# Should get another error response, not a crash | ||
second_response_message = write_receive_stream.receive_nowait() | ||
second_response = second_response_message.message.root | ||
|
||
assert isinstance(second_response, JSONRPCError) | ||
assert second_response.id == "test_id_2" | ||
assert second_response.error.code == INVALID_PARAMS | ||
|
||
except anyio.WouldBlock: | ||
pytest.fail("No response received - server likely crashed") | ||
finally: | ||
# Close all streams to ensure proper cleanup | ||
await read_send_stream.aclose() | ||
await write_send_stream.aclose() | ||
await read_receive_stream.aclose() | ||
await write_receive_stream.aclose() | ||
|
||
|
||
@pytest.mark.anyio | ||
async def test_multiple_concurrent_malformed_requests(): | ||
""" | ||
Test that multiple concurrent malformed requests don't crash the server. | ||
""" | ||
# Create in-memory streams for testing | ||
read_send_stream, read_receive_stream = anyio.create_memory_object_stream[ | ||
SessionMessage | Exception | ||
](100) | ||
write_send_stream, write_receive_stream = anyio.create_memory_object_stream[ | ||
SessionMessage | ||
](100) | ||
|
||
try: | ||
# Start a server session | ||
async with ServerSession( | ||
read_stream=read_receive_stream, | ||
write_stream=write_send_stream, | ||
init_options=InitializationOptions( | ||
server_name="test_server", | ||
server_version="1.0.0", | ||
capabilities=ServerCapabilities(), | ||
), | ||
): | ||
# Send multiple malformed requests concurrently | ||
malformed_requests = [] | ||
for i in range(10): | ||
malformed_request = JSONRPCRequest( | ||
jsonrpc="2.0", | ||
id=f"malformed_{i}", | ||
method="initialize", | ||
# params=None # Missing required params | ||
) | ||
request_message = SessionMessage( | ||
message=JSONRPCMessage(malformed_request) | ||
) | ||
malformed_requests.append(request_message) | ||
|
||
# Send all requests | ||
for request in malformed_requests: | ||
await read_send_stream.send(request) | ||
|
||
# Give time to process | ||
await anyio.sleep(0.2) | ||
|
||
# Verify we get error responses for all requests | ||
error_responses = [] | ||
try: | ||
while True: | ||
response_message = write_receive_stream.receive_nowait() | ||
error_responses.append(response_message.message.root) | ||
except anyio.WouldBlock: | ||
pass # No more messages | ||
|
||
# Should have received 10 error responses | ||
assert len(error_responses) == 10 | ||
|
||
for i, response in enumerate(error_responses): | ||
assert isinstance(response, JSONRPCError) | ||
assert response.id == f"malformed_{i}" | ||
assert response.err 50F2 or.code == INVALID_PARAMS | ||
finally: | ||
# Close all streams to ensure proper cleanup | ||
await read_send_stream.aclose() | ||
await write_send_stream.aclose() | ||
await read_receive_stream.aclose() | ||
await write_receive_stream.aclose() |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think is better to handle a specific exceptions before the general one, such as
RuntimeError
(Which mentioned in this issue)There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Given that the risk is the server becoming unresponsive, I believe catching all exceptions to isolate errors to a single request is correct.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think this is a valid answer from the package's point of view.
It makes it considerably hard to maintain the package if everything is
except Exception
.