APIGW NG: add LocalStack-only CORS handling for REST API (#11764)

bentsku · web-flow · commit fb7cde9cc16f · 2024-10-30T11:45:59.000+01:00
diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/gateway.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/gateway.py
@@ -32,15 +32,16 @@ def __init__(self):
                 handlers.method_response_handler,
             ]
         )
-        self.response_handlers.extend(
+        self.exception_handlers.extend(
             [
-                handlers.response_enricher
-                # add composite response handlers?
+                handlers.gateway_exception_handler,
             ]
         )
-        self.exception_handlers.extend(
+        self.response_handlers.extend(
             [
-                handlers.gateway_exception_handler,
+                handlers.response_enricher,
+                handlers.cors_response_enricher,
+                # add composite response handlers?
             ]
         )
 
diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/__init__.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/__init__.py
@@ -1,6 +1,7 @@
 from rolo.gateway import CompositeHandler
 
 from .api_key_validation import ApiKeyValidationHandler
+from .cors import CorsResponseEnricher
 from .gateway_exception import GatewayExceptionHandler
 from .integration import IntegrationHandler
 from .integration_request import IntegrationRequestHandler
@@ -23,3 +24,4 @@
 gateway_exception_handler = GatewayExceptionHandler()
 api_key_validation_handler = ApiKeyValidationHandler()
 response_enricher = InvocationResponseEnricher()
+cors_response_enricher = CorsResponseEnricher()
diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/cors.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/cors.py
@@ -0,0 +1,49 @@
+import logging
+from http import HTTPMethod
+
+from localstack import config
+from localstack.aws.handlers.cors import CorsEnforcer
+from localstack.aws.handlers.cors import CorsResponseEnricher as GlobalCorsResponseEnricher
+from localstack.http import Response
+
+from ..api import RestApiGatewayHandler, RestApiGatewayHandlerChain
+from ..context import RestApiInvocationContext
+from ..gateway_response import MissingAuthTokenError
+
+LOG = logging.getLogger(__name__)
+
+
+class CorsResponseEnricher(RestApiGatewayHandler):
+    def __call__(
+        self,
+        chain: RestApiGatewayHandlerChain,
+        context: RestApiInvocationContext,
+        response: Response,
+    ):
+        """
+        This is a LocalStack only handler, to allow users to override API Gateway CORS configuration and just use the
+        default LocalStack configuration instead, to ease the usage and reduce production code changes.
+        """
+        if not config.DISABLE_CUSTOM_CORS_APIGATEWAY:
+            return
+
+        if not context.invocation_request:
+            return
+
+        headers = context.invocation_request["headers"]
+
+        if "Origin" not in headers:
+            return
+
+        if context.request.method == HTTPMethod.OPTIONS:
+            # If the user did not configure an OPTIONS route, we still want LocalStack to properly respond to CORS
+            # requests
+            if context.invocation_exception:
+                if isinstance(context.invocation_exception, MissingAuthTokenError):
+                    response.data = b""
+                    response.status_code = 204
+                else:
+                    return
+
+        if CorsEnforcer.is_cors_origin_allowed(headers):
+            GlobalCorsResponseEnricher.add_cors_headers(headers, response.headers)

Original file line number	Diff line number	Diff line change
`@@ -32,15 +32,16 @@ def __init__(self):`
`32`	`32`	`handlers.method_response_handler,`
`33`	`33`	`]`
`34`	`34`	`)`
`35`		`- self.response_handlers.extend(`
	`35`	`+ self.exception_handlers.extend(`
`36`	`36`	`[`
`37`		`- handlers.response_enricher`
`38`		`- # add composite response handlers?`
	`37`	`+ handlers.gateway_exception_handler,`
`39`	`38`	`]`
`40`	`39`	`)`
`41`		`- self.exception_handlers.extend(`
	`40`	`+ self.response_handlers.extend(`
`42`	`41`	`[`
`43`		`- handlers.gateway_exception_handler,`
	`42`	`+ handlers.response_enricher,`
	`43`	`+ handlers.cors_response_enricher,`
	`44`	`+ # add composite response handlers?`
`44`	`45`	`]`
`45`	`46`	`)`
`46`	`47`