fix(Mutations): KMS Exception improvements

There are still more refactoring we could do in
MutationsErrorRefinement,
but this should be sufficient for now.

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/DefaultKeyStorageInterface.dfy

@@ -1679,6 +1679,7 @@ module DefaultKeyStorageInterface {
     // that can hold either Opaque or DDB Error
     match e {
       case Opaque(obj) => Types.Opaque(obj) //https://github.com/smithy-lang/smithy-dafny/issues/450#issuecomment-2322149920
+      case OpaqueWithText(obj, objMessage) => Types.OpaqueWithText(obj, objMessage)
       case IdempotentParameterMismatchException(Message) => Types.KeyStorageException(
         message :=
           "DDB through an exception for " +  storageOperation + "'s " + ddbOperation + ". Table Name: "

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/KMSKeystoreOperations.dfy

@@ -300,7 +300,7 @@ module {:options "/functionSyntax:4" } KMSKeystoreOperations {
     grantTokens: KMS.GrantTokenList,
     kmsClient: KMS.IKMSClient
   )
-    returns (res: Result<KMS.ReEncryptResponse, Types.Error>)
+    returns (res: Result<KMS.ReEncryptResponse, KmsError>)
     requires
       && Structure.BranchKeyContext?(sourceEncryptionContext)
       && Structure.BranchKeyContext?(destinationEncryptionContext)
@@ -349,7 +349,7 @@ module {:options "/functionSyntax:4" } KMSKeystoreOperations {
   {
     :- Need(
       KMS.IsValid_CiphertextType(ciphertext),
-      Types.KeyStoreException(
+      Types.KeyManagementException(
         message := "Invalid KMS ciphertext.")
     );
 
@@ -372,20 +372,20 @@ module {:options "/functionSyntax:4" } KMSKeystoreOperations {
     :- Need(
       && reEncryptResponse.SourceKeyId.Some?
       && reEncryptResponse.SourceKeyId.value == sourceKmsArn,  //kmsKeyArn
-      Types.KeyStoreException(
+      Types.KeyManagementException(
         message := "Invalid response from KMS ReEncrypt:: Invalid Source Key Id")
     );
     :- Need(
       && reEncryptResponse.KeyId.Some?
       && reEncryptResponse.KeyId.value == destinationKmsArn, // kmsKeyArn,
-      Types.KeyStoreException(
+      Types.KeyManagementException(
         message := "Invalid response from KMS ReEncrypt:: Invalid Destination Key Id")
     );
 
     :- Need(
       && reEncryptResponse.CiphertextBlob.Some?
       && KMS.IsValid_CiphertextType(reEncryptResponse.CiphertextBlob.value),
-      Types.KeyStoreException(
+      Types.KeyManagementException(
         message := "Invalid response from AWS KMS ReEncrypt: Invalid ciphertext.")
     );
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/src/MutationErrorRefinement.dfy

@@ -8,6 +8,8 @@ module {:options "/functionSyntax:4" } MutationErrorRefinement {
   import KeyStoreTypes = AwsCryptographyKeyStoreAdminTypes.AwsCryptographyKeyStoreTypes
   import KMSKeystoreOperations
   import KMS = Com.Amazonaws.Kms
+  import StandardLibrary.String
+  import Structure
 
   function ExtractKmsOpaque(
     error: KMSKeystoreOperations.KmsError
@@ -23,6 +25,7 @@ module {:options "/functionSyntax:4" } MutationErrorRefinement {
       case ComAmazonawsKms(comAmazonawsKms: KMS.Types.Error) =>
         match comAmazonawsKms {
           case Opaque(obj) => Some(comAmazonawsKms)
+          case OpaqueWithText(obj, objMessage) => Some(comAmazonawsKms)
           case _ => None
         }
     }
@@ -38,93 +41,130 @@ module {:options "/functionSyntax:4" } MutationErrorRefinement {
       case ComAmazonawsKms(comAmazonawsKms: KMS.Types.Error) =>
         match comAmazonawsKms {
           case Opaque(obj) => None
+          case OpaqueWithText(obj, objMessage) => Some(objMessage)
           case _ => comAmazonawsKms.message
         }
     }
   }
 
-  function DefaultKmsErrorMessage(
+  function ParsedErrorContext(
     nameonly localOperation: string,
     nameonly kmsOperation: string,
     nameonly identifier: string,
-    nameonly kmsArn: string,
-    nameonly while?: Option<string> := None,
+    nameonly itemType: string,
     nameonly errorMessage?: Option<string> := None
   ): (message: string)
   {
-    "KMS through an exception for "
-    + localOperation + "'s " + kmsOperation
-    + (if while?.Some? then " while " + while?.value else ".")
-    + " KMS ARN: " + kmsArn
-    + "\tBranch Key ID: " + identifier
+    "MPL Operation: " + localOperation + ";"
+    + " KMS Operation: "  + kmsOperation + ";"
+    + " Branch Key ID: " + identifier + ";"
+    + " Branch Key Type:  " + itemType + ";"
     + "\nKMS Message: " + errorMessage?.UnwrapOr("")
   }
 
   function VerifyActiveException(
-    branchKeyItem: KeyStoreTypes.EncryptedHierarchicalKey,
-    error: KMSKeystoreOperations.KmsError
+    nameonly branchKeyItem: KeyStoreTypes.EncryptedHierarchicalKey,
+    nameonly error: KMSKeystoreOperations.KmsError,
+    nameonly localOperation: string := "InitializeMutation",
+    nameonly kmsOperation: string := "ReEncrypt"
   ): (output: Types.Error)
     requires branchKeyItem.Type.ActiveHierarchicalSymmetricVersion?
   {
-    var message := DefaultKmsErrorMessage(
-                     localOperation := "InitializeMutation",
-                     kmsOperation := "ReEncrypt",
-                     identifier := branchKeyItem.Identifier,
-                     kmsArn := branchKeyItem.KmsArn,
-                     while? := Some("verifying the Active Branch Key. Do you have permission for the original KMS ARN?"),
-                     errorMessage? := ExtractMessageFromKmsError(error));
+    //TODO Mutations-FF :: Decrypt/Encrypt Strategy will need to refactor this
+    var opaqueKmsError? := ExtractKmsOpaque(error);
+    var kmsErrorMessage? := ExtractMessageFromKmsError(error);
+    var errorContext := ParsedErrorContext(
+                          localOperation := localOperation,
+                          kmsOperation := kmsOperation,
+                          identifier := branchKeyItem.Identifier,
+                          itemType := Structure.BRANCH_KEY_ACTIVE_TYPE,
+                          errorMessage? := kmsErrorMessage?);
+    var message :=
+      "Key Management denied access to the Active Branch Key."
+      + " Mutation is halted. Check access to KMS ARN: " + branchKeyItem.KmsArn  + " ."
+      + "\n" + errorContext;
     Types.MutationFromException(message := message)
   }
 
   function VerifyTerminalException(
     branchKeyItem: KeyStoreTypes.EncryptedHierarchicalKey,
-    error: KMSKeystoreOperations.KmsError
+    error: KMSKeystoreOperations.KmsError,
+    nameonly localOperation: string := "ApplyMutation",
+    nameonly kmsOperation: string := "ReEncrypt"
   ): (output: Types.Error)
     requires branchKeyItem.Type.HierarchicalSymmetricVersion?
   {
-    var message := DefaultKmsErrorMessage(
-                     localOperation := "ApplyMutation",
-                     kmsOperation := "ReEncrypt",
-                     identifier := branchKeyItem.Identifier,
-                     kmsArn := branchKeyItem.KmsArn,
-                     while? := Some("verifying a Version with terminal properities."
-                                    + " Do you have permission for the terminal KMS ARN?"
-                                    + " Version (Decrypt Only): " + branchKeyItem.Type.HierarchicalSymmetricVersion.Version),
-                     errorMessage? := ExtractMessageFromKmsError(error));
+    var opaqueKmsError? := ExtractKmsOpaque(error);
+    var kmsErrorMessage? := ExtractMessageFromKmsError(error);
+    var errorContext := ParsedErrorContext(
+                          localOperation := localOperation,
+                          kmsOperation := kmsOperation,
+                          identifier := branchKeyItem.Identifier,
+                          itemType := branchKeyItem.Type.HierarchicalSymmetricVersion.Version,
+                          errorMessage? := kmsErrorMessage?);
+    var message :=
+      "Key Management denied access to an already mutated item."
+      + " Mutation is halted. Check access to KMS ARN: " + branchKeyItem.KmsArn  + " ."
+      + "\n" + errorContext;
     Types.MutationToException(message := message)
   }
 
-  // https://github.com/smithy-lang/smithy-dafny/issues/609
-  // TODO-Mutations-GA :: Once we can get a string from KMS Oapque,
-  // we can check it for ReEncryptTo or ReEncryptFrom.
-  // Than, this function can return
-  // MutationToException or MutationFromException
-  function MutateException(
-    branchKeyItem: KeyStoreTypes.EncryptedHierarchicalKey,
-    error: KMSKeystoreOperations.KmsError,
-    terminalKmsArn: string
-  ): (output: Types.Error)
-    requires branchKeyItem.Type.HierarchicalSymmetricVersion? || branchKeyItem.Type.ActiveHierarchicalSymmetricBeacon?
+  // It would be nice if this was a function instead of a method,
+  // but the nested if logic in a function is a PITA...
+  // TODO-Mutations-DoNotVersion :: ActiveHierarchicalSymmetricVersion will need to be handled
+  method MutateExceptionParse(
+    nameonly branchKeyItem: KeyStoreTypes.EncryptedHierarchicalKey,
+    nameonly error: KMSKeystoreOperations.KmsError,
+    nameonly terminalKmsArn: string,
+    nameonly localOperation: string := "ApplyMutation",
+    nameonly kmsOperation: string := "ReEncrypt"
+  )
+    returns (output: Types.Error)
   {
-    var while? :=
-      if branchKeyItem.Type.HierarchicalSymmetricVersion?
-      then Some("mutating a Version."
-                + " Do you have permission for the original and terminal KMS ARN?"
-                + " Version (Decrypt Only): " + branchKeyItem.Type.HierarchicalSymmetricVersion.Version)
-      else Some("mutating the Beacon Key."
-                + " Do you have permission for the the original and terminal KMS ARN?");
-    // https://github.com/smithy-lang/smithy-dafny/issues/609
-    // If opaqueKmsError?.Some?, parse for ReEncryptTo or ReEncrytFrom
     var opaqueKmsError? := ExtractKmsOpaque(error);
-    var message := DefaultKmsErrorMessage(
-                     localOperation := "ApplyMutation",
-                     kmsOperation := "ReEncrypt",
-                     identifier := branchKeyItem.Identifier,
-                     kmsArn := "original: " + branchKeyItem.KmsArn + "\tterminal: " + terminalKmsArn,
-                     while? := while?,
-                     errorMessage? := ExtractMessageFromKmsError(error));
-    if opaqueKmsError?.Some?
-    then Types.ComAmazonawsKms(ComAmazonawsKms := opaqueKmsError?.value)
-    else Types.KeyStoreAdminException(message := message)
+    var kmsErrorMessage? := ExtractMessageFromKmsError(error);
+    var itemType := match branchKeyItem.Type {
+      case ActiveHierarchicalSymmetricVersion(version) => Structure.BRANCH_KEY_ACTIVE_TYPE
+      case HierarchicalSymmetricVersion(version) => Structure.BRANCH_KEY_TYPE_PREFIX + ":" + version.Version
+      case ActiveHierarchicalSymmetricBeacon(version) => Structure.BEACON_KEY_TYPE_VALUE
+    };
+    var errorContext := ParsedErrorContext(
+      localOperation := localOperation,
+      kmsOperation := kmsOperation,
+      identifier := branchKeyItem.Identifier,
+      itemType := itemType,
+      errorMessage? := kmsErrorMessage?);
+    // if it is an opaque KMS Error, and there is a message, it is KMS.Types.OpaqueWithText
+    if (opaqueKmsError?.Some? && kmsErrorMessage?.Some?) {
+      var hasReEncryptFrom? := String.HasSubString(kmsErrorMessage?.value, "ReEncryptFrom");
+      var hasReEncryptTo? := String.HasSubString(kmsErrorMessage?.value, "ReEncryptTo");
+      if (hasReEncryptFrom?.Some?) {
+        return Types.MutationFromException(
+            message := "Key Management denied access based on the original properties."
+            + " Mutation is halted. Check access to KMS ARN: " + branchKeyItem.KmsArn + "."
+            + "\n" + errorContext
+          );
+      }
+      if (hasReEncryptTo?.Some?) {
+        return Types.MutationToException(
+            message := "Key Management denied access based on the terminal properties."
+            + " Mutation is halted. Check access to KMS ARN: " + terminalKmsArn  + "."
+            + "\n" + errorContext
+          );
+      }
+    }
+    if (!branchKeyItem.Type.ActiveHierarchicalSymmetricVersion?) {
+      return Types.KeyStoreAdminException(
+          message := "Key Management through an exception."
+          + " Mutation is halted. Check access to KMS."
+          + "\n" + errorContext);
+    } else {
+      output := VerifyActiveException(
+        branchKeyItem := branchKeyItem,
+        error := error,
+        localOperation := localOperation,
+        kmsOperation := kmsOperation);
+    }
+    return output;
   }
 }

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/src/Mutations.dfy

@@ -304,13 +304,11 @@ module {:options "/functionSyntax:4" } Mutations {
     // --= Validate Active Branch Key
     var verifyActive? := VerifyEncryptedHierarchicalKey(
       item := activeItem,
-      keyManagerStrategy := input.keyManagerStrategy
+      keyManagerStrategy := input.keyManagerStrategy,
+      localOperation := "InitializeMutation"
     );
     if (verifyActive?.Fail?) {
-      return Failure(
-          MutationErrorRefinement.VerifyActiveException(
-            branchKeyItem := activeItem,
-            error := verifyActive?.error));
+      return Failure(verifyActive?.error);
     }
 
       // -= Assert Beacon Key is in Original
@@ -372,7 +370,8 @@ module {:options "/functionSyntax:4" } Mutations {
       originalKmsArn := MutationToApply.Terminal.kmsArn,
       terminalKmsArn := MutationToApply.Terminal.kmsArn,
       terminalEncryptionContext := ActiveEncryptionContext,
-      keyManagerStrategy := input.keyManagerStrategy
+      keyManagerStrategy := input.keyManagerStrategy,
+      localOperation := "InitializeMutation"
     );
 
     // -= Mutate Beacon Key
@@ -392,7 +391,8 @@ module {:options "/functionSyntax:4" } Mutations {
       originalKmsArn := MutationToApply.Original.kmsArn,
       terminalKmsArn := MutationToApply.Terminal.kmsArn,
       terminalEncryptionContext := BeaconEncryptionContext,
-      keyManagerStrategy := input.keyManagerStrategy
+      keyManagerStrategy := input.keyManagerStrategy,
+      localOperation := "InitializeMutation"
     );
 
     // -= Create Mutation Commitment
@@ -631,10 +631,7 @@ module {:options "/functionSyntax:4" } Mutations {
             keyManagerStrategy:= keyManagerStrategy
           );
           if (verify?.Fail?) {
-            return Failure(
-                MutationErrorRefinement.VerifyTerminalException(
-                  branchKeyItem := item,
-                  error := verify?.error));
+            return Failure(verify?.error);
           }
           logStatements := logStatements
           + [Types.MutatedBranchKeyItem(
@@ -730,14 +727,16 @@ module {:options "/functionSyntax:4" } Mutations {
 
   method VerifyEncryptedHierarchicalKey(
     nameonly item: Types.AwsCryptographyKeyStoreTypes.EncryptedHierarchicalKey,
-    nameonly keyManagerStrategy: KmsUtils.keyManagerStrat
+    nameonly keyManagerStrategy: KmsUtils.keyManagerStrat,
+    nameonly localOperation: string := "ApplyMutation"
   )
-    returns (output: Outcome<KMSKeystoreOperations.KmsError>)
+    returns (output: Outcome<Types.Error>)
     requires keyManagerStrategy.reEncrypt?
 
     requires Structure.EncryptedHierarchicalKey?(item)
     requires KmsArn.ValidKmsArn?(item.KmsArn)
     requires keyManagerStrategy.ValidState()
+    requires item.Type.ActiveHierarchicalSymmetricVersion? || item.Type.HierarchicalSymmetricVersion?
     modifies
       match keyManagerStrategy
       case reEncrypt(km) => km.kmsClient.Modifies
@@ -752,19 +751,41 @@ module {:options "/functionSyntax:4" } Mutations {
       grantTokens := keyManagerStrategy.reEncrypt.grantTokens,
       kmsClient := keyManagerStrategy.reEncrypt.kmsClient
     );
-
-    output := if throwAway?.Success? then
-      Pass
-    else
-      Fail(throwAway?.error);
+    if (
+        && throwAway?.Failure?
+        && keyManagerStrategy.reEncrypt?
+        && item.Type.ActiveHierarchicalSymmetricVersion?
+      ) {
+      var error := MutationErrorRefinement.VerifyActiveException(
+        branchKeyItem := item,
+        error := throwAway?.error,
+        localOperation := localOperation,
+        kmsOperation := "ReEncrypt");
+      return Fail(error);
+    }
+    if (
+        && throwAway?.Failure?
+        && keyManagerStrategy.reEncrypt?
+        && item.Type.HierarchicalSymmetricVersion?
+      ) {
+      var error := MutationErrorRefinement.VerifyTerminalException(
+        branchKeyItem := item,
+        error := throwAway?.error,
+        localOperation := localOperation,
+        kmsOperation := "ReEncrypt");
+      return Fail(error);
+    }
+    assert throwAway?.Success?;
+    return Pass;
   }
 
   method {:isolate_assertions} ReEncryptHierarchicalKey(
     nameonly item: Types.AwsCryptographyKeyStoreTypes.EncryptedHierarchicalKey,
     nameonly originalKmsArn: string,
     nameonly terminalKmsArn: string,
     nameonly terminalEncryptionContext: Structure.BranchKeyContext,
-    nameonly keyManagerStrategy: KmsUtils.keyManagerStrat
+    nameonly keyManagerStrategy: KmsUtils.keyManagerStrat,
+    nameonly localOperation: string := "ApplyMutation"
   )
     returns (output: Result<Types.AwsCryptographyKeyStoreTypes.EncryptedHierarchicalKey, Types.Error>)
     requires keyManagerStrategy.reEncrypt?
@@ -789,7 +810,15 @@ module {:options "/functionSyntax:4" } Mutations {
       grantTokens := keyManagerStrategy.reEncrypt.grantTokens,
       kmsClient := keyManagerStrategy.reEncrypt.kmsClient
     );
-
+    if (wrappedKey?.Failure? && !item.Type.ActiveHierarchicalSymmetricVersion?) {
+      var error := MutationErrorRefinement.MutateExceptionParse(
+        branchKeyItem := item,
+        error := wrappedKey?.error,
+        terminalKmsArn := terminalKmsArn,
+        localOperation := localOperation);
+      return Failure(error);
+    }
+    // TODO-Mutations-DoNotVersion :: ActiveHierarchicalSymmetricVersion will need to be handled
     var wrappedKey :- wrappedKey?
     .MapFailure(e => Types.Error.AwsCryptographyKeyStore(e));