fix(KSA): Describe Mutation bugs

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/src/DescribeMutation.dfy

@@ -3,6 +3,7 @@
 include "../Model/AwsCryptographyKeyStoreAdminTypes.dfy"
 include "MutationStateStructures.dfy"
 include "MutationsConstants.dfy"
+include "SystemKey/Handler.dfy"
 
 module {:options "/functionSyntax:4" } DescribeMutation {
   import opened StandardLibrary
@@ -12,6 +13,8 @@ module {:options "/functionSyntax:4" } DescribeMutation {
   import Types = AwsCryptographyKeyStoreAdminTypes
   import StateStrucs = MutationStateStructures
   import M_ErrorMessages = MutationsConstants.ErrorMessages
+  import SystemKeyHandler = SystemKey.Handler
+  import KMS = Com.Amazonaws.Kms
 
   datatype InternalDescribeMutationInput = | InternalDescribeMutationInput (
     nameonly Identifier: string ,
@@ -65,7 +68,7 @@ module {:options "/functionSyntax:4" } DescribeMutation {
     var Index := fetchMutation.MutationIndex.value;
     :- Need(
       // If custom storage is really bad
-      Commitment.Identifier == Index.Identifier,
+      Commitment.Identifier == Index.Identifier && 0 < |Commitment.Identifier|,
       Types.MutationInvalidException(
         message := "The Mutation Index read from storage and the Mutation Commitment are for different Branch Key IDs."
         + " The Storage implementation is wrong, or something terrible has happened to storage."
@@ -74,7 +77,7 @@ module {:options "/functionSyntax:4" } DescribeMutation {
         + " Mutation Index Branch Key ID: " + Index.Identifier + ";"
       ));
     :- Need(
-      Commitment.UUID == Index.UUID,
+      Commitment.UUID == Index.UUID && 0 < |Commitment.UUID| ,
       Types.MutationInvalidException(
         message := "The Mutation Index read from storage and the Mutation Commitment are for different Mutations."
         + " Branch Key ID: " + input.Identifier + ";"
@@ -85,7 +88,6 @@ module {:options "/functionSyntax:4" } DescribeMutation {
       Commitment := Commitment,
       Index := Index);
     assert CommitmentAndIndex.ValidState();
-    // TODO-Mutations-GA :: Use System Key to Verify Commitment and Index
     var MutationToApply :- StateStrucs.DeserializeMutation(CommitmentAndIndex);
     var original := Types.MutableBranchKeyProperties(
       KmsArn := MutationToApply.Original.kmsArn,
@@ -99,7 +101,7 @@ module {:options "/functionSyntax:4" } DescribeMutation {
       Original := original,
       Terminal := terminal,
       Input := MutationToApply.Input,
-      SystemKey := "TrustStorage",
+      SystemKey := SystemKeyDescription(Commitment),
       CreateTime := MutationToApply.CreateTime,
       UUID := MutationToApply.UUID
     );
@@ -110,4 +112,34 @@ module {:options "/functionSyntax:4" } DescribeMutation {
       Yes := description);
     return Success(Types.DescribeMutationOutput(MutationInFlight := inFlight));
   }
+
+  const TRUST_STORAGE_str := "Trust Storage"
+  const KMS_SYM_ENC_str := "KMS Symmetric Encryption"
+  const UNKOWN_str := "Unknown"
+
+  function SystemKeyDescription(
+    MutationCommitment: KeyStoreTypes.MutationCommitment
+  ): (output: string)
+    ensures
+      && MutationCommitment.CiphertextBlob == SystemKeyHandler.TRUST_STORAGE_UTF8_BYTES
+      ==>
+        output == TRUST_STORAGE_str
+    ensures
+      && MutationCommitment.CiphertextBlob != SystemKeyHandler.TRUST_STORAGE_UTF8_BYTES
+      && KMS.Types.IsValid_CiphertextType(MutationCommitment.CiphertextBlob)
+      ==>
+        output == KMS_SYM_ENC_str
+    ensures
+      && MutationCommitment.CiphertextBlob != SystemKeyHandler.TRUST_STORAGE_UTF8_BYTES
+      && !KMS.Types.IsValid_CiphertextType(MutationCommitment.CiphertextBlob)
+      ==>
+        output == UNKOWN_str
+  {
+    if MutationCommitment.CiphertextBlob == SystemKeyHandler.TRUST_STORAGE_UTF8_BYTES
+    then TRUST_STORAGE_str
+    else
+    if KMS.Types.IsValid_CiphertextType(MutationCommitment.CiphertextBlob)
+    then KMS_SYM_ENC_str
+    else UNKOWN_str
+  }
 }

AwsCryptographicMaterialProviders/runtimes/java/src/examples/java/software/amazon/cryptography/example/Fixtures.java

@@ -87,17 +87,29 @@ public static GetItemResponse getKeyStoreDdbItem(
   }
 
   public static void cleanUpBranchKeyId(
-    KeyStorageInterface storage,
-    String branchKeyId
+    @Nullable KeyStorageInterface storage,
+    String branchKeyId,
+    boolean alsoMutation
   ) {
-    QueryForVersionsOutput versions = storage.QueryForVersions(
+    final KeyStorageInterface _storage;
+    if (storage == null) {
+      _storage =
+        StorageCheater.create(
+          Fixtures.ddbClientWest2,
+          Fixtures.TEST_KEYSTORE_NAME,
+          Fixtures.TEST_LOGICAL_KEYSTORE_NAME
+        );
+    } else {
+      _storage = storage;
+    }
+    QueryForVersionsOutput versions = _storage.QueryForVersions(
       QueryForVersionsInput
         .builder()
         .Identifier(branchKeyId)
         .PageSize(99)
         .build()
     );
-    String physicalName = storage
+    String physicalName = _storage
       .GetKeyStorageInfo(GetKeyStorageInfoInput.builder().build())
       .Name();
     versions
@@ -126,5 +138,21 @@ public static void cleanUpBranchKeyId(
       ddbClientWest2,
       false
     );
+    if (alsoMutation) {
+      DdbHelper.deleteKeyStoreDdbItem(
+        branchKeyId,
+        "branch:MUTATION_COMMITMENT",
+        physicalName,
+        ddbClientWest2,
+        false
+      );
+      DdbHelper.deleteKeyStoreDdbItem(
+        branchKeyId,
+        "beacon:MUTATION_INDEX",
+        physicalName,
+        ddbClientWest2,
+        false
+      );
+    }
   }
 }

AwsCryptographicMaterialProviders/runtimes/java/src/examples/java/software/amazon/cryptography/example/hierarchy/mutations/DescribeMutationExample.java

@@ -9,6 +9,7 @@
 import javax.annotation.Nullable;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.cryptography.example.Fixtures;
 import software.amazon.cryptography.example.hierarchy.AdminProvider;
 import software.amazon.cryptography.example.hierarchy.CreateKeyExample;
 import software.amazon.cryptography.keystoreadmin.KeyStoreAdmin;
@@ -24,7 +25,7 @@
 public class DescribeMutationExample {
 
   @Nullable
-  public static MutationToken Example(
+  public static DescribeMutationOutput Example(
     String keyStoreTableName,
     String logicalKeyStoreName,
     String branchKeyId,
@@ -54,7 +55,7 @@ public static MutationToken Example(
       System.out.println(
         "Description: " + description.MutationDetails().UUID()
       );
-      return description.MutationToken();
+      return output;
     }
     throw new RuntimeException("Key Store Admin returned nonsensical response");
   }
@@ -127,13 +128,20 @@ public static void CompleteExample(
       kmsClient
     );
 
-    MutationToken fromDescribe = Example(
+    DescribeMutationOutput describeRes = Example(
       keyStoreTableName,
       logicalKeyStoreName,
       branchKeyId,
       dynamoDbClient
     );
+    assert Objects.requireNonNull(describeRes).MutationInFlight().Yes() !=
+    null : "No mutation in flight for Branch Key ID: " + branchKeyId;
+    MutationToken fromDescribe = describeRes
+      .MutationInFlight()
+      .Yes()
+      .MutationToken();
     assert fromDescribe != null;
     assert Objects.equals(fromInit.UUID(), fromDescribe.UUID());
+    Fixtures.cleanUpBranchKeyId(null, branchKeyId, true);
   }
 }

AwsCryptographicMaterialProviders/runtimes/java/src/examples/java/software/amazon/cryptography/example/hierarchy/mutations/MutationExample.java

@@ -6,6 +6,7 @@
 import javax.annotation.Nullable;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.cryptography.example.Fixtures;
 import software.amazon.cryptography.example.hierarchy.AdminProvider;
 import software.amazon.cryptography.keystoreadmin.KeyStoreAdmin;
 import software.amazon.cryptography.keystoreadmin.model.InitializeMutationInput;
@@ -93,5 +94,6 @@ public static void main(final String[] args) {
       null,
       null
     );
+    Fixtures.cleanUpBranchKeyId(null, branchKeyId, true);
   }
 }

AwsCryptographicMaterialProviders/runtimes/java/src/examples/java/software/amazon/cryptography/example/hierarchy/mutations/MutationResumeExample.java

@@ -7,6 +7,7 @@
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.kms.KmsClient;
 import software.amazon.cryptography.example.DdbHelper;
+import software.amazon.cryptography.example.Fixtures;
 import software.amazon.cryptography.example.hierarchy.AdminProvider;
 import software.amazon.cryptography.keystoreadmin.KeyStoreAdmin;
 import software.amazon.cryptography.keystoreadmin.model.ApplyMutationResult;
@@ -190,5 +191,6 @@ public static void main(final String[] args) {
       null,
       null
     );
+    Fixtures.cleanUpBranchKeyId(null, branchKeyId, true);
   }
 }

AwsCryptographicMaterialProviders/runtimes/java/src/examples/java/software/amazon/cryptography/example/hierarchy/mutations/MutationSystemKeyKMSExample.java

@@ -16,8 +16,6 @@
 import software.amazon.cryptography.keystoreadmin.model.ApplyMutationResult;
 import software.amazon.cryptography.keystoreadmin.model.InitializeMutationInput;
 import software.amazon.cryptography.keystoreadmin.model.KeyManagementStrategy;
-import software.amazon.cryptography.keystoreadmin.model.KeyStoreAdminException;
-import software.amazon.cryptography.keystoreadmin.model.MutationInvalidException;
 import software.amazon.cryptography.keystoreadmin.model.MutationToken;
 import software.amazon.cryptography.keystoreadmin.model.MutationVerificationException;
 import software.amazon.cryptography.keystoreadmin.model.Mutations;

AwsCryptographicMaterialProviders/runtimes/java/src/examples/java/software/amazon/cryptography/example/hierarchy/mutations/MutationsSystemKeyTrustExample.java

@@ -12,8 +12,6 @@
 import software.amazon.cryptography.keystoreadmin.model.ApplyMutationResult;
 import software.amazon.cryptography.keystoreadmin.model.InitializeMutationInput;
 import software.amazon.cryptography.keystoreadmin.model.KeyManagementStrategy;
-import software.amazon.cryptography.keystoreadmin.model.KeyStoreAdminException;
-import software.amazon.cryptography.keystoreadmin.model.MutationInvalidException;
 import software.amazon.cryptography.keystoreadmin.model.MutationToken;
 import software.amazon.cryptography.keystoreadmin.model.MutationVerificationException;
 import software.amazon.cryptography.keystoreadmin.model.Mutations;

AwsCryptographicMaterialProviders/runtimes/java/src/testExamples/java/software/amazon/cryptography/example/hierarchy/ExampleTests.java

@@ -85,6 +85,6 @@ public void End2EndTests() {
       Fixtures.TEST_KEYSTORE_NAME,
       Fixtures.TEST_LOGICAL_KEYSTORE_NAME
     );
-    Fixtures.cleanUpBranchKeyId(storage, branchKeyId);
+    Fixtures.cleanUpBranchKeyId(storage, branchKeyId, true);
   }
 }

AwsCryptographicMaterialProviders/runtimes/java/src/testExamples/java/software/amazon/cryptography/example/hierarchy/mutations/DescribeMutationTest.java

@@ -1,9 +1,17 @@
 package software.amazon.cryptography.example.hierarchy.mutations;
 
+import static software.amazon.cryptography.example.hierarchy.mutations.DescribeMutationExample.Example;
+import static software.amazon.cryptography.example.hierarchy.mutations.DescribeMutationExample.InitMutation;
+
+import java.util.Collections;
+import org.testng.Assert;
 import org.testng.annotations.Test;
 import software.amazon.cryptography.example.Fixtures;
 import software.amazon.cryptography.example.StorageCheater;
+import software.amazon.cryptography.example.hierarchy.CreateKeyExample;
 import software.amazon.cryptography.keystore.KeyStorageInterface;
+import software.amazon.cryptography.keystoreadmin.model.DescribeMutationOutput;
+import software.amazon.cryptography.keystoreadmin.model.MutationToken;
 import software.amazon.cryptography.keystoreadmin.model.SystemKey;
 import software.amazon.cryptography.keystoreadmin.model.TrustStorage;
 
@@ -34,6 +42,92 @@ public void test() {
       Fixtures.ddbClientWest2,
       Fixtures.kmsClientWest2
     );
-    Fixtures.cleanUpBranchKeyId(storage, branchKeyId);
+    Fixtures.cleanUpBranchKeyId(storage, branchKeyId, true);
+  }
+
+  @Test
+  public void TestTrustStorageDescription() {
+    SystemKey systemKey = SystemKey
+      .builder()
+      .trustStorage(TrustStorage.builder().build())
+      .build();
+    KeyStorageInterface storage = StorageCheater.create(
+      Fixtures.ddbClientWest2,
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME
+    );
+    final String branchKeyId =
+      testPrefix + java.util.UUID.randomUUID().toString();
+    CreateKeyExample.CreateKey(
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME,
+      Fixtures.KEYSTORE_KMS_ARN,
+      branchKeyId,
+      Fixtures.ddbClientWest2
+    );
+    MutationToken fromInit = InitMutation(
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME,
+      Fixtures.POSTAL_HORN_KEY_ARN,
+      branchKeyId,
+      systemKey,
+      Fixtures.ddbClientWest2,
+      Fixtures.kmsClientWest2
+    );
+    DescribeMutationOutput describeRes = Example(
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME,
+      branchKeyId,
+      Fixtures.ddbClientWest2
+    );
+    Assert.assertEquals(
+      describeRes.MutationInFlight().Yes().MutationDetails().SystemKey(),
+      "Trust Storage"
+    );
+    Fixtures.cleanUpBranchKeyId(storage, branchKeyId, true);
+  }
+
+  @Test
+  public void TestKmsSymEncDescription() {
+    //noinspection unchecked
+    SystemKey systemKey = MutationsProvider.KmsSystemKey(
+      Fixtures.POSTAL_HORN_KEY_ARN,
+      Fixtures.kmsClientWest2,
+      Collections.EMPTY_LIST
+    );
+    KeyStorageInterface storage = StorageCheater.create(
+      Fixtures.ddbClientWest2,
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME
+    );
+    final String branchKeyId =
+      testPrefix + java.util.UUID.randomUUID().toString();
+    CreateKeyExample.CreateKey(
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME,
+      Fixtures.KEYSTORE_KMS_ARN,
+      branchKeyId,
+      Fixtures.ddbClientWest2
+    );
+    MutationToken fromInit = InitMutation(
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME,
+      Fixtures.POSTAL_HORN_KEY_ARN,
+      branchKeyId,
+      systemKey,
+      Fixtures.ddbClientWest2,
+      Fixtures.kmsClientWest2
+    );
+    DescribeMutationOutput describeRes = Example(
+      Fixtures.TEST_KEYSTORE_NAME,
+      Fixtures.TEST_LOGICAL_KEYSTORE_NAME,
+      branchKeyId,
+      Fixtures.ddbClientWest2
+    );
+    Assert.assertEquals(
+      describeRes.MutationInFlight().Yes().MutationDetails().SystemKey(),
+      "KMS Symmetric Encryption"
+    );
+    Fixtures.cleanUpBranchKeyId(storage, branchKeyId, true);
   }
 }

AwsCryptographicMaterialProviders/runtimes/java/src/testExamples/java/software/amazon/cryptography/example/hierarchy/mutations/MutationKmsAccessOriginalInFlightTest.java

@@ -219,7 +219,7 @@ public void test() throws InterruptedException {
     }
 
     // Clean Up
-    Fixtures.cleanUpBranchKeyId(storage, branchKeyId);
+    Fixtures.cleanUpBranchKeyId(storage, branchKeyId, false);
     Assert.assertTrue(
       (exceptions.size() == 1),
       "Only 1 exceptions should have been thrown. But got " +

AwsCryptographicMaterialProviders/runtimes/java/src/testExamples/java/software/amazon/cryptography/example/hierarchy/mutations/MutationKmsAccessTerminalInFlightTest.java

@@ -210,7 +210,7 @@ public void test() {
     }
 
     // Clean Up
-    Fixtures.cleanUpBranchKeyId(storage, branchKeyId);
+    Fixtures.cleanUpBranchKeyId(storage, branchKeyId, false);
     Assert.assertTrue(
       (exceptions.size() == 2),
       "Only two exceptions should have been thrown. But got " +

AwsCryptographicMaterialProviders/runtimes/java/src/testExamples/java/software/amazon/cryptography/example/hierarchy/mutations/ScanForInFlightMutationsTest.java

@@ -16,11 +16,14 @@ public void testScanForInFlightMutations() {
     PageResult actual = underTest.scanForMutationLock(null);
     assert actual.lastEvaluatedKey() != null;
     assert actual.inFlightMutations().isEmpty();
-    while (actual.lastEvaluatedKey() != null) {
+    final short pageLimit = 5;
+    short pageIndex = 0;
+    while (actual.lastEvaluatedKey() != null && pageIndex < pageLimit) {
       actual = underTest.scanForMutationLock(actual.lastEvaluatedKey());
       if (!actual.inFlightMutations().isEmpty()) {
         System.out.println(actual.inFlightMutations());
       }
+      pageIndex++;
     }
   }
 }