chore(dafny): KSA Add test coverage for creating a hv-2 branch key.

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy

@@ -23,48 +23,64 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
   import AwsArnParsing
   import TestGetKeys
 
-  /*
-    // If you need to re-create the MRK Branch Keys
-
-    method {:test} TestCreateMRK()
-    {
-      var ddbClient :- expect DDB.DynamoDBClient();
-
-      var keyStoreConfigEast := Types.KeyStoreConfig(
-        id := None,
-        kmsConfiguration := KmsConfigEast,
-        logicalKeyStoreName := logicalKeyStoreName,
-        grantTokens := None,
-        ddbTableName := branchKeyStoreName,
-        ddbClient := Some(ddbClient)
-      );
-
-      var keyStoreConfigWest := Types.KeyStoreConfig(
-        id := None,
-        kmsConfiguration := KmsConfigWest,
-        logicalKeyStoreName := logicalKeyStoreName,
-        grantTokens := None,
-        ddbTableName := branchKeyStoreName,
-        ddbClient := Some(ddbClient)
-      );
-
-      var keyStoreEast :- expect KeyStore.KeyStore(keyStoreConfigEast);
-      var keyStoreWest :- expect KeyStore.KeyStore(keyStoreConfigWest);
-
-      var branchKeyIdWest :- expect keyStoreWest.CreateKey(Types.CreateKeyInput(
-                                                             branchKeyIdentifier := Some(WestBranchKey),
-                                                             encryptionContext := Some(KmsMrkEC)
-                                                           ));
-
-      var branchKeyIdEast :- expect keyStoreEast.CreateKey(Types.CreateKeyInput(
-                                                             branchKeyIdentifier := Some(EastBranchKey),
-                                                             encryptionContext := Some(KmsMrkEC)
-                                                           ));
-
-      expect branchKeyIdEast.branchKeyIdentifier == EastBranchKey;
-      expect branchKeyIdWest.branchKeyIdentifier == WestBranchKey;
-    }
-  */
+  const happyCaseId := "test-happy-case-create-key-hv-1"
+
+  method {:test} TestCreateMRKForHV1()
+  {
+    var ddbClient :- expect DDB.DynamoDBClient();
+
+    var keyStoreConfigEast := Types.KeyStoreConfig(
+      id := None,
+      kmsConfiguration := KmsConfigEast,
+      logicalKeyStoreName := logicalKeyStoreName,
+      grantTokens := None,
+      storage := Some(
+        Types.ddb(
+          Types.DynamoDBTable(
+            ddbTableName := branchKeyStoreName,
+            ddbClient := Some(ddbClient)
+          )))
+    );
+
+    var keyStoreConfigWest := Types.KeyStoreConfig(
+      id := None,
+      kmsConfiguration := KmsConfigWest,
+      logicalKeyStoreName := logicalKeyStoreName,
+      grantTokens := None,
+      storage := Some(
+        Types.ddb(
+          Types.DynamoDBTable(
+            ddbTableName := branchKeyStoreName,
+            ddbClient := Some(ddbClient)
+          )))
+    );
+
+    // Create key with Custom EC & Branch Key Identifier
+    var uuid :- expect UUID.GenerateUUID();
+    var branchKeyIdWest := happyCaseId + "-" + "west" + "-" + uuid;
+    // print branchKeyIdWest;
+    var branchKeyIdEast := happyCaseId + "-" + "east" + "-" + uuid;
+    // print branchKeyIdEast;
+
+    var keyStoreEast :- expect KeyStore.KeyStore(keyStoreConfigEast);
+    var keyStoreWest :- expect KeyStore.KeyStore(keyStoreConfigWest);
+
+    var branchKeyIdWest? :- expect keyStoreWest.CreateKey(Types.CreateKeyInput(
+                                                            branchKeyIdentifier := Some(branchKeyIdWest),
+                                                            encryptionContext := Some(KmsMrkEC)
+                                                          ));
+
+    var branchKeyIdEast? :- expect keyStoreEast.CreateKey(Types.CreateKeyInput(
+                                                            branchKeyIdentifier := Some(branchKeyIdEast),
+                                                            encryptionContext := Some(KmsMrkEC)
+                                                          ));
+
+    expect branchKeyIdEast?.branchKeyIdentifier == branchKeyIdEast;
+    expect branchKeyIdWest?.branchKeyIdentifier == branchKeyIdWest;
+
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyIdWest, ddbClient:=ddbClient);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyIdEast, ddbClient:=ddbClient);
+  }
 
   method {:test} TestCreateBranchAndBeaconKeys()
   {
@@ -81,12 +97,12 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
           Types.DynamoDBTable(
             ddbTableName := branchKeyStoreName,
             ddbClient := Some(ddbClient)
-          ))),
-      keyManagement := Some(
-        Types.kms(
-          Types.AwsKms(
-            kmsClient := Some(kmsClient)
           )))
+      // keyManagement := Some(
+      //   Types.kms(
+      //     Types.AwsKms(
+      //       kmsClient := Some(kmsClient)
+      //     )))
     );
 
     var keyStore :- expect KeyStore.KeyStore(keyStoreConfig);

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy

@@ -72,11 +72,29 @@ module TestGetKeys {
     testBranchKeyVersionHappyCase(keyStore, hv2BranchKeyId, hv2BranchKeyVersion, hv2BranchKeyIdActiveVersionUtf8Bytes);
   }
 
+  // This is a static test case to Get Branch Keys created with Mrk Keys
+  method {:test} TestGetActiveMrkKey()
+  {
+    VerifyGetActiveMrkKey(
+      KmsConfigEast := KmsConfigEast,
+      KmsConfigWest := KmsConfigWest,
+      KmsMrkConfigEast := KmsMrkConfigEast,
+      KmsMrkConfigWest := KmsMrkConfigWest
+    );
+  }
+
   // TODO-HV2-M1: Add MRK test for hv2
-  method {:test} {:isolate_assertions} TestGetActiveMrkKey()
+  method {:isolate_assertions} VerifyGetActiveMrkKey(
+    KmsConfigEast : Types.KMSConfiguration,
+    KmsConfigWest : Types.KMSConfiguration,
+    KmsMrkConfigEast : Types.KMSConfiguration,
+    KmsMrkConfigWest : Types.KMSConfiguration
+  )
   {
     var ddbClient :- expect DDB.DynamoDBClient();
 
+    print "Running MRK";
+
     var eastKeyStoreConfig := Types.KeyStoreConfig(
       id := None,
       kmsConfiguration := KmsConfigEast,

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/AdminFixtures.dfy

@@ -44,7 +44,8 @@ module {:options "/functionSyntax:4" } AdminFixtures {
       ddbClient := ddbClient,
       logicalKeyStoreName := logicalName,
       ddbTableNameUtf8 := physicalNameUtf8,
-      logicalKeyStoreNameUtf8 := logicalNameUtf8);
+      logicalKeyStoreNameUtf8 := logicalNameUtf8
+    );
 
     var underTestConfig := Types.KeyStoreAdminConfig(
       logicalKeyStoreName := logicalName,

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

@@ -82,6 +82,7 @@ module {:options "/functionSyntax:4" } TestDecryptEncryptStrat {
 
     var timestamp :- expect Time.GetCurrentTimeStamp();
     var newCustomEC: KeyStoreTypes.EncryptionContextString := map["Koda" := timestamp];
+    // TODO-HV-2-M2: We should Mutate with Mrk Key as well
     var mutationsRequest := Types.Mutations(TerminalEncryptionContext := Some(newCustomEC));
     var initInput := Types.InitializeMutationInput(
       Identifier := testId,

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy → AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy

@@ -8,9 +8,10 @@ include "../../AwsCryptographyKeyStore/test/TestGetKeys.dfy"
 include "../../AwsCryptographyKeyStore/Model/AwsCryptographyKeyStoreTypes.dfy"
 include "AdminFixtures.dfy"
 
-// TODO-HV-2-M1 : remove HV-1 restriction for CreateKey
 // TODO-HV-2-M2 : remove HV-1 restriction for Mutations
-module {:options "/functionSyntax:4" } TestAdminHV1Only {
+// TODO-HV-2-M3 : rename TestAdminHV2 -> TestAdminCreateKeyHV2
+//    and move HappyCases for version & mutations to it's own test modules
+module {:options "/functionSyntax:4" } TestAdminHV2 {
   import Types = AwsCryptographyKeyStoreAdminTypes
   import KeyStoreAdmin
   import KeyStore
@@ -26,9 +27,7 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
   import AdminFixtures
   import TestGetKeys
 
-  const happyCaseId := "test-create-key-hv-2-happy-case"
-
-  method {:test} TestCreateKeyForHV2HappyCase()
+  method {:test} TestCreateKeyHV2HappyCase()
   {
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
@@ -61,7 +60,8 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     var _ := CleanupItems.DeleteBranchKey(Identifier:=createKeyOutput?.Identifier, ddbClient:=ddbClient);
   }
 
-  method {:test} TestCreateKeyForHV2CreateOptions() {
+  const happyCaseId := "test-create-key-hv2-happy-case"
+  method {:test} TestCreateKeyHV2WithOptions() {
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
     var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
@@ -100,6 +100,109 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyId, ddbClient:=ddbClient);
   }
 
+  // TODO-HV-2-M2?: Document that BKSA ONLY calls KMS/DDB for the region of the supplied clients.
+  // or Ensure BKSA has similar behavior as BKS in terms of MRK region.
+  /*
+  method {:test} TestCreateMRKForHV2()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+
+    // Create key with Custom EC & Branch Key Identifier
+    var uuid :- expect UUID.GenerateUUID();
+    var branchKeyIdWest := happyCaseId + "-" + WestBranchKey + "-" + uuid;
+    var branchKeyIdEast := happyCaseId + "-" + EastBranchKey + "-" + uuid;
+
+    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
+
+    var westOutput? :- expect underTest.CreateKey(Types.CreateKeyInput(
+                                                    Identifier := Some(branchKeyIdWest),
+                                                    EncryptionContext := Some(customEC),
+                                                    KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnWest),
+                                                    Strategy := Some(strategy),
+                                                    HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+                                                  ));
+    expect westOutput?.Identifier == branchKeyIdWest;
+    expect westOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
+    print branchKeyIdWest + "\n";
+
+    var eastOutput? :- expect underTest.CreateKey(Types.CreateKeyInput(
+                                                    Identifier := Some(branchKeyIdEast),
+                                                    EncryptionContext := Some(customEC),
+                                                    KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnEast),
+                                                    Strategy := Some(strategy),
+                                                    HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+                                                  ));
+    expect eastOutput?.Identifier == branchKeyIdEast;
+    expect eastOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
+    print branchKeyIdEast + "\n";
+  }
+  */
+
+  const testCreateKeyWithBKIDandNoECFailsId := "create-key-hv2-bkid-no-ec-fail"
+  method {:test} TestCreateKeyBKIDNoECFails()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+
+    var uuid :- expect UUID.GenerateUUID();
+    var branchKeyId := testCreateKeyWithBKIDandNoECFailsId + "-" + uuid;
+
+    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
+
+    var input := Types.CreateKeyInput(
+      Identifier := Some(branchKeyId),
+      EncryptionContext := None,
+      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+      Strategy := Some(strategy),
+      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+    );
+
+    var createKeyOutput := underTest.CreateKey(input);
+    expect createKeyOutput.Failure?, "Expected a failure when provided with Branch Key Identifier with no Custom Encryption Context";
+  }
+
+  method {:test} TestCreateKeyCustomECOnly()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+
+    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
+
+    var input := Types.CreateKeyInput(
+      Identifier := None,
+      EncryptionContext := Some(customEC),
+      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+      Strategy := Some(strategy),
+      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+    );
+
+    var createKeyOutput? :- expect underTest.CreateKey(input);
+    expect createKeyOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
+
+    // Get branch key items from storage
+    TestGetKeys.VerifyGetKeys(
+      identifier := createKeyOutput?.Identifier,
+      keyStore := keyStore,
+      storage := storage
+    );
+
+    // Since this process uses a read DDB table,
+    // the number of records will forever increase.
+    // To avoid this, remove the items.
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=createKeyOutput?.Identifier, ddbClient:=ddbClient);
+  }
+
   // TODO-HV-2-M2 : Probably make this a happy test?
   const testMutateForHV2FailsCaseId := "dafny-initialize-mutation-hv-2-rejection"
   method {:test} TestMutateForHV2Fails()