Add streaming request errors

Throw errors for when a streaming request either provides too much data or not enough data

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."
         }
     }
 }

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

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
                     }

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)")
         }
     }