feat(event_handler): add opt-in ALB response size validation (#8490)

Author: leandrodamascena

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

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"""
 

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`.

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)