diff --git a/modules/kms-keyring-node/src/kms_hkeyring_node.ts b/modules/kms-keyring-node/src/kms_hkeyring_node.ts index f0a5ad8c2..8352e2957 100644 --- a/modules/kms-keyring-node/src/kms_hkeyring_node.ts +++ b/modules/kms-keyring-node/src/kms_hkeyring_node.ts @@ -15,6 +15,7 @@ import { KeyringNode, needs, NodeAlgorithmSuite, + NodeBranchKeyMaterial, NodeDecryptionMaterial, NodeEncryptionMaterial, readOnlyProperty, @@ -88,6 +89,7 @@ export interface IKmsHierarchicalKeyRingNode extends KeyringNode { encryptedDataKeys: EncryptedDataKey[] ): Promise cacheEntryHasExceededLimits(entry: BranchKeyMaterialEntry): boolean + _branchKeyMaterialsInFlight: Map> } export class KmsHierarchicalKeyRingNode @@ -104,6 +106,10 @@ export class KmsHierarchicalKeyRingNode public declare cacheLimitTtl: number public declare maxCacheSize?: number public declare _cmc: CryptographicMaterialsCache + public declare _branchKeyMaterialsInFlight: Map< + string, + Promise + > declare readonly _partition: Buffer public declare _utf8Sorting: boolean @@ -259,6 +265,8 @@ export class KmsHierarchicalKeyRingNode readOnlyProperty(this, 'maxCacheSize', maxCacheSize) readOnlyProperty(this, '_cmc', cache) + readOnlyProperty(this, '_branchKeyMaterialsInFlight', new Map()) + if (utf8Sorting === undefined) { readOnlyProperty(this, '_utf8Sorting', false) } else { diff --git a/modules/kms-keyring-node/src/kms_hkeyring_node_helpers.ts b/modules/kms-keyring-node/src/kms_hkeyring_node_helpers.ts index d142c9341..6a35227b4 100644 --- a/modules/kms-keyring-node/src/kms_hkeyring_node_helpers.ts +++ b/modules/kms-keyring-node/src/kms_hkeyring_node_helpers.ts @@ -201,48 +201,59 @@ export async function getBranchKeyMaterials( let branchKeyMaterials: NodeBranchKeyMaterial // if the cache entry is false, branch key materials were not found if (!cacheEntry || hKeyring.cacheEntryHasExceededLimits(cacheEntry)) { - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#onencrypt - //# If this is NOT true, then we MUST treat the cache entry as expired. + /* Concurrent misses for the same cache entry share one keystore request. + * Without this, N decrypts that start before the cache is populated each + * fire their own DynamoDB GetItem and KMS Decrypt. */ + branchKeyMaterials = await ensureBranchKeyMaterialsInFlight( + hKeyring._branchKeyMaterialsInFlight, + cacheEntryId, + async () => { + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#onencrypt + //# If this is NOT true, then we MUST treat the cache entry as expired. + + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#ondecrypt + //# If this is NOT true, then we MUST treat the cache entry as expired. + + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#onencrypt + //# If a cache entry is not found or the cache entry is expired, the hierarchical keyring MUST attempt to obtain the branch key materials + //# by querying the backing branch keystore specified in the [retrieve OnEncrypt branch key materials](#query-branch-keystore-onencrypt) section. + //# If the keyring is not able to retrieve [branch key materials](../structures.md#branch-key-materials) + //# through the underlying cryptographic materials cache or + //# it no longer has access to them through the backing keystore, OnEncrypt MUST fail. - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#ondecrypt - //# If this is NOT true, then we MUST treat the cache entry as expired. + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#query-branch-keystore-onencrypt + //# Otherwise, OnEncrypt MUST fail. + + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#getitem-branch-keystore-ondecrypt + //# Otherwise, OnDecrypt MUST fail. - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#onencrypt - //# If a cache entry is not found or the cache entry is expired, the hierarchical keyring MUST attempt to obtain the branch key materials - //# by querying the backing branch keystore specified in the [retrieve OnEncrypt branch key materials](#query-branch-keystore-onencrypt) section. - //# If the keyring is not able to retrieve [branch key materials](../structures.md#branch-key-materials) - //# through the underlying cryptographic materials cache or - //# it no longer has access to them through the backing keystore, OnEncrypt MUST fail. - - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#query-branch-keystore-onencrypt - //# Otherwise, OnEncrypt MUST fail. - - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#getitem-branch-keystore-ondecrypt - //# Otherwise, OnDecrypt MUST fail. - - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#query-branch-keystore-onencrypt - //# OnEncrypt MUST call the Keystore's [GetActiveBranchKey](../branch-key-store.md#getactivebranchkey) operation with the following inputs: - - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#getitem-branch-keystore-ondecrypt - //# OnDecrypt MUST call the Keystore's [GetBranchKeyVersion](../branch-key-store.md#getbranchkeyversion) operation with the following inputs: - branchKeyMaterials = branchKeyVersion - ? await keyStore.getBranchKeyVersion(branchKeyId, branchKeyVersion) - : // The complice needs a line //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#query-branch-keystore-onencrypt //# OnEncrypt MUST call the Keystore's [GetActiveBranchKey](../branch-key-store.md#getactivebranchkey) operation with the following inputs: - //# - the `branchKeyId` used in this operation - await keyStore.getActiveBranchKey(branchKeyId) - - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#query-branch-keystore-onencrypt - //# If the Keystore's GetActiveBranchKey operation succeeds - //# the keyring MUST put the returned branch key materials in the cache using the - //# formula defined in [Appendix A](#appendix-a-cache-entry-identifier-formulas). - - //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#getitem-branch-keystore-ondecrypt - //# If the Keystore's GetBranchKeyVersion operation succeeds - //# the keyring MUST put the returned branch key materials in the cache using the - //# formula defined in [Appendix A](#appendix-a-cache-entry-identifier-formulas). - cmc.putBranchKeyMaterial(cacheEntryId, branchKeyMaterials, cacheLimitTtl) + + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#getitem-branch-keystore-ondecrypt + //# OnDecrypt MUST call the Keystore's [GetBranchKeyVersion](../branch-key-store.md#getbranchkeyversion) operation with the following inputs: + const materials = branchKeyVersion + ? await keyStore.getBranchKeyVersion(branchKeyId, branchKeyVersion) + : // The complice needs a line + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#query-branch-keystore-onencrypt + //# OnEncrypt MUST call the Keystore's [GetActiveBranchKey](../branch-key-store.md#getactivebranchkey) operation with the following inputs: + //# - the `branchKeyId` used in this operation + await keyStore.getActiveBranchKey(branchKeyId) + + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#query-branch-keystore-onencrypt + //# If the Keystore's GetActiveBranchKey operation succeeds + //# the keyring MUST put the returned branch key materials in the cache using the + //# formula defined in [Appendix A](#appendix-a-cache-entry-identifier-formulas). + + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#getitem-branch-keystore-ondecrypt + //# If the Keystore's GetBranchKeyVersion operation succeeds + //# the keyring MUST put the returned branch key materials in the cache using the + //# formula defined in [Appendix A](#appendix-a-cache-entry-identifier-formulas). + cmc.putBranchKeyMaterial(cacheEntryId, materials, cacheLimitTtl) + + return materials + } + ) } else { //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#ondecrypt //# If a cache entry is found and the entry's TTL has not expired, the hierarchical keyring MUST use those branch key materials for key unwrapping. @@ -255,6 +266,31 @@ export async function getBranchKeyMaterials( return branchKeyMaterials } +// Coalesces concurrent misses for one cache entry onto a single in-flight +// request, evicted on settle so the cryptographic materials cache (not this +// map) keeps ownership of caching and TTL. A rejected request is evicted too, +// so the next call retries rather than sharing the failure. +async function ensureBranchKeyMaterialsInFlight( + inFlight: Map>, + cacheEntryId: string, + fetch: () => Promise +): Promise { + let pending = inFlight.get(cacheEntryId) + if (!pending) { + pending = fetch() + inFlight.set(cacheEntryId, pending) + } + + try { + const branchKeyMaterials = await pending + inFlight.delete(cacheEntryId) + return branchKeyMaterials + } catch (error) { + inFlight.delete(cacheEntryId) + throw error + } +} + //= aws-encryption-sdk-specification/framework/aws-kms/aws-kms-hierarchical-keyring.md#onencrypt //# If the input [encryption materials](../structures.md#encryption-materials) do not contain a plaintext data key, //# OnEncrypt MUST generate a random plaintext data key, according to the key length defined in the [algorithm suite](../algorithm-suites.md#encryption-key-length). diff --git a/modules/kms-keyring-node/test/kms_hkeyring_node.concurrency.test.ts b/modules/kms-keyring-node/test/kms_hkeyring_node.concurrency.test.ts new file mode 100644 index 000000000..6c0148af9 --- /dev/null +++ b/modules/kms-keyring-node/test/kms_hkeyring_node.concurrency.test.ts @@ -0,0 +1,54 @@ +// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { expect } from 'chai' +import { getBranchKeyMaterials } from '../src/kms_hkeyring_node_helpers' +import { getLocalCryptographicMaterialsCache } from '@aws-crypto/cache-material' +import { + NodeAlgorithmSuite, + NodeBranchKeyMaterial, +} from '@aws-crypto/material-management' +import { v4 } from 'uuid' + +const CONCURRENT_DECRYPTS = 3000 +const CACHE_LIMIT_TTL = 60_000 +const CACHE_ENTRY_ID = 'shared-cache-entry-id' + +function fixtureMaterial(): NodeBranchKeyMaterial { + return new NodeBranchKeyMaterial(Buffer.alloc(32), 'branchKeyId', v4(), {}) +} + +describe('KmsHierarchicalKeyRingNode: concurrent branch key retrieval', () => { + it(`coalesces ${CONCURRENT_DECRYPTS} concurrent cache misses into one keystore call`, async () => { + let getBranchKeyVersionCalls = 0 + const keyStore = { + async getBranchKeyVersion() { + getBranchKeyVersionCalls += 1 + await new Promise((resolve) => setTimeout(resolve, 5)) + return fixtureMaterial() + }, + } + + const cmc = getLocalCryptographicMaterialsCache(100) + const hKeyring = { + keyStore, + cacheLimitTtl: CACHE_LIMIT_TTL, + cacheEntryHasExceededLimits: () => false, + _branchKeyMaterialsInFlight: new Map(), + } as any + + await Promise.all( + Array.from({ length: CONCURRENT_DECRYPTS }, () => + getBranchKeyMaterials( + hKeyring, + cmc, + 'branchKeyId', + CACHE_ENTRY_ID, + 'branchKeyVersion' + ) + ) + ) + + expect(getBranchKeyVersionCalls).to.equal(1) + }) +})