make InsistentReaderBytesIO properly handle a stream that closes while being read from

Author: mattsb42-aws

src/aws_encryption_sdk/internal/utils/streams.py

@@ -81,7 +81,12 @@ def read(self, b=-1):
         remaining_bytes = b
         data = io.BytesIO()
         while True:
-            chunk = to_bytes(self.__wrapped__.read(remaining_bytes))
+            try:
+                chunk = to_bytes(self.__wrapped__.read(remaining_bytes))
+            except ValueError:
+                if self.__wrapped__.closed:
+                    break
+                raise
 
             if not chunk:
                 break

test/unit/test_util_streams.py

@@ -19,7 +19,7 @@
 from aws_encryption_sdk.internal.str_ops import to_bytes, to_str
 from aws_encryption_sdk.internal.utils.streams import InsistentReaderBytesIO, ROStream, TeeStream
 
-from .unit_test_utils import NothingButRead, SometimesIncompleteReaderIO
+from .unit_test_utils import ExactlyTwoReads, NothingButRead, SometimesIncompleteReaderIO
 
 pytestmark = [pytest.mark.unit, pytest.mark.local]
 
@@ -74,3 +74,16 @@ def test_insistent_stream(source_length, bytes_to_read, stream_type, converter):
     assert (source_length >= bytes_to_read and len(test) == bytes_to_read) or (
         source_length < bytes_to_read and len(test) == source_length
     )
+
+
+def test_insistent_stream_close_partway_through():
+    raw = data(length=100)
+    source = ExactlyTwoReads(raw.getvalue())
+
+    wrapped = InsistentReaderBytesIO(source)
+
+    test = b""
+    test += wrapped.read(10)  # actually reads 10 bytes
+    test += wrapped.read(10)  # reads 5 bytes, stream is closed before third read can complete, truncating the result
+
+    assert test == raw.getvalue()[:15]

test/unit/unit_test_utils.py

@@ -55,13 +55,13 @@ def build_valid_kwargs_list(base, optional_kwargs):
 
 class SometimesIncompleteReaderIO(io.BytesIO):
     def __init__(self, *args, **kwargs):
-        self.__read_counter = 0
+        self._read_counter = 0
         super(SometimesIncompleteReaderIO, self).__init__(*args, **kwargs)
 
     def read(self, size=-1):
         """Every other read request, return fewer than the requested number of bytes if more than one byte requested."""
-        self.__read_counter += 1
-        if size > 1 and self.__read_counter % 2 == 0:
+        self._read_counter += 1
+        if size > 1 and self._read_counter % 2 == 0:
             size //= 2
         return super(SometimesIncompleteReaderIO, self).read(size)
 
@@ -72,3 +72,10 @@ def __init__(self, data):
 
     def read(self, size=-1):
         return self._data.read(size)
+
+
+class ExactlyTwoReads(SometimesIncompleteReaderIO):
+    def read(self, size=-1):
+        if self._read_counter >= 2:
+            self.close()
+        return super(ExactlyTwoReads, self).read(size)