fix: Check if an InputStream is replayable when setting isOneShot (…

…#1187)
smithy-lang · Nov 22, 2024 · 1b5d3ce · 1b5d3ce
1 parent 5f288d6
commit 1b5d3ce
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/.changes/68aac1d5-64f9-4569-bf95-db8d5712a2df.json b/.changes/68aac1d5-64f9-4569-bf95-db8d5712a2df.json
@@ -0,0 +1,5 @@
+{
+    "id": "68aac1d5-64f9-4569-bf95-db8d5712a2df",
+    "type": "bugfix",
+    "description": "Infer the replayability of `InputStream` using `markSupported()` when setting `isOneShot`"
+}
diff --git a/runtime/runtime-core/jvm/src/aws/smithy/kotlin/runtime/content/ByteStreamJVM.kt b/runtime/runtime-core/jvm/src/aws/smithy/kotlin/runtime/content/ByteStreamJVM.kt
@@ -117,7 +117,7 @@ public fun InputStream.asByteStream(contentLength: Long? = null): ByteStream.Sou
     val source = source()
     return object : ByteStream.SourceStream() {
         override val contentLength: Long? = contentLength
-        override val isOneShot: Boolean = true
+        override val isOneShot: Boolean = !markSupported()
         override fun readFrom(): SdkSource = source
     }
 }

diff --git a/runtime/runtime-core/jvm/test/aws/smithy/kotlin/runtime/content/ByteStreamJVMTest.kt b/runtime/runtime-core/jvm/test/aws/smithy/kotlin/runtime/content/ByteStreamJVMTest.kt
@@ -8,6 +8,7 @@ package aws.smithy.kotlin.runtime.content
 import aws.smithy.kotlin.runtime.testing.RandomTempFile
 import kotlinx.coroutines.test.runTest
 import java.io.ByteArrayOutputStream
+import java.io.InputStream
 import java.io.OutputStream
 import java.nio.file.Files
 import kotlin.test.*
@@ -159,7 +160,7 @@ class ByteStreamJVMTest {
         binaryData.inputStream().use { inputStream ->
             val byteStream = inputStream.asByteStream()
             assertNull(byteStream.contentLength)
-            assertTrue(byteStream.isOneShot)
+            assertFalse(byteStream.isOneShot)
 
             val output = byteStream.toByteArray()
             assertContentEquals(binaryData, output)
@@ -169,6 +170,23 @@ class ByteStreamJVMTest {
     @Test
     fun testInputStreamAsByteStreamWithLength() = runTest {
         binaryData.inputStream().use { inputStream ->
+            val byteStream = inputStream.asByteStream(binaryData.size.toLong())
+            assertEquals(binaryData.size.toLong(), byteStream.contentLength)
+            assertFalse(byteStream.isOneShot)
+
+            val output = byteStream.toByteArray()
+            assertContentEquals(binaryData, output)
+        }
+    }
+
+    @Test
+    fun testOneShotInputStream() = runTest {
+        class NonReplayableInputStream(val delegate: InputStream) : InputStream() {
+            override fun read(): Int = delegate.read()
+            override fun markSupported(): Boolean = false
+        }
+
+        NonReplayableInputStream(binaryData.inputStream()).use { inputStream ->
             val byteStream = inputStream.asByteStream(binaryData.size.toLong())
             assertEquals(binaryData.size.toLong(), byteStream.contentLength)
             assertTrue(byteStream.isOneShot)