feat: add more attributes to encryption context

DynamoDbEncryption/dafny/DynamoDbItemEncryptor/src/AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations.dfy

@@ -100,11 +100,12 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
     else if action == CSE.ENCRYPT_AND_SIGN then
       "ENCRYPT_AND_SIGN"
     else
-      assert action != CSE.DO_NOTHING;
-      assert action != CSE.SIGN_ONLY;
-      assert action != CSE.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT;
-      assert action != CSE.ENCRYPT_AND_SIGN;
-      assert false;
+      assert false by {
+        assert action != CSE.DO_NOTHING;
+        assert action != CSE.SIGN_ONLY;
+        assert action != CSE.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT;
+        assert action != CSE.ENCRYPT_AND_SIGN;
+      }
       "internal error"
   }
 
@@ -731,8 +732,8 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
     //# [SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT](../structured-encryption/structures.md#contextandsign)
     //# then this item MUST include an Attribute with that name.
     ensures output.Success? ==>
-      forall k <- config.attributeActionsOnEncrypt :: (config.attributeActionsOnEncrypt[k] == CSE.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT) ==> (k in input.plaintextItem)
-      // ContextAttrsExist(config.attributeActionsOnEncrypt, input.plaintextItem)
+              forall k <- config.attributeActionsOnEncrypt :: (config.attributeActionsOnEncrypt[k] == CSE.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT) ==> (k in input.plaintextItem)
+    // ContextAttrsExist(config.attributeActionsOnEncrypt, input.plaintextItem)
 
     ensures
       && output.Success?
@@ -782,6 +783,16 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
               && SE.IsAuthAttr(parsedHeaderMap[k].content.Action))
         && var maybeCryptoSchema := ConvertCryptoSchemaToAttributeActions(config, structuredEncParsed.cryptoSchema);
         && maybeCryptoSchema.Success?
+        && ConvertContextForSelector(structuredEncParsed.encryptionContext).Success?
+        && var selectorContext := ConvertContextForSelector(structuredEncParsed.encryptionContext).value;
+        && output.value.parsedHeader.value == ParsedHeader(
+                                                attributeActionsOnEncrypt := maybeCryptoSchema.value,
+                                                algorithmSuiteId := structuredEncParsed.algorithmSuiteId,
+                                                storedEncryptionContext := structuredEncParsed.storedEncryptionContext,
+                                                encryptedDataKeys := structuredEncParsed.encryptedDataKeys,
+                                                encryptionContext := structuredEncParsed.encryptionContext,
+                                                selectorContext := selectorContext
+                                              )
 
     //= specification/dynamodb-encryption-client/encrypt-item.md#behavior
     //= type=implication
@@ -993,6 +1004,16 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
               && SE.IsAuthAttr(structuredEncParsed.cryptoSchema.content.SchemaMap[k].content.Action))
         && var maybeCryptoSchema := ConvertCryptoSchemaToAttributeActions(config, structuredEncParsed.cryptoSchema);
         && maybeCryptoSchema.Success?
+        && ConvertContextForSelector(structuredEncParsed.encryptionContext).Success?
+        && var selectorContext := ConvertContextForSelector(structuredEncParsed.encryptionContext).value;
+        && output.value.parsedHeader.value == ParsedHeader(
+                                                attributeActionsOnEncrypt := maybeCryptoSchema.value,
+                                                algorithmSuiteId := structuredEncParsed.algorithmSuiteId,
+                                                storedEncryptionContext := structuredEncParsed.storedEncryptionContext,
+                                                encryptedDataKeys := structuredEncParsed.encryptedDataKeys,
+                                                encryptionContext := structuredEncParsed.encryptionContext,
+                                                selectorContext := selectorContext
+                                              )
 
     //= specification/dynamodb-encryption-client/decrypt-item.md#behavior
     //= type=implication

DynamoDbEncryption/dafny/DynamoDbItemEncryptor/src/Util.dfy

@@ -43,22 +43,31 @@ module DynamoDbItemEncryptorUtil {
   const SELECTOR_PARTITION_NAME : DDB.AttributeName := "aws_dbe_partition_name"
   const SELECTOR_SORT_NAME : DDB.AttributeName := "aws_dbe_sort_name"
 
-  method GetLiteralValue(x : seq<uint8>) returns (ret : Result<DDB.AttributeValue, string>)
+  function method GetLiteralValue(x : seq<uint8>) : (ret : Result<DDB.AttributeValue, string>)
   {
     var str :- UTF8.Decode(x);
-    if str == SE.TRUE_STR {
-      return Success(DDB.AttributeValue.BOOL(true));
-    } else if str == SE.FALSE_STR {
-      return Success(DDB.AttributeValue.BOOL(false));
-    } else if str == SE.NULL_STR {
-      return Success(DDB.AttributeValue.NULL(true));
-    } else {
-      return Failure("Encryption Context literal value has unexpected value : '" + str + "'.");
-    }
+    if str == SE.TRUE_STR then
+      Success(DDB.AttributeValue.BOOL(true))
+    else if str == SE.FALSE_STR then
+      Success(DDB.AttributeValue.BOOL(false))
+    else if str == SE.NULL_STR then
+      Success(DDB.AttributeValue.NULL(true))
+    else
+      Failure("Encryption Context literal value has unexpected value : '" + str + "'.")
   }
 
-  method ConvertContextForSelector(context : MPL.EncryptionContext)
-    returns (output: Result<DDB.Key, string>)
+  function method GetSortKey(context : MPL.EncryptionContext) : Result<Option<UTF8.ValidUTF8Bytes>, string>
+  {
+    if SORT_NAME in context.Keys then
+      var sortName := SE.EC_ATTR_PREFIX + context[SORT_NAME];
+      :- Need(UTF8.ValidUTF8Seq(sortName), "Internal Error : bad utf8 in sortName");
+      Success(Some(sortName))
+    else
+      Success(None)
+  }
+
+  function method ConvertContextForSelector(context : MPL.EncryptionContext)
+    : (output: Result<DDB.Key, string>)
     //= specification/dynamodb-encryption-client/ddb-encryption-branch-key-id-supplier.md#behavior
     //= type=implication
     //# - If the partition name does not exist in the encryption context, this operation MUST fail.
@@ -111,75 +120,86 @@ module DynamoDbItemEncryptorUtil {
     var partitionValueKey := SE.EC_ATTR_PREFIX + partitionName;
     :- Need(partitionValueKey in context.Keys, "Invalid encryption context: Missing partition value");
 
-    if SORT_NAME in context.Keys {
-      var sortName := context[SORT_NAME];
-      var sortValueKey := SE.EC_ATTR_PREFIX + sortName;
-      :- Need(sortValueKey in context.Keys, "Invalid encryption context: Missing sort value");
-    }
+    var sortValueKey :- GetSortKey(context);
+    :- Need(sortValueKey.None? || sortValueKey.value in context, "Invalid encryption context: Missing sort value");
 
-    var attrMap: DDB.AttributeMap := map[];
     var keys : seq<UTF8.ValidUTF8Bytes> := SortedSets.ComputeSetToOrderedSequence2(context.Keys, SE.ByteLess);
 
     //= specification/dynamodb-encryption-client/ddb-encryption-branch-key-id-supplier.md#behavior
     //# - If the field "aws-crypto-legend" exists in the encryption context,
     //# it MUST [deserialize](./ddb-attribute-serialization.md), all values with keys beginning "aws-crypto-attr.",
     //# and create a Key with these values, using names with the "aws-crypto-attr." removed.
-    if SE.LEGEND_UTF8 in context {
+    if SE.LEGEND_UTF8 in context then
       var legend :- UTF8.Decode(context[SE.LEGEND_UTF8]);
-
-      var attrPos := 0;
-      for i := 0 to |keys|
-        invariant attrPos <= i
-      {
-        var key : UTF8.ValidUTF8Bytes := keys[i];
-        if SE.EC_ATTR_PREFIX < key {
-          :- Need(|legend| > attrPos, "Encryption Context Legend is too short.");
-          var ddbAttrName :- GetAttributeName(key);
-          var ch := legend[attrPos];
-          if ch == SE.LEGEND_STRING {
-            var value :- UTF8.Decode(context[key]);
-            attrMap := attrMap[ddbAttrName := DDB.AttributeValue.S(value)];
-          } else if ch == SE.LEGEND_NUMBER {
-            var value :- UTF8.Decode(context[key]);
-            attrMap := attrMap[ddbAttrName := DDB.AttributeValue.N(value)];
-          } else if ch == SE.LEGEND_LITERAL {
-            var value :- GetLiteralValue(context[key]);
-            attrMap := attrMap[ddbAttrName := value];
-          } else if ch == SE.LEGEND_BINARY {
-            attrMap :- AddAttributeToMap(key, context[key], attrMap);
-          } else {
-            return Failure("Encryption Context Legend has unexpected character : '" + [legend[attrPos]] + "'.");
-          }
-          attrPos := attrPos + 1;
-        }
-      }
-      :- Need(|legend| == attrPos, "Encryption Context Legend is too long.");
+      var attrMap :- GetV2AttrMap(keys, context, legend);
 
       :- Need(TABLE_NAME in context, "Internal error, table name not in encryption context.");
       var tableName :- UTF8.Decode(context[TABLE_NAME]);
-      attrMap := attrMap[SELECTOR_TABLE_NAME := DDB.AttributeValue.S(tableName)];
+      var attrMap2 := attrMap[SELECTOR_TABLE_NAME := DDB.AttributeValue.S(tableName)];
 
       :- Need(PARTITION_NAME in context, "Internal error, table name not in encryption context.");
       var partitionName :- UTF8.Decode(context[PARTITION_NAME]);
-      attrMap := attrMap[SELECTOR_PARTITION_NAME := DDB.AttributeValue.S(partitionName)];
+      var attrMap3 := attrMap2[SELECTOR_PARTITION_NAME := DDB.AttributeValue.S(partitionName)];
 
-      if SORT_NAME in context {
+      if SORT_NAME in context then
         var sortName :- UTF8.Decode(context[SORT_NAME]);
-        attrMap := attrMap[SELECTOR_SORT_NAME := DDB.AttributeValue.S(sortName)];
-      }
-      //= specification/dynamodb-encryption-client/ddb-encryption-branch-key-id-supplier.md#behavior
-      //# - If the field "aws-crypto-legend" does not exist in the encryption context, it MUST [deserialize the partition (and optionally sort) value](./ddb-attribute-serialization.md), and create a Key with these values.
-    } else {
-      for i := 0 to |keys| {
-        var key : UTF8.ValidUTF8Bytes := keys[i];
-        if SE.EC_ATTR_PREFIX < key {
-          attrMap :- AddAttributeToMap(key, context[key], attrMap);
-        }
-      }
-    }
-    return Success(attrMap);
+        Success(attrMap3[SELECTOR_SORT_NAME := DDB.AttributeValue.S(sortName)])
+      else
+        Success(attrMap3)
+
+    //= specification/dynamodb-encryption-client/ddb-encryption-branch-key-id-supplier.md#behavior
+    //# - If the field "aws-crypto-legend" does not exist in the encryption context, it MUST [deserialize the partition (and optionally sort) value](./ddb-attribute-serialization.md), and create a Key with these values.
+    else if sortValueKey.None? then
+      AddAttributeToMap(partitionValueKey, context[partitionValueKey], map[])
+    else
+      var attrMap :- AddAttributeToMap(partitionValueKey, context[partitionValueKey], map[]);
+      AddAttributeToMap(sortValueKey.value, context[sortValueKey.value], attrMap)
   }
 
+  function method GetAttrValue(ecValue : UTF8.ValidUTF8Bytes, legend : char)
+    : Result<DDB.AttributeValue, string>
+  {
+    if legend == SE.LEGEND_STRING then
+      var value :- UTF8.Decode(ecValue);
+      Success(DDB.AttributeValue.S(value))
+    else if legend == SE.LEGEND_NUMBER then
+      var value :- UTF8.Decode(ecValue);
+      Success(DDB.AttributeValue.N(value))
+    else if legend == SE.LEGEND_LITERAL then
+      var value :- GetLiteralValue(ecValue);
+      Success(value)
+    else if legend == SE.LEGEND_BINARY then
+      var terminal :- SE.DecodeTerminal(ecValue);
+      var ddbAttrValue :- DynamoToStruct.BytesToAttr(terminal.value, terminal.typeId, false);
+      Success(ddbAttrValue.val)
+    else
+      Failure("Encryption Context Legend has unexpected character : '" + [legend] + "'.")
+  }
+
+  function method GetV2AttrMap(
+    keys : seq<UTF8.ValidUTF8Bytes>,
+    context : MPL.EncryptionContext,
+    legend : string,
+    attrMap : DDB.AttributeMap := map[]
+  )
+    : Result<DDB.AttributeMap, string>
+    requires forall k <- keys :: k in context
+  {
+    if |keys| == 0 then
+      if |legend| == 0 then
+        Success(attrMap)
+      else
+        Failure("Encryption Context Legend is too long.")
+    else
+      var key : UTF8.ValidUTF8Bytes := keys[0];
+      if SE.EC_ATTR_PREFIX < key then
+        :- Need(0 < |legend|, "Encryption Context Legend is too short.");
+        var attrName :- GetAttributeName(key);
+        var attrValue :- GetAttrValue(context[key], legend[0]);
+        GetV2AttrMap(keys[1..], context, legend[1..], attrMap[attrName := attrValue])
+      else
+        GetV2AttrMap(keys[1..], context, legend, attrMap)
+  }
 
   function method GetAttributeName(ddbAttrKey: UTF8.ValidUTF8Bytes)
     : (res: Result<DDB.AttributeName, string>)
@@ -191,6 +211,7 @@ module DynamoDbItemEncryptorUtil {
     :- Need(DDB.IsValid_AttributeName(ddbAttrName), "Invalid serialization of DDB Attribute in encryption context.");
     Success(ddbAttrName)
   }
+
   function method AddAttributeToMap(ddbAttrKey: UTF8.ValidUTF8Bytes, encodedAttrValue: UTF8.ValidUTF8Bytes, attrMap: DDB.AttributeMap)
     : (res: Result<DDB.AttributeMap, string>)
     requires |ddbAttrKey| >= |SE.EC_ATTR_PREFIX|