Repository: http4s/http4s
Pull request #6756 (2022-10-15 18:33:52)

Encode trailer headers in Ember
Fixes another issue discovered while working on https://github.com/http4s/http4s/pull/6753.


Patch:
diff --git a/ember-core/shared/src/main/scala/org/http4s/ember/core/ChunkedEncoding.scala b/ember-core/shared/src/main/scala/org/http4s/ember/core/ChunkedEncoding.scala
index 3f9d382871..936dc29308 100644
--- a/ember-core/shared/src/main/scala/org/http4s/ember/core/ChunkedEncoding.scala
+++ b/ember-core/shared/src/main/scala/org/http4s/ember/core/ChunkedEncoding.scala
@@ -32,8 +32,10 @@ import cats.effect.kernel.Deferred
 import cats.effect.kernel.Ref
 import cats.syntax.all._
 import fs2._
+import org.http4s.internal.appendSanitized
 import scodec.bits.ByteVector
 
+import java.nio.charset.StandardCharsets
 import scala.util.control.NonFatal
 
 import Shared._
@@ -120,15 +122,15 @@ private[ember] object ChunkedEncoding {
   private def parseTrailers[F[_]: MonadThrow](
       maxHeaderSize: Int
   )(buffer: Array[Byte], read: F[Option[Chunk[Byte]]]): F[Trailers] =
-    if (buffer.startsWith(Shared.crlf.toArray)) {
-      Trailers(Headers.empty, buffer.drop(crlf.size.toInt)).pure[F]
-    } else if (buffer.length < 2) {
+    if (buffer.length < 2) {
       read.flatMap {
         case None =>
           MonadThrow[F].raiseError(EmberException.ReachedEndOfStream())
         case Some(chunk) =>
-          parseTrailers(maxHeaderSize)(buffer ++ chunk.toArray[Byte], read)
+          parseTrailers(maxHeaderSize)(Util.concatBytes(buffer, chunk), read)
       }
+    } else if (buffer(0) == '\r' && buffer(1) == '\n') {
+      Trailers(Headers.empty, buffer.drop(2)).pure[F]
     } else {
       Parser.MessageP
         .recurseFind(buffer, read, maxHeaderSize, Parser.HeaderP.ParserState.initial)(
@@ -140,20 +142,44 @@ private[ember] object ChunkedEncoding {
     }
 
   private[this] val lastChunk: Stream[fs2.Pure, Byte] = {
-    val bytes = Array[Byte]('0', '\r', '\n', '\r', '\n')
+    val bytes = Array[Byte]('0', '\r', '\n')
+    Stream.chunk(Chunk.array(bytes))
+  }
+
+  private[this] val finalCrlf: Stream[fs2.Pure, Byte] = {
+    val bytes = Array[Byte]('\r', '\n')
     Stream.chunk(Chunk.array(bytes))
   }
 
   /** Encodes chunk of bytes to http chunked encoding.
     */
-  def encode[F[_]]: Pipe[F, Byte, Byte] = {
+  def encode[F[_]: Functor](trailers: F[Headers]): Pipe[F, Byte, Byte] = {
     def encodeChunk(bv: ByteVector): Chunk[Byte] =
-      Chunk.byteVector(
-        ByteVector.view(bv.size.toHexString.toUpperCase.getBytes) ++ crlf ++ bv ++ crlf
-      )
-    _.mapChunks { ch =>
-      encodeChunk(ch.toByteVector)
-    } ++ lastChunk
+      if (bv.isEmpty) Chunk.empty
+      else
+        Chunk.byteVector(
+          ByteVector.view(bv.size.toHexString.toUpperCase.getBytes) ++ crlf ++ bv ++ crlf
+        )
+    _.mapChunks(ch => encodeChunk(ch.toByteVector)) ++
+      lastChunk ++
+      Stream.evalUnChunk(trailers.map(encodeTrailers(_))) ++
+      finalCrlf
+  }
+
+  private def encodeTrailers(trailers: Headers): Chunk[Byte] = {
+    val stringBuilder = new StringBuilder()
+
+    trailers.foreach { h =>
+      if (h.isNameValid) {
+        stringBuilder
+          .append(h.name)
+          .append(": ")
+        appendSanitized(stringBuilder, h.value)
+        stringBuilder.append("\r\n")
+      }
+    }
+
+    Chunk.array(stringBuilder.toString.getBytes(StandardCharsets.ISO_8859_1))
   }
 
   /** yields to size of header in case the chunked header was succesfully parsed, else yields to None */
diff --git a/ember-core/shared/src/main/scala/org/http4s/ember/core/Encoder.scala b/ember-core/shared/src/main/scala/org/http4s/ember/core/Encoder.scala
index 193cf34521..0d25205c52 100644
--- a/ember-core/shared/src/main/scala/org/http4s/ember/core/Encoder.scala
+++ b/ember-core/shared/src/main/scala/org/http4s/ember/core/Encoder.scala
@@ -16,6 +16,7 @@
 
 package org.http4s.ember.core
 
+import cats.Applicative
 import cats.ApplicativeThrow
 import fs2._
 import org.http4s._
@@ -33,7 +34,10 @@ private[ember] object Encoder {
   private[this] final val chunkedTransferEncodingHeaderRaw = "Transfer-Encoding: chunked"
   private[this] final val zeroContentLengthRaw = "Content-Length: 0"
 
-  def respToBytes[F[_]](resp: Response[F], writeBufferSize: Int = 32 * 1024): Stream[F, Byte] = {
+  def respToBytes[F[_]: Applicative](
+      resp: Response[F],
+      writeBufferSize: Int = 32 * 1024,
+  ): Stream[F, Byte] = {
     var chunked = resp.isChunked
     // resp.status.isEntityAllowed TODO
     val initSection = {
@@ -73,7 +77,9 @@ private[ember] object Encoder {
     }
 
     if (chunked)
-      Stream.chunk(Chunk.array(initSection)) ++ resp.body.through(ChunkedEncoding.encode[F])
+      Stream.chunk(Chunk.array(initSection)) ++ resp.body.through(
+        ChunkedEncoding.encode[F](resp.trailerHeaders)
+      )
     else
       (Stream.chunk(Chunk.array(initSection)) ++ resp.body)
         .chunkMin(writeBufferSize)
@@ -141,7 +147,9 @@ private[ember] object Encoder {
         stringBuilder.toString.getBytes(StandardCharsets.ISO_8859_1)
       }
       if (chunked)
-        Stream.chunk(Chunk.array(initSection)) ++ req.body.through(ChunkedEncoding.encode[F])
+        Stream.chunk(Chunk.array(initSection)) ++ req.body.through(
+          ChunkedEncoding.encode[F](req.trailerHeaders)
+        )
       else
         (Stream.chunk(Chunk.array(initSection)) ++ req.body)
           .chunkMin(writeBufferSize)


Test patch:
diff --git a/ember-core/shared/src/test/scala/org/http4s/ember/core/EncoderSuite.scala b/ember-core/shared/src/test/scala/org/http4s/ember/core/EncoderSuite.scala
index 3e4eca4ffe..8f9b63948f 100644
--- a/ember-core/shared/src/test/scala/org/http4s/ember/core/EncoderSuite.scala
+++ b/ember-core/shared/src/test/scala/org/http4s/ember/core/EncoderSuite.scala
@@ -177,6 +177,32 @@ class EncoderSuite extends Http4sSuite {
     Helpers.encodeResponseRig(resp).assertEquals(expected)
   }
 
+  test("respToBytes should encode trailer headers") {
+    val resp = Response[IO](Status.Ok)
+      .withEntity(Stream[IO, String]("Mozilla", "Developer", "Network"))
+      .putHeaders("Trailer" -> "Expires")
+      .withTrailerHeaders(IO(Headers("Expires" -> "Wed, 21 Oct 2015 07:28:00 GMT")))
+
+    val expected =
+      """HTTP/1.1 200 OK
+      |Content-Type: text/plain; charset=UTF-8
+      |Transfer-Encoding: chunked
+      |Trailer: Expires
+      |
+      |7
+      |Mozilla
+      |9
+      |Developer
+      |7
+      |Network
+      |0
+      |Expires: Wed, 21 Oct 2015 07:28:00 GMT
+      |
+      |""".stripMargin
+
+    Helpers.encodeResponseRig(resp).assertEquals(expected)
+  }
+
   test("encoder a response where entity is not allowed correctly") {
     val resp = Response[IO](Status.NoContent)
     val expected =