Add support for creating stream from byteBuffer, only allow one iterate of HBStreamedRequestBody

Author: adam-fowler

Sources/HummingbirdCore/Request/RequestBody.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import NIOConcurrencyHelpers
 import NIOCore
 import NIOHTTPTypes
 
@@ -27,10 +28,11 @@ public enum HBRequestBody: Sendable, AsyncSequence {
 
     public func makeAsyncIterator() -> HBStreamedRequestBody.AsyncIterator {
         switch self {
-        case .byteBuffer:
-            /// The server always creates the HBRequestBody as a stream. If it is converted
-            /// into a single ByteBuffer it cannot be treated as a stream after that
-            preconditionFailure("Cannot convert collapsed request body back into a sequence")
+        case .byteBuffer(let buffer):
+            let (stream, source) = NIOAsyncChannelInboundStream<HTTPRequestPart>.makeTestingStream()
+            source.yield(.body(buffer))
+            source.finish()
+            return HBStreamedRequestBody(iterator: stream.makeAsyncIterator()).makeAsyncIterator()
         case .stream(let streamer):
             return streamer.makeAsyncIterator()
         }
@@ -49,13 +51,19 @@ public enum HBRequestBody: Sendable, AsyncSequence {
 }
 
 /// Request body that is a stream of ByteBuffers.
-public struct HBStreamedRequestBody: Sendable, AsyncSequence {
+///
+/// This is a unicast async sequence that allows a single iterator to be created.
+public final class HBStreamedRequestBody: Sendable, AsyncSequence {
     public typealias Element = ByteBuffer
     public typealias InboundStream = NIOAsyncChannelInboundStream<HTTPRequestPart>
 
+    private let underlyingIterator: UnsafeTransfer<NIOAsyncChannelInboundStream<HTTPRequestPart>.AsyncIterator>
+    private let alreadyIterated: NIOLockedValueBox<Bool>
+
     /// Initialize HBStreamedRequestBody from AsyncIterator of a NIOAsyncChannelInboundStream
     public init(iterator: InboundStream.AsyncIterator) {
         self.underlyingIterator = .init(iterator)
+        self.alreadyIterated = .init(false)
     }
 
     /// Async Iterator for HBStreamedRequestBody
@@ -65,9 +73,9 @@ public struct HBStreamedRequestBody: Sendable, AsyncSequence {
         private var underlyingIterator: InboundStream.AsyncIterator
         private var done: Bool
 
-        init(underlyingIterator: InboundStream.AsyncIterator) {
+        init(underlyingIterator: InboundStream.AsyncIterator, done: Bool = false) {
             self.underlyingIterator = underlyingIterator
-            self.done = false
+            self.done = done
         }
 
         public mutating func next() async throws -> ByteBuffer? {
@@ -88,8 +96,15 @@ public struct HBStreamedRequestBody: Sendable, AsyncSequence {
     }
 
     public func makeAsyncIterator() -> AsyncIterator {
-        AsyncIterator(underlyingIterator: self.underlyingIterator.wrappedValue)
+        // verify if an iterator has already been created. If it has then create an
+        // iterator that returns nothing. This could be a precondition failure (currently
+        // an assert) as you should not be allowed to do this.
+        let done = self.alreadyIterated.withLockedValue {
+            assert($0 == false, "Can only create iterator from request body once")
+            let done = $0
+            $0 = true
+            return done
+        }
+        return AsyncIterator(underlyingIterator: self.underlyingIterator.wrappedValue, done: done)
     }
-
-    private var underlyingIterator: UnsafeTransfer<NIOAsyncChannelInboundStream<HTTPRequestPart>.AsyncIterator>
 }

Sources/HummingbirdXCT/HBXCTRouter.swift

@@ -112,7 +112,6 @@ struct HBXCTRouter<Responder: HBResponder>: HBXCTApplication where Responder.Con
                     }
                     let responseWriter = RouterResponseWriter()
                     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, trailerHeaders: trailerHeaders)
                     }

Tests/HummingbirdTests/ApplicationTests.swift

@@ -282,6 +282,27 @@ final class ApplicationTests: XCTestCase {
         }
     }
 
+    func testDoubleStreaming() async throws {
+        let router = HBRouter()
+        router.post("size") { request, context -> String in
+            var request = request
+            _ = try await request.collateBody(context: context)
+            var size = 0
+            for try await buffer in request.body {
+                size += buffer.readableBytes
+            }
+            return size.description
+        }
+        let app = HBApplication(responder: router.buildResponder())
+
+        try await app.test(.router) { client in
+            let buffer = self.randomBuffer(size: 100_000)
+            try await client.XCTExecute(uri: "/size", method: .post, body: buffer) { response in
+                XCTAssertEqual(String(buffer: response.body), "100000")
+            }
+        }
+    }
+
     func testOptional() async throws {
         let router = HBRouter()
         router