Revert "Back the JWKS cache with the Workers edge cache (#1479)" (#1480)

This reverts commit 75f16edd319813dbd2bc7b0eb1b3edab32ba3210.

Rhys Sullivan committed 3 months ago

Commit ddbbe861 parent 167d8992 files · +21 −255Browse files

Unified · Split

2 files

apps/cloud/src/auth/jwks-cache.node.test.ts +1 −144
apps/cloud/src/auth/jwks-cache.ts +20 −111

apps/cloud/src/auth/jwks-cache.node.test.ts+1−144

@@ -9,7 +9,7 @@ import { · 8 unmodified lines
9	type KeyLike,	9	type KeyLike,
10	} from "jose";	10	} from "jose";
11		11
12	import { createCachedRemoteJWKSet, type JwksEdgeCache } from "./jwks-cache";	12	import { createCachedRemoteJWKSet } from "./jwks-cache";
13		13
14	const issuer = "https://test-authkit.example.com";	14	const issuer = "https://test-authkit.example.com";
15	const audience = "client_test_fixture";	15	const audience = "client_test_fixture";
@@ -169,146 +169,3 @@ describe("createCachedRemoteJWKSet", () => { · 153 unmodified lines
169	expect(harness.callCount()).toBe(2);	169	expect(harness.callCount()).toBe(2);
170	});	170	});
171	});	171	});
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+20−111

@@ -11,18 +11,10 @@ · 10 unmodified lines
11	//	11	//
12	// * Caches the JSON Web Key Set in module-scope memory for a configurable	12	// * Caches the JSON Web Key Set in module-scope memory for a configurable
13	// TTL (default 1 hour).	13	// 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.
20	// * Single-flights concurrent fetches so a stampede of verifies during a	14	// * Single-flights concurrent fetches so a stampede of verifies during a
21	// cache miss only fires one upstream request.	15	// cache miss only fires one upstream request.
22	// * Force-refreshes once when verification fails with a cached key, so	16	// * Force-refreshes once when verification fails with a cached key, so
23	// genuine key rotation isn't blocked by the TTL. A forced refresh also	17	// genuine key rotation isn't blocked by the TTL.
24	// bypasses the shared edge copy (it is just as stale as ours) and
25	// overwrites it with the freshly fetched set.
26	//	18	//
27	// The returned function is a `JWTVerifyGetKey` and slots directly into	19	// The returned function is a `JWTVerifyGetKey` and slots directly into
28	// `jose.jwtVerify`. It also exposes `forceRefresh()` so the verify path can	20	// `jose.jwtVerify`. It also exposes `forceRefresh()` so the verify path can
@@ -40,21 +32,6 @@ import { · 11 unmodified lines
40	import { Schema } from "effect";	32	import { Schema } from "effect";
41	import { JWKSNoMatchingKey } from "jose/errors";	33	import { JWKSNoMatchingKey } from "jose/errors";
42		34
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
58	export interface CachedRemoteJWKSetOptions {	35	export interface CachedRemoteJWKSetOptions {
59	/**	36	/**
60	* How long a successful fetch is considered fresh. Defaults to 1 hour —	37	* How long a successful fetch is considered fresh. Defaults to 1 hour —
@@ -66,12 +43,6 @@ export interface CachedRemoteJWKSetOptions { · 5 unmodified lines
66	readonly fetch?: typeof globalThis.fetch;	43	readonly fetch?: typeof globalThis.fetch;
67	/** HTTP request timeout. Defaults to 5s, matching jose. */	44	/** HTTP request timeout. Defaults to 5s, matching jose. */
68	readonly timeoutMs?: number;	45	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;
75	}	46	}
76		47
77	export interface CachedRemoteJWKSet extends JWTVerifyGetKey {	48	export interface CachedRemoteJWKSet extends JWTVerifyGetKey {
@@ -151,92 +122,31 @@ export const createCachedRemoteJWKSet = ( · 73 unmodified lines
151	// (tests do this) without us snapshotting a stale reference.	122	// (tests do this) without us snapshotting a stale reference.
152	const fetchImpl = (): typeof globalThis.fetch =>	123	const fetchImpl = (): typeof globalThis.fetch =>
153	options.fetch ?? globalThis.fetch.bind(globalThis);	124	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;
163
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	},
191	});
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		125
202	let entry: CacheEntry \| null = null;	126	let entry: CacheEntry \| null = null;
203	let inflight: { promise: Promise<CacheEntry>; bypassedEdge: boolean } \| null = null;	127	let inflight: Promise<CacheEntry> \| null = null;
204	// Set by `forceRefresh()`: the shared edge copy is as suspect as our	128
205	// in-memory one (same staleness), so the next refresh must go upstream and	129	const refresh = (): Promise<CacheEntry> => {
206	// overwrite it rather than read it back.	130	if (inflight) return inflight;
207	let bypassEdgeOnNextRefresh = false;	131	inflight = (async () => {
208		132	const jwks = await fetchJwksOnce(url, fetchImpl(), timeoutMs);
209	const refresh = (bypassEdge: boolean): Promise<CacheEntry> => {	133	const next: CacheEntry = {
210	const effectiveBypass = bypassEdge \|\| bypassEdgeOnNextRefresh;	134	jwks,
211	// Join an inflight refresh only when it satisfies this request: a	135	fetchedAt: Date.now(),
212	// bypassing caller must not receive the result of an edge-cache read.	136	resolver: createLocalJWKSet(jwks),
213	if (inflight && (inflight.bypassedEdge \|\| !effectiveBypass)) return inflight.promise;	137	};
214	bypassEdgeOnNextRefresh = false;	138	entry = next;
215	const record: { promise: Promise<CacheEntry>; bypassedEdge: boolean } = {	139	return next;
216	bypassedEdge: effectiveBypass,	140	})().finally(() => {
217	promise: (async () => {	141	inflight = null;
218	const jwks = effectiveBypass	142	});
219	? await fetchUpstream()	143	return inflight;
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;
234	};	144	};
235		145
236	const ensureFresh = async (forceRefresh: boolean): Promise<CacheEntry> => {	146	const ensureFresh = async (forceRefresh: boolean): Promise<CacheEntry> => {
237	if (forceRefresh) return refresh(true);	147	if (forceRefresh) return refresh();
238	if (entry && Date.now() - entry.fetchedAt < ttlMs) return entry;	148	if (entry && Date.now() - entry.fetchedAt < ttlMs) return entry;
239	return refresh(false);	149	return refresh();
240	};	150	};
241		151
242	const get: JWTVerifyGetKey = async (protectedHeader, token) => {	152	const get: JWTVerifyGetKey = async (protectedHeader, token) => {
@@ -261,7 +171,6 @@ export const createCachedRemoteJWKSet = ( · 18 unmodified lines
261	Object.defineProperty(result, "forceRefresh", {	171	Object.defineProperty(result, "forceRefresh", {
262	value: () => {	172	value: () => {
263	entry = null;	173	entry = null;
264	bypassEdgeOnNextRefresh = true;
265	},	174	},
266	});	175	});
267	Object.defineProperty(result, "inspect", {	176	Object.defineProperty(result, "inspect", {