Fix headers in ApiGatewayProxyHandler

andre-bisa · Mar 12, 2024 · 4a9c35e · 4a9c35e
1 parent 5a9f0d1
commit 4a9c35e
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 16 deletions.
diff --git a/lambda-http4s/src/main/scala/feral/lambda/http4s/ApiGatewayProxyHandler.scala b/lambda-http4s/src/main/scala/feral/lambda/http4s/ApiGatewayProxyHandler.scala
@@ -32,6 +32,7 @@ import org.http4s.HttpRoutes
 import org.http4s.Method
 import org.http4s.Request
 import org.http4s.Uri
+import org.typelevel.ci.CIString
 
 object ApiGatewayProxyHandler {
 
@@ -63,24 +64,13 @@ object ApiGatewayProxyHandler {
       event: ApiGatewayProxyEvent): F[Request[F]] = {
     val queryString: String = getMultiValueQueryStringParameters(
       event.multiValueQueryStringParameters)
+
     val uriString: String = event.path + (if (queryString.nonEmpty) s"?$queryString" else "")
 
     for {
       method <- Method.fromString(event.httpMethod).liftTo[F]
       uri <- Uri.fromString(uriString).liftTo[F]
-      headers = {
-        val builder = List.newBuilder[Header.Raw]
-        event.headers.foreach { h => h.foreachEntry(builder += Header.Raw(_, _)) }
-        event.multiValueHeaders.foreach { hMap =>
-          hMap.foreach {
-            case (key, values) =>
-              if (!event.headers.exists(_.contains(key))) {
-                values.foreach(value => builder += Header.Raw(key, value))
-              }
-          }
-        }
-        Headers(builder.result())
-      }
+      headers = Headers(getMultiValueHeaders(event.multiValueHeaders))
       readBody =
         if (event.isBase64Encoded)
           fs2.text.base64.decode[F]
@@ -105,6 +95,15 @@ object ApiGatewayProxyHandler {
         .mkString("&")
     }
 
+  private def getMultiValueHeaders(
+      multiValueHeaders: Option[Map[CIString, List[String]]]): List[Header.Raw] = {
+    for {
+      headerMap <- multiValueHeaders.toList
+      (key, values) <- headerMap
+      headerRaw <- values.map(value => Header.Raw(key, value))
+    } yield headerRaw
+  }
+
   @deprecated("Use ApiGatewayProxyHandlerV2", "0.3.0")
   def apply[F[_]: ApiGatewayProxyInvocationV2: Concurrent](
       routes: HttpRoutes[F]): F[Option[ApiGatewayProxyStructuredResultV2]] = httpRoutes(routes)

diff --git a/lambda-http4s/src/test/scala/feral/lambda/http4s/ApiGatewayProxyHandlerSuite.scala b/lambda-http4s/src/test/scala/feral/lambda/http4s/ApiGatewayProxyHandlerSuite.scala
@@ -47,7 +47,9 @@ class ApiGatewayProxyHandlerSuite extends CatsEffectSuite {
     "CloudFront-Forwarded-Proto" -> "https",
     "X-Forwarded-For" -> "127.0.0.1, 127.0.0.2",
     "CloudFront-Is-Tablet-Viewer" -> "false",
-    "Accept" -> "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"
+    "Accept" -> "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
+    "X-MultiHeader" -> "foo",
+    "X-MultiHeader" -> "bar"
   )
 
   val expectedBody: String = """{"test":"body"}"""

diff --git a/lambda/shared/src/test/scala/feral/lambda/events/ApiGatewayProxyEventSuite.scala b/lambda/shared/src/test/scala/feral/lambda/events/ApiGatewayProxyEventSuite.scala
@@ -43,7 +43,8 @@ object ApiGatewayProxyEventSuite {
     },
     "multiValueQueryStringParameters": {
       "foo": [
-        "bar"
+        "bar",
+        "baz"
       ]
     },
     "pathParameters": {
@@ -70,7 +71,8 @@ object ApiGatewayProxyEventSuite {
       "X-Amz-Cf-Id": "cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA==",
       "X-Forwarded-For": "127.0.0.1, 127.0.0.2",
       "X-Forwarded-Port": "443",
-      "X-Forwarded-Proto": "https"
+      "X-Forwarded-Proto": "https",
+      "X-MultiHeader": "foo"
     },
     "multiValueHeaders": {
       "Accept": [
@@ -126,6 +128,10 @@ object ApiGatewayProxyEventSuite {
       ],
       "X-Forwarded-Proto": [
         "https"
+      ],
+      "X-MultiHeader": [
+        "foo",
+        "bar"
       ]
     },
     "requestContext": {