Handle ResponseAccumulator not being able to buffer large response in memory

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -356,7 +356,7 @@ extension HTTPClient {
 ///
 /// This ``HTTPClientResponseDelegate`` buffers a complete HTTP response in memory. It does not stream the response body in.
 /// The resulting ``Response`` type is ``HTTPClient/Response``.
-public class ResponseAccumulator: HTTPClientResponseDelegate {
+public final class ResponseAccumulator: HTTPClientResponseDelegate {
     public typealias Response = HTTPClient.Response
 
     enum State {
@@ -367,16 +367,48 @@ public class ResponseAccumulator: HTTPClientResponseDelegate {
         case error(Error)
     }
 
+    public struct ResponseTooBigError: Error, CustomStringConvertible {
+        var maxBodySize: Int
+
+        public var description: String {
+            return "ResponseTooBigError: received response body exceeds maximum accepted size of \(self.maxBodySize) bytes"
+        }
+    }
+
     var state = State.idle
     let request: HTTPClient.Request
 
+    static let maxByteBufferSize = Int(UInt32.max)
+
+    /// Maximum size in bytes of the HTTP response body that ``ResponseAccumulator`` will accept
+    /// until it will abort the request and throw an ``ResponseTooBigError``.
+    ///
+    /// Default is 2^32.
+    /// - precondition: not allowed to exceed 2^32 because ``ByteBuffer`` can not store more bytes
+    public var maxBodySize: Int = maxByteBufferSize {
+        didSet {
+            precondition(
+                self.maxBodySize <= Self.maxByteBufferSize,
+                "maxBodyLength is not allowed to exceed 2^32 because ByteBuffer can not store more bytes"
+            )
+        }
+    }
+
     public init(request: HTTPClient.Request) {
         self.request = request
     }
 
     public func didReceiveHead(task: HTTPClient.Task<Response>, _ head: HTTPResponseHead) -> EventLoopFuture<Void> {
         switch self.state {
         case .idle:
+            if let contentLength = head.headers.first(name: "Content-Length"),
+               let announcedBodySize = Int(contentLength),
+               announcedBodySize > self.maxBodySize {
+                let error = ResponseTooBigError(maxBodySize: maxBodySize)
+                self.state = .error(error)
+                return task.eventLoop.makeFailedFuture(error)
+            }
+
             self.state = .head(head)
         case .head:
             preconditionFailure("head already set")
@@ -395,8 +427,20 @@ public class ResponseAccumulator: HTTPClientResponseDelegate {
         case .idle:
             preconditionFailure("no head received before body")
         case .head(let head):
+            guard part.readableBytes <= self.maxBodySize else {
+                let error = ResponseTooBigError(maxBodySize: self.maxBodySize)
+                self.state = .error(error)
+                return task.eventLoop.makeFailedFuture(error)
+            }
             self.state = .body(head, part)
         case .body(let head, var body):
+            let newBufferSize = body.writerIndex + part.readableBytes
+            guard newBufferSize <= self.maxBodySize else {
+                let error = ResponseTooBigError(maxBodySize: self.maxBodySize)
+                self.state = .error(error)
+                return task.eventLoop.makeFailedFuture(error)
+            }
+
             // The compiler can't prove that `self.state` is dead here (and it kinda isn't, there's
             // a cross-module call in the way) so we need to drop the original reference to `body` in
             // `self.state` or we'll get a CoW. To fix that we temporarily set the state to `.end` (which

Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift

@@ -366,6 +366,18 @@ internal final class HTTPBin<RequestHandler: ChannelInboundHandler> where
         return self.serverChannel.localAddress!
     }
 
+    var baseURL: String {
+        let scheme: String = {
+            switch mode {
+            case .http1_1, .refuse:
+                return "http"
+            case .http2:
+                return "https"
+            }
+        }()
+        return "\(scheme)://localhost:\(self.port)/"
+    }
+
     private let mode: Mode
     private let sslContext: NIOSSLContext?
     private var serverChannel: Channel!

Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift

@@ -132,6 +132,10 @@ extension HTTPClientTests {
             ("testSSLHandshakeErrorPropagationDelayedClose", testSSLHandshakeErrorPropagationDelayedClose),
             ("testWeCloseConnectionsWhenConnectionCloseSetByServer", testWeCloseConnectionsWhenConnectionCloseSetByServer),
             ("testBiDirectionalStreaming", testBiDirectionalStreaming),
+            ("testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength),
+            ("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength),
+            ("testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked),
+            ("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked),
             ("testBiDirectionalStreamingEarly200", testBiDirectionalStreamingEarly200),
             ("testBiDirectionalStreamingEarly200DoesntPreventUsFromSendingMoreRequests", testBiDirectionalStreamingEarly200DoesntPreventUsFromSendingMoreRequests),
             ("testCloseConnectionAfterEarly2XXWhenStreaming", testCloseConnectionAfterEarly2XXWhenStreaming),

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -2677,8 +2677,8 @@ class HTTPClientTests: XCTestCase {
         let delegate = TestDelegate()
 
         XCTAssertThrowsError(try httpClient.execute(request: request, delegate: delegate).wait()) {
-            XCTAssertEqual(.connectTimeout, $0 as? HTTPClientError)
-            XCTAssertEqual(.connectTimeout, delegate.error as? HTTPClientError)
+            XCTAssertEqualTypeAndValue($0, HTTPClientError.connectTimeout)
+            XCTAssertEqualTypeAndValue(delegate.error, HTTPClientError.connectTimeout)
         }
     }
 
@@ -3092,6 +3092,82 @@ class HTTPClientTests: XCTestCase {
         XCTAssertNil(try delegate.next().wait())
     }
 
+    func testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength() throws {
+        let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
+        defer { XCTAssertNoThrow(try httpBin.shutdown()) }
+
+        let body = ByteBuffer(bytes: 0..<11)
+
+        var request = try Request(url: httpBin.baseURL)
+        request.body = .byteBuffer(body)
+        let delegate = ResponseAccumulator(request: request)
+        delegate.maxBodySize = 10
+        XCTAssertThrowsError(try self.defaultClient.execute(
+            request: request,
+            delegate: delegate
+        ).wait()) { error in
+            XCTAssertTrue(error is ResponseAccumulator.ResponseTooBigError, "unexpected error \(error)")
+        }
+    }
+
+    func testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength() throws {
+        let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
+        defer { XCTAssertNoThrow(try httpBin.shutdown()) }
+
+        let body = ByteBuffer(bytes: 0..<10)
+
+        var request = try Request(url: httpBin.baseURL)
+        request.body = .byteBuffer(body)
+        let delegate = ResponseAccumulator(request: request)
+        delegate.maxBodySize = 10
+        let response = try self.defaultClient.execute(
+            request: request,
+            delegate: delegate
+        ).wait()
+
+        XCTAssertEqual(response.body, body)
+    }
+
+    func testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked() throws {
+        let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
+        defer { XCTAssertNoThrow(try httpBin.shutdown()) }
+
+        let body = ByteBuffer(bytes: 0..<11)
+
+        var request = try Request(url: httpBin.baseURL)
+        request.body = .stream { writer in
+            writer.write(.byteBuffer(body))
+        }
+        let delegate = ResponseAccumulator(request: request)
+        delegate.maxBodySize = 10
+        XCTAssertThrowsError(try self.defaultClient.execute(
+            request: request,
+            delegate: delegate
+        ).wait()) { error in
+            XCTAssertTrue(error is ResponseAccumulator.ResponseTooBigError, "unexpected error \(error)")
+        }
+    }
+
+    func testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked() throws {
+        let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
+        defer { XCTAssertNoThrow(try httpBin.shutdown()) }
+
+        let body = ByteBuffer(bytes: 0..<10)
+
+        var request = try Request(url: httpBin.baseURL)
+        request.body = .stream { writer in
+            writer.write(.byteBuffer(body))
+        }
+        let delegate = ResponseAccumulator(request: request)
+        delegate.maxBodySize = 10
+        let response = try self.defaultClient.execute(
+            request: request,
+            delegate: delegate
+        ).wait()
+
+        XCTAssertEqual(response.body, body)
+    }
+
     // In this test, we test that a request can continue to stream its body after the response head and end
     // was received where the end is a 200.
     func testBiDirectionalStreamingEarly200() {

Tests/AsyncHTTPClientTests/RequestBagTests+XCTest.swift

@@ -35,6 +35,7 @@ extension RequestBagTests {
             ("testFailsTaskWhenTaskIsWaitingForMoreFromServer", testFailsTaskWhenTaskIsWaitingForMoreFromServer),
             ("testChannelBecomingWritableDoesntCrashCancelledTask", testChannelBecomingWritableDoesntCrashCancelledTask),
             ("testDidReceiveBodyPartFailedPromise", testDidReceiveBodyPartFailedPromise),
+            ("testDidReceiveBodyPartFailedPromise", testDidReceiveBodyPartFailedPromise),
             ("testHTTPUploadIsCancelledEvenThoughRequestSucceeds", testHTTPUploadIsCancelledEvenThoughRequestSucceeds),
             ("testRaceBetweenConnectionCloseAndDemandMoreData", testRaceBetweenConnectionCloseAndDemandMoreData),
             ("testRedirectWith3KBBody", testRedirectWith3KBBody),

Tests/AsyncHTTPClientTests/RequestBagTests.swift

@@ -522,6 +522,73 @@ final class RequestBagTests: XCTestCase {
         }
     }
 
+    func testDidReceiveBodyPartFailedPromise() {
+        let embeddedEventLoop = EmbeddedEventLoop()
+        defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }
+        let logger = Logger(label: "test")
+
+        var maybeRequest: HTTPClient.Request?
+
+        XCTAssertNoThrow(maybeRequest = try HTTPClient.Request(
+            url: "https://swift.org",
+            method: .POST,
+            body: .byteBuffer(.init(bytes: [1]))
+        ))
+        guard let request = maybeRequest else { return XCTFail("Expected to have a request") }
+
+        struct MyError: Error, Equatable {}
+        final class Delegate: HTTPClientResponseDelegate {
+            typealias Response = Void
+            let didFinishPromise: EventLoopPromise<Void>
+            init(didFinishPromise: EventLoopPromise<Void>) {
+                self.didFinishPromise = didFinishPromise
+            }
+
+            func didReceiveBodyPart(task: HTTPClient.Task<Void>, _ buffer: ByteBuffer) -> EventLoopFuture<Void> {
+                task.eventLoop.makeFailedFuture(MyError())
+            }
+
+            func didReceiveError(task: HTTPClient.Task<Void>, _ error: Error) {
+                self.didFinishPromise.fail(error)
+            }
+
+            func didFinishRequest(task: AsyncHTTPClient.HTTPClient.Task<Void>) throws {
+                XCTFail("\(#function) should not be called")
+                self.didFinishPromise.succeed(())
+            }
+        }
+        let delegate = Delegate(didFinishPromise: embeddedEventLoop.makePromise())
+        var maybeRequestBag: RequestBag<Delegate>?
+        XCTAssertNoThrow(maybeRequestBag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embeddedEventLoop),
+            task: .init(eventLoop: embeddedEventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: delegate
+        ))
+        guard let bag = maybeRequestBag else { return XCTFail("Expected to be able to create a request bag.") }
+
+        let executor = MockRequestExecutor(eventLoop: embeddedEventLoop)
+
+        executor.runRequest(bag)
+
+        bag.resumeRequestBodyStream()
+        XCTAssertNoThrow(try executor.receiveRequestBody { XCTAssertEqual($0, ByteBuffer(bytes: [1])) })
+
+        bag.receiveResponseHead(.init(version: .http1_1, status: .ok))
+
+        bag.succeedRequest([ByteBuffer([1])])
+
+        XCTAssertThrowsError(try delegate.didFinishPromise.futureResult.wait()) { error in
+            XCTAssertEqualTypeAndValue(error, MyError())
+        }
+        XCTAssertThrowsError(try bag.task.futureResult.wait()) { error in
+            XCTAssertEqualTypeAndValue(error, MyError())
+        }
+    }
+
     func testHTTPUploadIsCancelledEvenThoughRequestSucceeds() {
         let embeddedEventLoop = EmbeddedEventLoop()
         defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }