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

CHANGELOG.rst

@@ -2,6 +2,11 @@
 Changelog
 *********
 
+1.3.2
+=====
+* Addressed `issue #13 </~https://github.com/awslabs/aws-encryption-sdk-python/issues/13>`_
+  to properly handle non-seekable source streams.
+
 1.3.1
 =====
 

requirements.txt

@@ -1,3 +1,4 @@
 boto3>=1.4.4
 cryptography>=1.8.1
 attrs>=16.3.0
+wrapt>=1.10.11

src/aws_encryption_sdk/identifiers.py

@@ -21,7 +21,7 @@
 
 from aws_encryption_sdk.exceptions import InvalidAlgorithmError
 
-__version__ = '1.3.1'
+__version__ = '1.3.2'
 USER_AGENT_SUFFIX = 'AwsEncryptionSdkPython-KMSMasterKey/{}'.format(__version__)
 
 

src/aws_encryption_sdk/internal/formatting/deserialize.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """Components for handling AWS Encryption SDK message deserialization."""
 from __future__ import division
+import io
 import logging
 import struct
 
@@ -29,38 +30,34 @@
     EncryptedData, MessageFooter,
     MessageFrameBody, MessageHeaderAuthentication
 )
+from aws_encryption_sdk.internal.utils.streams import TeeStream
 from aws_encryption_sdk.structures import EncryptedDataKey, MasterKeyInfo, MessageHeader
 
 _LOGGER = logging.getLogger(__name__)
 
 
-def validate_header(header, header_auth, stream, header_start, header_end, data_key):
+def validate_header(header, header_auth, raw_header, data_key):
     """Validates the header using the header authentication data.
 
     :param header: Deserialized header
     :type header: aws_encryption_sdk.structures.MessageHeader
     :param header_auth: Deserialized header auth
     :type header_auth: aws_encryption_sdk.internal.structures.MessageHeaderAuthentication
-    :param stream: Stream containing serialized message
     :type stream: io.BytesIO
-    :param int header_start: Position in stream of start of serialized header
-    :param int header_end: Position in stream of end of serialized header
+    :param bytes raw_header: Raw header bytes
     :param bytes data_key: Data key with which to perform validation
     :raises SerializationError: if header authorization fails
     """
     _LOGGER.debug('Starting header validation')
-    current_position = stream.tell()
-    stream.seek(header_start)
     try:
         decrypt(
             algorithm=header.algorithm,
             key=data_key,
             encrypted_data=EncryptedData(header_auth.iv, b'', header_auth.tag),
-            associated_data=stream.read(header_end - header_start)
+            associated_data=raw_header
         )
     except InvalidTag:
         raise SerializationError('Header authorization failed')
-    stream.seek(current_position)
 
 
 def deserialize_header(stream):
@@ -69,13 +66,15 @@ def deserialize_header(stream):
     :param stream: Source data stream
     :type stream: io.BytesIO
     :returns: Deserialized MessageHeader object
-    :rtype: aws_encryption_sdk.structures.MessageHeader
+    :rtype: :class:`aws_encryption_sdk.structures.MessageHeader` and bytes
     :raises NotSupportedError: if unsupported data types are found
     :raises UnknownIdentityError: if unknown data types are found
     :raises SerializationError: if IV length does not match algorithm
     """
     _LOGGER.debug('Starting header deserialization')
-    version_id, message_type_id = unpack_values('>BB', stream)
+    tee = io.BytesIO()
+    tee_stream = TeeStream(stream, tee)
+    version_id, message_type_id = unpack_values('>BB', tee_stream)
     try:
         message_type = ObjectType(message_type_id)
     except ValueError as error:
@@ -89,7 +88,7 @@ def deserialize_header(stream):
         raise NotSupportedError('Unsupported version {}'.format(version_id), error)
     header = {'version': version, 'type': message_type}
 
-    algorithm_id, message_id, ser_encryption_context_length = unpack_values('>H16sH', stream)
+    algorithm_id, message_id, ser_encryption_context_length = unpack_values('>H16sH', tee_stream)
 
     try:
         alg = Algorithm.get_by_id(algorithm_id)
@@ -101,24 +100,24 @@ def deserialize_header(stream):
     header['message_id'] = message_id
 
     header['encryption_context'] = deserialize_encryption_context(
-        stream.read(ser_encryption_context_length)
+        tee_stream.read(ser_encryption_context_length)
     )
-    (encrypted_data_key_count,) = unpack_values('>H', stream)
+    (encrypted_data_key_count,) = unpack_values('>H', tee_stream)
 
     encrypted_data_keys = set([])
     for _ in range(encrypted_data_key_count):
-        (key_provider_length,) = unpack_values('>H', stream)
+        (key_provider_length,) = unpack_values('>H', tee_stream)
         (key_provider_identifier,) = unpack_values(
             '>{}s'.format(key_provider_length),
-            stream
+            tee_stream
         )
-        (key_provider_information_length,) = unpack_values('>H', stream)
+        (key_provider_information_length,) = unpack_values('>H', tee_stream)
         (key_provider_information,) = unpack_values(
             '>{}s'.format(key_provider_information_length),
-            stream
+            tee_stream
         )
-        (encrypted_data_key_length,) = unpack_values('>H', stream)
-        encrypted_data_key = stream.read(encrypted_data_key_length)
+        (encrypted_data_key_length,) = unpack_values('>H', tee_stream)
+        encrypted_data_key = tee_stream.read(encrypted_data_key_length)
         encrypted_data_keys.add(EncryptedDataKey(
             key_provider=MasterKeyInfo(
                 provider_id=to_str(key_provider_identifier),
@@ -128,7 +127,7 @@ def deserialize_header(stream):
         ))
     header['encrypted_data_keys'] = encrypted_data_keys
 
-    (content_type_id,) = unpack_values('>B', stream)
+    (content_type_id,) = unpack_values('>B', tee_stream)
     try:
         content_type = ContentType(content_type_id)
     except ValueError as error:
@@ -138,14 +137,14 @@ def deserialize_header(stream):
         )
     header['content_type'] = content_type
 
-    (content_aad_length,) = unpack_values('>I', stream)
+    (content_aad_length,) = unpack_values('>I', tee_stream)
     if content_aad_length != 0:
         raise SerializationError(
             'Content AAD length field is currently unused, its value must be always 0'
         )
     header['content_aad_length'] = 0
 
-    (iv_length,) = unpack_values('>B', stream)
+    (iv_length,) = unpack_values('>B', tee_stream)
     if iv_length != alg.iv_len:
         raise SerializationError(
             'Specified IV length ({length}) does not match algorithm IV length ({alg})'.format(
@@ -155,7 +154,7 @@ def deserialize_header(stream):
         )
     header['header_iv_length'] = iv_length
 
-    (frame_length,) = unpack_values('>I', stream)
+    (frame_length,) = unpack_values('>I', tee_stream)
     if content_type == ContentType.FRAMED_DATA and frame_length > MAX_FRAME_SIZE:
         raise SerializationError('Specified frame length larger than allowed maximum: {found} > {max}'.format(
             found=frame_length,
@@ -165,7 +164,7 @@ def deserialize_header(stream):
         raise SerializationError('Non-zero frame length found for non-framed message')
     header['frame_length'] = frame_length
 
-    return MessageHeader(**header)
+    return MessageHeader(**header), tee.getvalue()
 
 
 def deserialize_header_auth(stream, algorithm, verifier=None):

src/aws_encryption_sdk/internal/utils.py → src/aws_encryption_sdk/internal/utils/__init__.py

@@ -17,10 +17,7 @@
 
 import six
 
-from aws_encryption_sdk.exceptions import (
-    ActionNotAllowedError, InvalidDataKeyError,
-    SerializationError, UnknownIdentityError
-)
+from aws_encryption_sdk.exceptions import InvalidDataKeyError, SerializationError, UnknownIdentityError
 from aws_encryption_sdk.identifiers import ContentAADString, ContentType
 import aws_encryption_sdk.internal.defaults
 from aws_encryption_sdk.internal.str_ops import to_bytes
@@ -95,37 +92,6 @@ def get_aad_content_string(content_type, is_final_frame):
     return aad_content_string
 
 
-class ROStream(object):
-    """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 __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))
-
-    def write(self, b):  # pylint: disable=unused-argument
-        """Blocks calls to write.
-
-        :raises ActionNotAllowedError: when called
-        """
-        raise ActionNotAllowedError('Write not allowed on ROStream objects')
-
-
 def prepare_data_keys(primary_master_key, master_keys, algorithm, encryption_context):
     """Prepares a DataKey to be used for encrypting message and list
     of EncryptedDataKey objects to be serialized into header.

src/aws_encryption_sdk/internal/utils/streams.py

@@ -0,0 +1,58 @@
+# 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."""
+from wrapt import ObjectProxy
+
+from aws_encryption_sdk.exceptions import ActionNotAllowedError
+
+
+class ROStream(ObjectProxy):
+    """Provides a read-only interface on top of a file-like object.
+
+    Used to provide MasterKeyProviders with read-only access to plaintext.
+
+    :param wrapped: 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(ObjectProxy):
+    """Provides a ``tee``-like interface on top of a file-like object, which collects read bytes
+    into a local :class:`io.BytesIO`.
+
+    :param wrapped: File-like object
+    :param tee: Stream to copy read bytes into.
+    :type tee: io.BaseIO
+    """
+
+    __tee = None  # Prime ObjectProxy's attributes to allow setting in init.
+
+    def __init__(self, wrapped, tee):
+        """Creates the local tee stream."""
+        super(TeeStream, self).__init__(wrapped)
+        self.__tee = tee
+
+    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.__wrapped__.read(b)
+        self.__tee.write(data)
+        return data

src/aws_encryption_sdk/key_providers/base.py

@@ -110,7 +110,7 @@ def master_keys_for_encryption(self, encryption_context, plaintext_rostream, pla
 
         :param dict encryption_context: Encryption context passed to client
         :param plaintext_rostream: Source plaintext read-only stream
-        :type plaintext_rostream: aws_encryption_sdk.internal.utils.ROStream
+        :type plaintext_rostream: aws_encryption_sdk.internal.utils.streams.ROStream
         :param int plaintext_length: Length of source plaintext (optional)
         :returns: Tuple containing Primary Master Key and List of all Master Keys added to
             this Provider and any member Providers
@@ -387,7 +387,7 @@ def master_keys_for_encryption(self, encryption_context, plaintext_rostream, pla
 
         :param dict encryption_context: Encryption context passed to client
         :param plaintext_rostream: Source plaintext read-only stream
-        :type plaintext_rostream: aws_encryption_sdk.internal.utils.ROStream
+        :type plaintext_rostream: aws_encryption_sdk.internal.utils.streams.ROStream
         :param int plaintext_length: Length of source plaintext (optional)
         :returns: Tuple containing self and a list of self
         :rtype: tuple containing :class:`aws_encryption_sdk.key_providers.base.MasterKey`

src/aws_encryption_sdk/materials_managers/__init__.py

@@ -18,7 +18,7 @@
 import six
 
 from ..identifiers import Algorithm
-from ..internal.utils import ROStream
+from ..internal.utils.streams import ROStream
 from ..structures import DataKey
 
 
@@ -34,7 +34,7 @@ class EncryptionMaterialsRequest(object):
     :param dict encryption_context: Encryption context passed to underlying master key provider and master keys
     :param int frame_length: Frame length to be used while encrypting stream
     :param plaintext_rostream: Source plaintext read-only stream (optional)
-    :type plaintext_rostream: aws_encryption_sdk.internal.utils.ROStream
+    :type plaintext_rostream: aws_encryption_sdk.internal.utils.streams.ROStream
     :param algorithm: Algorithm passed to underlying master key provider and master keys (optional)
     :type algorithm: aws_encryption_sdk.identifiers.Algorithm
     :param int plaintext_length: Length of source plaintext (optional)