feat(lambda-at-edge): support custom headers (with caveats)

packages/e2e-tests/next-app/cypress/integration/headers.test.ts

@@ -0,0 +1,45 @@
+describe("Headers Tests", () => {
+  describe("Custom headers defined in next.config.js", () => {
+    [
+      {
+        path: "/ssr-page",
+        expectedHeaders: { "x-custom-header-ssr-page": "custom" }
+      },
+      {
+        path: "/ssg-page",
+        expectedHeaders: { "x-custom-header-ssg-page": "custom" }
+      },
+      {
+        path: "/",
+        expectedHeaders: { "x-custom-header-all": "custom" }
+      },
+      {
+        path: "/not-found",
+        expectedHeaders: { "x-custom-header-all": "custom" }
+      },
+      {
+        path: "/api/basic-api",
+        expectedHeaders: { "x-custom-header-api": "custom" }
+      },
+      {
+        path: "/app-store-badge.png",
+        expectedHeaders: { "x-custom-header-public-file": "custom" }
+      }
+    ].forEach(({ path, expectedHeaders }) => {
+      it(`add headers ${JSON.stringify(
+        expectedHeaders
+      )} for path ${path}`, () => {
+        cy.request({
+          url: path,
+          failOnStatusCode: false
+        }).then((response) => {
+          for (const expectedHeader in expectedHeaders) {
+            expect(response.headers[expectedHeader]).to.equal(
+              expectedHeaders[expectedHeader]
+            );
+          }
+        });
+      });
+    });
+  });
+});

packages/e2e-tests/next-app/next.config.js

@@ -90,5 +90,60 @@ module.exports = {
         destination: "/api/basic-api"
       }
     ];
+  },
+  async headers() {
+    return [
+      {
+        source: "/:path*",
+        headers: [
+          {
+            key: "x-custom-header-all",
+            value: "custom"
+          }
+        ]
+      },
+      {
+        source: "/ssr-page",
+        headers: [
+          {
+            key: "x-custom-header-ssr-page",
+            value: "custom"
+          }
+        ]
+      },
+      {
+        /**
+         * TODO: we need to specify S3 key here for SSG page (ssg-page.html) because of how things currently work.
+         * Request URI is rewritten to the S3 key, so in origin response handler we have no easy way to determine the original page path.
+         * In the future, we may bypass S3 origin + remove origin response handler so origin request handler directly calls S3, making this easier.
+         */
+        source: "/ssg-page.html",
+        headers: [
+          {
+            key: "x-custom-header-ssg-page",
+            value: "custom"
+          }
+        ]
+      },
+      {
+        // For public files, the original path matches the S3 key
+        source: "/app-store-badge.png",
+        headers: [
+          {
+            key: "x-custom-header-public-file",
+            value: "custom"
+          }
+        ]
+      },
+      {
+        source: "/api/basic-api",
+        headers: [
+          {
+            key: "x-custom-header-api",
+            value: "custom"
+          }
+        ]
+      }
+    ];
   }
 };

packages/libs/lambda-at-edge/src/api-handler.ts

@@ -15,6 +15,7 @@ import {
   getRedirectPath
 } from "./routing/redirector";
 import { getRewritePath } from "./routing/rewriter";
+import { addHeadersToResponse } from "./headers/addHeaders";
 
 const basePath = RoutesManifestJson.basePath;
 
@@ -52,7 +53,7 @@ const router = (
 
 export const handler = async (
   event: OriginRequestEvent
-): Promise<CloudFrontResultResponse | CloudFrontRequest> => {
+): Promise<CloudFrontResultResponse> => {
   const request = event.Records[0].cf.request;
   const routesManifest: RoutesManifest = RoutesManifestJson;
   const buildManifest: OriginRequestApiHandlerManifest = manifest;
@@ -95,5 +96,10 @@ export const handler = async (
 
   page.default(req, res);
 
-  return responsePromise;
+  const response = await responsePromise;
+
+  // Add custom headers before returning response
+  addHeadersToResponse(request.uri, response, routesManifest);
+
+  return response;
 };

packages/libs/lambda-at-edge/src/default-handler.ts

@@ -30,6 +30,7 @@ import {
   getRedirectPath
 } from "./routing/redirector";
 import { getRewritePath } from "./routing/rewriter";
+import { addHeadersToResponse } from "./headers/addHeaders";
 
 const basePath = RoutesManifestJson.basePath;
 const NEXT_PREVIEW_DATA_COOKIE = "__next_preview_data";
@@ -195,6 +196,18 @@ export const handler = async (
     });
   }
 
+  // Add custom headers to responses only.
+  // TODO: for paths that hit S3 origin, it will match on the rewritten URI, i.e it may be rewritten to S3 key.
+  if (response.hasOwnProperty("status")) {
+    const request = event.Records[0].cf.request;
+
+    addHeadersToResponse(
+      request.uri,
+      response as CloudFrontResultResponse,
+      routesManifest
+    );
+  }
+
   const tHandlerEnd = now();
 
   log("handler execution time", tHandlerBegin, tHandlerEnd);

packages/libs/lambda-at-edge/src/headers/addHeaders.ts

@@ -0,0 +1,30 @@
+import { CloudFrontResultResponse } from "aws-lambda";
+import { RoutesManifest } from "../../types";
+import { matchPath } from "../routing/matcher";
+
+export function addHeadersToResponse(
+  path: string,
+  response: CloudFrontResultResponse,
+  routesManifest: RoutesManifest
+): void {
+  // Add custom headers to response
+  if (response.headers) {
+    for (const headerData of routesManifest.headers) {
+      const match = matchPath(path, headerData.source);
+
+      if (match) {
+        for (const header of headerData.headers) {
+          if (header.key && header.value) {
+            const headerLowerCase = header.key.toLowerCase();
+            response.headers[headerLowerCase] = [
+              {
+                key: headerLowerCase,
+                value: header.value
+              }
+            ];
+          }
+        }
+      }
+    }
+  }
+}

packages/libs/lambda-at-edge/src/routing/redirector.ts

@@ -7,6 +7,7 @@ import {
 } from "../../types";
 import * as http from "http";
 import { CloudFrontRequest } from "aws-lambda";
+import { CloudFrontResultResponse } from "aws-lambda";
 
 /**
  * Whether this is the default trailing slash redirect.
@@ -90,7 +91,7 @@ export function createRedirectResponse(
   uri: string,
   querystring: string,
   statusCode: number
-) {
+): CloudFrontResultResponse {
   const location = querystring ? `${uri}?${querystring}` : uri;
 
   const status = statusCode.toString();

packages/libs/lambda-at-edge/tests/api-handler/api-basepath-routes-manifest.json

@@ -17,6 +17,17 @@
       "regex": "^/basepath/api/rewrite-getCustomers$"
     }
   ],
-  "headers": [],
+  "headers": [
+    {
+      "source": "/basepath/api/getCustomers",
+      "headers": [
+        {
+          "key": "x-custom-header",
+          "value": "custom"
+        }
+      ],
+      "regex": "^/basepath/api/basic-api$"
+    }
+  ],
   "dynamicRoutes": []
 }

packages/libs/lambda-at-edge/tests/api-handler/api-handler-with-basepath.test.ts

@@ -99,4 +99,33 @@ describe("API lambda handler with basePath configured", () => {
       }
     );
   });
+
+  describe("Custom Headers", () => {
+    it.each`
+      path                            | expectedHeaders                    | expectedJs
+      ${"/basepath/api/getCustomers"} | ${{ "x-custom-header": "custom" }} | ${"pages/customers/[customer].js"}
+    `(
+      "has custom headers $expectedHeaders and expectedPage $expectedPage for path $path",
+      async ({ path, expectedHeaders, expectedJs }) => {
+        const event = createCloudFrontEvent({
+          uri: path,
+          host: "mydistribution.cloudfront.net"
+        });
+
+        mockPageRequire(expectedJs);
+
+        const response = await handler(event);
+
+        expect(response.headers).not.toBeUndefined();
+
+        for (const header in expectedHeaders) {
+          const headerEntry = response.headers![header][0];
+          expect(headerEntry).toEqual({
+            key: header,
+            value: expectedHeaders[header]
+          });
+        }
+      }
+    );
+  });
 });

packages/libs/lambda-at-edge/tests/api-handler/api-handler.test.ts

@@ -159,4 +159,33 @@ describe("API lambda handler", () => {
       }
     );
   });
+
+  describe("Custom Headers", () => {
+    it.each`
+      path                   | expectedHeaders                    | expectedJs
+      ${"/api/getCustomers"} | ${{ "x-custom-header": "custom" }} | ${"pages/customers/[customer].js"}
+    `(
+      "has custom headers $expectedHeaders and expectedPage $expectedPage for path $path",
+      async ({ path, expectedHeaders, expectedJs }) => {
+        const event = createCloudFrontEvent({
+          uri: path,
+          host: "mydistribution.cloudfront.net"
+        });
+
+        mockPageRequire(expectedJs);
+
+        const response = await handler(event);
+
+        expect(response.headers).not.toBeUndefined();
+
+        for (const header in expectedHeaders) {
+          const headerEntry = response.headers![header][0];
+          expect(headerEntry).toEqual({
+            key: header,
+            value: expectedHeaders[header]
+          });
+        }
+      }
+    );
+  });
 });

packages/libs/lambda-at-edge/tests/api-handler/api-routes-manifest.json

@@ -17,6 +17,17 @@
       "regex": "^/api/rewrite-getCustomers$"
     }
   ],
-  "headers": [],
+  "headers": [
+    {
+      "source": "/api/getCustomers",
+      "headers": [
+        {
+          "key": "x-custom-header",
+          "value": "custom"
+        }
+      ],
+      "regex": "^/api/basic-api$"
+    }
+  ],
   "dynamicRoutes": []
 }

packages/libs/lambda-at-edge/tests/default-handler/default-basepath-routes-manifest-with-trailing-slash.json

@@ -45,6 +45,17 @@
       "regex": "^/basepath/path-rewrite(?:/([^/]+?))/$"
     }
   ],
-  "headers": [],
+  "headers": [
+    {
+      "source": "/basepath/customers/another/",
+      "headers": [
+        {
+          "key": "x-custom-header",
+          "value": "custom"
+        }
+      ],
+      "regex": "^/basepath/customers/another/$"
+    }
+  ],
   "dynamicRoutes": []
 }

packages/libs/lambda-at-edge/tests/default-handler/default-basepath-routes-manifest.json

@@ -45,6 +45,17 @@
       "regex": "^/basepath/path-rewrite(?:/([^/]+?))$"
     }
   ],
-  "headers": [],
+  "headers": [
+    {
+      "source": "/basepath/customers/another",
+      "headers": [
+        {
+          "key": "x-custom-header",
+          "value": "custom"
+        }
+      ],
+      "regex": "^/basepath/customers/another$"
+    }
+  ],
   "dynamicRoutes": []
 }

packages/libs/lambda-at-edge/tests/default-handler/default-handler-with-basepath.test.ts

@@ -723,5 +723,37 @@ describe("Lambda@Edge", () => {
         }
       );
     });
+
+    describe("Custom Headers", () => {
+      it.each`
+        path                             | expectedHeaders                    | expectedPage
+        ${"/basepath/customers/another"} | ${{ "x-custom-header": "custom" }} | ${"pages/customers/[customer].js"}
+      `(
+        "has custom headers $expectedHeaders and expectedPage $expectedPage for path $path",
+        async ({ path, expectedHeaders, expectedPage }) => {
+          // If trailingSlash = true, append "/" to get the non-redirected path
+          if (trailingSlash && !path.endsWith("/")) {
+            path += "/";
+          }
+
+          const event = createCloudFrontEvent({
+            uri: path,
+            host: "mydistribution.cloudfront.net"
+          });
+
+          mockPageRequire(expectedPage);
+
+          const response = await handler(event);
+
+          for (const header in expectedHeaders) {
+            const headerEntry = response.headers[header][0];
+            expect(headerEntry).toEqual({
+              key: header,
+              value: expectedHeaders[header]
+            });
+          }
+        }
+      );
+    });
   });
 });