From b6afa7710172f6f13a923c5a3de5dbef4d735487 Mon Sep 17 00:00:00 2001 From: meraklbz Date: Mon, 10 Aug 2026 20:23:43 +0800 Subject: [PATCH 1/3] fix: harden SSE transport and enforce request body size limit --- src/a2a/server/routes/common.py | 51 ++++++++++++++ src/a2a/server/routes/jsonrpc_dispatcher.py | 9 ++- src/a2a/server/routes/rest_dispatcher.py | 40 +++++++++-- src/a2a/utils/constants.py | 18 +++++ .../server/routes/test_jsonrpc_dispatcher.py | 66 +++++++++++++++++++ tests/server/routes/test_rest_dispatcher.py | 39 +++++++++++ 6 files changed, 215 insertions(+), 8 deletions(-) diff --git a/src/a2a/server/routes/common.py b/src/a2a/server/routes/common.py index 18b6865c5..4dc1aabdc 100644 --- a/src/a2a/server/routes/common.py +++ b/src/a2a/server/routes/common.py @@ -13,12 +13,63 @@ Request = Any BaseUser = Any + try: + from starlette.exceptions import HTTPException + except ImportError: + HTTPException = Any + from a2a.auth.user import UnauthenticatedUser, User from a2a.extensions.common import ( HTTP_EXTENSION_HEADER, get_requested_extensions, ) from a2a.server.context import ServerCallContext +from a2a.utils.constants import MAX_REQUEST_BODY_SIZE + + +try: + from starlette.status import HTTP_413_CONTENT_TOO_LARGE +except ImportError: + HTTP_413_CONTENT_TOO_LARGE = Any + + +async def read_request_body_with_limit(request: Request) -> bytes: + """Reads a request body, rejecting bodies over ``MAX_REQUEST_BODY_SIZE``. + + The content-length header is checked first (fast reject), and the body + is then streamed in chunks with an incremental cap so oversized + chunked bodies are rejected while being read instead of buffered + unboundedly. The body is cached on the request (``request._body``) so + later ``request.body()`` calls return the same bytes. + + Raises: + starlette.exceptions.HTTPException: With status 413 (content too + large) when the body exceeds the limit. + """ + content_length = request.headers.get('content-length') + if content_length: + try: + if int(content_length) > MAX_REQUEST_BODY_SIZE: + raise HTTPException( + status_code=HTTP_413_CONTENT_TOO_LARGE + ) + except ValueError: + pass + + chunks: list[bytes] = [] + size = 0 + async for chunk in request.stream(): + size += len(chunk) + if size > MAX_REQUEST_BODY_SIZE: + raise HTTPException(status_code=HTTP_413_CONTENT_TOO_LARGE) + chunks.append(chunk) + + body = b''.join(chunks) + # Cache the body exactly like Request.body() does (Starlette stores it on + # `_body`), so subsequent request.body()/stream() calls return the same + # bytes instead of re-reading an already-consumed stream. + request._body = body # noqa: SLF001 + return body class StarletteUser(User): diff --git a/src/a2a/server/routes/jsonrpc_dispatcher.py b/src/a2a/server/routes/jsonrpc_dispatcher.py index b59ed7551..e3a642bc4 100644 --- a/src/a2a/server/routes/jsonrpc_dispatcher.py +++ b/src/a2a/server/routes/jsonrpc_dispatcher.py @@ -26,6 +26,7 @@ from a2a.server.routes.common import ( DefaultServerCallContextBuilder, ServerCallContextBuilder, + read_request_body_with_limit, ) from a2a.types.a2a_pb2 import ( CancelTaskRequest, @@ -224,7 +225,7 @@ async def handle_requests(self, request: Request) -> Response: # noqa: PLR0911, body = None try: - body = await request.json() + body = json.loads(await read_request_body_with_limit(request)) if isinstance(body, dict): request_id = body.get('id') # Ensure request_id is valid for JSON-RPC response (str/int/None only) @@ -595,7 +596,11 @@ async def event_generator( 'data': json_utils.dumps(error_response), } - return EventSourceResponse(event_generator(handler_result)) # ty:ignore[invalid-argument-type] + return EventSourceResponse( # ty:ignore[invalid-argument-type] + event_generator(handler_result), + ping=constants.SSE_PING_INTERVAL_SECONDS, + send_timeout=constants.SSE_SEND_TIMEOUT_SECONDS, + ) # handler_result is a dict (JSON-RPC response) return JSONResponse(handler_result) diff --git a/src/a2a/server/routes/rest_dispatcher.py b/src/a2a/server/routes/rest_dispatcher.py index 06c73c592..46dc7ab7a 100644 --- a/src/a2a/server/routes/rest_dispatcher.py +++ b/src/a2a/server/routes/rest_dispatcher.py @@ -10,6 +10,7 @@ from a2a.server.routes.common import ( DefaultServerCallContextBuilder, ServerCallContextBuilder, + read_request_body_with_limit, ) from a2a.types import a2a_pb2 from a2a.types.a2a_pb2 import ( @@ -34,16 +35,20 @@ if TYPE_CHECKING: from sse_starlette.event import ServerSentEvent from sse_starlette.sse import EventSourceResponse + from starlette.exceptions import HTTPException from starlette.requests import Request from starlette.responses import JSONResponse, Response + from starlette.status import HTTP_413_CONTENT_TOO_LARGE _package_starlette_installed = True else: try: from sse_starlette.event import ServerSentEvent from sse_starlette.sse import EventSourceResponse + from starlette.exceptions import HTTPException from starlette.requests import Request from starlette.responses import JSONResponse, Response + from starlette.status import HTTP_413_CONTENT_TOO_LARGE _package_starlette_installed = True except ImportError: @@ -52,6 +57,8 @@ Request = Any JSONResponse = Any Response = Any + HTTPException = Any + HTTP_413_CONTENT_TOO_LARGE = Any _package_starlette_installed = False @@ -98,6 +105,19 @@ def _build_call_context(self, request: Request) -> ServerCallContext: call_context.tenant = request.path_params['tenant'] return call_context + async def _read_request_body(self, request: Request) -> bytes: + """Reads a request body with the size limit enforced. + + Raises: + InvalidRequestError: If the body exceeds the configured limit. + """ + try: + return await read_request_body_with_limit(request) + except HTTPException as e: + if e.status_code == HTTP_413_CONTENT_TOO_LARGE: + raise InvalidRequestError(message='Payload too large') from e + raise + async def _handle_non_streaming( self, request: Request, @@ -117,7 +137,7 @@ async def _handle_streaming( # This is required because Starlette's request.body() can only be consumed once, # and attempting to consume it after EventSourceResponse starts causes deadlock try: - await request.body() + await self._read_request_body(request) except (ValueError, RuntimeError, OSError) as e: raise InvalidRequestError( message=f'Failed to pre-consume request body: {e}' @@ -136,7 +156,11 @@ async def _handle_streaming( try: first_item = await anext(stream) except StopAsyncIteration: - return EventSourceResponse(iter([])) + return EventSourceResponse( + iter([]), + ping=constants.SSE_PING_INTERVAL_SECONDS, + send_timeout=constants.SSE_SEND_TIMEOUT_SECONDS, + ) async def event_generator() -> AsyncIterator[ServerSentEvent]: yield ServerSentEvent(data=json_utils.dumps(first_item)) @@ -150,7 +174,11 @@ async def event_generator() -> AsyncIterator[ServerSentEvent]: event='error', ) - return EventSourceResponse(event_generator()) + return EventSourceResponse( + event_generator(), + ping=constants.SSE_PING_INTERVAL_SECONDS, + send_timeout=constants.SSE_SEND_TIMEOUT_SECONDS, + ) @rest_error_handler async def on_message_send(self, request: Request) -> Response: @@ -160,7 +188,7 @@ async def on_message_send(self, request: Request) -> Response: async def _handler( context: ServerCallContext, ) -> a2a_pb2.SendMessageResponse: - body = await request.body() + body = await self._read_request_body(request) params = a2a_pb2.SendMessageRequest() Parse(body, params) task_or_message = await self.request_handler.on_message_send( @@ -183,7 +211,7 @@ async def on_message_send_stream( async def _handler( context: ServerCallContext, ) -> AsyncIterator[dict[str, Any]]: - body = await request.body() + body = await self._read_request_body(request) params = a2a_pb2.SendMessageRequest() Parse(body, params) async for event in self.request_handler.on_message_send_stream( @@ -295,7 +323,7 @@ async def set_push_notification(self, request: Request) -> Response: async def _handler( context: ServerCallContext, ) -> a2a_pb2.TaskPushNotificationConfig: - body = await request.body() + body = await self._read_request_body(request) params = a2a_pb2.TaskPushNotificationConfig() Parse(body, params) params.task_id = request.path_params['id'] diff --git a/src/a2a/utils/constants.py b/src/a2a/utils/constants.py index 5497d8a24..822595d06 100644 --- a/src/a2a/utils/constants.py +++ b/src/a2a/utils/constants.py @@ -11,6 +11,24 @@ MAX_LIST_TASKS_PAGE_SIZE = 100 """Maximum page size for the `tasks/list` method.""" +MAX_REQUEST_BODY_SIZE = 10 * 1024 * 1024 +"""Maximum accepted HTTP request body size in bytes (10 MiB). + +Requests with a body larger than this are rejected with HTTP 413 +(payload too large) instead of being buffered unboundedly. +""" + +SSE_PING_INTERVAL_SECONDS = 15 +"""Heartbeat interval for SSE streams, in seconds. + +A comment-only ``: ping`` frame is sent on every interval to keep the +connection alive and detect dead clients. +""" + +SSE_SEND_TIMEOUT_SECONDS = 300 +"""Maximum time in seconds without a successful send before an SSE +stream is torn down, preventing zombie streams from accumulating.""" + class TransportProtocol(str, Enum): """Transport protocol string constants.""" diff --git a/tests/server/routes/test_jsonrpc_dispatcher.py b/tests/server/routes/test_jsonrpc_dispatcher.py index 3bde4fc2e..8fa2a506d 100644 --- a/tests/server/routes/test_jsonrpc_dispatcher.py +++ b/tests/server/routes/test_jsonrpc_dispatcher.py @@ -653,3 +653,69 @@ async def stream_generator(): if __name__ == '__main__': pytest.main([__file__]) + + +# --- Transport hardening (BUG-09 / BUG-10) --- + + +class TestTransportHardening: + def test_oversized_request_body_rejected(self, client, mock_handler): + """Bodies over the size limit must be rejected with 413 handling.""" + from a2a.utils.constants import MAX_REQUEST_BODY_SIZE + + oversized = 'x' * (MAX_REQUEST_BODY_SIZE + 1) + response = client.post( + '/', + content=oversized, + headers={ + 'A2A-Version': '1.0', + 'Content-Type': 'application/json', + }, + ) + data = response.json() + assert data['error']['code'] == -32600 # InvalidRequestError + assert 'Payload too large' in data['error']['message'] + + def test_oversized_request_body_rejected_via_content_length( + self, client, mock_handler + ): + """A content-length header over the limit is rejected fast.""" + from a2a.utils.constants import MAX_REQUEST_BODY_SIZE + + response = client.post( + '/', + content=b'{}', + headers={ + 'A2A-Version': '1.0', + 'Content-Type': 'application/json', + 'Content-Length': str(MAX_REQUEST_BODY_SIZE + 1), + }, + ) + data = response.json() + assert data['error']['code'] == -32600 + assert 'Payload too large' in data['error']['message'] + + @pytest.mark.asyncio + async def test_sse_response_explicit_ping_and_send_timeout(self): + """EventSourceResponse must carry explicit ping/send_timeout.""" + from a2a.server.context import ServerCallContext + from a2a.server.request_handlers.request_handler import RequestHandler + from a2a.server.routes.jsonrpc_dispatcher import JsonRpcDispatcher + from a2a.utils.constants import ( + SSE_PING_INTERVAL_SECONDS, + SSE_SEND_TIMEOUT_SECONDS, + ) + + dispatcher = JsonRpcDispatcher( + request_handler=AsyncMock(spec=RequestHandler) + ) + + async def stream(): + yield {'result': 'ok'} + + response = dispatcher._create_response( + ServerCallContext(state={'request_id': '1'}), + stream(), + ) + assert response.ping_interval == SSE_PING_INTERVAL_SECONDS + assert response.send_timeout == SSE_SEND_TIMEOUT_SECONDS diff --git a/tests/server/routes/test_rest_dispatcher.py b/tests/server/routes/test_rest_dispatcher.py index 00cc45f7b..74f714cdf 100644 --- a/tests/server/routes/test_rest_dispatcher.py +++ b/tests/server/routes/test_rest_dispatcher.py @@ -96,6 +96,11 @@ def make_mock_request( mock_req.headers = Headers(default_headers) mock_req.body = AsyncMock(return_value=body) + async def _stream_body(): + yield body + + mock_req.stream = _stream_body + # Needs to be able to build ServerCallContext, so provide .user and .auth etc. if needed mock_req.user = MagicMock(is_authenticated=False) mock_req.auth = None @@ -323,3 +328,37 @@ async def stream_with_non_ascii( payload = payload.decode('utf-8') assert non_ascii_text in payload assert '\\u4f60\\u597d' not in payload + + +@pytest.mark.asyncio +class TestRestDispatcherHardening: + async def test_oversized_request_body_rejected( + self, rest_dispatcher_instance, mock_handler + ): + """A body over the limit is rejected with 400 Payload too large.""" + from a2a.utils.constants import MAX_REQUEST_BODY_SIZE + + oversized = b'x' * (MAX_REQUEST_BODY_SIZE + 1) + req = make_mock_request(method='POST', body=oversized) + response = await rest_dispatcher_instance.on_message_send(req) + + assert response.status_code == 400 + import json + + payload = json.loads(response.body) + assert 'Payload too large' in payload['error']['message'] + + async def test_sse_response_explicit_ping_and_send_timeout( + self, rest_dispatcher_instance + ): + """REST EventSourceResponse must carry explicit ping/send_timeout.""" + from a2a.utils.constants import ( + SSE_PING_INTERVAL_SECONDS, + SSE_SEND_TIMEOUT_SECONDS, + ) + + req = make_mock_request(method='POST') + response = await rest_dispatcher_instance.on_message_send_stream(req) + + assert response.ping_interval == SSE_PING_INTERVAL_SECONDS + assert response.send_timeout == SSE_SEND_TIMEOUT_SECONDS From 025f3c3d999941633f319d1837a9ac0b32030d3d Mon Sep 17 00:00:00 2001 From: meraklbz Date: Tue, 11 Aug 2026 00:39:45 +0800 Subject: [PATCH 2/3] style: apply ruff formatting --- src/a2a/server/routes/common.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/a2a/server/routes/common.py b/src/a2a/server/routes/common.py index 4dc1aabdc..cb5e899ee 100644 --- a/src/a2a/server/routes/common.py +++ b/src/a2a/server/routes/common.py @@ -50,9 +50,7 @@ async def read_request_body_with_limit(request: Request) -> bytes: if content_length: try: if int(content_length) > MAX_REQUEST_BODY_SIZE: - raise HTTPException( - status_code=HTTP_413_CONTENT_TOO_LARGE - ) + raise HTTPException(status_code=HTTP_413_CONTENT_TOO_LARGE) except ValueError: pass From 4f6af77884dd706fcd9564101a5eb16da7a38fd2 Mon Sep 17 00:00:00 2001 From: meraklbz Date: Tue, 11 Aug 2026 01:41:50 +0800 Subject: [PATCH 3/3] fix: correct type fallbacks for HTTP 413 constant and SSE response Use the literal 413 as the ImportError fallback instead of the Any special form, and move the ty ignore comment to the line that actually errors so the type checker suppression takes effect. --- src/a2a/server/routes/common.py | 2 +- src/a2a/server/routes/jsonrpc_dispatcher.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/a2a/server/routes/common.py b/src/a2a/server/routes/common.py index cb5e899ee..5cddb048b 100644 --- a/src/a2a/server/routes/common.py +++ b/src/a2a/server/routes/common.py @@ -30,7 +30,7 @@ try: from starlette.status import HTTP_413_CONTENT_TOO_LARGE except ImportError: - HTTP_413_CONTENT_TOO_LARGE = Any + HTTP_413_CONTENT_TOO_LARGE = 413 async def read_request_body_with_limit(request: Request) -> bytes: diff --git a/src/a2a/server/routes/jsonrpc_dispatcher.py b/src/a2a/server/routes/jsonrpc_dispatcher.py index e3a642bc4..53b453c81 100644 --- a/src/a2a/server/routes/jsonrpc_dispatcher.py +++ b/src/a2a/server/routes/jsonrpc_dispatcher.py @@ -596,8 +596,8 @@ async def event_generator( 'data': json_utils.dumps(error_response), } - return EventSourceResponse( # ty:ignore[invalid-argument-type] - event_generator(handler_result), + return EventSourceResponse( + event_generator(handler_result), # ty: ignore[invalid-argument-type] ping=constants.SSE_PING_INTERVAL_SECONDS, send_timeout=constants.SSE_SEND_TIMEOUT_SECONDS, )