From 2d080ea2aaab9c4280b386719be042e0d8a38e0c Mon Sep 17 00:00:00 2001
From: Sebastian Stenzel <sebastian.stenzel@skymatic.de>
Date: Sat, 2 Mar 2024 12:10:28 +0100
Subject: [PATCH] add `"alg": "A256KW"` support for JWEs

---
 frontend/src/common/jwe.ts       | 40 +++++++++++++++++++++++++++++---
 frontend/test/common/jwe.spec.ts | 23 ++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/frontend/src/common/jwe.ts b/frontend/src/common/jwe.ts
index 6ea42ade..ec0c33cd 100644
--- a/frontend/src/common/jwe.ts
+++ b/frontend/src/common/jwe.ts
@@ -36,8 +36,8 @@ export class ConcatKDF {
 }
 
 export type JWEHeader = {
-  readonly alg: 'ECDH-ES' | 'PBES2-HS512+A256KW',
-  readonly enc: 'A256GCM' | 'A128GCM',
+  readonly alg: 'ECDH-ES' | 'PBES2-HS512+A256KW' | 'A256KW',
+  readonly enc: 'A256GCM' | 'A128GCM', // A128GCM for testing only, as we use test vectors with 128 bit keys
   readonly apu?: string,
   readonly apv?: string,
   readonly epk?: JsonWebKey,
@@ -97,7 +97,7 @@ export class JWEParser {
    * @throws {UnwrapKeyError} if decryption failed (wrong password?)
    */
   public async decryptPbes2(password: string): Promise<any> {
-    if (this.header.alg != 'PBES2-HS512+A256KW' || /* this.header.enc != 'A256GCM' || */ !this.header.p2s || !this.header.p2c) {
+    if (this.header.alg != 'PBES2-HS512+A256KW' || this.header.enc != 'A256GCM' || !this.header.p2s || !this.header.p2c) {
       throw new Error('unsupported alg or enc');
     }
     const saltInput = base64url.parse(this.header.p2s, { loose: true });
@@ -110,6 +110,24 @@ export class JWEParser {
     }
   }
 
+  /**
+   * Decrypts the JWE, assuming alg == A256KW and enc == A256GCM.
+   * @param kek The key used to wrap the CEK
+   * @returns Decrypted payload
+   * @throws {UnwrapKeyError} if decryption failed (wrong kek?)
+   */
+  public async decryptA256kw(kek: CryptoKey): Promise<any> {
+    if (this.header.alg != 'A256KW' || this.header.enc != 'A256GCM') {
+      throw new Error('unsupported alg or enc');
+    }
+    try {
+      const cek = crypto.subtle.unwrapKey('raw', this.encryptedKey, kek, 'AES-KW', { name: 'AES-GCM', length: 256 }, false, ['decrypt']);
+      return this.decrypt(await cek);
+    } catch (error) {
+      throw new UnwrapKeyError(error);
+    }
+  }
+
   private async decrypt(cek: CryptoKey): Promise<any> {
     const utf8enc = new TextEncoder();
     const m = new Uint8Array(this.ciphertext.length + this.tag.length);
@@ -180,6 +198,22 @@ export class JWEBuilder {
     return new JWEBuilder(header, encryptedKey, cek);
   }
 
+  /**
+   * Prepares a new JWE using alg: A256KW and enc: A256GCM.
+   * 
+   * @param kek The key used to wrap the CEK
+   * @returns A new JWEBuilder ready to encrypt the payload
+   */
+  public static a256kw(kek: CryptoKey): JWEBuilder {
+    const header = (async () => <JWEHeader>{
+      alg: 'A256KW',
+      enc: 'A256GCM'
+    })();
+    const cek = crypto.subtle.generateKey({ name: 'AES-GCM', length: 256 }, true, ['encrypt']);
+    const encryptedKey = (async () => new Uint8Array(await crypto.subtle.wrapKey('raw', await cek, kek, 'AES-KW')))();
+    return new JWEBuilder(header, encryptedKey, cek);
+  }
+
   /**
    * Builds the JWE.
    * @param payload Payload to be encrypted
diff --git a/frontend/test/common/jwe.spec.ts b/frontend/test/common/jwe.spec.ts
index 62c65586..0d0b73e1 100644
--- a/frontend/test/common/jwe.spec.ts
+++ b/frontend/test/common/jwe.spec.ts
@@ -74,6 +74,29 @@ describe('JWE', () => {
     // TODO: add some more decrypt-only tests with JWE from 3rd party
   });
 
+  describe('JWE using alg: A256KW', () => {
+    it('x = decrypt(encrypt(x, kek), kek)', async () => {
+      const kek = await crypto.subtle.generateKey({ name: 'AES-KW', length: 256 }, false, ['wrapKey', 'unwrapKey']);
+      const orig = { hello: 'world' };
+
+      const jwe = await JWEBuilder.a256kw(kek).encrypt(orig);
+
+      const decrypted = await JWEParser.parse(jwe).decryptA256kw(kek);
+      expect(decrypted).to.deep.eq(orig);
+    });
+
+    it('decrypt', async () => {
+      // JWE generated by https://dinochiesa.github.io/jwt/
+      const jwe = 'eyJhbGciOiJBMjU2S1ciLCJlbmMiOiJBMjU2R0NNIn0.JTSrGbw4XEKXYFTC7siTT7DIZUX2SogThcLKXgxe0FPK3Fi8ckjr9A.zQx0t4qoTVIc-h5f.cmqzZ-md3cvdTNH9FWbKOsw.DCdGhmdwjoYKIuNC5zgQJQ';
+      const rawKek = base64url.parse('y_uxz8iAtcOXlqMYpm2jASvDWokpCYMtwkthFSK6IF0', { loose: true });
+      const kek = await crypto.subtle.importKey('raw', rawKek, 'AES-KW', false, ['unwrapKey']);
+      const orig = { hello: 'world' };
+
+      const decrypted = await JWEParser.parse(jwe).decryptA256kw(kek);
+      expect(decrypted).to.deep.eq(orig);
+    });
+  });
+
   describe('PBES2', () => {
     /**
      * Test vectors from https://www.rfc-editor.org/rfc/rfc7517#appendix-C.4