feat: Improvements to the message decryption process

.gitmodules

@@ -1,3 +1,4 @@
 [submodule "aws-encryption-sdk-test-vectors"]
 	path = aws-encryption-sdk-test-vectors
-	url = https://github.com/awslabs/aws-encryption-sdk-test-vectors.git
+	url = https://github.com/awslabs/private-aws-encryption-sdk-test-vectors-staging.git
+	branch = known-invalid-test-vectors

modules/client-browser/src/index.ts

@@ -10,16 +10,19 @@ export * from '@aws-crypto/raw-aes-keyring-browser'
 export * from '@aws-crypto/raw-rsa-keyring-browser'
 export * from '@aws-crypto/web-crypto-backend'
 
-import { CommitmentPolicy } from '@aws-crypto/material-management-browser'
+import {
+  CommitmentPolicy,
+  ClientOptions,
+} from '@aws-crypto/material-management-browser'
 
 import { buildEncrypt } from '@aws-crypto/encrypt-browser'
 import { buildDecrypt } from '@aws-crypto/decrypt-browser'
 
 export function buildClient(
-  commitmentPolicy: CommitmentPolicy
+  options: CommitmentPolicy | ClientOptions
 ): ReturnType<typeof buildEncrypt> & ReturnType<typeof buildDecrypt> {
   return {
-    ...buildEncrypt(commitmentPolicy),
-    ...buildDecrypt(commitmentPolicy),
+    ...buildEncrypt(options),
+    ...buildDecrypt(options),
   }
 }

modules/client-node/src/index.ts

@@ -9,16 +9,20 @@ export * from '@aws-crypto/kms-keyring-node'
 export * from '@aws-crypto/raw-aes-keyring-node'
 export * from '@aws-crypto/raw-rsa-keyring-node'
 
-import { CommitmentPolicy } from '@aws-crypto/material-management-node'
+import {
+  CommitmentPolicy,
+  ClientOptions,
+} from '@aws-crypto/material-management-node'
 
 import { buildEncrypt } from '@aws-crypto/encrypt-node'
-import { buildDecrypt } from '@aws-crypto/decrypt-node'
+import { buildDecrypt, DecryptOutput } from '@aws-crypto/decrypt-node'
 
 export function buildClient(
-  commitmentPolicy: CommitmentPolicy
+  options: CommitmentPolicy | ClientOptions
 ): ReturnType<typeof buildEncrypt> & ReturnType<typeof buildDecrypt> {
   return {
-    ...buildEncrypt(commitmentPolicy),
-    ...buildDecrypt(commitmentPolicy),
+    ...buildEncrypt(options),
+    ...buildDecrypt(options),
   }
 }
+export { DecryptOutput }

modules/decrypt-browser/src/decrypt.ts

@@ -11,6 +11,7 @@ import {
   WebCryptoMaterialsManager,
   CommitmentPolicy,
   CommitmentPolicySuites,
+  ClientOptions,
 } from '@aws-crypto/material-management-browser'
 import {
   deserializeSignature,
@@ -34,20 +35,29 @@ export interface DecryptResult {
 }
 
 export async function _decrypt(
-  commitmentPolicy: CommitmentPolicy,
+  { commitmentPolicy, maxEncryptedDataKeys }: ClientOptions,
   cmm: KeyringWebCrypto | WebCryptoMaterialsManager,
   ciphertext: Uint8Array
 ): Promise<DecryptResult> {
   /* Precondition: _decrypt needs a valid commitmentPolicy. */
   needs(CommitmentPolicy[commitmentPolicy], 'Invalid commitment policy.')
 
+  // buildDecrypt defaults this to false for backwards compatibility, so this is satisfied
+  /* Precondition: _decrypt needs a valid maxEncryptedDataKeys. */
+  needs(
+    maxEncryptedDataKeys === false || maxEncryptedDataKeys >= 1,
+    'Invalid maxEncryptedDataKeys value.'
+  )
+
   /* If the cmm is a Keyring, wrap it with WebCryptoDefaultCryptographicMaterialsManager. */
   cmm =
     cmm instanceof KeyringWebCrypto
       ? new WebCryptoDefaultCryptographicMaterialsManager(cmm)
       : cmm
 
-  const headerInfo = deserialize.deserializeMessageHeader(ciphertext)
+  const headerInfo = deserialize.deserializeMessageHeader(ciphertext, {
+    maxEncryptedDataKeys,
+  })
   if (headerInfo === false) throw new Error('Unable to parse Header')
   const { messageHeader, algorithmSuite } = headerInfo
   const { rawHeader, headerAuth } = headerInfo

modules/decrypt-browser/src/decrypt_client.ts

@@ -4,6 +4,7 @@
 import { _decrypt } from './decrypt'
 import {
   CommitmentPolicy,
+  ClientOptions,
   needs,
 } from '@aws-crypto/material-management-browser'
 
@@ -15,13 +16,26 @@ type CurryFirst<fn extends (...a: any[]) => any> = fn extends (
   : []
 
 export function buildDecrypt(
-  commitmentPolicy: CommitmentPolicy
+  options: CommitmentPolicy | ClientOptions
 ): {
   decrypt: (...args: CurryFirst<typeof _decrypt>) => ReturnType<typeof _decrypt>
 } {
+  const { commitmentPolicy, maxEncryptedDataKeys = false } =
+    typeof options === 'string' ? { commitmentPolicy: options } : options
+
   /* Precondition: browser buildDecrypt needs a valid commitmentPolicy. */
   needs(CommitmentPolicy[commitmentPolicy], 'Invalid commitment policy.')
+  /* Precondition: browser buildDecrypt needs a valid maxEncryptedDataKeys. */
+  needs(
+    maxEncryptedDataKeys === false || maxEncryptedDataKeys >= 1,
+    'Invalid maxEncryptedDataKeys value.'
+  )
+
+  const clientOptions: ClientOptions = {
+    commitmentPolicy,
+    maxEncryptedDataKeys,
+  }
   return {
-    decrypt: _decrypt.bind({}, commitmentPolicy),
+    decrypt: _decrypt.bind({}, clientOptions),
   }
 }

modules/decrypt-browser/src/index.ts

@@ -3,6 +3,7 @@
 
 import { CommitmentPolicy } from '@aws-crypto/material-management-browser'
 import { buildDecrypt } from './decrypt_client'
+export { DecryptResult } from './decrypt'
 export { MessageHeader } from '@aws-crypto/serialize'
 /** @deprecated Use `buildDecrypt(CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT)` for migration. */
 export const { decrypt } = buildDecrypt(

modules/decrypt-browser/test/decrypt.test.ts

@@ -5,7 +5,7 @@
 
 import * as chai from 'chai'
 import chaiAsPromised from 'chai-as-promised'
-import { decrypt } from '../src/index'
+import { decrypt, buildDecrypt } from '../src/index'
 import { _decrypt } from '../src/decrypt'
 import {
   AlgorithmSuiteIdentifier,
@@ -14,6 +14,7 @@ import {
   KeyringWebCrypto,
   WebCryptoDecryptionMaterial,
   WebCryptoEncryptionMaterial,
+  CommitmentPolicy,
 } from '@aws-crypto/material-management-browser'
 import * as fixtures from './fixtures'
 import { MessageFormat } from '@aws-crypto/material-management'
@@ -38,10 +39,27 @@ describe('decrypt', () => {
 
   it('Precondition: _decrypt needs a valid commitmentPolicy.', async () => {
     await expect(
-      _decrypt('fake_policy' as any, {} as any, {} as any)
+      _decrypt(
+        { commitmentPolicy: 'fake_policy' as any, maxEncryptedDataKeys: false },
+        {} as any,
+        {} as any
+      )
     ).to.rejectedWith(Error, 'Invalid commitment policy.')
   })
 
+  it('Precondition: _decrypt needs a valid maxEncryptedDataKeys.', async () => {
+    await expect(
+      _decrypt(
+        {
+          commitmentPolicy: CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+          maxEncryptedDataKeys: 0,
+        },
+        {} as any,
+        {} as any
+      )
+    ).to.rejectedWith(Error, 'Invalid maxEncryptedDataKeys value.')
+  })
+
   // it('The parsed header algorithmSuite in _decrypt must be supported by the commitmentPolicy.', async () => {
   //   await expect(
   //     _decrypt(
@@ -136,6 +154,40 @@ describe('decrypt', () => {
       await expect(decrypt(keyring, data.slice(0, i))).to.rejectedWith(Error)
     }
   })
+
+  it('can decrypt data with less than maxEncryptedDataKeys', async () => {
+    const { decrypt } = buildDecrypt({
+      commitmentPolicy: CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+      maxEncryptedDataKeys: 3,
+    })
+    const { plaintext } = await decrypt(
+      fixtures.decryptKeyring(),
+      fixtures.twoEdksMessage()
+    )
+    expect(plaintext).to.deep.equal(Buffer.from('asdf'))
+  })
+
+  it('can decrypt data with exactly maxEncryptedDataKeys', async () => {
+    const { decrypt } = buildDecrypt({
+      commitmentPolicy: CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+      maxEncryptedDataKeys: 3,
+    })
+    const { plaintext } = await decrypt(
+      fixtures.decryptKeyring(),
+      fixtures.threeEdksMessage()
+    )
+    expect(plaintext).to.deep.equal(Buffer.from('asdf'))
+  })
+
+  it('will not decrypt data with more than maxEncryptedDataKeys', async () => {
+    const { decrypt } = buildDecrypt({
+      commitmentPolicy: CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+      maxEncryptedDataKeys: 3,
+    })
+    await expect(
+      decrypt(fixtures.decryptKeyring(), fixtures.fourEdksMessage())
+    ).to.rejectedWith(Error, 'maxEncryptedDataKeys exceeded.')
+  })
 })
 
 // prettier-ignore

modules/decrypt-browser/test/decrypt_client.test.ts

@@ -12,12 +12,39 @@ chai.use(chaiAsPromised)
 const { expect } = chai
 
 describe('buildDecrypt', () => {
-  it('can build a client', () => {
+  it('can build a client with a commitment policy', () => {
     const test = buildDecrypt(CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT)
     expect(test).to.have.property('decrypt').and.to.be.a('function')
   })
 
+  it('can build a client with max encrypted data keys', () => {
+    for (const numKeys of [1, 10, Math.pow(2, 16) - 1, Math.pow(2, 16)]) {
+      const test = buildDecrypt({
+        commitmentPolicy: CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+        maxEncryptedDataKeys: numKeys,
+      })
+      expect(test).to.have.property('decrypt').and.to.be.a('function')
+    }
+  })
+
   it('Precondition: browser buildDecrypt needs a valid commitmentPolicy.', () => {
-    expect(() => buildDecrypt({} as any)).to.throw('Invalid commitment policy.')
+    expect(() => buildDecrypt('BAD_POLICY' as any)).to.throw(
+      'Invalid commitment policy.'
+    )
+  })
+
+  it('Precondition: browser buildDecrypt needs a valid maxEncryptedDataKeys.', () => {
+    expect(() =>
+      buildDecrypt({
+        commitmentPolicy: CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+        maxEncryptedDataKeys: 0,
+      })
+    ).to.throw('Invalid maxEncryptedDataKeys value.')
+    expect(() =>
+      buildDecrypt({
+        commitmentPolicy: CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+        maxEncryptedDataKeys: -1,
+      })
+    ).to.throw('Invalid maxEncryptedDataKeys value.')
   })
 })

modules/decrypt-browser/test/fixtures.ts

@@ -250,6 +250,27 @@ export function frameSequenceOutOfOrder() {
   ])
 }
 
+export function twoEdksMessage() {
+  // prettier-ignore
+  return new Uint8Array([
+    1,128,0,20,81,159,165,226,136,253,2,91,121,212,103,155,181,101,192,42,0,0,0,2,0,4,116,101,109,112,0,24,107,101,121,49,0,0,0,128,0,0,0,12,253,51,172,10,239,164,139,246,179,46,90,151,0,32,112,49,205,185,161,227,110,112,214,205,129,210,255,139,70,54,55,191,225,87,28,109,71,168,124,168,187,90,185,220,183,111,0,4,116,101,109,112,0,24,107,101,121,48,0,0,0,128,0,0,0,12,130,105,57,211,107,129,63,61,28,113,29,245,0,32,223,169,9,70,195,85,36,140,164,225,228,190,70,137,36,66,153,118,45,227,30,48,41,64,4,38,172,46,186,38,100,204,2,0,0,0,0,12,0,0,16,0,0,0,0,0,0,0,0,0,0,0,0,0,241,48,199,72,160,81,240,122,45,194,9,196,238,161,95,137,255,255,255,255,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,4,15,114,189,86,167,87,157,63,128,15,95,53,55,99,185,59,74,12,150,52
+  ])
+}
+
+export function threeEdksMessage() {
+  // prettier-ignore
+  return new Uint8Array([
+    1,128,0,20,6,147,84,28,195,247,246,16,92,66,254,35,91,134,21,113,0,0,0,3,0,4,116,101,109,112,0,24,107,101,121,50,0,0,0,128,0,0,0,12,243,10,221,12,30,48,159,193,178,76,143,111,0,32,164,227,105,185,29,231,91,130,47,214,222,135,16,165,55,47,97,135,251,159,33,232,143,25,70,73,55,217,248,216,69,44,0,4,116,101,109,112,0,24,107,101,121,48,0,0,0,128,0,0,0,12,146,167,240,23,105,252,154,26,125,18,74,76,0,32,243,219,32,78,83,237,209,122,150,92,241,94,91,201,171,7,110,218,233,15,20,56,246,15,48,131,41,225,114,157,197,252,0,4,116,101,109,112,0,24,107,101,121,49,0,0,0,128,0,0,0,12,215,162,33,39,252,66,105,179,78,16,42,232,0,32,30,167,255,38,200,18,212,153,115,229,175,89,11,101,31,244,6,146,25,12,80,32,22,53,53,63,239,107,104,139,65,167,2,0,0,0,0,12,0,0,16,0,0,0,0,0,0,0,0,0,0,0,0,0,25,251,17,157,236,34,27,100,13,10,155,194,218,166,195,57,255,255,255,255,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,4,15,114,189,86,231,188,115,162,175,197,222,3,249,202,81,98,227,231,67,146
+  ])
+}
+
+export function fourEdksMessage() {
+  // prettier-ignore
+  return new Uint8Array([
+    1,128,0,20,152,195,218,126,82,140,170,75,107,29,97,41,37,39,140,194,0,0,0,4,0,4,116,101,109,112,0,24,107,101,121,50,0,0,0,128,0,0,0,12,232,117,224,17,241,139,48,60,112,45,223,30,0,32,90,160,47,236,89,127,74,188,210,95,186,61,117,166,82,207,243,171,211,192,190,29,63,83,198,122,54,25,74,198,123,0,0,4,116,101,109,112,0,24,107,101,121,48,0,0,0,128,0,0,0,12,179,67,229,65,157,148,209,70,133,50,200,225,0,32,114,101,189,108,89,151,6,211,84,211,236,72,253,45,230,183,9,40,79,146,106,109,224,46,254,102,50,129,174,123,49,247,0,4,116,101,109,112,0,24,107,101,121,49,0,0,0,128,0,0,0,12,122,245,182,159,26,134,245,191,73,101,240,68,0,32,6,247,51,137,165,160,143,119,82,111,242,247,21,58,45,184,62,75,131,72,192,43,205,106,120,75,40,128,102,165,73,128,0,4,116,101,109,112,0,24,107,101,121,51,0,0,0,128,0,0,0,12,58,191,161,65,150,32,99,20,197,143,126,208,0,32,179,18,132,26,48,25,198,187,236,47,191,24,24,153,205,76,158,181,61,82,249,110,192,95,153,90,0,113,113,142,36,173,2,0,0,0,0,12,0,0,16,0,0,0,0,0,0,0,0,0,0,0,0,0,215,170,96,229,229,32,220,210,179,153,146,71,139,224,7,208,255,255,255,255,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,4,15,114,189,86,12,123,234,170,45,72,187,11,199,119,21,216,211,21,95,51
+  ])
+}
+
 class TestKeyring extends KeyringWebCrypto {
   async _onEncrypt(): Promise<WebCryptoEncryptionMaterial> {
     throw new Error('I should never see this error')

modules/decrypt-node/src/decipher_stream.ts

@@ -104,7 +104,7 @@ export function getDecipherStream() {
       decipherState.contentLength -= chunk.length
       /* Precondition: Only content should be transformed, so the lengths must always match.
        * The BodyHeader and AuthTag are striped in the VerifyStream and passed in
-       * through events.  This means that if I receive a chunk without havening reset
+       * through events.  This means that if I receive a chunk without having reset
        * the content accumulation events are out of order.  Panic.
        */
       needs(decipherState.contentLength >= 0, 'Lengths do not match')

modules/decrypt-node/src/decrypt.ts

@@ -4,9 +4,9 @@
 import {
   NodeMaterialsManager,
   KeyringNode,
-  CommitmentPolicy,
 } from '@aws-crypto/material-management-node'
 import { _decryptStream } from './decrypt_stream'
+import { DecryptParameters } from './types'
 
 // @ts-ignore
 import { finished } from 'readable-stream'
@@ -24,13 +24,12 @@ export interface DecryptOptions {
 }
 
 export async function _decrypt(
-  commitmentPolicy: CommitmentPolicy,
+  decryptParameters: DecryptParameters,
   cmm: NodeMaterialsManager | KeyringNode,
   ciphertext: Buffer | Uint8Array | Readable | string | NodeJS.ReadableStream,
   { encoding, maxBodySize }: DecryptOptions = {}
 ): Promise<DecryptOutput> {
-  const stream = _decryptStream(commitmentPolicy, cmm, { maxBodySize })
-
+  const stream = _decryptStream(decryptParameters, cmm, { maxBodySize })
   const plaintext: Buffer[] = []
   let messageHeader: MessageHeader | false = false
   stream