refactor: weak cache normalized CryptoKey instances

panva · Jun 27, 2024 · 32b25a5 · 32b25a5
1 parent b7751f5
commit 32b25a5
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 16 deletions.
diff --git a/src/runtime/browser/jwk_to_key.ts b/src/runtime/browser/jwk_to_key.ts
@@ -1,5 +1,4 @@
 import crypto from './webcrypto.js'
-import type { JWKImportFunction } from '../interfaces.d'
 import { JOSENotSupported } from '../../util/errors.js'
 import type { JWK } from '../../types.d'
 
@@ -91,7 +90,7 @@ function subtleMapping(jwk: JWK): {
   return { algorithm, keyUsages }
 }
 
-const parse: JWKImportFunction = async (jwk: JWK): Promise<CryptoKey> => {
+const parse = async (jwk: JWK): Promise<CryptoKey> => {
   if (!jwk.alg) {
     throw new TypeError('"alg" argument is required when "jwk.alg" is not present')
   }

diff --git a/src/runtime/browser/normalize_key.ts b/src/runtime/browser/normalize_key.ts
@@ -1,14 +1,41 @@
-import type { KeyLike } from '../../types.d'
+import type { JWK, KeyLike } from '../../types.d'
 import { decode } from '../base64url.js'
 import importJWK from '../jwk_to_key.js'
 
 const normalizeSecretKey = (k: string) => decode(k)
 
-const normalizePublicKey = (key: KeyLike | Uint8Array | unknown, alg: string) => {
+let privCache: WeakMap<object, Record<string, CryptoKey>>
+let pubCache: WeakMap<object, Record<string, CryptoKey>>
+
+const isKeyObject = (key: unknown): key is KeyLike => {
   // @ts-expect-error
-  if (key?.[Symbol.toStringTag] === 'KeyObject') {
+  return key?.[Symbol.toStringTag] === 'KeyObject'
+}
+
+const importAndCache = async (
+  cache: typeof privCache | typeof pubCache,
+  key: KeyLike,
+  jwk: JWK,
+  alg: string,
+) => {
+  let cached = cache.get(key)
+  if (cached?.[alg]) {
+    return cached[alg]
+  }
+
+  const cryptoKey = await importJWK({ ...jwk, alg })
+  if (!cached) {
+    cache.set(key, { [alg]: cryptoKey })
+  } else {
+    cached[alg] = cryptoKey
+  }
+  return cryptoKey
+}
+
+const normalizePublicKey = (key: KeyLike | Uint8Array | unknown, alg: string) => {
+  if (isKeyObject(key)) {
     // @ts-expect-error
-    let jwk = key.export({ format: 'jwk' })
+    let jwk: JWK = key.export({ format: 'jwk' })
     delete jwk.d
     delete jwk.dp
     delete jwk.dq
@@ -19,22 +46,23 @@ const normalizePublicKey = (key: KeyLike | Uint8Array | unknown, alg: string) =>
       return normalizeSecretKey(jwk.k)
     }
 
-    return importJWK({ ...jwk, alg })
+    pubCache ||= new WeakMap()
+    return importAndCache(pubCache, key, jwk, alg)
   }
 
   return <KeyLike | Uint8Array>key
 }
 
 const normalizePrivateKey = (key: KeyLike | Uint8Array | unknown, alg: string) => {
-  // @ts-expect-error
-  if (key?.[Symbol.toStringTag] === 'KeyObject') {
+  if (isKeyObject(key)) {
     // @ts-expect-error
-    let jwk = key.export({ format: 'jwk' })
+    let jwk: JWK = key.export({ format: 'jwk' })
     if (jwk.k) {
       return normalizeSecretKey(jwk.k)
     }
 
-    return importJWK({ ...jwk, alg })
+    privCache ||= new WeakMap()
+    return importAndCache(privCache, key, jwk, alg)
   }
 
   return <KeyLike | Uint8Array>key

diff --git a/src/runtime/interfaces.d.ts b/src/runtime/interfaces.d.ts
@@ -75,9 +75,6 @@ export interface FetchFunction {
 export interface DigestFunction {
   (digest: 'sha256' | 'sha384' | 'sha512', data: Uint8Array): AsyncOrSync<Uint8Array>
 }
-export interface JWKImportFunction {
-  (jwk: JWK): AsyncOrSync<KeyLike>
-}
 export interface PEMImportFunction {
   (pem: string, alg: string, options?: PEMImportOptions): AsyncOrSync<KeyLike>
 }

diff --git a/src/runtime/node/jwk_to_key.ts b/src/runtime/node/jwk_to_key.ts
@@ -1,10 +1,9 @@
 import { createPrivateKey, createPublicKey } from 'node:crypto'
 import type { KeyObject } from 'node:crypto'
 
-import type { JWKImportFunction } from '../interfaces.d'
 import type { JWK } from '../../types.d'
 
-const parse: JWKImportFunction = (jwk: JWK): KeyObject => {
+const parse = (jwk: JWK): KeyObject => {
   return (jwk.d ? createPrivateKey : createPublicKey)({ format: 'jwk', key: jwk })
 }
 export default parse