Implement warmer and key rotation

Author: devksingh4

src/api/lambda.ts

@@ -10,7 +10,41 @@ const realHandler = awsLambdaFastify(app, {
   callbackWaitsForEmptyEventLoop: false,
   binaryMimeTypes: ["application/octet-stream", "application/vnd.apple.pkpass"],
 });
+
 type WarmerEvent = { action: "warmer" };
+
+/**
+ * Validates the origin verification header against the current and previous keys.
+ * @returns {boolean} `true` if the request is valid, otherwise `false`.
+ */
+const validateOriginHeader = (
+  originHeader: string | undefined,
+  currentKey: string,
+  previousKey: string | undefined,
+  previousKeyExpiresAt: string | undefined,
+): boolean => {
+  // 1. A header must exist to be valid.
+  if (!originHeader) {
+    return false;
+  }
+
+  // 2. Check against the current key first for an early return on the happy path.
+  if (originHeader === currentKey) {
+    return true;
+  }
+
+  // 3. If it's not the current key, check the previous key during the rotation window.
+  if (previousKey && previousKeyExpiresAt) {
+    const isExpired = new Date() >= new Date(previousKeyExpiresAt);
+    if (originHeader === previousKey && !isExpired) {
+      return true;
+    }
+  }
+
+  // 4. If all checks fail, the header is invalid.
+  return false;
+};
+
 const handler = async (
   event: APIGatewayEvent | WarmerEvent,
   context: Context,
@@ -19,12 +53,32 @@ const handler = async (
     return { instanceId };
   }
   event = event as APIGatewayEvent;
-  if (process.env.ORIGIN_VERIFY_KEY) {
-    // check that the request has the right header (coming from cloudfront)
-    if (
-      !event.headers ||
-      !(event.headers["x-origin-verify"] === process.env.ORIGIN_VERIFY_KEY)
-    ) {
+
+  const currentKey = process.env.ORIGIN_VERIFY_KEY;
+  const previousKey = process.env.PREVIOUS_ORIGIN_VERIFY_KEY;
+  const previousKeyExpiresAt =
+    process.env.PREVIOUS_ORIGIN_VERIFY_KEY_EXPIRES_AT;
+
+  // Log an error if the previous key has expired but is still configured.
+  if (previousKey && previousKeyExpiresAt) {
+    if (new Date() >= new Date(previousKeyExpiresAt)) {
+      console.error(
+        "Expired previous origin verify key is still present in the environment. Expired at:",
+        previousKeyExpiresAt,
+      );
+    }
+  }
+
+  // Proceed with verification only if a current key is set.
+  if (currentKey) {
+    const isValid = validateOriginHeader(
+      event.headers?.["x-origin-verify"],
+      currentKey,
+      previousKey,
+      previousKeyExpiresAt,
+    );
+
+    if (!isValid) {
       const newError = new ValidationError({
         message: "Request is not valid.",
       });
@@ -38,9 +92,11 @@ const handler = async (
         isBase64Encoded: false,
       };
     }
+
     delete event.headers["x-origin-verify"];
   }
-  // else proceed with handler logic
+
+  // If verification is disabled or passed, proceed with the real handler logic.
   return await realHandler(event, context).catch((e) => {
     console.error(e);
     const newError = new InternalServerError({
@@ -58,5 +114,5 @@ const handler = async (
   });
 };
 
-await app.ready(); // needs to be placed after awsLambdaFastify call because of the decoration: https://github.com/fastify/aws-lambda-fastify/blob/master/index.js#L9
+await app.ready();
 export { handler };

terraform/envs/prod/.terraform.lock.hcl

@@ -45,10 +45,6 @@ module "sqs_queues" {
   core_sqs_consumer_lambda_name = module.lambdas.core_sqs_consumer_lambda_name
 }
 
-module "lambda_warmer" {
-  source           = "github.com/acm-uiuc/terraform-modules/lambda-warmer?ref=v1.0.1"
-  function_to_warm = module.lambdas.core_api_lambda_name
-}
 module "dynamo" {
   source    = "../../modules/dynamo"
   ProjectId = var.ProjectId
@@ -74,14 +70,15 @@ module "alarms" {
 }
 
 module "lambdas" {
-  source                  = "../../modules/lambdas"
-  ProjectId               = var.ProjectId
-  RunEnvironment          = "prod"
-  LinkryKvArn             = aws_cloudfront_key_value_store.linkry_kv.arn
-  CurrentOriginVerifyKey  = module.origin_verify.current_origin_verify_key
-  PreviousOriginVerifyKey = module.origin_verify.previous_origin_verify_key
-  LogRetentionDays        = 30
-  EmailDomain             = var.EmailDomain
+  source                           = "../../modules/lambdas"
+  ProjectId                        = var.ProjectId
+  RunEnvironment                   = "prod"
+  LinkryKvArn                      = aws_cloudfront_key_value_store.linkry_kv.arn
+  CurrentOriginVerifyKey           = module.origin_verify.current_origin_verify_key
+  PreviousOriginVerifyKey          = module.origin_verify.previous_origin_verify_key
+  PreviousOriginVerifyKeyExpiresAt = module.origin_verify.previous_invalid_time
+  LogRetentionDays                 = 30
+  EmailDomain                      = var.EmailDomain
 }
 
 module "frontend" {

terraform/envs/prod/main.tf

@@ -46,44 +46,39 @@ locals {
   }
 }
 
-
-module "lambda_warmer" {
-  source           = "github.com/acm-uiuc/terraform-modules/lambda-warmer?ref=v1.0.1"
-  function_to_warm = module.lambdas.core_api_lambda_name
-}
 module "dynamo" {
   source    = "../../modules/dynamo"
   ProjectId = var.ProjectId
 }
 
-resource "random_password" "origin_verify_key" {
-  length  = 16
-  special = false
-  keepers = {
-    force_recreation = formatdate("DD-MMM-YYYY", plantimestamp())
-  }
+module "origin_verify" {
+  source    = "../../modules/origin_verify"
+  ProjectId = var.ProjectId
 }
+
 resource "aws_cloudfront_key_value_store" "linkry_kv" {
   name = "${var.ProjectId}-cloudfront-linkry-kv"
 }
 
 
 module "lambdas" {
-  source           = "../../modules/lambdas"
-  ProjectId        = var.ProjectId
-  RunEnvironment   = "dev"
-  LinkryKvArn      = aws_cloudfront_key_value_store.linkry_kv.arn
-  OriginVerifyKey  = random_password.origin_verify_key.result
-  LogRetentionDays = 30
-  EmailDomain      = var.EmailDomain
+  source                           = "../../modules/lambdas"
+  ProjectId                        = var.ProjectId
+  RunEnvironment                   = "dev"
+  LinkryKvArn                      = aws_cloudfront_key_value_store.linkry_kv.arn
+  CurrentOriginVerifyKey           = module.origin_verify.current_origin_verify_key
+  PreviousOriginVerifyKey          = module.origin_verify.previous_origin_verify_key
+  PreviousOriginVerifyKeyExpiresAt = module.origin_verify.previous_invalid_time
+  LogRetentionDays                 = 30
+  EmailDomain                      = var.EmailDomain
 }
 
 module "frontend" {
   source             = "../../modules/frontend"
   BucketPrefix       = local.bucket_prefix
   CoreLambdaHost     = module.lambdas.core_function_url
   CoreSlowLambdaHost = module.lambdas.core_slow_function_url
-  OriginVerifyKey    = random_password.origin_verify_key.result
+  OriginVerifyKey    = module.origin_verify.current_origin_verify_key
   ProjectId          = var.ProjectId
   CoreCertificateArn = var.CoreCertificateArn
   CorePublicDomain   = var.CorePublicDomain

terraform/envs/qa/.terraform.lock.hcl

@@ -21,7 +21,7 @@ locals {
   api_cache_behaviors = {
     "/api/v1/syncIdentity" = {
       target_origin_id = "SlowLambdaFunction"
-      cache_policy_id  = "658327ea-f89d-4fab-a63d-7e88639e58f6"
+      cache_policy_id  = aws_cloudfront_cache_policy.no_cache.id
     },
     "/api/v1/organizations" = {
       target_origin_id = "LambdaFunction"

terraform/envs/qa/main.tf

@@ -341,6 +341,7 @@ resource "aws_lambda_function" "api_lambda" {
       "AWS_CRT_NODEJS_BINARY_RELATIVE_PATH" = "node_modules/aws-crt/dist/bin/linux-arm64-glibc/aws-crt-nodejs.node"
       ORIGIN_VERIFY_KEY                     = var.CurrentOriginVerifyKey
       PREVIOUS_ORIGIN_VERIFY_KEY            = var.PreviousOriginVerifyKey
+      PREVIOUS_ORIGIN_VERIFY_KEY_EXPIRES_AT = var.PreviousOriginVerifyKeyExpiresAt
       EntraRoleArn                          = aws_iam_role.entra_role.arn
       LinkryKvArn                           = var.LinkryKvArn
       "NODE_OPTIONS"                        = "--enable-source-maps"
@@ -401,6 +402,7 @@ resource "aws_lambda_function" "slow_lambda" {
       "AWS_CRT_NODEJS_BINARY_RELATIVE_PATH" = "node_modules/aws-crt/dist/bin/linux-arm64-glibc/aws-crt-nodejs.node"
       ORIGIN_VERIFY_KEY                     = var.CurrentOriginVerifyKey
       PREVIOUS_ORIGIN_VERIFY_KEY            = var.PreviousOriginVerifyKey
+      PREVIOUS_ORIGIN_VERIFY_KEY_EXPIRES_AT = var.PreviousOriginVerifyKeyExpiresAt
       EntraRoleArn                          = aws_iam_role.entra_role.arn
       LinkryKvArn                           = var.LinkryKvArn
       "NODE_OPTIONS"                        = "--enable-source-maps"
@@ -413,6 +415,19 @@ resource "aws_lambda_function_url" "slow_api_lambda_function_url" {
   authorization_type = "NONE"
 }
 
+module "lambda_warmer_main" {
+  source           = "github.com/acm-uiuc/terraform-modules/lambda-warmer?ref=v1.0.1"
+  function_to_warm = local.core_api_lambda_name
+}
+
+module "lambda_warmer_slow" {
+  source           = "github.com/acm-uiuc/terraform-modules/lambda-warmer?ref=v1.0.1"
+  function_to_warm = local.core_api_slow_lambda_name
+}
+
+
+// Outputs
+
 output "core_function_url" {
   value = replace(replace(aws_lambda_function_url.api_lambda_function_url.function_url, "https://", ""), "/", "")
 }

terraform/envs/qa/tfplan

@@ -26,6 +26,10 @@ variable "PreviousOriginVerifyKey" {
   sensitive = true
 }
 
+variable "PreviousOriginVerifyKeyExpiresAt" {
+  type = string
+}
+
 variable "LogRetentionDays" {
   type = number
 }