Collect function fix

Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse.swift

@@ -32,19 +32,8 @@ public struct HTTPClientResponse: Sendable {
     /// The body of this HTTP response.
     public var body: Body
 
-    init(
-        bag: Transaction,
-        version: HTTPVersion,
-        status: HTTPResponseStatus,
-        headers: HTTPHeaders
-    ) {
-        self.version = version
-        self.status = status
-        self.headers = headers
-        self.body = Body(TransactionBody(bag))
-    }
 
-    @inlinable public init(
+    @inlinable init(
         version: HTTPVersion = .http1_1,
         status: HTTPResponseStatus = .ok,
         headers: HTTPHeaders = [:],
@@ -55,6 +44,17 @@ public struct HTTPClientResponse: Sendable {
         self.headers = headers
         self.body = body
     }
+
+    init(
+        bag: Transaction,
+        version: HTTPVersion,
+        status: HTTPResponseStatus,
+        headers: HTTPHeaders,
+        requestMethod: HTTPMethod
+    ) {
+        let contentLength = HTTPClientResponse.expectedContentLength(requestMethod: requestMethod, headers: headers, status: status)
+        self.init(version: version, status: status, headers: headers, body: .init(TransactionBody(bag, expextedContentLength: contentLength)))
+    }
 }
 
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
@@ -83,6 +83,50 @@ extension HTTPClientResponse {
         @inlinable public func makeAsyncIterator() -> AsyncIterator {
             .init(storage: self.storage.makeAsyncIterator())
         }
+
+        @inlinable init(storage: Storage) {
+            self.storage = storage
+        }
+
+        /// Accumulates `Body` of ``ByteBuffer``s into a single ``ByteBuffer``.
+        /// - Parameters:
+        ///   - maxBytes: The maximum number of bytes this method is allowed to accumulate
+        /// - Throws: `NIOTooManyBytesError` if the the sequence contains more than `maxBytes`.
+        /// - Returns: the number of bytes collected over time
+        @inlinable public func collect(upTo maxBytes: Int) async throws -> ByteBuffer {
+            switch storage {
+            case .transaction(let transactionBody):
+                if let contentLength = transactionBody.expectedContentLength {
+                    if contentLength > maxBytes {
+                        throw NIOTooManyBytesError()
+                    }
+                }
+            case .anyAsyncSequence:
+                break
+            }
+
+            /// calling collect function within here in order to ensure the correct nested type
+            func collect<Body: AsyncSequence>(_ body: Body, maxBytes: Int) async throws -> ByteBuffer where Body.Element == ByteBuffer {
+                try await body.collect(upTo: maxBytes)
+            }
+            return try await collect(self, maxBytes: maxBytes)
+
+        }
+
+    }
+}
+
+@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+extension HTTPClientResponse {
+    static func expectedContentLength(requestMethod: HTTPMethod, headers: HTTPHeaders, status: HTTPResponseStatus) -> Int? {
+        if status == .notModified {
+            return 0
+        } else if requestMethod == .HEAD {
+            return 0
+        } else {
+            let contentLength = headers["content-length"].first.flatMap({Int($0, radix: 10)})
+            return contentLength
+        }
     }
 }
 
@@ -132,10 +176,10 @@ extension HTTPClientResponse.Body.Storage.AsyncIterator: AsyncIteratorProtocol {
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
 extension HTTPClientResponse.Body {
     init(_ body: TransactionBody) {
-        self.init(.transaction(body))
+        self.init(storage: .transaction(body))
     }
 
-    @usableFromInline init(_ storage: Storage) {
+    @inlinable init(_ storage: Storage) {
         self.storage = storage
     }
 
@@ -146,7 +190,7 @@ extension HTTPClientResponse.Body {
     @inlinable public static func stream<SequenceOfBytes>(
         _ sequenceOfBytes: SequenceOfBytes
     ) -> Self where SequenceOfBytes: AsyncSequence & Sendable, SequenceOfBytes.Element == ByteBuffer {
-        self.init(.anyAsyncSequence(AnyAsyncSequence(sequenceOfBytes.singleIteratorPrecondition)))
+        Self.init(storage: .anyAsyncSequence(AnyAsyncSequence(sequenceOfBytes.singleIteratorPrecondition)))
     }
 
     public static func bytes(_ byteBuffer: ByteBuffer) -> Self {

Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift

@@ -236,7 +236,8 @@ extension Transaction: HTTPExecutableRequest {
                 bag: self,
                 version: head.version,
                 status: head.status,
-                headers: head.headers
+                headers: head.headers,
+                requestMethod: requestHead.method
             )
             continuation.resume(returning: asyncResponse)
         }

Sources/AsyncHTTPClient/AsyncAwait/TransactionBody.swift

@@ -20,9 +20,11 @@ import NIOCore
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
 @usableFromInline final class TransactionBody: Sendable {
     @usableFromInline let transaction: Transaction
+    @usableFromInline let expectedContentLength: Int?
 
-    init(_ transaction: Transaction) {
+    init(_ transaction: Transaction, expextedContentLength: Int?) {
         self.transaction = transaction
+        self.expectedContentLength = expextedContentLength
     }
 
     deinit {

Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift

@@ -114,7 +114,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response.headers["content-length"], ["4"])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response.body.collect()
+                try await response.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
         }
@@ -137,7 +137,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response.headers["content-length"], [])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response.body.collect()
+                try await response.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
         }
@@ -160,7 +160,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response.headers["content-length"], [])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response.body.collect()
+                try await response.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
         }
@@ -183,7 +183,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response.headers["content-length"], ["4"])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response.body.collect()
+                try await response.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
         }
@@ -210,7 +210,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response.headers["content-length"], [])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response.body.collect()
+                try await response.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
         }
@@ -233,7 +233,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response.headers["content-length"], [])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response.body.collect()
+                try await response.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
         }
@@ -522,7 +522,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else {
                 return
             }
-            guard let body = await XCTAssertNoThrowWithResult(try await response.body.collect()) else { return }
+            guard let body = await XCTAssertNoThrowWithResult(try await response.body.collect(upTo: 1024)) else { return }
             var maybeRequestInfo: RequestInfo?
             XCTAssertNoThrow(maybeRequestInfo = try JSONDecoder().decode(RequestInfo.self, from: body))
             guard let requestInfo = maybeRequestInfo else { return }
@@ -583,7 +583,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response1.headers["content-length"], [])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response1.body.collect()
+                try await response1.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
 
@@ -592,12 +592,11 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             ) else { return }
             XCTAssertEqual(response2.headers["content-length"], [])
             guard let body = await XCTAssertNoThrowWithResult(
-                try await response2.body.collect()
+                try await response2.body.collect(upTo: 1024)
             ) else { return }
             XCTAssertEqual(body, ByteBuffer(string: "1234"))
         }
     }
-
     func testRejectsInvalidCharactersInHeaderFieldNames_http1() {
         self._rejectsInvalidCharactersInHeaderFieldNames(mode: .http1_1(ssl: true))
     }
@@ -745,17 +744,32 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             XCTAssertEqual(response.version, .http2)
         }
     }
-}
-
-extension AsyncSequence where Element == ByteBuffer {
-    func collect() async rethrows -> ByteBuffer {
-        try await self.reduce(into: ByteBuffer()) { accumulatingBuffer, nextBuffer in
-            var nextBuffer = nextBuffer
-            accumulatingBuffer.writeBuffer(&nextBuffer)
+
+    func testSimpleContentLengthError() {
+        guard #available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) else { return }
+        XCTAsyncTest {
+            let bin = HTTPBin(.http2(compress: false)) { _ in HTTPEchoHandler() }
+            defer { XCTAssertNoThrow(try bin.shutdown()) }
+            let client = makeDefaultHTTPClient()
+            defer { XCTAssertNoThrow(try client.syncShutdown()) }
+            let logger = Logger(label: "HTTPClient", factory: StreamLogHandler.standardOutput(label:))
+            var request = HTTPClientRequest(url: "https://localhost:\(bin.port)/")
+            request.method = .GET
+            request.body = .bytes(ByteBuffer(string: "1234"))
+
+            guard var response = await XCTAssertNoThrowWithResult(
+                try await client.execute(request, deadline: .now() + .seconds(10), logger: logger)
+            ) else { return }
+            await XCTAssertThrowsError(
+                try await response.body.collect(upTo: 3)
+            ) {
+                XCTAssertEqualTypeAndValue($0, NIOTooManyBytesError())
+            }
         }
     }
 }
 
+
 struct AnySendableSequence<Element>: @unchecked Sendable {
     private let wrapped: AnySequence<Element>
     init<WrappedSequence: Sequence & Sendable>(

Tests/AsyncHTTPClientTests/HTTPClientResponseTests.swift

@@ -0,0 +1,39 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2023 Apple Inc. and the AsyncHTTPClient project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+
+
+@testable import AsyncHTTPClient
+import Logging
+import NIOCore
+import XCTest
+
+
+final class HTTPClientResponseTests: XCTestCase {
+
+    func testSimpleResponse() {
+        let response = HTTPClientResponse.expectedContentLength(requestMethod: .GET, headers: ["content-length": "1025"], status: .ok)
+        XCTAssertEqual(response, 1025)
+    }
+
+    func testSimpleResponseNotModified() {
+        let response = HTTPClientResponse.expectedContentLength(requestMethod: .GET, headers: ["content-length": "1025"], status: .notModified)
+        XCTAssertEqual(response, 0)
+    }
+
+    func testSimpleResponseHeadRequestMethod() {
+        let response = HTTPClientResponse.expectedContentLength(requestMethod: .HEAD, headers: ["content-length": "1025"], status: .ok)
+        XCTAssertEqual(response, 0)
+    }
+}