History

Back the JWKS cache with the Workers edge cache

Rhys Sullivan committed 3 months ago

Commit 8aa3e6c1 parent fc1e5892 files · +255 −21Browse files

Unified · Split

2 files
  1. apps/cloud/src/auth/jwks-cache.node.test.ts +144 −1
  2. apps/cloud/src/auth/jwks-cache.ts +111 −20
apps/cloud/src/auth/jwks-cache.node.test.ts+144−1
@@ -9,7 +9,7 @@ import { · 8 unmodified lines
99 type KeyLike,
1010 } from "jose";
1111
12−import { createCachedRemoteJWKSet } from "./jwks-cache";
12+import { createCachedRemoteJWKSet, type JwksEdgeCache } from "./jwks-cache";
1313
1414 const issuer = "https://test-authkit.example.com";
1515 const audience = "client_test_fixture";
@@ -169,3 +169,146 @@ describe("createCachedRemoteJWKSet", () => { · 153 unmodified lines
169169 expect(harness.callCount()).toBe(2);
170170 });
171171 });
172+
173+interface EdgeCacheHarness {
174+ readonly cache: JwksEdgeCache;
175+ readonly matchCount: () => number;
176+ readonly putCount: () => number;
177+ readonly stored: () => Response | null;
178+ readonly seed: (keys: ReadonlyArray<JWK>) => void;
179+}
180+
181+const makeEdgeCacheHarness = (): EdgeCacheHarness => {
182+ let stored: Response | null = null;
183+ let matches = 0;
184+ let puts = 0;
185+
186+ return {
187+ cache: {
188+ match: async () => {
189+ matches++;
190+ return stored ? stored.clone() : undefined;
191+ },
192+ put: async (_url, response) => {
193+ puts++;
194+ stored = response;
195+ },
196+ },
197+ matchCount: () => matches,
198+ putCount: () => puts,
199+ stored: () => stored,
200+ seed: (keys) => {
201+ const body: JSONWebKeySet = { keys: keys.map((k) => ({ ...k })) };
202+ stored = new Response(JSON.stringify(body), {
203+ headers: { "content-type": "application/json" },
204+ });
205+ },
206+ };
207+};
208+
209+describe("createCachedRemoteJWKSet edge cache", () => {
210+ it("FAILING-WITHOUT-EDGE-CACHE: a cold isolate with a warm edge cache never fetches upstream", async () => {
211+ const kp = await generateRotatableKeypair("k1");
212+ const upstream = makeFetchHarness([kp.publicJwk]);
213+ const edge = makeEdgeCacheHarness();
214+ edge.seed([kp.publicJwk]);
215+
216+ // Fresh resolver = fresh isolate: empty in-memory cache.
217+ const jwks = createCachedRemoteJWKSet(jwksUrl, {
218+ fetch: upstream.fetch,
219+ edgeCache: edge.cache,
220+ });
221+
222+ const token = await sign(kp);
223+ const { payload } = await jwtVerify(token, jwks, { issuer, audience });
224+ expect(payload.sub).toBe("user_test");
225+ expect(edge.matchCount()).toBe(1);
226+ expect(upstream.callCount()).toBe(0);
227+ });
228+
229+ it("an edge-cache miss falls through to upstream and populates the edge copy", async () => {
230+ const kp = await generateRotatableKeypair("k1");
231+ const upstream = makeFetchHarness([kp.publicJwk]);
232+ const edge = makeEdgeCacheHarness();
233+
234+ const jwks = createCachedRemoteJWKSet(jwksUrl, {
235+ fetch: upstream.fetch,
236+ edgeCache: edge.cache,
237+ });
238+
239+ const token = await sign(kp);
240+ await jwtVerify(token, jwks, { issuer, audience });
241+ expect(upstream.callCount()).toBe(1);
242+ expect(edge.putCount()).toBe(1);
243+
244+ // A second cold isolate now reads the populated edge copy.
245+ const jwks2 = createCachedRemoteJWKSet(jwksUrl, {
246+ fetch: upstream.fetch,
247+ edgeCache: edge.cache,
248+ });
249+ await jwtVerify(token, jwks2, { issuer, audience });
250+ expect(upstream.callCount()).toBe(1);
251+ });
252+
253+ it("a malformed edge-cache entry is a miss, not a failure", async () => {
254+ const kp = await generateRotatableKeypair("k1");
255+ const upstream = makeFetchHarness([kp.publicJwk]);
256+ const edge = makeEdgeCacheHarness();
257+ await edge.cache.put(jwksUrl.toString(), new Response("not json"));
258+
259+ const jwks = createCachedRemoteJWKSet(jwksUrl, {
260+ fetch: upstream.fetch,
261+ edgeCache: edge.cache,
262+ });
263+
264+ const token = await sign(kp);
265+ const { payload } = await jwtVerify(token, jwks, { issuer, audience });
266+ expect(payload.sub).toBe("user_test");
267+ expect(upstream.callCount()).toBe(1);
268+ });
269+
270+ it("key rotation bypasses the stale edge copy and overwrites it", async () => {
271+ const oldKey = await generateRotatableKeypair("k_old");
272+ const newKey = await generateRotatableKeypair("k_new");
273+ const upstream = makeFetchHarness([newKey.publicJwk]);
274+ const edge = makeEdgeCacheHarness();
275+ // The edge copy predates the rotation: it only has the old key.
276+ edge.seed([oldKey.publicJwk]);
277+
278+ const jwks = createCachedRemoteJWKSet(jwksUrl, {
279+ fetch: upstream.fetch,
280+ edgeCache: edge.cache,
281+ });
282+
283+ // First verify warms in-memory from the (stale) edge copy; the new-key
284+ // token then misses, which must force an UPSTREAM refetch — re-reading
285+ // the same stale edge copy would loop the miss forever.
286+ const token = await sign(newKey);
287+ const { payload } = await jwtVerify(token, jwks, { issuer, audience });
288+ expect(payload.sub).toBe("user_test");
289+ expect(upstream.callCount()).toBe(1);
290+ // And the rotated set replaced the stale edge copy for other isolates.
291+ expect(edge.putCount()).toBe(1);
292+ });
293+
294+ it("degrades to upstream-only when the edge cache itself fails", async () => {
295+ const kp = await generateRotatableKeypair("k1");
296+ const upstream = makeFetchHarness([kp.publicJwk]);
297+ // oxlint-disable-next-line executor/no-promise-reject, executor/no-error-constructor -- boundary: fakes the Cache API's own failure mode, a rejected promise
298+ const rejectingCacheCall = () => Promise.reject(new TypeError("cache backend unavailable"));
299+ const broken: JwksEdgeCache = {
300+ match: rejectingCacheCall,
301+ put: rejectingCacheCall,
302+ };
303+
304+ const jwks = createCachedRemoteJWKSet(jwksUrl, {
305+ fetch: upstream.fetch,
306+ edgeCache: broken,
307+ });
308+
309+ const token = await sign(kp);
310+ const { payload } = await jwtVerify(token, jwks, { issuer, audience });
311+ expect(payload.sub).toBe("user_test");
312+ expect(upstream.callCount()).toBe(1);
313+ });
314+});
apps/cloud/src/auth/jwks-cache.ts+111−20
@@ -11,10 +11,18 @@ · 10 unmodified lines
1111 //
1212 // * Caches the JSON Web Key Set in module-scope memory for a configurable
1313 // TTL (default 1 hour).
14+// * Backs the in-memory layer with the Workers Cache API (`caches.default`)
15+// so a COLD isolate reads the JWKS from its datacenter's shared cache
16+// instead of round-tripping to the upstream endpoint. The in-memory layer
17+// alone left every fresh isolate paying that full fetch, which is exactly
18+// when users notice: an app-open burst fans out across many new isolates
19+// and each one stalled ~1-2s inside session verification.
1420 // * Single-flights concurrent fetches so a stampede of verifies during a
1521 // cache miss only fires one upstream request.
1622 // * Force-refreshes once when verification fails with a cached key, so
17−// genuine key rotation isn't blocked by the TTL.
23+// genuine key rotation isn't blocked by the TTL. A forced refresh also
24+// bypasses the shared edge copy (it is just as stale as ours) and
25+// overwrites it with the freshly fetched set.
1826 //
1927 // The returned function is a `JWTVerifyGetKey` and slots directly into
2028 // `jose.jwtVerify`. It also exposes `forceRefresh()` so the verify path can
@@ -32,6 +40,21 @@ import { · 11 unmodified lines
3240 import { Schema } from "effect";
3341 import { JWKSNoMatchingKey } from "jose/errors";
3442
43+/** The slice of the Workers Cache API this module uses, expressed
44+ * structurally: `Cache` from workers-types satisfies it, and node tests can
45+ * supply a conforming fake without the full `Cache` surface. */
46+export interface JwksEdgeCache {
47+ match(url: string): Promise<Response | undefined>;
48+ put(url: string, response: Response): Promise<void>;
49+}
50+
51+// The `caches` global exists in the Workers runtime but not under node (the
52+// test environment). Typed structurally (the ASSETS precedent in
53+// env-augment.d.ts) rather than via workers-types: its `Cache` type drags in
54+// a second `Response` that intersects with the ambient one and rejects
55+// responses this module constructs.
56+declare const caches: { readonly default: JwksEdgeCache } | undefined;
57+
3558 export interface CachedRemoteJWKSetOptions {
3659 /**
3760 * How long a successful fetch is considered fresh. Defaults to 1 hour —
@@ -43,6 +66,12 @@ export interface CachedRemoteJWKSetOptions { · 5 unmodified lines
4366 readonly fetch?: typeof globalThis.fetch;
4467 /** HTTP request timeout. Defaults to 5s, matching jose. */
4568 readonly timeoutMs?: number;
69+ /**
70+ * Datacenter-shared cache behind the in-memory layer. Defaults to
71+ * `caches.default` when the runtime provides it (Workers) and to disabled
72+ * when it does not (node tests). Pass explicitly to fake it in tests.
73+ */
74+ readonly edgeCache?: JwksEdgeCache | null;
4675 }
4776
4877 export interface CachedRemoteJWKSet extends JWTVerifyGetKey {
@@ -122,31 +151,92 @@ export const createCachedRemoteJWKSet = ( · 73 unmodified lines
122151 // (tests do this) without us snapshotting a stale reference.
123152 const fetchImpl = (): typeof globalThis.fetch =>
124153 options.fetch ?? globalThis.fetch.bind(globalThis);
154+ // `undefined` = not specified, use the runtime default; `null` = disabled.
155+ // Node (tests) has no `caches` global, so the default degrades to disabled
156+ // there without configuration.
157+ const edgeCache =
158+ options.edgeCache !== undefined
159+ ? options.edgeCache
160+ : typeof caches !== "undefined"
161+ ? caches.default
162+ : null;
125163
126− let entry: CacheEntry | null = null;
127− let inflight: Promise<CacheEntry> | null = null;
128−
129− const refresh = (): Promise<CacheEntry> => {
130− if (inflight) return inflight;
131− inflight = (async () => {
132− const jwks = await fetchJwksOnce(url, fetchImpl(), timeoutMs);
133− const next: CacheEntry = {
134− jwks,
135− fetchedAt: Date.now(),
136− resolver: createLocalJWKSet(jwks),
137− };
138− entry = next;
139− return next;
140− })().finally(() => {
141− inflight = null;
164+ // The edge cache is best-effort: a failed read is a miss, a failed write is
165+ // dropped. Verification correctness never depends on it — the upstream
166+ // fetch remains the source of truth.
167+ const readEdgeCache = async (): Promise<JSONWebKeySet | null> => {
168+ if (!edgeCache) return null;
169+ // oxlint-disable-next-line executor/no-promise-catch -- boundary: Cache API failure degrades to a cache miss inside a jose Promise resolver
170+ const cached = await edgeCache.match(url.toString()).catch(() => undefined);
171+ if (!cached) return null;
172+ // oxlint-disable-next-line executor/no-promise-catch -- boundary: a non-JSON cached body degrades to a cache miss inside a jose Promise resolver
173+ const body: unknown = await cached.json().catch(() => null);
174+ if (body === null) return null;
175+ return decodeJsonWebKeySetPayload(body).then(
176+ () => body as JSONWebKeySet,
177+ () => null,
178+ );
179+ };
180+
181+ const writeEdgeCache = async (jwks: JSONWebKeySet): Promise<void> => {
182+ if (!edgeCache) return;
183+ // The Cache API owns expiry: `match` stops returning the entry once
184+ // `max-age` elapses, so the shared copy has the same freshness window as
185+ // the in-memory layer.
186+ const response = new Response(JSON.stringify(jwks), {
187+ headers: {
188+ "content-type": "application/json",
189+ "cache-control": `public, max-age=${Math.floor(ttlMs / 1000)}`,
190+ },
142191 });
143− return inflight;
192+ // oxlint-disable-next-line executor/no-promise-catch -- boundary: a failed shared-cache write is dropped; the fetched JWKS is already in hand
193+ await edgeCache.put(url.toString(), response).catch(() => undefined);
194+ };
195+
196+ const fetchUpstream = async (): Promise<JSONWebKeySet> => {
197+ const jwks = await fetchJwksOnce(url, fetchImpl(), timeoutMs);
198+ await writeEdgeCache(jwks);
199+ return jwks;
200+ };
201+
202+ let entry: CacheEntry | null = null;
203+ let inflight: { promise: Promise<CacheEntry>; bypassedEdge: boolean } | null = null;
204+ // Set by `forceRefresh()`: the shared edge copy is as suspect as our
205+ // in-memory one (same staleness), so the next refresh must go upstream and
206+ // overwrite it rather than read it back.
207+ let bypassEdgeOnNextRefresh = false;
208+
209+ const refresh = (bypassEdge: boolean): Promise<CacheEntry> => {
210+ const effectiveBypass = bypassEdge || bypassEdgeOnNextRefresh;
211+ // Join an inflight refresh only when it satisfies this request: a
212+ // bypassing caller must not receive the result of an edge-cache read.
213+ if (inflight && (inflight.bypassedEdge || !effectiveBypass)) return inflight.promise;
214+ bypassEdgeOnNextRefresh = false;
215+ const record: { promise: Promise<CacheEntry>; bypassedEdge: boolean } = {
216+ bypassedEdge: effectiveBypass,
217+ promise: (async () => {
218+ const jwks = effectiveBypass
219+ ? await fetchUpstream()
220+ : ((await readEdgeCache()) ?? (await fetchUpstream()));
221+ const next: CacheEntry = {
222+ jwks,
223+ fetchedAt: Date.now(),
224+ resolver: createLocalJWKSet(jwks),
225+ };
226+ entry = next;
227+ return next;
228+ })().finally(() => {
229+ if (inflight === record) inflight = null;
230+ }),
231+ };
232+ inflight = record;
233+ return record.promise;
144234 };
145235
146236 const ensureFresh = async (forceRefresh: boolean): Promise<CacheEntry> => {
147− if (forceRefresh) return refresh();
237+ if (forceRefresh) return refresh(true);
148238 if (entry && Date.now() - entry.fetchedAt < ttlMs) return entry;
149− return refresh();
239+ return refresh(false);
150240 };
151241
152242 const get: JWTVerifyGetKey = async (protectedHeader, token) => {
@@ -171,6 +261,7 @@ export const createCachedRemoteJWKSet = ( · 18 unmodified lines
171261 Object.defineProperty(result, "forceRefresh", {
172262 value: () => {
173263 entry = null;
264+ bypassEdgeOnNextRefresh = true;
174265 },
175266 });
176267 Object.defineProperty(result, "inspect", {