feat: Add readFully extension method to SdkSource which reads an exact amount of bytes (#799)

Author: lauzadis

.changes/55bff483-d8ae-4872-9d64-6187c28ba923.json

@@ -0,0 +1,5 @@
+{
+    "id": "55bff483-d8ae-4872-9d64-6187c28ba923",
+    "type": "feature",
+    "description": "Add readFully extension method to SdkSource"
+}

gradle.properties

@@ -44,4 +44,4 @@ kotlinLoggingVersion=3.0.0
 slf4jVersion=2.0.6
 
 # crt
-crtKotlinVersion=0.6.8-SNAPSHOT
+crtKotlinVersion=0.6.9-SNAPSHOT

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

@@ -53,3 +53,26 @@ public expect suspend fun SdkSource.readToByteArray(): ByteArray
  */
 @InternalApi
 public expect fun SdkSource.toSdkByteReadChannel(coroutineScope: CoroutineScope? = null): SdkByteReadChannel
+
+/**
+ * Remove exactly [byteCount] bytes from this source and appends them to [sink].
+ * @param sink The sink to append bytes to
+ * @param byteCount the number of bytes to read from the source
+ * @throws [IllegalArgumentException] when [byteCount] is less than zero
+ * @throws [EOFException] when the source is exhausted before [byteCount] bytes could be read
+ */
+@Throws(IOException::class)
+public fun SdkSource.readFully(sink: SdkBuffer, byteCount: Long) {
+    require(byteCount >= 0L) { "Invalid length ($byteCount) must be >= 0L" }
+
+    var totalBytesRead = 0L
+    while (totalBytesRead != byteCount) {
+        val rc = read(sink, byteCount - totalBytesRead)
+
+        if (rc == -1L) {
+            throw EOFException("Unexpected EOF: expected ${byteCount - totalBytesRead} more bytes; consumed: $totalBytesRead")
+        }
+
+        totalBytesRead += rc
+    }
+}

runtime/runtime-core/common/test/aws/smithy/kotlin/runtime/io/SdkBufferedSourceTest.kt

@@ -315,4 +315,35 @@ abstract class BufferedSourceTest(
             source.require(1024 * 9 + 1)
         }
     }
+
+    @Test
+    fun testReadFully() {
+        val data = "123456789".repeat(1024)
+        sink.writeUtf8(data)
+        sink.flush()
+
+        val dest = SdkBuffer()
+        source.readFully(dest, data.length.toLong())
+        assertEquals(data, dest.readUtf8())
+    }
+
+    @Test
+    fun testReadFullyIllegalArgumentException() {
+        val dest = SdkBuffer()
+        assertFailsWith<IllegalArgumentException> {
+            source.readFully(dest, -1)
+        }
+    }
+
+    @Test
+    fun testReadFullyEOFException() {
+        val data = "123456789".repeat(1024)
+        sink.writeUtf8(data)
+        sink.flush()
+
+        val dest = SdkBuffer()
+        assertFailsWith<EOFException> {
+            source.readFully(dest, data.length.toLong() + 1)
+        }
+    }
 }