fix: Using discovery filter while constructing keyring

modules/kms-keyring-browser/src/kms_keyring_browser.ts

@@ -46,8 +46,16 @@ export class KmsKeyringBrowser extends KmsKeyringClass<
     generatorKeyId,
     grantTokens,
     discovery,
+    discoveryFilter,
   }: KmsKeyringWebCryptoInput = {}) {
-    super({ clientProvider, keyIds, generatorKeyId, grantTokens, discovery })
+    super({
+      clientProvider,
+      keyIds,
+      generatorKeyId,
+      grantTokens,
+      discovery,
+      discoveryFilter,
+    })
   }
 
   async _onEncrypt(material: WebCryptoEncryptionMaterial) {

modules/kms-keyring-browser/test/kms_keyring_browser.test.ts

@@ -48,6 +48,102 @@ describe('KmsKeyringBrowser::constructor', () => {
     const test = new KmsKeyringBrowser({ discovery: true })
     expect(test instanceof KeyringWebCrypto).to.equal(true)
   })
+
+  it('forwards discoveryFilter to base KmsKeyring', () => {
+    const discoveryFilter = { accountIDs: ['123456789012'], partition: 'aws' }
+    const test = new KmsKeyringBrowser({ discovery: true, discoveryFilter })
+    expect(test.isDiscovery).to.equal(true)
+    expect(test.discoveryFilter).to.deep.equal(discoveryFilter)
+  })
+
+  it('discoveryFilter excludes EDKs from non-allowed accounts on decrypt', async () => {
+    const allowedAccount = '111111111111'
+    const otherAccount = '222222222222'
+    const allowedArn = `arn:aws:kms:us-east-1:${allowedAccount}:key/12345678-1234-1234-1234-123456789012`
+    const otherArn = `arn:aws:kms:us-east-1:${otherAccount}:key/12345678-1234-1234-1234-123456789012`
+
+    const decryptCalls: string[] = []
+    const mockClientProvider: any = () => ({
+      decrypt: ({ KeyId }: any) => {
+        decryptCalls.push(KeyId)
+        return {
+          Plaintext: new Uint8Array(16),
+          KeyId,
+        }
+      },
+    })
+
+    const keyring = new KmsKeyringBrowser({
+      clientProvider: mockClientProvider,
+      discovery: true,
+      discoveryFilter: { accountIDs: [allowedAccount], partition: 'aws' },
+    })
+
+    const suite = new WebCryptoAlgorithmSuite(
+      AlgorithmSuiteIdentifier.ALG_AES128_GCM_IV12_TAG16
+    )
+    const material = new WebCryptoDecryptionMaterial(suite, {})
+    const edks = [
+      new EncryptedDataKey({
+        providerId: 'aws-kms',
+        providerInfo: otherArn,
+        encryptedDataKey: new Uint8Array(Buffer.from(otherArn)),
+      }),
+      new EncryptedDataKey({
+        providerId: 'aws-kms',
+        providerInfo: allowedArn,
+        encryptedDataKey: new Uint8Array(Buffer.from(allowedArn)),
+      }),
+    ]
+
+    await keyring.onDecrypt(material, edks)
+
+    expect(decryptCalls).to.deep.equal([allowedArn])
+  })
+
+  it('throws when discoveryFilter has empty accountIDs', () => {
+    expect(
+      () =>
+        new KmsKeyringBrowser({
+          discovery: true,
+          discoveryFilter: { accountIDs: [], partition: 'aws' },
+        })
+    ).to.throw('A discovery filter must be able to match something.')
+  })
+
+  it('throws when discoveryFilter has empty partition', () => {
+    expect(
+      () =>
+        new KmsKeyringBrowser({
+          discovery: true,
+          discoveryFilter: { accountIDs: ['123456789012'], partition: '' },
+        })
+    ).to.throw('A discovery filter must be able to match something.')
+  })
+
+  it('throws when discoveryFilter accountIDs contains an empty string', () => {
+    expect(
+      () =>
+        new KmsKeyringBrowser({
+          discovery: true,
+          discoveryFilter: { accountIDs: [''], partition: 'aws' },
+        })
+    ).to.throw('A discovery filter must be able to match something.')
+  })
+
+  it('throws when discoveryFilter is set without discovery=true', () => {
+    expect(
+      () =>
+        new KmsKeyringBrowser({
+          keyIds: [
+            'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012',
+          ],
+          discoveryFilter: { accountIDs: ['123456789012'], partition: 'aws' },
+        })
+    ).to.throw(
+      'Account and partition decrypt filtering are only supported when discovery === true'
+    )
+  })
 })
 
 describe('KmsKeyringBrowser can encrypt/decrypt with AWS SDK v3 client', () => {

modules/kms-keyring-node/src/kms_keyring_node.ts

@@ -39,8 +39,16 @@ export class KmsKeyringNode extends KmsKeyringClass<
     generatorKeyId,
     grantTokens,
     discovery,
+    discoveryFilter,
   }: KmsKeyringNodeInput = {}) {
-    super({ clientProvider, keyIds, generatorKeyId, grantTokens, discovery })
+    super({
+      clientProvider,
+      keyIds,
+      generatorKeyId,
+      grantTokens,
+      discovery,
+      discoveryFilter,
+    })
   }
 }
 immutableClass(KmsKeyringNode)

modules/kms-keyring-node/test/kms_keyring_node.test.ts

@@ -37,6 +37,104 @@ describe('KmsKeyringNode::constructor', () => {
     const test = new KmsKeyringNode({ discovery: true })
     expect(test instanceof KeyringNode).to.equal(true)
   })
+
+  it('forwards discoveryFilter to base KmsKeyring', () => {
+    const discoveryFilter = { accountIDs: ['123456789012'], partition: 'aws' }
+    const test = new KmsKeyringNode({ discovery: true, discoveryFilter })
+    expect(test.isDiscovery).to.equal(true)
+    expect(test.discoveryFilter).to.deep.equal(discoveryFilter)
+  })
+
+  it('discoveryFilter excludes EDKs from non-allowed accounts on decrypt', async () => {
+    const allowedAccount = '111111111111'
+    const otherAccount = '222222222222'
+    const allowedArn = `arn:aws:kms:us-east-1:${allowedAccount}:key/12345678-1234-1234-1234-123456789012`
+    const otherArn = `arn:aws:kms:us-east-1:${otherAccount}:key/12345678-1234-1234-1234-123456789012`
+
+    const decryptCalls: string[] = []
+    const clientProvider: any = () => ({
+      decrypt: ({ KeyId }: any) => {
+        decryptCalls.push(KeyId)
+        // Always succeed for the keys the keyring chooses to call.
+        return {
+          Plaintext: new Uint8Array(16),
+          KeyId,
+        }
+      },
+    })
+
+    const keyring = new KmsKeyringNode({
+      clientProvider,
+      discovery: true,
+      discoveryFilter: { accountIDs: [allowedAccount], partition: 'aws' },
+    })
+
+    const suite = new NodeAlgorithmSuite(
+      AlgorithmSuiteIdentifier.ALG_AES128_GCM_IV12_TAG16
+    )
+    const material = new NodeDecryptionMaterial(suite, {})
+    const edks = [
+      new EncryptedDataKey({
+        providerId: 'aws-kms',
+        providerInfo: otherArn,
+        encryptedDataKey: Buffer.from(otherArn),
+      }),
+      new EncryptedDataKey({
+        providerId: 'aws-kms',
+        providerInfo: allowedArn,
+        encryptedDataKey: Buffer.from(allowedArn),
+      }),
+    ]
+
+    await keyring.onDecrypt(material, edks)
+
+    // Only the allowed-account EDK should reach KMS.
+    expect(decryptCalls).to.deep.equal([allowedArn])
+  })
+
+  it('throws when discoveryFilter has empty accountIDs', () => {
+    expect(
+      () =>
+        new KmsKeyringNode({
+          discovery: true,
+          discoveryFilter: { accountIDs: [], partition: 'aws' },
+        })
+    ).to.throw('A discovery filter must be able to match something.')
+  })
+
+  it('throws when discoveryFilter has empty partition', () => {
+    expect(
+      () =>
+        new KmsKeyringNode({
+          discovery: true,
+          discoveryFilter: { accountIDs: ['123456789012'], partition: '' },
+        })
+    ).to.throw('A discovery filter must be able to match something.')
+  })
+
+  it('throws when discoveryFilter accountIDs contains an empty string', () => {
+    expect(
+      () =>
+        new KmsKeyringNode({
+          discovery: true,
+          discoveryFilter: { accountIDs: [''], partition: 'aws' },
+        })
+    ).to.throw('A discovery filter must be able to match something.')
+  })
+
+  it('throws when discoveryFilter is set without discovery=true', () => {
+    expect(
+      () =>
+        new KmsKeyringNode({
+          keyIds: [
+            'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012',
+          ],
+          discoveryFilter: { accountIDs: ['123456789012'], partition: 'aws' },
+        })
+    ).to.throw(
+      'Account and partition decrypt filtering are only supported when discovery === true'
+    )
+  })
 })
 
 describe('KmsKeyringNode can encrypt/decrypt with AWS SDK v3 client', () => {