feat: Add ability to add IAM Auth to Lambda Function URLs (#109)

* feat: add 'rsc', 'next-router-prefetch', 'next-router-state-tree' to lambda cache policy

* feat: add ability to use IAM auth for lambda function url

* fix: clearer comments; use node18 for lambda@edge; use jest for tests

* fix: _next/image routes sigv4 mismatch

* fix: try to add updated jest config via projen

.projenrc.js

@@ -36,8 +36,15 @@ const project = new awscdk.AwsCdkConstructLibrary({
     'jszip',
     'glob',
     'node-fetch',
+    '@aws-sdk/signature-v4',
+    '@aws-crypto/sha256-js',
   ] /* Runtime dependencies of this module. */,
   devDeps: ['open-next', 'aws-sdk', 'constructs'] /* Build dependencies for this module. */,
+  jestOptions: {
+    jestConfig: {
+      testMatch: ['<rootDir>/src/**/__tests__/**/*.ts?(x)', '<rootDir>/(test|src)/**/*(*.)@(spec|test|assets).ts?(x)'],
+    },
+  },
 
   // do not generate sample test files
   sampleCode: false,

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+  "editor.codeActionsOnSave": {
+    "source.fixAll.eslint": true
+  },
+  "eslint.packageManager": "yarn",
+}

assets/lambda@edge/LambdaOriginRequest.ts

@@ -1,53 +1,16 @@
-import url from 'url';
-import type { CloudFrontRequest, CloudFrontRequestHandler } from 'aws-lambda';
+import type { CloudFrontRequestHandler } from 'aws-lambda';
+import { fixHostHeader, handleS3Request } from './common';
 
 /**
  * This fixes the "host" header to be the host of the origin.
  * The origin is the lambda server function URL.
  * If we don't provide its expected "host", it will not know how to route the request.
  */
-export const handler: CloudFrontRequestHandler = (event, _context, callback) => {
+export const handler: CloudFrontRequestHandler = async (event) => {
   const request = event.Records[0].cf.request;
   // console.log('request', JSON.stringify(request, null, 2));
 
-  // remove cookies from requests to S3
-  if (request.origin?.s3) {
-    request.headers.cookie = [];
-    request.headers.host = [
-      {
-        key: 'host',
-        value: request.origin.s3.domainName,
-      },
-    ]; // sending the wrong host header to S3 will cause a 403
-  }
-
-  // get config (only for custom lambda HTTP origin)
-  const originUrlHeader = getCustomHeaderValue(request, 'x-origin-url');
-  if (!originUrlHeader) return callback(null, request);
-  const originUrl = url.parse(originUrlHeader, true);
-  if (!originUrl.host) throw new Error('Origin url host is missing');
-
-  // fix host header and pass along the original host header
-  const originalHost = request.headers.host[0].value;
-  request.headers['x-forwarded-host'] = [{ key: 'x-forwarded-host', value: originalHost }];
-  request.headers.host = [{ key: 'host', value: originUrl.host }];
-  callback(null, request);
+  handleS3Request(request);
+  fixHostHeader(request);
+  return request;
 };
-
-/**
- * We can't use environment variables in the lambda@edge function.
- * We have to use custom headers for passing configuration to the lambda@edge function.
- */
-function getCustomHeaderValue(request: CloudFrontRequest, headerName: string): string | undefined {
-  const originUrlHeader = request.origin?.custom?.customHeaders?.[headerName];
-
-  if (!originUrlHeader || !originUrlHeader[0]) {
-    if (request.origin?.custom) {
-      // we should have an origin url header for custom lambda HTTP origin
-      console.error('Origin header wasn"t set correctly, cannot get origin url');
-    }
-    return undefined;
-  }
-
-  return originUrlHeader[0].value;
-}

assets/lambda@edge/LambdaOriginRequestIamAuth.test.ts

@@ -0,0 +1,117 @@
+import type { CloudFrontRequestEvent } from 'aws-lambda';
+import { signRequest } from './LambdaOriginRequestIamAuth';
+
+describe('LambdaOriginRequestIamAuth', () => {
+  test('signRequest should add x-amz headers', async () => {
+    // dummy AWS credentials
+    process.env = { ...process.env, ...getFakeAwsCreds() };
+    const event = getFakeCloudFrontRequest();
+    const request = event.Records[0].cf.request;
+    await signRequest(request);
+    const securityHeaders = ['x-amz-date', 'x-amz-security-token', 'x-amz-content-sha256', 'authorization'];
+    const hasSignedHeaders = securityHeaders.every((h) => h in request.headers);
+    expect(hasSignedHeaders).toBe(true);
+  });
+});
+
+function getFakeCloudFrontRequest(): CloudFrontRequestEvent {
+  return {
+    Records: [
+      {
+        cf: {
+          config: {
+            distributionDomainName: 'd6b8brjqfujeb.cloudfront.net',
+            distributionId: 'EHX2SDUU61T7U',
+            eventType: 'origin-request',
+            requestId: '',
+          },
+          request: {
+            clientIp: '1.1.1.1',
+            headers: {
+              host: [
+                {
+                  key: 'Host',
+                  value: 'd6b8brjqfujeb.cloudfront.net',
+                },
+              ],
+              'accept-language': [
+                {
+                  key: 'Accept-Language',
+                  value: 'en-US,en;q=0.9',
+                },
+              ],
+              referer: [
+                {
+                  key: 'Referer',
+                  value: 'https://d6b8brjqfujeb.cloudfront.net/batches/2/overview',
+                },
+              ],
+              'x-forwarded-for': [
+                {
+                  key: 'X-Forwarded-For',
+                  value: '1.1.1.1',
+                },
+              ],
+              'user-agent': [
+                {
+                  key: 'User-Agent',
+                  value:
+                    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/113.0.0.0 Safari/537.36',
+                },
+              ],
+              via: [
+                {
+                  key: 'Via',
+                  value: '2.0 8bf94e29f889f8d0076c4502ae008b58.cloudfront.net (CloudFront)',
+                },
+              ],
+              'accept-encoding': [
+                {
+                  key: 'Accept-Encoding',
+                  value: 'br,gzip',
+                },
+              ],
+              'sec-ch-ua': [
+                {
+                  key: 'sec-ch-ua',
+                  value: '"Google Chrome";v="113", "Chromium";v="113", "Not-A.Brand";v="24"',
+                },
+              ],
+            },
+            method: 'GET',
+            querystring: '',
+            uri: '/some/path',
+            origin: {
+              custom: {
+                customHeaders: {},
+                domainName: 'kjtbbx7u533q7p7n5font6gpci0phrng.lambda-url.us-east-1.on.aws',
+                keepaliveTimeout: 5,
+                path: '',
+                port: 443,
+                protocol: 'https',
+                readTimeout: 30,
+                sslProtocols: ['TLSv1.2'],
+              },
+            },
+            body: {
+              action: 'read-only',
+              data: '',
+              encoding: 'base64',
+              inputTruncated: false,
+            },
+          },
+        },
+      },
+    ],
+  };
+}
+
+function getFakeAwsCreds() {
+  return {
+    AWS_REGION: 'us-east-1',
+    AWS_ACCESS_KEY_ID: 'ZSBAT5GENDHC3XYRH36I',
+    AWS_SECRET_ACCESS_KEY: 'jpWfApw1AO0xzGZeeT1byQq1zqfQITVqVhTkkql4',
+    AWS_SESSION_TOKEN:
+      'ZQoJb3JpZ2luX2VjEFgaCXVzLWVhc3QtMSJGMEQCIHijzdTXh59aSe2hRfCWpFd2/jacPUC+8rCq3qBIiuG2AiAGX8jqld+p04nPYfuShi1lLN/Z1hEXG9QSNEmEFLTxGSqmAgiR//////////8BEAIaDDI2ODkxNDQ2NTIzMSIMrAMO5/GTvMgoG+chKvoB4f4V1TfkZiHOlmeMK6Ep58mav65A0WU3K9WPzdrJojnGqqTuS85zTlKhm3lfmMxCOtwS/OlOuiBQ1MZNlksK2je1FazgbXN46fNSi+iHiY9VfyRAd0wSLmXB8FFrCGsU92QOy/+deji0qIVadsjEyvBRxzQj5oIUI5sb74Yt7uNvka9fVZcT4s4IndYda0N7oZwIrApCuzzBMuoMAhabmgVrZTbiLmvOiFHS2XZWBySABdygqaIzfV7G4hjckvcXhtxpkw+HJUZTNzVUlspghzte1UG6VvIRV8ax3kWA3zqm8nA/1gHkl40DubJIXz1AJbg5Cps5moE1pjD7vNijBjqeAZh0Q/e0awIHnV4dXMfXUu5mWJ7Db9K1eUlSSL9FyiKeKd94HEdrbIrnPuIWVT/I/5RjNm7NgPYiqmpyx3fSpVcq9CKws0oEfBw6J9Hxk0IhV8yWFZYNMWIarUUZdmL9vVeJmFZmwyL4JjY1s/SZIU/oa8DtvkmP4RG4tTJfpyyhoKL0wJOevkYyoigNllBlLN59SZAT8CCADpN/B+sK',
+  };
+}

assets/lambda@edge/LambdaOriginRequestIamAuth.ts

@@ -0,0 +1,133 @@
+import qs from 'node:querystring';
+import { Sha256 } from '@aws-crypto/sha256-js';
+import { SignatureV4 } from '@aws-sdk/signature-v4';
+import type { CloudFrontHeaders, CloudFrontRequest, CloudFrontRequestHandler } from 'aws-lambda';
+import { fixHostHeader, handleS3Request } from './common';
+
+/**
+ * This Lambda@Edge handler fixes s3 requests, fixes the host header, and
+ * signs requests as they're destined for Lambda Function URL that requires
+ * IAM Auth.
+ */
+export const handler: CloudFrontRequestHandler = async (event) => {
+  const request = event.Records[0].cf.request;
+  // console.log('request', JSON.stringify(request, null, 2));
+
+  handleS3Request(request);
+  fixHostHeader(request);
+  await signRequest(request);
+  // console.log(JSON.stringify(request), null, 2);
+  return request;
+};
+
+let sigv4: SignatureV4;
+
+/**
+ * When `NextjsDistributionProps.functionUrlAuthType` is set to
+ * `lambda.FunctionUrlAuthType.AWS_IAM` we need to sign the `CloudFrontRequest`s
+ * with AWS IAM SigV4 so that CloudFront can invoke the Nextjs server and image
+ * optimization functions via function URLs. When configured, this lambda@edge
+ * function has the permission, lambda:InvokeFunctionUrl, to invoke both
+ * functions.
+ * @link https://medium.com/@dario_26152/restrict-access-to-lambda-functionurl-to-cloudfront-using-aws-iam-988583834705
+ */
+export async function signRequest(request: CloudFrontRequest) {
+  if (!sigv4) {
+    sigv4 = getSigV4();
+  }
+  const headerBag = cfHeadersToHeaderBag(request);
+  let body: string | undefined;
+  if (request.body?.data) {
+    body = Buffer.from(request.body.data, 'base64').toString();
+  }
+  const params = queryStringToParams(request);
+  const signed = await sigv4.sign({
+    method: request.method,
+    headers: headerBag,
+    hostname: headerBag.host,
+    path: request.uri,
+    body,
+    query: params,
+    protocol: 'https',
+  });
+  request.headers = headerBagToCfHeaders(signed.headers);
+}
+
+function getSigV4(): SignatureV4 {
+  const region = process.env.AWS_REGION;
+  const accessKeyId = process.env.AWS_ACCESS_KEY_ID;
+  const secretAccessKey = process.env.AWS_SECRET_ACCESS_KEY;
+  const sessionToken = process.env.AWS_SESSION_TOKEN;
+  if (!region) throw new Error('AWS_REGION missing');
+  if (!accessKeyId) throw new Error('AWS_ACCESS_KEY_ID missing');
+  if (!secretAccessKey) throw new Error('AWS_SECRET_ACCESS_KEY missing');
+  if (!sessionToken) throw new Error('AWS_SESSION_TOKEN missing');
+  return new SignatureV4({
+    service: 'lambda',
+    region,
+    credentials: {
+      accessKeyId,
+      secretAccessKey,
+      sessionToken,
+    },
+    sha256: Sha256,
+  });
+}
+
+type HeaderBag = Record<string, string>;
+/**
+ * Converts CloudFront headers (can have array of header values) to simple
+ * header bag (object) required by `sigv4.sign`
+ *
+ * NOTE: only includes headers allowed by origin policy to prevent signature
+ * mismatch
+ */
+export function cfHeadersToHeaderBag(request: CloudFrontRequest): HeaderBag {
+  let headerBag: HeaderBag = {};
+  for (const [header, values] of Object.entries(request.headers)) {
+    // don't sign 'x-forwarded-for' b/c it changes from hop to hop
+    if (header === 'x-forwarded-for') continue;
+    if (request.uri === '_next/image') {
+      // _next/image origin policy only allows accept
+      if (header === 'accept') {
+        headerBag[header] = values[0].value;
+      }
+    } else {
+      headerBag[header] = values[0].value;
+    }
+  }
+  return headerBag;
+}
+
+/**
+ * Converts simple header bag (object) to CloudFront headers
+ */
+export function headerBagToCfHeaders(headerBag: HeaderBag): CloudFrontHeaders {
+  const cfHeaders: CloudFrontHeaders = {};
+  for (const [header, value] of Object.entries(headerBag)) {
+    cfHeaders[header] = [{ key: header, value }];
+  }
+  return cfHeaders;
+}
+
+/**
+ * Converts CloudFront querystring to `HttpRequest.query` for IAM Sig V4
+ *
+ * NOTE: only includes query parameters allowed at origin to prevent signature
+ * mismatch errors
+ */
+export function queryStringToParams(request: CloudFrontRequest) {
+  const params: Record<string, string> = {};
+  const _params = new URLSearchParams(request.querystring);
+  for (const [k, v] of _params) {
+    if (request.uri === '_next/image') {
+      // _next/image origin policy only allows these querystrings
+      if (['url', 'q', 'w'].includes(k)) {
+        params[k] = v;
+      }
+    } else {
+      params[k] = v;
+    }
+  }
+  return params;
+}