fix: correctly return number of bytes read from chunked streams

Author: ianbotsf

.changes/7251f5e7-0e9a-4ea4-b4c7-30dad31f4622.json

@@ -0,0 +1,9 @@
+{
+    "id": "7251f5e7-0e9a-4ea4-b4c7-30dad31f4622",
+    "type": "bugfix",
+    "description": "⚠️ **IMPORTANT**: Correctly return number of bytes read from chunked streams",
+    "issues": [
+        "awslabs/smithy-kotlin#1285"
+    ],
+    "requiresMinorVersionBump": true
+}

runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/AwsChunkedByteReadChannel.kt

@@ -48,7 +48,8 @@ public class AwsChunkedByteReadChannel(
     override suspend fun read(sink: SdkBuffer, limit: Long): Long {
         require(limit >= 0L) { "Invalid limit ($limit) must be >= 0L" }
         if (!chunkReader.ensureValidChunk()) return -1L
-        return chunkReader.chunk.read(sink, limit)
+        chunkReader.chunk.read(sink, limit)
+        return chunkReader.readCountAndReset()
     }
 }
 

runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/AwsChunkedSource.kt

@@ -47,7 +47,8 @@ public class AwsChunkedSource(
             chunkReader.ensureValidChunk()
         }
         if (!isChunkValid) return -1L
-        return chunkReader.chunk.read(sink, limit)
+        chunkReader.chunk.read(sink, limit)
+        return chunkReader.readCountAndReset()
     }
 
     override fun close() {

runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/internal/AwsChunkedReader.kt

@@ -57,6 +57,15 @@ internal class AwsChunkedReader(
      */
     internal var hasLastChunkBeenSent: Boolean = false
 
+    private var read = 0L
+
+    /**
+     * Gets the most recent read count and then resets the counter. This is meant to be checked after every caller
+     * invocation of AwsChunkedReader.chunk.read(...). Note that because of buffering, this value may return _more_ than
+     * the total number of bytes written to the sink.
+     */
+    fun readCountAndReset(): Long = read.also { read = 0L }
+
     /**
      * Ensures that the internal [chunk] is valid for reading. If it's not valid, try to load the next chunk. Note that
      * this function will suspend until the whole chunk has been loaded.
@@ -117,6 +126,7 @@ internal class AwsChunkedReader(
         while (remaining > 0L) {
             val rc = read(sink, remaining)
             if (rc == -1L) break
+            read += rc
             remaining -= rc
         }
 

runtime/auth/aws-signing-tests/common/src/aws/smithy/kotlin/runtime/auth/awssigning/tests/AwsChunkedByteReadChannelTestBase.kt

@@ -47,11 +47,12 @@ abstract class AwsChunkedByteReadChannelTestBase : AwsChunkedTestBase(AwsChunked
         val sink = SdkBuffer()
 
         val bytesRead = awsChunked.readAll(sink)
-        writeJob.join()
+        // writeJob.join()
+        assertEquals(dataLengthBytes.toLong(), bytesRead)
+        assertEquals(totalBytesExpected.toLong(), sink.size)
 
         val bytesAsString = sink.readUtf8()
 
-        assertEquals(totalBytesExpected.toLong(), bytesRead)
         assertTrue(awsChunked.isClosedForRead)
 
         val chunkSignatures = getChunkSignatures(bytesAsString)

runtime/auth/aws-signing-tests/common/src/aws/smithy/kotlin/runtime/auth/awssigning/tests/AwsChunkedTestBase.kt

@@ -22,6 +22,7 @@ interface AwsChunkedTestReader {
     // This may modify the chunked reader state and cause loss of data!
     fun isClosedForRead(): Boolean
     suspend fun read(sink: SdkBuffer, limit: Long): Long
+    suspend fun readAll(sink: SdkBuffer): Long
 }
 
 fun interface AwsChunkedReaderFactory {
@@ -32,6 +33,7 @@ fun interface AwsChunkedReaderFactory {
             object : AwsChunkedTestReader {
                 override fun isClosedForRead(): Boolean = chunked.isClosedForRead
                 override suspend fun read(sink: SdkBuffer, limit: Long): Long = chunked.read(sink, limit)
+                override suspend fun readAll(sink: SdkBuffer): Long = chunked.readAll(sink)
             }
         }
     }
@@ -198,6 +200,7 @@ abstract class AwsChunkedTestBase(
         // need to make 2 successive calls because there are two chunks -- read will only fetch the first one due to limit
         var bytesRead = awsChunked.read(sink, readLimit.toLong())
         bytesRead += awsChunked.read(sink, readLimit - bytesRead)
+        assertEquals(readLimit.toLong(), sink.size)
 
         val bytesAsString = sink.readUtf8()
 
@@ -213,7 +216,7 @@ abstract class AwsChunkedTestBase(
         assertEquals(CHUNK_SIZE_BYTES, chunkSizes[0])
         assertEquals(0, chunkSizes[1])
 
-        assertEquals(readLimit, bytesRead.toInt())
+        assertEquals(dataLengthBytes, bytesRead.toInt())
         assertTrue(awsChunked.isClosedForRead())
     }
 
@@ -257,7 +260,8 @@ abstract class AwsChunkedTestBase(
 
         val sink = SdkBuffer()
         val bytesRead = awsChunked.read(sink, readLimit.toLong())
-        assertEquals(readLimit.toLong(), bytesRead)
+        assertEquals(CHUNK_SIZE_BYTES.toLong(), bytesRead)
+        assertEquals(readLimit.toLong(), sink.size)
 
         val bytesAsString = sink.readUtf8()
         val chunkSignatures = getChunkSignatures(bytesAsString)
@@ -289,7 +293,8 @@ abstract class AwsChunkedTestBase(
             bytesRead += awsChunked.read(sink, readLimit.toLong())
         }
         bytesRead += awsChunked.read(sink, readLimit.toLong())
-        assertEquals(totalBytesExpected.toLong(), bytesRead)
+        assertEquals(dataLengthBytes.toLong(), bytesRead)
+        assertEquals(totalBytesExpected.toLong(), sink.size)
 
         val bytesAsString = sink.readUtf8()
 
@@ -345,9 +350,11 @@ abstract class AwsChunkedTestBase(
             }
         }
 
+        assertEquals(dataLengthBytes.toLong(), bytesRead)
+        assertEquals(totalBytesExpected.toLong(), sink.size)
+
         val bytesAsString = sink.readUtf8()
 
-        assertEquals(totalBytesExpected.toLong(), bytesRead)
         assertTrue(awsChunked.isClosedForRead())
 
         val chunkSignatures = getChunkSignatures(bytesAsString)
@@ -402,13 +409,9 @@ abstract class AwsChunkedTestBase(
         val totalBytesExpected = encodedChunkLength(CHUNK_SIZE_BYTES) + encodedChunkLength(0) + trailingHeadersLength + "\r\n".length
         val sink = SdkBuffer()
 
-        var bytesRead = 0L
-
-        while (bytesRead < totalBytesExpected.toLong()) {
-            bytesRead += awsChunked.read(sink, Long.MAX_VALUE)
-        }
-
-        assertEquals(totalBytesExpected.toLong(), bytesRead)
+        val bytesRead = awsChunked.readAll(sink)
+        assertEquals(dataLengthBytes.toLong(), bytesRead)
+        assertEquals(totalBytesExpected.toLong(), sink.size)
         assertTrue(awsChunked.isClosedForRead())
 
         val bytesAsString = sink.readUtf8()
@@ -445,13 +448,9 @@ abstract class AwsChunkedTestBase(
         val totalBytesExpected = encodedUnsignedChunkLength(CHUNK_SIZE_BYTES) + encodedUnsignedChunkLength(0) + "\r\n".length
         val sink = SdkBuffer()
 
-        var bytesRead = 0L
-
-        while (bytesRead < totalBytesExpected.toLong()) {
-            bytesRead += awsChunked.read(sink, Long.MAX_VALUE)
-        }
-
-        assertEquals(totalBytesExpected.toLong(), bytesRead)
+        val bytesRead = awsChunked.readAll(sink)
+        assertEquals(dataLengthBytes.toLong(), bytesRead)
+        assertEquals(totalBytesExpected.toLong(), sink.size)
         assertTrue(awsChunked.isClosedForRead())
 
         val bytesAsString = sink.readUtf8()
@@ -482,13 +481,9 @@ abstract class AwsChunkedTestBase(
         val totalBytesExpected = encodedUnsignedChunkLength(CHUNK_SIZE_BYTES) + encodedUnsignedChunkLength(0) + trailingHeadersLength + "\r\n".length
         val sink = SdkBuffer()
 
-        var bytesRead = 0L
-
-        while (bytesRead < totalBytesExpected.toLong()) {
-            bytesRead += awsChunked.read(sink, Long.MAX_VALUE)
-        }
-
-        assertEquals(totalBytesExpected.toLong(), bytesRead)
+        val bytesRead = awsChunked.readAll(sink)
+        assertEquals(dataLengthBytes.toLong(), bytesRead)
+        assertEquals(totalBytesExpected.toLong(), sink.size)
         assertTrue(awsChunked.isClosedForRead())
 
         val bytesAsString = sink.readUtf8()

runtime/auth/aws-signing-tests/jvm/src/aws/smithy/kotlin/runtime/auth/awssigning/tests/AwsChunkedSourceTestBase.kt

@@ -6,6 +6,7 @@
 package aws.smithy.kotlin.runtime.auth.awssigning.tests
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsChunkedSource
 import aws.smithy.kotlin.runtime.io.SdkBuffer
+import aws.smithy.kotlin.runtime.io.readAll
 import aws.smithy.kotlin.runtime.io.source
 
 val AwsChunkedReaderFactory.Companion.Source: AwsChunkedReaderFactory
@@ -19,6 +20,7 @@ val AwsChunkedReaderFactory.Companion.Source: AwsChunkedReaderFactory
                 return rc == -1L
             }
             override suspend fun read(sink: SdkBuffer, limit: Long): Long = chunked.read(sink, limit)
+            override suspend fun readAll(sink: SdkBuffer): Long = chunked.readAll(sink)
         }
     }
 

runtime/runtime-core/api/runtime-core.api

@@ -1075,6 +1075,7 @@ public final class aws/smithy/kotlin/runtime/io/SdkSourceJVMKt {
 }
 
 public final class aws/smithy/kotlin/runtime/io/SdkSourceKt {
+	public static final fun readAll (Laws/smithy/kotlin/runtime/io/SdkSource;Laws/smithy/kotlin/runtime/io/SdkSink;)J
 	public static final fun readFully (Laws/smithy/kotlin/runtime/io/SdkSource;Laws/smithy/kotlin/runtime/io/SdkBuffer;J)V
 }
 

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/io/SdkSource.kt

@@ -76,3 +76,19 @@ public fun SdkSource.readFully(sink: SdkBuffer, byteCount: Long) {
         totalBytesRead += rc
     }
 }
+
+/**
+ * Read all bytes from this source into [sink]. Returns the total number of bytes written.
+ */
+public fun SdkSource.readAll(sink: SdkSink): Long {
+    val bufferedSink = sink.buffer()
+    var totalWritten = 0L
+    while (true) {
+        val rc = read(bufferedSink.buffer, DEFAULT_BYTE_CHANNEL_MAX_BUFFER_SIZE.toLong())
+        if (rc == -1L) break
+        totalWritten += rc
+        bufferedSink.emit()
+    }
+    bufferedSink.emit()
+    return totalWritten
+}