From 433e56d518dfc715619700e60bc35213d46386b8 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Mon, 28 Sep 2026 13:41:21 +0000 Subject: [PATCH] feat(event_handler): add opt-in ALB response size validation --- .../event_handler/api_gateway.py | 53 ++- .../event_handler/exceptions.py | 19 + docs/core/event_handler/api_gateway.md | 16 + .../src/alb_response_size_validation.py | 30 ++ .../test_alb_response_size_validation.py | 413 ++++++++++++++++++ 5 files changed, 528 insertions(+), 3 deletions(-) create mode 100644 examples/event_handler_rest/src/alb_response_size_validation.py create mode 100644 tests/functional/event_handler/required_dependencies/test_alb_response_size_validation.py diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py index d5af0b33a98..25b34613f44 100644 --- a/aws_lambda_powertools/event_handler/api_gateway.py +++ b/aws_lambda_powertools/event_handler/api_gateway.py @@ -3,6 +3,7 @@ import base64 import json import logging +import os import re import traceback import warnings @@ -18,7 +19,7 @@ from aws_lambda_powertools.event_handler import content_types from aws_lambda_powertools.event_handler.exception_handling import ExceptionHandlerManager -from aws_lambda_powertools.event_handler.exceptions import NotFoundError, ServiceError +from aws_lambda_powertools.event_handler.exceptions import NotFoundError, ResponseSizeExceededError, ServiceError from aws_lambda_powertools.event_handler.openapi.config import OpenAPIConfig from aws_lambda_powertools.event_handler.openapi.constants import ( DEFAULT_API_VERSION, @@ -72,6 +73,7 @@ _NAMED_GROUP_BOUNDARY_PATTERN = rf"(?P\1[{_SAFE_URI}{_UNSAFE_URI}\\w]+)" _ROUTE_REGEX = "^{}$" _JSON_DUMP_CALL = partial(json.dumps, separators=(",", ":"), cls=Encoder) +_ALB_MAX_RESPONSE_SIZE = 1024 * 1024 ResponseEventT = TypeVar("ResponseEventT", bound=BaseProxyEvent) ResponseT = TypeVar("ResponseT") @@ -2558,7 +2560,7 @@ def resolve(self, event: Mapping[str, Any], context: LambdaContext) -> dict[str, BaseRouter.current_event = self._to_proxy_event(cast(dict, event)) BaseRouter.lambda_context = context - response = self._resolve().build(self.current_event, self._cors) + response = self._build_response(self._resolve()) # Debug print Processed Middlewares if self._debug: @@ -2619,7 +2621,7 @@ def lambda_handler(event, context): BaseRouter.current_event = self._to_proxy_event(cast(dict, event)) BaseRouter.lambda_context = context - response = (await self._resolve_async()).build(self.current_event, self._cors) + response = self._build_response(await self._resolve_async()) if self._debug: print("\nProcessed Middlewares:") @@ -2631,6 +2633,10 @@ def lambda_handler(event, context): return response + def _build_response(self, response_builder: ResponseBuilder) -> dict[str, Any]: + """Build the final integration response after route and middleware processing.""" + return response_builder.build(self.current_event, self._cors) + async def _resolve_async(self) -> ResponseBuilder: method = self.current_event.http_method.upper() path = self._remove_prefix(self.current_event.path) @@ -3342,6 +3348,7 @@ def __init__( response_validation_error_http_code: HTTPStatus | int | None = None, json_body_deserializer: Callable[[str], dict] | None = None, decode_query_parameters: bool = False, + enable_response_size_validation: bool = False, ): """Amazon Application Load Balancer (ALB) resolver @@ -3368,6 +3375,10 @@ def __init__( by default json.loads when integrating with EventSource data class decode_query_parameters: bool | None Enables URL-decoding of query parameters (both keys and values), by default False. + enable_response_size_validation: bool + Validate the complete serialized response against ALB's 1 MB limit, by default False. + Raises ResponseSizeExceededError when the limit is exceeded. Registered exception handlers + can return a smaller response, which is also validated. """ super().__init__( cors=cors, @@ -3379,6 +3390,42 @@ def __init__( json_body_deserializer=json_body_deserializer, ) self.decode_query_parameters = decode_query_parameters + self._enable_response_size_validation = enable_response_size_validation + + @override + def _build_response(self, response_builder: ResponseBuilder) -> dict[str, Any]: + response = super()._build_response(response_builder) + if not self._enable_response_size_validation: + return response + + try: + self._validate_response_size(response) + except ResponseSizeExceededError as exc: + try: + # Resolved responses retain their route, including not-found and preflight responses. + handled_response = self._call_exception_handler(exc, cast(Route, response_builder.route)) + if handled_response is None: + raise + + handled_response.response = cast(Response, self._to_response(handled_response.response)) + response = super()._build_response(handled_response) + # Validate once more without invoking an exception handler recursively. + self._validate_response_size(response) + except Exception: + self.clear_context() + raise + + return response + + @staticmethod + def _validate_response_size(response: dict[str, Any]) -> None: + # Match the Lambda Python runtime's outer JSON encoding, including its default separators. + # Managed Python 3.12+ runtimes return Unicode directly; older/custom runtimes escape it. + runtime = re.fullmatch(r"AWS_Lambda_python3\.(\d+)", os.getenv("AWS_EXECUTION_ENV", "")) + ensure_ascii = runtime is None or int(runtime[1]) < 12 + actual_size = len(json.dumps(response, ensure_ascii=ensure_ascii).encode("utf-8")) + if actual_size > _ALB_MAX_RESPONSE_SIZE: + raise ResponseSizeExceededError(actual_size=actual_size, max_size=_ALB_MAX_RESPONSE_SIZE) def _get_base_path(self) -> str: # ALB doesn't have a stage variable, so we just return an empty string diff --git a/aws_lambda_powertools/event_handler/exceptions.py b/aws_lambda_powertools/event_handler/exceptions.py index e870727cfba..ca30bd3da1f 100644 --- a/aws_lambda_powertools/event_handler/exceptions.py +++ b/aws_lambda_powertools/event_handler/exceptions.py @@ -3,6 +3,25 @@ from http import HTTPStatus +class ResponseSizeExceededError(Exception): + """The serialized Lambda response exceeds the integration's size limit.""" + + def __init__(self, *, actual_size: int, max_size: int): + """ + Parameters + ---------- + actual_size: int + Size of the complete serialized response in bytes. + max_size: int + Maximum allowed serialized response size in bytes. + """ + self.actual_size = actual_size + self.max_size = max_size + super().__init__( + f"Serialized response size ({actual_size} bytes) exceeds the maximum allowed size ({max_size} bytes)", + ) + + class ServiceError(Exception): """Powertools class HTTP Service Error""" diff --git a/docs/core/event_handler/api_gateway.md b/docs/core/event_handler/api_gateway.md index 67bb07a25cc..c0bd69372ef 100644 --- a/docs/core/event_handler/api_gateway.md +++ b/docs/core/event_handler/api_gateway.md @@ -151,6 +151,22 @@ When using Amazon Application Load Balancer (ALB) to front your Lambda functions --8<-- "examples/event_handler_rest/src/getting_started_alb_api_resolver.py" ``` +##### Validate response size + +ALB limits the complete response JSON returned by Lambda to [1 MB](https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html){target="_blank"}. An oversized response can produce a 502 error even when the Lambda handler finishes successfully. + +Set `enable_response_size_validation=True` to check the response against the 1,048,576-byte limit before returning it to ALB. This option is disabled by default and applies only to `ALBResolver`, through both `resolve()` and `resolve_async()`. + +The check includes the JSON envelope, serialized body, headers, and cookies after compression and base64 encoding. It follows the Lambda Python runtime's JSON encoding, including escaping and UTF-8 byte lengths. + +When the limit is exceeded, `ResponseSizeExceededError` exposes `actual_size` and `max_size`, both in bytes. Without a registered exception handler, the exception propagates and fails the invocation. You can register a handler to log the sizes and return a small 500 response: + +```python hl_lines="7 10-22" title="Validate ALB response size" +--8<-- "examples/event_handler_rest/src/alb_response_size_validation.py" +``` + +The handler's response is also validated. If it exceeds the limit, the exception propagates without invoking the handler again. + #### Lambda Function URL When using [AWS Lambda Function URL](https://docs.aws.amazon.com/lambda/latest/dg/urls-configuration.html){target="_blank"}, you can use `LambdaFunctionUrlResolver`. diff --git a/examples/event_handler_rest/src/alb_response_size_validation.py b/examples/event_handler_rest/src/alb_response_size_validation.py new file mode 100644 index 00000000000..9c080455a33 --- /dev/null +++ b/examples/event_handler_rest/src/alb_response_size_validation.py @@ -0,0 +1,30 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.event_handler import ALBResolver, Response, content_types +from aws_lambda_powertools.event_handler.exceptions import ResponseSizeExceededError +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() +app = ALBResolver(enable_response_size_validation=True) + + +@app.exception_handler(ResponseSizeExceededError) +def handle_large_response(error: ResponseSizeExceededError) -> Response: + logger.error( + "Response exceeds the ALB limit", + actual_size=error.actual_size, + max_size=error.max_size, + ) + return Response( + status_code=500, + content_type=content_types.APPLICATION_JSON, + body={"message": "Unable to return the response"}, + ) + + +@app.get("/reports") +def get_report(): + return {"report": "Report contents"} + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/tests/functional/event_handler/required_dependencies/test_alb_response_size_validation.py b/tests/functional/event_handler/required_dependencies/test_alb_response_size_validation.py new file mode 100644 index 00000000000..e6762aca8fc --- /dev/null +++ b/tests/functional/event_handler/required_dependencies/test_alb_response_size_validation.py @@ -0,0 +1,413 @@ +from __future__ import annotations + +import asyncio +import base64 +import gzip +import json +import random + +import pytest + +from aws_lambda_powertools.event_handler import ( + ALBResolver, + APIGatewayHttpResolver, + APIGatewayRestResolver, + CORSConfig, + Response, + content_types, +) +from aws_lambda_powertools.event_handler.exceptions import InternalServerError, ResponseSizeExceededError +from aws_lambda_powertools.shared.cookies import Cookie +from tests.functional.utils import load_event + +MAX_RESPONSE_SIZE = 1024 * 1024 + + +@pytest.fixture(autouse=True) +def lambda_runtime_environment(monkeypatch): + monkeypatch.setenv("AWS_EXECUTION_ENV", "AWS_Lambda_python3.14") + monkeypatch.delenv("POWERTOOLS_DEV", raising=False) + + +@pytest.fixture +def alb_event(): + return load_event("albEvent.json") + + +@pytest.fixture(params=["sync", "async"]) +def resolve_response(request): + def resolve(app, event): + if request.param == "async": + return asyncio.run(app.resolve_async(event, {})) + return app.resolve(event, {}) + + return resolve + + +def make_app(body, *, enabled=True, compress=False, **response_options): + app = ALBResolver(enable_response_size_validation=enabled) + + @app.get("/lambda", compress=compress) + def get_response(): + return Response(200, content_types.TEXT_PLAIN, body, **response_options) + + return app + + +def response_size(response, *, ensure_ascii=False): + return len(json.dumps(response, ensure_ascii=ensure_ascii).encode("utf-8")) + + +def test_response_size_validation_is_disabled_by_default(alb_event, resolve_response): + app = ALBResolver() + + @app.get("/lambda") + def get_response(): + return Response(200, content_types.TEXT_PLAIN, "x" * MAX_RESPONSE_SIZE) + + response = resolve_response(app, alb_event) + assert response["statusCode"] == 200 + assert response_size(response) > MAX_RESPONSE_SIZE + + +def test_response_size_validation_can_be_disabled_explicitly(alb_event, resolve_response): + app = make_app("x" * MAX_RESPONSE_SIZE, enabled=False) + assert response_size(resolve_response(app, alb_event)) > MAX_RESPONSE_SIZE + + +@pytest.mark.parametrize("offset", [-1, 0, 1], ids=["below-limit", "at-limit", "above-limit"]) +def test_response_size_boundary_includes_the_envelope(alb_event, resolve_response, offset): + envelope = { + "statusCode": 200, + "body": "", + "isBase64Encoded": False, + "headers": {"Content-Type": content_types.TEXT_PLAIN}, + } + body = "x" * (MAX_RESPONSE_SIZE - response_size(envelope) + offset) + app = make_app(body) + + if offset <= 0: + response = resolve_response(app, alb_event) + assert response_size(response) == MAX_RESPONSE_SIZE + offset + else: + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(app, alb_event) + assert raised.value.actual_size == MAX_RESPONSE_SIZE + offset + assert raised.value.max_size == MAX_RESPONSE_SIZE + assert app.context == {} + + +@pytest.mark.parametrize( + "execution_env,ensure_ascii", + [ + (None, True), + ("AWS_Lambda_python3.10", True), + ("AWS_Lambda_python3.11", True), + ("AWS_Lambda_python3.12", False), + ("AWS_Lambda_python3.14", False), + ], +) +def test_response_size_matches_runtime_unicode_encoding( + alb_event, + resolve_response, + monkeypatch, + execution_env, + ensure_ascii, +): + if execution_env is None: + monkeypatch.delenv("AWS_EXECUTION_ENV") + else: + monkeypatch.setenv("AWS_EXECUTION_ENV", execution_env) + + body = "é" * 200_000 + expected = resolve_response(make_app(body, enabled=False), alb_event) + size = response_size(expected, ensure_ascii=ensure_ascii) + + if ensure_ascii: + assert size > MAX_RESPONSE_SIZE + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(make_app(body), alb_event) + assert raised.value.actual_size == size + else: + assert size < MAX_RESPONSE_SIZE + assert resolve_response(make_app(body), alb_event) == expected + + +def test_json_escaping_is_included_in_response_size(alb_event, resolve_response): + body = '"\\\n' * (MAX_RESPONSE_SIZE // 4) + assert len(body.encode("utf-8")) < MAX_RESPONSE_SIZE + expected = resolve_response(make_app(body, enabled=False), alb_event) + + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(make_app(body), alb_event) + + assert raised.value.actual_size == response_size(expected) + assert raised.value.actual_size > MAX_RESPONSE_SIZE + + +@pytest.mark.parametrize("multi_value_headers", [False, True]) +def test_response_size_includes_serialized_headers_and_cookies(alb_event, resolve_response, multi_value_headers): + if multi_value_headers: + alb_event["multiValueHeaders"] = {name: [value] for name, value in alb_event.pop("headers").items()} + body = "x" * (MAX_RESPONSE_SIZE - 500) + options = {"headers": {"X-Extra": "h" * 300}, "cookies": [Cookie("session", "c" * 300)]} + expected = resolve_response(make_app(body, enabled=False, **options), alb_event) + + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(make_app(body, **options), alb_event) + + assert len(body.encode("utf-8")) < MAX_RESPONSE_SIZE + assert raised.value.actual_size == response_size(expected) + assert raised.value.actual_size > MAX_RESPONSE_SIZE + headers_key = "multiValueHeaders" if multi_value_headers else "headers" + assert "Set-Cookie" in expected[headers_key] + + +def test_response_size_is_checked_after_base64_encoding(alb_event, resolve_response): + body = b"\x00" * (MAX_RESPONSE_SIZE * 3 // 4) + expected = resolve_response(make_app(body, enabled=False), alb_event) + assert expected["isBase64Encoded"] is True + assert len(body) < MAX_RESPONSE_SIZE + + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(make_app(body), alb_event) + + assert raised.value.actual_size == response_size(expected) + + +def test_compression_can_bring_a_response_below_the_limit(alb_event, resolve_response): + body = "x" * (MAX_RESPONSE_SIZE * 2) + response = resolve_response(make_app(body, compress=True), alb_event) + + assert response["headers"]["Content-Encoding"] == "gzip" + assert response["isBase64Encoded"] is True + assert gzip.decompress(base64.b64decode(response["body"])).decode() == body + assert response_size(response) < MAX_RESPONSE_SIZE + + +def test_compressed_responses_are_still_validated(alb_event, resolve_response): + body = random.Random(42).randbytes(MAX_RESPONSE_SIZE * 3 // 4) + expected = resolve_response(make_app(body, enabled=False, compress=True), alb_event) + assert expected["headers"]["Content-Encoding"] == "gzip" + + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(make_app(body, compress=True), alb_event) + + assert raised.value.actual_size == response_size(expected) + assert raised.value.actual_size > MAX_RESPONSE_SIZE + + +def test_custom_body_serializer_runs_before_response_size_validation(alb_event, resolve_response): + serialized_bodies = [] + + def serialize(body): + serialized_bodies.append(body) + return "x" * MAX_RESPONSE_SIZE + + app = ALBResolver(enable_response_size_validation=True, serializer=serialize) + + @app.get("/lambda") + def get_response(): + return {"small": "body"} + + with pytest.raises(ResponseSizeExceededError): + resolve_response(app, alb_event) + + assert serialized_bodies == [{"small": "body"}] + + +@pytest.mark.parametrize("exception_type", [ResponseSizeExceededError, Exception]) +def test_registered_handler_receives_size_error_with_route_context(alb_event, resolve_response, exception_type): + origin = "https://example.com" + alb_event["headers"]["origin"] = origin + app = ALBResolver(cors=CORSConfig(allow_origin=origin), enable_response_size_validation=True) + handled = [] + + @app.get("/lambda", cors=True, cache_control="max-age=300") + def get_response(): + return Response(200, content_types.TEXT_PLAIN, "x" * MAX_RESPONSE_SIZE) + + @app.exception_handler(exception_type) + def handle_large_response(error): + handled.append(error) + assert app.context["_path"] == "/lambda" + assert app.current_event.path == "/lambda" + return Response(500, content_types.APPLICATION_JSON, {"message": "Response too large"}) + + response = resolve_response(app, alb_event) + + assert response["statusCode"] == 500 + assert json.loads(response["body"]) == {"message": "Response too large"} + assert response["headers"]["Access-Control-Allow-Origin"] == origin + assert response["headers"]["Cache-Control"] == "no-cache" + assert len(handled) == 1 + assert isinstance(handled[0], ResponseSizeExceededError) + assert handled[0].actual_size > handled[0].max_size == MAX_RESPONSE_SIZE + assert app.context == {} + + +@pytest.mark.parametrize("debug", [False, True]) +def test_unhandled_size_errors_propagate_even_in_debug_mode(alb_event, resolve_response, debug): + app = ALBResolver(enable_response_size_validation=True, debug=debug) + + @app.get("/lambda") + def get_response(): + return Response(200, content_types.TEXT_PLAIN, "customer-secret" + "x" * MAX_RESPONSE_SIZE) + + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(app, alb_event) + + assert str(raised.value.actual_size) in str(raised.value) + assert str(MAX_RESPONSE_SIZE) in str(raised.value) + assert "customer-secret" not in str(raised.value) + assert app.context == {} + + +def test_oversized_handler_response_does_not_recurse(alb_event, resolve_response): + app = make_app("x" * MAX_RESPONSE_SIZE) + handled = [] + + @app.exception_handler(ResponseSizeExceededError) + def handle_large_response(error): + handled.append(error) + return Response(500, content_types.TEXT_PLAIN, "y" * (MAX_RESPONSE_SIZE + 10)) + + with pytest.raises(ResponseSizeExceededError) as raised: + resolve_response(app, alb_event) + + assert len(handled) == 1 + assert raised.value.actual_size == handled[0].actual_size + 10 + assert app.context == {} + + +def test_exceptions_raised_by_size_handler_propagate(alb_event, resolve_response): + app = make_app("x" * MAX_RESPONSE_SIZE) + + @app.exception_handler(ResponseSizeExceededError) + def handle_large_response(error): + raise ValueError("handler failure") + + with pytest.raises(ValueError, match="handler failure"): + resolve_response(app, alb_event) + + assert app.context == {} + + +def test_size_handler_can_raise_a_service_error(alb_event, resolve_response): + app = make_app("x" * MAX_RESPONSE_SIZE) + + @app.exception_handler(ResponseSizeExceededError) + def handle_large_response(error): + raise InternalServerError("Response too large") + + response = resolve_response(app, alb_event) + + assert response["statusCode"] == 500 + assert json.loads(response["body"]) == {"statusCode": 500, "message": "Response too large"} + + +def test_size_handler_empty_body_is_normalized_for_alb(alb_event, resolve_response): + app = make_app("x" * MAX_RESPONSE_SIZE) + + @app.exception_handler(ResponseSizeExceededError) + def handle_large_response(error): + return Response(status_code=500) + + response = resolve_response(app, alb_event) + + assert response["statusCode"] == 500 + assert response["body"] == "" + + +def test_not_found_responses_are_validated(alb_event, resolve_response): + app = ALBResolver(enable_response_size_validation=True) + + @app.not_found + def handle_not_found(error): + return Response(404, content_types.TEXT_PLAIN, "x" * MAX_RESPONSE_SIZE) + + with pytest.raises(ResponseSizeExceededError): + resolve_response(app, alb_event) + + +def test_responses_from_other_exception_handlers_are_validated(alb_event, resolve_response): + app = ALBResolver(enable_response_size_validation=True) + handled = [] + + @app.get("/lambda") + def get_response(): + raise ValueError("route failure") + + @app.exception_handler(ValueError) + def handle_route_error(error): + return Response(500, content_types.TEXT_PLAIN, "x" * MAX_RESPONSE_SIZE) + + @app.exception_handler(ResponseSizeExceededError) + def handle_large_response(error): + handled.append(error) + return Response(500, content_types.TEXT_PLAIN, "Small error response") + + response = resolve_response(app, alb_event) + + assert response["body"] == "Small error response" + assert len(handled) == 1 + + +def test_middleware_response_changes_are_validated(alb_event, resolve_response): + app = make_app("small") + + def enlarge_response(app, next_middleware): + response = next_middleware(app) + response.body = "x" * MAX_RESPONSE_SIZE + return response + + app.use([enlarge_response]) + + with pytest.raises(ResponseSizeExceededError): + resolve_response(app, alb_event) + + +def test_native_async_handler_size_error_can_be_handled(alb_event): + app = ALBResolver(enable_response_size_validation=True) + + @app.get("/lambda") + async def get_response(): + await asyncio.sleep(0) + return Response(200, content_types.TEXT_PLAIN, "x" * MAX_RESPONSE_SIZE) + + @app.exception_handler(ResponseSizeExceededError) + def handle_large_response(error): + return Response(500, content_types.TEXT_PLAIN, "Small error response") + + response = asyncio.run(app.resolve_async(alb_event, {})) + + assert response["statusCode"] == 500 + assert response["body"] == "Small error response" + + +@pytest.mark.parametrize( + "resolver,event_name", + [ + (APIGatewayRestResolver, "apiGatewayProxyEvent.json"), + (APIGatewayHttpResolver, "apiGatewayProxyV2Event_GET.json"), + ], +) +def test_alb_response_limit_does_not_apply_to_api_gateway(resolver, event_name, resolve_response): + app = resolver() + event = load_event(event_name) + if "rawPath" in event: + event["rawPath"] = "/large" + event["requestContext"]["http"]["path"] = "/large" + event["requestContext"]["http"]["method"] = "GET" + else: + event["path"] = "/large" + event["httpMethod"] = "GET" + + @app.get("/large") + def get_response(): + return Response(200, content_types.TEXT_PLAIN, "x" * MAX_RESPONSE_SIZE) + + response = resolve_response(app, event) + + assert response["statusCode"] == 200 + assert response_size(response) > MAX_RESPONSE_SIZE