From 26f388fa5384b5fbcbeb536c35e004948353d25d Mon Sep 17 00:00:00 2001
From: Daniel Van Der Ploeg <daniel.vanderploeg@aligent.com.au>
Date: Thu, 11 Jul 2024 15:07:36 +0930
Subject: [PATCH] feat DO-1683: cache 404s

---
 .../lib/prerender-fargate-options.ts                  |  5 +++++
 packages/prerender-fargate/lib/prerender-fargate.ts   |  2 ++
 packages/prerender-fargate/lib/prerender/server.js    | 11 ++++++++++-
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/packages/prerender-fargate/lib/prerender-fargate-options.ts b/packages/prerender-fargate/lib/prerender-fargate-options.ts
index 5572ed0d..9e9f688a 100644
--- a/packages/prerender-fargate/lib/prerender-fargate-options.ts
+++ b/packages/prerender-fargate/lib/prerender-fargate-options.ts
@@ -72,6 +72,11 @@ export interface PrerenderFargateOptions {
    * @default - false
    */
   enableRedirectCache?: boolean;
+  /**
+   * Whether to enable caching of HTTP 404s.
+   * @default - false
+   */
+  enableNotFoundCache?: boolean;
   /**
    * Whether to enable the S3 endpoint for the VPC.
    * @default - false
diff --git a/packages/prerender-fargate/lib/prerender-fargate.ts b/packages/prerender-fargate/lib/prerender-fargate.ts
index 2cd41e28..9ec537a5 100644
--- a/packages/prerender-fargate/lib/prerender-fargate.ts
+++ b/packages/prerender-fargate/lib/prerender-fargate.ts
@@ -83,6 +83,7 @@ export class PrerenderFargate extends Construct {
       expirationDays,
       enableS3Endpoint,
       enableRedirectCache,
+      enableNotFoundCache,
       desiredInstanceCount,
       bucketName,
       domainName,
@@ -142,6 +143,7 @@ export class PrerenderFargate extends Construct {
       S3_BUCKET_NAME: this.bucket.bucketName,
       AWS_REGION: Stack.of(this).region,
       ENABLE_REDIRECT_CACHE: enableRedirectCache?.toString() || "false",
+      ENABLE_NOTFOUND_CACHE: enableNotFoundCache?.toString() || "false",
       ENABLE_PRERENDER_HEADER: enablePrerenderHeader?.toString() || "true",
     };
 
diff --git a/packages/prerender-fargate/lib/prerender/server.js b/packages/prerender-fargate/lib/prerender/server.js
index 63e09a01..81b1b4ae 100644
--- a/packages/prerender-fargate/lib/prerender/server.js
+++ b/packages/prerender-fargate/lib/prerender/server.js
@@ -160,7 +160,16 @@ server.use(prerender.removeScriptTags());
 
 server.use({
     pageLoaded: function(req, res, next) {
-        const statusCodesToCache = process.env.ENABLE_REDIRECT_CACHE.toLowerCase() === 'true' ? ['200', '301', '302', '308'] : ['200'];
+        const statusCodesToCache = ['200'];
+
+        if (process.env.ENABLE_REDIRECT_CACHE.toLowerCase() === 'true') {
+            statusCodesToCache.push('301', '302', '308');
+        }
+
+        if (process.env.ENABLE_NOTFOUND_CACHE.toLowerCase() === 'true') {
+            statusCodesToCache.push('404');
+        }
+
         var s3Metadata = {}
         const cacheObject = function (err, result) {
             if (!err && result) {