Add auth psk example

src/jose-hpke/jwe/compact.ts

@@ -3,35 +3,62 @@ import { base64url } from "jose";
 import { privateKeyFromJwk, publicKeyFromJwk  } from "../../crypto/keys";
 import { isKeyAlgorithmSupported  } from "../jwk";
 
-import { AeadId, CipherSuite, KdfId, KemId } from "hpke-js";
+import { AeadId, CipherSuite, KdfId, KemId, RecipientContextParams, SenderContextParams } from "hpke-js";
 
-import { HPKE_JWT_ENCRYPT_OPTIONS  } from '../types'
+import { HPKE_JWT_DECRYPT_OPTIONS, HPKE_JWT_ENCRYPT_OPTIONS  } from '../types'
+
+const decoder = new TextDecoder()
 
 export const encrypt = async (plaintext: Uint8Array, publicKeyJwk: any, options?: HPKE_JWT_ENCRYPT_OPTIONS): Promise<string> => {
   if (!isKeyAlgorithmSupported(publicKeyJwk)) {
     throw new Error('Public key algorithm is not supported')
   }
+
+  const senderParams = {
+    recipientPublicKey: await publicKeyFromJwk(publicKeyJwk),
+  } as SenderContextParams
+
   const suite = new CipherSuite({
     kem: KemId.DhkemP256HkdfSha256,
     kdf: KdfId.HkdfSha256,
     aead: AeadId.Aes128Gcm,
   })
-  const sender = await suite.createSenderContext({
-    recipientPublicKey: await publicKeyFromJwk(publicKeyJwk),
-  });
-  const encodedEncapsulatedKey = base64url.encode(new Uint8Array(sender.enc))
+
   const headerParams = {
     alg: publicKeyJwk.alg,
     enc: publicKeyJwk.alg.split('-').pop() // HPKE algorithms always end in an AEAD.
   } as Record<string, any>
+
   if (options?.keyManagementParameters){
-    if (options?.keyManagementParameters.apu){
-      headerParams.apu = base64url.encode(options?.keyManagementParameters.apu)
+    const { keyManagementParameters } = options
+    if (keyManagementParameters.apu){
+      headerParams.apu = base64url.encode(keyManagementParameters.apu)
+    }
+    if (keyManagementParameters.apv){
+      headerParams.apv = base64url.encode(keyManagementParameters.apv)
     }
-    if (options?.keyManagementParameters.apv){
-      headerParams.apv = base64url.encode(options?.keyManagementParameters.apv)
+    if (keyManagementParameters.psk){
+      // in JOSE kid is known to be a string
+      headerParams.psk_id = decoder.decode(keyManagementParameters.psk.id) 
+      if (!keyManagementParameters.psk.key){
+        throw new Error('psk key required when id present.')
+      }
+      senderParams.psk = {
+        id: keyManagementParameters.psk.id,
+        key: keyManagementParameters.psk.key
+      }
     }
   }
+
+  // auth mode
+  if (options?.senderPrivateKey){
+    headerParams.auth_kid = options.senderPrivateKey.kid
+    senderParams.senderKey = await privateKeyFromJwk(options.senderPrivateKey)
+  }
+
+  const sender = await suite.createSenderContext(senderParams);
+  const encodedEncapsulatedKey = base64url.encode(new Uint8Array(sender.enc))
+
   const protectedHeader = base64url.encode(JSON.stringify(headerParams))
   const aad = new TextEncoder().encode(protectedHeader)
   // apu / apv are protected by aad, not as part of kdf
@@ -44,20 +71,40 @@ export const encrypt = async (plaintext: Uint8Array, publicKeyJwk: any, options?
   return jwe
 }
 
-export const decrypt = async (compact: string, privateKeyJwk: any): Promise<Uint8Array> => {
-  if (!isKeyAlgorithmSupported(privateKeyJwk)) {
+export const decrypt = async (compact: string, options: HPKE_JWT_DECRYPT_OPTIONS): Promise<Uint8Array> => {
+  if (!isKeyAlgorithmSupported(options.recipientPrivateKey)) {
     throw new Error('Public key algorithm is not supported')
   }
+  const [protectedHeader, encrypted_key, iv, ciphertext, tag] = compact.split('.');
+  const encapsulated_key = base64url.decode(encrypted_key)
+
+  const recipientParams = {
+    recipientKey: await privateKeyFromJwk(options.recipientPrivateKey),
+    enc: encapsulated_key
+  } as RecipientContextParams
+
+  if (options.keyManagementParameters){
+    const { keyManagementParameters } = options
+    if (keyManagementParameters.psk){
+      recipientParams.psk = {
+        id: keyManagementParameters.psk.id as any, 
+        key: keyManagementParameters.psk.key as any
+      }
+    }
+  }
+
+  if (options.senderPublicKey){
+    recipientParams.senderPublicKey = await publicKeyFromJwk(options.senderPublicKey)
+  }
+
   const suite = new CipherSuite({
     kem: KemId.DhkemP256HkdfSha256,
     kdf: KdfId.HkdfSha256,
     aead: AeadId.Aes128Gcm,
   })
-  const [protectedHeader, encrypted_key, iv, ciphertext, tag] = compact.split('.');
-  const recipient = await suite.createRecipientContext({
-    recipientKey: await privateKeyFromJwk(privateKeyJwk),
-    enc: base64url.decode(encrypted_key)
-  })
+
+
+  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)

src/jose-hpke/jwk.ts

@@ -2,7 +2,7 @@
 import { generateKeyPair, exportJWK, calculateJwkThumbprintUri } from "jose"
 
 export const isKeyAlgorithmSupported = (recipient: Record<string, any>) => {
-  return recipient.alg === 'HPKE-Base-P256-SHA256-A128GCM'
+  return ['HPKE-Base-P256-SHA256-A128GCM', 'HPKE-AuthPSK-P256-SHA256-A128GCM'].includes(recipient.alg)
 }
 
 export const formatJWK = (jwk: any) => {
@@ -19,7 +19,7 @@ export const publicFromPrivate = (privateKeyJwk: any) => {
   }
 }
 
-export const generate = async (alg: 'HPKE-Base-P256-SHA256-A128GCM') => {
+export const generate = async (alg: 'HPKE-Base-P256-SHA256-A128GCM' | 'HPKE-AuthPSK-P256-SHA256-A128GCM') => {
   let kp;
   if (alg.includes('P256')){
     kp = await generateKeyPair('ECDH-ES+A256KW', { crv: 'P-256', extractable: true })

src/jose-hpke/jwt/index.ts

@@ -12,8 +12,8 @@ export const encryptJWT = async (claims: Record<string, any>, options?: HPKE_JWT
   return jwe.compact.encrypt(plaintext, options?.recipientPublicKey, options)
 }
 
-export const decryptJWT = async (jwt: string, options?: HPKE_JWT_DECRYPT_OPTIONS) => {
-  const plaintext = await jwe.compact.decrypt(jwt, options?.recipientPrivateKey)
+export const decryptJWT = async (jwt: string, options: HPKE_JWT_DECRYPT_OPTIONS) => {
+  const plaintext = await jwe.compact.decrypt(jwt, options)
   const [protectedHeader] = jwt.split('.')
   return {
     protectedHeader: JSON.parse(decoder.decode(base64url.decode(protectedHeader))),

src/jose-hpke/types.ts

@@ -3,15 +3,24 @@ export type HPKE_JWT_ENCRYPT_OPTIONS = {
   senderPrivateKey?: Record<string, any>, 
   recipientPublicKey?: Record<string, any>, 
   keyManagementParameters?: {
-    id?: Uint8Array,
-    key?: Uint8Array
     apu?: Uint8Array,
     apv?: Uint8Array
+
+    psk?: {
+      id: Uint8Array,
+      key: Uint8Array
+    }
   }
 }
 
 export type HPKE_JWT_DECRYPT_OPTIONS = {
-  senderPublicKey: Record<string, any>,
+  senderPublicKey?: Record<string, any>,
   recipientPrivateKey: Record<string, any>,
+  keyManagementParameters?: {
+    psk?: {
+      id: Uint8Array,
+      key: Uint8Array
+    }
+  }
 }
 

tests/jwt-auth-psk.test.ts

@@ -4,11 +4,22 @@ import { jose as hpke  } from '../src'
 
 const claims = { 'urn:example:claim': true }
 
-it('Encrypted JWT with HPKE-Base-P256-SHA256-A128GCM, and pre shared key', async () => {
-  const privateKey = await hpke.jwk.generate('HPKE-Base-P256-SHA256-A128GCM')
+it('Encrypted JWT with HPKE-AuthPSK-P256-SHA256-A128GCM (auth and psk)', async () => {
+  const privateKey = {
+    "kid": "urn:ietf:params:oauth:jwk-thumbprint:sha-256:S6AXfdU_6Yfzvu0KDDJb0sFuwnIWPk6LMTErYhPb32s",
+    "alg": "HPKE-AuthPSK-P256-SHA256-A128GCM",
+    "kty": "EC",
+    "crv": "P-256",
+    "x": "wt36K06T4T4APWfGtioqDBXCvRN9evqkZjNydib9MaM",
+    "y": "eupgedeE_HAmVJ62kpSt2_EOoXb6e0y2YF1JPlfr1-I",
+    "d": "O3KznUTAxw-ov-9ZokwNaJ289RgP9VxQc7GJthaXzWY"
+  }
   const publicKey = await hpke.jwk.publicFromPrivate(privateKey)
   const iat = moment().unix()
   const exp = moment().add(2, 'hours').unix()
+  const pskid = new TextEncoder().encode("our-pre-shared-key-id")
+  // a PSK MUST have at least 32 bytes.
+  const psk = new TextEncoder().encode("jugemujugemugokounosurikirekaija")
   const jwe = await hpke.jwt.encryptJWT({
     ...claims,
     iss: 'urn:example:issuer',
@@ -20,22 +31,30 @@ it('Encrypted JWT with HPKE-Base-P256-SHA256-A128GCM, and pre shared key', async
     senderPrivateKey:  privateKey, 
     recipientPublicKey:  publicKey, 
     keyManagementParameters: {
-      id: new TextEncoder().encode("our-pre-shared-key-id"),
-      // a PSK MUST have at least 32 bytes.
-      key: new TextEncoder().encode("jugemujugemugokounosurikirekaija"),
+      psk: {
+        id: pskid,
+        key: psk,
+      }
     }
   })
   // protected.encapsulated_key.<no iv>.ciphertext.<no tag>
   const result = await hpke.jwt.decryptJWT(jwe, {
     senderPublicKey: publicKey,
-    recipientPrivateKey: privateKey
+    recipientPrivateKey: privateKey,
+    keyManagementParameters: {
+      psk: {
+        id: pskid,
+        key: psk,
+      }
+    }
   })
   expect(result.payload['urn:example:claim']).toBe(true)
   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-Base-P256-SHA256-A128GCM')
+  expect(result.protectedHeader['alg']).toBe('HPKE-AuthPSK-P256-SHA256-A128GCM')
   expect(result.protectedHeader['enc']).toBe('A128GCM')
-
+  expect(result.protectedHeader['psk_id']).toBe('our-pre-shared-key-id')
+  expect(result.protectedHeader['auth_kid']).toBe('urn:ietf:params:oauth:jwk-thumbprint:sha-256:S6AXfdU_6Yfzvu0KDDJb0sFuwnIWPk6LMTErYhPb32s')
 })