diff --git a/src/aws_encryption_sdk/internal/formatting/encryption_context.py b/src/aws_encryption_sdk/internal/formatting/encryption_context.py index 4d3a5e773..61756f663 100644 --- a/src/aws_encryption_sdk/internal/formatting/encryption_context.py +++ b/src/aws_encryption_sdk/internal/formatting/encryption_context.py @@ -47,6 +47,38 @@ def assemble_content_aad(message_id, aad_content_string, seq_num, length): return struct.pack(fmt, message_id, aad_content_string.value, seq_num, length) +def _serialize_encryption_context_list(encryption_context_list, serialized_context): + """Helper function to serialize the list of encryption context key-value pairs.""" + for key, value in sorted(encryption_context_list, key=lambda x: x[0]): + try: + if len(key) > 2 ** 16 - 1 or len(value) > 2 ** 16 - 1: + raise SerializationError("Key or Value are too large. Maximum length is 65535") + + serialized_context.extend( + struct.pack( + ">H{key_size}sH{value_size}s".format(key_size=len(key), value_size=len(value)), + len(key), + key, + len(value), + value, + ) + ) + except struct.error as struct_error: + message = str(struct_error) + if message == "'H' format requires 0 <= number <= 65535": + raise SerializationError("Key or Value are too large. Maximum length is 65535") + + # else unknown + raise SerializationError( + [SerializationError("Unknown Error with struct pack. Could not serialize key or value"), struct_error] + ) + + if len(serialized_context) > aws_encryption_sdk.internal.defaults.MAX_BYTE_ARRAY_SIZE: + raise SerializationError("The serialized context is too large.") + + return bytes(serialized_context) + + def serialize_encryption_context(encryption_context): """Serializes the contents of a dictionary into a byte string. @@ -80,20 +112,7 @@ def serialize_encryption_context(encryption_context): raise SerializationError( "Cannot encode dictionary key or value using {}.".format(aws_encryption_sdk.internal.defaults.ENCODING) ) - - for key, value in sorted(encryption_context_list, key=lambda x: x[0]): - serialized_context.extend( - struct.pack( - ">H{key_size}sH{value_size}s".format(key_size=len(key), value_size=len(value)), - len(key), - key, - len(value), - value, - ) - ) - if len(serialized_context) > aws_encryption_sdk.internal.defaults.MAX_BYTE_ARRAY_SIZE: - raise SerializationError("The serialized context is too large.") - return bytes(serialized_context) + return _serialize_encryption_context_list(encryption_context_list, serialized_context) def read_short(source, offset): diff --git a/test/unit/test_encryption_context.py b/test/unit/test_encryption_context.py index 187365783..f25ce62b8 100644 --- a/test/unit/test_encryption_context.py +++ b/test/unit/test_encryption_context.py @@ -11,6 +11,8 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Unit test suite for aws_encryption_sdk.internal.formatting.encryption_context""" +import struct + import pytest import aws_encryption_sdk.internal.defaults @@ -23,6 +25,23 @@ pytestmark = [pytest.mark.unit, pytest.mark.local] +def test_struct_behavior_pack_too_large(): + """ + If a key or value in the encryption context are too large we need to throw + a serialization error instead of the implicit struct.error caused by trying to + >H pack a key or value that is too long. The current implementation requires + the error to be thrown and checks for some conditions involving the size of a + key or value and the error message thrown. For it to work properly we need + struct to fail in a certain manner, thus this test. + """ + value = struct.pack(">H", 2 ** 16 - 1) + assert value + + with pytest.raises(struct.error) as excinfo: + struct.pack(">H", 2 ** 16) + excinfo.match(r"'H' format requires 0 <= number <= 65535") + + class TestEncryptionContext(object): def test_assemble_content_aad(self): """Validate that the assemble_content_aad function @@ -190,3 +209,21 @@ def test_deserialize_encryption_context_empty(self): serialized_encryption_context=b"" ) assert test == {} + + def test_serialize_encryption_context_key_value_too_long(self): + dictionary = {"0" * 2 ** 16: "0" * 2 ** 16} + with pytest.raises(SerializationError) as excinfo: + aws_encryption_sdk.internal.formatting.encryption_context.serialize_encryption_context(dictionary) + excinfo.match(r"Key or Value are too large. Maximum length is 65535") + + def test_serialize_encryption_context_value_too_long(self): + dictionary = {"short": "0" * 2 ** 16} + with pytest.raises(SerializationError) as excinfo: + aws_encryption_sdk.internal.formatting.encryption_context.serialize_encryption_context(dictionary) + excinfo.match(r"Key or Value are too large. Maximum length is 65535") + + def test_serialize_encryption_context_key_too_long(self): + dictionary = {"0" * 2 ** 16: "short"} + with pytest.raises(SerializationError) as excinfo: + aws_encryption_sdk.internal.formatting.encryption_context.serialize_encryption_context(dictionary) + excinfo.match(r"Key or Value are too large. Maximum length is 65535")