From 9da75045a68758bdb63fe746b10604700ed625a6 Mon Sep 17 00:00:00 2001
From: Frederic Germain <frederic.germain@gmail.com>
Date: Thu, 25 Mar 2021 17:13:54 +0100
Subject: [PATCH] feat: support both enable/disable attribute for multi value
 headers on ALB

Setting found in target groups of lambda type.
---
 src/event-sources/aws/alb.js | 12 ++++++++++--
 src/transport.js             |  3 +++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/event-sources/aws/alb.js b/src/event-sources/aws/alb.js
index 9edf1fc7..235db3f0 100644
--- a/src/event-sources/aws/alb.js
+++ b/src/event-sources/aws/alb.js
@@ -2,16 +2,24 @@ const { getRequestValuesFromEvent, getMultiValueHeaders } = require('../utils')
 
 const getRequestValuesFromAlbEvent = ({ event }) => getRequestValuesFromEvent({ event })
 const getResponseToAlb = ({
+  event,
   statusCode,
   body,
-  headers,
+  headers: responseHeaders,
   isBase64Encoded
 }) => {
-  const multiValueHeaders = getMultiValueHeaders({ headers })
+  const multiValueHeaders = !event.headers ? getMultiValueHeaders({ headers: responseHeaders }) : undefined
+  const headers = event.headers
+    ? Object.entries(responseHeaders).reduce((acc, [k, v]) => {
+      acc[k] = Array.isArray(v) ? v[0] : v
+      return acc
+    }, {})
+    : undefined
 
   return {
     statusCode,
     body,
+    headers,
     multiValueHeaders,
     isBase64Encoded
   }
diff --git a/src/transport.js b/src/transport.js
index a647065d..f378bc8d 100644
--- a/src/transport.js
+++ b/src/transport.js
@@ -10,6 +10,7 @@ function forwardResponse ({
   response,
   resolver,
   eventSource,
+  event,
   log
 }) {
   const statusCode = response.statusCode
@@ -30,6 +31,7 @@ function forwardResponse ({
   })
 
   const successResponse = eventSource.getResponse({
+    event,
     statusCode,
     body,
     headers,
@@ -141,6 +143,7 @@ async function forwardRequestToNodeServer ({
     response,
     resolver,
     eventSource,
+    event,
     log
   })
   return response