update the anotations

Author: seebees

modules/branch-keystore-node/src/branch_keystore.ts

@@ -1,7 +1,7 @@
 // Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { isKmsConfig, KmsConfig, RegionalKmsConfig } from './kms_config'
+import { KmsConfig, KmsKeyConfig } from './kms_config'
 import { KMSClient } from '@aws-sdk/client-kms'
 import { DynamoDBClient } from '@aws-sdk/client-dynamodb'
 import {
@@ -62,12 +62,12 @@ export interface KeyStoreInfoOutput {
 }
 
 export class BranchKeyStoreNode implements IBranchKeyStoreNode {
-  public declare logicalKeyStoreName: string
-  public declare kmsConfiguration: Readonly<KmsConfig>
-  public declare kmsClient: KMSClient
-  public declare keyStoreId: string
-  public declare grantTokens?: ReadonlyArray<string>
-  public declare storage: IBranchKeyStorage
+  public declare readonly logicalKeyStoreName: string
+  public declare readonly kmsConfiguration: Readonly<KmsKeyConfig>
+  public declare readonly kmsClient: KMSClient
+  public declare readonly keyStoreId: string
+  public declare readonly grantTokens?: ReadonlyArray<string>
+  public declare readonly storage: IBranchKeyStorage
 
   constructor({
     logicalKeyStoreName,
@@ -82,8 +82,12 @@ export class BranchKeyStoreNode implements IBranchKeyStoreNode {
       'Logical keystore name must be a string'
     )
 
-    /* Precondition: KMS Configuration must be SRK */
-    needs(isKmsConfig(kmsConfiguration), 'KMS Configuration must be SRK')
+    needs(kmsConfiguration, 'AWS KMS Configuration required')
+    readOnlyProperty(
+      this,
+      'kmsConfiguration',
+      new KmsKeyConfig(kmsConfiguration)
+    )
 
     /* Precondition: KMS client must be a KMSClient */
     if (keyManagement?.kmsClient) {
@@ -131,7 +135,16 @@ export class BranchKeyStoreNode implements IBranchKeyStoreNode {
           storage.ddbClient instanceof DynamoDBClient
             ? storage.ddbClient
             : new DynamoDBClient({
-                region: (kmsConfiguration as RegionalKmsConfig).getRegion(),
+                //= aws-encryption-sdk-specification/framework/branch-key-store.md#initialization
+                //# If a DDB client needs to be constructed and the AWS KMS Configuration is KMS Key ARN or KMS MRKey ARN,
+                //# a new DynamoDb client MUST be created with the region of the supplied KMS ARN.
+                //#
+                //# If a DDB client needs to be constructed and the AWS KMS Configuration is Discovery,
+                //# a new DynamoDb client MUST be created with the default configuration.
+                //#
+                //# If a DDB client needs to be constructed and the AWS KMS Configuration is MRDiscovery,
+                //# a new DynamoDb client MUST be created with the region configured in the MRDiscovery.
+                region: this.kmsConfiguration.getRegion(),
               }),
       })
     }
@@ -176,9 +189,6 @@ export class BranchKeyStoreNode implements IBranchKeyStoreNode {
     )
     /* Postcondition: If unprovided, the grant tokens are undefined */
 
-    needs(kmsConfiguration, 'AWS KMS Configuration required')
-    readOnlyProperty(this, 'kmsConfiguration', Object.freeze(kmsConfiguration))
-
     // TODO: when other KMS configuration types/classes are supported for the keystore,
     // verify the configuration object type to determine how we instantiate the
     // KMS client. This will ensure safe type casting.
@@ -189,7 +199,16 @@ export class BranchKeyStoreNode implements IBranchKeyStoreNode {
       //# If no AWS KMS client is provided one MUST be constructed.
       keyManagement?.kmsClient ||
         new KMSClient({
-          region: (this.kmsConfiguration as RegionalKmsConfig).getRegion(),
+          //= aws-encryption-sdk-specification/framework/branch-key-store.md#initialization
+          //# If AWS KMS client needs to be constructed and the AWS KMS Configuration is KMS Key ARN or KMS MRKey ARN,
+          //# a new AWS KMS client MUST be created with the region of the supplied KMS ARN.
+          //#
+          //# If AWS KMS client needs to be constructed and the AWS KMS Configuration is Discovery,
+          //# a new AWS KMS client MUST be created with the default configuration.
+          //#
+          //# If AWS KMS client needs to be constructed and the AWS KMS Configuration is MRDiscovery,
+          //# a new AWS KMS client MUST be created with the region configured in the MRDiscovery.
+          region: this.kmsConfiguration.getRegion(),
           //= aws-encryption-sdk-specification/framework/branch-key-store.md#initialization
           //# On initialization the KeyStore SHOULD
           //# append a user agent string to the AWS KMS SDK Client with
@@ -375,7 +394,7 @@ export class BranchKeyStoreNode implements IBranchKeyStoreNode {
       keystoreTableName: this.storage.getKeyStorageInfo().name,
       logicalKeyStoreName: this.logicalKeyStoreName,
       grantTokens: !!this.grantTokens ? this.grantTokens.slice() : [],
-      kmsConfiguration: this.kmsConfiguration,
+      kmsConfiguration: this.kmsConfiguration._config,
     }
   }
 }
@@ -650,25 +669,25 @@ export function isIBranchKeyStoreNode(
 //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
 //= type=exception
 //# On invocation, the caller:
-//# 
+//#
 //# - MUST supply a `branch-key-id`
-//# 
+//#
 //# GetBeaconKey MUST get the requested beacon key from the keystore
 //# by calling the configured [KeyStorage interface's](./key-store/key-storage.md#interface)
 //# [GetEncryptedBeaconKey](./key-store/key-storage.md#getencryptedbeaconkey)
 //# using the supplied `branch-key-id`.
-//# 
+//#
 //# Because the storage interface can be a custom implementation the key store needs to verify correctness.
-//# 
+//#
 //# GetBeaconKey MUST verify that the returned EncryptedHierarchicalKey MUST have the requested `branch-key-id`.
 //# GetBeaconKey MUST verify that the returned EncryptedHierarchicalKey is an ActiveHierarchicalSymmetricBeacon.
 //# GetBeaconKey MUST verify that the returned EncryptedHierarchicalKey MUST have a logical table name equal to the configured logical table name.
-//# 
+//#
 //# The operation MUST decrypt the beacon key according to the [AWS KMS Branch Key Decryption](#aws-kms-branch-key-decryption) section.
-//# 
+//#
 //# If the beacon key fails to decrypt, this operation MUST fail.
-//# 
+//#
 //# This GetBeaconKey MUST construct [beacon key materials](./structures.md#beacon-key-materials) from the decrypted branch key material
 //# and the `branchKeyId` from the returned `branch-key-id` field.
-//# 
-//# This operation MUST return the constructed [beacon key materials](./structures.md#beacon-key-materials).
+//#
+//# This operation MUST return the constructed [beacon key materials](./structures.md#beacon-key-materials).

modules/branch-keystore-node/src/branch_keystore_helpers.ts

@@ -13,7 +13,7 @@ import { BranchKeyItem, BranchKeyRecord } from './branch_keystore_structures'
 import { EncryptedHierarchicalKey, BranchKeyEncryptionContext } from './types'
 // import { IBranchKeyStoreNode } from './branch_keystore'
 import { DecryptCommand } from '@aws-sdk/client-kms'
-import { KmsKeyArnConfig, KmsConfig } from './kms_config'
+import { KmsKeyConfig } from './kms_config'
 import {
   PARTITION_KEY,
   SORT_KEY,
@@ -31,6 +31,7 @@ import {
   BEACON_KEY_TYPE_VALUE,
   POTENTIAL_BRANCH_KEY_RECORD_FIELDS,
 } from './constants'
+import { parseAwsKmsKeyArn } from '@aws-crypto/kms-keyring'
 
 /**
  * This utility function uses a partition and sort key to query for a single branch
@@ -254,32 +255,33 @@ export async function decryptBranchKey(
     kmsClient,
   }: {
     kmsClient: KMSClient
-    kmsConfiguration: Readonly<KmsConfig>
+    kmsConfiguration: Readonly<KmsKeyConfig>
     grantTokens?: ReadonlyArray<string>
   },
   encryptedHierarchicalKey: EncryptedHierarchicalKey
 ): Promise<Buffer> {
-  //= aws-encryption-sdk-specification/framework/branch-key-store.md#aws-kms-branch-key-decryption
-  //# If the Keystore's [AWS KMS Configuration](#aws-kms-configuration) is `KMS Key ARN` or `KMS MRKey ARN`,
-  //# the `kms-arn` field of the DDB response item MUST be
-  //# [compatible with](#aws-key-arn-compatibility) the configured KMS Key in
-  //# the [AWS KMS Configuration](#aws-kms-configuration) for this keystore,
-  //# or the operation MUST fail.
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#discovery
+  //# The Keystore MAY use the KMS Key ARNs already
+  //# persisted to the backing DynamoDB table,
+  //# provided they are in records created
+  //# with an identical Logical Keystore Name.
 
-  //# If the Keystore's [AWS KMS Configuration](#aws-kms-configuration) is `Discovery` or `MRDiscovery`,
-  //# the `kms-arn` field of DDB response item MUST NOT be an Alias
-  //# or the operation MUST fail.
-  needs(
-    kmsConfiguration.isCompatibleWithArn(encryptedHierarchicalKey.kmsArn),
-    'KMS ARN from DDB response item MUST be compatible with the configured KMS Key in the AWS KMS Configuration for this keystore'
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#mrdiscovery
+  //# The Keystore MAY use the KMS Key ARNs already
+  //# persisted to the backing DynamoDB table,
+  //# provided they are in records created
+  //# with an identical Logical Keystore Name.
+
+  const KeyId = kmsConfiguration.getCompatibleArnArn(
+    encryptedHierarchicalKey.kmsArn
   )
 
   //= aws-encryption-sdk-specification/framework/branch-key-store.md#aws-kms-branch-key-decryption
   //# When calling [AWS KMS Decrypt](https://docs.aws.amazon.com/kms/latest/APIReference/API_Decrypt.html),
   //# the keystore operation MUST call with a request constructed as follows:
   const response = await kmsClient.send(
     new DecryptCommand({
-      KeyId: (kmsConfiguration as KmsKeyArnConfig).getArn(), // make this type casting assumption since only SRK Compatibility is supported currently
+      KeyId,
       //= aws-encryption-sdk-specification/framework/branch-key-store.md#aws-kms-branch-key-decryption
       //# - `CiphertextBlob` MUST be the `CiphertextBlob` attribute value on the provided EncryptedHierarchicalKey
       CiphertextBlob: encryptedHierarchicalKey.ciphertextBlob,