chore(Mutations): respond to comments on PR #720

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/DefaultKeyStorageInterface.dfy

@@ -3,6 +3,7 @@
 include "../Model/AwsCryptographyKeyStoreTypes.dfy"
 include "Structure.dfy"
 include "ErrorMessages.dfy"
+include "KmsArn.dfy"
 
 module DefaultKeyStorageInterface {
   import opened Wrappers
@@ -13,6 +14,7 @@ module DefaultKeyStorageInterface {
   import UTF8
   import Structure
   import String = StandardLibrary.String
+  import KmsArn
 
   const ToAttributeMap := Structure.ToAttributeMap
   const ToEncryptedHierarchicalKey := Structure.ToEncryptedHierarchicalKey
@@ -33,7 +35,7 @@ module DefaultKeyStorageInterface {
   // To use these values in a match Dafny needs to match these as local variables.
   // This means that Dafny can not use `Structure.MUTATION_LOCK_TYPE`
   // in the case statement to evaluate a literal.
-  const MUTATION_LOCK_TYPE := "MUTATION_LOCK" // Structure.MUTATION_LOCK_TYPE
+  const MUTATION_LOCK_TYPE := "branch:MUTATION_LOCK" // Structure.MUTATION_LOCK_TYPE
   const BRANCH_KEY_ACTIVE_TYPE := "branch:ACTIVE" // Structure.BRANCH_KEY_ACTIVE_TYPE
   const BEACON_KEY_TYPE_VALUE :=  "beacon:ACTIVE" // Structure.BEACON_KEY_TYPE_VALUE
   const VERSION_TYPE_PREFIX := "branch:version:" // Structure.BRANCH_KEY_TYPE_PREFIX
@@ -828,11 +830,10 @@ module DefaultKeyStorageInterface {
         Types.KeyStorageException(
           message:="DDB returned no items from table: " + ddbTableName));
 
-      // SDKs/Smithy-Dafny/Custom Implementations of Storage MAY respond with None or an Empty Map.
-      // .NET returns an empty map, Java returns None.
-      var thereAreNoUnprocessedKeys: bool := ddbResponse.UnprocessedKeys.None? || |ddbResponse.UnprocessedKeys.value| == 0;
+        // SDKs/Smithy-Dafny/Custom Implementations of Storage MAY respond with None or an Empty Map.
+        // .NET returns an empty map, Java returns None.
       :- Need(
-        thereAreNoUnprocessedKeys,
+        ddbResponse.UnprocessedKeys.None? || |ddbResponse.UnprocessedKeys.value| == 0,
         Types.KeyStorageException(
           message:=
             "DDB returned UnprocessedKeys for a BatchGetItem request of 3 items"
@@ -898,7 +899,8 @@ module DefaultKeyStorageInterface {
       :- Need(
         && Structure.ActiveHierarchicalSymmetricKey?(activeItem.value)
         && activeItem.value.Identifier == input.Identifier
-        && activeItem.value.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName,
+        && activeItem.value.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName
+        && KmsArn.ValidKmsArn?(activeItem.value.KmsArn),
         Types.KeyStorageException(
           message:=
             "Item returned for the ACTIVE is malformed. TableName: " + ddbTableName + "\tBranch Key ID: " + input.Identifier
@@ -912,7 +914,8 @@ module DefaultKeyStorageInterface {
       :- Need(
         && Structure.ActiveHierarchicalSymmetricBeaconKey?(beaconItem.value)
         && beaconItem.value.Identifier == input.Identifier
-        && beaconItem.value.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName,
+        && beaconItem.value.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName
+        && KmsArn.ValidKmsArn?(beaconItem.value.KmsArn),
         Types.KeyStorageException(
           message:=
             "Item returned for Beacon Key is malformed. TableName: " + ddbTableName + "\tBranch Key ID: " + input.Identifier
@@ -1040,7 +1043,7 @@ module DefaultKeyStorageInterface {
       );
 
       var transactWriteItemsResponse? := ddbClient.TransactWriteItems(transactRequest);
-      // TODO: Post Beta, we need to check the cancellation reason for
+      // TODO-Mutations-FF: we need to check the cancellation reason for
       // ConditionalCheckFailed on the Active item (VersionRaceException)
       // OR the Mutation Lock (MutationLockException)
       // OR something else.
@@ -1051,6 +1054,7 @@ module DefaultKeyStorageInterface {
                       ddbOperation:="TransactWriteItems",
                       identifier:=input.active.Identifier,
                       tableName:=ddbTableName));
+      // This is a Smithy Modeled Operation; the output MUST be a Structure
       output := Success(Types.WriteInitializeMutationOutput());
     }
 
@@ -1122,23 +1126,27 @@ module DefaultKeyStorageInterface {
                        ));
       }
 
-      var items: seq<Types.EncryptedHierarchicalKey> :- Seq.MapWithResult(
+      /* Map DDB items to Branch Keys.*/
+      var branchKeys: seq<Types.EncryptedHierarchicalKey> :- Seq.MapWithResult(
+        // Dafny requires the type of the element being mapped over, or it feaks out.
         (item: DDB.AttributeMap)
         =>
-          var key :- encryptedHierarchicalKeyFromItem(item, logicalKeyStoreName, input.Identifier);
+          /* Convert DDB Item to Branch Key. */
+          var branchKey :- encryptedHierarchicalKeyFromItem(item, logicalKeyStoreName, input.Identifier);
+          /* Validate that Branch Key is a Version, or Decrypt Only, Branch Key Type. */
           :- Need(
-               key.Type.HierarchicalSymmetricVersion?,
+               branchKey.Type.HierarchicalSymmetricVersion?,
                Types.KeyStorageException(
                  message:="Unexpected item returned by DDB. TableName: " + ddbTableName + "\tBranch Key ID: " + input.Identifier
                ));
-          Success(key),
+          Success(branchKey),
         ddbResponse.Items.value
       );
 
       return Success(
           Types.QueryForVersionsOutput(
             exclusiveStartKey := lastKeyBlob,
-            items := items
+            items := branchKeys
           ));
     }
 
@@ -1155,11 +1163,8 @@ module DefaultKeyStorageInterface {
       // From that, we can infer that Partition Key is just going to be the Identifier.
       // Thus, we only need to store the Type value.
       // This will be the full "branch:version:<uuidv4>"
-      var versionStr :- UTF8.Decode(blob)
-                        .MapFailure(
-                          eString
-                          =>
-                            Types.KeyStorageException(
+      var versionStr :- UTF8.Decode(blob).MapFailure(
+                          eString => Types.KeyStorageException(
                               message:="Could not UTF8 Decode Exclusive Start Key. " + eString));
       :- Need(
            // I elected to require len > 15, rather than len == 51, in case we or someone else ever uses not-UUIDv4 for version.
@@ -1230,23 +1235,27 @@ module DefaultKeyStorageInterface {
 
       /** Convert Items to DDB */
       var items?: seq<DDB.TransactWriteItem> :- Seq.MapWithResult(
-        (key: Types.EncryptedHierarchicalKey)
+        (branchKey: Types.EncryptedHierarchicalKey)
         =>
+          /* All Attribute Names MUST comply with DDB's limits.*/
+          /* Attribute Names are the "keys" of the Encryption Context.*/
           :- Need(
-               && (forall k <- key.EncryptionContext.Keys :: DDB.IsValid_AttributeName(k)),
+               && (forall k <- branchKey.EncryptionContext.Keys :: DDB.IsValid_AttributeName(k)),
                Types.KeyStorageException( message := ErrorMessages.ENCRYPTION_CONTEXT_EXCEEDS_DDB_LIMIT)
              );
+          /* Only Version, or Decrypt Only, items are permitted.*/
           :- Need(
-               key.Type.HierarchicalSymmetricVersion?,
+               branchKey.Type.HierarchicalSymmetricVersion?,
                Types.KeyStorageException(
                  message :=
                    "WriteMutatedVersions of DynamoDB Encrypted Key Storage ONLY writes Decrypt Only Items to Storage. "
                    + "Encountered a non-Decrypt Only Item."
                ));
+          /* The branch key is valid for DDB; create a Put request.*/
           var item := CreateTransactWritePutItem(
-                        key,
+                        branchKey,
                         ddbTableName,
-                        BRANCH_KEY_EXISTS);
+                        BRANCH_KEY_EXISTS); // The Branch Key Item already exists in the table.
           Success(item),
         input.items);
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/KMSKeystoreOperations.dfy

@@ -363,8 +363,8 @@ module {:options "/functionSyntax:4" } KMSKeystoreOperations {
       GrantTokens := Some(grantTokens)
     );
 
-    var maybeReEncryptResponse := kmsClient.ReEncrypt(reEncryptRequest);
-    var reEncryptResponse :- maybeReEncryptResponse
+    var reEncryptResponse? := kmsClient.ReEncrypt(reEncryptRequest);
+    var reEncryptResponse :- reEncryptResponse?
     .MapFailure(e => Types.ComAmazonawsKms(ComAmazonawsKms := e));
 
     :- Need(

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/Structure.dfy

@@ -42,7 +42,7 @@ module {:options "/functionSyntax:4" } Structure {
   const BRANCH_KEY_TYPE_PREFIX := "branch:version:"
   const BRANCH_KEY_ACTIVE_TYPE := "branch:ACTIVE"
   const BEACON_KEY_TYPE_VALUE := "beacon:ACTIVE"
-  const MUTATION_LOCK_TYPE := "MUTATION_LOCK"
+  const MUTATION_LOCK_TYPE := "branch:MUTATION_LOCK"
   //= aws-encryption-sdk-specification/framework/branch-key-store.md#custom-encryption-context
   //= type=exception
   //# Across all versions of a Branch Key, the custom encryption context MUST be equal.
@@ -458,7 +458,7 @@ module {:options "/functionSyntax:4" } Structure {
       KEY_CREATE_TIME := timestamp,
       TABLE_FIELD := logicalKeyStoreName,
       KMS_FIELD := kmsKeyArn,
-      HIERARCHY_VERSION := "1"
+      HIERARCHY_VERSION := HIERARCHY_VERSION_VALUE
     ] + map k <- customEncryptionContext :: ENCRYPTION_CONTEXT_PREFIX + k := customEncryptionContext[k]
   }
 
@@ -494,45 +494,45 @@ module {:options "/functionSyntax:4" } Structure {
   }
 
   function ReplaceMutableContext(
-    item: map<string, string>,
+    branchKeyContext: map<string, string>,
     terminalKmsArn: string,
     terminalCustomEncryptionContext: map<string, string>
   ) : (output: map<string, string>)
 
-    requires BranchKeyContext?(item)
+    requires BranchKeyContext?(branchKeyContext)
     requires BRANCH_KEY_RESTRICTED_FIELD_NAMES !! terminalCustomEncryptionContext.Keys
 
     ensures BranchKeyContext?(output)
     ensures output[KMS_FIELD] == terminalKmsArn
     ensures
-      && item[BRANCH_KEY_IDENTIFIER_FIELD] == output[BRANCH_KEY_IDENTIFIER_FIELD]
-      && item[TYPE_FIELD] == output[TYPE_FIELD]
-      && item[KEY_CREATE_TIME] == output[KEY_CREATE_TIME]
-      && item[HIERARCHY_VERSION] == output[HIERARCHY_VERSION]
-      && item[TABLE_FIELD] == output[TABLE_FIELD]
-      && (BRANCH_KEY_ACTIVE_VERSION_FIELD in item
+      && branchKeyContext[BRANCH_KEY_IDENTIFIER_FIELD] == output[BRANCH_KEY_IDENTIFIER_FIELD]
+      && branchKeyContext[TYPE_FIELD] == output[TYPE_FIELD]
+      && branchKeyContext[KEY_CREATE_TIME] == output[KEY_CREATE_TIME]
+      && branchKeyContext[HIERARCHY_VERSION] == output[HIERARCHY_VERSION]
+      && branchKeyContext[TABLE_FIELD] == output[TABLE_FIELD]
+      && (BRANCH_KEY_ACTIVE_VERSION_FIELD in branchKeyContext
           <==>
           && BRANCH_KEY_ACTIVE_VERSION_FIELD in output
-          && item[BRANCH_KEY_ACTIVE_VERSION_FIELD] == output[BRANCH_KEY_ACTIVE_VERSION_FIELD])
+          && branchKeyContext[BRANCH_KEY_ACTIVE_VERSION_FIELD] == output[BRANCH_KEY_ACTIVE_VERSION_FIELD])
   {
     terminalCustomEncryptionContext
-    + if BRANCH_KEY_ACTIVE_VERSION_FIELD in item then
+    + if BRANCH_KEY_ACTIVE_VERSION_FIELD in branchKeyContext then
       map[
-        BRANCH_KEY_IDENTIFIER_FIELD := item[BRANCH_KEY_IDENTIFIER_FIELD],
-        TYPE_FIELD := item[TYPE_FIELD],
-        KEY_CREATE_TIME := item[KEY_CREATE_TIME],
-        HIERARCHY_VERSION := item[HIERARCHY_VERSION],
-        TABLE_FIELD := item[TABLE_FIELD],
+        BRANCH_KEY_IDENTIFIER_FIELD := branchKeyContext[BRANCH_KEY_IDENTIFIER_FIELD],
+        TYPE_FIELD := branchKeyContext[TYPE_FIELD],
+        KEY_CREATE_TIME := branchKeyContext[KEY_CREATE_TIME],
+        HIERARCHY_VERSION := branchKeyContext[HIERARCHY_VERSION],
+        TABLE_FIELD := branchKeyContext[TABLE_FIELD],
         KMS_FIELD := terminalKmsArn,
-        BRANCH_KEY_ACTIVE_VERSION_FIELD := item[BRANCH_KEY_ACTIVE_VERSION_FIELD]
+        BRANCH_KEY_ACTIVE_VERSION_FIELD := branchKeyContext[BRANCH_KEY_ACTIVE_VERSION_FIELD]
       ]
     else
       map[
-        BRANCH_KEY_IDENTIFIER_FIELD := item[BRANCH_KEY_IDENTIFIER_FIELD],
-        TYPE_FIELD := item[TYPE_FIELD],
-        KEY_CREATE_TIME := item[KEY_CREATE_TIME],
-        HIERARCHY_VERSION := item[HIERARCHY_VERSION],
-        TABLE_FIELD := item[TABLE_FIELD],
+        BRANCH_KEY_IDENTIFIER_FIELD := branchKeyContext[BRANCH_KEY_IDENTIFIER_FIELD],
+        TYPE_FIELD := branchKeyContext[TYPE_FIELD],
+        KEY_CREATE_TIME := branchKeyContext[KEY_CREATE_TIME],
+        HIERARCHY_VERSION := branchKeyContext[HIERARCHY_VERSION],
+        TABLE_FIELD := branchKeyContext[TABLE_FIELD],
         KMS_FIELD := terminalKmsArn
       ]
   }
@@ -724,7 +724,7 @@ module {:options "/functionSyntax:4" } Structure {
     && KEY_CREATE_TIME in m && m[KEY_CREATE_TIME].S?
     && TYPE_FIELD in m && m[TYPE_FIELD].S?
     && M_LOCK_UUID in m && m[M_LOCK_UUID].S?
-    && HIERARCHY_VERSION in m && m[HIERARCHY_VERSION].N? && m[HIERARCHY_VERSION].N == "1"
+    && HIERARCHY_VERSION in m && m[HIERARCHY_VERSION].N? && m[HIERARCHY_VERSION].N == HIERARCHY_VERSION_VALUE
 
     && 0 < |m[BRANCH_KEY_IDENTIFIER_FIELD].S|
     && 0 < |m[TYPE_FIELD].S|
@@ -734,6 +734,9 @@ module {:options "/functionSyntax:4" } Structure {
 
     && m[TYPE_FIELD].S == MUTATION_LOCK_TYPE
 
+    // Structure & DefaultKeyStorage do not care about the Byte structure of the original or terminal.
+    // That is the concern of Mutation State Structures.
+    // Structure & DefaultKeyStorage care that these are non-empty Byte Fields.
     && M_LOCK_ORIGINAL in m && m[M_LOCK_ORIGINAL].B? && 0 < |m[M_LOCK_ORIGINAL].B|
     && M_LOCK_TERMINAL in m && m[M_LOCK_TERMINAL].B? && 0 < |m[M_LOCK_TERMINAL].B|
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/src/MutationStateStructures.dfy

@@ -37,6 +37,11 @@ module {:options "/functionSyntax:4" } MutationStateStructures {
 
   const MUTABLE_PROPERTY_COUNT: int := 2
   const MUTABLE_PROPERTY_COUNT_str := "2"
+  // We use "aws-crypto-ec" instead of "aws-crypto-ec:" out of a paranoia
+  // that the JSON serialization implementation is worng and does not escape
+  // keys correctly.
+  const MUTABLE_PROPERTY_EC_WORD := "aws-crypto-ec"
+  const MUTABLE_PROPERTY_KMS_ARN_WORD := Structure.KMS_FIELD
   // Ensures
   // - if KMS ARN, Valid KMS ARN
   // - if EC, Valid non-empty EC
@@ -50,13 +55,11 @@ module {:options "/functionSyntax:4" } MutationStateStructures {
           && |input.terminalEncryptionContext.value| > 0
           &&  forall k <- input.terminalEncryptionContext.value ::
                && |k| > 0 && |input.terminalEncryptionContext.value[k]| > 0
-                  // && |Structure.SelectCustomEncryptionContextAsString(input.terminalEncryptionContext.value)| == 0
                && input.terminalEncryptionContext.value.Keys !! Structure.BRANCH_KEY_RESTRICTED_FIELD_NAMES)
     && !(input.terminalKmsArn.None? && input.terminalEncryptionContext.None?)
   }
 
   datatype MutableProperties = | MutableProperties (
-    // nameonly kmsArn: KeyStoreTypes.KMSConfiguration,
     nameonly kmsArn: validKmsArn,
     nameonly customEncryptionContext: KeyStoreTypes.EncryptionContextString
   )
@@ -151,23 +154,23 @@ module {:options "/functionSyntax:4" } MutationStateStructures {
   {
     var OriginalJson
       := JSONValues.Object([
-                             ("aws-crypto-ec", EncryptionContextStringToJSON(MutationToApply.Original.customEncryptionContext)),
-                             ("kms-arn", JSONValues.JSON.String(MutationToApply.Original.kmsArn))
+                             (MUTABLE_PROPERTY_EC_WORD, EncryptionContextStringToJSON(MutationToApply.Original.customEncryptionContext)),
+                             (MUTABLE_PROPERTY_KMS_ARN_WORD, JSONValues.JSON.String(MutationToApply.Original.kmsArn))
                            ]);
     var TerminalJson
       := JSONValues.Object([
-                             ("aws-crypto-ec", EncryptionContextStringToJSON(MutationToApply.Terminal.customEncryptionContext)),
-                             ("kms-arn", JSONValues.JSON.String(MutationToApply.Terminal.kmsArn))
+                             (MUTABLE_PROPERTY_EC_WORD, EncryptionContextStringToJSON(MutationToApply.Terminal.customEncryptionContext)),
+                             (MUTABLE_PROPERTY_KMS_ARN_WORD, JSONValues.JSON.String(MutationToApply.Terminal.kmsArn))
                            ]);
 
     var originalBytes :- JSON.Serialize(OriginalJson).MapFailure(
                            (e: JSONErrors.SerializationError)
                            => Types.KeyStoreAdminException(
-                               message := "Could JSON Serialize state: original properties. " + e.ToString()));
+                               message := "Could not JSON Serialize state: original properties. " + e.ToString()));
     var terminalBytes :- JSON.Serialize(TerminalJson).MapFailure(
                            (e: JSONErrors.SerializationError)
                            => Types.KeyStoreAdminException(
-                               message := "Could JSON Serialize state: terminal properties. " + e.ToString()));
+                               message := "Could not JSON Serialize state: terminal properties. " + e.ToString()));
     Success(
       Types.MutationToken(
         Identifier := MutationToApply.Identifier,
@@ -179,8 +182,6 @@ module {:options "/functionSyntax:4" } MutationStateStructures {
       ))
   }
 
-  const ERROR_PRFX := "Serialized State properties is malformed!"
-
   function DeserializeMutationToken(
     Token: Types.MutationToken
   )
@@ -217,45 +218,50 @@ module {:options "/functionSyntax:4" } MutationStateStructures {
       ))
   }
 
+  const ERROR_PRFX := "Serialized State properties is malformed! "
+
   function MutablePropertiesJson?(
     MutableProperties: JSONValues.JSON
   ): (output: Outcome<Types.Error>)
   {
     :- NeedOutcome(
          MutableProperties.Object? && |MutableProperties.obj| == 2,
-         () => Types.KeyStoreAdminException( message := "WIP")
+         () => Types.KeyStoreAdminException( message := ERROR_PRFX + "There should be two objects.")
        );
     :- NeedOutcome(
-         MutableProperties.obj[0].0 == "aws-crypto-ec",
-         () => Types.KeyStoreAdminException( message := "WIP")
+         MutableProperties.obj[0].0 == MUTABLE_PROPERTY_EC_WORD,
+         () => Types.KeyStoreAdminException( message := ERROR_PRFX + "First Key MUST be Encryption Context.")
        );
     :- NeedOutcome(
-         MutableProperties.obj[1].0 == "kms-arn",
-         () => Types.KeyStoreAdminException( message := "WIP")
+         MutableProperties.obj[1].0 == MUTABLE_PROPERTY_KMS_ARN_WORD,
+         () => Types.KeyStoreAdminException( message := ERROR_PRFX + "Second Key MUST be KMS ARN.")
        );
     :- NeedOutcome(
          MutableProperties.obj[0].1.Object?,
-         () => Types.KeyStoreAdminException( message := "WIP")
+         () => Types.KeyStoreAdminException(
+             message := ERROR_PRFX + "Value for `" + MUTABLE_PROPERTY_EC_WORD + "` MUST be an object.")
        );
     :- NeedOutcome(
          MutableProperties.obj[1].1.String?,
-         () => Types.KeyStoreAdminException( message := "WIP")
+         () => Types.KeyStoreAdminException(
+             message := ERROR_PRFX + "Value for `" + MUTABLE_PROPERTY_KMS_ARN_WORD + "` MUST be a string.")
        );
     :- NeedOutcome(
          KmsArn.ValidKmsArn?(MutableProperties.obj[1].1.str),
-         () => Types.KeyStoreAdminException( message := "WIP")
+         () => Types.KeyStoreAdminException( message := ERROR_PRFX + "KMS ARN that has been deserialized is invalid.")
        );
 
     var EncryptionContext := MutableProperties.obj[0].1;
     :- NeedOutcome(
          forall p <- EncryptionContext.obj :: p.1.String?,
-         () => Types.KeyStoreAdminException( message := "WIP")
+         () => Types.KeyStoreAdminException( message := ERROR_PRFX + "Member of Encryption Context cannot be deserialized.")
        );
 
     var EncryptionContextKeys := set p <- EncryptionContext.obj :: p.0;
     :- NeedOutcome(
          |EncryptionContextKeys| == |EncryptionContext.obj|,
-         () => Types.KeyStoreAdminException( message := "WIP")
+         () => Types.KeyStoreAdminException(
+             message := ERROR_PRFX + "Size of Encryption Context keys is not equal to size of Encryption Context values. ")
        );
 
     :- NeedOutcome(