ZIO HTTP Server should not fold multiple Set-Cookie headers (#3723)

softwaremill · Apr 26, 2024 · 9e819cd · 9e819cd
1 parent a2a4f76
commit 9e819cd
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 9 deletions.
diff --git a/server/tests/src/main/scala/sttp/tapir/server/tests/ServerBasicTests.scala b/server/tests/src/main/scala/sttp/tapir/server/tests/ServerBasicTests.scala
@@ -224,6 +224,24 @@ class ServerBasicTests[F[_], OPTIONS, ROUTE](
         IO.pure(succeed)
       }
     },
+    testServerLogic(
+      endpoint.in("cookies-test").get.out(setCookies).serverLogic { _ =>
+        pureResult(
+          List(
+            CookieWithMeta.unsafeApply("name1", "value1", path = Some("/path1")),
+            CookieWithMeta.unsafeApply("name2", "value2", path = Some("/path2"))
+          ).asRight[Unit]
+        )
+      },
+      "should not fold multiple Set-Cookie headers"
+    ) { (backend, baseUri) =>
+      basicRequest.response(asStringAlways).get(uri"$baseUri/cookies-test").send(backend).map { r =>
+        r.headers should contain allOf(
+          Header.setCookie(CookieWithMeta.unsafeApply("name1", "value1", path = Some("/path1"))),
+          Header.setCookie(CookieWithMeta.unsafeApply("name2", "value2", path = Some("/path2")))
+        )
+      }
+    },
     // Reproduces https://github.com/http4s/http4s/security/advisories/GHSA-5vcm-3xc3-w7x3
     testServer(in_query_out_cookie_raw, "should be invulnerable to response splitting from unsanitized headers")((q: String) =>
       pureResult((q, "test").asRight[Unit])

diff --git a/server/zio-http-server/src/main/scala/sttp/tapir/server/ziohttp/ZioHttpInterpreter.scala b/server/zio-http-server/src/main/scala/sttp/tapir/server/ziohttp/ZioHttpInterpreter.scala
@@ -2,19 +2,15 @@ package sttp.tapir.server.ziohttp
 
 import sttp.capabilities.WebSockets
 import sttp.capabilities.zio.ZioStreams
-import sttp.model.Method
 import sttp.model.{Header => SttpHeader}
 import sttp.monad.MonadError
 import sttp.tapir.server.interceptor.RequestResult
 import sttp.tapir.server.interceptor.reject.RejectInterceptor
-import sttp.tapir.server.interpreter.FilterServerEndpoints
-import sttp.tapir.server.interpreter.ServerInterpreter
+import sttp.tapir.server.interpreter.{FilterServerEndpoints, ServerInterpreter}
 import sttp.tapir.server.model.ServerResponse
 import sttp.tapir.ztapir._
 import zio._
-import zio.http.{Header => ZioHttpHeader}
-import zio.http.{Headers => ZioHttpHeaders}
-import zio.http._
+import zio.http.{Header => ZioHttpHeader, Headers => ZioHttpHeaders, _}
 
 trait ZioHttpInterpreter[R] {
   def zioHttpServerOptions: ZioHttpServerOptions[R] = ZioHttpServerOptions.default
@@ -88,7 +84,7 @@ trait ZioHttpInterpreter[R] {
       resp: ServerResponse[ZioResponseBody],
       body: Option[ZioHttpResponseBody]
   ): UIO[Response] = {
-    val baseHeaders = resp.headers.groupBy(_.name).map(sttpToZioHttpHeader).toList
+    val baseHeaders = resp.headers.groupBy(_.name).flatMap(sttpToZioHttpHeader).toList
     val allHeaders = body.flatMap(_.contentLength) match {
       case Some(contentLength) if resp.contentLength.isEmpty => ZioHttpHeader.ContentLength(contentLength) :: baseHeaders
       case _                                                 => baseHeaders
@@ -110,8 +106,17 @@ trait ZioHttpInterpreter[R] {
     )
   }
 
-  private def sttpToZioHttpHeader(hl: (String, Seq[SttpHeader])): ZioHttpHeader =
-    ZioHttpHeader.Custom(hl._1, hl._2.map(_.value).mkString(", "))
+  private def sttpToZioHttpHeader(hl: (String, Seq[SttpHeader])): Seq[ZioHttpHeader] = {
+    hl._1.toLowerCase match {
+      case "set-cookie" =>
+        hl._2.map(_.value).map { rawValue =>
+          ZioHttpHeader.SetCookie.parse(rawValue).toOption.getOrElse {
+            ZioHttpHeader.Custom(hl._1, rawValue)
+          }
+        }
+      case _ => List(ZioHttpHeader.Custom(hl._1, hl._2.map(_.value).mkString(", ")))
+    }
+  }
 }
 
 object ZioHttpInterpreter {