From f088c5bb37dde25c476b99c7332dfbc0c4e87404 Mon Sep 17 00:00:00 2001 From: Adam Fowler Date: Mon, 9 Nov 2020 18:28:45 +0000 Subject: [PATCH] Add streaming request errors Throw errors for when a streaming request either provides too much data or not enough data --- Sources/SotoCore/AWSClient.swift | 5 ++ Sources/SotoCore/HTTP/S3StreamReader.swift | 12 ++++- .../SotoCore/HTTP/StreamWriter+write.swift | 8 +++ Tests/SotoCoreTests/AWSClientTests.swift | 54 +++++++++++++------ 4 files changed, 62 insertions(+), 17 deletions(-) diff --git a/Sources/SotoCore/AWSClient.swift b/Sources/SotoCore/AWSClient.swift index 7f108d162..d7e6a2956 100644 --- a/Sources/SotoCore/AWSClient.swift +++ b/Sources/SotoCore/AWSClient.swift @@ -35,6 +35,7 @@ public final class AWSClient { case alreadyShutdown case invalidURL case tooMuchData + case notEnoughData } let error: Error @@ -45,6 +46,8 @@ public final class AWSClient { public static var invalidURL: ClientError { .init(error: .invalidURL) } /// Too much data has been supplied for the Request public static var tooMuchData: ClientError { .init(error: .tooMuchData) } + /// Not enough data has been supplied for the Request + public static var notEnoughData: ClientError { .init(error: .notEnoughData) } } /// Specifies how `HTTPClient` will be created and establishes lifecycle ownership. @@ -549,6 +552,8 @@ extension AWSClient.ClientError: CustomStringConvertible { """ case .tooMuchData: return "You have supplied too much data for the Request." + case .notEnoughData: + return "You have not supplied enough data for the Request." } } } diff --git a/Sources/SotoCore/HTTP/S3StreamReader.swift b/Sources/SotoCore/HTTP/S3StreamReader.swift index 710e031bb..b3bebd96c 100644 --- a/Sources/SotoCore/HTTP/S3StreamReader.swift +++ b/Sources/SotoCore/HTTP/S3StreamReader.swift @@ -89,11 +89,19 @@ class S3ChunkedStreamReader: StreamReader { self.read(eventLoop).map { (result) -> Void in // check if a byte buffer was returned. If not then it must have been `.end` guard case .byteBuffer(var buffer) = result else { - // if we have `.end` then return what we have in the working buffer - promise.succeed(self.workingBuffer) + if self.bytesLeftToRead == self.workingBuffer.readableBytes { + promise.succeed(self.workingBuffer) + } else { + promise.fail(AWSClient.ClientError.notEnoughData) + } return } self.bytesLeftToRead -= buffer.readableBytes + + guard self.bytesLeftToRead >= 0 else { + promise.fail(AWSClient.ClientError.tooMuchData) + return + } // if working buffer is empty and this buffer is the chunk buffer size or there is no data // left to read and this buffer is less than the size of the chunk buffer then just return // this buffer. This allows us to avoid the buffer copy diff --git a/Sources/SotoCore/HTTP/StreamWriter+write.swift b/Sources/SotoCore/HTTP/StreamWriter+write.swift index 9f5312893..b96f5fc75 100644 --- a/Sources/SotoCore/HTTP/StreamWriter+write.swift +++ b/Sources/SotoCore/HTTP/StreamWriter+write.swift @@ -37,8 +37,16 @@ extension AsyncHTTPClient.HTTPClient.Body.StreamWriter { // calculate amount left to write let newAmountLeft: Int? if let amountLeft = amountLeft { + guard byteBuffers.count > 0 else { + promise.fail(AWSClient.ClientError.notEnoughData) + return + } let bytesToWrite = byteBuffers.reduce(0) { $0 + $1.readableBytes } newAmountLeft = amountLeft - bytesToWrite + guard newAmountLeft! >= 0 else { + promise.fail(AWSClient.ClientError.tooMuchData) + return + } } else { newAmountLeft = nil } diff --git a/Tests/SotoCoreTests/AWSClientTests.swift b/Tests/SotoCoreTests/AWSClientTests.swift index 541f46afc..ad92c3ca9 100644 --- a/Tests/SotoCoreTests/AWSClientTests.swift +++ b/Tests/SotoCoreTests/AWSClientTests.swift @@ -232,14 +232,16 @@ class AWSClientTests: XCTestCase { let payload = AWSPayload.stream(size: bufferSize) { eventLoop in let size = min(blockSize, byteBuffer.readableBytes) // don't ask for 0 bytes - XCTAssertNotEqual(size, 0) + if size == 0 { + return eventLoop.makeSucceededFuture(.end) + } let buffer = byteBuffer.readSlice(length: size)! return eventLoop.makeSucceededFuture(.byteBuffer(buffer)) } let input = Input(payload: payload) let response = client.execute(operation: "test", path: "/", httpMethod: .POST, serviceConfig: config, input: input, logger: TestEnvironment.logger) - try server.processRaw { request in + try? server.processRaw { request in let bytes = request.body.getBytes(at: 0, length: request.body.readableBytes) XCTAssertEqual(bytes, data) let response = AWSTestServer.Response(httpStatus: .ok, headers: [:], body: nil) @@ -287,7 +289,7 @@ class AWSClientTests: XCTestCase { XCTAssertNoThrow(try self.testRequestStreaming(config: config, client: client, server: awsServer, bufferSize: 65552, blockSize: 65552)) } - func testRequestStreamingTooMuchData() { + func testRequestStreamingWithPayload(_ payload: AWSPayload) throws { struct Input: AWSEncodableShape & AWSShapeWithPayload { static var _payloadPath: String = "payload" static var _payloadOptions: AWSShapePayloadOptions = [.allowStreaming] @@ -305,19 +307,41 @@ class AWSClientTests: XCTestCase { XCTAssertNoThrow(try client.syncShutdown()) XCTAssertNoThrow(try httpClient.syncShutdown()) } - do { - // set up stream of 8 bytes but supply more than that - let payload = AWSPayload.stream(size: 8) { eventLoop in - var buffer = ByteBufferAllocator().buffer(capacity: 0) - buffer.writeString("String longer than 8 bytes") - return eventLoop.makeSucceededFuture(.byteBuffer(buffer)) + let input = Input(payload: payload) + let response = client.execute(operation: "test", path: "/", httpMethod: .POST, serviceConfig: config, input: input, logger: TestEnvironment.logger) + try response.wait() + } + + func testRequestStreamingTooMuchData() { + // set up stream of 8 bytes but supply more than that + let payload = AWSPayload.stream(size: 8) { eventLoop in + var buffer = ByteBufferAllocator().buffer(capacity: 0) + buffer.writeString("String longer than 8 bytes") + return eventLoop.makeSucceededFuture(.byteBuffer(buffer)) + } + XCTAssertThrowsError(try testRequestStreamingWithPayload(payload)) { error in + guard let error = error as? AWSClient.ClientError, error == .tooMuchData else { + XCTFail() + return + } + } + } + + func testRequestStreamingNotEnoughData() { + var byteBuffer = ByteBufferAllocator().buffer(staticString: "Buffer") + let payload = AWSPayload.stream(size: byteBuffer.readableBytes+1) { eventLoop in + let size = byteBuffer.readableBytes + if size == 0 { + return eventLoop.makeSucceededFuture(.end) + } + let buffer = byteBuffer.readSlice(length: size)! + return eventLoop.makeSucceededFuture(.byteBuffer(buffer)) + } + XCTAssertThrowsError(try testRequestStreamingWithPayload(payload)) { error in + guard let error = error as? AWSClient.ClientError, error == .notEnoughData else { + XCTFail() + return } - let input = Input(payload: payload) - let response = client.execute(operation: "test", path: "/", httpMethod: .POST, serviceConfig: config, input: input, logger: TestEnvironment.logger) - try response.wait() - } catch let error as HTTPClientError where error == .bodyLengthMismatch { - } catch { - XCTFail("Unexpected error: \(error)") } }