From 78ba32caf553fc79b4cfbcb900e8f0d477b78a25 Mon Sep 17 00:00:00 2001 From: Bikram Sharma Date: Tue, 10 Mar 2026 23:12:00 -0700 Subject: [PATCH] feat: Adds CreateKey API to create a branch key --- .../src/branch_keystore.ts | 67 ++++++- .../branch-keystore-node/src/key_helpers.ts | 187 ++++++++++++++++++ .../test/branch_keystore.test.ts | 147 ++++++++++++++ 3 files changed, 400 insertions(+), 1 deletion(-) diff --git a/modules/branch-keystore-node/src/branch_keystore.ts b/modules/branch-keystore-node/src/branch_keystore.ts index 13cfcc5c..8f07d08f 100644 --- a/modules/branch-keystore-node/src/branch_keystore.ts +++ b/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 + //= aws-encryption-sdk-specification/framework/branch-key-store.md#operations + //= type=implication + //# - [CreateKey](#createkey) + createKey(input?: CreateKeyInput): Promise } //= 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 @@ -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 { + //= 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, diff --git a/modules/branch-keystore-node/src/key_helpers.ts b/modules/branch-keystore-node/src/key_helpers.ts index 43f613c5..4aaf9405 100644 --- a/modules/branch-keystore-node/src/key_helpers.ts +++ b/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 + grantTokens?: ReadonlyArray + 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 { + 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` diff --git a/modules/branch-keystore-node/test/branch_keystore.test.ts b/modules/branch-keystore-node/test/branch_keystore.test.ts index d4add7bf..20d05540 100644 --- a/modules/branch-keystore-node/test/branch_keystore.test.ts +++ b/modules/branch-keystore-node/test/branch_keystore.test.ts @@ -871,4 +871,151 @@ describe('Test Branch keystore', () => { expect(oldMaterial.branchKeyIdentifier).to.equal(BRANCH_KEY_ID_WITH_EC) }) }) + + describe('CreateKey', () => { + //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey + //= type=test + //# If the Keystore's KMS Configuration is `Discovery` or `MRDiscovery`, + //# this operation MUST fail. + it('MUST fail with Discovery KMS Configuration', async () => { + const keyStore = new BranchKeyStoreNode({ + kmsConfiguration: 'discovery', + logicalKeyStoreName: LOGICAL_KEYSTORE_NAME, + storage: { ddbTableName: DDB_TABLE_NAME }, + }) + + await expect(keyStore.createKey()).to.be.rejectedWith( + 'CreateKey is not supported with Discovery or MRDiscovery KMS Configuration' + ) + }) + + it('MUST fail with MRDiscovery KMS Configuration', async () => { + const keyStore = new BranchKeyStoreNode({ + kmsConfiguration: { region: 'us-west-2' }, + logicalKeyStoreName: LOGICAL_KEYSTORE_NAME, + storage: { ddbTableName: DDB_TABLE_NAME }, + }) + + await expect(keyStore.createKey()).to.be.rejectedWith( + 'CreateKey is not supported with Discovery or MRDiscovery KMS Configuration' + ) + }) + + //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey + //= type=test + //# If an optional branch key id is provided and no encryption context is provided + //# this operation MUST fail. + it('MUST fail if branch key id provided without encryption context', async () => { + const keyStore = new BranchKeyStoreNode({ + kmsConfiguration: { identifier: KEY_ARN }, + logicalKeyStoreName: LOGICAL_KEYSTORE_NAME, + storage: { ddbTableName: DDB_TABLE_NAME }, + }) + + await expect( + keyStore.createKey({ branchKeyIdentifier: 'some-id' }) + ).to.be.rejectedWith( + 'If branch key identifier is provided, encryption context must also be provided' + ) + + await expect( + keyStore.createKey({ + branchKeyIdentifier: 'some-id', + encryptionContext: {}, + }) + ).to.be.rejectedWith( + 'If branch key identifier is provided, encryption context must also be provided' + ) + }) + + //= aws-encryption-sdk-specification/framework/branch-key-store.md#createkey + //= type=test + //# If no branch key id is provided, then this operation MUST create a + //# version 4 UUID to be used as the branch key id. + it('Test create key with auto-generated ID', async () => { + const kmsClient = new KMSClient({}) + const ddbClient = new DynamoDBClient({}) + const keyStore = new BranchKeyStoreNode({ + kmsConfiguration: { identifier: KEY_ARN }, + logicalKeyStoreName: LOGICAL_KEYSTORE_NAME, + storage: { ddbTableName: DDB_TABLE_NAME, ddbClient }, + keyManagement: { kmsClient }, + }) + + const result = await keyStore.createKey() + + // Must return a valid v4 UUID + expect(result.branchKeyIdentifier).to.be.a('string') + expect(validate(result.branchKeyIdentifier)).to.be.true + expect(version(result.branchKeyIdentifier)).to.equal(4) + + // Must be retrievable + const material = await keyStore.getActiveBranchKey( + result.branchKeyIdentifier + ) + expect(material.branchKey().length).to.equal(32) + expect(material.branchKeyIdentifier).to.equal(result.branchKeyIdentifier) + }) + + it('Test create key with custom ID and encryption context', async () => { + const kmsClient = new KMSClient({}) + const ddbClient = new DynamoDBClient({}) + const keyStore = new BranchKeyStoreNode({ + kmsConfiguration: { identifier: KEY_ARN }, + logicalKeyStoreName: LOGICAL_KEYSTORE_NAME, + storage: { ddbTableName: DDB_TABLE_NAME, ddbClient }, + keyManagement: { kmsClient }, + }) + + const customId = v4() + const result = await keyStore.createKey({ + branchKeyIdentifier: customId, + encryptionContext: { department: 'test' }, + }) + + expect(result.branchKeyIdentifier).to.equal(customId) + + // Active key must be retrievable + const material = await keyStore.getActiveBranchKey(customId) + expect(material.branchKey().length).to.equal(32) + }) + }) + + describe('CreateKey + VersionKey lifecycle', () => { + it('Create, retrieve, version, retrieve new, retrieve old', async () => { + const kmsClient = new KMSClient({}) + const ddbClient = new DynamoDBClient({}) + const keyStore = new BranchKeyStoreNode({ + kmsConfiguration: { identifier: KEY_ARN }, + logicalKeyStoreName: LOGICAL_KEYSTORE_NAME, + storage: { ddbTableName: DDB_TABLE_NAME, ddbClient }, + keyManagement: { kmsClient }, + }) + + // 1. Create a new branch key + const { branchKeyIdentifier } = await keyStore.createKey() + + // 2. Retrieve the active key + const v1 = await keyStore.getActiveBranchKey(branchKeyIdentifier) + const v1Version = v1.branchKeyVersion.toString('utf8') + expect(v1.branchKey().length).to.equal(32) + + // 3. Version the key + await keyStore.versionKey({ branchKeyIdentifier }) + + // 4. Retrieve the new active key — must be different version + const v2 = await keyStore.getActiveBranchKey(branchKeyIdentifier) + const v2Version = v2.branchKeyVersion.toString('utf8') + expect(v2.branchKey().length).to.equal(32) + expect(v2Version).to.not.equal(v1Version) + + // 5. Old version is still retrievable + const oldMaterial = await keyStore.getBranchKeyVersion( + branchKeyIdentifier, + v1Version + ) + expect(oldMaterial.branchKey().length).to.equal(32) + expect(oldMaterial.branchKeyIdentifier).to.equal(branchKeyIdentifier) + }) + }) })