Proper support of the trailing headers; redact sensitive headers in the log output (#20)

Author: igor-vovk

README.md

@@ -82,8 +82,8 @@ Run the following command to run Connect-RPC conformance tests:
 docker build . --output "out" --progress=plain
 ```
 
-Execution results are output in STDOUT.
-Diagnostic data from the server itself is output in the `out/out.log` file.
+Execution results are output to STDOUT.
+Diagnostic data from the server itself is written to the log file `out/out.log`.
 
 ### Conformance tests status
 
@@ -92,7 +92,6 @@ Current status: 6/79 tests pass
 Known issues:
 
 * fs2-grpc server implementation doesn't support setting response headers
-* Trailers that are set in http4s aren’t being sent to the client
 * `google.protobuf.Any` serialization doesn't follow Connect-RPC spec
 
 ## Future improvements

conformance/src/main/scala/org/ivovk/connect_rpc_scala/conformance/ConformanceServiceImpl.scala

@@ -12,6 +12,7 @@ import java.util.concurrent.TimeUnit
 import scala.concurrent.duration.Duration
 import scala.jdk.CollectionConverters.*
 
+case class UnaryHandlerResponse(payload: ConformancePayload, trailers: Metadata)
 
 class ConformanceServiceImpl[F[_] : Async] extends ConformanceServiceFs2GrpcTrailers[F, Metadata] {
 
@@ -24,38 +25,38 @@ class ConformanceServiceImpl[F[_] : Async] extends ConformanceServiceFs2GrpcTrai
     ctx: Metadata
   ): F[(UnaryResponse, Metadata)] = {
     for
-      payload <- handleUnaryRequest(
+      res <- handleUnaryRequest(
         request.getResponseDefinition,
         Seq(request.toProtoAny),
         ctx
       )
-    yield (UnaryResponse(payload.some), new Metadata())
+    yield (
+      UnaryResponse(res.payload.some),
+      res.trailers
+    )
   }
 
   override def idempotentUnary(
     request: IdempotentUnaryRequest,
     ctx: Metadata,
   ): F[(IdempotentUnaryResponse, Metadata)] = {
     for
-      payload <- handleUnaryRequest(
+      res <- handleUnaryRequest(
         request.getResponseDefinition,
         Seq(request.toProtoAny),
         ctx
       )
-    yield (IdempotentUnaryResponse(payload.some), new Metadata())
+    yield (
+      IdempotentUnaryResponse(res.payload.some),
+      res.trailers
+    )
   }
 
   private def handleUnaryRequest(
     responseDefinition: UnaryResponseDefinition,
     requests: Seq[com.google.protobuf.any.Any],
     ctx: Metadata,
-  ): F[ConformancePayload] = {
-    //val trailers = new Metadata()
-    //responseDefinition.responseTrailers.foreach { h =>
-    //  val key = Metadata.Key.of(h.name, Metadata.ASCII_STRING_MARSHALLER)
-    //  h.value.foreach(v => trailers.put(key, v))
-    //}
-
+  ): F[UnaryHandlerResponse] = {
     val requestInfo = ConformancePayload.RequestInfo(
       requestHeaders = mkConformanceHeaders(ctx),
       timeoutMs = extractTimeout(ctx),
@@ -77,10 +78,15 @@ class ConformanceServiceImpl[F[_] : Async] extends ConformanceServiceFs2GrpcTrai
         throw new StatusRuntimeException(status)
     }
 
+    val trailers = mkMetadata(responseDefinition.responseTrailers)
+
     Async[F].sleep(Duration(responseDefinition.responseDelayMs, TimeUnit.MILLISECONDS)) *>
-      Async[F].pure(ConformancePayload(
-        responseData.getOrElse(ByteString.EMPTY),
-        requestInfo.some
+      Async[F].pure(UnaryHandlerResponse(
+        payload = ConformancePayload(
+          responseData.getOrElse(ByteString.EMPTY),
+          requestInfo.some
+        ),
+        trailers = trailers
       ))
   }
 
@@ -93,6 +99,16 @@ class ConformanceServiceImpl[F[_] : Async] extends ConformanceServiceFs2GrpcTrai
     }.toSeq
   }
 
+  private def mkMetadata(headers: Seq[Header]): Metadata = {
+    val metadata = new Metadata()
+    headers.foreach { h =>
+      h.value.foreach { v =>
+        metadata.put(keyof(h.name), v)
+      }
+    }
+    metadata
+  }
+
   private def extractTimeout(metadata: Metadata): Option[Long] = {
     Option(metadata.get(keyof("grpc-timeout")))
       .map(v => v.substring(0, v.length - 1).toLong / 1000)

core/src/main/scala/org/ivovk/connect_rpc_scala/ConnectRpcHttpRoutes.scala

@@ -12,7 +12,7 @@ import io.grpc.stub.MetadataUtils
 import org.http4s.*
 import org.http4s.dsl.Http4sDsl
 import org.ivovk.connect_rpc_scala.http.*
-import org.ivovk.connect_rpc_scala.http.Headers.*
+import org.ivovk.connect_rpc_scala.http.Headers.`X-Test-Case-Name`
 import org.ivovk.connect_rpc_scala.http.MessageCodec.given
 import org.ivovk.connect_rpc_scala.http.QueryParams.*
 import org.slf4j.{Logger, LoggerFactory}
@@ -77,15 +77,15 @@ object ConnectRpcHttpRoutes {
 
           _ <- EitherT.cond[F](
             // Support GET-requests for all methods until https://github.com/scalapb/ScalaPB/pull/1774 is merged
-            httpMethod == Method.POST || (httpMethod == Method.GET && method.methodDescriptor.isSafe) || true,
+            httpMethod == Method.POST || (httpMethod == Method.GET && method.descriptor.isSafe) || true,
             (),
             Forbidden(connectrpc.Error(
               code = io.grpc.Status.PERMISSION_DENIED.toConnectCode,
               message = s"Only POST-requests are allowed for method: $grpcMethod".some
             ))
           ).leftSemiflatMap(identity)
 
-          response <- method.methodDescriptor.getType match
+          response <- method.descriptor.getType match
             case MethodType.UNARY =>
               EitherT.right(handleUnary(dsl, method, entity, ipChannel))
             case unsupported =>
@@ -115,7 +115,7 @@ object ConnectRpcHttpRoutes {
 
   private def handleUnary[F[_] : Async](
     dsl: Http4sDsl[F],
-    entry: RegistryEntry,
+    method: MethodRegistry.Entry,
     req: RequestEntity[F],
     channel: Channel
   )(using codec: MessageCodec[F]): F[Response[F]] = {
@@ -130,15 +130,15 @@ object ConnectRpcHttpRoutes {
       }
     }
 
-    given GeneratedMessageCompanion[GeneratedMessage] = entry.requestMessageCompanion
+    given GeneratedMessageCompanion[GeneratedMessage] = method.requestMessageCompanion
 
     req.as[GeneratedMessage]
       .flatMap { message =>
         val responseHeaderMetadata  = new AtomicReference[Metadata]()
         val responseTrailerMetadata = new AtomicReference[Metadata]()
 
         if (logger.isTraceEnabled) {
-          logger.trace(s">>> Method: ${entry.methodDescriptor.getFullMethodName}, Entity: $message")
+          logger.trace(s">>> Method: ${method.descriptor.getFullMethodName}, Entity: $message")
         }
 
         Async[F].fromFuture(Async[F].delay {
@@ -148,26 +148,25 @@ object ConnectRpcHttpRoutes {
               MetadataUtils.newAttachHeadersInterceptor(req.headers.toMetadata),
               MetadataUtils.newCaptureMetadataInterceptor(responseHeaderMetadata, responseTrailerMetadata),
             ),
-            entry.methodDescriptor,
+            method.descriptor,
             CallOptions.DEFAULT
               .pipe(
-                req.headers.get[`Connect-Timeout-Ms`].fold[Endo[CallOptions]](identity) { header =>
-                  _.withDeadlineAfter(header.value, MILLISECONDS)
+                req.timeout.fold[Endo[CallOptions]](identity) { timeout =>
+                  _.withDeadlineAfter(timeout, MILLISECONDS)
                 }
               ),
             message
           )
         }).map { response =>
-          val headers  = responseHeaderMetadata.get().toHeaders
-          val trailers = responseTrailerMetadata.get().toHeaders
+          val headers = org.http4s.Headers.empty ++
+            responseHeaderMetadata.get().toHeaders ++
+            responseTrailerMetadata.get().toTrailingHeaders
 
           if (logger.isTraceEnabled) {
-            logger.trace(s"<<< Headers: $headers, Trailers: $trailers")
+            logger.trace(s"<<< Headers: ${headers.redactSensitive}")
           }
 
-          Response(Ok, headers = headers)
-            .withEntity(response)
-            .withTrailerHeaders(Async[F].pure(trailers))
+          Response(Ok, headers = headers).withEntity(response)
         }
       }
       .recover { case e =>

core/src/main/scala/org/ivovk/connect_rpc_scala/Mappings.scala

@@ -5,8 +5,6 @@ import org.http4s.{Header, Headers}
 import org.typelevel.ci.CIString
 import scalapb.GeneratedMessage
 
-import scala.jdk.CollectionConverters.*
-
 object Mappings extends HeaderMappings, StatusCodeMappings, AnyMappings
 
 trait HeaderMappings {
@@ -25,17 +23,26 @@ trait HeaderMappings {
   }
 
   extension (metadata: Metadata) {
-    def toHeaders: Headers = {
-      val headers = metadata.keys()
-        .asScala.toList
-        .flatMap { key =>
-          metadata.getAll(asciiKey(key)).asScala.map { value =>
-            Header.Raw(CIString(key), value)
-          }
+    private def headers(prefix: String = ""): Headers = {
+      val keys = metadata.keys()
+      if (keys.isEmpty) return Headers.empty
+
+      val b = List.newBuilder[Header.Raw]
+
+      keys.forEach { key =>
+        val name = CIString(prefix + key)
+
+        metadata.getAll(asciiKey(key)).forEach { value =>
+          b += Header.Raw(name, value)
         }
+      }
 
-      Headers(headers)
+      new Headers(b.result())
     }
+
+    def toHeaders: Headers = headers()
+
+    def toTrailingHeaders: Headers = headers("trailer-")
   }
 
 }

core/src/main/scala/org/ivovk/connect_rpc_scala/MethodRegistry.scala

@@ -5,13 +5,13 @@ import scalapb.{GeneratedMessage, GeneratedMessageCompanion}
 
 import scala.jdk.CollectionConverters.*
 
-case class RegistryEntry(
-  requestMessageCompanion: GeneratedMessageCompanion[GeneratedMessage],
-  methodDescriptor: MethodDescriptor[GeneratedMessage, GeneratedMessage],
-)
-
 object MethodRegistry {
 
+  case class Entry(
+    requestMessageCompanion: GeneratedMessageCompanion[GeneratedMessage],
+    descriptor: MethodDescriptor[GeneratedMessage, GeneratedMessage],
+  )
+
   def apply(services: Seq[ServerServiceDefinition]): MethodRegistry = {
     val entries = services
       .flatMap(_.getMethods.asScala)
@@ -30,12 +30,12 @@ object MethodRegistry {
         val requestCompanion = companionField.get(requestMarshaller)
           .asInstanceOf[GeneratedMessageCompanion[GeneratedMessage]]
 
-        val entry = RegistryEntry(
+        val methodEntry = Entry(
           requestMessageCompanion = requestCompanion,
-          methodDescriptor = methodDescriptor,
+          descriptor = methodDescriptor,
         )
 
-        methodDescriptor.getFullMethodName -> entry
+        methodDescriptor.getFullMethodName -> methodEntry
       }
       .toMap
 
@@ -44,9 +44,9 @@ object MethodRegistry {
 
 }
 
-class MethodRegistry private(entries: Map[String, RegistryEntry]) {
+class MethodRegistry private(entries: Map[String, MethodRegistry.Entry]) {
 
-  def get(fullMethodName: String): Option[RegistryEntry] =
+  def get(fullMethodName: String): Option[MethodRegistry.Entry] =
     entries.get(fullMethodName)
 
 }

core/src/main/scala/org/ivovk/connect_rpc_scala/http/MessageCodec.scala

@@ -56,7 +56,7 @@ class JsonMessageCodec[F[_] : Sync : Compression](printer: Printer) extends Mess
     val f = string
       .flatMap { str =>
         if (logger.isTraceEnabled) {
-          logger.trace(s">>> Headers: ${entity.headers}")
+          logger.trace(s">>> Headers: ${entity.headers.redactSensitive}")
           logger.trace(s">>> JSON: $str")
         }
 
@@ -98,7 +98,7 @@ class ProtoMessageCodec[F[_] : Async : Compression] extends MessageCodec[F] {
 
     EitherT.right(f.map { message =>
       if (logger.isTraceEnabled) {
-        logger.trace(s">>> Headers: ${entity.headers}")
+        logger.trace(s">>> Headers: ${entity.headers.redactSensitive}")
         logger.trace(s">>> Proto: ${message.toProtoString}")
       }
 

core/src/main/scala/org/ivovk/connect_rpc_scala/http/RequestEntity.scala

@@ -2,8 +2,9 @@ package org.ivovk.connect_rpc_scala.http
 
 import cats.MonadThrow
 import fs2.Stream
-import org.http4s.{Charset, Headers, Media}
 import org.http4s.headers.`Content-Type`
+import org.http4s.{Charset, Headers, Media}
+import org.ivovk.connect_rpc_scala.http.Headers.`Connect-Timeout-Ms`
 import scalapb.{GeneratedMessage as Message, GeneratedMessageCompanion as Companion}
 
 object RequestEntity {
@@ -26,6 +27,9 @@ case class RequestEntity[F[_]](
   def charset: Charset =
     contentType.flatMap(_.charset).getOrElse(Charset.`UTF-8`)
 
+  def timeout: Option[Long] =
+    headers.get[`Connect-Timeout-Ms`].map(_.value)
+
   def as[A <: Message](using M: MonadThrow[F], codec: MessageCodec[F], cmp: Companion[A]): F[A] =
     M.rethrow(codec.decode(this).value)