diff --git a/Sources/AsyncHTTPClient/HTTPHandler.swift b/Sources/AsyncHTTPClient/HTTPHandler.swift index 744139f68..26880ef8b 100644 --- a/Sources/AsyncHTTPClient/HTTPHandler.swift +++ b/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,63 @@ public class ResponseAccumulator: HTTPClientResponseDelegate { case error(Error) } + public struct ResponseTooBigError: Error, CustomStringConvertible { + public var maxBodySize: Int + public init(maxBodySize: Int) { + self.maxBodySize = maxBodySize + } + + 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 - public init(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 let maxBodySize: Int + + public convenience init(request: HTTPClient.Request) { + self.init(request: request, maxBodySize: Self.maxByteBufferSize) + } + + /// - Parameters: + /// - request: The corresponding request of the response this delegate will be accumulating. + /// - maxBodySize: 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: maxBodySize is not allowed to exceed 2^32 because ``ByteBuffer`` can not store more bytes + /// - warning: You can use ``ResponseAccumulator`` for just one request. + /// If you start another request, you need to initiate another ``ResponseAccumulator``. + public init(request: HTTPClient.Request, maxBodySize: Int) { + precondition(maxBodySize >= 0, "maxBodyLength is not allowed to be negative") + precondition( + maxBodySize <= Self.maxByteBufferSize, + "maxBodyLength is not allowed to exceed 2^32 because ByteBuffer can not store more bytes" + ) self.request = request + self.maxBodySize = maxBodySize } public func didReceiveHead(task: HTTPClient.Task, _ head: HTTPResponseHead) -> EventLoopFuture { switch self.state { case .idle: + if self.request.method != .HEAD, + 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 +442,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 diff --git a/Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift b/Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift index 91358a361..884681123 100644 --- a/Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift +++ b/Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift @@ -366,6 +366,18 @@ internal final class HTTPBin 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! @@ -1319,6 +1331,25 @@ class HTTPEchoHandler: ChannelInboundHandler { } } +final class HTTPEchoHeaders: ChannelInboundHandler { + typealias InboundIn = HTTPServerRequestPart + typealias OutboundOut = HTTPServerResponsePart + + func channelRead(context: ChannelHandlerContext, data: NIOAny) { + let request = self.unwrapInboundIn(data) + switch request { + case .head(let requestHead): + context.writeAndFlush(self.wrapOutboundOut(.head(.init(version: .http1_1, status: .ok, headers: requestHead.headers))), promise: nil) + case .body: + break + case .end: + context.writeAndFlush(self.wrapOutboundOut(.end(nil))).whenSuccess { + context.close(promise: nil) + } + } + } +} + final class HTTP200DelayedHandler: ChannelInboundHandler { typealias InboundIn = HTTPServerRequestPart typealias OutboundOut = HTTPServerResponsePart diff --git a/Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift b/Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift index 2427d6cbf..ef6690c00 100644 --- a/Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift +++ b/Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift @@ -132,6 +132,11 @@ extension HTTPClientTests { ("testSSLHandshakeErrorPropagationDelayedClose", testSSLHandshakeErrorPropagationDelayedClose), ("testWeCloseConnectionsWhenConnectionCloseSetByServer", testWeCloseConnectionsWhenConnectionCloseSetByServer), ("testBiDirectionalStreaming", testBiDirectionalStreaming), + ("testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength), + ("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength), + ("testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead", testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead), + ("testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked), + ("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked), ("testBiDirectionalStreamingEarly200", testBiDirectionalStreamingEarly200), ("testBiDirectionalStreamingEarly200DoesntPreventUsFromSendingMoreRequests", testBiDirectionalStreamingEarly200DoesntPreventUsFromSendingMoreRequests), ("testCloseConnectionAfterEarly2XXWhenStreaming", testCloseConnectionAfterEarly2XXWhenStreaming), diff --git a/Tests/AsyncHTTPClientTests/HTTPClientTests.swift b/Tests/AsyncHTTPClientTests/HTTPClientTests.swift index 3c275c5eb..d5108bb27 100644 --- a/Tests/AsyncHTTPClientTests/HTTPClientTests.swift +++ b/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,90 @@ 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) + XCTAssertThrowsError(try self.defaultClient.execute( + request: request, + delegate: ResponseAccumulator(request: request, maxBodySize: 10) + ).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 response = try self.defaultClient.execute( + request: request, + delegate: ResponseAccumulator(request: request, maxBodySize: 10) + ).wait() + + XCTAssertEqual(response.body, body) + } + + func testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead() throws { + let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHeaders() } + defer { XCTAssertNoThrow(try httpBin.shutdown()) } + + let body = ByteBuffer(bytes: 0..<11) + + var request = try Request(url: httpBin.baseURL, method: .HEAD) + request.body = .byteBuffer(body) + let response = try self.defaultClient.execute( + request: request, + delegate: ResponseAccumulator(request: request, maxBodySize: 10) + ).wait() + + XCTAssertEqual(response.body ?? ByteBuffer(), ByteBuffer()) + } + + 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)) + } + XCTAssertThrowsError(try self.defaultClient.execute( + request: request, + delegate: ResponseAccumulator(request: request, maxBodySize: 10) + ).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 response = try self.defaultClient.execute( + request: request, + delegate: ResponseAccumulator(request: request, maxBodySize: 10) + ).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() {