remove source_stream.closed from StreamEncryptor._read_bytes()

Author: mattsb42-aws

src/aws_encryption_sdk/streaming_client.py

@@ -404,6 +404,7 @@ 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()
+        self.__message_complete = False
 
     def ciphertext_length(self):
         """Returns the length of the resulting ciphertext message in bytes.
@@ -556,6 +557,7 @@ def _read_bytes_to_non_framed_body(self, b):
 
             if self.signer is not None:
                 closing += serialize_footer(self.signer)
+            self.__message_complete = True
             return ciphertext + closing
 
         return ciphertext
@@ -622,6 +624,7 @@ def _read_bytes_to_framed_body(self, b):
             _LOGGER.debug("Writing footer")
             if self.signer is not None:
                 output += serialize_footer(self.signer)
+            self.__message_complete = True
             self.source_stream.close()
         return output
 
@@ -632,7 +635,7 @@ def _read_bytes(self, b):
         :raises NotSupportedError: if content type is not supported
         """
         _LOGGER.debug("%d bytes requested from stream with content type: %s", b, self.content_type)
-        if 0 <= b <= len(self.output_buffer) or self.source_stream.closed:
+        if 0 <= b <= len(self.output_buffer) or self.__message_complete:
             _LOGGER.debug("No need to read from source stream or source stream closed")
             return
 

test/functional/test_f_aws_encryption_sdk_client.py

@@ -761,9 +761,11 @@ def tell(self):
         raise NotImplementedError("NoTell does not tell().")
 
 
-class NoClose(ObjectProxy):
+class NoClosed(ObjectProxy):
     closed = NotImplemented
 
+
+class NoClose(ObjectProxy):
     def close(self):
         raise NotImplementedError("NoClose does not close().")
 
@@ -772,6 +774,7 @@ def close(self):
     "wrapping_class",
     (
         NoTell,
+        NoClosed,
         pytest.param(NoClose, marks=pytest.mark.xfail(strict=True)),
         pytest.param(NothingButRead, marks=pytest.mark.xfail(strict=True)),
     ),
@@ -793,6 +796,7 @@ def test_cycle_minimal_source_stream_api(frame_length, wrapping_class):
     "wrapping_class",
     (
         NoTell,
+        NoClosed,
         pytest.param(NoClose, marks=pytest.mark.xfail(strict=True)),
         pytest.param(NothingButRead, marks=pytest.mark.xfail(strict=True)),
     ),
@@ -813,6 +817,7 @@ def test_encrypt_minimal_source_stream_api(frame_length, wrapping_class):
     "wrapping_class",
     (
         NoTell,
+        NoClosed,
         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

@@ -451,10 +451,10 @@ def test_read_bytes_less_than_buffer(self, mock_read_non_framed, mock_read_frame
 
     @patch("aws_encryption_sdk.streaming_client.StreamEncryptor._read_bytes_to_framed_body")
     @patch("aws_encryption_sdk.streaming_client.StreamEncryptor._read_bytes_to_non_framed_body")
-    def test_read_bytes_closed(self, mock_read_non_framed, mock_read_framed):
+    def test_read_bytes_completed(self, mock_read_non_framed, mock_read_framed):
         pt_stream = io.BytesIO(self.plaintext)
         test_encryptor = StreamEncryptor(source=pt_stream, key_provider=self.mock_key_provider)
-        test_encryptor.source_stream.close()
+        test_encryptor._StreamEncryptor__message_complete = True
         test_encryptor._read_bytes(5)
         assert not mock_read_non_framed.called
         assert not mock_read_framed.called