feat(BKS & BKSA)!: Smithy Model for HV-2 (#1350)

BREAKING CHANGE: Branch Key Store Admin has the following breaking changes for MRK support:
- `KmsSymmetricKeyArn` has changed the MRK type from String to a `KmsMRKey`
- `KmsMRKey` is a new type that allows for proper MRK logic

Note that this MRK support is not completed yet.
The `CreateKey` and `VersionKey` operations will
throw an `UnsupportedFeatureException` if the MRK option is passed.

Author: texastony

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/Model/AwsCryptographyKeyStoreTypes.dfy

@@ -151,6 +151,9 @@ module {:extern "software.amazon.cryptography.keystore.internaldafny.types" } Aw
   datatype HierarchicalSymmetric = | HierarchicalSymmetric (
     nameonly Version: string
   )
+  datatype HierarchyVersion =
+    | v1
+    | v2
   type HmacKeyMap = map<string, Secret>
   datatype KeyManagement =
     | kms(kms: AwsKms)
@@ -924,6 +927,9 @@ module {:extern "software.amazon.cryptography.keystore.internaldafny.types" } Aw
     | BranchKeyCiphertextException (
         nameonly message: string
       )
+    | HierarchyVersionException (
+        nameonly message: string
+      )
     | KeyManagementException (
         nameonly message: string
       )

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/Model/KeyStore.smithy

@@ -62,9 +62,28 @@ service KeyStore {
   errors: [
     KeyStoreException
     VersionRaceException
+    KeyManagementException
+    BranchKeyCiphertextException
+    HierarchyVersionException
   ]
 }
 
+@documentation(
+  "The hierarchy-version of a Branch Key;
+  all items of the same Branch Key SHOULD
+  have the same hierarchy-version.
+  The hierarchy-version determines how the Branch Key Store
+  protects and validates the branch key context (BKC).")
+@enum([
+  { name: "v1", value: "1" },
+  { name: "v2", value: "2" }
+])
+string HierarchyVersion
+
+// Tony does not think that this shape will be useable, but conceptually, this is what we want
+// @documentation("For MPL 1.x, this HierarchyVersion is optional and defaults to \"1\".")
+//string HierarchyVersionDefault = "1"
+
 structure KeyStoreConfig {
 
   //= aws-encryption-sdk-specification/framework/branch-key-store.md#initialization
@@ -107,7 +126,7 @@ structure KeyStoreConfig {
   @javadoc("The DynamoDB client this Key Store uses to call Amazon DynamoDB. If None is provided and the KMS ARN is, the KMS ARN is used to determine the Region of the default client.")
   ddbClient: DdbClientReference,
   @javadoc("The KMS client this Key Store uses to call AWS KMS.  If None is provided and the KMS ARN is, the KMS ARN is used to determine the Region of the default client.")
-  kmsClient: KmsClientReference,
+  kmsClient: KmsClientReference
 }
 
 union Storage {
@@ -226,14 +245,33 @@ structure CreateKeyStoreOutput {
   tableArn: com.amazonaws.dynamodb#TableArn
 }
 
+// TODO : Dev Guide for BKS needs refactoring
 // CreateKey will create two keys to add to the key store
 // One is the branch key, which is used in the hierarchical keyring
 // The second is a beacon key that is used as a root key to
 // derive different beacon keys per beacon.
-@javadoc("Create a new Branch Key in the Key Store. Additionally create a Beacon Key that is tied to this Branch Key.")
+@javadoc(
+"Create a new Branch Key in the Branch Key Store.
+This method ONLY creates hierarchy-version-1 branch keys.
+This creates 3 items: the ACTIVE branch key item, the DECRYPT_ONLY for the ACTIVE branch key item, and the beacon key.
+In DynamoDB, the sort-key for the ACTIVE branch key is 'branch:ACTIVE`;
+the sort-key for the decrypt_only is 'branch:version:<uuid>';
+the sort-key for the beacon key is `beacon:ACTIVE'.
+The active branch key and the decrypt_only items have the same plain-text data key.
+The beacon key plain-text data key is unqiue.
+KMS is called 3 times; GenerateDataKeyWithoutPlaintext is called twice, ReEncrypt is called once.
+All three items are written to DDB by a TransactionWriteItems, conditioned on the absence of a conflicting Branch Key ID.
+See Branch Key Store Developer Guide's 'Create Branch Keys': https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/create-branch-keys.html")
 operation CreateKey {
   input: CreateKeyInput,
   output: CreateKeyOutput
+  errors: [
+    // KmsException https://sdk.amazonaws.com/java/api/2.0.0-preview-11/software/amazon/awssdk/services/kms/model/KmsException.html
+    // DynamoDbException https://sdk.amazonaws.com/java/api/2.0.0-preview-11/software/amazon/awssdk/services/dynamodb/model/DynamoDbException.html
+    KeyStoreException
+    KeyManagementException
+    HierarchyVersionException
+  ]
 }
 
 //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey
@@ -256,14 +294,35 @@ structure CreateKeyOutput {
   branchKeyIdentifier: String
 }
 
+// TODO : Dev Guide for BKS needs refactoring 
 // VersionKey will create a new branch key under the 
 // provided branchKeyIdentifier and rotate the "older" material 
 // on the key store under the branchKeyIdentifier. This operation MUST NOT
 // rotate the beacon key under the branchKeyIdentifier.
-@javadoc("Create a new ACTIVE version of an existing Branch Key in the Key Store, and set the previously ACTIVE version to DECRYPT_ONLY.")
+@javadoc(
+"Rotates an exsisting Branch Key;
+this generates a fresh AES-256 key which all future encrypts will use
+for the Key Derivation Function,
+until VersionKey is executed again.
+This method ONLY works with hierarchy-version-1 Branch Keys;
+if a hierarchy-version-2 Branch Key is encountered, the operation fails before calling KMS.
+Rotation is accomplished by first authenticating the ACTIVE branch key item via 'kms:ReEncrypt'.
+'kms:GenerateDataKeyWithoutPlaintext', followed by 'kms:ReEncrypt' is used to create a new ACTIVE and matching DECRYPT_ONLY.
+These two items are then writen to the Branch Key Store via a TransactionWriteItems;
+this only overwrites the ACTIVE item, the DECRYPT_ONLY is a new item.
+This leaves all the previous DECRYPT_ONLY items avabile to service decryption of previous rotations.
+See Branch Key Store Developer Guide's 'Rotate your active branch key': https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/rotate-branch-key.html")
 operation VersionKey {
   input: VersionKeyInput,
   output: VersionKeyOutput
+  errors: [
+    // KmsException https://sdk.amazonaws.com/java/api/2.0.0-preview-11/software/amazon/awssdk/services/kms/model/KmsException.html
+    // DynamoDbException https://sdk.amazonaws.com/java/api/2.0.0-preview-11/software/amazon/awssdk/services/dynamodb/model/DynamoDbException.html
+    KeyStoreException
+    VersionRaceException
+    KeyManagementException
+    HierarchyVersionException
+  ]  
 }
 
 @javadoc("Inputs for versioning a Branch Key.")
@@ -457,18 +516,29 @@ structure KeyManagementException {
 
 @error("client")
 @documentation("
-The cipher-text or additional authenticated data incorporated into the cipher-text,
+The cipher-text or branch key context incorporated into the cipher-text,
 such as the encryption context, is corrupted, missing, or otherwise invalid.
-For Branch Keys,
-the Encryption Context is a combination of:
-- the custom encryption context
+For branch keys,
+the branch key context (BKC) is a combination of:
+- the encryption context
 - storage identifiers (partition key, sort key, logical name)
 - metadata that binds the Branch Key to encrypted data (version)
+- create-time
+- hierarchy-version
 
 If any of the above are modified without calling KMS,
-the Branch Key's cipher-text becomes invalid.
+the branch key's cipher-text becomes invalid.
 ")
 structure BranchKeyCiphertextException {
   @required
-  message: String,
+  message: String
+}
+
+@error("client")
+@documentation(
+"The HierarchyVersion of the Branch Key is not supported by the operation.
+'hierarchy-version-2' Branch Keys can only be created or versioned (rotated) via the Branch Key Store Admin.")
+structure HierarchyVersionException {
+  @required
+  message: String
 }

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/CreateKeys.dfy

@@ -495,7 +495,12 @@ module {:options "/functionSyntax:4" } CreateKeys {
       Types.KeyStoreException(
         message := ErrorMessages.INVALID_ACTIVE_BRANCH_KEY_FROM_STORAGE)
     );
-
+    // TODO-HV-2-M3: Support Version in HV-2
+    :- Need(
+      oldActiveItem.EncryptionContext[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_VALUE_1,
+      Types.KeyStoreException(
+        message := "At this time, VersionKey ONLY supports HV-1; BK's Active Item is HV-2.")
+    );
     :- Need(
       && KMSKeystoreOperations.AttemptKmsOperation?(kmsConfiguration, oldActiveItem.EncryptionContext),
       Types.KeyStoreException(

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/Index.dfy

@@ -465,6 +465,13 @@ module {:extern "software.amazon.cryptography.keystore.internaldafny"} KeyStore
     assert Operations.ValidInternalConfig?(internalConfig);
     var client := new KeyStoreClient(internalConfig);
 
+    assume {:axiom} fresh(client.Modifies
+                          - ( if config.ddbClient.Some? then config.ddbClient.value.Modifies else {})
+                          - ( if config.kmsClient.Some? then config.kmsClient.value.Modifies else {})
+                          - ( if (config.storage.Some? && config.storage.value.custom?) then config.storage.value.custom.Modifies else {})
+                          - ( if (config.keyManagement.Some? && config.keyManagement.value.kms? && config.keyManagement.value.kms.kmsClient.Some?) then config.keyManagement.value.kms.kmsClient.value.Modifies else {})
+                          - ( if (config.storage.Some? && config.storage.value.ddb? && config.storage.value.ddb.ddbClient.Some?) then config.storage.value.ddb.ddbClient.value.Modifies else {}));
+
     return Success(client);
   }
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/Structure.dfy

@@ -41,8 +41,14 @@ module {:options "/functionSyntax:4" } Structure {
   }
 
   //Attribute Values
-  const HIERARCHY_VERSION_VALUE := "1"
+  const HIERARCHY_VERSION_VALUE_1 := "1"
+  const HIERARCHY_VERSION_VALUE_2 := "2"
+  // TODO-HV-2-M1 : Find all HIERARCHY_VERSION_VALUE and replace them with something
+  const HIERARCHY_VERSION_VALUE := HIERARCHY_VERSION_VALUE_1
+  // TODO-HV-2-M1 : Find all HIERARCHY_VERSION_ATTRIBUTE_VALUE and replace them with something
   const HIERARCHY_VERSION_ATTRIBUTE_VALUE := DDB.AttributeValue.N(HIERARCHY_VERSION_VALUE)
+  const HIERARCHY_VERSION_ATTRIBUTE_VALUE_1 := DDB.AttributeValue.N(HIERARCHY_VERSION_VALUE_1)
+  const HIERARCHY_VERSION_ATTRIBUTE_VALUE_2 := DDB.AttributeValue.N(HIERARCHY_VERSION_VALUE_2)
   const BRANCH_KEY_TYPE_PREFIX := "branch:version:"
   const BRANCH_KEY_ACTIVE_TYPE := "branch:ACTIVE"
   const BEACON_KEY_TYPE_VALUE := "beacon:ACTIVE"

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/Model/AwsCryptographyKeyStoreAdminTypes.dfy

@@ -41,10 +41,12 @@ module {:extern "software.amazon.cryptography.keystoreadmin.internaldafny.types"
     nameonly Identifier: Option<string> := Option.None ,
     nameonly EncryptionContext: Option<AwsCryptographyKeyStoreTypes.EncryptionContext> := Option.None ,
     nameonly KmsArn: KmsSymmetricKeyArn ,
-    nameonly Strategy: Option<KeyManagementStrategy> := Option.None
+    nameonly Strategy: Option<KeyManagementStrategy> := Option.None ,
+    nameonly HierarchyVersion: Option<AwsCryptographyKeyStoreTypes.HierarchyVersion> := Option.None
   )
   datatype CreateKeyOutput = | CreateKeyOutput (
-    nameonly Identifier: string
+    nameonly Identifier: string ,
+    nameonly HierarchyVersion: AwsCryptographyKeyStoreTypes.HierarchyVersion
   )
   datatype DescribeMutationInput = | DescribeMutationInput (
     nameonly Identifier: string
@@ -71,6 +73,7 @@ module {:extern "software.amazon.cryptography.keystoreadmin.internaldafny.types"
   datatype KeyManagementStrategy =
     | AwsKmsReEncrypt(AwsKmsReEncrypt: AwsCryptographyKeyStoreTypes.AwsKms)
     | AwsKmsDecryptEncrypt(AwsKmsDecryptEncrypt: AwsKmsDecryptEncrypt)
+    | AwsKmsSimple(AwsKmsSimple: AwsCryptographyKeyStoreTypes.AwsKms)
   class IKeyStoreAdminClientCallHistory {
     ghost constructor() {
       CreateKey := [];
@@ -192,16 +195,21 @@ module {:extern "software.amazon.cryptography.keystoreadmin.internaldafny.types"
     nameonly logicalKeyStoreName: string ,
     nameonly storage: AwsCryptographyKeyStoreTypes.Storage
   )
+  datatype KmsMRKey = | KmsMRKey (
+    nameonly KeyArn: ComAmazonawsKmsTypes.KeyIdType ,
+    nameonly Region: ComAmazonawsKmsTypes.RegionType
+  )
   datatype KmsSymmetricEncryption = | KmsSymmetricEncryption (
     nameonly KmsArn: ComAmazonawsKmsTypes.KeyIdType ,
     nameonly AwsKms: AwsCryptographyKeyStoreTypes.AwsKms
   )
   datatype KmsSymmetricKeyArn =
     | KmsKeyArn(KmsKeyArn: string)
-    | KmsMRKeyArn(KmsMRKeyArn: string)
-  datatype MutableBranchKeyProperties = | MutableBranchKeyProperties (
+    | KmsMRKey(KmsMRKey: KmsMRKey)
+  datatype MutableBranchKeyContext = | MutableBranchKeyContext (
     nameonly KmsArn: string ,
-    nameonly CustomEncryptionContext: AwsCryptographyKeyStoreTypes.EncryptionContextString
+    nameonly EncryptionContext: AwsCryptographyKeyStoreTypes.EncryptionContextString ,
+    nameonly HierarchyVersion: AwsCryptographyKeyStoreTypes.HierarchyVersion
   )
   datatype MutatedBranchKeyItem = | MutatedBranchKeyItem (
     nameonly ItemType: string ,
@@ -216,8 +224,8 @@ module {:extern "software.amazon.cryptography.keystoreadmin.internaldafny.types"
     nameonly MutationToken: MutationToken
   )
   datatype MutationDetails = | MutationDetails (
-    nameonly Original: MutableBranchKeyProperties ,
-    nameonly Terminal: MutableBranchKeyProperties ,
+    nameonly Original: MutableBranchKeyContext ,
+    nameonly Terminal: MutableBranchKeyContext ,
     nameonly Input: Mutations ,
     nameonly SystemKey: string ,
     nameonly CreateTime: string ,
@@ -228,7 +236,8 @@ module {:extern "software.amazon.cryptography.keystoreadmin.internaldafny.types"
     | No(No: string)
   datatype Mutations = | Mutations (
     nameonly TerminalKmsArn: Option<string> := Option.None ,
-    nameonly TerminalEncryptionContext: Option<AwsCryptographyKeyStoreTypes.EncryptionContextString> := Option.None
+    nameonly TerminalEncryptionContext: Option<AwsCryptographyKeyStoreTypes.EncryptionContextString> := Option.None ,
+    nameonly TerminalHierarchyVersion: Option<AwsCryptographyKeyStoreTypes.HierarchyVersion> := Option.None
   )
   datatype MutationToken = | MutationToken (
     nameonly Identifier: string ,