Move CORS implementation into lambda (#1)

Signed-off-by: Dan R <[email protected]>

deploy/dkms_api.py

@@ -82,6 +82,7 @@ def deploy_dkms_lambda(self) -> lambda_.Function:
             environment={
                 "DKMS_KMS_KEY_ID": self.kms_key.key_id,
                 "JWKS_URL": self.jwks_url,
+                "CORS_ALLOW_ORIGINS": self.cors_allow_origins,
             },
         )
 
@@ -110,31 +111,20 @@ def deploy_dkms_api(self) -> apigwv2.HttpApi:
             "dkms-customer-api",
             api_name=f"dkms-customer-api-{self.env_name}",
             default_domain_mapping=default_domain_mapping,
-            cors_preflight=apigwv2.CorsPreflightOptions(
-                allow_headers=["Authorization"],
-                allow_methods=[
-                    apigwv2.CorsHttpMethod.GET,
-                    apigwv2.CorsHttpMethod.HEAD,
-                    apigwv2.CorsHttpMethod.OPTIONS,
-                    apigwv2.CorsHttpMethod.POST,
-                ],
-                allow_origins=[self.cors_allow_origins],
-                max_age=Duration.days(10),
-            ),
         )
         dkms_api.add_routes(
             path="/healthz",
-            methods=[apigwv2.HttpMethod.GET],
+            methods=[apigwv2.HttpMethod.GET, apigwv2.HttpMethod.OPTIONS],
             integration=dkms_default_integration,
         )
         dkms_api.add_routes(
             path="/encrypt",
-            methods=[apigwv2.HttpMethod.POST],
+            methods=[apigwv2.HttpMethod.POST, apigwv2.HttpMethod.OPTIONS],
             integration=dkms_default_integration,
         )
         dkms_api.add_routes(
             path="/decrypt",
-            methods=[apigwv2.HttpMethod.POST],
+            methods=[apigwv2.HttpMethod.POST, apigwv2.HttpMethod.OPTIONS],
             integration=dkms_default_integration,
         )
         return dkms_api

lambdas/dkms_handler/index.py

@@ -19,6 +19,18 @@
 assert jwks_url is not None, "JWKS_URL environment variable must be set"
 jwks_client = jwt.PyJWKClient(jwks_url)
 
+cors_allow_origins = os.getenv("CORS_ALLOW_ORIGINS", None)
+assert (
+    cors_allow_origins is not None
+), "CORS_ALLOW_ORIGINS environment variable must be set"
+
+cors_headers = {
+    "Access-Control-Allow-Origin": cors_allow_origins,
+    "Access-Control-Allow-Credentials": True,
+    "Access-Control-Allow-Headers": "Content-Type,Authorization",
+    "Access-Control-Allow-Methods": "OPTIONS,GET,POST",
+}
+
 
 class AuthenticationError(Exception):
     """Raised when authentication fails."""
@@ -69,7 +81,9 @@ def router(event) -> dict:
     http_method = event["requestContext"]["http"]["method"]
     path = event["rawPath"]
 
-    if http_method == "GET" and path == "/healthz":
+    if http_method == "OPTIONS":
+        return return_options_handler()
+    elif http_method == "GET" and path == "/healthz":
         return return_handler(status=HTTPStatus.OK)
     elif http_method == "POST" and path == "/encrypt":
         payload = authenticate(event)
@@ -154,6 +168,21 @@ def decrypt(body: str, kms_key_id: str, encryption_context: dict) -> dict:
     return return_handler(status=HTTPStatus.OK, data={"plaintext": decrypted_data})
 
 
+def return_options_handler() -> dict:
+    """Return an OPTIONS request."""
+    status = HTTPStatus.OK
+    logger.info(
+        {
+            "status": status.name,
+            "statusCode": status.value,
+        }
+    )
+    return {
+        "statusCode": status.value,
+        "headers": {**cors_headers},
+    }
+
+
 def return_handler(
     data: dict = {},
     error_code: str = "",
@@ -172,7 +201,7 @@ def return_handler(
     )
     return {
         "statusCode": status.value,
-        "headers": {"Content-Type": "application/json"},
+        "headers": {"Content-Type": "application/json", **cors_headers},
         "body": json.dumps(
             {
                 "data": data,

lambdas/dkms_handler/tests/conftest.py

@@ -10,6 +10,7 @@
 os.environ[
     "JWKS_URL"
 ] = "https://example.com/.well-known/jwks.json"  # this can be anything because we mock the response in a fixture
+os.environ["CORS_ALLOW_ORIGINS"] = "*"
 
 # openssl genrsa -out private_key.pem 2048
 test_private_key = """-----BEGIN PRIVATE KEY-----

lambdas/dkms_handler/tests/test_dkms_handler.py

@@ -17,7 +17,13 @@ def test_lambda_dkms_healthz():
     response = index.handler(event, context)
     assert response == {
         "statusCode": 200,
-        "headers": {"Content-Type": "application/json"},
+        "headers": {
+            "Content-Type": "application/json",
+            "Access-Control-Allow-Credentials": True,
+            "Access-Control-Allow-Headers": "Content-Type,Authorization",
+            "Access-Control-Allow-Methods": "OPTIONS,GET,POST",
+            "Access-Control-Allow-Origin": "*",
+        },
         "body": '{"data": {}, "error_code": "", "message": "", "status": "OK"}',
     }
 
@@ -246,3 +252,41 @@ def test_authenticate_jwt(user_jwt):
     }
     payload = index.authenticate(event)
     assert payload == {"sub": "test_user", "ewi": "abcd1234"}
+
+
+def test_options_request_healthz():
+    event = {
+        "rawPath": "/healthz",
+        "requestContext": {"http": {"method": "OPTIONS"}},
+    }
+    context = None
+
+    response = index.handler(event, context)
+    assert response == {
+        "headers": {
+            "Access-Control-Allow-Credentials": True,
+            "Access-Control-Allow-Headers": "Content-Type,Authorization",
+            "Access-Control-Allow-Methods": "OPTIONS,GET,POST",
+            "Access-Control-Allow-Origin": "*",
+        },
+        "statusCode": 200,
+    }
+
+
+def test_options_request_encrypt():
+    event = {
+        "rawPath": "/encrypt",
+        "requestContext": {"http": {"method": "OPTIONS"}},
+    }
+    context = None
+
+    response = index.handler(event, context)
+    assert response == {
+        "headers": {
+            "Access-Control-Allow-Credentials": True,
+            "Access-Control-Allow-Headers": "Content-Type,Authorization",
+            "Access-Control-Allow-Methods": "OPTIONS,GET,POST",
+            "Access-Control-Allow-Origin": "*",
+        },
+        "statusCode": 200,
+    }