From a1ae224dcf19c1f9bc020bcf422a59600c408240 Mon Sep 17 00:00:00 2001
From: Kent Kwee <79371980+kentkwee@users.noreply.github.com>
Date: Fri, 31 Jan 2025 07:23:28 +0100
Subject: [PATCH] fix(lambda-edge-adapter): Support Alternate Domain Names and
 Multiple Cookies in Lambda@Edge (#3858)

* fix(lambda-edge): improve domain name detection

Previously, only the CloudFront distribution name was used. Now, the Host header value is utilized for more accurate detection.

* fix(lambda-edge): handle multiple Set-Cookie headers correctly

Previously, only the last Set-Cookie header was returned. This update ensures all Set-Cookie headers are processed and returned as expected.

* fix(lambda-edge): add host fallback detection for unusual request headers
---
 src/adapter/lambda-edge/handler.test.ts | 86 ++++++++++++++++++++++++-
 src/adapter/lambda-edge/handler.ts      | 10 ++-
 2 files changed, 93 insertions(+), 3 deletions(-)

diff --git a/src/adapter/lambda-edge/handler.test.ts b/src/adapter/lambda-edge/handler.test.ts
index d9bbba863..232a5eea5 100644
--- a/src/adapter/lambda-edge/handler.test.ts
+++ b/src/adapter/lambda-edge/handler.test.ts
@@ -1,5 +1,9 @@
+import { describe } from 'vitest'
+import { setCookie } from '../../helper/cookie'
+import { Hono } from '../../hono'
 import { encodeBase64 } from '../../utils/encode'
-import { createBody, isContentTypeBinary } from './handler'
+import type { CloudFrontEdgeEvent } from './handler'
+import { createBody, handle, isContentTypeBinary } from './handler'
 
 describe('isContentTypeBinary', () => {
   it('Should determine whether it is binary', () => {
@@ -36,3 +40,83 @@ describe('createBody', () => {
     expect(createBody('POST', body)).not.toEqual(undefined)
   })
 })
+
+describe('handle', () => {
+  const cloudFrontEdgeEvent: CloudFrontEdgeEvent = {
+    Records: [
+      {
+        cf: {
+          config: {
+            distributionDomainName: 'd111111abcdef8.cloudfront.net',
+            distributionId: 'EDFDVBD6EXAMPLE',
+            eventType: 'viewer-request',
+            requestId: '4TyzHTaYWb1GX1qTfsHhEqV6HUDd_BzoBZnwfnvQc_1oF26ClkoUSEQ==',
+          },
+          request: {
+            clientIp: '1.2.3.4',
+            headers: {
+              host: [
+                {
+                  key: 'Host',
+                  value: 'hono.dev',
+                },
+              ],
+              accept: [
+                {
+                  key: 'accept',
+                  value: '*/*',
+                },
+              ],
+            },
+            method: 'GET',
+            querystring: '',
+            uri: '/test-path',
+          },
+        },
+      },
+    ],
+  }
+
+  it('Should support alternate domain names', async () => {
+    const app = new Hono()
+    app.get('/test-path', (c) => {
+      return c.text(c.req.url)
+    })
+    const handler = handle(app)
+
+    const res = await handler(cloudFrontEdgeEvent)
+
+    expect(res.body).toBe('https://hono.dev/test-path')
+  })
+
+  it('Should support multiple cookies', async () => {
+    const app = new Hono()
+    app.get('/test-path', (c) => {
+      setCookie(c, 'cookie1', 'value1')
+      setCookie(c, 'cookie2', 'value2')
+      return c.text('')
+    })
+    const handler = handle(app)
+
+    const res = await handler(cloudFrontEdgeEvent)
+
+    expect(res.headers).toEqual({
+      'content-type': [
+        {
+          key: 'content-type',
+          value: 'text/plain; charset=UTF-8',
+        },
+      ],
+      'set-cookie': [
+        {
+          key: 'set-cookie',
+          value: 'cookie1=value1; Path=/',
+        },
+        {
+          key: 'set-cookie',
+          value: 'cookie2=value2; Path=/',
+        },
+      ],
+    })
+  })
+})
diff --git a/src/adapter/lambda-edge/handler.ts b/src/adapter/lambda-edge/handler.ts
index 3f512d7b8..3aa7ea014 100644
--- a/src/adapter/lambda-edge/handler.ts
+++ b/src/adapter/lambda-edge/handler.ts
@@ -105,7 +105,10 @@ interface CloudFrontResult {
 const convertHeaders = (headers: Headers): CloudFrontHeaders => {
   const cfHeaders: CloudFrontHeaders = {}
   headers.forEach((value, key) => {
-    cfHeaders[key.toLowerCase()] = [{ key: key.toLowerCase(), value }]
+    cfHeaders[key.toLowerCase()] = [
+      ...(cfHeaders[key.toLowerCase()] || []),
+      { key: key.toLowerCase(), value },
+    ]
   })
   return cfHeaders
 }
@@ -147,7 +150,10 @@ const createResult = async (res: Response): Promise<CloudFrontResult> => {
 
 const createRequest = (event: CloudFrontEdgeEvent): Request => {
   const queryString = event.Records[0].cf.request.querystring
-  const urlPath = `https://${event.Records[0].cf.config.distributionDomainName}${event.Records[0].cf.request.uri}`
+  const host =
+    event.Records[0].cf.request.headers?.host?.[0]?.value ||
+    event.Records[0].cf.config.distributionDomainName
+  const urlPath = `https://${host}${event.Records[0].cf.request.uri}`
   const url = queryString ? `${urlPath}?${queryString}` : urlPath
 
   const headers = new Headers()