feat: allow observing remote JWKS resolver state and its manual reload

panva · May 10, 2024 · fa8b639 · fa8b639
1 parent 2595451
commit fa8b639
Show file tree

Hide file tree

Showing 6 changed files with 148 additions and 4 deletions.
diff --git a/src/jwks/local.ts b/src/jwks/local.ts
@@ -59,7 +59,7 @@ function clone<T>(obj: T): T {
 
 /** @private */
 export class LocalJWKSet<KeyLikeType extends KeyLike = KeyLike> {
-  protected _jwks?: JSONWebKeySet
+  private _jwks?: JSONWebKeySet
 
   private _cached: WeakMap<JWK, Cache<KeyLikeType>> = new WeakMap()
 
@@ -252,8 +252,21 @@ async function importWithAlgCache<KeyLikeType extends KeyLike = KeyLike>(
  */
 export function createLocalJWKSet<KeyLikeType extends KeyLike = KeyLike>(jwks: JSONWebKeySet) {
   const set = new LocalJWKSet<KeyLikeType>(jwks)
-  return async (
+
+  const localJWKSet = async (
     protectedHeader?: JWSHeaderParameters,
     token?: FlattenedJWSInput,
   ): Promise<KeyLikeType> => set.getKey(protectedHeader, token)
+
+  Object.defineProperties(localJWKSet, {
+    jwks: {
+      // @ts-expect-error
+      value: () => clone(set._jwks),
+      enumerable: true,
+      configurable: false,
+      writable: false,
+    },
+  })
+
+  return localJWKSet
 }
diff --git a/src/jwks/remote.ts b/src/jwks/remote.ts
@@ -225,10 +225,58 @@ class RemoteJWKSet<KeyLikeType extends KeyLike = KeyLike> {
 export function createRemoteJWKSet<KeyLikeType extends KeyLike = KeyLike>(
   url: URL,
   options?: RemoteJWKSetOptions,
-) {
+): {
+  (protectedHeader?: JWSHeaderParameters, token?: FlattenedJWSInput): Promise<KeyLikeType>
+  /** @ignore */
+  coolingDown: boolean
+  /** @ignore */
+  fresh: boolean
+  /** @ignore */
+  reloading: boolean
+  /** @ignore */
+  reload: () => Promise<void>
+  /** @ignore */
+  jwks: () => JSONWebKeySet | undefined
+} {
   const set = new RemoteJWKSet<KeyLikeType>(url, options)
-  return async (
+
+  const remoteJWKSet = async (
     protectedHeader?: JWSHeaderParameters,
     token?: FlattenedJWSInput,
   ): Promise<KeyLikeType> => set.getKey(protectedHeader, token)
+
+  Object.defineProperties(remoteJWKSet, {
+    coolingDown: {
+      get: () => set.coolingDown(),
+      enumerable: true,
+      configurable: false,
+    },
+    fresh: {
+      get: () => set.fresh(),
+      enumerable: true,
+      configurable: false,
+    },
+    reload: {
+      value: () => set.reload(),
+      enumerable: true,
+      configurable: false,
+      writable: false,
+    },
+    reloading: {
+      // @ts-expect-error
+      get: () => !!set._pendingFetch,
+      enumerable: true,
+      configurable: false,
+    },
+    jwks: {
+      // @ts-expect-error
+      value: () => set._local?.jwks(),
+      enumerable: true,
+      configurable: false,
+      writable: false,
+    },
+  })
+
+  // @ts-expect-error
+  return remoteJWKSet
 }
diff --git a/tap/jwks.ts b/tap/jwks.ts
@@ -10,13 +10,22 @@ export default (QUnit: QUnit, lib: typeof jose) => {
   test('[createRemoteJWKSet] fetches the JWKSet', async (t: typeof QUnit.assert) => {
     const response = await fetch(jwksUri).then((r) => r.json())
     const { alg, kid } = response.keys[0]
+
     const jwks = lib.createRemoteJWKSet(new URL(jwksUri))
+    t.false(jwks.coolingDown)
+    t.false(jwks.fresh)
+    t.equal(jwks.jwks(), undefined)
+
     await t.rejects(jwks({ alg: 'RS256' }), 'multiple matching keys found in the JSON Web Key Set')
     await t.rejects(
       jwks({ kid: 'foo', alg: 'RS256' }),
       'no applicable key found in the JSON Web Key Set',
     )
     t.ok(await Promise.all([jwks({ alg, kid }), jwks({ alg, kid })]))
+
+    t.true(jwks.coolingDown)
+    t.true(jwks.fresh)
+    t.ok(jwks.jwks())
   })
 
   test('[createLocalJWKSet] establishes local JWKSet', async (t: typeof QUnit.assert) => {

diff --git a/test/jwks/local.test.mjs b/test/jwks/local.test.mjs
@@ -17,4 +17,13 @@ test('LocalJWKSet', async (t) => {
   ]) {
     t.throws(() => createLocalJWKSet(f), { code: 'ERR_JWKS_INVALID' })
   }
+
+  const jwks = { keys: [] }
+  const set = createLocalJWKSet(jwks)
+
+  const clone = set.jwks()
+  t.false(clone === jwks)
+  t.false(clone === set.jwks())
+  t.deepEqual(clone, jwks)
+  t.deepEqual(clone, set.jwks())
 })
diff --git a/test/jwks/remote.test.mjs b/test/jwks/remote.test.mjs
@@ -267,6 +267,62 @@ test.serial('refreshes the JWKS once off cooldown', async (t) => {
   }
 })
 
+test.serial('createRemoteJWKSet manual reload', async (t) => {
+  timekeeper.freeze(now * 1000)
+  const jwk = {
+    crv: 'P-256',
+    x: 'fqCXPnWs3sSfwztvwYU9SthmRdoT4WCXxS8eD8icF6U',
+    y: 'nP6GIc42c61hoKqPcZqkvzhzIJkBV3Jw3g8sGG7UeP8',
+    d: 'XikZvoy8ayRpOnuz7ont2DkgMxp_kmmg1EKcuIJWX_E',
+    kty: 'EC',
+  }
+  const jwks = {
+    keys: [
+      {
+        crv: 'P-256',
+        x: 'fqCXPnWs3sSfwztvwYU9SthmRdoT4WCXxS8eD8icF6U',
+        y: 'nP6GIc42c61hoKqPcZqkvzhzIJkBV3Jw3g8sGG7UeP8',
+        kty: 'EC',
+        kid: 'one',
+      },
+    ],
+  }
+
+  const scope = nock('https://as.example.com').get('/jwks').once().reply(200, jwks)
+
+  const url = new URL('https://as.example.com/jwks')
+  const JWKS = createRemoteJWKSet(url)
+  t.false(JWKS.coolingDown)
+  t.false(JWKS.fresh)
+  t.false(JWKS.reloading)
+  t.is(JWKS.jwks(), undefined)
+  const key = await importJWK({ ...jwk, alg: 'ES256' })
+  {
+    const jwt = await new SignJWT().setProtectedHeader({ alg: 'ES256', kid: 'two' }).sign(key)
+    await t.throwsAsync(jwtVerify(jwt, JWKS), {
+      code: 'ERR_JWKS_NO_MATCHING_KEY',
+      message: 'no applicable key found in the JSON Web Key Set',
+    })
+    jwks.keys[0].kid = 'two'
+    scope.get('/jwks').once().reply(200, jwks)
+    t.true(JWKS.coolingDown)
+    t.true(JWKS.fresh)
+    t.false(JWKS.reloading)
+    t.notDeepEqual(JWKS.jwks(), jwks)
+    const reload = JWKS.reload()
+    t.true(JWKS.reloading)
+    await reload
+    t.true(JWKS.coolingDown)
+    t.true(JWKS.fresh)
+    t.false(JWKS.reloading)
+    t.deepEqual(JWKS.jwks(), jwks)
+    await t.notThrowsAsync(jwtVerify(jwt, JWKS))
+    JWKS.jwks().keys = []
+    t.deepEqual(JWKS.jwks(), jwks)
+    await t.notThrowsAsync(jwtVerify(jwt, JWKS))
+  }
+})
+
 test.serial('refreshes the JWKS once stale', async (t) => {
   timekeeper.freeze(now * 1000)
   const jwk = {

diff --git a/test/types/index.test-d.ts b/test/types/index.test-d.ts
@@ -205,6 +205,15 @@ expectType<lib.KeyLike>(await lib.createRemoteJWKSet(new URL(''))())
 expectType<CryptoKey>(await lib.createRemoteJWKSet<CryptoKey>(new URL(''))())
 expectType<KeyObject>(await lib.createRemoteJWKSet<KeyObject>(new URL(''))())
 
+{
+  const jwks = lib.createRemoteJWKSet(new URL(''))
+  expectType<boolean>(jwks.fresh)
+  expectType<boolean>(jwks.coolingDown)
+  expectType<boolean>(jwks.reloading)
+  expectType<Promise<void>>(jwks.reload())
+  expectType<undefined | lib.JSONWebKeySet>(jwks.jwks())
+}
+
 expectType<lib.KeyLike>(await lib.EmbeddedJWK())
 expectType<CryptoKey>(await lib.EmbeddedJWK())
 expectType<KeyObject>(await lib.EmbeddedJWK())