Back the JWKS cache with the Workers edge cache

apps/cloud/src/auth/jwks-cache.node.test.ts

@@ -9,7 +9,7 @@ import {
   type KeyLike,
 } from "jose";
 
-import { createCachedRemoteJWKSet } from "./jwks-cache";
+import { createCachedRemoteJWKSet, type JwksEdgeCache } from "./jwks-cache";
 
 const issuer = "https://test-authkit.example.com";
 const audience = "client_test_fixture";
@@ -169,3 +169,146 @@ describe("createCachedRemoteJWKSet", () => {
     expect(harness.callCount()).toBe(2);
   });
 });
+
+interface EdgeCacheHarness {
+  readonly cache: JwksEdgeCache;
+  readonly matchCount: () => number;
+  readonly putCount: () => number;
+  readonly stored: () => Response | null;
+  readonly seed: (keys: ReadonlyArray<JWK>) => void;
+}
+
+const makeEdgeCacheHarness = (): EdgeCacheHarness => {
+  let stored: Response | null = null;
+  let matches = 0;
+  let puts = 0;
+
+  return {
+    cache: {
+      match: async () => {
+        matches++;
+        return stored ? stored.clone() : undefined;
+      },
+      put: async (_url, response) => {
+        puts++;
+        stored = response;
+      },
+    },
+    matchCount: () => matches,
+    putCount: () => puts,
+    stored: () => stored,
+    seed: (keys) => {
+      const body: JSONWebKeySet = { keys: keys.map((k) => ({ ...k })) };
+      stored = new Response(JSON.stringify(body), {
+        headers: { "content-type": "application/json" },
+      });
+    },
+  };
+};
+
+describe("createCachedRemoteJWKSet edge cache", () => {
+  it("FAILING-WITHOUT-EDGE-CACHE: a cold isolate with a warm edge cache never fetches upstream", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const upstream = makeFetchHarness([kp.publicJwk]);
+    const edge = makeEdgeCacheHarness();
+    edge.seed([kp.publicJwk]);
+
+    // Fresh resolver = fresh isolate: empty in-memory cache.
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: upstream.fetch,
+      edgeCache: edge.cache,
+    });
+
+    const token = await sign(kp);
+    const { payload } = await jwtVerify(token, jwks, { issuer, audience });
+    expect(payload.sub).toBe("user_test");
+    expect(edge.matchCount()).toBe(1);
+    expect(upstream.callCount()).toBe(0);
+  });
+
+  it("an edge-cache miss falls through to upstream and populates the edge copy", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const upstream = makeFetchHarness([kp.publicJwk]);
+    const edge = makeEdgeCacheHarness();
+
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: upstream.fetch,
+      edgeCache: edge.cache,
+    });
+
+    const token = await sign(kp);
+    await jwtVerify(token, jwks, { issuer, audience });
+    expect(upstream.callCount()).toBe(1);
+    expect(edge.putCount()).toBe(1);
+
+    // A second cold isolate now reads the populated edge copy.
+    const jwks2 = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: upstream.fetch,
+      edgeCache: edge.cache,
+    });
+    await jwtVerify(token, jwks2, { issuer, audience });
+    expect(upstream.callCount()).toBe(1);
+  });
+
+  it("a malformed edge-cache entry is a miss, not a failure", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const upstream = makeFetchHarness([kp.publicJwk]);
+    const edge = makeEdgeCacheHarness();
+    await edge.cache.put(jwksUrl.toString(), new Response("not json"));
+
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: upstream.fetch,
+      edgeCache: edge.cache,
+    });
+
+    const token = await sign(kp);
+    const { payload } = await jwtVerify(token, jwks, { issuer, audience });
+    expect(payload.sub).toBe("user_test");
+    expect(upstream.callCount()).toBe(1);
+  });
+
+  it("key rotation bypasses the stale edge copy and overwrites it", async () => {
+    const oldKey = await generateRotatableKeypair("k_old");
+    const newKey = await generateRotatableKeypair("k_new");
+    const upstream = makeFetchHarness([newKey.publicJwk]);
+    const edge = makeEdgeCacheHarness();
+    // The edge copy predates the rotation: it only has the old key.
+    edge.seed([oldKey.publicJwk]);
+
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: upstream.fetch,
+      edgeCache: edge.cache,
+    });
+
+    // First verify warms in-memory from the (stale) edge copy; the new-key
+    // token then misses, which must force an UPSTREAM refetch — re-reading
+    // the same stale edge copy would loop the miss forever.
+    const token = await sign(newKey);
+    const { payload } = await jwtVerify(token, jwks, { issuer, audience });
+    expect(payload.sub).toBe("user_test");
+    expect(upstream.callCount()).toBe(1);
+    // And the rotated set replaced the stale edge copy for other isolates.
+    expect(edge.putCount()).toBe(1);
+  });
+
+  it("degrades to upstream-only when the edge cache itself fails", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const upstream = makeFetchHarness([kp.publicJwk]);
+    // oxlint-disable-next-line executor/no-promise-reject, executor/no-error-constructor -- boundary: fakes the Cache API's own failure mode, a rejected promise
+    const rejectingCacheCall = () => Promise.reject(new TypeError("cache backend unavailable"));
+    const broken: JwksEdgeCache = {
+      match: rejectingCacheCall,
+      put: rejectingCacheCall,
+    };
+
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: upstream.fetch,
+      edgeCache: broken,
+    });
+
+    const token = await sign(kp);
+    const { payload } = await jwtVerify(token, jwks, { issuer, audience });
+    expect(payload.sub).toBe("user_test");
+    expect(upstream.callCount()).toBe(1);
+  });
+});

apps/cloud/src/auth/jwks-cache.ts

@@ -11,10 +11,18 @@
 //
 //   * Caches the JSON Web Key Set in module-scope memory for a configurable
 //     TTL (default 1 hour).
+//   * Backs the in-memory layer with the Workers Cache API (`caches.default`)
+//     so a COLD isolate reads the JWKS from its datacenter's shared cache
+//     instead of round-tripping to the upstream endpoint. The in-memory layer
+//     alone left every fresh isolate paying that full fetch, which is exactly
+//     when users notice: an app-open burst fans out across many new isolates
+//     and each one stalled ~1-2s inside session verification.
 //   * Single-flights concurrent fetches so a stampede of verifies during a
 //     cache miss only fires one upstream request.
 //   * Force-refreshes once when verification fails with a cached key, so
-//     genuine key rotation isn't blocked by the TTL.
+//     genuine key rotation isn't blocked by the TTL. A forced refresh also
+//     bypasses the shared edge copy (it is just as stale as ours) and
+//     overwrites it with the freshly fetched set.
 //
 // The returned function is a `JWTVerifyGetKey` and slots directly into
 // `jose.jwtVerify`. It also exposes `forceRefresh()` so the verify path can
@@ -32,6 +40,21 @@ import {
 import { Schema } from "effect";
 import { JWKSNoMatchingKey } from "jose/errors";
 
+/** The slice of the Workers Cache API this module uses, expressed
+ *  structurally: `Cache` from workers-types satisfies it, and node tests can
+ *  supply a conforming fake without the full `Cache` surface. */
+export interface JwksEdgeCache {
+  match(url: string): Promise<Response | undefined>;
+  put(url: string, response: Response): Promise<void>;
+}
+
+// The `caches` global exists in the Workers runtime but not under node (the
+// test environment). Typed structurally (the ASSETS precedent in
+// env-augment.d.ts) rather than via workers-types: its `Cache` type drags in
+// a second `Response` that intersects with the ambient one and rejects
+// responses this module constructs.
+declare const caches: { readonly default: JwksEdgeCache } | undefined;
+
 export interface CachedRemoteJWKSetOptions {
   /**
    * How long a successful fetch is considered fresh. Defaults to 1 hour —
@@ -43,6 +66,12 @@ export interface CachedRemoteJWKSetOptions {
   readonly fetch?: typeof globalThis.fetch;
   /** HTTP request timeout. Defaults to 5s, matching jose. */
   readonly timeoutMs?: number;
+  /**
+   * Datacenter-shared cache behind the in-memory layer. Defaults to
+   * `caches.default` when the runtime provides it (Workers) and to disabled
+   * when it does not (node tests). Pass explicitly to fake it in tests.
+   */
+  readonly edgeCache?: JwksEdgeCache | null;
 }
 
 export interface CachedRemoteJWKSet extends JWTVerifyGetKey {
@@ -122,31 +151,92 @@ export const createCachedRemoteJWKSet = (
   // (tests do this) without us snapshotting a stale reference.
   const fetchImpl = (): typeof globalThis.fetch =>
     options.fetch ?? globalThis.fetch.bind(globalThis);
+  // `undefined` = not specified, use the runtime default; `null` = disabled.
+  // Node (tests) has no `caches` global, so the default degrades to disabled
+  // there without configuration.
+  const edgeCache =
+    options.edgeCache !== undefined
+      ? options.edgeCache
+      : typeof caches !== "undefined"
+        ? caches.default
+        : null;
 
-  let entry: CacheEntry | null = null;
-  let inflight: Promise<CacheEntry> | null = null;
-
-  const refresh = (): Promise<CacheEntry> => {
-    if (inflight) return inflight;
-    inflight = (async () => {
-      const jwks = await fetchJwksOnce(url, fetchImpl(), timeoutMs);
-      const next: CacheEntry = {
-        jwks,
-        fetchedAt: Date.now(),
-        resolver: createLocalJWKSet(jwks),
-      };
-      entry = next;
-      return next;
-    })().finally(() => {
-      inflight = null;
+  // The edge cache is best-effort: a failed read is a miss, a failed write is
+  // dropped. Verification correctness never depends on it — the upstream
+  // fetch remains the source of truth.
+  const readEdgeCache = async (): Promise<JSONWebKeySet | null> => {
+    if (!edgeCache) return null;
+    // oxlint-disable-next-line executor/no-promise-catch -- boundary: Cache API failure degrades to a cache miss inside a jose Promise resolver
+    const cached = await edgeCache.match(url.toString()).catch(() => undefined);
+    if (!cached) return null;
+    // oxlint-disable-next-line executor/no-promise-catch -- boundary: a non-JSON cached body degrades to a cache miss inside a jose Promise resolver
+    const body: unknown = await cached.json().catch(() => null);
+    if (body === null) return null;
+    return decodeJsonWebKeySetPayload(body).then(
+      () => body as JSONWebKeySet,
+      () => null,
+    );
+  };
+
+  const writeEdgeCache = async (jwks: JSONWebKeySet): Promise<void> => {
+    if (!edgeCache) return;
+    // The Cache API owns expiry: `match` stops returning the entry once
+    // `max-age` elapses, so the shared copy has the same freshness window as
+    // the in-memory layer.
+    const response = new Response(JSON.stringify(jwks), {
+      headers: {
+        "content-type": "application/json",
+        "cache-control": `public, max-age=${Math.floor(ttlMs / 1000)}`,
+      },
     });
-    return inflight;
+    // oxlint-disable-next-line executor/no-promise-catch -- boundary: a failed shared-cache write is dropped; the fetched JWKS is already in hand
+    await edgeCache.put(url.toString(), response).catch(() => undefined);
+  };
+
+  const fetchUpstream = async (): Promise<JSONWebKeySet> => {
+    const jwks = await fetchJwksOnce(url, fetchImpl(), timeoutMs);
+    await writeEdgeCache(jwks);
+    return jwks;
+  };
+
+  let entry: CacheEntry | null = null;
+  let inflight: { promise: Promise<CacheEntry>; bypassedEdge: boolean } | null = null;
+  // Set by `forceRefresh()`: the shared edge copy is as suspect as our
+  // in-memory one (same staleness), so the next refresh must go upstream and
+  // overwrite it rather than read it back.
+  let bypassEdgeOnNextRefresh = false;
+
+  const refresh = (bypassEdge: boolean): Promise<CacheEntry> => {
+    const effectiveBypass = bypassEdge || bypassEdgeOnNextRefresh;
+    // Join an inflight refresh only when it satisfies this request: a
+    // bypassing caller must not receive the result of an edge-cache read.
+    if (inflight && (inflight.bypassedEdge || !effectiveBypass)) return inflight.promise;
+    bypassEdgeOnNextRefresh = false;
+    const record: { promise: Promise<CacheEntry>; bypassedEdge: boolean } = {
+      bypassedEdge: effectiveBypass,
+      promise: (async () => {
+        const jwks = effectiveBypass
+          ? await fetchUpstream()
+          : ((await readEdgeCache()) ?? (await fetchUpstream()));
+        const next: CacheEntry = {
+          jwks,
+          fetchedAt: Date.now(),
+          resolver: createLocalJWKSet(jwks),
+        };
+        entry = next;
+        return next;
+      })().finally(() => {
+        if (inflight === record) inflight = null;
+      }),
+    };
+    inflight = record;
+    return record.promise;
   };
 
   const ensureFresh = async (forceRefresh: boolean): Promise<CacheEntry> => {
-    if (forceRefresh) return refresh();
+    if (forceRefresh) return refresh(true);
     if (entry && Date.now() - entry.fetchedAt < ttlMs) return entry;
-    return refresh();
+    return refresh(false);
   };
 
   const get: JWTVerifyGetKey = async (protectedHeader, token) => {
@@ -171,6 +261,7 @@ export const createCachedRemoteJWKSet = (
   Object.defineProperty(result, "forceRefresh", {
     value: () => {
       entry = null;
+      bypassEdgeOnNextRefresh = true;
     },
   });
   Object.defineProperty(result, "inspect", {