Fixed issue #55

src/aws_encryption_sdk/internal/formatting/encryption_context.py

@@ -47,9 +47,15 @@ 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)
 
 
-# help function to simplify code
-# this is necessary to pass flake8
+def _key_or_value_too_large(encryption_context_list):
+    """Helper method to check whether a key or value are too long."""
+    longest_key = max(encryption_context_list, key=lambda x: len(x[0]))[0]
+    longest_value = max(encryption_context_list, key=lambda x: len(x[1]))[1]
+    return len(longest_key) > 2 ** 16 - 1 or len(longest_value) > 2 ** 16 - 1
+
+
 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:
             serialized_context.extend(
@@ -64,11 +70,14 @@ def _serialize_encryption_context_list(encryption_context_list, serialized_conte
         # we check to make sure that we return the right type of error message for an overly long key or value
         except struct.error as struct_error:
             message = str(struct_error)
+            # a key or value were too long and struct gave us the expected error
             if message == "'H' format requires 0 <= number <= 65535":
-                # the key or value were too large
+                raise SerializationError("Key or Value are too large.  Maximum length is 65535")
+            # a key or value were too long and struct's error message is different (maybe an update)
+            if _key_or_value_too_large(encryption_context_list):
                 raise SerializationError("Key or Value are too large.  Maximum length is 65535")
             # else we can just raise the struct error as it was (unknown)
-            raise struct_error
+            raise SerializationError("Unknown Error with struct pack.  Could not serialize key or 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)

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
@@ -191,11 +210,20 @@ def test_deserialize_encryption_context_empty(self):
         )
         assert test == {}
 
-    # these three tests (in one function) make sure that whether a key or value or both are too long
-    # we throw the right type of serialization error and not the struct.error
-    # from struct.pack()
     def test_serialize_encryption_context_key_value_too_long(self):
-        for dictionary in [{"0" * 2 ** 16: "short"}, {"short": "0" * 2 ** 16}, {"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")
+        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")