Skip to content
This repository was archived by the owner on Jan 28, 2025. It is now read-only.

Commit 8b9e822

Browse files
dphangdphang
authored
feat(lambda-at-edge): support custom headers (with caveats) (#662)
1 parent 4798ddd commit 8b9e822

17 files changed

+358
-9
lines changed

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
describe("Headers Tests", () => {
2+
describe("Custom headers defined in next.config.js", () => {
3+
[
4+
{
5+
path: "/ssr-page",
6+
expectedHeaders: { "x-custom-header-ssr-page": "custom" }
7+
},
8+
{
9+
path: "/ssg-page",
10+
expectedHeaders: { "x-custom-header-ssg-page": "custom" }
11+
},
12+
{
13+
path: "/",
14+
expectedHeaders: { "x-custom-header-all": "custom" }
15+
},
16+
{
17+
path: "/not-found",
18+
expectedHeaders: { "x-custom-header-all": "custom" }
19+
},
20+
{
21+
path: "/api/basic-api",
22+
expectedHeaders: { "x-custom-header-api": "custom" }
23+
},
24+
{
25+
path: "/app-store-badge.png",
26+
expectedHeaders: { "x-custom-header-public-file": "custom" }
27+
}
28+
].forEach(({ path, expectedHeaders }) => {
29+
it(`add headers ${JSON.stringify(
30+
expectedHeaders
31+
)} for path ${path}`, () => {
32+
cy.request({
33+
url: path,
34+
failOnStatusCode: false
35+
}).then((response) => {
36+
for (const expectedHeader in expectedHeaders) {
37+
expect(response.headers[expectedHeader]).to.equal(
38+
expectedHeaders[expectedHeader]
39+
);
40+
}
41+
});
42+
});
43+
});
44+
});
45+
});

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,5 +90,60 @@ module.exports = {
9090
destination: "/api/basic-api"
9191
}
9292
];
93+
},
94+
async headers() {
95+
return [
96+
{
97+
source: "/:path*",
98+
headers: [
99+
{
100+
key: "x-custom-header-all",
101+
value: "custom"
102+
}
103+
]
104+
},
105+
{
106+
source: "/ssr-page",
107+
headers: [
108+
{
109+
key: "x-custom-header-ssr-page",
110+
value: "custom"
111+
}
112+
]
113+
},
114+
{
115+
/**
116+
* TODO: we need to specify S3 key here for SSG page (ssg-page.html) because of how things currently work.
117+
* Request URI is rewritten to the S3 key, so in origin response handler we have no easy way to determine the original page path.
118+
* In the future, we may bypass S3 origin + remove origin response handler so origin request handler directly calls S3, making this easier.
119+
*/
120+
source: "/ssg-page.html",
121+
headers: [
122+
{
123+
key: "x-custom-header-ssg-page",
124+
value: "custom"
125+
}
126+
]
127+
},
128+
{
129+
// For public files, the original path matches the S3 key
130+
source: "/app-store-badge.png",
131+
headers: [
132+
{
133+
key: "x-custom-header-public-file",
134+
value: "custom"
135+
}
136+
]
137+
},
138+
{
139+
source: "/api/basic-api",
140+
headers: [
141+
{
142+
key: "x-custom-header-api",
143+
value: "custom"
144+
}
145+
]
146+
}
147+
];
93148
}
94149
};

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import {
1515
getRedirectPath
1616
} from "./routing/redirector";
1717
import { getRewritePath } from "./routing/rewriter";
18+
import { addHeadersToResponse } from "./headers/addHeaders";
1819

1920
const basePath = RoutesManifestJson.basePath;
2021

@@ -52,7 +53,7 @@ const router = (
5253

5354
export const handler = async (
5455
event: OriginRequestEvent
55-
): Promise<CloudFrontResultResponse | CloudFrontRequest> => {
56+
): Promise<CloudFrontResultResponse> => {
5657
const request = event.Records[0].cf.request;
5758
const routesManifest: RoutesManifest = RoutesManifestJson;
5859
const buildManifest: OriginRequestApiHandlerManifest = manifest;
@@ -95,5 +96,10 @@ export const handler = async (
9596

9697
page.default(req, res);
9798

98-
return responsePromise;
99+
const response = await responsePromise;
100+
101+
// Add custom headers before returning response
102+
addHeadersToResponse(request.uri, response, routesManifest);
103+
104+
return response;
99105
};

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import {
3030
getRedirectPath
3131
} from "./routing/redirector";
3232
import { getRewritePath } from "./routing/rewriter";
33+
import { addHeadersToResponse } from "./headers/addHeaders";
3334

3435
const basePath = RoutesManifestJson.basePath;
3536
const NEXT_PREVIEW_DATA_COOKIE = "__next_preview_data";
@@ -195,6 +196,18 @@ export const handler = async (
195196
});
196197
}
197198

199+
// Add custom headers to responses only.
200+
// TODO: for paths that hit S3 origin, it will match on the rewritten URI, i.e it may be rewritten to S3 key.
201+
if (response.hasOwnProperty("status")) {
202+
const request = event.Records[0].cf.request;
203+
204+
addHeadersToResponse(
205+
request.uri,
206+
response as CloudFrontResultResponse,
207+
routesManifest
208+
);
209+
}
210+
198211
const tHandlerEnd = now();
199212

200213
log("handler execution time", tHandlerBegin, tHandlerEnd);

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import { CloudFrontResultResponse } from "aws-lambda";
2+
import { RoutesManifest } from "../../types";
3+
import { matchPath } from "../routing/matcher";
4+
5+
export function addHeadersToResponse(
6+
path: string,
7+
response: CloudFrontResultResponse,
8+
routesManifest: RoutesManifest
9+
): void {
10+
// Add custom headers to response
11+
if (response.headers) {
12+
for (const headerData of routesManifest.headers) {
13+
const match = matchPath(path, headerData.source);
14+
15+
if (match) {
16+
for (const header of headerData.headers) {
17+
if (header.key && header.value) {
18+
const headerLowerCase = header.key.toLowerCase();
19+
response.headers[headerLowerCase] = [
20+
{
21+
key: headerLowerCase,
22+
value: header.value
23+
}
24+
];
25+
}
26+
}
27+
}
28+
}
29+
}
30+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
} from "../../types";
88
import * as http from "http";
99
import { CloudFrontRequest } from "aws-lambda";
10+
import { CloudFrontResultResponse } from "aws-lambda";
1011

1112
/**
1213
* Whether this is the default trailing slash redirect.
@@ -90,7 +91,7 @@ export function createRedirectResponse(
9091
uri: string,
9192
querystring: string,
9293
statusCode: number
93-
) {
94+
): CloudFrontResultResponse {
9495
const location = querystring ? `${uri}?${querystring}` : uri;
9596

9697
const status = statusCode.toString();

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,17 @@
1717
"regex": "^/basepath/api/rewrite-getCustomers$"
1818
}
1919
],
20-
"headers": [],
20+
"headers": [
21+
{
22+
"source": "/basepath/api/getCustomers",
23+
"headers": [
24+
{
25+
"key": "x-custom-header",
26+
"value": "custom"
27+
}
28+
],
29+
"regex": "^/basepath/api/basic-api$"
30+
}
31+
],
2132
"dynamicRoutes": []
2233
}

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,4 +99,33 @@ describe("API lambda handler with basePath configured", () => {
9999
}
100100
);
101101
});
102+
103+
describe("Custom Headers", () => {
104+
it.each`
105+
path | expectedHeaders | expectedJs
106+
${"/basepath/api/getCustomers"} | ${{ "x-custom-header": "custom" }} | ${"pages/customers/[customer].js"}
107+
`(
108+
"has custom headers $expectedHeaders and expectedPage $expectedPage for path $path",
109+
async ({ path, expectedHeaders, expectedJs }) => {
110+
const event = createCloudFrontEvent({
111+
uri: path,
112+
host: "mydistribution.cloudfront.net"
113+
});
114+
115+
mockPageRequire(expectedJs);
116+
117+
const response = await handler(event);
118+
119+
expect(response.headers).not.toBeUndefined();
120+
121+
for (const header in expectedHeaders) {
122+
const headerEntry = response.headers![header][0];
123+
expect(headerEntry).toEqual({
124+
key: header,
125+
value: expectedHeaders[header]
126+
});
127+
}
128+
}
129+
);
130+
});
102131
});

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,4 +159,33 @@ describe("API lambda handler", () => {
159159
}
160160
);
161161
});
162+
163+
describe("Custom Headers", () => {
164+
it.each`
165+
path | expectedHeaders | expectedJs
166+
${"/api/getCustomers"} | ${{ "x-custom-header": "custom" }} | ${"pages/customers/[customer].js"}
167+
`(
168+
"has custom headers $expectedHeaders and expectedPage $expectedPage for path $path",
169+
async ({ path, expectedHeaders, expectedJs }) => {
170+
const event = createCloudFrontEvent({
171+
uri: path,
172+
host: "mydistribution.cloudfront.net"
173+
});
174+
175+
mockPageRequire(expectedJs);
176+
177+
const response = await handler(event);
178+
179+
expect(response.headers).not.toBeUndefined();
180+
181+
for (const header in expectedHeaders) {
182+
const headerEntry = response.headers![header][0];
183+
expect(headerEntry).toEqual({
184+
key: header,
185+
value: expectedHeaders[header]
186+
});
187+
}
188+
}
189+
);
190+
});
162191
});

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,17 @@
1717
"regex": "^/api/rewrite-getCustomers$"
1818
}
1919
],
20-
"headers": [],
20+
"headers": [
21+
{
22+
"source": "/api/getCustomers",
23+
"headers": [
24+
{
25+
"key": "x-custom-header",
26+
"value": "custom"
27+
}
28+
],
29+
"regex": "^/api/basic-api$"
30+
}
31+
],
2132
"dynamicRoutes": []
2233
}

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,17 @@
4545
"regex": "^/basepath/path-rewrite(?:/([^/]+?))/$"
4646
}
4747
],
48-
"headers": [],
48+
"headers": [
49+
{
50+
"source": "/basepath/customers/another/",
51+
"headers": [
52+
{
53+
"key": "x-custom-header",
54+
"value": "custom"
55+
}
56+
],
57+
"regex": "^/basepath/customers/another/$"
58+
}
59+
],
4960
"dynamicRoutes": []
5061
}

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,17 @@
4545
"regex": "^/basepath/path-rewrite(?:/([^/]+?))$"
4646
}
4747
],
48-
"headers": [],
48+
"headers": [
49+
{
50+
"source": "/basepath/customers/another",
51+
"headers": [
52+
{
53+
"key": "x-custom-header",
54+
"value": "custom"
55+
}
56+
],
57+
"regex": "^/basepath/customers/another$"
58+
}
59+
],
4960
"dynamicRoutes": []
5061
}

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -723,5 +723,37 @@ describe("Lambda@Edge", () => {
723723
}
724724
);
725725
});
726+
727+
describe("Custom Headers", () => {
728+
it.each`
729+
path | expectedHeaders | expectedPage
730+
${"/basepath/customers/another"} | ${{ "x-custom-header": "custom" }} | ${"pages/customers/[customer].js"}
731+
`(
732+
"has custom headers $expectedHeaders and expectedPage $expectedPage for path $path",
733+
async ({ path, expectedHeaders, expectedPage }) => {
734+
// If trailingSlash = true, append "/" to get the non-redirected path
735+
if (trailingSlash && !path.endsWith("/")) {
736+
path += "/";
737+
}
738+
739+
const event = createCloudFrontEvent({
740+
uri: path,
741+
host: "mydistribution.cloudfront.net"
742+
});
743+
744+
mockPageRequire(expectedPage);
745+
746+
const response = await handler(event);
747+
748+
for (const header in expectedHeaders) {
749+
const headerEntry = response.headers[header][0];
750+
expect(headerEntry).toEqual({
751+
key: header,
752+
value: expectedHeaders[header]
753+
});
754+
}
755+
}
756+
);
757+
});
726758
});
727759
});

0 commit comments

Comments
 (0)