feat: Adds CreateKey API to create a branch key

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

@@ -16,6 +16,10 @@ import {
   decryptBranchKey,
 } from './branch_keystore_helpers'
 import { KMS_CLIENT_USER_AGENT, TABLE_FIELD } from './constants'
+import {
+  createBranchAndBeaconKeys,
+  versionActiveBranchKey,
+} from './key_helpers'
 
 import {
   IBranchKeyStorage,
@@ -49,6 +53,10 @@ interface IBranchKeyStoreNode {
   //= type=implication
   //# - [VersionKey](#versionkey)
   versionKey(input: VersionKeyInput): Promise<void>
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#operations
+  //= type=implication
+  //# - [CreateKey](#createkey)
+  createKey(input?: CreateKeyInput): Promise<CreateKeyOutput>
 }
 //= aws-encryption-sdk-specification/framework/branch-key-store.md#getkeystoreinfo
 //= type=implication
@@ -71,6 +79,15 @@ export interface VersionKeyInput {
   branchKeyIdentifier: string
 }
 
+export interface CreateKeyInput {
+  branchKeyIdentifier?: string
+  encryptionContext?: { [key: string]: string }
+}
+
+export interface CreateKeyOutput {
+  branchKeyIdentifier: string
+}
+
 export class BranchKeyStoreNode implements IBranchKeyStoreNode {
   public declare readonly logicalKeyStoreName: string
   public declare readonly kmsConfiguration: Readonly<KmsKeyConfig>
@@ -390,6 +407,55 @@ export class BranchKeyStoreNode implements IBranchKeyStoreNode {
     return branchKeyMaterials
   }
 
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey
+  //# The CreateKey caller MUST provide:
+  //# - An optional branch key id
+  //# - An optional encryption context
+  async createKey(input?: CreateKeyInput): Promise<CreateKeyOutput> {
+    //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey
+    //# If the Keystore's KMS Configuration is `Discovery` or `MRDiscovery`,
+    //# this operation MUST fail.
+    needs(
+      typeof this.kmsConfiguration._config === 'object' &&
+        ('identifier' in this.kmsConfiguration._config ||
+          'mrkIdentifier' in this.kmsConfiguration._config),
+      'CreateKey is not supported with Discovery or MRDiscovery KMS Configuration'
+    )
+
+    //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey
+    //# If an optional branch key id is provided and no encryption context is provided
+    //# this operation MUST fail.
+    if (input?.branchKeyIdentifier) {
+      needs(
+        input.encryptionContext &&
+          Object.keys(input.encryptionContext).length > 0,
+        'If branch key identifier is provided, encryption context must also be provided'
+      )
+    }
+
+    //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey
+    //# If no branch key id is provided, then this operation MUST create a
+    //# version 4 UUID to be used as the branch key id.
+    const branchKeyIdentifier = input?.branchKeyIdentifier || v4()
+    const customEncryptionContext = input?.encryptionContext || {}
+
+    await createBranchAndBeaconKeys({
+      branchKeyIdentifier,
+      customEncryptionContext,
+      logicalKeyStoreName: this.logicalKeyStoreName,
+      kmsConfiguration: this.kmsConfiguration,
+      grantTokens: this.grantTokens,
+      kmsClient: this.kmsClient,
+      ddbClient: (this.storage as any).ddbClient,
+      ddbTableName: (this.storage as any).ddbTableName,
+    })
+
+    //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey
+    //# If writing to the keystore succeeds,
+    //# the operation MUST return the branch-key-id that maps to both the branch key and the beacon key.
+    return { branchKeyIdentifier }
+  }
+
   //= aws-encryption-sdk-specification/framework/branch-key-store.md#versionkey
   //# On invocation, the caller:
   //# - MUST supply a `branch-key-id`
@@ -406,7 +472,6 @@ export class BranchKeyStoreNode implements IBranchKeyStoreNode {
       'VersionKey is not supported with Discovery or MRDiscovery KMS Configuration'
     )
 
-    const { versionActiveBranchKey } = await import('./key_helpers')
     await versionActiveBranchKey({
       branchKeyIdentifier: input.branchKeyIdentifier,
       logicalKeyStoreName: this.logicalKeyStoreName,

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

@@ -23,9 +23,23 @@ import {
   BRANCH_KEY_TYPE_PREFIX,
   BRANCH_KEY_ACTIVE_TYPE,
   BRANCH_KEY_ACTIVE_VERSION_FIELD,
+  BEACON_KEY_TYPE_VALUE,
+  CUSTOM_ENCRYPTION_CONTEXT_FIELD_PREFIX,
+  KMS_FIELD,
 } from './constants'
 import { IBranchKeyStorage } from './types'
 
+interface CreateKeyParams {
+  branchKeyIdentifier: string
+  customEncryptionContext: { [key: string]: string }
+  logicalKeyStoreName: string
+  kmsConfiguration: Readonly<KmsKeyConfig>
+  grantTokens?: ReadonlyArray<string>
+  kmsClient: KMSClient
+  ddbClient: DynamoDBClient
+  ddbTableName: string
+}
+
 interface VersionKeyParams {
   branchKeyIdentifier: string
   logicalKeyStoreName: string
@@ -95,6 +109,179 @@ function getKmsKeyArn(
     : undefined
 }
 
+//= aws-encryption-sdk-specification/framework/branch-key-store.md#decrypt_only-encryption-context
+//# The DECRYPT_ONLY encryption context MUST NOT have a `version` attribute.
+//# The `type` attribute MUST stores the branch key version formatted like `"branch:version:"` + `version`.
+function buildDecryptOnlyEncryptionContext(
+  branchKeyIdentifier: string,
+  branchKeyVersion: string,
+  timestamp: string,
+  logicalKeyStoreName: string,
+  kmsArn: string,
+  customEncryptionContext: { [key: string]: string }
+): { [key: string]: string } {
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#encryption-context
+  //# - MUST have a `branch-key-id` attribute
+  //# - MUST have a `type` attribute
+  //# - MUST have a `create-time` attribute
+  //# - MUST have a `tablename` attribute to store the logicalKeyStoreName
+  //# - MUST have a `kms-arn` attribute
+  //# - MUST have a `hierarchy-version`
+  //# - MUST NOT have a `enc` attribute
+  const context: { [key: string]: string } = {
+    [BRANCH_KEY_IDENTIFIER_FIELD]: branchKeyIdentifier,
+    [TYPE_FIELD]: `${BRANCH_KEY_TYPE_PREFIX}${branchKeyVersion}`,
+    [KEY_CREATE_TIME_FIELD]: timestamp,
+    [TABLE_FIELD]: logicalKeyStoreName,
+    [KMS_FIELD]: kmsArn,
+    [HIERARCHY_VERSION_FIELD]: '1',
+  }
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#custom-encryption-context
+  //# To avoid name collisions each added attribute from the custom encryption context
+  //# MUST be prefixed with `aws-crypto-ec:`.
+  for (const [key, value] of Object.entries(customEncryptionContext)) {
+    context[`${CUSTOM_ENCRYPTION_CONTEXT_FIELD_PREFIX}${key}`] = value
+  }
+
+  return context
+}
+
+//= aws-encryption-sdk-specification/framework/branch-key-store.md#beacon-key-encryption-context
+//# The Beacon key encryption context value of the `type` attribute MUST equal to `"beacon:ACTIVE"`.
+//# The Beacon key encryption context MUST NOT have a `version` attribute.
+function buildBeaconEncryptionContext(decryptOnlyContext: {
+  [key: string]: string
+}): { [key: string]: string } {
+  const beaconContext = { ...decryptOnlyContext }
+  beaconContext[TYPE_FIELD] = BEACON_KEY_TYPE_VALUE
+  return beaconContext
+}
+
+//= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
+//# This operation MUST create a branch key and a beacon key
+//# according to the Branch Key and Beacon Key Creation section.
+export async function createBranchAndBeaconKeys(
+  params: CreateKeyParams
+): Promise<void> {
+  const {
+    branchKeyIdentifier,
+    customEncryptionContext,
+    logicalKeyStoreName,
+    kmsConfiguration,
+    grantTokens,
+    kmsClient,
+    ddbClient,
+    ddbTableName,
+  } = params
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
+  //# - `version`: a new guid. This guid MUST be version 4 UUID
+  const branchKeyVersion = v4()
+  const timestamp = getCurrentTimestamp()
+
+  const kmsKeyArn = getKmsKeyArn(kmsConfiguration)
+  needs(kmsKeyArn, 'KMS Key ARN is required')
+
+  const decryptOnlyContext = buildDecryptOnlyEncryptionContext(
+    branchKeyIdentifier,
+    branchKeyVersion,
+    timestamp,
+    logicalKeyStoreName,
+    kmsKeyArn,
+    customEncryptionContext
+  )
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#wrapped-branch-key-creation
+  //# The operation MUST call AWS KMS API GenerateDataKeyWithoutPlaintext
+  const decryptOnlyResponse = await kmsClient.send(
+    new GenerateDataKeyWithoutPlaintextCommand({
+      KeyId: kmsKeyArn,
+      NumberOfBytes: 32,
+      EncryptionContext: decryptOnlyContext,
+      GrantTokens: grantTokens ? [...grantTokens] : undefined,
+    })
+  )
+
+  needs(
+    decryptOnlyResponse.CiphertextBlob,
+    'Failed to generate DECRYPT_ONLY branch key'
+  )
+
+  const activeContext = buildActiveEncryptionContext(decryptOnlyContext)
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#wrapped-branch-key-creation
+  //# The operation MUST call AWS KMS API ReEncrypt
+  const activeResponse = await kmsClient.send(
+    new ReEncryptCommand({
+      SourceKeyId: kmsKeyArn,
+      SourceEncryptionContext: decryptOnlyContext,
+      CiphertextBlob: decryptOnlyResponse.CiphertextBlob,
+      DestinationKeyId: kmsKeyArn,
+      DestinationEncryptionContext: activeContext,
+      GrantTokens: grantTokens ? [...grantTokens] : undefined,
+    })
+  )
+
+  needs(activeResponse.CiphertextBlob, 'Failed to generate ACTIVE branch key')
+
+  const beaconContext = buildBeaconEncryptionContext(decryptOnlyContext)
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
+  //# The operation MUST call AWS KMS GenerateDataKeyWithoutPlaintext for beacon key
+  const beaconResponse = await kmsClient.send(
+    new GenerateDataKeyWithoutPlaintextCommand({
+      KeyId: kmsKeyArn,
+      NumberOfBytes: 32,
+      EncryptionContext: beaconContext,
+      GrantTokens: grantTokens ? [...grantTokens] : undefined,
+    })
+  )
+
+  needs(beaconResponse.CiphertextBlob, 'Failed to generate beacon key')
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#writing-branch-key-and-beacon-key-to-keystore
+  //# The call to Amazon DynamoDB TransactWriteItems MUST use the configured Amazon DynamoDB Client to make the call.
+  await ddbClient.send(
+    new TransactWriteItemsCommand({
+      TransactItems: [
+        {
+          Put: {
+            TableName: ddbTableName,
+            Item: toAttributeMap(
+              decryptOnlyContext,
+              decryptOnlyResponse.CiphertextBlob
+            ),
+            ConditionExpression: 'attribute_not_exists(#bkid)',
+            ExpressionAttributeNames: {
+              '#bkid': BRANCH_KEY_IDENTIFIER_FIELD,
+            },
+          },
+        },
+        {
+          Put: {
+            TableName: ddbTableName,
+            Item: toAttributeMap(activeContext, activeResponse.CiphertextBlob),
+            ConditionExpression: 'attribute_not_exists(#bkid)',
+            ExpressionAttributeNames: {
+              '#bkid': BRANCH_KEY_IDENTIFIER_FIELD,
+            },
+          },
+        },
+        {
+          Put: {
+            TableName: ddbTableName,
+            Item: toAttributeMap(beaconContext, beaconResponse.CiphertextBlob),
+            ConditionExpression: 'attribute_not_exists(#bkid)',
+            ExpressionAttributeNames: {
+              '#bkid': BRANCH_KEY_IDENTIFIER_FIELD,
+            },
+          },
+        },
+      ],
+    })
+  )
+}
 //= aws-encryption-sdk-specification/framework/branch-key-store.md#versionkey
 //# On invocation, the caller:
 //# - MUST supply a `branch-key-id`