Skip to content

Commit 9937d87

Browse files
dnadobadnadoba
authored
Handle ResponseAccumulator not being able to buffer large response in memory (swift-server#637)
* Handle ResponseAccumulator not being able to buffer large response in memory Check content length header for early exit * Add test which currently hangs indefinitely * Run `generate_linux_tests.rb` and `SwiftFormat` * Print type and value if assert fails * Run `generate_linux_tests.rb` and `SwiftFormat` * Remove duplicate test due too merge conflict * Validate that maxBodySize is positive * Address review comments
1 parent f7a84af commit 9937d87

File tree

4 files changed

+183
-4
lines changed

4 files changed

+183
-4
lines changed

‎Sources/AsyncHTTPClient/HTTPHandler.swift

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,7 @@ extension HTTPClient {
356356
///
357357
/// This ``HTTPClientResponseDelegate`` buffers a complete HTTP response in memory. It does not stream the response body in.
358358
/// The resulting ``Response`` type is ``HTTPClient/Response``.
359-
public class ResponseAccumulator: HTTPClientResponseDelegate {
359+
public final class ResponseAccumulator: HTTPClientResponseDelegate {
360360
public typealias Response = HTTPClient.Response
361361

362362
enum State {
@@ -367,16 +367,63 @@ public class ResponseAccumulator: HTTPClientResponseDelegate {
367367
case error(Error)
368368
}
369369

370+
public struct ResponseTooBigError: Error, CustomStringConvertible {
371+
public var maxBodySize: Int
372+
public init(maxBodySize: Int) {
373+
self.maxBodySize = maxBodySize
374+
}
375+
376+
public var description: String {
377+
return "ResponseTooBigError: received response body exceeds maximum accepted size of \(self.maxBodySize) bytes"
378+
}
379+
}
380+
370381
var state = State.idle
371382
let request: HTTPClient.Request
372383

373-
public init(request: HTTPClient.Request) {
384+
static let maxByteBufferSize = Int(UInt32.max)
385+
386+
/// Maximum size in bytes of the HTTP response body that ``ResponseAccumulator`` will accept
387+
/// until it will abort the request and throw an ``ResponseTooBigError``.
388+
///
389+
/// Default is 2^32.
390+
/// - precondition: not allowed to exceed 2^32 because ``ByteBuffer`` can not store more bytes
391+
public let maxBodySize: Int
392+
393+
public convenience init(request: HTTPClient.Request) {
394+
self.init(request: request, maxBodySize: Self.maxByteBufferSize)
395+
}
396+
397+
/// - Parameters:
398+
/// - request: The corresponding request of the response this delegate will be accumulating.
399+
/// - maxBodySize: Maximum size in bytes of the HTTP response body that ``ResponseAccumulator`` will accept
400+
/// until it will abort the request and throw an ``ResponseTooBigError``.
401+
/// Default is 2^32.
402+
/// - precondition: maxBodySize is not allowed to exceed 2^32 because ``ByteBuffer`` can not store more bytes
403+
/// - warning: You can use ``ResponseAccumulator`` for just one request.
404+
/// If you start another request, you need to initiate another ``ResponseAccumulator``.
405+
public init(request: HTTPClient.Request, maxBodySize: Int) {
406+
precondition(maxBodySize >= 0, "maxBodyLength is not allowed to be negative")
407+
precondition(
408+
maxBodySize <= Self.maxByteBufferSize,
409+
"maxBodyLength is not allowed to exceed 2^32 because ByteBuffer can not store more bytes"
410+
)
374411
self.request = request
412+
self.maxBodySize = maxBodySize
375413
}
376414

377415
public func didReceiveHead(task: HTTPClient.Task<Response>, _ head: HTTPResponseHead) -> EventLoopFuture<Void> {
378416
switch self.state {
379417
case .idle:
418+
if self.request.method != .HEAD,
419+
let contentLength = head.headers.first(name: "Content-Length"),
420+
let announcedBodySize = Int(contentLength),
421+
announcedBodySize > self.maxBodySize {
422+
let error = ResponseTooBigError(maxBodySize: maxBodySize)
423+
self.state = .error(error)
424+
return task.eventLoop.makeFailedFuture(error)
425+
}
426+
380427
self.state = .head(head)
381428
case .head:
382429
preconditionFailure("head already set")
@@ -395,8 +442,20 @@ public class ResponseAccumulator: HTTPClientResponseDelegate {
395442
case .idle:
396443
preconditionFailure("no head received before body")
397444
case .head(let head):
445+
guard part.readableBytes <= self.maxBodySize else {
446+
let error = ResponseTooBigError(maxBodySize: self.maxBodySize)
447+
self.state = .error(error)
448+
return task.eventLoop.makeFailedFuture(error)
449+
}
398450
self.state = .body(head, part)
399451
case .body(let head, var body):
452+
let newBufferSize = body.writerIndex + part.readableBytes
453+
guard newBufferSize <= self.maxBodySize else {
454+
let error = ResponseTooBigError(maxBodySize: self.maxBodySize)
455+
self.state = .error(error)
456+
return task.eventLoop.makeFailedFuture(error)
457+
}
458+
400459
// The compiler can't prove that `self.state` is dead here (and it kinda isn't, there's
401460
// a cross-module call in the way) so we need to drop the original reference to `body` in
402461
// `self.state` or we'll get a CoW. To fix that we temporarily set the state to `.end` (which

‎Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,18 @@ internal final class HTTPBin<RequestHandler: ChannelInboundHandler> where
366366
return self.serverChannel.localAddress!
367367
}
368368

369+
var baseURL: String {
370+
let scheme: String = {
371+
switch mode {
372+
case .http1_1, .refuse:
373+
return "http"
374+
case .http2:
375+
return "https"
376+
}
377+
}()
378+
return "\(scheme)://localhost:\(self.port)/"
379+
}
380+
369381
private let mode: Mode
370382
private let sslContext: NIOSSLContext?
371383
private var serverChannel: Channel!
@@ -1319,6 +1331,25 @@ class HTTPEchoHandler: ChannelInboundHandler {
13191331
}
13201332
}
13211333

1334+
final class HTTPEchoHeaders: ChannelInboundHandler {
1335+
typealias InboundIn = HTTPServerRequestPart
1336+
typealias OutboundOut = HTTPServerResponsePart
1337+
1338+
func channelRead(context: ChannelHandlerContext, data: NIOAny) {
1339+
let request = self.unwrapInboundIn(data)
1340+
switch request {
1341+
case .head(let requestHead):
1342+
context.writeAndFlush(self.wrapOutboundOut(.head(.init(version: .http1_1, status: .ok, headers: requestHead.headers))), promise: nil)
1343+
case .body:
1344+
break
1345+
case .end:
1346+
context.writeAndFlush(self.wrapOutboundOut(.end(nil))).whenSuccess {
1347+
context.close(promise: nil)
1348+
}
1349+
}
1350+
}
1351+
}
1352+
13221353
final class HTTP200DelayedHandler: ChannelInboundHandler {
13231354
typealias InboundIn = HTTPServerRequestPart
13241355
typealias OutboundOut = HTTPServerResponsePart

‎Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,11 @@ extension HTTPClientTests {
132132
("testSSLHandshakeErrorPropagationDelayedClose", testSSLHandshakeErrorPropagationDelayedClose),
133133
("testWeCloseConnectionsWhenConnectionCloseSetByServer", testWeCloseConnectionsWhenConnectionCloseSetByServer),
134134
("testBiDirectionalStreaming", testBiDirectionalStreaming),
135+
("testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength),
136+
("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength),
137+
("testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead", testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead),
138+
("testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked),
139+
("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked),
135140
("testBiDirectionalStreamingEarly200", testBiDirectionalStreamingEarly200),
136141
("testBiDirectionalStreamingEarly200DoesntPreventUsFromSendingMoreRequests", testBiDirectionalStreamingEarly200DoesntPreventUsFromSendingMoreRequests),
137142
("testCloseConnectionAfterEarly2XXWhenStreaming", testCloseConnectionAfterEarly2XXWhenStreaming),

‎Tests/AsyncHTTPClientTests/HTTPClientTests.swift

Lines changed: 86 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2677,8 +2677,8 @@ class HTTPClientTests: XCTestCase {
26772677
let delegate = TestDelegate()
26782678

26792679
XCTAssertThrowsError(try httpClient.execute(request: request, delegate: delegate).wait()) {
2680-
XCTAssertEqual(.connectTimeout, $0 as? HTTPClientError)
2681-
XCTAssertEqual(.connectTimeout, delegate.error as? HTTPClientError)
2680+
XCTAssertEqualTypeAndValue($0, HTTPClientError.connectTimeout)
2681+
XCTAssertEqualTypeAndValue(delegate.error, HTTPClientError.connectTimeout)
26822682
}
26832683
}
26842684

@@ -3092,6 +3092,90 @@ class HTTPClientTests: XCTestCase {
30923092
XCTAssertNil(try delegate.next().wait())
30933093
}
30943094

3095+
func testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength() throws {
3096+
let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
3097+
defer { XCTAssertNoThrow(try httpBin.shutdown()) }
3098+
3099+
let body = ByteBuffer(bytes: 0..<11)
3100+
3101+
var request = try Request(url: httpBin.baseURL)
3102+
request.body = .byteBuffer(body)
3103+
XCTAssertThrowsError(try self.defaultClient.execute(
3104+
request: request,
3105+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
3106+
).wait()) { error in
3107+
XCTAssertTrue(error is ResponseAccumulator.ResponseTooBigError, "unexpected error \(error)")
3108+
}
3109+
}
3110+
3111+
func testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength() throws {
3112+
let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
3113+
defer { XCTAssertNoThrow(try httpBin.shutdown()) }
3114+
3115+
let body = ByteBuffer(bytes: 0..<10)
3116+
3117+
var request = try Request(url: httpBin.baseURL)
3118+
request.body = .byteBuffer(body)
3119+
let response = try self.defaultClient.execute(
3120+
request: request,
3121+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
3122+
).wait()
3123+
3124+
XCTAssertEqual(response.body, body)
3125+
}
3126+
3127+
func testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead() throws {
3128+
let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHeaders() }
3129+
defer { XCTAssertNoThrow(try httpBin.shutdown()) }
3130+
3131+
let body = ByteBuffer(bytes: 0..<11)
3132+
3133+
var request = try Request(url: httpBin.baseURL, method: .HEAD)
3134+
request.body = .byteBuffer(body)
3135+
let response = try self.defaultClient.execute(
3136+
request: request,
3137+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
3138+
).wait()
3139+
3140+
XCTAssertEqual(response.body ?? ByteBuffer(), ByteBuffer())
3141+
}
3142+
3143+
func testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked() throws {
3144+
let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
3145+
defer { XCTAssertNoThrow(try httpBin.shutdown()) }
3146+
3147+
let body = ByteBuffer(bytes: 0..<11)
3148+
3149+
var request = try Request(url: httpBin.baseURL)
3150+
request.body = .stream { writer in
3151+
writer.write(.byteBuffer(body))
3152+
}
3153+
XCTAssertThrowsError(try self.defaultClient.execute(
3154+
request: request,
3155+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
3156+
).wait()) { error in
3157+
XCTAssertTrue(error is ResponseAccumulator.ResponseTooBigError, "unexpected error \(error)")
3158+
}
3159+
}
3160+
3161+
func testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked() throws {
3162+
let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
3163+
defer { XCTAssertNoThrow(try httpBin.shutdown()) }
3164+
3165+
let body = ByteBuffer(bytes: 0..<10)
3166+
3167+
var request = try Request(url: httpBin.baseURL)
3168+
request.body = .stream { writer in
3169+
writer.write(.byteBuffer(body))
3170+
}
3171+
let response = try self.defaultClient.execute(
3172+
request: request,
3173+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
3174+
).wait()
3175+
3176+
XCTAssertEqual(response.body, body)
3177+
}
3178+
30953179
// In this test, we test that a request can continue to stream its body after the response head and end
30963180
// was received where the end is a 200.
30973181
func testBiDirectionalStreamingEarly200() {

0 commit comments

Comments
 (0)