APIGW NG implement AWS bug for INTEGRATION_FAILURE GatewayResponse (#11183)

bentsku · web-flow · commit d3438159dfeb · 2024-07-12T12:40:51.000+02:00
diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/gateway_response.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/gateway_response.py
@@ -8,14 +8,38 @@
 )
 
 
+class GatewayResponseCode(StatusCode, Enum):
+    REQUEST_TOO_LARGE = "413"
+    RESOURCE_NOT_FOUND = "404"
+    AUTHORIZER_CONFIGURATION_ERROR = "500"
+    MISSING_AUTHENTICATION_TOKEN = "403"
+    BAD_REQUEST_BODY = "400"
+    INVALID_SIGNATURE = "403"
+    INVALID_API_KEY = "403"
+    BAD_REQUEST_PARAMETERS = "400"
+    AUTHORIZER_FAILURE = "500"
+    UNAUTHORIZED = "401"
+    INTEGRATION_TIMEOUT = "504"
+    ACCESS_DENIED = "403"
+    DEFAULT_4XX = ""
+    DEFAULT_5XX = ""
+    WAF_FILTERED = "403"
+    QUOTA_EXCEEDED = "429"
+    THROTTLED = "429"
+    API_CONFIGURATION_ERROR = "500"
+    UNSUPPORTED_MEDIA_TYPE = "415"
+    INTEGRATION_FAILURE = "504"
+    EXPIRED_TOKEN = "403"
+
+
 class BaseGatewayException(Exception):
     """
     Base class for all Gateway exceptions
     Do not raise from this class directly. Instead, raise the specific Exception
     """
 
     message: str = "Unimplemented Response"
-    type: GatewayResponseType
+    type: GatewayResponseType = None
     status_code: int | str = None
     code: str = ""
 
@@ -24,14 +48,17 @@ def __init__(self, message: str = None, status_code: int | str = None):
             self.message = message
         if status_code is not None:
             self.status_code = status_code
+        elif self.status_code is None and self.type:
+            # Fallback to the default value
+            self.status_code = GatewayResponseCode[self.type]
 
 
 class Default4xxError(BaseGatewayException):
     """Do not raise from this class directly.
     Use one of the subclasses instead, as they contain the appropriate header
     """
 
-    type: GatewayResponseType.DEFAULT_4XX
+    type = GatewayResponseType.DEFAULT_4XX
     status_code = 400
 
 
@@ -40,7 +67,7 @@ class Default5xxError(BaseGatewayException):
     Use one of the subclasses instead, as they contain the appropriate header
     """
 
-    type: GatewayResponseType.DEFAULT_5XX
+    type = GatewayResponseType.DEFAULT_5XX
     status_code = 500
 
 
@@ -98,8 +125,9 @@ class ExpiredTokenError(BaseGatewayException):
 
 class IntegrationFailureError(BaseGatewayException):
     type = GatewayResponseType.INTEGRATION_FAILURE
-    # TODO validate this header with aws validated tests
-    code = "IntegrationFailureException"
+    # TODO: tested manually for now
+    code = "InternalServerErrorException"
+    status_code = 500
 
 
 class IntegrationTimeoutError(BaseGatewayException):
@@ -161,30 +189,6 @@ class WafFilteredError(BaseGatewayException):
     code = "WafFilteredException"
 
 
-class GatewayResponseCode(StatusCode, Enum):
-    REQUEST_TOO_LARGE = "413"
-    RESOURCE_NOT_FOUND = "404"
-    AUTHORIZER_CONFIGURATION_ERROR = "500"
-    MISSING_AUTHENTICATION_TOKEN = "403"
-    BAD_REQUEST_BODY = "400"
-    INVALID_SIGNATURE = "403"
-    INVALID_API_KEY = "403"
-    BAD_REQUEST_PARAMETERS = "400"
-    AUTHORIZER_FAILURE = "500"
-    UNAUTHORIZED = "401"
-    INTEGRATION_TIMEOUT = "504"
-    ACCESS_DENIED = "403"
-    DEFAULT_4XX = ""
-    DEFAULT_5XX = ""
-    WAF_FILTERED = "403"
-    QUOTA_EXCEEDED = "429"
-    THROTTLED = "429"
-    API_CONFIGURATION_ERROR = "500"
-    UNSUPPORTED_MEDIA_TYPE = "415"
-    INTEGRATION_FAILURE = "504"
-    EXPIRED_TOKEN = "403"
-
-
 def build_gateway_response(
     response_type: GatewayResponseType,
     status_code: StatusCode = None,
@@ -199,10 +203,9 @@ def build_gateway_response(
         or {"application/json": '{"message":$context.error.messageString}'},
         responseType=response_type,
         defaultResponse=default_response,
+        statusCode=status_code,
     )
-    if status_code or (status_code := GatewayResponseCode[response_type]):
-        # DEFAULT_4XX and DEFAULT_5XX do not have `statusCode` present in the response
-        response["statusCode"] = status_code
+
     return response
 
 
diff --git a/localstack-core/localstack/services/apigateway/next_gen/provider.py b/localstack-core/localstack/services/apigateway/next_gen/provider.py
@@ -30,6 +30,7 @@
 from ..models import apigateway_stores
 from .execute_api.gateway_response import (
     DEFAULT_GATEWAY_RESPONSES,
+    GatewayResponseCode,
     build_gateway_response,
     get_gateway_response_or_default,
 )
@@ -163,8 +164,18 @@ def put_gateway_response(
             response_type=response_type,
             default_response=False,
         )
+
         rest_api_container.gateway_responses[response_type] = gateway_response
-        return gateway_response
+
+        # The CRUD provider has a weird behavior: for some responses (for now, INTEGRATION_FAILURE), it sets the default
+        # status code to `504`. However, in the actual invocation logic, it returns 500. To deal with the inconsistency,
+        # we need to set the value to None if not provided by the user, so that the invocation logic can properly return
+        # 500, and the CRUD layer can still return 504 even though it is technically wrong.
+        response = gateway_response.copy()
+        if response.get("statusCode") is None:
+            response["statusCode"] = GatewayResponseCode[response_type]
+
+        return response
 
     def get_gateway_response(
         self,
@@ -185,7 +196,7 @@ def get_gateway_response(
                 message=f"1 validation error detected: Value '{response_type}' at 'responseType' failed to satisfy constraint: Member must satisfy enum value set: [{', '.join(DEFAULT_GATEWAY_RESPONSES)}]",
             )
 
-        gateway_response = get_gateway_response_or_default(
+        gateway_response = _get_gateway_response_or_default(
             response_type, rest_api_container.gateway_responses
         )
         # TODO: add validation with the parameters? seems like it validated client side? how to try?
@@ -207,7 +218,23 @@ def get_gateway_responses(
 
         user_gateway_resp = rest_api_container.gateway_responses
         gateway_responses = [
-            get_gateway_response_or_default(response_type, user_gateway_resp)
+            _get_gateway_response_or_default(response_type, user_gateway_resp)
             for response_type in DEFAULT_GATEWAY_RESPONSES
         ]
         return GatewayResponses(items=gateway_responses)
+
+
+def _get_gateway_response_or_default(
+    response_type: GatewayResponseType,
+    gateway_responses: dict[GatewayResponseType, GatewayResponse],
+) -> GatewayResponse:
+    """
+    Utility function that overrides the behavior of `get_gateway_response_or_default` by setting a default status code
+    from the `GatewayResponseCode` values. In reality, some default values in the invocation layer are different from
+    what the CRUD layer of API Gateway is returning.
+    """
+    response = get_gateway_response_or_default(response_type, gateway_responses)
+    if response.get("statusCode") is None and (status_code := GatewayResponseCode[response_type]):
+        response["statusCode"] = status_code
+
+    return response
