From 4423a4716a333b6c509d2db3247ec9bd869ac3c5 Mon Sep 17 00:00:00 2001
From: Jan Varho <jan@varho.org>
Date: Wed, 28 Apr 2021 20:13:09 +0300
Subject: [PATCH] Upload dynamic SSG pages with revalidate-based expiry

Again, not useful now, prerequisite for revalidation.
---
 .../lambda-at-edge/src/default-handler.ts     | 12 +++-
 .../default-handler-origin-response.test.ts   | 67 ++++++++++++++-----
 .../pages/fallback-blocking/[slug].js         |  3 +-
 3 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/packages/libs/lambda-at-edge/src/default-handler.ts b/packages/libs/lambda-at-edge/src/default-handler.ts
index 0d79be48ce..b305f82056 100644
--- a/packages/libs/lambda-at-edge/src/default-handler.ts
+++ b/packages/libs/lambda-at-edge/src/default-handler.ts
@@ -655,6 +655,12 @@ const handleOriginResponse = async ({
       "passthrough"
     );
     if (isSSG) {
+      const cacheControl = renderOpts.revalidate
+        ? undefined
+        : "public, max-age=0, s-maxage=2678400, must-revalidate";
+      const expires = renderOpts.revalidate
+        ? new Date(new Date().getTime() + 1000 * renderOpts.revalidate)
+        : undefined;
       const baseKey = uri
         .replace(/^\//, "")
         .replace(/\.(json|html)$/, "")
@@ -666,14 +672,16 @@ const handleOriginResponse = async ({
         Key: `${s3BasePath}${jsonKey}`,
         Body: JSON.stringify(renderOpts.pageData),
         ContentType: "application/json",
-        CacheControl: "public, max-age=0, s-maxage=2678400, must-revalidate"
+        CacheControl: cacheControl,
+        Expires: expires
       };
       const s3HtmlParams = {
         Bucket: bucketName,
         Key: `${s3BasePath}${htmlKey}`,
         Body: html,
         ContentType: "text/html",
-        CacheControl: "public, max-age=0, s-maxage=2678400, must-revalidate"
+        CacheControl: cacheControl,
+        Expires: expires
       };
       const { PutObjectCommand } = await import(
         "@aws-sdk/client-s3/commands/PutObjectCommand"
diff --git a/packages/libs/lambda-at-edge/tests/default-handler/default-handler-origin-response.test.ts b/packages/libs/lambda-at-edge/tests/default-handler/default-handler-origin-response.test.ts
index c0749623df..7ea4cf2c97 100644
--- a/packages/libs/lambda-at-edge/tests/default-handler/default-handler-origin-response.test.ts
+++ b/packages/libs/lambda-at-edge/tests/default-handler/default-handler-origin-response.test.ts
@@ -166,24 +166,57 @@ describe("Lambda@Edge origin response", () => {
       expect(decodedBody).toEqual("<div>Rendered Page</div>");
       expect(cfResponse.status).toEqual(200);
 
-      expect(s3Client.send).toHaveBeenNthCalledWith(1, {
-        Command: "PutObjectCommand",
-        Bucket: "my-bucket.s3.amazonaws.com",
-        Key: "_next/data/build-id/fallback-blocking/not-yet-built.json",
-        Body: JSON.stringify({
-          page: "pages/fallback-blocking/[slug].js"
-        }),
-        ContentType: "application/json",
-        CacheControl: "public, max-age=0, s-maxage=2678400, must-revalidate"
-      });
-      expect(s3Client.send).toHaveBeenNthCalledWith(2, {
-        Command: "PutObjectCommand",
-        Bucket: "my-bucket.s3.amazonaws.com",
-        Key: "static-pages/build-id/fallback-blocking/not-yet-built.html",
-        Body: "<div>Rendered Page</div>",
-        ContentType: "text/html",
-        CacheControl: "public, max-age=0, s-maxage=2678400, must-revalidate"
+      expect(s3Client.send).toHaveBeenNthCalledWith(
+        1,
+        expect.objectContaining({
+          Command: "PutObjectCommand",
+          Bucket: "my-bucket.s3.amazonaws.com",
+          Key: "_next/data/build-id/fallback-blocking/not-yet-built.json",
+          Body: JSON.stringify({
+            page: "pages/fallback-blocking/[slug].js"
+          }),
+          ContentType: "application/json"
+        })
+      );
+      expect(s3Client.send).toHaveBeenNthCalledWith(
+        2,
+        expect.objectContaining({
+          Command: "PutObjectCommand",
+          Bucket: "my-bucket.s3.amazonaws.com",
+          Key: "static-pages/build-id/fallback-blocking/not-yet-built.html",
+          Body: "<div>Rendered Page</div>",
+          ContentType: "text/html"
+        })
+      );
+    });
+
+    it("uploads with revalidate-based expires", async () => {
+      const event = createCloudFrontEvent({
+        uri: "/fallback-blocking/not-yet-built.html",
+        host: "mydistribution.cloudfront.net",
+        config: { eventType: "origin-response" } as any,
+        response: {
+          headers: {},
+          status: "403"
+        } as any
       });
+
+      mockPageRequire("pages/fallback-blocking/[slug].js");
+
+      await handler(event);
+
+      expect(
+        (s3Client.send as jest.Mock).mock.calls[0][0].Expires.getTime()
+      ).toBeGreaterThan(new Date().getTime());
+      expect(
+        (s3Client.send as jest.Mock).mock.calls[0][0].Expires.getTime()
+      ).toBeLessThan(new Date().getTime() + 300000);
+      expect(
+        (s3Client.send as jest.Mock).mock.calls[1][0].Expires.getTime()
+      ).toBeGreaterThan(new Date().getTime());
+      expect(
+        (s3Client.send as jest.Mock).mock.calls[1][0].Expires.getTime()
+      ).toBeLessThan(new Date().getTime() + 300000);
     });
 
     it("renders and uploads HTML and JSON for fallback SSG data requests", async () => {
diff --git a/packages/libs/lambda-at-edge/tests/shared-fixtures/built-artifact/pages/fallback-blocking/[slug].js b/packages/libs/lambda-at-edge/tests/shared-fixtures/built-artifact/pages/fallback-blocking/[slug].js
index 5768347d05..25289044d3 100644
--- a/packages/libs/lambda-at-edge/tests/shared-fixtures/built-artifact/pages/fallback-blocking/[slug].js
+++ b/packages/libs/lambda-at-edge/tests/shared-fixtures/built-artifact/pages/fallback-blocking/[slug].js
@@ -9,7 +9,8 @@ module.exports = {
       renderOpts: {
         pageData: {
           page: "pages/fallback-blocking/[slug].js"
-        }
+        },
+        revalidate: 300
       }
     });
   }