Cleanup ResponseBodyWriter, add ResponseBody.map (#526)

Sources/HummingbirdCore/Response/ResponseBody.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Hummingbird server framework project
 //
-// Copyright (c) 2021-2021 the Hummingbird authors
+// Copyright (c) 2021-2024 the Hummingbird authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -15,10 +15,6 @@
 import HTTPTypes
 import NIOCore
 
-public protocol ResponseBodyWriter {
-    func write(_ buffer: ByteBuffer) async throws
-}
-
 /// Response body
 public struct ResponseBody: Sendable {
     public let write: @Sendable (any ResponseBodyWriter) async throws -> HTTPFields?
@@ -71,6 +67,18 @@ public struct ResponseBody: Sendable {
         self.init(contentLength: contentLength, write: write)
     }
 
+    /// Returns a ResponseBody containing the results of mapping the given closure over the sequence of
+    /// ByteBuffers written.
+    /// - Parameter transform: A mapping closure applied to every ByteBuffer in ResponseBody
+    /// - Returns: The transformed ResponseBody
+    public consuming func map(_ transform: @escaping @Sendable (ByteBuffer) async throws -> ByteBuffer) -> ResponseBody {
+        let body = self
+        return Self.withTrailingHeaders { writer in
+            let tailHeaders = try await body.write(writer.map(transform))
+            return tailHeaders
+        }
+    }
+
     /// Create new response body that call a callback once original response body has been written
     /// to the channel
     ///

Sources/HummingbirdCore/Response/ResponseBodyWriter.swift

@@ -0,0 +1,62 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Hummingbird server framework project
+//
+// Copyright (c) 2024 the Hummingbird authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See hummingbird/CONTRIBUTORS.txt for the list of Hummingbird authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import NIOCore
+
+/// HTTP Response Body part writer
+public protocol ResponseBodyWriter {
+    /// Write a single ByteBuffer
+    /// - Parameter buffer: single buffer to write
+    func write(_ buffer: ByteBuffer) async throws
+    /// Write a sequence of ByteBuffers
+    /// - Parameter buffers: Sequence of buffers
+    func write(contentsOf buffers: some Sequence<ByteBuffer>) async throws
+}
+
+extension ResponseBodyWriter {
+    /// Default implementation of writing a sequence of ByteBuffers
+    @inlinable
+    public func write(contentsOf buffers: some Sequence<ByteBuffer>) async throws {
+        for part in buffers {
+            try await self.write(part)
+        }
+    }
+}
+
+struct MappedResponseBodyWriter<ParentWriter: ResponseBodyWriter>: ResponseBodyWriter {
+    fileprivate let parentWriter: ParentWriter
+    fileprivate let transform: @Sendable (ByteBuffer) async throws -> ByteBuffer
+
+    /// Write a single ByteBuffer
+    /// - Parameter buffer: single buffer to write
+    func write(_ buffer: ByteBuffer) async throws {
+        try await self.parentWriter.write(self.transform(buffer))
+    }
+
+    /// Write a sequence of ByteBuffers
+    /// - Parameter buffers: Sequence of buffers
+    func write(contentsOf buffers: some Sequence<ByteBuffer>) async throws {
+        for part in buffers {
+            try await self.parentWriter.write(self.transform(part))
+        }
+    }
+}
+
+extension ResponseBodyWriter {
+    /// Return ResponseBodyWriter that applies transform to all ByteBuffers written to it
+    /// ResponseBodyWriter.
+    public consuming func map(_ transform: @escaping @Sendable (ByteBuffer) async throws -> ByteBuffer) -> some ResponseBodyWriter {
+        MappedResponseBodyWriter(parentWriter: self, transform: transform)
+    }
+}

Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Hummingbird server framework project
 //
-// Copyright (c) 2023 the Hummingbird authors
+// Copyright (c) 2023-2024 the Hummingbird authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -93,26 +93,23 @@ extension HTTPChannelHandler {
     }
 }
 
-/// Writes ByteBuffers to AsyncChannel outbound writer
+/// ResponseBodyWriter that writes ByteBuffers to AsyncChannel outbound writer
 struct HTTPServerBodyWriter: Sendable, ResponseBodyWriter {
     typealias Out = HTTPResponsePart
     /// The components of a HTTP response from the view of a HTTP server.
     public typealias OutboundWriter = NIOAsyncChannelOutboundWriter<Out>
 
     let outbound: OutboundWriter
 
+    /// Write a single ByteBuffer
+    /// - Parameter buffer: single buffer to write
     func write(_ buffer: ByteBuffer) async throws {
         try await self.outbound.write(.body(buffer))
     }
-}
 
-extension NIOLockedValueBox {
-    /// Exchange stored value for new value and return the old stored value
-    func exchange(_ newValue: Value) -> Value {
-        self.withLockedValue { value in
-            let prevValue = value
-            value = newValue
-            return prevValue
-        }
+    /// Write a sequence of ByteBuffers
+    /// - Parameter buffers: Sequence of buffers
+    func write(contentsOf buffers: some Sequence<ByteBuffer>) async throws {
+        try await self.outbound.write(contentsOf: buffers.map { .body($0) })
     }
 }

Tests/HummingbirdTests/MiddlewareTests.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Hummingbird server framework project
 //
-// Copyright (c) 2021-2023 the Hummingbird authors
+// Copyright (c) 2021-2024 the Hummingbird authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -181,6 +181,42 @@ final class MiddlewareTests: XCTestCase {
         }
     }
 
+    func testMappedResponseBodyWriter() async throws {
+        struct TransformWriter: ResponseBodyWriter {
+            let parentWriter: any ResponseBodyWriter
+
+            func write(_ buffer: ByteBuffer) async throws {
+                let output = ByteBuffer(bytes: buffer.readableBytesView.map { $0 ^ 255 })
+                try await self.parentWriter.write(output)
+            }
+        }
+        struct TransformMiddleware<Context: RequestContext>: RouterMiddleware {
+            public func handle(_ request: Request, context: Context, next: (Request, Context) async throws -> Response) async throws -> Response {
+                let response = try await next(request, context)
+                var editedResponse = response
+                editedResponse.body = editedResponse.body.map {
+                    return ByteBuffer(bytes: $0.readableBytesView.map { $0 ^ 255 })
+                }
+                return editedResponse
+            }
+        }
+        let router = Router()
+        router.group()
+            .add(middleware: TransformMiddleware())
+            .get("test") { request, _ in
+                return Response(status: .ok, body: .init(asyncSequence: request.body))
+            }
+        let app = Application(responder: router.buildResponder())
+
+        try await app.test(.router) { client in
+            let buffer = Self.randomBuffer(size: 64000)
+            try await client.execute(uri: "/test", method: .get, body: buffer) { response in
+                let expectedOutput = ByteBuffer(bytes: buffer.readableBytesView.map { $0 ^ 255 })
+                XCTAssertEqual(expectedOutput, response.body)
+            }
+        }
+    }
+
     func testCORSUseOrigin() async throws {
         let router = Router()
         router.add(middleware: CORSMiddleware())