From ea21bbf65e8563d8def22b314d1fa48a5259f4c1 Mon Sep 17 00:00:00 2001
From: Tim Holm <timothy.j.holm@gmail.com>
Date: Tue, 7 Oct 2025 10:00:29 +1100
Subject: [PATCH] fix: add 404 default behaviour is not default origin
 specified

---
 cloudfront/module/main.tf                  | 22 +++++++++++-
 cloudfront/module/outputs.tf               |  4 +++
 cloudfront/module/scripts/default-route.js | 12 +++++++
 cloudfront/module/shim.tf                  | 42 ++++++++++++++++++++++
 4 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 cloudfront/module/scripts/default-route.js
 create mode 100644 cloudfront/module/shim.tf

diff --git a/cloudfront/module/main.tf b/cloudfront/module/main.tf
index b7348fa..861d927 100644
--- a/cloudfront/module/main.tf
+++ b/cloudfront/module/main.tf
@@ -187,6 +187,8 @@ resource "aws_lambda_function" "origin_request" {
   }
 }
 
+
+
 resource "aws_lambda_permission" "allow_cloudfront_origin_request" {
   count         = length(local.lambda_origins) > 0 ? 1 : 0
   region        = "us-east-1"
@@ -435,6 +437,24 @@ resource "aws_cloudfront_distribution" "distribution" {
     }
   }
 
+  dynamic "origin" {
+    for_each = aws_lambda_function_url.default_origin_shim
+
+    content {
+      domain_name              = split("/", origin.value.function_url)[2]
+      origin_id                = "default-lambda-shim"
+      origin_access_control_id = aws_cloudfront_origin_access_control.lambda_oac[0].id
+
+      custom_origin_config {
+        origin_read_timeout    = 30
+        origin_protocol_policy = "https-only"
+        origin_ssl_protocols   = ["TLSv1.2", "SSLv3"]
+        http_port              = 80
+        https_port             = 443
+      }
+    }
+  }
+
   dynamic "origin" {
     for_each = local.vpc_origins
 
@@ -560,7 +580,7 @@ resource "aws_cloudfront_distribution" "distribution" {
   default_cache_behavior {
     allowed_methods        = ["GET", "HEAD", "OPTIONS", "PUT", "POST", "PATCH", "DELETE"]
     cached_methods         = ["GET", "HEAD", "OPTIONS"]
-    target_origin_id       = keys(local.default_origin)[0]
+    target_origin_id       = keys(local.actual_default_origin)[0]
     viewer_protocol_policy = "redirect-to-https"
 
     # Add Lambda@Edge for auth preservation and webhook signing (only if not a Lambda origin and Lambda@Edge exists)
diff --git a/cloudfront/module/outputs.tf b/cloudfront/module/outputs.tf
index e69de29..1e9ed21 100644
--- a/cloudfront/module/outputs.tf
+++ b/cloudfront/module/outputs.tf
@@ -0,0 +1,4 @@
+output "domain_name" {
+  value       = aws_cloudfront_distribution.distribution.domain_name
+  description = "The domain name of the cloudfront distribution."
+}
\ No newline at end of file
diff --git a/cloudfront/module/scripts/default-route.js b/cloudfront/module/scripts/default-route.js
new file mode 100644
index 0000000..237c61c
--- /dev/null
+++ b/cloudfront/module/scripts/default-route.js
@@ -0,0 +1,12 @@
+'use strict';
+
+exports.handler = async () => {
+    return {
+        statusCode: 404,
+        headers: {
+            'Content-Type': 'text/plain',
+            'Cache-Control': 'no-cache, no-store, must-revalidate'
+        },
+        body: '404 - Not Found'
+    };
+};
\ No newline at end of file
diff --git a/cloudfront/module/shim.tf b/cloudfront/module/shim.tf
new file mode 100644
index 0000000..a55a188
--- /dev/null
+++ b/cloudfront/module/shim.tf
@@ -0,0 +1,42 @@
+data "archive_file" "default_route_lambda" {
+  type        = "zip"
+  output_path = "${path.module}/default-route.zip"
+
+  source {
+    content  = file("${path.module}/scripts/default-route.js")
+    filename = "index.js"
+  }
+}
+
+resource "aws_lambda_function" "default_origin_shim" {
+  count            = length(local.default_origin) < 1 ? 1 : 0
+  filename         = data.archive_file.default_route_lambda.output_path
+  function_name    = "${var.suga.stack_id}-cloudfront-default-origin-shim"
+  role             = aws_iam_role.lambda_edge_origin_request[0].arn
+  handler          = "index.handler"
+  source_code_hash = data.archive_file.default_route_lambda.output_base64sha256
+  runtime          = "nodejs22.x"
+  timeout          = 5
+  memory_size      = 128
+  publish          = true
+}
+
+resource "aws_lambda_function_url" "default_origin_shim" {
+  count              = length(aws_lambda_function.default_origin_shim) > 0 ? 1 : 0
+  function_name      = aws_lambda_function.default_origin_shim[0].function_name
+  authorization_type = "AWS_IAM"
+  depends_on         = [aws_lambda_function.default_origin_shim[0]]
+}
+
+resource "aws_lambda_permission" "allow_cloudfront_to_execute_lambda_shim" {
+  count         = length(aws_lambda_function.default_origin_shim) > 0 ? 1 : 0
+  function_name = aws_lambda_function.default_origin_shim[0].function_name
+  principal     = "cloudfront.amazonaws.com"
+  action        = "lambda:InvokeFunctionUrl"
+  source_arn    = aws_cloudfront_distribution.distribution.arn
+  depends_on    = [aws_lambda_function.default_origin_shim[0], aws_cloudfront_distribution.distribution]
+}
+
+locals {
+  actual_default_origin = length(local.default_origin) > 0 ? local.default_origin : { "default-lambda-shim" : aws_lambda_function_url.default_origin_shim[0].function_url }
+}
\ No newline at end of file