fix merge

Author: mattsb42-aws

src/aws_encryption_sdk/streaming_client.py

@@ -209,16 +209,20 @@ def read(self, b=-1):
         :returns: Processed (encrypted or decrypted) bytes from source stream
         :rtype: bytes
         """
-        # NoneType value for b is interpretted as a full read for legacy compatibility
-        if b is None:
+        # Any negative value for b is interpreted as a full read
+        # None is also accepted for legacy compatibility
+        if b is None or b < 0:
             b = -1
 
         _LOGGER.debug("Stream read called, requesting %s bytes", b)
         output = io.BytesIO()
+
         if not self._message_prepped:
             self._prep_message()
+
         if self.closed:
             raise ValueError("I/O operation on closed file")
+
         if b:
             self._read_bytes(b)
             output.write(self.output_buffer[:b])
@@ -228,6 +232,7 @@ def read(self, b=-1):
                 self._read_bytes(LINE_LENGTH)
                 output.write(self.output_buffer)
                 self.output_buffer = b""
+
         self.bytes_read += output.tell()
         _LOGGER.debug("Returning %s bytes of %s bytes requested", output.tell(), b)
         return output.getvalue()
@@ -511,14 +516,18 @@ def _read_bytes_to_non_framed_body(self, b):
             _LOGGER.debug("Closing encryptor after receiving only %s bytes of %s bytes requested", plaintext, b)
             self.source_stream.close()
             closing = self.encryptor.finalize()
+
             if self.signer is not None:
                 self.signer.update(closing)
+
             closing += aws_encryption_sdk.internal.formatting.serialize.serialize_non_framed_close(
                 tag=self.encryptor.tag, signer=self.signer
             )
+
             if self.signer is not None:
                 closing += aws_encryption_sdk.internal.formatting.serialize.serialize_footer(self.signer)
             return ciphertext + closing
+
         return ciphertext
 
     def _read_bytes_to_framed_body(self, b):
@@ -535,9 +544,11 @@ def _read_bytes_to_framed_body(self, b):
         plaintext = self.source_stream.read(b)
         _LOGGER.debug("%s bytes read from source", len(plaintext))
         finalize = False
+
         if len(plaintext) < b:
             _LOGGER.debug("Final plaintext read from source")
             finalize = True
+
         output = b""
         final_frame_written = False
 
@@ -776,10 +787,13 @@ def _read_bytes_from_non_framed_body(self, b):
         bytes_to_read = self.body_end - self.source_stream.tell()
         _LOGGER.debug("%s bytes requested; reading %s bytes", b, bytes_to_read)
         ciphertext = self.source_stream.read(bytes_to_read)
+
         if len(self.output_buffer) + len(ciphertext) < self.body_length:
             raise SerializationError("Total message body contents less than specified in body description")
+
         if self.verifier is not None:
             self.verifier.update(ciphertext)
+
         plaintext = self.decryptor.update(ciphertext)
         plaintext += self.decryptor.finalize()
         aws_encryption_sdk.internal.formatting.deserialize.update_verifier_with_tag(

test/functional/test_f_aws_encryption_sdk_client.py

@@ -598,3 +598,75 @@ def test_stream_decryptor_readable():
         assert handler.readable()
         handler.read()
     assert not handler.readable()
+
+
+def exact_length_plaintext(length):
+    plaintext = b""
+    while len(plaintext) < length:
+        plaintext += VALUES["plaintext_128"]
+    return plaintext[:length]
+
+
+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)
+
+
+@pytest.mark.parametrize(
+    "frame_length",
+    (
+        0,  # 0: unframed
+        128,  # 128: framed with exact final frame size match
+        256,  # 256: framed with inexact final frame size match
+    ),
+)
+def test_incomplete_read_stream_cycle(caplog, frame_length):
+    caplog.set_level(logging.DEBUG)
+    chunk_size = 21  # Will never be an exact match for the frame size
+    key_provider = fake_kms_key_provider()
+
+    plaintext = exact_length_plaintext(384)
+    ciphertext = b""
+    cycle_count = 0
+    with aws_encryption_sdk.stream(
+        mode="encrypt",
+        source=SometimesIncompleteReaderIO(plaintext),
+        key_provider=key_provider,
+        frame_length=frame_length,
+    ) as encryptor:
+        while True:
+            cycle_count += 1
+            chunk = encryptor.read(chunk_size)
+            if not chunk:
+                break
+            ciphertext += chunk
+            if cycle_count > len(VALUES["plaintext_128"]):
+                raise aws_encryption_sdk.exceptions.AWSEncryptionSDKClientError(
+                    "Unexpected error encrypting message: infinite loop detected."
+                )
+
+    decrypted = b""
+    cycle_count = 0
+    with aws_encryption_sdk.stream(
+        mode="decrypt", source=SometimesIncompleteReaderIO(ciphertext), key_provider=key_provider
+    ) as decryptor:
+        while True:
+            cycle_count += 1
+            chunk = decryptor.read(chunk_size)
+            if not chunk:
+                break
+            decrypted += chunk
+            if cycle_count > len(VALUES["plaintext_128"]):
+                raise aws_encryption_sdk.exceptions.AWSEncryptionSDKClientError(
+                    "Unexpected error encrypting message: infinite loop detected."
+                )
+
+    assert ciphertext != decrypted == VALUES["plaintext_128"]