remove requirement for source_stream.tell() on encrypt

Author: mattsb42-aws

src/aws_encryption_sdk/streaming_client.py

@@ -399,6 +399,8 @@ def __init__(self, **kwargs):  # pylint: disable=unused-argument,super-init-not-
         ):
             raise SerializationError("Source too large for non-framed message")
 
+        self.__unframed_plaintext_cache = io.BytesIO()
+
     def ciphertext_length(self):
         """Returns the length of the resulting ciphertext message in bytes.
 
@@ -486,14 +488,25 @@ def _write_header(self):
 
     def _prep_non_framed(self):
         """Prepare the opening data for a non-framed message."""
+        try:
+            plaintext_length = self.stream_length
+            self.__unframed_plaintext_cache = self.source_stream
+        except NotSupportedError:
+            # We need to know the plaintext length before we can start processing the data.
+            # If we cannot seek on the source then we need to read the entire source into memory.
+            self.__unframed_plaintext_cache = io.BytesIO()
+            self.__unframed_plaintext_cache.write(self.source_stream.read())
+            plaintext_length = self.__unframed_plaintext_cache.tell()
+            self.__unframed_plaintext_cache.seek(0)
+
         aad_content_string = aws_encryption_sdk.internal.utils.get_aad_content_string(
             content_type=self.content_type, is_final_frame=True
         )
         associated_data = assemble_content_aad(
             message_id=self._header.message_id,
             aad_content_string=aad_content_string,
             seq_num=1,
-            length=self.stream_length,
+            length=plaintext_length,
         )
         self.encryptor = Encryptor(
             algorithm=self._encryption_materials.algorithm,
@@ -504,7 +517,7 @@ def _prep_non_framed(self):
         self.output_buffer += serialize_non_framed_open(
             algorithm=self._encryption_materials.algorithm,
             iv=self.encryptor.iv,
-            plaintext_length=self.stream_length,
+            plaintext_length=plaintext_length,
             signer=self.signer,
         )
 
@@ -516,7 +529,7 @@ def _read_bytes_to_non_framed_body(self, b):
         :rtype: bytes
         """
         _LOGGER.debug("Reading %d bytes", b)
-        plaintext = self.source_stream.read(b)
+        plaintext = self.__unframed_plaintext_cache.read(b)
         plaintext_length = len(plaintext)
         if self.tell() + len(plaintext) > MAX_NON_FRAMED_SIZE:
             raise SerializationError("Source too large for non-framed message")
@@ -529,6 +542,7 @@ def _read_bytes_to_non_framed_body(self, b):
         if len(plaintext) < b:
             _LOGGER.debug("Closing encryptor after receiving only %d bytes of %d bytes requested", plaintext_length, b)
             self.source_stream.close()
+            self.__unframed_plaintext_cache.close()
             closing = self.encryptor.finalize()
 
             if self.signer is not None:

test/functional/test_f_aws_encryption_sdk_client.py

@@ -771,7 +771,7 @@ def close(self):
 @pytest.mark.parametrize(
     "wrapping_class",
     (
-        pytest.param(NoTell, marks=pytest.mark.xfail),
+        NoTell,
         pytest.param(NoClose, marks=pytest.mark.xfail(strict=True)),
         pytest.param(NothingButRead, marks=pytest.mark.xfail(strict=True)),
     ),
@@ -792,7 +792,7 @@ def test_cycle_minimal_source_stream_api(frame_length, wrapping_class):
 @pytest.mark.parametrize(
     "wrapping_class",
     (
-        pytest.param(NoTell, marks=pytest.mark.xfail),
+        NoTell,
         pytest.param(NoClose, marks=pytest.mark.xfail(strict=True)),
         pytest.param(NothingButRead, marks=pytest.mark.xfail(strict=True)),
     ),

test/unit/test_streaming_client_stream_encryptor.py

@@ -382,7 +382,10 @@ def test_read_bytes_to_non_framed_body(self):
         test_encryptor.encryptor = MagicMock()
         test_encryptor._encryption_materials = self.mock_encryption_materials
         test_encryptor.encryptor.update.return_value = sentinel.ciphertext
+        test_encryptor._StreamEncryptor__unframed_plaintext_cache = pt_stream
+
         test = test_encryptor._read_bytes_to_non_framed_body(5)
+
         test_encryptor.encryptor.update.assert_called_once_with(self.plaintext[:5])
         test_encryptor.signer.update.assert_called_once_with(sentinel.ciphertext)
         assert not test_encryptor.source_stream.closed
@@ -392,6 +395,8 @@ def test_read_bytes_to_non_framed_body_too_large(self):
         pt_stream = io.BytesIO(self.plaintext)
         test_encryptor = StreamEncryptor(source=pt_stream, key_provider=self.mock_key_provider)
         test_encryptor.bytes_read = aws_encryption_sdk.internal.defaults.MAX_NON_FRAMED_SIZE
+        test_encryptor._StreamEncryptor__unframed_plaintext_cache = pt_stream
+
         with six.assertRaisesRegex(self, SerializationError, "Source too large for non-framed message"):
             test_encryptor._read_bytes_to_non_framed_body(5)