feat(dafny): Improving the Encryption Context experience of Branch Keys

AwsCryptographicMaterialProviders/codegen-patches/AwsCryptographyKeyStore/dotnet/dafny-4.2.0.patch

@@ -12,7 +12,7 @@ index 95c9eba1..bcd537fb 100644
 +          // END MANUAL EDIT
              dafnyVal._ComAmazonawsKms
            );
-         case software.amazon.cryptography.keystore.internaldafny.types.Error_KeyStoreException dafnyVal:
+         case software.amazon.cryptography.keystore.internaldafny.types.Error_BranchKeyCiphertextException dafnyVal:
 @@ -755,7 +757,9 @@ namespace AWS.Cryptography.KeyStore
        {
          case "Com.Amazonaws.KMS":

AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/CanonicalEncryptionContext.dfy

@@ -11,6 +11,7 @@ module CanonicalEncryptionContext {
   import opened Wrappers
   import Seq
   import SortedSets
+  import AtomicPrimitives
 
   //= aws-encryption-sdk-specification/framework/raw-aes-keyring.md#onencrypt
   //# The keyring MUST attempt to serialize the [encryption materials']
@@ -47,4 +48,51 @@ module CanonicalEncryptionContext {
       var allBytes := UInt16ToSeq(|keys| as uint16) + Seq.Flatten(pairsBytes);
       Success(allBytes)
   }
+
+  method EncryptionContextDigest(
+    crypto: AtomicPrimitives.AtomicPrimitivesClient,
+    encryptionContext: Types.EncryptionContext
+  )
+    returns (output: Result<seq<uint8>, string>)
+    requires crypto.ValidState()
+    modifies crypto.Modifies
+    ensures crypto.ValidState()
+    // We tried to collapse all of this post condition into a common predicate.
+    // We failed; we will have to copy and paste it.
+    // @texastony thinks reason is that is very difficult to keep track of the input.
+    ensures
+      && output.Success?
+      ==>
+        && var content? := EncryptionContextToAAD(encryptionContext);
+        && content?.Success?
+        && |crypto.History.Digest| == |old(crypto.History.Digest)| + 1
+        && old(crypto.History.Digest) < crypto.History.Digest
+        && var digestEvent := Seq.Last(crypto.History.Digest);
+        && digestEvent.input.digestAlgorithm == AtomicPrimitives.Types.SHA_384
+        && digestEvent.input.message == content?.value
+        && digestEvent.output.Success?
+        && |digestEvent.output.value| == 48 // 384 bits / 8 bits per byte == 48 bytes
+        && digestEvent.output.value == output.value
+  {
+    var canonicalEC :- EncryptionContextToAAD(encryptionContext).MapFailure(e => "Could not SHA-384 Content.");
+
+    var DigestInput := AtomicPrimitives.Types.DigestInput(
+      digestAlgorithm := AtomicPrimitives.Types.SHA_384,
+      message := canonicalEC
+    );
+    var maybeDigest := crypto.Digest(DigestInput);
+    var digest :- maybeDigest.MapFailure(e => "Could not SHA-384 Content.");
+
+    // The digest is not truncated.
+    // There is an impact on the key size.
+    // See: https://docs.aws.amazon.com/kms/latest/developerguide/asymmetric-key-specs.html
+    // This is not safe to do for 1024 keys,
+    // but AWS KMS does not support these keys.
+    // Further we use SHA_384 to save a little on size
+    // and avoid even the possibility of length extension.
+    // Though length extension does not matter in this situation,
+    // because a decryptor already has access to the key.
+    return Success(digest);
+  }
+
 }

AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/test/CMCs/LocalCMC.dfy

@@ -18,7 +18,10 @@ module  {:options "/functionSyntax:4"} TestLocalCMC {
                 branchKeyIdentifier := "spoo",
                 branchKeyVersion := data,
                 branchKey := data,
-                encryptionContext := map[]
+                encryptionContext := map[],
+                kmsArn := "KeyId",
+                createTime := "CreateTime",
+                hierarchyVersion := HierarchyVersion.v1
               ))
   }
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/test/CMCs/StormTracker.dfy

@@ -18,7 +18,10 @@ module  {:options "/functionSyntax:4"} TestStormTracker {
                 branchKeyIdentifier := "spoo",
                 branchKeyVersion := data,
                 branchKey := data,
-                encryptionContext := map[]
+                encryptionContext := map[],
+                kmsArn := "KeyId",
+                createTime := "CreateTime",
+                hierarchyVersion := HierarchyVersion.v1
               ))
   }
 

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/Model/AwsCryptographyKeyStoreTypes.dfy

@@ -2,13 +2,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // Do not modify this file. This file is machine generated, and any changes to it will be overwritten.
 include "../../../../StandardLibrary/src/Index.dfy"
+include "../../../../AwsCryptographyPrimitives/src/Index.dfy"
 include "../../../../ComAmazonawsDynamodb/src/Index.dfy"
 include "../../../../ComAmazonawsKms/src/Index.dfy"
 module {:extern "software.amazon.cryptography.keystore.internaldafny.types" } AwsCryptographyKeyStoreTypes
 {
   import opened Wrappers
   import opened StandardLibrary.UInt
   import opened UTF8
+  import AwsCryptographyPrimitivesTypes
   import ComAmazonawsDynamodbTypes
   import ComAmazonawsKmsTypes
     // Generic helpers for verification of mock/unit tests.
@@ -20,17 +22,25 @@ module {:extern "software.amazon.cryptography.keystore.internaldafny.types" } Aw
     nameonly beaconKeyIdentifier: string ,
     nameonly encryptionContext: EncryptionContext ,
     nameonly beaconKey: Option<Secret> := Option.None ,
-    nameonly hmacKeys: Option<HmacKeyMap> := Option.None
+    nameonly hmacKeys: Option<HmacKeyMap> := Option.None ,
+    nameonly kmsArn: ComAmazonawsKmsTypes.KeyIdType ,
+    nameonly createTime: CreateTime ,
+    nameonly hierarchyVersion: HierarchyVersion
   )
+  type BranchKeyIdentifier = string
   datatype BranchKeyMaterials = | BranchKeyMaterials (
     nameonly branchKeyIdentifier: string ,
     nameonly branchKeyVersion: Utf8Bytes ,
     nameonly encryptionContext: EncryptionContext ,
-    nameonly branchKey: Secret
+    nameonly branchKey: Secret ,
+    nameonly kmsArn: ComAmazonawsKmsTypes.KeyIdType ,
+    nameonly createTime: CreateTime ,
+    nameonly hierarchyVersion: HierarchyVersion
   )
   datatype CreateKeyInput = | CreateKeyInput (
     nameonly branchKeyIdentifier: Option<string> := Option.None ,
-    nameonly encryptionContext: Option<EncryptionContext> := Option.None
+    nameonly encryptionContext: Option<EncryptionContext> := Option.None ,
+    nameonly hierarchyVersion: Option<HierarchyVersion> := Option.None
   )
   datatype CreateKeyOutput = | CreateKeyOutput (
     nameonly branchKeyIdentifier: string
@@ -41,6 +51,7 @@ module {:extern "software.amazon.cryptography.keystore.internaldafny.types" } Aw
   datatype CreateKeyStoreOutput = | CreateKeyStoreOutput (
     nameonly tableArn: ComAmazonawsDynamodbTypes.TableArn
   )
+  type CreateTime = string
   datatype Discovery = | Discovery (
 
                        )
@@ -72,6 +83,9 @@ module {:extern "software.amazon.cryptography.keystore.internaldafny.types" } Aw
     nameonly kmsConfiguration: KMSConfiguration
   )
   type GrantTokenList = seq<string>
+  datatype HierarchyVersion =
+    | v1
+    | v2
   type HmacKeyMap = map<string, Secret>
   class IKeyStoreClientCallHistory {
     ghost constructor() {
@@ -251,10 +265,14 @@ module {:extern "software.amazon.cryptography.keystore.internaldafny.types" } Aw
                               )
   datatype Error =
       // Local Error structures are listed here
+    | BranchKeyCiphertextException (
+        nameonly message: string
+      )
     | KeyStoreException (
         nameonly message: string
       )
       // Any dependent models are listed here
+    | AwsCryptographyPrimitives(AwsCryptographyPrimitives: AwsCryptographyPrimitivesTypes.Error)
     | ComAmazonawsDynamodb(ComAmazonawsDynamodb: ComAmazonawsDynamodbTypes.Error)
     | ComAmazonawsKms(ComAmazonawsKms: ComAmazonawsKmsTypes.Error)
       // The Collection error is used to collect several errors together

AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/Model/KeyStore.smithy

@@ -18,6 +18,9 @@ use com.amazonaws.dynamodb#DynamoDB_20120810
 
 use com.amazonaws.kms#TrentService
 
+use aws.cryptography.primitives#AwsCryptographicPrimitives
+use aws.cryptography.materialProviders#AwsCryptographicMaterialProviders
+
 
 @dafnyUtf8Bytes
 string Utf8Bytes
@@ -28,12 +31,16 @@ structure KmsClientReference {}
 @reference(service: DynamoDB_20120810)
 structure DdbClientReference {}
 
+@reference(service: AwsCryptographicPrimitives)
+structure PrimitivesReference {}
+
 @localService(
   sdkId: "KeyStore",
   config: KeyStoreConfig,
   dependencies: [
     DynamoDB_20120810,
-    TrentService
+    TrentService,
+    AwsCryptographicPrimitives
   ] 
 )
 service KeyStore {
@@ -59,9 +66,25 @@ service KeyStore {
     GetBranchKeyVersion,
     GetBeaconKey
   ],
-  errors: [KeyStoreException]
+  errors: [
+    KeyStoreException
+    BranchKeyCiphertextException
+  ]
 }
 
+@documentation("The identifier for the Branch Key.")
+string BranchKeyIdentifier
+
+@documentation("Timestamp in ISO 8601 format in UTC, to microsecond precision, that this Branch Key Item's Material was generated.")
+string CreateTime
+
+@documentation("Schema Version of the Branch Key. All items of the same Branch Key Identifier SHOULD have the same hierarchy-version. The hierarchy-version determines how the Branch Key Store protects and validates the branch key with KMS.")
+@enum([
+  { name: "v1", value: "1" },
+  { name: "v2", value: "2" }
+])
+string HierarchyVersion
+
 structure KeyStoreConfig {
 
   //= aws-encryption-sdk-specification/framework/branch-key-store.md#initialization
@@ -180,7 +203,15 @@ structure CreateKeyStoreOutput {
 // 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.")
+@documentation(
+  "Create a new Branch Key in the Key Store. Additionally create a Beacon Key that is tied to this Branch Key.
+  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 AES-256 key.
+  The beacon key AES-256 is unqiue.
+  For Hierarchy Version 1, KMS is called 3 times; GenerateDataKeyWithoutPlaintext is called twice, ReEncrypt is called once.
+  For Hierarchy Version 2, KMS is called 5 times; GenerateDataKey follwed by Encrypt twice to create the ACTIVE branch key and decrypt_only. Another GenerateDataKey follwed by Encrypt creates the beacon key.
+  ")
 operation CreateKey {
   input: CreateKeyInput,
   output: CreateKeyOutput
@@ -197,6 +228,9 @@ structure CreateKeyInput {
 
   @javadoc("Custom encryption context for the Branch Key. Required if branchKeyIdentifier is set.")
   encryptionContext: EncryptionContext
+
+  @documentation("Optional. Defaults to v1.")
+  hierarchyVersion: HierarchyVersion
 }
 
 @javadoc("Outputs for Branch Key creation.")
@@ -210,7 +244,15 @@ structure CreateKeyOutput {
 // 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.
+   Rotation is accomplished by first authenticating the ACTIVE branch key item according to it's hierarchy-version with KMS.
+   Then, again using KMS, new material is generated and encrypted, creating a new ACTIVE and corresponding decrypt_only.
+   These two items are then writen to the DynamoDB table via a TransactionWriteItems;
+   this only overwrites the ACTIVE item, the corresponding 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
@@ -312,14 +354,17 @@ list GrantTokenList {
   member: String
 }
 
-//= aws-encryption-sdk-specification/framework/structures.md#structure-3
-//= type=implication
-//# This structure MUST include all of the following fields:
-//# 
-//# - [Branch Key](#branch-key)
-//# - [Branch Key Id](#branch-key-id)
-//# - [Branch Key Version](#branch-key-version)
-//# - [Encryption Context](#encryption-context-3)
+//XX= aws-encryption-sdk-specification/framework/structures.md#structure-3
+//XX= type=implication
+//XX# This structure MUST include all of the following fields:
+//XX# 
+//XX# - [Branch Key](#branch-key)
+//XX# - [Branch Key Id](#branch-key-id)
+//XX# - [Branch Key Version](#branch-key-version)
+//XX# - [Encryption Context](#encryption-context-3)
+//XX# - [KMS ARN](#kms-arn)
+//XX# - [Create Time](#create-time)
+//XX# - [Hierarchy Version](#hierarchy-version)
 structure BranchKeyMaterials {
     @required
     branchKeyIdentifier: String,
@@ -332,6 +377,16 @@ structure BranchKeyMaterials {
 
     @required
     branchKey: Secret,
+
+    @required
+    @documentation("The AWS KMS Key ARN used to protect this Branch Key.")
+    kmsArn: com.amazonaws.kms#KeyIdType
+
+    @required
+    createTime: CreateTime
+
+    @required
+    hierarchyVersion: HierarchyVersion
 }
 
 structure BeaconKeyMaterials {
@@ -355,6 +410,16 @@ structure BeaconKeyMaterials {
   beaconKey: Secret,
 
   hmacKeys: HmacKeyMap
+
+  @required
+  @documentation("The AWS KMS Key ARN used to protect this Branch Key.")
+  kmsArn: com.amazonaws.kms#KeyIdType
+
+  @required
+  createTime: CreateTime
+
+  @required
+  hierarchyVersion: HierarchyVersion
 }
 
 map HmacKeyMap {
@@ -378,3 +443,23 @@ structure KeyStoreException {
   @required
   message: String,
 }
+
+@error("client")
+@documentation("
+The cipher-text or branch key materials incorporated into the cipher-text,
+such as the encryption context, is corrupted, missing, or otherwise invalid.
+For branch keys,
+the branch key materials 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.
+")
+structure BranchKeyCiphertextException {
+  @required
+  message: String
+}