Handle ResponseAccumulator not being able to buffer large response in memory

Check content length header for early exit

Author: dnadoba

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 {
+final public class ResponseAccumulator: HTTPClientResponseDelegate {
     public typealias Response = HTTPClient.Response
 
     enum State {
@@ -366,9 +366,33 @@ public class ResponseAccumulator: HTTPClientResponseDelegate {
         case end
         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 \(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(
+                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
@@ -377,6 +401,14 @@ public class ResponseAccumulator: HTTPClientResponseDelegate {
     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

@@ -365,6 +365,19 @@ internal final class HTTPBin<RequestHandler: ChannelInboundHandler> where
     var socketAddress: SocketAddress {
         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?

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -3091,6 +3091,82 @@ class HTTPClientTests: XCTestCase {
         XCTAssertNoThrow(try future.wait())
         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.