Add example implementation of ML-KEM-68 for COSE HPKE

package.json

@@ -47,6 +47,7 @@
     "typescript": "^4.9.4"
   },
   "dependencies": {
+    "@noble/post-quantum": "^0.1.0",
     "@peculiar/x509": "^1.9.7",
     "@transmute/cose": "^0.1.0",
     "@transmute/rfc9162": "^0.0.5",

src/cose/Params.ts

@@ -1,6 +1,9 @@
 // This module is just just a limited set of the IANA registries, 
 // exposed to make Map initialization more readable
 
+import { IANACOSEKeyCommonParameters } from "./key-common-parameters"
+import * as requested from './requested-assignment'
+
 export type HeaderMapEntry = [number, any]
 export type HeaderMap = Map<number, any>
 
@@ -77,14 +80,12 @@ export const KeyWrap = {
   A128KW: -3
 }
 
-export const Direct = {
-  'HPKE-Base-P256-SHA256-AES128GCM': 35
-}
 
 export const EC2 = 2
 
 export const KeyType = {
-  EC2
+  EC2,
+  ['ML-KEM']: requested.KeyTypes['ML-KEM']
 }
 
 export const Epk = {
@@ -97,7 +98,39 @@ export const Curve = {
   P256: 1,
 }
 
+export const Key = {
+  Type: parseInt(IANACOSEKeyCommonParameters['1'].Label, 10),
+  Algorithm: parseInt(IANACOSEKeyCommonParameters['3'].Label, 10)
+}
+
+export const KeyTypeAlgorithms = {
+  ['ML-KEM']: {
+    ['HPKE-Base-ML-KEM-768-SHA256-AES128GCM']: requested.Algorithms['HPKE-Base-ML-KEM-768-SHA256-AES128GCM'],
+    ['ML-KEM-768']: requested.Algorithms['ML-KEM-768']
+  }
+}
+
+export const KeyTypeParameters = {
+  ['ML-KEM']: {
+    Public: -1,
+    Secret: -2,
+  },
+  ['EC2']: {
+    Curve: -1,
+    PublicX: -2,
+    PublicY: -3,
+    Secret: -4,
+  }
+}
+
+
+export const Direct = {
+  'HPKE-Base-P256-SHA256-AES128GCM': 35,
+  'HPKE-Base-ML-KEM-768-SHA256-AES128GCM': KeyTypeAlgorithms['ML-KEM']['HPKE-Base-ML-KEM-768-SHA256-AES128GCM']
+}
 
 export const COSE_Encrypt0 = 16
 export const COSE_Sign1 = 18
-export const COSE_Encrypt = 96
+export const COSE_Encrypt = 96
+
+

src/cose/encrypt/direct.ts

@@ -8,7 +8,7 @@ import { EMPTY_BUFFER } from "../../cbor"
 import * as aes from './aes'
 import * as ecdh from './ecdh'
 
-import { COSE_Encrypt, Epk, KeyAgreement, Protected, ProtectedHeader, Unprotected } from "../Params"
+import { COSE_Encrypt, Direct, Epk, KeyAgreement, Protected, ProtectedHeader, Unprotected } from "../Params"
 
 import { createAAD } from './utils'
 
@@ -33,10 +33,10 @@ export const encrypt = async (req: RequestDirectEncryption) => {
     throw new Error('Direct encryption requires a single recipient')
   }
   const recipientPublicKeyJwk = req.recipients.keys[0]
-  if (recipientPublicKeyJwk.crv !== 'P-256') {
-    throw new Error('Only P-256 is supported currently')
+  if (recipientPublicKeyJwk.crv !== 'P-256' && recipientPublicKeyJwk.kty !== 'ML-KEM') {
+    throw new Error('Only P-256 DHKEM and ML-KEM-768 are currently supported')
   }
-  if (recipientPublicKeyJwk.alg === hpke.primaryAlgorithm.label) {
+  if (Object.keys(Direct).includes(recipientPublicKeyJwk.alg)) {
     return hpke.encrypt.direct(req)
   }
   const alg = req.protectedHeader.get(Protected.Alg)
@@ -71,7 +71,7 @@ export const encrypt = async (req: RequestDirectEncryption) => {
 
 export const decrypt = async (req: RequestDirectDecryption) => {
   const receiverPrivateKeyJwk = req.recipients.keys[0]
-  if (receiverPrivateKeyJwk.alg === hpke.primaryAlgorithm.label) {
+  if (Object.keys(Direct).includes(receiverPrivateKeyJwk.alg)) {
     return hpke.decrypt.direct(req)
   }
   const decoded = await decodeFirst(req.ciphertext)

src/cose/encrypt/hpke/direct.ts

@@ -1,22 +1,35 @@
 
 
-import { COSE_Encrypt0, Direct, Protected, Unprotected, UnprotectedHeader } from '../../Params'
+
+import { ml_kem768 } from '@noble/post-quantum/ml-kem';
+
+import { COSE_Encrypt0, Direct, KeyTypeAlgorithms, Protected, Unprotected, UnprotectedHeader } from '../../Params'
 import { RequestDirectEncryption, RequestDirectDecryption } from '../types'
 import { Tagged, decodeFirst, encodeAsync } from "cbor-web"
 
 import { computeInfo } from './computeInfo'
 import { suites, JOSE_HPKE_ALG } from './suites'
 
 import { publicKeyFromJwk, privateKeyFromJwk, computeHPKEAad } from './common'
+import { base64url } from 'jose'
 
-export const encryptDirect = async (req: RequestDirectEncryption) => {
+import * as aes from '../aes'
+
+import { CipherSuite, KemId, KdfId, AeadId } from 'hpke-js'
+import { EMPTY_BUFFER, toArrayBuffer } from '../../../cbor';
+
+import { createAAD } from '../utils';
+
+const dhkemsuite = new CipherSuite({
+  kem: KemId.DhkemP256HkdfSha256,
+  kdf: KdfId.HkdfSha256,
+  aead: AeadId.Aes128Gcm,
+});
+
+const handleDHKemEncrypt = async (req: RequestDirectEncryption) => {
   if (req.unprotectedHeader === undefined) {
     req.unprotectedHeader = UnprotectedHeader([])
   }
-  const alg = req.protectedHeader.get(Protected.Alg)
-  if (alg !== Direct['HPKE-Base-P256-SHA256-AES128GCM']) {
-    throw new Error('Only alg 35 is supported')
-  }
   const protectedHeader = await encodeAsync(req.protectedHeader)
   const unprotectedHeader = req.unprotectedHeader;
   const [recipientPublicKeyJwk] = req.recipients.keys
@@ -44,7 +57,74 @@ export const encryptDirect = async (req: RequestDirectEncryption) => {
   ]), { canonical: true })
 }
 
-export const decryptDirect = async (req: RequestDirectDecryption) => {
+
+const sharedSecretToContentEncryptionKey = async (sharedSecret: Uint8Array) => {
+  const ikm = sharedSecret;
+  // https://datatracker.ietf.org/doc/html/rfc9180#section-4-10
+  // labeled_ikm = concat("HPKE-v1", suite_id, label, ikm)
+  // 🔥 this is ALL WRONG.... 🔥
+  // fake
+  const suite_id = Buffer.from('0xFFFF', 'hex') // unassigned kem id https://www.iana.org/assignments/hpke/hpke.xhtml
+  // should be:
+  // suite_id = concat(
+  //   "HPKE",
+  //   I2OSP(kem_id, 2),
+  //   I2OSP(kdf_id, 2),
+  //   I2OSP(aead_id, 2)
+  // )
+  const labeled_ikm = Buffer.concat([
+    new TextEncoder().encode('HPKE-v1'),
+    suite_id,
+    Buffer.from(''),
+    ikm
+  ])
+  const salt = new TextEncoder().encode('') // empty string?
+  return dhkemsuite.kdf.extract(salt, labeled_ikm)
+}
+
+const handleMLKemEncrypt = async (req: RequestDirectEncryption) => {
+  const protectedHeader = await encodeAsync(req.protectedHeader)
+  const unprotectedHeader = req.unprotectedHeader || new Map<any, any>();
+  const [recipientPublicKeyJwk] = req.recipients.keys
+  const publicKey = base64url.decode(recipientPublicKeyJwk.x)
+  const { cipherText, sharedSecret } = ml_kem768.encapsulate(publicKey);
+  const kemCt = cipherText;
+  const aeadContentEncryptionKey = await sharedSecretToContentEncryptionKey(sharedSecret)
+  const aeadAlg = 1; // AES 128 GCM
+  const iv = await aes.getIv(aeadAlg) // random for each direct encryption
+  const externalAad = EMPTY_BUFFER
+  // const hpkeSealAad = computeHPKEAad(protectedHeader) // confused why I don't need this...
+  const aad = await createAAD(protectedHeader, 'Encrypt', externalAad)
+  const ct = await aes.encrypt(aeadAlg, new Uint8Array(req.plaintext), new Uint8Array(iv), new Uint8Array(aad), new Uint8Array(aeadContentEncryptionKey))
+  // HPKE direct mode prefix iv?
+  const ctWithIv = Buffer.concat([
+    Buffer.from(iv),
+    Buffer.from(ct)
+  ])
+  if (recipientPublicKeyJwk.kid) {
+    unprotectedHeader.set(Unprotected.Kid, recipientPublicKeyJwk.kid)
+  }
+  unprotectedHeader.set(Unprotected.Ek, toArrayBuffer(kemCt))
+  return encodeAsync(new Tagged(COSE_Encrypt0, [
+    protectedHeader,
+    unprotectedHeader,
+    ctWithIv,
+  ]), { canonical: true })
+}
+
+export const encryptDirect = async (req: RequestDirectEncryption) => {
+  const alg = req.protectedHeader.get(Protected.Alg)
+  if (alg === Direct['HPKE-Base-P256-SHA256-AES128GCM']) {
+    return handleDHKemEncrypt(req)
+  }
+  if (alg === Direct['HPKE-Base-ML-KEM-768-SHA256-AES128GCM']) {
+    return handleMLKemEncrypt(req)
+  }
+  throw new Error('Unsupported HPKE algorithm')
+}
+
+
+const handleDHKemDecrypt = async (req: RequestDirectDecryption) => {
   const decoded = await decodeFirst(req.ciphertext)
   if (decoded.tag !== COSE_Encrypt0) {
     throw new Error('Only tag 16 cose encrypt are supported')
@@ -73,3 +153,45 @@ export const decryptDirect = async (req: RequestDirectDecryption) => {
   const plaintext = await hpkeRecipient.open(ciphertext, hpkeSealAad)
   return plaintext
 }
+
+
+const handleMLKemDecrypt = async (req: RequestDirectDecryption) => {
+  const decoded = await decodeFirst(req.ciphertext)
+  if (decoded.tag !== COSE_Encrypt0) {
+    throw new Error('Only tag 16 cose encrypt are supported')
+  }
+  const [protectedHeader, unprotectedHeader, ctWithIv] = decoded.value
+  const kid = unprotectedHeader.get(Unprotected.Kid).toString();
+  const receiverPrivateKeyJwk = req.recipients.keys.find((k) => {
+    return k.kid === kid
+  })
+  const ek = unprotectedHeader.get(Unprotected.Ek)
+  const iv = ctWithIv.slice(0, 16) // AES-128-GCM iv length
+  const encryptedContent = ctWithIv.slice(16, ctWithIv.length)
+  const secretKey = base64url.decode(receiverPrivateKeyJwk.d)
+  const sharedSecret = ml_kem768.decapsulate(ek, secretKey);
+  const aeadContentEncryptionKey = await sharedSecretToContentEncryptionKey(sharedSecret)
+  const aeadAlg = 1; // AES 128 GCM
+  const externalAad = EMPTY_BUFFER
+  // const hpkeSealAad = computeHPKEAad(protectedHeader) // confused why I don't need this...
+  const aad = await createAAD(protectedHeader, 'Encrypt', externalAad)
+  return aes.decrypt(aeadAlg, encryptedContent, new Uint8Array(iv), new Uint8Array(aad), new Uint8Array(aeadContentEncryptionKey))
+}
+
+
+export const decryptDirect = async (req: RequestDirectDecryption) => {
+  const decoded = await decodeFirst(req.ciphertext)
+  if (decoded.tag !== COSE_Encrypt0) {
+    throw new Error('Only tag 16 cose encrypt are supported')
+  }
+  const [protectedHeader] = decoded.value
+  const decodedProtectedHeader = await decodeFirst(protectedHeader)
+  const alg = decodedProtectedHeader.get(Protected.Alg)
+  if (alg === Direct['HPKE-Base-P256-SHA256-AES128GCM']) {
+    return handleDHKemDecrypt(req)
+  }
+  if (alg === Direct['HPKE-Base-ML-KEM-768-SHA256-AES128GCM']) {
+    return handleMLKemDecrypt(req)
+  }
+  throw new Error('Unsupported HPKE algorithm')
+}

src/cose/key/convertCoseKeyToJsonWebKey.ts

@@ -1,5 +1,5 @@
 import { base64url, calculateJwkThumbprint } from "jose";
-import { CoseKey } from ".";
+import { CoseKey, thumbprint } from ".";
 
 
 import { IANACOSEAlgorithms } from '../algorithms';
@@ -9,12 +9,24 @@ const algorithms = Object.values(IANACOSEAlgorithms)
 const curves = Object.values(IANACOSEEllipticCurves)
 
 import { formatJwk } from "./formatJwk";
+import { KeyType, KeyTypeParameters } from "../Params";
 
 export const convertCoseKeyToJsonWebKey = async <T>(coseKey: CoseKey): Promise<T> => {
   const kty = coseKey.get(1) as number
   const kid = coseKey.get(2)
   const alg = coseKey.get(3)
   const crv = coseKey.get(-1)
+
+  if (kty === KeyType["ML-KEM"]) {
+    //short circuit.
+    return formatJwk({
+      "kid": await thumbprint.calculateCoseKeyThumbprintUri(coseKey),
+      "alg": "HPKE-Base-ML-KEM-768-SHA256-AES128GCM",
+      "kty": "ML-KEM",
+      "x": base64url.encode(coseKey.get(KeyTypeParameters['ML-KEM'].Public) as Uint8Array),
+      "d": coseKey.get(KeyTypeParameters['ML-KEM'].Secret) ? base64url.encode(coseKey.get(KeyTypeParameters['ML-KEM'].Secret) as Uint8Array) : undefined,
+    }) as T
+  }
   // kty EC, kty: EK
   if (![2, 5].includes(kty)) {
     throw new Error('This library requires does not support the given key type')