fix(local-apigw): apply GatewayResponses headers to authorizer and Lambda-exception error responses

samcli/lib/providers/api_collector.py

@@ -34,6 +34,7 @@ def __init__(self) -> None:
         self.stage_name: Optional[str] = None
         self.stage_variables: Optional[Dict] = None
         self.cors: Optional[Cors] = None
+        self.gateway_responses: Optional[Dict[str, Dict[str, str]]] = None
 
     def __iter__(self) -> Iterator[Tuple[str, List[Route]]]:
         """
@@ -200,6 +201,7 @@ def get_api(self) -> Api:
         api.stage_name = self.stage_name
         api.stage_variables = self.stage_variables
         api.cors = self.cors
+        api.gateway_responses = self.gateway_responses
 
         for authorizers in self._authorizers_per_resources.values():
             if len(authorizers):

samcli/lib/providers/provider.py

@@ -499,6 +499,9 @@ def __eq__(self, other: object) -> bool:
 
 
 class Api:
+    _HTTP_STATUS_CODE_4XX_LOWER_BOUND = 400
+    _HTTP_STATUS_CODE_4XX_UPPER_BOUND = 500
+
     def __init__(self, routes: Optional[Union[List["Route"], Set[str]]] = None) -> None:
         if routes is None:
             routes = []
@@ -517,6 +520,11 @@ def __init__(self, routes: Optional[Union[List["Route"], Set[str]]] = None) -> N
         self.stage_name: Optional[str] = None
         self.stage_variables: Optional[Dict] = None
 
+        # Headers configured through the API's GatewayResponses property, keyed by response type
+        # (e.g. "UNAUTHORIZED", "ACCESS_DENIED", "DEFAULT_5XX"). Only REST APIs (AWS::Serverless::Api)
+        # support GatewayResponses.
+        self.gateway_responses: Optional[Dict[str, Dict[str, str]]] = None
+
     def __hash__(self) -> int:
         # Other properties are not a part of the hash
         return hash(self.routes) * hash(self.cors) * hash(self.binary_media_types_set)
@@ -525,6 +533,32 @@ def __hash__(self) -> int:
     def binary_media_types(self) -> List[str]:
         return list(self.binary_media_types_set)
 
+    def get_gateway_response_headers(self, response_type: str, status_code: int) -> Dict[str, str]:
+        """
+        Returns the headers configured through this API's GatewayResponses property for the given
+        response type. If the specific type isn't configured, falls back to the generic DEFAULT_4XX
+        or DEFAULT_5XX response (matching the fallback behavior of API Gateway itself).
+
+        Parameters
+        ----------
+        response_type : str
+            The API Gateway response type this response corresponds to (e.g. "UNAUTHORIZED")
+        status_code : int
+            The HTTP status code being returned, used to select the DEFAULT_4XX/DEFAULT_5XX fallback
+
+        Returns
+        -------
+        dict
+            Headers configured for this response type, or an empty dict if none are configured
+        """
+        if not self.gateway_responses:
+            return {}
+
+        is_4xx = self._HTTP_STATUS_CODE_4XX_LOWER_BOUND <= status_code < self._HTTP_STATUS_CODE_4XX_UPPER_BOUND
+        default_response_type = "DEFAULT_4XX" if is_4xx else "DEFAULT_5XX"
+        headers = self.gateway_responses.get(response_type) or self.gateway_responses.get(default_response_type)
+        return dict(headers) if headers else {}
+
 
 _CorsTuple = namedtuple(
     "_CorsTuple", ["allow_origin", "allow_methods", "allow_headers", "allow_credentials", "max_age"]

samcli/lib/providers/sam_api_provider.py

@@ -140,6 +140,7 @@ def _extract_from_serverless_api(
         cors = self.extract_cors(properties.get("Cors", {}))
         stage_name = properties.get("StageName")
         stage_variables = properties.get("Variables")
+        gateway_responses = self._extract_gateway_responses(logical_id, properties.get("GatewayResponses", {}))
         if not body and not uri:
             # Swagger is not found anywhere.
             LOG.debug(
@@ -152,6 +153,7 @@ def _extract_from_serverless_api(
         collector.stage_name = stage_name
         collector.stage_variables = stage_variables
         collector.cors = cors
+        collector.gateway_responses = gateway_responses
 
         auth = properties.get(SamApiProvider._AUTH, {})
         if not auth or disable_authorizer:
@@ -164,6 +166,65 @@ def _extract_from_serverless_api(
 
         self._extract_authorizers_from_props(logical_id, auth, collector, Route.API)
 
+    @staticmethod
+    def _extract_gateway_responses(logical_id: str, gateway_responses: Dict) -> Dict[str, Dict[str, str]]:
+        """
+        Extracts the Headers configured under each GatewayResponse's ResponseParameters so they can be
+        applied to local error responses (e.g. Lambda Authorizer failures, unhandled Lambda exceptions)
+        the same way API Gateway applies them.
+
+        Parameters
+        ----------
+        logical_id: str
+            Logical ID of the AWS::Serverless::Api resource, used for warning messages
+        gateway_responses: dict
+            The GatewayResponses property dictionary, keyed by response type (e.g. "UNAUTHORIZED")
+
+        Returns
+        -------
+        dict
+            Dictionary of response type to a dictionary of header name/value pairs
+        """
+        if not gateway_responses or not isinstance(gateway_responses, dict):
+            return {}
+
+        extracted_responses: Dict[str, Dict[str, str]] = {}
+
+        for response_type, response in gateway_responses.items():
+            if not isinstance(response, dict):
+                LOG.warning(
+                    "Ignoring GatewayResponse '%s' on resource '%s', must be an object", response_type, logical_id
+                )
+                continue
+
+            response_parameters = response.get("ResponseParameters") or {}
+            if not isinstance(response_parameters, dict):
+                continue
+
+            response_headers = response_parameters.get("Headers") or {}
+            if not isinstance(response_headers, dict):
+                continue
+
+            extracted_headers = {}
+            for header_name, header_value in response_headers.items():
+                if not isinstance(header_value, str) or not (
+                    header_value.startswith("'") and header_value.endswith("'")
+                ):
+                    LOG.warning(
+                        "Ignoring GatewayResponses header '%s' for '%s' on resource '%s'. Header values must be "
+                        'a quoted string (i.e. "\'value\'" is correct, but "value" is not).',
+                        header_name,
+                        response_type,
+                        logical_id,
+                    )
+                    continue
+                extracted_headers[header_name] = header_value.strip("'")
+
+            if extracted_headers:
+                extracted_responses[response_type] = extracted_headers
+
+        return extracted_responses
+
     @staticmethod
     def _extract_request_lambda_authorizer(
         auth_name: str, function_name: str, prefix: str, properties: dict, event_type: str

samcli/local/apigw/local_apigw_service.py

@@ -678,7 +678,9 @@ def _request_handler(self, **kwargs):
 
         # check for LambdaAuthorizer since that is the only authorizer we currently support
         if isinstance(lambda_authorizer, LambdaAuthorizer) and not self._valid_identity_sources(request, route):
-            return ServiceErrorResponses.missing_lambda_auth_identity_sources()
+            return ServiceErrorResponses.missing_lambda_auth_identity_sources(
+                headers=self.api.get_gateway_response_headers("UNAUTHORIZED", 401)
+            )
 
         try:
             route_lambda_event = self._generate_lambda_event(request, route, method, endpoint)
@@ -688,7 +690,9 @@ def _request_handler(self, **kwargs):
                 auth_lambda_event = self._generate_lambda_authorizer_event(request, route, lambda_authorizer)
         except UnicodeDecodeError as error:
             LOG.error("UnicodeDecodeError while processing HTTP request: %s", error)
-            return ServiceErrorResponses.lambda_failure_response()
+            return ServiceErrorResponses.lambda_failure_response(
+                headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 502)
+            )
 
         lambda_authorizer_exception = None
         try:
@@ -697,10 +701,14 @@ def _request_handler(self, **kwargs):
             if lambda_authorizer:
                 self._invoke_parse_lambda_authorizer(lambda_authorizer, auth_lambda_event, route_lambda_event, route)
         except AuthorizerUnauthorizedRequest as ex:
-            auth_service_error = ServiceErrorResponses.lambda_authorizer_unauthorized()
+            auth_service_error = ServiceErrorResponses.lambda_authorizer_unauthorized(
+                headers=self.api.get_gateway_response_headers("ACCESS_DENIED", 403)
+            )
             lambda_authorizer_exception = ex
         except InvalidLambdaAuthorizerResponse as ex:
-            auth_service_error = ServiceErrorResponses.lambda_failure_response()
+            auth_service_error = ServiceErrorResponses.lambda_failure_response(
+                headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 502)
+            )
             lambda_authorizer_exception = ex
         except FunctionNotFound as ex:
             lambda_authorizer_exception = ex
@@ -742,20 +750,30 @@ def _request_handler(self, **kwargs):
             # invoke the route's Lambda function
             lambda_response = self._invoke_lambda_function(route.function_name, route_lambda_event, tenant_id)
         except TenantIdValidationError as e:
-            endpoint_service_error = ServiceErrorResponses.tenant_id_validation_error(str(e))
+            endpoint_service_error = ServiceErrorResponses.tenant_id_validation_error(
+                str(e), headers=self.api.get_gateway_response_headers("DEFAULT_4XX", 400)
+            )
         except FunctionNotFound:
-            endpoint_service_error = ServiceErrorResponses.lambda_not_found_response()
+            endpoint_service_error = ServiceErrorResponses.lambda_not_found_response(
+                headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 502)
+            )
         except UnsupportedInlineCodeError:
             endpoint_service_error = ServiceErrorResponses.not_implemented_locally(
-                "Inline code is not supported for sam local commands. Please write your code in a separate file."
+                "Inline code is not supported for sam local commands. Please write your code in a separate file.",
+                headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 501),
             )
         except LambdaResponseParseException:
-            endpoint_service_error = ServiceErrorResponses.lambda_body_failure_response()
+            endpoint_service_error = ServiceErrorResponses.lambda_body_failure_response(
+                headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 500)
+            )
         except DockerContainerCreationFailedException as ex:
-            endpoint_service_error = ServiceErrorResponses.container_creation_failed(ex.message)
+            endpoint_service_error = ServiceErrorResponses.container_creation_failed(
+                ex.message, headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 501)
+            )
         except MissingFunctionNameException as ex:
             endpoint_service_error = ServiceErrorResponses.lambda_failure_response(
                 f"Failed to execute endpoint. Got an invalid function name ({str(ex)})",
+                headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 502),
             )
 
         if endpoint_service_error:
@@ -774,7 +792,9 @@ def _request_handler(self, **kwargs):
                 )
         except LambdaResponseParseException as ex:
             LOG.error("Invalid lambda response received: %s", ex)
-            return ServiceErrorResponses.lambda_failure_response()
+            return ServiceErrorResponses.lambda_failure_response(
+                headers=self.api.get_gateway_response_headers("DEFAULT_5XX", 502)
+            )
 
         # Add CORS headers to the response
         headers.update(cors_headers)

samcli/local/apigw/service_error_responses.py

@@ -1,6 +1,7 @@
 """Class container to hold common Service Responses"""
 
 import logging
+from typing import Optional
 
 from flask import Response, jsonify, make_response
 
@@ -22,7 +23,21 @@ class ServiceErrorResponses:
     HTTP_STATUS_CODE_400 = 400
 
     @staticmethod
-    def lambda_authorizer_unauthorized() -> Response:
+    def _add_headers(response: Response, headers: Optional[dict]) -> Response:
+        """
+        Adds the given headers (e.g. from a template's GatewayResponses configuration) to a Flask response
+
+        Returns
+        -------
+        Response
+            The same Flask Response object, with headers added
+        """
+        if headers:
+            response.headers.update(headers)
+        return response
+
+    @staticmethod
+    def lambda_authorizer_unauthorized(headers: Optional[dict] = None) -> Response:
         """
         Constructs a Flask response for when a route invokes a Lambda Authorizer, but
         is the identity sources provided are not authorized for that method
@@ -33,10 +48,11 @@ def lambda_authorizer_unauthorized() -> Response:
             A Flask Response object
         """
         response_data = jsonify(ServiceErrorResponses._LAMBDA_AUTHORIZER_NOT_AUTHORIZED)
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_403)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_403)
+        return ServiceErrorResponses._add_headers(response, headers)
 
     @staticmethod
-    def missing_lambda_auth_identity_sources() -> Response:
+    def missing_lambda_auth_identity_sources(headers: Optional[dict] = None) -> Response:
         """
         Constructs a Flask response for when a route contains a Lambda Authorizer
         but is missing the required identity services
@@ -47,31 +63,34 @@ def missing_lambda_auth_identity_sources() -> Response:
             A Flask Response object
         """
         response_data = jsonify(ServiceErrorResponses._MISSING_LAMBDA_AUTH_IDENTITY_SOURCES)
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_401)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_401)
+        return ServiceErrorResponses._add_headers(response, headers)
 
     @staticmethod
-    def lambda_failure_response(*args):
+    def lambda_failure_response(*args, headers: Optional[dict] = None):
         """
         Helper function to create a Lambda Failure Response
 
         :return: A Flask Response
         """
         LOG.debug("Lambda execution failed %s", args)
         response_data = jsonify(ServiceErrorResponses._LAMBDA_FAILURE)
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_502)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_502)
+        return ServiceErrorResponses._add_headers(response, headers)
 
     @staticmethod
-    def lambda_body_failure_response(*args):
+    def lambda_body_failure_response(*args, headers: Optional[dict] = None):
         """
         Helper function to create a Lambda Body Failure Response
 
         :return: A Flask Response
         """
         response_data = jsonify(ServiceErrorResponses._LAMBDA_FAILURE)
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_500)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_500)
+        return ServiceErrorResponses._add_headers(response, headers)
 
     @staticmethod
-    def not_implemented_locally(message):
+    def not_implemented_locally(message, headers: Optional[dict] = None):
         """
         Constructs a Flask Response for for when a Lambda function functionality is
         not implemented
@@ -80,17 +99,19 @@ def not_implemented_locally(message):
         """
         exception_dict = {"message": message}
         response_data = jsonify(exception_dict)
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_501)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_501)
+        return ServiceErrorResponses._add_headers(response, headers)
 
     @staticmethod
-    def lambda_not_found_response(*args):
+    def lambda_not_found_response(*args, headers: Optional[dict] = None):
         """
         Constructs a Flask Response for when a Lambda function is not found for an endpoint
 
         :return: a Flask Response
         """
         response_data = jsonify(ServiceErrorResponses._NO_LAMBDA_INTEGRATION)
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_502)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_502)
+        return ServiceErrorResponses._add_headers(response, headers)
 
     @staticmethod
     def route_not_found(*args):
@@ -104,16 +125,17 @@ def route_not_found(*args):
         return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_403)
 
     @staticmethod
-    def container_creation_failed(message):
+    def container_creation_failed(message, headers: Optional[dict] = None):
         """
         Constuct a Flask Response for when container creation fails for a Lambda Function
         :return: a Flask Response
         """
         response_data = jsonify({"message": message})
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_501)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_501)
+        return ServiceErrorResponses._add_headers(response, headers)
 
     @staticmethod
-    def tenant_id_validation_error(message: str) -> Response:
+    def tenant_id_validation_error(message: str, headers: Optional[dict] = None) -> Response:
         """
         Constructs a Flask Response for when tenant ID validation fails
 
@@ -128,4 +150,5 @@ def tenant_id_validation_error(message: str) -> Response:
             A Flask Response object with HTTP 400 status
         """
         response_data = jsonify({"message": message})
-        return make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_400)
+        response = make_response(response_data, ServiceErrorResponses.HTTP_STATUS_CODE_400)
+        return ServiceErrorResponses._add_headers(response, headers)