From 3e77feb7cf5b1f139a1e4b78bda29a93283fd1fe Mon Sep 17 00:00:00 2001
From: Andrew Jewell <107044381+ajewellamz@users.noreply.github.com>
Date: Thu, 30 May 2024 15:28:05 -0400
Subject: [PATCH] chore: spec and test for ARN unchanged in VersionKey (#361)

* chore: spec and test for ARN unchanged in VersionKey
---
 .../AwsCryptographyKeyStore/test/Fixtures.dfy |  2 ++
 .../test/TestVersionKey.dfy                   | 33 ++++++++++++++-----
 aws-encryption-sdk-specification              |  2 +-
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
index e7e88f8a1..46416ee20 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
@@ -82,6 +82,8 @@ module Fixtures {
   const KmsConfigWest : Types.KMSConfiguration := Types.KMSConfiguration.kmsKeyArn(MrkArnWest)
   const KmsMrkConfigEast : Types.KMSConfiguration := Types.KMSConfiguration.kmsMRKeyArn(MrkArnEast)
   const KmsMrkConfigWest : Types.KMSConfiguration := Types.KMSConfiguration.kmsMRKeyArn(MrkArnWest)
+  const KmsSrkConfigEast : Types.KMSConfiguration := Types.KMSConfiguration.kmsKeyArn(MrkArnEast)
+  const KmsSrkConfigWest : Types.KMSConfiguration := Types.KMSConfiguration.kmsKeyArn(MrkArnWest)
   const KmsMrkConfigAP : Types.KMSConfiguration := Types.KMSConfiguration.kmsMRKeyArn(MrkArnAP)
   const KmsMrkEC : Types.EncryptionContext := map[UTF8.EncodeAscii("abc") := UTF8.EncodeAscii("123")]
   const EastBranchKey : string := "MyEastBranch2"
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestVersionKey.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestVersionKey.dfy
index 0711720d9..a521aeba7 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestVersionKey.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestVersionKey.dfy
@@ -20,6 +20,7 @@ module TestVersionKey {
   import Structure
   import DDBKeystoreOperations
   import ComAmazonawsDynamodbTypes
+  import KeyStoreErrorMessages
 
   method {:test} TestVersionKey()
   {
@@ -207,17 +208,14 @@ module TestVersionKey {
       ddbClient := Some(ddbClient)
     );
 
-    var westKeyStoreConfig := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsMrkConfigWest,
-      logicalKeyStoreName := logicalKeyStoreName,
-      grantTokens := None,
-      ddbTableName := branchKeyStoreName,
-      ddbClient := Some(ddbClient)
-    );
+    var westKeyStoreConfig    := eastKeyStoreConfig.(kmsConfiguration := KmsMrkConfigWest);
+    var eastSrkKeyStoreConfig := eastKeyStoreConfig.(kmsConfiguration := KmsSrkConfigEast);
+    var westSrkKeyStoreConfig := eastKeyStoreConfig.(kmsConfiguration := KmsSrkConfigWest);
 
     var eastKeyStore :- expect KeyStore.KeyStore(eastKeyStoreConfig);
     var westKeyStore :- expect KeyStore.KeyStore(westKeyStoreConfig);
+    var eastSrkKeyStore :- expect KeyStore.KeyStore(eastSrkKeyStoreConfig);
+    var westSrkKeyStore :- expect KeyStore.KeyStore(westSrkKeyStoreConfig);
 
     // Create a new key with the WEST key store
     // We will create a use this new key per run to avoid tripping up
@@ -269,6 +267,25 @@ module TestVersionKey {
 
     expect newActiveResultWest == newActiveResultEast;
 
+    //= aws-encryption-sdk-specification/framework/branch-key-store.md#versionkey
+    // = type=test
+    // # The `kms-arn` stored in the DDB table MUST NOT change as a result of this operation,
+    //# even if the KeyStore is configured with a `KMS MRKey ARN` that does not exactly match the stored ARN.
+    var newActiveResultSrkWest :- expect westSrkKeyStore.GetActiveBranchKey(
+      Types.GetActiveBranchKeyInput(
+        branchKeyIdentifier := branchKeyId.branchKeyIdentifier
+      ));
+    // westSrkKeyStore succeeds, because ARN is still west
+    expect newActiveResultSrkWest == newActiveResultEast;
+    var newActiveResultSrkEastResult := eastSrkKeyStore.GetActiveBranchKey(
+      Types.GetActiveBranchKeyInput(
+        branchKeyIdentifier := branchKeyId.branchKeyIdentifier
+      ));
+    // eastSrkKeyStore fails, because ARN is still west
+    expect newActiveResultSrkEastResult.Failure?;
+    expect newActiveResultSrkEastResult.error ==
+           Types.KeyStoreException(message :=  KeyStoreErrorMessages.GET_KEY_ARN_DISAGREEMENT);
+
     var newActiveVersionWest :- expect UTF8.Decode(newActiveResultWest.branchKeyMaterials.branchKeyVersion);
     var newActiveVersionEast :- expect UTF8.Decode(newActiveResultEast.branchKeyMaterials.branchKeyVersion);
     expect newActiveVersionWest == newActiveVersionEast;
diff --git a/aws-encryption-sdk-specification b/aws-encryption-sdk-specification
index abdf4ea5a..82911ea39 160000
--- a/aws-encryption-sdk-specification
+++ b/aws-encryption-sdk-specification
@@ -1 +1 @@
-Subproject commit abdf4ea5a9ffaf0f03686203c99c13febf8bb1ad
+Subproject commit 82911ea396c72a71098eead11d3aca3bf058acd0