From 5ab5dc0c63f495ad244a87991a7b1c2bbc648f32 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lu=C3=ADs=20Duarte?= <lduarte@luisduarte.net>
Date: Wed, 30 Aug 2023 15:23:52 +0000
Subject: [PATCH] Fix [Bug] API GW V2 Payload setting same header on Single and
 MultiValue #1858

---
 ...faultServletToAwsProxyResponseAdapter.java |  3 ++-
 .../test/AwsApiProxyTestServerSpec.groovy     | 22 +++++++++++++++++++
 .../function/aws/proxy/MapCollapseUtils.java  | 10 +++++++--
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/function-aws-api-proxy-test/src/main/java/io/micronaut/function/aws/proxy/test/DefaultServletToAwsProxyResponseAdapter.java b/function-aws-api-proxy-test/src/main/java/io/micronaut/function/aws/proxy/test/DefaultServletToAwsProxyResponseAdapter.java
index 84ad1be503..7e978b2c5b 100644
--- a/function-aws-api-proxy-test/src/main/java/io/micronaut/function/aws/proxy/test/DefaultServletToAwsProxyResponseAdapter.java
+++ b/function-aws-api-proxy-test/src/main/java/io/micronaut/function/aws/proxy/test/DefaultServletToAwsProxyResponseAdapter.java
@@ -73,7 +73,8 @@ private void populateHeaders(@NonNull ConversionService conversionService,
         MutableMapListOfStringAndMapStringConvertibleMultiValue entries = new MutableMapListOfStringAndMapStringConvertibleMultiValue(conversionService, multiValueHeaders, singleHeaders);
 
         for (String name: entries.names()) {
-            response.addHeader(name, String.join(",", entries.getAll(name)));
+            entries.getAll(name)
+                .forEach(headerValue -> response.addHeader(name, headerValue));
         }
     }
 
diff --git a/function-aws-api-proxy-test/src/test/groovy/io/micronaut/function/aws/proxy/test/AwsApiProxyTestServerSpec.groovy b/function-aws-api-proxy-test/src/test/groovy/io/micronaut/function/aws/proxy/test/AwsApiProxyTestServerSpec.groovy
index fbd43183bc..8b0ad14b53 100644
--- a/function-aws-api-proxy-test/src/test/groovy/io/micronaut/function/aws/proxy/test/AwsApiProxyTestServerSpec.groovy
+++ b/function-aws-api-proxy-test/src/test/groovy/io/micronaut/function/aws/proxy/test/AwsApiProxyTestServerSpec.groovy
@@ -1,5 +1,6 @@
 package io.micronaut.function.aws.proxy.test
 
+import io.micronaut.http.HttpHeaders
 import io.micronaut.http.HttpRequest
 import io.micronaut.http.HttpResponse
 import io.micronaut.http.HttpStatus
@@ -79,6 +80,19 @@ class AwsApiProxyTestServerSpec extends Specification {
         response.body.get() == (1..256).collect { it as byte } as byte[]
     }
 
+    @Issue('https://github.com/micronaut-projects/micronaut-aws/issues/1858')
+    void 'can return a ByteArray'() {
+        when:
+        HttpResponse<?> response = client.toBlocking()
+                .exchange(HttpRequest.GET('/multiple-headers'))
+
+        then:
+        response.status == HttpStatus.OK
+        response.headers.getAll(HttpHeaders.LINK) ==
+                ['<https://example1.com>; rel="next"', '<https://example2.com>; rel="prev"']
+        response.headers.get(HttpHeaders.CACHE_CONTROL) == "no-cache"
+    }
+
     @Controller
     static class TestController {
         @Get(value = '/test', produces = MediaType.TEXT_PLAIN)
@@ -105,5 +119,13 @@ class AwsApiProxyTestServerSpec extends Specification {
         HttpResponse<byte[]> byteArray() {
             return HttpResponse.ok((1..256).collect { it as byte } as byte[])
         }
+
+        @Get(value = '/multiple-headers', processes = MediaType.TEXT_PLAIN)
+        HttpResponse emptyBody() {
+            return HttpResponse.ok()
+                        .header(HttpHeaders.LINK, '<https://example1.com>; rel="next"')
+                        .header(HttpHeaders.LINK, '<https://example2.com>; rel="prev"')
+                        .header(HttpHeaders.CACHE_CONTROL,  "no-cache")
+        }
     }
 }
diff --git a/function-aws-api-proxy/src/main/java/io/micronaut/function/aws/proxy/MapCollapseUtils.java b/function-aws-api-proxy/src/main/java/io/micronaut/function/aws/proxy/MapCollapseUtils.java
index e1534240be..678b442671 100644
--- a/function-aws-api-proxy/src/main/java/io/micronaut/function/aws/proxy/MapCollapseUtils.java
+++ b/function-aws-api-proxy/src/main/java/io/micronaut/function/aws/proxy/MapCollapseUtils.java
@@ -43,7 +43,10 @@ private MapCollapseUtils() {
     public static Map<String, String> getSingleValueHeaders(MutableHttpHeaders headers) {
         Map<String, String> result = new HashMap<>();
         for (String paramName : headers.names()) {
-            result.put(paramName, headers.get(paramName));
+            List<String> values = headers.getAll(paramName);
+            if (values.size() == 1) {
+                result.put(paramName, values.get(0));
+            }
         }
         return result;
     }
@@ -57,7 +60,10 @@ public static Map<String, String> getSingleValueHeaders(MutableHttpHeaders heade
     public static Map<String, List<String>> getMultiHeaders(MutableHttpHeaders headers) {
         Map<String, List<String>> result = new HashMap<>();
         for (String paramName : headers.names()) {
-            result.put(paramName, headers.getAll(paramName));
+            List<String> values = headers.getAll(paramName);
+            if (values.size() > 1) {
+                result.put(paramName, values);
+            }
         }
         return result;
     }