From 03d6d013bf6e070e85adfe5731f526978e3e8e4d Mon Sep 17 00:00:00 2001
From: Filip Skokan <panva.ip@gmail.com>
Date: Wed, 31 Aug 2022 15:47:55 +0200
Subject: [PATCH] fix: limit default PBES2 alg's computational expense

Also adds an option to opt-in to higher computation expense.
---
 src/jwe/flattened/decrypt.ts        |  4 ++--
 src/lib/decrypt_key_management.ts   |  8 +++++++-
 src/types.d.ts                      |  7 +++++++
 test/jwe/flattened.decrypt.test.mjs | 18 +++++++++++++++---
 4 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/jwe/flattened/decrypt.ts b/src/jwe/flattened/decrypt.ts
index 777e9e5756..afb621a1dc 100644
--- a/src/jwe/flattened/decrypt.ts
+++ b/src/jwe/flattened/decrypt.ts
@@ -188,9 +188,9 @@ export async function flattenedDecrypt(
 
   let cek: KeyLike | Uint8Array
   try {
-    cek = await decryptKeyManagement(alg, key, encryptedKey, joseHeader)
+    cek = await decryptKeyManagement(alg, key, encryptedKey, joseHeader, options)
   } catch (err) {
-    if (err instanceof TypeError) {
+    if (err instanceof TypeError || err instanceof JWEInvalid || err instanceof JOSENotSupported) {
       throw err
     }
     // https://www.rfc-editor.org/rfc/rfc7516#section-11.5
diff --git a/src/lib/decrypt_key_management.ts b/src/lib/decrypt_key_management.ts
index 3e4f9068ff..7279306559 100644
--- a/src/lib/decrypt_key_management.ts
+++ b/src/lib/decrypt_key_management.ts
@@ -4,7 +4,7 @@ import { decrypt as pbes2Kw } from '../runtime/pbes2kw.js'
 import { decrypt as rsaEs } from '../runtime/rsaes.js'
 import { decode as base64url } from '../runtime/base64url.js'
 
-import type { JWEHeaderParameters, KeyLike, JWK } from '../types.d'
+import type { DecryptOptions, JWEHeaderParameters, KeyLike, JWK } from '../types.d'
 import { JOSENotSupported, JWEInvalid } from '../util/errors.js'
 import { bitLength as cekLength } from '../lib/cek.js'
 import { importJWK } from '../key/import.js'
@@ -17,6 +17,7 @@ async function decryptKeyManagement(
   key: KeyLike | Uint8Array,
   encryptedKey: Uint8Array | undefined,
   joseHeader: JWEHeaderParameters,
+  options?: DecryptOptions,
 ): Promise<KeyLike | Uint8Array> {
   checkKeyType(alg, key, 'decrypt')
 
@@ -96,6 +97,11 @@ async function decryptKeyManagement(
       if (typeof joseHeader.p2c !== 'number')
         throw new JWEInvalid(`JOSE Header "p2c" (PBES2 Count) missing or invalid`)
 
+      const p2cLimit = options?.maxPBES2Count || 10_000
+
+      if (joseHeader.p2c > p2cLimit)
+        throw new JWEInvalid(`JOSE Header "p2c" (PBES2 Count) out is of acceptable bounds`)
+
       if (typeof joseHeader.p2s !== 'string')
         throw new JWEInvalid(`JOSE Header "p2s" (PBES2 Salt) missing or invalid`)
 
diff --git a/src/types.d.ts b/src/types.d.ts
index cc7f2c2dcf..2cc937eb4c 100644
--- a/src/types.d.ts
+++ b/src/types.d.ts
@@ -401,6 +401,13 @@ export interface DecryptOptions extends CritOption {
    * with compressed plaintext.
    */
   inflateRaw?: InflateFunction
+
+  /**
+   * (PBES2 Key Management Algorithms only) Maximum allowed "p2c" (PBES2 Count) Header Parameter
+   * value. The PBKDF2 iteration count defines the algorithm's computational expense. By default
+   * this value is set to 10000.
+   */
+  maxPBES2Count?: number
 }
 
 /** JWE Deflate option. */
diff --git a/test/jwe/flattened.decrypt.test.mjs b/test/jwe/flattened.decrypt.test.mjs
index a00350d580..e3f09604c4 100644
--- a/test/jwe/flattened.decrypt.test.mjs
+++ b/test/jwe/flattened.decrypt.test.mjs
@@ -1,6 +1,6 @@
 import test from 'ava'
 import * as crypto from 'crypto'
-import { root } from '../dist.mjs'
+import { root, conditional } from '../dist.mjs'
 
 const { FlattenedEncrypt, flattenedDecrypt, base64url } = await import(root)
 
@@ -177,8 +177,8 @@ test('JWE format validation', async (t) => {
     jwe.encrypted_key = 'foo'
 
     await t.throwsAsync(flattenedDecrypt(jwe, t.context.secret), {
-      message: 'decryption operation failed',
-      code: 'ERR_JWE_DECRYPTION_FAILED',
+      message: 'Encountered unexpected JWE Encrypted Key',
+      code: 'ERR_JWE_INVALID',
     })
   }
 })
@@ -239,3 +239,15 @@ test('decrypt empty data (CBC)', async (t) => {
   const { plaintext } = await flattenedDecrypt(jwe, new Uint8Array(32))
   t.is(plaintext.byteLength, 0)
 })
+
+conditional({ electron: 0 })('decrypt PBES2 p2c limit', async (t) => {
+  const jwe = await new FlattenedEncrypt(new Uint8Array(0))
+    .setProtectedHeader({ alg: 'PBES2-HS256+A128KW', enc: 'A128CBC-HS256' })
+    .setKeyManagementParameters({ p2c: 2049 })
+    .encrypt(new Uint8Array(32))
+
+  await t.throwsAsync(flattenedDecrypt(jwe, new Uint8Array(32), { maxPBES2Count: 2048 }), {
+    message: 'JOSE Header "p2c" (PBES2 Count) out is of acceptable bounds',
+    code: 'ERR_JWE_INVALID',
+  })
+})