From 378ef710273ee65a015163f23f7355bbfbec4dba Mon Sep 17 00:00:00 2001
From: Corwin Marsh <corwinm@users.noreply.github.com>
Date: Wed, 22 May 2024 12:04:57 -0700
Subject: [PATCH] fix: Support ALB Event response headers (#1756)

* fix: Support ALB Event response headers

* fix: Run prettier:fix
---
 examples/events/alb/serverless.yml            |  9 +++++
 .../events/alb/src/handler-response-header.js | 19 ++++++++++
 src/events/alb/HttpServer.js                  |  2 ++
 .../alb-headers/alb-headers.test.js           | 36 +++++++++++++++++++
 tests/integration/alb-headers/serverless.yml  | 27 ++++++++++++++
 tests/integration/alb-headers/src/handler.js  | 10 ++++++
 .../integration/alb-headers/src/package.json  |  3 ++
 7 files changed, 106 insertions(+)
 create mode 100644 examples/events/alb/src/handler-response-header.js
 create mode 100644 tests/integration/alb-headers/alb-headers.test.js
 create mode 100644 tests/integration/alb-headers/serverless.yml
 create mode 100644 tests/integration/alb-headers/src/handler.js
 create mode 100644 tests/integration/alb-headers/src/package.json

diff --git a/examples/events/alb/serverless.yml b/examples/events/alb/serverless.yml
index 44d29ec11..67db2815c 100644
--- a/examples/events/alb/serverless.yml
+++ b/examples/events/alb/serverless.yml
@@ -26,3 +26,12 @@ functions:
           listenerArn: arn:aws:elasticloadbalancing:us-east-1:12345:listener/app/my-load-balancer/50dc6c495c0c9188/50dc6c495c0c9188
           priority: 1
     handler: src/handler.hello
+  hello-headers:
+    events:
+      - alb:
+          conditions:
+            method: GET
+            path: hello-headers
+          listenerArn: arn:aws:elasticloadbalancing:us-east-1:12345:listener/app/my-load-balancer/50dc6c495c0c9188/50dc6c495c0c9188
+          priority: 2
+    handler: src/handler-response-header.responseHeader
diff --git a/examples/events/alb/src/handler-response-header.js b/examples/events/alb/src/handler-response-header.js
new file mode 100644
index 000000000..d469cf8f4
--- /dev/null
+++ b/examples/events/alb/src/handler-response-header.js
@@ -0,0 +1,19 @@
+const { stringify } = JSON
+
+export async function responseHeader(event, context) {
+  return {
+    body: stringify(
+      {
+        context,
+        event,
+      },
+      null,
+      2,
+    ),
+    headers: {
+      "Content-Type": "text/plain",
+      "Set-Cookie": "alb-cookie=works",
+    },
+    statusCode: 200,
+  }
+}
diff --git a/src/events/alb/HttpServer.js b/src/events/alb/HttpServer.js
index 57b42ed5f..172509c06 100644
--- a/src/events/alb/HttpServer.js
+++ b/src/events/alb/HttpServer.js
@@ -267,6 +267,8 @@ export default class HttpServer {
         override: false,
       })
 
+      response.headers = headers
+
       if (typeof result === "string") {
         response.source = stringify(result)
       } else if (result && result.body !== undefined) {
diff --git a/tests/integration/alb-headers/alb-headers.test.js b/tests/integration/alb-headers/alb-headers.test.js
new file mode 100644
index 000000000..b10d7d6ba
--- /dev/null
+++ b/tests/integration/alb-headers/alb-headers.test.js
@@ -0,0 +1,36 @@
+import assert from "node:assert"
+import { join } from "desm"
+import { setup, teardown } from "../../_testHelpers/index.js"
+import { BASE_URL } from "../../config.js"
+
+describe("ALB Headers Tests", function desc() {
+  beforeEach(() =>
+    setup({
+      noPrependStageInUrl: false,
+      servicePath: join(import.meta.url),
+    }),
+  )
+
+  afterEach(() => teardown())
+
+  //
+  ;["GET", "POST"].forEach((method) => {
+    it(`${method} response headers`, async () => {
+      const url = new URL("/dev/response-headers", BASE_URL)
+      url.port = url.port ? "3003" : url.port
+      const options = {
+        method,
+      }
+
+      const response = await fetch(url, options)
+
+      assert.equal(response.status, 200)
+
+      const body = await response.text()
+
+      assert.equal(body, "Plain text")
+      assert.equal(response.headers.get("Content-Type"), "text/plain")
+      assert.equal(response.headers.get("Set-Cookie"), "alb-cookie=works")
+    })
+  })
+})
diff --git a/tests/integration/alb-headers/serverless.yml b/tests/integration/alb-headers/serverless.yml
new file mode 100644
index 000000000..0ebfb72a9
--- /dev/null
+++ b/tests/integration/alb-headers/serverless.yml
@@ -0,0 +1,27 @@
+service: alb-headers
+
+configValidationMode: error
+deprecationNotificationMode: error
+
+plugins:
+  - ../../../src/index.js
+
+provider:
+  architecture: arm64
+  deploymentMethod: direct
+  memorySize: 1024
+  name: aws
+  region: us-east-1
+  runtime: nodejs18.x
+  stage: dev
+  versionFunctions: false
+
+functions:
+  response-headers:
+    events:
+      - alb:
+          conditions:
+            path: response-headers
+          listenerArn: arn:aws:elasticloadbalancing:us-east-1:12345:listener/app/my-load-balancer/50dc6c495c0c9188/50dc6c495c0c9188
+          priority: 1
+    handler: src/handler.responseHeader
diff --git a/tests/integration/alb-headers/src/handler.js b/tests/integration/alb-headers/src/handler.js
new file mode 100644
index 000000000..f73756c69
--- /dev/null
+++ b/tests/integration/alb-headers/src/handler.js
@@ -0,0 +1,10 @@
+export async function responseHeader() {
+  return {
+    body: "Plain text",
+    headers: {
+      "Content-Type": "text/plain",
+      "Set-Cookie": "alb-cookie=works",
+    },
+    statusCode: 200,
+  }
+}
diff --git a/tests/integration/alb-headers/src/package.json b/tests/integration/alb-headers/src/package.json
new file mode 100644
index 000000000..3dbc1ca59
--- /dev/null
+++ b/tests/integration/alb-headers/src/package.json
@@ -0,0 +1,3 @@
+{
+  "type": "module"
+}