handle partial reads from source streams

* Add InsistentReaderBytesIO wrapper that will keep reading from the wrapped stream until empty or enough bytes have been collected
* update prep_stream_data to make sure that all input data is wrapped in InsistentReaderBytesIO

Author: mattsb42-aws

src/aws_encryption_sdk/internal/utils/__init__.py

@@ -23,6 +23,8 @@
 from aws_encryption_sdk.internal.str_ops import to_bytes
 from aws_encryption_sdk.structures import EncryptedDataKey
 
+from .streams import InsistentReaderBytesIO
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -132,12 +134,14 @@ def prep_stream_data(data):
 
     :param data: Input data
     :returns: Prepared stream
-    :rtype: io.BytesIO
+    :rtype: InsistentReaderBytesIO
     """
     if isinstance(data, (six.string_types, six.binary_type)):
-        return io.BytesIO(to_bytes(data))
+        stream = io.BytesIO(to_bytes(data))
+    else:
+        stream = data
 
-    return data
+    return InsistentReaderBytesIO(stream)
 
 
 def source_data_key_length_check(source_data_key, algorithm):

src/aws_encryption_sdk/internal/utils/streams.py

@@ -11,9 +11,12 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Helper stream utility objects for AWS Encryption SDK."""
+import io
+
 from wrapt import ObjectProxy
 
 from aws_encryption_sdk.exceptions import ActionNotAllowedError
+from aws_encryption_sdk.internal.str_ops import to_bytes
 
 
 class ROStream(ObjectProxy):
@@ -56,3 +59,36 @@ def read(self, b=None):
         data = self.__wrapped__.read(b)
         self.__tee.write(data)
         return data
+
+
+class InsistentReaderBytesIO(ObjectProxy):
+    """Wrapper around a readable stream that insists on reading exactly the requested
+    number of bytes. It will keep trying to read bytes from the wrapped stream until
+    either the requested number of bytes are available or the wrapped stream has
+    nothing more to return.
+
+    :param wrapped: File-like object
+    """
+
+    def read(self, b=-1):
+        """Keep reading from source stream until either the source stream is done
+        or the requested number of bytes have been obtained.
+
+        :param int b: number of bytes to read
+        :return: All bytes read from wrapped stream
+        :rtype: bytes
+        """
+        remaining_bytes = b
+        data = io.BytesIO()
+        while True:
+            chunk = to_bytes(self.__wrapped__.read(remaining_bytes))
+
+            if not chunk:
+                break
+
+            data.write(chunk)
+            remaining_bytes -= len(chunk)
+
+            if remaining_bytes <= 0:
+                break
+        return data.getvalue()

test/functional/test_f_aws_encryption_sdk_client.py

@@ -669,4 +669,4 @@ def test_incomplete_read_stream_cycle(caplog, frame_length):
                     "Unexpected error encrypting message: infinite loop detected."
                 )
 
-    assert ciphertext != decrypted == VALUES["plaintext_128"]
+    assert ciphertext != decrypted == plaintext

test/unit/test_util_streams.py

@@ -16,13 +16,25 @@
 import pytest
 
 from aws_encryption_sdk.exceptions import ActionNotAllowedError
+from aws_encryption_sdk.internal.str_ops import to_bytes, to_str
 from aws_encryption_sdk.internal.utils.streams import ROStream, TeeStream
 
+from .unit_test_utils import NothingButRead, SometimesIncompleteReaderIO
+
 pytestmark = [pytest.mark.unit, pytest.mark.local]
 
 
-def data():
-    return io.BytesIO(b"asdijfhoaisjdfoiasjdfoijawef")
+def data(length=None, stream_type=io.BytesIO, converter=to_bytes):
+    source = b"asdijfhoaisjdfoiasjdfoijawef"
+    chunk_length = 100
+
+    if length is None:
+        length = len(source)
+
+    while len(source) < length:
+        source += source[:chunk_length]
+
+    return stream_type(converter(source[:length]))
 
 
 def test_rostream():
@@ -41,3 +53,24 @@ def test_teestream_full():
     raw_read = test_tee.read()
 
     assert data().getvalue() == raw_read == new_tee.getvalue()
+
+
+@pytest.mark.parametrize(
+    "stream_type, converter",
+    (
+        (io.BytesIO, to_bytes),
+        (SometimesIncompleteReaderIO, to_bytes),
+        (io.StringIO, to_str),
+        (NothingButRead, to_bytes),
+    ),
+)
+@pytest.mark.parametrize("bytes_to_read", (1, 4, 7, 99, 500))
+@pytest.mark.parametrize("source_length", (1, 11, 100))
+def test_insistent_stream(source_length, bytes_to_read, stream_type, converter):
+    source = data(length=source_length, stream_type=stream_type, converter=converter)
+
+    test = source.read(bytes_to_read)
+
+    assert (source_length >= bytes_to_read and len(test) == bytes_to_read) or (
+        source_length < bytes_to_read and len(test) == source_length
+    )

test/unit/unit_test_utils.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """Utility functions to handle common test framework functions."""
 import copy
+import io
 import itertools
 
 
@@ -50,3 +51,24 @@ def build_valid_kwargs_list(base, optional_kwargs):
             _kwargs.update(dict(valid_options))
             valid_kwargs.append(_kwargs)
     return valid_kwargs
+
+
+class SometimesIncompleteReaderIO(io.BytesIO):
+    def __init__(self, *args, **kwargs):
+        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:
+            size //= 2
+        return super(SometimesIncompleteReaderIO, self).read(size)
+
+
+class NothingButRead(object):
+    def __init__(self, data):
+        self._data = io.BytesIO(data)
+
+    def read(self, size=-1):
+        return self._data.read(size)