fix(kms-keyring-node): coalesce concurrent branch key cache misses

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<NodeDecryptionMaterial>
   cacheEntryHasExceededLimits(entry: BranchKeyMaterialEntry): boolean
+  _branchKeyMaterialsInFlight: Map<string, Promise<NodeBranchKeyMaterial>>
 }
 
 export class KmsHierarchicalKeyRingNode
@@ -104,6 +106,10 @@ export class KmsHierarchicalKeyRingNode
   public declare cacheLimitTtl: number
   public declare maxCacheSize?: number
   public declare _cmc: CryptographicMaterialsCache<NodeAlgorithmSuite>
+  public declare _branchKeyMaterialsInFlight: Map<
+    string,
+    Promise<NodeBranchKeyMaterial>
+  >
   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 {

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<string, Promise<NodeBranchKeyMaterial>>,
+  cacheEntryId: string,
+  fetch: () => Promise<NodeBranchKeyMaterial>
+): Promise<NodeBranchKeyMaterial> {
+  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).

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<NodeAlgorithmSuite>(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)
+  })
+})