diff --git a/src/a2a/server/routes/common.py b/src/a2a/server/routes/common.py index 18b6865c5..5cddb048b 100644 --- a/src/a2a/server/routes/common.py +++ b/src/a2a/server/routes/common.py @@ -13,12 +13,61 @@ 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 = 413 + + +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..53b453c81 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( + event_generator(handler_result), # ty: ignore[invalid-argument-type] + 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..303c8435a 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 --- + + +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