Fixes issue with non-seekable sources breaking StreamEncryptor and StreamDecryptor handling

src/aws_encryption_sdk/internal/utils/streams.py

@@ -0,0 +1,76 @@
+# Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# 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 aws_encryption_sdk.exceptions import ActionNotAllowedError
+
+
+class PassThroughStream(object):
+    """Provides a pass-through interface on top of a stream object.
+
+    :param source_stream: File-like object
+    """
+
+    def __init__(self, source_stream):
+        """Prepares the passthroughs."""
+        self._source_stream = source_stream
+        self._duplicate_api()
+
+    def _duplicate_api(self):
+        """Maps the source stream API onto this object."""
+        source_attributes = set([
+            method for method in dir(self._source_stream)
+            if not method.startswith('_')
+        ])
+        self_attributes = set(dir(self))
+        for attribute in source_attributes.difference(self_attributes):
+            setattr(self, attribute, getattr(self._source_stream, attribute))
+
+
+class ROStream(PassThroughStream):
+    """Provides a read-only interface on top of a stream object.
+
+    Used to provide MasterKeyProviders with read-only access to plaintext.
+
+    :param source_stream: File-like object
+    """
+
+    def write(self, b):  # pylint: disable=unused-argument
+        """Blocks calls to write.
+
+        :raises ActionNotAllowedError: when called
+        """
+        raise ActionNotAllowedError('Write not allowed on ROStream objects')
+
+
+class TeeStream(PassThroughStream):
+    """Provides a ``tee``-like interface on top of a stream object, which collects read bytes
+    into a local :class:`io.BytesIO`.
+
+    :param source_stream: File-like object
+    """
+
+    def __init__(self, source_stream):
+        """Creates the local tee stream."""
+        self.tee = io.BytesIO()
+        super(TeeStream, self).__init__(source_stream)
+
+    def read(self, b=None):
+        """Reads data from source, copying it into ``tee`` before returning.
+
+        :param int b: number of bytes to read
+        """
+        data = self._source_stream.read(b)
+        self.tee.write(data)
+        return data

test/functional/test_f_aws_encryption_sdk_client.py

@@ -30,6 +30,7 @@
 from aws_encryption_sdk.identifiers import Algorithm, EncryptionKeyType, WrappingAlgorithm
 from aws_encryption_sdk.internal.crypto.wrapping_keys import WrappingKey
 from aws_encryption_sdk.internal.formatting.encryption_context import serialize_encryption_context
+from aws_encryption_sdk.internal.utils import ROStream
 from aws_encryption_sdk.key_providers.base import MasterKeyProviderConfig
 from aws_encryption_sdk.key_providers.raw import RawMasterKeyProvider
 from aws_encryption_sdk.materials_managers import DecryptionMaterialsRequest, EncryptionMaterialsRequest
@@ -534,3 +535,31 @@ def test_encrypt_source_length_enforcement_legacy_support():
         key_provider=key_provider,
         source_length=int(len(VALUES['plaintext_128']) / 2)
     )
+
+
+class NoSeekBytesIO(io.BytesIO):
+    """``io.BytesIO`` that blocks ``seek()`` and ``tell()``."""
+
+    def seekable(self):
+        return False
+
+    def seek(self, offset, whence=0):
+        raise NotImplementedError('seek is blocked')
+
+    def tell(self):
+        raise NotImplementedError('tell is blocked')
+
+
+def test_decrypt_no_seek_input():
+    """Test that StreamDecryptor can decrypt an input stream that is not seekable."""
+    key_provider = fake_kms_key_provider()
+    ciphertext, _header = aws_encryption_sdk.encrypt(
+        source=VALUES['plaintext_128'],
+        key_provider=key_provider,
+        encryption_context=VALUES['encryption_context']
+    )
+    ciphertext_no_seek = NoSeekBytesIO(ciphertext)
+    aws_encryption_sdk.decrypt(
+        source=ciphertext_no_seek,
+        key_provider=key_provider
+    )