From 9e1c922784e8751400a0b4abeaa7b25f311818dd Mon Sep 17 00:00:00 2001
From: Andy Jewell <ajewell@amazon.com>
Date: Tue, 28 Jan 2025 07:25:09 -0500
Subject: [PATCH] m

---
 .../src/Keyrings/AwsKms/AwsKmsMrkKeyring.dfy                    | 2 ++
 .../src/Keyrings/AwsKms/AwsKmsRsaKeyring.dfy                    | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsMrkKeyring.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsMrkKeyring.dfy
index ba5af7728..e837ff616 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsMrkKeyring.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsMrkKeyring.dfy
@@ -478,6 +478,8 @@ module AwsKmsMrkKeyring {
         Materials.DecryptionMaterialsWithoutPlaintextDataKey(materials),
         Types.AwsCryptographicMaterialProvidersException( message := "Keyring received decryption materials that already contain a plaintext data key."));
 
+      :- OkForDecrypt(awsKmsArn, awsKmsKey);
+
       //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-mrk-keyring.md#ondecrypt
       //# The set of encrypted data keys MUST first be filtered to match this
       //# keyring’s configuration.
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsRsaKeyring.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsRsaKeyring.dfy
index 44374aa23..1a8a8448c 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsRsaKeyring.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/AwsKms/AwsKmsRsaKeyring.dfy
@@ -223,6 +223,8 @@ module AwsKmsRsaKeyring {
         Types.AwsCryptographicMaterialProvidersException(
           message := "Keyring received decryption materials that already contain a plaintext data key."));
 
+      :- AwsKmsUtils.OkForDecrypt(awsKmsArn, awsKmsKey);
+
       :- Need(
         input.materials.algorithmSuite.signature.None?,
         Types.AwsCryptographicMaterialProvidersException(