unit tests

Author: lucasmcdonald3

src/aws_encryption_sdk/streaming_client.py

@@ -954,8 +954,72 @@ def _prep_message(self):
             self._prep_non_framed()
         self._message_prepped = True
 
-    # TODO-MPL: Refactor this function, remove linter disablers
-    def _read_header(self):  # noqa pylint: disable=too-many-branches
+    def _create_decrypt_materials_request(self, header):
+        """
+        Create a DecryptionMaterialsRequest based on whether
+        the StreamDecryptor was provided encryption_context on decrypt
+        (i.e. expects to use required encryption context CMM from the MPL).
+        """
+        # If encryption_context is provided on decrypt,
+        # pass it to the DecryptionMaterialsRequest as reproduced_encryption_context
+        if hasattr(self.config, "encryption_context"):
+            return DecryptionMaterialsRequest(
+                encrypted_data_keys=header.encrypted_data_keys,
+                algorithm=header.algorithm,
+                encryption_context=header.encryption_context,
+                commitment_policy=self.config.commitment_policy,
+                reproduced_encryption_context=self.config.encryption_context
+            )
+        return DecryptionMaterialsRequest(
+            encrypted_data_keys=header.encrypted_data_keys,
+            algorithm=header.algorithm,
+            encryption_context=header.encryption_context,
+            commitment_policy=self.config.commitment_policy,
+        )
+    
+    def _validate_parsed_header(
+        self,
+        header,
+        header_auth,
+        raw_header,
+    ):
+        """
+        Pass arguments from this StreamDecryptor to validate_header based on whether
+        the StreamDecryptor has the _required_encryption_context attribute
+        (i.e. is using the required encryption context CMM from the MPL).
+        """
+        # If _required_encryption_context is present,
+        # serialize it and pass it to validate_header.
+        if hasattr(self, "_required_encryption_context") \
+                and self._required_encryption_context is not None:
+            # The authenticated only encryption context is all encryption context key-value pairs where the
+            # key exists in Required Encryption Context Keys. It is then serialized according to the
+            # message header Key Value Pairs.
+            required_ec_serialized = \
+                aws_encryption_sdk.internal.formatting.encryption_context.serialize_encryption_context(
+                    self._required_encryption_context
+                )
+
+            validate_header(
+                header=header,
+                header_auth=header_auth,
+                # When verifying the header, the AAD input to the authenticated encryption algorithm
+                # specified by the algorithm suite is the message header body and the serialized
+                # authenticated only encryption context.
+                raw_header=raw_header + required_ec_serialized,
+                data_key=self._derived_data_key
+            )
+        else:
+            validate_header(
+                header=header,
+                header_auth=header_auth,
+                raw_header=raw_header,
+                data_key=self._derived_data_key
+            )
+
+        return header, header_auth
+
+    def _read_header(self):
         """Reads the message header from the input stream.
 
         :returns: tuple containing deserialized header and header_auth objects
@@ -981,24 +1045,7 @@ def _read_header(self):  # noqa pylint: disable=too-many-branches
                 )
             )
 
-        # If encryption_context is provided on decrypt,
-        # pass it to the DecryptionMaterialsRequest
-        if hasattr(self.config, "encryption_context"):
-            decrypt_materials_request = DecryptionMaterialsRequest(
-                encrypted_data_keys=header.encrypted_data_keys,
-                algorithm=header.algorithm,
-                encryption_context=header.encryption_context,
-                commitment_policy=self.config.commitment_policy,
-                reproduced_encryption_context=self.config.encryption_context
-            )
-        else:
-            decrypt_materials_request = DecryptionMaterialsRequest(
-                encrypted_data_keys=header.encrypted_data_keys,
-                algorithm=header.algorithm,
-                encryption_context=header.encryption_context,
-                commitment_policy=self.config.commitment_policy,
-            )
-
+        decrypt_materials_request = self._create_decrypt_materials_request(header)
         decryption_materials = self.config.materials_manager.decrypt_materials(request=decrypt_materials_request)
 
         # If the materials_manager passed required_encryption_context_keys,
@@ -1049,36 +1096,12 @@ def _read_header(self):  # noqa pylint: disable=too-many-branches
                     "Key commitment validation failed. Key identity does not match the identity asserted in the "
                     "message. Halting processing of this message."
                 )
-
-        # If _required_encryption_context is present,
-        # serialize it and pass it to validate_header.
-        if self._required_encryption_context is not None:
-            # The authenticated only encryption context is all encryption context key-value pairs where the
-            # key exists in Required Encryption Context Keys. It is then serialized according to the
-            # message header Key Value Pairs.
-            required_ec_serialized = \
-                aws_encryption_sdk.internal.formatting.encryption_context.serialize_encryption_context(
-                    self._required_encryption_context
-                )
-
-            validate_header(
-                header=header,
-                header_auth=header_auth,
-                # When verifying the header, the AAD input to the authenticated encryption algorithm
-                # specified by the algorithm suite is the message header body and the serialized
-                # authenticated only encryption context.
-                raw_header=raw_header + required_ec_serialized,
-                data_key=self._derived_data_key
-            )
-        else:
-            validate_header(
-                header=header,
-                header_auth=header_auth,
-                raw_header=raw_header,
-                data_key=self._derived_data_key
-            )
-
-        return header, header_auth
+            
+        return self._validate_parsed_header(
+            header=header,
+            header_auth=header_auth,
+            raw_header=raw_header,
+        )
 
     def _prep_non_framed(self):
         """Prepare the opening data for a non-framed message."""

test/mpl/unit/test_material_managers_mpl_cmm.py

@@ -38,6 +38,7 @@
 mock_mpl_cmm = MagicMock(__class__=MPL_ICryptographicMaterialsManager)
 mock_mpl_encryption_materials = MagicMock(__class__=MPL_EncryptionMaterials)
 mock_mpl_decrypt_materials = MagicMock(__class__=MPL_DecryptionMaterials)
+mock_reproduced_encryption_context = MagicMock(__class_=dict)
 
 
 mock_edk = MagicMock(__class__=Native_EncryptedDataKey)
@@ -259,6 +260,7 @@ def test_GIVEN_valid_request_WHEN_create_mpl_decrypt_materials_input_from_reques
     for mock_edks in [no_mock_edks, one_mock_edk, two_mock_edks]:
 
         mock_decryption_materials_request.encrypted_data_keys = mock_edks
+        mock_decryption_materials_request.reproduced_encryption_context = mock_reproduced_encryption_context
 
         # When: _create_mpl_decrypt_materials_input_from_request
         output = CryptoMaterialsManagerFromMPL._create_mpl_decrypt_materials_input_from_request(
@@ -271,6 +273,7 @@ def test_GIVEN_valid_request_WHEN_create_mpl_decrypt_materials_input_from_reques
         assert output.algorithm_suite_id == mock_algorithm_id
         assert output.commitment_policy == mock_commitment_policy
         assert output.encryption_context == mock_decryption_materials_request.encryption_context
+        assert output.reproduced_encryption_context == mock_reproduced_encryption_context
 
         assert len(output.encrypted_data_keys) == len(mock_edks)
         for i in range(len(output.encrypted_data_keys)):

test/mpl/unit/test_material_managers_mpl_materials.py

@@ -160,6 +160,19 @@ def test_GIVEN_valid_signing_key_WHEN_EncryptionMaterials_get_signing_key_THEN_r
     assert output == mock_signing_key
 
 
+def test_GIVEN_valid_required_encryption_context_keys_WHEN_EncryptionMaterials_get_required_encryption_context_keys_THEN_returns_required_encryption_context_keys():
+    # Given: valid required encryption context keys
+    mock_required_encryption_context_keys = MagicMock(__class__=bytes)
+    mock_mpl_encryption_materials.required_encryption_context_keys = mock_required_encryption_context_keys
+
+    # When: get required encryption context keys
+    mpl_encryption_materials = EncryptionMaterialsFromMPL(mpl_materials=mock_mpl_encryption_materials)
+    output = mpl_encryption_materials.required_encryption_context_keys
+
+    # Then: returns required encryption context keys
+    assert output == mock_required_encryption_context_keys
+
+
 def test_GIVEN_valid_data_key_WHEN_DecryptionMaterials_get_data_key_THEN_returns_data_key():
     # Given: valid MPL data key
     mock_data_key = MagicMock(__class__=bytes)
@@ -187,3 +200,29 @@ def test_GIVEN_valid_verification_key_WHEN_DecryptionMaterials_get_verification_
 
     # Then: returns verification key
     assert output == mock_verification_key
+
+
+def test_GIVEN_valid_encryption_context_WHEN_DecryptionMaterials_get_encryption_context_THEN_returns_encryption_context():
+    # Given: valid encryption context
+    mock_encryption_context = MagicMock(__class__=Dict[str, str])
+    mock_mpl_decrypt_materials.encryption_context = mock_encryption_context
+
+    # When: get encryption context
+    mpl_decryption_materials = DecryptionMaterialsFromMPL(mpl_materials=mock_mpl_decrypt_materials)
+    output = mpl_decryption_materials.encryption_context
+
+    # Then: returns valid encryption context
+    assert output == mock_encryption_context
+
+
+def test_GIVEN_valid_required_encryption_context_keys_WHEN_DecryptionMaterials_get_required_encryption_context_keys_THEN_returns_required_encryption_context_keys():
+    # Given: valid required encryption context keys
+    mock_required_encryption_context_keys = MagicMock(__class__=bytes)
+    mock_mpl_decrypt_materials.required_encryption_context_keys = mock_required_encryption_context_keys
+
+    # When: get required encryption context keys
+    mpl_decryption_materials = DecryptionMaterialsFromMPL(mpl_materials=mock_mpl_decrypt_materials)
+    output = mpl_decryption_materials.required_encryption_context_keys
+
+    # Then: returns required encryption context keys
+    assert output == mock_required_encryption_context_keys

test/unit/test_serialize.py

@@ -79,6 +79,7 @@ def apply_fixtures(self):
             "aws_encryption_sdk.internal.formatting.serialize.aws_encryption_sdk.internal.utils.validate_frame_length"
         )
         self.mock_valid_frame_length = self.mock_valid_frame_length_patcher.start()
+        self.mock_required_ec_bytes = MagicMock()
         # Set up mock signer
         self.mock_signer = MagicMock()
         self.mock_signer.update.return_value = None
@@ -167,6 +168,31 @@ def test_serialize_header_auth_v1_no_signer(self):
             data_encryption_key=VALUES["data_key_obj"],
         )
 
+    @patch("aws_encryption_sdk.internal.formatting.serialize.header_auth_iv")
+    def test_GIVEN_required_ec_bytes_WHEN_serialize_header_auth_v1_THEN_aad_has_required_ec_bytes(self, mock_header_auth_iv):
+        """Validate that the _create_header_auth function
+        behaves as expected for SerializationVersion.V1
+        when required_ec_bytes are provided.
+        """
+        self.mock_encrypt.return_value = VALUES["header_auth_base"]
+        test = aws_encryption_sdk.internal.formatting.serialize.serialize_header_auth(
+            version=SerializationVersion.V1,
+            algorithm=self.mock_algorithm,
+            header=VALUES["serialized_header"],
+            data_encryption_key=sentinel.encryption_key,
+            signer=self.mock_signer,
+            required_ec_bytes=self.mock_required_ec_bytes,
+        )
+        self.mock_encrypt.assert_called_once_with(
+            algorithm=self.mock_algorithm,
+            key=sentinel.encryption_key,
+            plaintext=b"",
+            associated_data=VALUES["serialized_header"] + self.mock_required_ec_bytes,
+            iv=mock_header_auth_iv.return_value,
+        )
+        self.mock_signer.update.assert_called_once_with(VALUES["serialized_header_auth"])
+        assert test == VALUES["serialized_header_auth"]
+
     @patch("aws_encryption_sdk.internal.formatting.serialize.header_auth_iv")
     def test_serialize_header_auth_v2(self, mock_header_auth_iv):
         """Validate that the _create_header_auth function
@@ -203,6 +229,30 @@ def test_serialize_header_auth_v2_no_signer(self):
             data_encryption_key=VALUES["data_key_obj"],
         )
 
+    @patch("aws_encryption_sdk.internal.formatting.serialize.header_auth_iv")
+    def test_GIVEN_required_ec_bytes_WHEN_serialize_header_auth_v2_THEN_aad_has_required_ec_bytes(self, mock_header_auth_iv):
+        """Validate that the _create_header_auth function
+        behaves as expected for SerializationVersion.V2.
+        """
+        self.mock_encrypt.return_value = VALUES["header_auth_base"]
+        test = aws_encryption_sdk.internal.formatting.serialize.serialize_header_auth(
+            version=SerializationVersion.V2,
+            algorithm=self.mock_algorithm,
+            header=VALUES["serialized_header_v2_committing"],
+            data_encryption_key=sentinel.encryption_key,
+            signer=self.mock_signer,
+            required_ec_bytes=self.mock_required_ec_bytes,
+        )
+        self.mock_encrypt.assert_called_once_with(
+            algorithm=self.mock_algorithm,
+            key=sentinel.encryption_key,
+            plaintext=b"",
+            associated_data=VALUES["serialized_header_v2_committing"] + self.mock_required_ec_bytes,
+            iv=mock_header_auth_iv.return_value,
+        )
+        self.mock_signer.update.assert_called_once_with(VALUES["serialized_header_auth_v2"])
+        assert test == VALUES["serialized_header_auth_v2"]
+
     def test_serialize_non_framed_open(self):
         """Validate that the serialize_non_framed_open
         function behaves as expected.

test/unit/test_streaming_client_configs.py

@@ -15,7 +15,7 @@
 
 import pytest
 import six
-from mock import patch
+from mock import MagicMock, patch
 
 from aws_encryption_sdk import CommitmentPolicy
 from aws_encryption_sdk.internal.defaults import ALGORITHM, FRAME_LENGTH, LINE_LENGTH
@@ -33,7 +33,10 @@
 # Ideally, this logic would be based on mocking imports and testing logic,
 # but doing that introduces errors that cause other tests to fail.
 try:
-    from aws_cryptographic_materialproviders.mpl.references import IKeyring
+    from aws_cryptographic_materialproviders.mpl.references import (
+        ICryptographicMaterialsManager,
+        IKeyring,
+    )
     HAS_MPL = True
 
     from aws_encryption_sdk.materials_managers.mpl.cmm import CryptoMaterialsManagerFromMPL
@@ -236,24 +239,21 @@ def test_client_configs_with_mpl(
     assert test.materials_manager is not None
 
     # If materials manager was provided, it should be directly used
-    if hasattr(kwargs, "materials_manager"):
+    if "materials_manager" in kwargs:
         assert kwargs["materials_manager"] == test.materials_manager
 
-    # If MPL keyring was provided, it should be wrapped in MPL materials manager
-    if hasattr(kwargs, "keyring"):
-        assert test.keyring is not None
-        assert test.keyring == kwargs["keyring"]
-        assert isinstance(test.keyring, IKeyring)
-        assert isinstance(test.materials_manager, CryptoMaterialsManagerFromMPL)
-
     # If native key_provider was provided, it should be wrapped in native materials manager
-    if hasattr(kwargs, "key_provider"):
+    elif "key_provider" in kwargs:
         assert test.key_provider is not None
         assert test.key_provider == kwargs["key_provider"]
         assert isinstance(test.materials_manager, DefaultCryptoMaterialsManager)
 
+    else:
+        raise ValueError(f"Test did not find materials_manager or key_provider. {kwargs}")
+
 
-# This needs its own test; pytest parametrize cannot use a conditionally-loaded type
+# This is an addition to test_client_configs_with_mpl;
+# This needs its own test; pytest's parametrize cannot use a conditionally-loaded type (IKeyring)
 @pytest.mark.skipif(not HAS_MPL, reason="Test should only be executed with MPL in installation")
 def test_keyring_client_config_with_mpl(
 ):
@@ -265,16 +265,30 @@ def test_keyring_client_config_with_mpl(
 
     test = _ClientConfig(**kwargs)
 
-    # In all cases, config should have a materials manager
     assert test.materials_manager is not None
 
-    # If materials manager was provided, it should be directly used
-    if hasattr(kwargs, "materials_manager"):
-        assert kwargs["materials_manager"] == test.materials_manager
+    assert test.keyring is not None
+    assert test.keyring == kwargs["keyring"]
+    assert isinstance(test.keyring, IKeyring)
+    assert isinstance(test.materials_manager, CryptoMaterialsManagerFromMPL)
+
+
+# This is an addition to test_client_configs_with_mpl;
+# This needs its own test; pytest's parametrize cannot use a conditionally-loaded type (MPL CMM)
+@pytest.mark.skipif(not HAS_MPL, reason="Test should only be executed with MPL in installation")
+def test_mpl_cmm_client_config_with_mpl(
+):
+    mock_mpl_cmm = MagicMock(__class__=ICryptographicMaterialsManager)
+    kwargs = {
+        "source": b"",
+        "materials_manager": mock_mpl_cmm,
+        "commitment_policy": CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
+    }
+
+    test = _ClientConfig(**kwargs)
 
-    # If MPL keyring was provided, it should be wrapped in MPL materials manager
-    if hasattr(kwargs, "keyring"):
-        assert test.keyring is not None
-        assert test.keyring == kwargs["keyring"]
-        assert isinstance(test.keyring, IKeyring)
-        assert isinstance(test.materials_manager, CryptoMaterialsManagerFromMPL)
+    assert test.materials_manager is not None
+    # Assert that the MPL CMM is wrapped in the native interface
+    assert isinstance(test.materials_manager, CryptoMaterialsManagerFromMPL)
+    # Assert the MPL CMM is used by the native interface
+    assert test.materials_manager.mpl_cmm == mock_mpl_cmm

test/unit/test_streaming_client_stream_decryptor.py

@@ -924,3 +924,213 @@ def test_close_no_footer(self, mock_close):
         with pytest.raises(SerializationError) as excinfo:
             test_decryptor.close()
         excinfo.match("Footer not read")
+
+    @patch("aws_encryption_sdk.streaming_client.validate_header")
+    def test_GIVEN_does_not_have_required_EC_WHEN_validate_parsed_header_THEN_validate_header(
+        self,
+        mock_validate_header
+    ):
+        self.mock_header.content_type = ContentType.FRAMED_DATA
+        test_decryptor = StreamDecryptor(
+            materials_manager=self.mock_materials_manager,
+            source=self.mock_input_stream,
+            commitment_policy=self.mock_commitment_policy,
+        )
+        test_decryptor._derived_data_key = sentinel.derived_data_key
+        # Given: test_decryptor does not have _required_encryption_context attribute
+        # When: _validate_parsed_header
+        test_decryptor._validate_parsed_header(
+            header=self.mock_header,
+            header_auth=sentinel.header_auth,
+            raw_header=self.mock_raw_header
+        )
+        # Then: validate_header
+        mock_validate_header.assert_called_once_with(
+            header=self.mock_header,
+            header_auth=sentinel.header_auth,
+            raw_header=self.mock_raw_header,
+            data_key=sentinel.derived_data_key,
+        )
+
+    @patch("aws_encryption_sdk.internal.formatting.encryption_context.serialize_encryption_context")
+    @patch("aws_encryption_sdk.streaming_client.validate_header")
+    def test_GIVEN_has_required_EC_WHEN_validate_parsed_header_THEN_validate_header_with_serialized_required_EC(
+        self,
+        mock_validate_header,
+        mock_serialize_encryption_context,
+    ):
+        self.mock_header.content_type = ContentType.FRAMED_DATA
+        test_decryptor = StreamDecryptor(
+            materials_manager=self.mock_materials_manager,
+            source=self.mock_input_stream,
+            commitment_policy=self.mock_commitment_policy,
+        )
+        test_decryptor._derived_data_key = sentinel.derived_data_key
+        # Given: test_decryptor has _required_encryption_context attribute
+        mock_required_ec = MagicMock(__class__=dict)
+        test_decryptor._required_encryption_context = mock_required_ec
+        mock_serialized_required_ec = MagicMock(__class__=bytes)
+        mock_serialize_encryption_context.return_value = mock_serialized_required_ec
+        # When: _validate_parsed_header
+        test_decryptor._validate_parsed_header(
+            header=self.mock_header,
+            header_auth=sentinel.header_auth,
+            raw_header=self.mock_raw_header
+        )
+        # Then: call validate_header with serialized required EC
+        mock_validate_header.assert_called_once_with(
+            header=self.mock_header,
+            header_auth=sentinel.header_auth,
+            raw_header=self.mock_raw_header + mock_serialized_required_ec,
+            data_key=sentinel.derived_data_key,
+        )
+
+    def test_GIVEN_config_has_EC_WHEN_create_decrypt_materials_request_THEN_provide_reproduced_EC(
+        self,
+    ):
+        self.mock_header.content_type = ContentType.FRAMED_DATA
+        test_decryptor = StreamDecryptor(
+            materials_manager=self.mock_materials_manager,
+            source=self.mock_input_stream,
+            commitment_policy=self.mock_commitment_policy,
+        )
+
+        # Given: StreamDecryptor.config has encryption_context attribute
+        mock_reproduced_encryption_context = MagicMock(__class__=dict)
+        test_decryptor.config.encryption_context = mock_reproduced_encryption_context
+        # Type checking on header encryption context seems to require concrete instance,
+        # neither MagicMock nor sentinel value work
+        self.mock_header.encryption_context = {"some_key_to_pass_type_validation": "some_value"}
+
+        # When: _create_decrypt_materials_request
+        output = test_decryptor._create_decrypt_materials_request(
+            header=self.mock_header,
+        )
+
+        # Then: decrypt_materials_request has reproduced_encryption_context attribute
+        assert hasattr(output, "reproduced_encryption_context")
+        assert output.reproduced_encryption_context == mock_reproduced_encryption_context
+
+    def test_GIVEN_config_does_not_have_EC_WHEN_create_decrypt_materials_request_THEN_request_does_not_have_reproduced_EC(
+        self,
+    ):
+        self.mock_header.content_type = ContentType.FRAMED_DATA
+        test_decryptor = StreamDecryptor(
+            materials_manager=self.mock_materials_manager,
+            source=self.mock_input_stream,
+            commitment_policy=self.mock_commitment_policy,
+        )
+
+        # Given: StreamDecryptor.config does not have an encryption_context attribute
+        del test_decryptor.config.encryption_context
+        # Type checking on header encryption context seems to require concrete instance,
+        # neither MagicMock nor sentinel value work
+        self.mock_header.encryption_context = {"some_key_to_pass_type_validation": "some_value"}
+
+        # When: _create_decrypt_materials_request
+        output = test_decryptor._create_decrypt_materials_request(
+            header=self.mock_header,
+        )
+
+        # Then: decrypt_materials_request.reproduced_encryption_context is None
+        assert output.reproduced_encryption_context is None
+
+    @patch("aws_encryption_sdk.streaming_client.derive_data_encryption_key")
+    @patch("aws_encryption_sdk.streaming_client.DecryptionMaterialsRequest")
+    @patch("aws_encryption_sdk.streaming_client.Verifier")
+    def test_GIVEN_materials_has_no_required_encryption_context_keys_attr_WHEN_read_header_THEN_required_EC_is_None(
+        self,
+        mock_verifier,
+        *_
+    ):
+
+        mock_verifier_instance = MagicMock()
+        mock_verifier.from_key_bytes.return_value = mock_verifier_instance
+
+        self.mock_header.content_type = ContentType.FRAMED_DATA
+        test_decryptor = StreamDecryptor(
+            materials_manager=self.mock_materials_manager,
+            source=self.mock_input_stream,
+            commitment_policy=self.mock_commitment_policy,
+        )
+
+        # Given: decryption_materials does not have a required_encryption_context_keys attribute
+        del self.mock_decrypt_materials.required_encryption_context_keys
+
+        # When: _read_header
+        test_decryptor._read_header()
+
+        # Then: StreamDecryptor._required_encryption_context is None
+        assert test_decryptor._required_encryption_context is None
+
+    @patch("aws_encryption_sdk.streaming_client.derive_data_encryption_key")
+    @patch("aws_encryption_sdk.streaming_client.DecryptionMaterialsRequest")
+    @patch("aws_encryption_sdk.streaming_client.Verifier")
+    def test_GIVEN_materials_has_required_encryption_context_keys_attr_WHEN_read_header_THEN_creates_correct_required_EC(
+        self,
+        mock_verifier,
+        *_
+    ): 
+        required_encryption_context_keys_values = [
+            # Case of empty encryption context list is not allowed;
+            # if a list is provided, it must be non-empty.
+            # The MPL enforces this behavior on construction.
+            ["one_key"],
+            ["one_key", "two_key"],
+            ["one_key", "two_key", "red_key"],
+            ["one_key", "two_key", "red_key", "blue_key"],
+        ]
+
+        encryption_context_values = [
+            {},
+            {"one_key": "some_value"},
+            {
+                "one_key": "some_value",
+                "two_key": "some_other_value",
+            },
+            {
+                "one_key": "some_value",
+                "two_key": "some_other_value",
+                "red_key": "some_red_value",
+            },
+            {
+                "one_key": "some_value",
+                "two_key": "some_other_value",
+                "red_key": "some_red_value",
+                "blue_key": "some_blue_value",
+            }
+        ]
+
+        for required_encryption_context_keys in required_encryption_context_keys_values:
+
+            # Given: decryption_materials has required_encryption_context_keys
+            self.mock_decrypt_materials.required_encryption_context_keys = \
+                required_encryption_context_keys
+            
+            for encryption_context in encryption_context_values:
+                
+                self.mock_decrypt_materials.encryption_context = encryption_context
+
+                mock_verifier_instance = MagicMock()
+                mock_verifier.from_key_bytes.return_value = mock_verifier_instance
+
+                self.mock_header.content_type = ContentType.FRAMED_DATA
+                test_decryptor = StreamDecryptor(
+                    materials_manager=self.mock_materials_manager,
+                    source=self.mock_input_stream,
+                    commitment_policy=self.mock_commitment_policy,
+                )
+
+                # When: _read_header
+                test_decryptor._read_header()
+
+                # Then: Assert correctness of partitioned EC
+                for k in encryption_context:
+                    # If a key is in required_encryption_context_keys, then ...
+                    if k in required_encryption_context_keys:
+                        # ... its EC is in the StreamEncryptor._required_encryption_context
+                        assert k in test_decryptor._required_encryption_context
+                    # If a key is NOT in required_encryption_context_keys, then ...
+                    else:
+                        # ... its EC is NOT in the StreamEncryptor._required_encryption_context
+                        assert k not in test_decryptor._required_encryption_context

test/unit/test_streaming_client_stream_encryptor.py

@@ -451,6 +451,113 @@ def test_GIVEN_has_mpl_AND_has_MPLCMM_AND_uses_signer_WHEN_prep_message_THEN_sig
                 encoding=serialization.Encoding.PEM
             )
 
+    # Given: has MPL
+    @pytest.mark.skipif(not HAS_MPL, reason="Test should only be executed with MPL in installation")
+    def test_GIVEN_has_mpl_AND_encryption_materials_has_required_EC_keys_WHEN_prep_message_THEN_paritions_stored_and_required_EC(self):
+        # Create explicit values to explicitly test logic in smaller cases
+        required_encryption_context_keys_values = [
+            # Case of empty encryption context list is not allowed;
+            # if a list is provided, it must be non-empty.
+            # The MPL enforces this behavior on construction.
+            ["one_key"],
+            ["one_key", "two_key"],
+            ["one_key", "two_key", "red_key"],
+            ["one_key", "two_key", "red_key", "blue_key"],
+        ]
+
+        encryption_context_values = [
+            {},
+            {"one_key": "some_value"},
+            {
+                "one_key": "some_value",
+                "two_key": "some_other_value",
+            },
+            {
+                "one_key": "some_value",
+                "two_key": "some_other_value",
+                "red_key": "some_red_value",
+            },
+            {
+                "one_key": "some_value",
+                "two_key": "some_other_value",
+                "red_key": "some_red_value",
+                "blue_key": "some_blue_value",
+            }
+        ]
+
+        self.mock_encryption_materials.algorithm = Algorithm.AES_128_GCM_IV12_TAG16
+
+        for required_encryption_context_keys in required_encryption_context_keys_values:
+
+            # Given: encryption context has required_encryption_context_keys
+            self.mock_encryption_materials.required_encryption_context_keys = \
+                required_encryption_context_keys
+            
+            for encryption_context in encryption_context_values:
+                self.mock_encryption_materials.encryption_context = encryption_context
+
+                test_encryptor = StreamEncryptor(
+                    source=VALUES["data_128"],
+                    materials_manager=self.mock_mpl_materials_manager,
+                    frame_length=self.mock_frame_length,
+                    algorithm=Algorithm.AES_128_GCM_IV12_TAG16,
+                    commitment_policy=self.mock_commitment_policy,
+                    signature_policy=self.mock_signature_policy,
+                )
+                test_encryptor.content_type = ContentType.FRAMED_DATA
+                # When: prep_message
+                test_encryptor._prep_message()
+
+                # Then: Assert correctness of partitioned EC
+                for k in encryption_context:
+                    # If a key is in required_encryption_context_keys, then
+                    if k in required_encryption_context_keys:
+                        # 1) Its EC is in the StreamEncryptor._required_encryption_context
+                        assert k in test_encryptor._required_encryption_context
+                        # 2) Its EC is NOT in the StreamEncryptor._stored_encryption_context
+                        assert k not in test_encryptor._stored_encryption_context
+                    # If a key is NOT in required_encryption_context_keys, then
+                    else:
+                        # 1) Its EC is NOT in the StreamEncryptor._required_encryption_context
+                        assert k not in test_encryptor._required_encryption_context
+                        # 2) Its EC is in the StreamEncryptor._stored_encryption_context
+                        assert k in test_encryptor._stored_encryption_context
+                
+                # Assert size(stored_EC) + size(required_EC) == size(EC)
+                # (i.e. every EC was sorted into one or the other)
+                assert len(test_encryptor._required_encryption_context) \
+                    + len(test_encryptor._stored_encryption_context) \
+                    == len(encryption_context)
+    
+    # Given: has MPL
+    @pytest.mark.skipif(not HAS_MPL, reason="Test should only be executed with MPL in installation")
+    def test_GIVEN_has_mpl_AND_encryption_materials_does_not_have_required_EC_keys_WHEN_prep_message_THEN_stored_EC_is_EC(self):
+
+        self.mock_encryption_materials.algorithm = Algorithm.AES_128_GCM_IV12_TAG16
+
+        mock_encryption_context = MagicMock(__class__=dict)
+        self.mock_encryption_materials.encryption_context = mock_encryption_context
+        # Given: encryption materials does not have required encryption context keys
+        # (MagicMock default is to "make up" "Some" value here; this deletes that value)
+        del self.mock_encryption_materials.required_encryption_context_keys
+
+        test_encryptor = StreamEncryptor(
+            source=VALUES["data_128"],
+            materials_manager=self.mock_mpl_materials_manager,
+            frame_length=self.mock_frame_length,
+            algorithm=Algorithm.AES_128_GCM_IV12_TAG16,
+            commitment_policy=self.mock_commitment_policy,
+            signature_policy=self.mock_signature_policy,
+        )
+        test_encryptor.content_type = ContentType.FRAMED_DATA
+        # When: prep_message
+        test_encryptor._prep_message()
+
+        # Then: _stored_encryption_context is the provided encryption_context
+        assert test_encryptor._stored_encryption_context == mock_encryption_context
+        # Then: _required_encryption_context is None
+        assert test_encryptor._required_encryption_context is None
+                
     def test_prep_message_no_signer(self):
         self.mock_encryption_materials.algorithm = Algorithm.AES_128_GCM_IV12_TAG16
         test_encryptor = StreamEncryptor(
@@ -575,6 +682,53 @@ def test_write_header(self):
         )
         assert test_encryptor.output_buffer == b"1234567890"
 
+    @patch("aws_encryption_sdk.internal.formatting.encryption_context.serialize_encryption_context")
+    # Given: has MPL
+    @pytest.mark.skipif(not HAS_MPL, reason="Test should only be executed with MPL in installation")
+    def test_GIVEN_has_mpl_AND_has_required_EC_WHEN_write_header_THEN_adds_serialized_required_ec_to_header_auth(
+        self,
+        serialize_encryption_context
+    ):
+        self.mock_serialize_header.return_value = b"12345"
+        self.mock_serialize_header_auth.return_value = b"67890"
+        pt_stream = io.BytesIO(self.plaintext)
+        test_encryptor = StreamEncryptor(
+            source=pt_stream,
+            materials_manager=self.mock_materials_manager,
+            algorithm=aws_encryption_sdk.internal.defaults.ALGORITHM,
+            frame_length=self.mock_frame_length,
+            commitment_policy=self.mock_commitment_policy,
+            signature_policy=self.mock_signature_policy,
+        )
+        test_encryptor.signer = sentinel.signer
+        test_encryptor.content_type = sentinel.content_type
+        test_encryptor._header = sentinel.header
+        sentinel.header.version = SerializationVersion.V1
+        test_encryptor.output_buffer = b""
+        test_encryptor._encryption_materials = self.mock_encryption_materials
+        test_encryptor._derived_data_key = sentinel.derived_data_key
+
+        # Given: StreamEncryptor has _required_encryption_context
+        mock_required_ec = MagicMock(__class__=dict)
+        test_encryptor._required_encryption_context = mock_required_ec
+        mock_serialized_required_ec = MagicMock(__class__=bytes)
+        serialize_encryption_context.return_value = mock_serialized_required_ec
+
+        # When: _write_header()
+        test_encryptor._write_header()
+
+        self.mock_serialize_header.assert_called_once_with(header=test_encryptor._header, signer=sentinel.signer)
+        self.mock_serialize_header_auth.assert_called_once_with(
+            version=sentinel.header.version,
+            algorithm=self.mock_encryption_materials.algorithm,
+            header=b"12345",
+            data_encryption_key=sentinel.derived_data_key,
+            signer=sentinel.signer,
+            # Then: Pass serialized required EC to serialize_header_auth
+            required_ec_bytes=mock_serialized_required_ec,
+        )
+        assert test_encryptor.output_buffer == b"1234567890"
+
     @patch("aws_encryption_sdk.streaming_client.non_framed_body_iv")
     def test_prep_non_framed(self, mock_non_framed_iv):
         self.mock_serialize_non_framed_open.return_value = b"1234567890"