2.x.x - Trailer headers (#330)

* Add support for trailer headers * Add support for testing trailer headers * Fix benchmarks * Fix benchmarks v2 * Add HBResponseBody.withTrailingHeaders * Comment
hummingbird-project · Jan 11, 2024 · dde5dfa · dde5dfa
1 parent 1e3b93d
commit dde5dfa
Show file tree

Hide file tree

Showing 13 changed files with 70 additions and 23 deletions.
diff --git a/Benchmarks/Benchmarks/Router/RouterBenchmarks.swift b/Benchmarks/Benchmarks/Router/RouterBenchmarks.swift
@@ -65,7 +65,7 @@ extension Benchmark {
                         let hbRequest = HBRequest(head: request, body: requestBody)
                         group.addTask {
                             let response = try await responder.respond(to: hbRequest, context: context)
-                            try await response.body.write(BenchmarkBodyWriter())
+                            _ = try await response.body.write(BenchmarkBodyWriter())
                         }
                         try await writeBody(requestBodyStream)
                         requestBodyStream.finish()
@@ -126,6 +126,7 @@ func routerBenchmarks() {
                 for try await buffer in request.body {
                     try await writer.write(buffer)
                 }
+                return nil
             })
         }
     }

diff --git a/Sources/HummingbirdCore/Response/ResponseBody.swift b/Sources/HummingbirdCore/Response/ResponseBody.swift
@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import HTTPTypes
 import NIOCore
 
 public protocol HBResponseBodyWriter {
@@ -20,15 +21,15 @@ public protocol HBResponseBodyWriter {
 
 /// Response body
 public struct HBResponseBody: Sendable {
-    public let write: @Sendable (any HBResponseBodyWriter) async throws -> Void
+    public let write: @Sendable (any HBResponseBodyWriter) async throws -> HTTPFields?
     public let contentLength: Int?
 
     /// Initialise HBResponseBody with closure writing body contents
     /// - Parameters:
     ///   - contentLength: Optional length of body
     ///   - write: closure provided with `writer` type that can be used to write to response body
     public init(contentLength: Int? = nil, _ write: @Sendable @escaping (any HBResponseBodyWriter) async throws -> Void) {
-        self.write = write
+        self.write = { try await write($0); return nil }
         self.contentLength = contentLength
     }
 
@@ -40,7 +41,9 @@ public struct HBResponseBody: Sendable {
     /// Initialise HBResponseBody that contains a single ByteBuffer
     /// - Parameter byteBuffer: ByteBuffer to write
     public init(byteBuffer: ByteBuffer) {
-        self.init(contentLength: byteBuffer.readableBytes) { writer in try await writer.write(byteBuffer) }
+        self.init(contentLength: byteBuffer.readableBytes) { writer in
+            try await writer.write(byteBuffer)
+        }
     }
 
     /// Initialise HBResponseBody with an AsyncSequence of ByteBuffers
@@ -50,6 +53,30 @@ public struct HBResponseBody: Sendable {
             for try await buffer in asyncSequence {
                 try await writer.write(buffer)
             }
+            return
         }
     }
+
+    /// Create HBResponseBody that returns trailing headers from its closure once all the 
+    /// body parts have been written 
+    /// - Parameters:
+    ///   - contentLength: Optional length of body
+    ///   - write: closure provided with `writer` type that can be used to write to response body
+    ///         trailing headers are returned from the closure after all the body parts have been
+    ///         written
+    public static func withTrailingHeaders(
+        contentLength: Int? = nil, 
+        _ write: @Sendable @escaping (any HBResponseBodyWriter) async throws -> HTTPFields?
+    ) -> Self {
+        self.init(contentLength: contentLength, write: write)
+    }
+
+    /// Initialise HBResponseBody with closure writing body contents
+    /// 
+    /// This version of init is private and only available via ``withTrailingHeaders`` because
+    /// if it is public the compiler gets confused when a complex closure is provided.
+    private init(contentLength: Int? = nil, write: @Sendable @escaping (any HBResponseBodyWriter) async throws -> HTTPFields?) {
+        self.write = { return try await write($0) }
+        self.contentLength = contentLength
+    }
 }
diff --git a/Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift b/Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift
@@ -66,8 +66,8 @@ extension HTTPChannelHandler {
                                 }
                                 do {
                                     try await outbound.write(.head(response.head))
-                                    try await response.body.write(responseWriter)
-                                    try await outbound.write(.end(nil))
+                                    let tailHeaders = try await response.body.write(responseWriter)
+                                    try await outbound.write(.end(tailHeaders))
                                     // flush request body
                                     for try await _ in request.body {}
                                 } catch {

diff --git a/Sources/HummingbirdXCT/Application+XCT.swift b/Sources/HummingbirdXCT/Application+XCT.swift
@@ -33,9 +33,9 @@ public struct XCTTestingSetup {
 
     /// Test writing requests directly to router.
     public static var router: XCTTestingSetup { .init(value: .router) }
-    /// Sets up a live server and execute tests using a HTTP client.
+    /// Sets up a live server and execute tests using a HTTP client. Only supports HTTP1
     public static var live: XCTTestingSetup { .init(value: .live) }
-    /// Sets up a live server and execute tests using a HTTP client.
+    /// Sets up a live server and execute tests using a HTTP client. Does not support trailer headers
     public static func ahc(_ scheme: XCTScheme) -> XCTTestingSetup { .init(value: .ahc(scheme)) }
 }
 

diff --git a/Sources/HummingbirdXCT/HBXCTApplication.swift b/Sources/HummingbirdXCT/HBXCTApplication.swift
@@ -26,6 +26,8 @@ public struct HBXCTResponse: Sendable {
     public var headers: HTTPFields { self.head.headerFields }
     /// response body
     public let body: ByteBuffer?
+    /// trailer headers
+    public let trailerHeaders: HTTPFields?
 }
 
 /// Errors thrown by XCT framework.

diff --git a/Sources/HummingbirdXCT/HBXCTAsyncHTTPClient.swift b/Sources/HummingbirdXCT/HBXCTAsyncHTTPClient.swift
@@ -45,7 +45,7 @@ final class HBXCTAsyncHTTPClient<App: HBApplicationProtocol>: HBXCTApplication {
             request.body = body.map { .bytes($0) }
             let response = try await client.execute(request, deadline: .now() + self.timeout)
             let responseHead = HTTPResponseHead(version: response.version, status: response.status, headers: response.headers)
-            return try await .init(head: .init(responseHead), body: response.body.collect(upTo: .max))
+            return try await .init(head: .init(responseHead), body: response.body.collect(upTo: .max), trailerHeaders: nil)
         }
     }
 

diff --git a/Sources/HummingbirdXCT/HBXCTLive.swift b/Sources/HummingbirdXCT/HBXCTLive.swift
@@ -38,7 +38,7 @@ final class HBXCTLive<App: HBApplicationProtocol>: HBXCTApplication {
             headers[.connection] = "keep-alive"
             let request = HBXCTClient.Request(uri, method: method, authority: "localhost", headers: headers, body: body)
             let response = try await client.execute(request)
-            return .init(head: response.head, body: response.body)
+            return .init(head: response.head, body: response.body, trailerHeaders: response.trailerHeaders)
         }
     }
 

diff --git a/Sources/HummingbirdXCT/HBXCTRouter.swift b/Sources/HummingbirdXCT/HBXCTRouter.swift
@@ -91,10 +91,10 @@ struct HBXCTRouter<Responder: HBResponder>: HBXCTApplication where Responder.Con
                         response = HBResponse(status: .internalServerError)
                     }
                     let responseWriter = RouterResponseWriter()
-                    try await response.body.write(responseWriter)
+                    let trailerHeaders = try await response.body.write(responseWriter)
                     for try await _ in request.body {}
                     return responseWriter.collated.withLockedValue { collated in
-                        HBXCTResponse(head: response.head, body: collated)
+                        HBXCTResponse(head: response.head, body: collated, trailerHeaders: trailerHeaders)
                     }
                 }
 

diff --git a/Sources/HummingbirdXCT/XCTClient+types.swift b/Sources/HummingbirdXCT/XCTClient+types.swift
@@ -51,10 +51,12 @@ extension HBXCTClient {
     public struct Response: Sendable {
         public var head: HTTPResponse
         public var body: ByteBuffer?
+        public var trailerHeaders: HTTPFields?
 
-        public init(head: HTTPResponse, body: ByteBuffer? = nil) {
+        public init(head: HTTPResponse, body: ByteBuffer? = nil, trailerHeaders: HTTPFields? = nil) {
             self.head = head
             self.body = body
+            self.trailerHeaders = trailerHeaders
         }
 
         public var status: HTTPResponse.Status {

diff --git a/Sources/HummingbirdXCT/XCTClient.swift b/Sources/HummingbirdXCT/XCTClient.swift
@@ -234,21 +234,21 @@ public struct HBXCTClient: Sendable {
             case (.body(var part), .body(let head, var body)):
                 body.writeBuffer(&part)
                 self.state = .body(head, body)
-            case (.end(let tailHeaders), .body(let head, let body)):
-                assert(tailHeaders == nil, "Unexpected tail headers")
+            case (.end(let trailerHeaders), .body(let head, let body)):
                 let response = HBXCTClient.Response(
                     head: head,
-                    body: body
+                    body: body,
+                    trailerHeaders: trailerHeaders
                 )
                 if context.channel.isActive {
                     context.fireChannelRead(wrapInboundOut(response))
                 }
                 self.state = .idle
-            case (.end(let tailHeaders), .head(let head)):
-                assert(tailHeaders == nil, "Unexpected tail headers")
+            case (.end(let trailerHeaders), .head(let head)):
                 let response = HBXCTClient.Response(
                     head: head,
-                    body: nil
+                    body: nil,
+                    trailerHeaders: trailerHeaders
                 )
                 if context.channel.isActive {
                     context.fireChannelRead(wrapInboundOut(response))

diff --git a/Tests/HummingbirdCoreTests/CoreTests.swift b/Tests/HummingbirdCoreTests/CoreTests.swift
@@ -184,6 +184,19 @@ class HummingBirdCoreTests: XCTestCase {
         }
     }
 
+    func testTrailerHeaders() async throws {
+        try await testServer(
+            responder: { _, _ in .init(status: .ok, body: .withTrailingHeaders { _ in return [.contentType: "text"] }) },
+            httpChannelSetup: .http1(),
+            configuration: .init(address: .hostname(port: 0)),
+            eventLoopGroup: Self.eventLoopGroup,
+            logger: Logger(label: "HB")
+        ) { client in
+            let response = try await client.get("/")
+            XCTAssertEqual(response.trailerHeaders?[.contentType], "text")
+        }
+    }
+
     func testChannelHandlerErrorPropagation() async throws {
         class CreateErrorHandler: ChannelInboundHandler, RemovableChannelHandler {
             typealias InboundIn = HTTPRequestPart

diff --git a/Tests/HummingbirdRouterTests/MiddlewareTests.swift b/Tests/HummingbirdRouterTests/MiddlewareTests.swift
@@ -135,9 +135,10 @@ final class MiddlewareTests: XCTestCase {
             func handle(_ request: HBRequest, context: Context, next: (HBRequest, Context) async throws -> HBResponse) async throws -> HBResponse {
                 let response = try await next(request, context)
                 var editedResponse = response
-                editedResponse.body = .init { writer in
+                editedResponse.body = .withTrailingHeaders { writer in
                     let transformWriter = TransformWriter(parentWriter: writer, allocator: context.allocator)
-                    try await response.body.write(transformWriter)
+                    let tailHeaders = try await response.body.write(transformWriter)
+                    return tailHeaders
                 }
                 return editedResponse
             }

diff --git a/Tests/HummingbirdTests/MiddlewareTests.swift b/Tests/HummingbirdTests/MiddlewareTests.swift
@@ -147,9 +147,10 @@ final class MiddlewareTests: XCTestCase {
             public func handle(_ request: HBRequest, context: Context, next: (HBRequest, Context) async throws -> HBResponse) async throws -> HBResponse {
                 let response = try await next(request, context)
                 var editedResponse = response
-                editedResponse.body = .init { writer in
+                editedResponse.body = .withTrailingHeaders { writer in
                     let transformWriter = TransformWriter(parentWriter: writer, allocator: context.allocator)
-                    try await response.body.write(transformWriter)
+                    let tailHeaders = try await response.body.write(transformWriter)
+                    return tailHeaders
                 }
                 return editedResponse
             }