Skip to content

Commit 986bb05

Browse files
dnadobadnadoba
committed
Address review comments
1 parent 3d8cac8 commit 986bb05

File tree

4 files changed

+65
-23
lines changed

4 files changed

+65
-23
lines changed

‎Sources/AsyncHTTPClient/HTTPHandler.swift‎

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,10 @@ public final class ResponseAccumulator: HTTPClientResponseDelegate {
368368
}
369369

370370
public struct ResponseTooBigError: Error, CustomStringConvertible {
371-
var maxBodySize: Int
371+
public var maxBodySize: Int
372+
public init(maxBodySize: Int) {
373+
self.maxBodySize = maxBodySize
374+
}
372375

373376
public var description: String {
374377
return "ResponseTooBigError: received response body exceeds maximum accepted size of \(self.maxBodySize) bytes"
@@ -385,24 +388,35 @@ public final class ResponseAccumulator: HTTPClientResponseDelegate {
385388
///
386389
/// Default is 2^32.
387390
/// - precondition: not allowed to exceed 2^32 because ``ByteBuffer`` can not store more bytes
388-
public var maxBodySize: Int = maxByteBufferSize {
389-
didSet {
390-
precondition(self.maxBodySize >= 0, "maxBodyLength is not allowed to be negative")
391-
precondition(
392-
self.maxBodySize <= Self.maxByteBufferSize,
393-
"maxBodyLength is not allowed to exceed 2^32 because ByteBuffer can not store more bytes"
394-
)
395-
}
391+
public let maxBodySize: Int
392+
393+
public convenience init(request: HTTPClient.Request) {
394+
self.init(request: request, maxBodySize: Self.maxByteBufferSize)
396395
}
397396

398-
public init(request: HTTPClient.Request) {
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+
)
399411
self.request = request
412+
self.maxBodySize = maxBodySize
400413
}
401414

402415
public func didReceiveHead(task: HTTPClient.Task<Response>, _ head: HTTPResponseHead) -> EventLoopFuture<Void> {
403416
switch self.state {
404417
case .idle:
405-
if let contentLength = head.headers.first(name: "Content-Length"),
418+
if self.request.method != .HEAD,
419+
let contentLength = head.headers.first(name: "Content-Length"),
406420
let announcedBodySize = Int(contentLength),
407421
announcedBodySize > self.maxBodySize {
408422
let error = ResponseTooBigError(maxBodySize: maxBodySize)

‎Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1331,6 +1331,25 @@ class HTTPEchoHandler: ChannelInboundHandler {
13311331
}
13321332
}
13331333

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+
13341353
final class HTTP200DelayedHandler: ChannelInboundHandler {
13351354
typealias InboundIn = HTTPServerRequestPart
13361355
typealias OutboundOut = HTTPServerResponsePart

‎Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ extension HTTPClientTests {
134134
("testBiDirectionalStreaming", testBiDirectionalStreaming),
135135
("testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLength),
136136
("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithContentLength),
137+
("testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead", testResponseAccumulatorMaxBodySizeLimitExceedingWithContentLengthButMethodIsHead),
137138
("testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked),
138139
("testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked", testResponseAccumulatorMaxBodySizeLimitNotExceedingWithTransferEncodingChuncked),
139140
("testBiDirectionalStreamingEarly200", testBiDirectionalStreamingEarly200),

‎Tests/AsyncHTTPClientTests/HTTPClientTests.swift‎

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3100,11 +3100,9 @@ class HTTPClientTests: XCTestCase {
31003100

31013101
var request = try Request(url: httpBin.baseURL)
31023102
request.body = .byteBuffer(body)
3103-
let delegate = ResponseAccumulator(request: request)
3104-
delegate.maxBodySize = 10
31053103
XCTAssertThrowsError(try self.defaultClient.execute(
31063104
request: request,
3107-
delegate: delegate
3105+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
31083106
).wait()) { error in
31093107
XCTAssertTrue(error is ResponseAccumulator.ResponseTooBigError, "unexpected error \(error)")
31103108
}
@@ -3118,16 +3116,30 @@ class HTTPClientTests: XCTestCase {
31183116

31193117
var request = try Request(url: httpBin.baseURL)
31203118
request.body = .byteBuffer(body)
3121-
let delegate = ResponseAccumulator(request: request)
3122-
delegate.maxBodySize = 10
31233119
let response = try self.defaultClient.execute(
31243120
request: request,
3125-
delegate: delegate
3121+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
31263122
).wait()
31273123

31283124
XCTAssertEqual(response.body, body)
31293125
}
31303126

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+
31313143
func testResponseAccumulatorMaxBodySizeLimitExceedingWithTransferEncodingChuncked() throws {
31323144
let httpBin = HTTPBin(.http1_1(ssl: false, compress: false)) { _ in HTTPEchoHandler() }
31333145
defer { XCTAssertNoThrow(try httpBin.shutdown()) }
@@ -3138,11 +3150,9 @@ class HTTPClientTests: XCTestCase {
31383150
request.body = .stream { writer in
31393151
writer.write(.byteBuffer(body))
31403152
}
3141-
let delegate = ResponseAccumulator(request: request)
3142-
delegate.maxBodySize = 10
31433153
XCTAssertThrowsError(try self.defaultClient.execute(
31443154
request: request,
3145-
delegate: delegate
3155+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
31463156
).wait()) { error in
31473157
XCTAssertTrue(error is ResponseAccumulator.ResponseTooBigError, "unexpected error \(error)")
31483158
}
@@ -3158,11 +3168,9 @@ class HTTPClientTests: XCTestCase {
31583168
request.body = .stream { writer in
31593169
writer.write(.byteBuffer(body))
31603170
}
3161-
let delegate = ResponseAccumulator(request: request)
3162-
delegate.maxBodySize = 10
31633171
let response = try self.defaultClient.execute(
31643172
request: request,
3165-
delegate: delegate
3173+
delegate: ResponseAccumulator(request: request, maxBodySize: 10)
31663174
).wait()
31673175

31683176
XCTAssertEqual(response.body, body)

0 commit comments

Comments
 (0)