add example X-Wing JWT

transmute-industries · Nov 18, 2024 · 6aef1c2 · 6aef1c2
1 parent f260cb2
commit 6aef1c2
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 16 deletions.
diff --git a/src/jose-hpke/jwe/compact.ts b/src/jose-hpke/jwe/compact.ts
@@ -1,9 +1,16 @@
 import { base64url } from "jose";
 
+import { XWing } from "@hpke/hybridkem-x-wing";
+
+
+import { Aes128Gcm, CipherSuite, HkdfSha256, DhkemP256HkdfSha256,} from "@hpke/core";
+
+
+
 import { privateKeyFromJwk, publicKeyFromJwk  } from "../../crypto/keys";
 import { isKeyAlgorithmSupported  } from "../jwk";
 
-import { AeadId, CipherSuite, KdfId, KemId, RecipientContextParams, SenderContextParams } from "hpke-js";
+import {  RecipientContextParams } from "hpke-js";
 
 import { HPKE_JWT_DECRYPT_OPTIONS, HPKE_JWT_ENCRYPT_OPTIONS  } from '../types'
 import { modes } from "..";
@@ -28,12 +35,31 @@ export const decrypt = async (compact: string, options: HPKE_JWT_DECRYPT_OPTIONS
   const [protectedHeader, encrypted_key, iv, ciphertext, tag] = compact.split('.');
   const encapsulated_key = base64url.decode(encrypted_key)
 
-  const recipientParams = {
+  let recipientParams = {
     recipientKey: await privateKeyFromJwk(options.recipientPrivateKey),
     enc: encapsulated_key,
     info: options.hpke_info
   } as RecipientContextParams
 
+  let suite = new CipherSuite({
+    kem: new DhkemP256HkdfSha256(),
+    kdf: new HkdfSha256(),
+    aead: new Aes128Gcm(),
+  });
+
+  if (options.recipientPrivateKey.alg === 'HPKE-X-Wing-SHA256-A128GCM'){
+    suite = new CipherSuite({
+      kem: new XWing(),
+      kdf: new HkdfSha256(),
+      aead: new Aes128Gcm(),
+    });
+   recipientParams = {
+      recipientKey: await privateKeyFromJwk(options.recipientPrivateKey) ,
+      enc: encapsulated_key
+    } as RecipientContextParams
+   await suite.kem.importKey('jwk', { ...options.recipientPrivateKey, alg:'X-Wing' }, false)
+  }
+
   if (options.keyManagementParameters){
     const { keyManagementParameters } = options
     if (keyManagementParameters.psk){
@@ -47,21 +73,13 @@ export const decrypt = async (compact: string, options: HPKE_JWT_DECRYPT_OPTIONS
   if (options.senderPublicKey){
     recipientParams.senderPublicKey = await publicKeyFromJwk(options.senderPublicKey)
   }
-
-  const suite = new CipherSuite({
-    kem: KemId.DhkemP256HkdfSha256,
-    kdf: KdfId.HkdfSha256,
-    aead: AeadId.Aes128Gcm,
-  })
-
-
+
   const recipient = await suite.createRecipientContext(recipientParams)
   const aad = new TextEncoder().encode(protectedHeader)
   const plaintext = await recipient.open(base64url.decode(ciphertext), aad)
   return new Uint8Array(plaintext)
 }
 
-
 // https://datatracker.ietf.org/doc/html/rfc7516#section-3.2
 // "protected", with the value BASE64URL(UTF8(JWE Protected Header))
 // "unprotected", with the value JWE Shared Unprotected Header

diff --git a/src/jose-hpke/jwk.ts b/src/jose-hpke/jwk.ts
@@ -1,11 +1,11 @@
 
-import { Aes256Gcm, CipherSuite, HkdfSha256 } from "@hpke/core";
+import { Aes128Gcm, CipherSuite, HkdfSha256 } from "@hpke/core";
 import { XWing } from "@hpke/hybridkem-x-wing";
 
 import { generateKeyPair, exportJWK, calculateJwkThumbprintUri, base64url } from "jose"
 
 export const isKeyAlgorithmSupported = (recipient: Record<string, any>) => {
-  return ['HPKE-P256-SHA256-A128GCM', 'HPKE-P256-SHA256-A128GCM'].includes(recipient.alg)
+  return ['HPKE-P256-SHA256-A128GCM', 'HPKE-P256-SHA256-A128GCM', `HPKE-X-Wing-SHA256-A128GCM`].includes(recipient.alg)
 }
 
 export const formatJWK = (jwk: any) => {
@@ -32,7 +32,7 @@ export const generate = async (alg: 'HPKE-P256-SHA256-A128GCM' | 'HPKE-P256-SHA2
     const suite = new CipherSuite({
       kem: new XWing(),
       kdf: new HkdfSha256(),
-      aead: new Aes256Gcm(),
+      aead: new Aes128Gcm(),
     });
     const rkp:any = await suite.kem.generateKeyPair()
     return {

diff --git a/tests/ml-kem.akp.test.ts b/tests/ml-kem.akp.test.ts
@@ -2,14 +2,14 @@ import moment from 'moment'
 import { jose as hpke } from '../src'
 
 
-import { Aes256Gcm, CipherSuite, HkdfSha256 } from "@hpke/core";
+import { Aes128Gcm , CipherSuite, HkdfSha256 } from "@hpke/core";
 import { XWing } from "@hpke/hybridkem-x-wing";
 
 it('X-Wing Sanity', async () => {
     const suite = new CipherSuite({
       kem: new XWing(),
       kdf: new HkdfSha256(),
-      aead: new Aes256Gcm(),
+      aead: new Aes128Gcm(),
     });
     // NOTE: The following support for JWKs with the AKP key type is experimental.
     // Please be aware that the specifications are subject to change without notice.
@@ -70,3 +70,34 @@ it('X-Wing JWE', async () => {
   expect(new TextDecoder().decode(decrypted.additionalAuthenticatedData)).toBe('🏴‍☠️ beware the aad!')
 })
 
+
+
+it('X-Wing JWT', async () => {
+  const privateKey = await hpke.jwk.generate('HPKE-X-Wing-SHA256-A128GCM')
+  const publicKey = await hpke.jwk.publicFromPrivate(privateKey)
+  const iat = moment().unix()
+  const exp = moment().add(2, 'hours').unix()
+  const jwe = await hpke.jwt.encryptJWT({
+    iss: 'urn:example:issuer',
+    aud: 'urn:example:audience',
+    iat,
+    exp,
+  }, {
+    recipientPublicKey: publicKey
+  })
+  // console.log(jwe)
+  const result = await hpke.jwt.decryptJWT(
+    jwe,
+    {
+      recipientPrivateKey: privateKey
+    })
+  expect(result.payload['iss']).toBe('urn:example:issuer')
+  expect(result.payload['aud']).toBe('urn:example:audience')
+  expect(result.payload.iat).toBeDefined()
+  expect(result.payload.exp).toBeDefined()
+  expect(result.protectedHeader['alg']).toBe('HPKE-X-Wing-SHA256-A128GCM')
+  expect(result.protectedHeader['enc']).toBe('dir')
+  // protected header does not contain epk
+  expect(result.protectedHeader.epk).toBeUndefined()
+  // encapsulated key is transported through "encrypted_key"
+})