fix: Reading from a DataStreamReader with rewind: true properly rewinds the stream allowing for subsequent reads. (#539)

* Fixes an issue where reading with rewind true didn't work

* Adds test for writing+read-with-rewind

Author: Ed Paulosky

Sources/ClientRuntime/Networking/Streaming/DataStreamReader.swift

@@ -36,7 +36,7 @@ public class DataStreamReader: StreamReader {
 
     init(byteBuffer: ByteBuffer = ByteBuffer(size: 0)) {
         self.byteBuffer = byteBuffer
-        self._availableForRead = 0
+        self._availableForRead = UInt(byteBuffer.length())
         self.offset = 0
         self._hasFinishedWriting = false
     }
@@ -51,9 +51,16 @@ public class DataStreamReader: StreamReader {
                 bytesRead = byteBuffer.read(buffer: typedBuffer)
             }
 
-            if !rewind, let bytesRead = bytesRead {
-                _availableForRead -= UInt(bytesRead)
-                offset += UInt(bytesRead)
+            if let bytesRead = bytesRead {
+                if rewind {
+                    try? byteBuffer.seek(
+                        offset: Int64(offset),
+                        streamSeekType: .begin
+                    )
+                } else {
+                    _availableForRead -= UInt(bytesRead)
+                    offset += UInt(bytesRead)
+                }
             }
         }
         return ByteBuffer(data: data)

Tests/ClientRuntimeTests/StreamingTests/DataStreamReaderTests.swift

@@ -0,0 +1,54 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+@testable import ClientRuntime
+import AwsCommonRuntimeKit
+import XCTest
+
+class DataStreamReaderTests: XCTestCase {
+    let testData = Data([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A])
+    
+    func testRead() {
+        let byteBuffer = ByteBuffer(data: testData)
+        let dataStreamReader = DataStreamReader(byteBuffer: byteBuffer)
+        
+        let firstRead = dataStreamReader.read().getData()
+        XCTAssertEqual(firstRead, testData)
+        
+        let secondRead = dataStreamReader.read().getData()
+        XCTAssertNotEqual(secondRead, testData)
+        XCTAssertEqual(secondRead, Data())
+    }
+    
+    func testReadWithRewind() {
+        let byteBuffer = ByteBuffer(data: testData)
+        let dataStreamReader = DataStreamReader(byteBuffer: byteBuffer)
+        
+        let firstRead = dataStreamReader.read(rewind: true).getData()
+        XCTAssertEqual(firstRead, testData)
+        
+        let secondRead = dataStreamReader.read().getData()
+        XCTAssertEqual(secondRead, testData)
+    }
+    
+    func testWriteThenReadWithRewind() {
+        let dataStreamReader = DataStreamReader()
+        
+        let byteBuffer = ByteBuffer(data: testData)
+        dataStreamReader.write(buffer: byteBuffer)
+        
+        let firstRead = dataStreamReader.read(rewind: true).getData()
+        XCTAssertEqual(firstRead, testData)
+        
+        let secondRead = dataStreamReader.read().getData()
+        XCTAssertEqual(secondRead, testData)
+        
+        let thirdRead = dataStreamReader.read().getData()
+        XCTAssertNotEqual(thirdRead, testData)
+        XCTAssertEqual(thirdRead, Data())
+    }
+}