feat(dafny): mutate HV-2 to HV-2 without new version

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/Structure.dfy

@@ -510,10 +510,7 @@ module {:options "/functionSyntax:4" } Structure {
     ensures BRANCH_KEY_ACTIVE_VERSION_FIELD !in output
     ensures output[KMS_FIELD] == kmsKeyArn
     ensures output[TABLE_FIELD] == logicalKeyStoreName
-    ensures output[HIERARCHY_VERSION] == match hierarchyVersion {
-                                           case v1 => HIERARCHY_VERSION_VALUE_1
-                                           case v2 => HIERARCHY_VERSION_VALUE_2
-                                         }
+    ensures output[HIERARCHY_VERSION] == HierarchyVersionToString(hierarchyVersion)
     ensures forall k <- encryptionContext
               ::
                 && ENCRYPTION_CONTEXT_PREFIX + k in output
@@ -544,10 +541,7 @@ module {:options "/functionSyntax:4" } Structure {
       KEY_CREATE_TIME := timestamp,
       TABLE_FIELD := logicalKeyStoreName,
       KMS_FIELD := kmsKeyArn,
-      HIERARCHY_VERSION := match hierarchyVersion {
-      case v1 => HIERARCHY_VERSION_VALUE_1
-      case v2 => HIERARCHY_VERSION_VALUE_2
-    }
+      HIERARCHY_VERSION := HierarchyVersionToString(hierarchyVersion)
     ]
   }
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/src/InitializeMutation.dfy

@@ -22,6 +22,7 @@ module {:options "/functionSyntax:4" } InternalInitializeMutation {
     // KMS & MPL Imports
   import KMS = ComAmazonawsKmsTypes
   import AwsKmsUtils
+  import AtomicPrimitives
     // KeyStore Imports
   import KeyStoreTypes = AwsCryptographyKeyStoreAdminTypes.AwsCryptographyKeyStoreTypes
   import HvUtils = HierarchicalVersionUtils
@@ -231,12 +232,6 @@ module {:options "/functionSyntax:4" } InternalInitializeMutation {
       Types.KeyStoreAdminException(
         message := "Active Branch Key Item read from storage is malformed!")
     );
-    // TODO-HV-2-M3: Support mutations on HV-2 item (mutation starting with hv-2 item)
-    :- Need(
-      readItems.ActiveItem.EncryptionContext[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_VALUE_1,
-      Types.KeyStoreAdminException(
-        message := "At this time, Mutations ONLY support HV-1; BK's Active Item is HV-2.")
-    );
     var isTerminalHv2 := input.Mutations.TerminalHierarchyVersion.Some? &&
                          input.Mutations.TerminalHierarchyVersion.value.v2?;
     :- Need(
@@ -704,14 +699,13 @@ module {:options "/functionSyntax:4" } InternalInitializeMutation {
     .MapFailure(e => Types.KeyStoreAdminException(
                     message := "Could not generate UUID for new Decrypt Only. " + e));
 
-    // TODO-HV-2-M2: Refactor to allow HV-2 for Mutations
     var decryptOnlyEncryptionContext := Mutations.DecryptOnlyBranchKeyEncryptionContextForMutation(
       localInput.input.Identifier,
       newVersion,
       localInput.timestamp,
       localInput.input.logicalKeyStoreName,
       localInput.mutationToApply.Terminal.kmsArn,
-      // localInput.mutationToApply.Terminal.hierarchyVersion,
+      localInput.mutationToApply.Terminal.hierarchyVersion,
       localInput.mutationToApply.Terminal.customEncryptionContext
     );
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/src/Mutations.dfy

@@ -107,7 +107,6 @@ module {:options "/functionSyntax:4" } Mutations {
     var success?: bool := false;
     var throwAwayError;
     var kmsOperation: string;
-
     match keyManagerStrategy {
       case reEncrypt(kms) =>
         kmsOperation := "ReEncrypt";
@@ -437,7 +436,7 @@ module {:options "/functionSyntax:4" } Mutations {
     timestamp: string,
     logicalKeyStoreName: string,
     kmsKeyArn: string,
-    // hierachyVersion: HierarchyVersion,
+    hierarchyVersion: KeyStoreTypes.HierarchyVersion,
     prefixedCustomEncryptionContext: map<string, string>
   ): (output: map<string, string>)
     requires 0 < |branchKeyId|
@@ -448,6 +447,7 @@ module {:options "/functionSyntax:4" } Mutations {
     ensures Structure.BRANCH_KEY_ACTIVE_VERSION_FIELD !in output
     ensures output[Structure.KMS_FIELD] == kmsKeyArn
     ensures output[Structure.TABLE_FIELD] == logicalKeyStoreName
+    ensures output[Structure.HIERARCHY_VERSION] == HierarchyVersionToString(hierarchyVersion)
     ensures forall k <- prefixedCustomEncryptionContext
               ::
                 && k in output
@@ -459,7 +459,7 @@ module {:options "/functionSyntax:4" } Mutations {
       Structure.KEY_CREATE_TIME := timestamp,
       Structure.TABLE_FIELD := logicalKeyStoreName,
       Structure.KMS_FIELD := kmsKeyArn,
-      Structure.HIERARCHY_VERSION := Structure.HIERARCHY_VERSION_VALUE
+      Structure.HIERARCHY_VERSION := HierarchyVersionToString(hierarchyVersion)
     ] + prefixedCustomEncryptionContext
   }
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/AdminFixtures.dfy

@@ -271,7 +271,7 @@ module {:options "/functionSyntax:4" } AdminFixtures {
     nameonly id: string,
     nameonly kmsId: string := Fixtures.keyArn,
     nameonly hierarchyVersion: KeyStoreTypes.HierarchyVersion := KeyStoreTypes.HierarchyVersion.v1,
-    nameonly strategy: Types.KeyManagementStrategy,
+    nameonly strategy?: Option<Types.KeyManagementStrategy> := None,
     nameonly admin?: Option<Types.IKeyStoreAdminClient> := None,
     // nameonly versionCount: nat := 3,
     nameonly customEC: KeyStoreTypes.EncryptionContext := Fixtures.RobbieEC
@@ -289,6 +289,18 @@ module {:options "/functionSyntax:4" } AdminFixtures {
     } else {
       admin := admin?.value;
     }
+    var strategy;
+    if strategy?.None? {
+      // If no strategy is provided, set default based on hierarchy version
+      if hierarchyVersion.v1? {
+        strategy :- expect DefaultKeyManagerStrategy();
+      } else {
+        strategy :- expect SimpleKeyManagerStrategy();
+      }
+    } else {
+      strategy := strategy?.value;
+    }
+
     assume {:axiom} fresh(admin) && fresh(admin.Modifies);
     var input := Types.CreateKeyInput(
       Identifier := Some(id),

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/Mutations/TestMutateToHV2FromHV1.dfy

@@ -8,6 +8,7 @@ include "../../src/KeyStoreAdminErrorMessages.dfy"
 include "../AdminFixtures.dfy"
 include "TestMutationHappyPath.dfy"
 
+// TODO-HV-2-M2-FF: Rename to TestMutations and Abstract all mutations (v1 -> v1, v1 -> v2, v2 -> v2)
 module {:options "/functionSyntax:4" } TestMutateToHV2FromHV1 {
   import UUID
   import AdminFixtures
@@ -104,10 +105,9 @@ module {:options "/functionSyntax:4" } TestMutateToHV2FromHV1 {
   }
 
   method TestHV1toHV2HappyCase(strategy: Types.KeyManagementStrategy, ddbClient: DDB.Types.IDynamoDBClient, kmsClient: KMS.Types.IKMSClient)
-    requires ddbClient.ValidState()
-    requires kmsClient.ValidState()
-    modifies ddbClient.Modifies
-    modifies kmsClient.Modifies
+    requires ddbClient.ValidState() && kmsClient.ValidState()
+    modifies ddbClient.Modifies, kmsClient.Modifies
+    ensures ddbClient.ValidState() && kmsClient.ValidState()
   {
     var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
     var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
@@ -197,4 +197,99 @@ module {:options "/functionSyntax:4" } TestMutateToHV2FromHV1 {
       doNotVersion := true
     );
   }
+
+  method {:test} TestHV2toHV2HappyCaseKMSSimple() {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(
+      kmsClient?:=Some(kmsClient)
+    );
+    TestHV2toHV2HappyCase(strategy, ddbClient, kmsClient);
+  }
+
+  method {:test} TestHV2toHV2HappyCaseDecryptEncrypt() {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var strategy :- expect AdminFixtures.DecryptEncrypKeyManagerStrategy(
+      decryptKmsClient?:=Some(kmsClient),
+      encryptKmsClient?:=Some(kmsClient)
+    );
+    TestHV2toHV2HappyCase(strategy, ddbClient, kmsClient);
+  }
+
+  const happyCaseIdHV2toHV2 := "test-mutate-hv2-to-hv2"
+  method TestHV2toHV2HappyCase(strategy: Types.KeyManagementStrategy, ddbClient: DDB.Types.IDynamoDBClient, kmsClient: KMS.Types.IKMSClient)
+    requires ddbClient.ValidState() && kmsClient.ValidState()
+    modifies ddbClient.Modifies, kmsClient.Modifies
+  {
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+    var keyStoreOriginal :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
+
+    // Test HV2 mutating for EC only
+    var uuidForHV2AndECMutationTest :- expect UUID.GenerateUUID();
+    var testIdForHV2AndECMutation := happyCaseIdHV2toHV2 + "-" + uuidForHV2AndECMutationTest;
+    var mutationsRequestChangeHVAndEC := Types.Mutations(
+      TerminalEncryptionContext := Some(terminalEC)
+    );
+    TestMutationHappyPath.MutationRoundTripTest(
+      ddbClient := ddbClient,
+      storage := storage,
+      keyStoreAdminUnderTest := underTest,
+      strategy := strategy,
+      keyStoreTerminal := keyStoreOriginal,
+      branchKeyIdentifier := testIdForHV2AndECMutation,
+      mutationsRequest := mutationsRequestChangeHVAndEC,
+      versionCount := 1,
+      initialHV := KeyStoreTypes.HierarchyVersion.v2,
+      doNotVersion := true
+    );
+
+    // Test HV2 mutating for kmsArn
+    var uuidForHV2AndKmsArnMutationTest :- expect UUID.GenerateUUID();
+    var testIdForHV2AndKmsArnMutation := happyCaseIdHV2toHV2 + "-" + uuidForHV2AndKmsArnMutationTest;
+    var mutationsRequestChangeHVAndKmsArn := Types.Mutations(
+      TerminalKmsArn := Some(terminalKmsId)
+    );
+    var keyStoreTerminalForHV2AndKmsArnMutation :- expect Fixtures.DefaultKeyStore(
+      kmsId:=terminalKmsId,
+      ddbClient?:=Some(ddbClient),
+      kmsClient?:=Some(kmsClient));
+    TestMutationHappyPath.MutationRoundTripTest(
+      ddbClient := ddbClient,
+      storage := storage,
+      keyStoreAdminUnderTest := underTest,
+      strategy := strategy,
+      keyStoreTerminal := keyStoreTerminalForHV2AndKmsArnMutation,
+      branchKeyIdentifier := testIdForHV2AndKmsArnMutation,
+      mutationsRequest := mutationsRequestChangeHVAndKmsArn,
+      versionCount := 1,
+      initialHV := KeyStoreTypes.HierarchyVersion.v2,
+      doNotVersion := true
+    );
+
+    // Test HV2 mutating for EC and kmsArn
+    var uuidForHV2KmsArnAndECMutationTest :- expect UUID.GenerateUUID();
+    var testIdForHV2KmsArnAndECMutation := happyCaseIdHV2toHV2 + "-" + uuidForHV2KmsArnAndECMutationTest;
+    var mutationsRequestChangeHVKmsArnAndEC := Types.Mutations(
+      TerminalEncryptionContext := Some(terminalEC),
+      TerminalKmsArn := Some(terminalKmsId)
+    );
+    var keyStoreTerminalForHV2KmsArnAndECMutation :- expect Fixtures.DefaultKeyStore(
+      kmsId:=terminalKmsId,
+      ddbClient?:=Some(ddbClient),
+      kmsClient?:=Some(kmsClient));
+    TestMutationHappyPath.MutationRoundTripTest(
+      ddbClient := ddbClient,
+      storage := storage,
+      keyStoreAdminUnderTest := underTest,
+      strategy := strategy,
+      keyStoreTerminal := keyStoreTerminalForHV2KmsArnAndECMutation,
+      branchKeyIdentifier := testIdForHV2KmsArnAndECMutation,
+      mutationsRequest := mutationsRequestChangeHVKmsArnAndEC,
+      versionCount := 1,
+      initialHV := KeyStoreTypes.HierarchyVersion.v2,
+      doNotVersion := true
+    );
+  }
 }

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/Mutations/TestMutationHappyPath.dfy

@@ -47,17 +47,41 @@ module {:options "/functionSyntax:4" } TestMutationHappyPath {
     modifies keyStoreTerminal.Modifies
   {
     match initialHV {
-      // TODO-HV-2-FF: Both v1 and v2 item could use AdminFixtures to create key
-      case v1 => Fixtures.CreateHappyCaseId(
-        id := branchKeyIdentifier,
-        versionCount := versionCount
-      );
-      case v2 => AdminFixtures.CreateHappyCaseId(
-        id := branchKeyIdentifier,
-        // versionCount := versionCount,
-        strategy := strategy,
-        hierarchyVersion := initialHV
-      );
+      case v1 =>
+        Fixtures.CreateHappyCaseId(
+          id := branchKeyIdentifier,
+          versionCount := versionCount
+        );
+      case v2 =>
+        // TODO-HV-2-Version: Once, we support version key for HV2 keys.
+        // Uncomment this to directly create HV2 keys with versions.
+        // AdminFixtures.CreateHappyCaseId(
+        //   id := branchKeyIdentifier,
+        //   // versionCount := versionCount,
+        //   hierarchyVersion := initialHV
+        // );
+        Fixtures.CreateHappyCaseId(
+          id := branchKeyIdentifier,
+          versionCount := versionCount
+        );
+        // TODO-HV-2-Version: Below is Temporary Logic to have HV-2 Key with multiple decrypt-only/version items
+        // Add Mutation Logic Here for Mutating from HV-1 to HV-2
+        var initInput := Types.InitializeMutationInput(
+          Identifier := branchKeyIdentifier,
+          Mutations := Types.Mutations(
+            TerminalHierarchyVersion := Some(initialHV)
+          ),
+          Strategy := Some(strategy),
+          SystemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage()),
+          DoNotVersion := Some(doNotVersion));
+        var initializeOutput :- expect keyStoreAdminUnderTest.InitializeMutation(initInput);
+        var initializeToken := initializeOutput.MutationToken;
+        var testInput := Types.ApplyMutationInput(
+          MutationToken := initializeToken,
+          PageSize := Some(24),
+          Strategy := Some(strategy),
+          SystemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage()));
+        var applyOutput :- expect keyStoreAdminUnderTest.ApplyMutation(testInput);
     }
 
     var initInput := Types.InitializeMutationInput(
@@ -101,10 +125,8 @@ module {:options "/functionSyntax:4" } TestMutationHappyPath {
     mutationsRequest: Types.Mutations,
     expectedDecryptOnlyItems: int
   )
-    requires storage.ValidState()
-    requires keyStoreTerminal.ValidState()
-    modifies storage.Modifies
-    modifies keyStoreTerminal.Modifies
+    requires storage.ValidState() && keyStoreTerminal.ValidState()
+    modifies storage.Modifies, keyStoreTerminal.Modifies
   {
     var versionQuery := KeyStoreTypes.QueryForVersionsInput(
       Identifier := branchKeyIdentifier,

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestHv2KmsExceptions.dfy

@@ -58,7 +58,7 @@ module {:options "/functionSyntax:4" } TestHv2KmsExceptions {
 
     var uuid :- expect UUID.GenerateUUID();
     var testId := testHV2MutationKMSDeniedId + "-" + uuid;
-    AdminFixtures.CreateHappyCaseId(id := testId, hierarchyVersion := KeyStoreTypes.HierarchyVersion.v2, strategy := strategy, admin? := Some(underTest));
+    AdminFixtures.CreateHappyCaseId(id := testId, hierarchyVersion := KeyStoreTypes.HierarchyVersion.v2, strategy? := Some(strategy), admin? := Some(underTest));
 
     // Mutating HV2 Branch Key with a Kms Key with ReEncyrpt Only KMS Permissions
     var mutationsRequest := Types.Mutations(TerminalKmsArn := Some(Fixtures.kmsKeyForHV1));