History

cache JWKS to cut MCP auth p99 from 1.7s to ~5ms (#448)

Production telemetry (Axiom, 8h window) showed `mcp.auth.jwt_verify` fetched signin.executor.sh/oauth2/jwks 222 times with p99 1.71s. JWKS rotates on the order of hours/days, so this fetch was being repeated unnecessarily across short-lived Cloudflare Worker isolates. Replace `createRemoteJWKSet` with a custom `createCachedRemoteJWKSet` that: - caches the JSON Web Key Set in module-scope memory with a 1h TTL - single-flights concurrent verifies during a cache miss into one upstream fetch - force-refreshes on `JWKSNoMatchingKey` so genuine key rotations remain unblocked Tests cover N-verifies-one-fetch, single-flight stampede, forced refresh on rotation, TTL expiry, and explicit `forceRefresh()`.

Rhys Sullivan committed 6 months ago

Commit 899a99e1 parent ef45d8a3 files · +345 −2Browse files

Unified · Split

3 files
  1. apps/cloud/src/jwks-cache.node.test.ts +173 −0
  2. apps/cloud/src/jwks-cache.ts +163 −0
  3. apps/cloud/src/mcp.ts +9 −2
apps/cloud/src/jwks-cache.node.test.ts+173−0
@@ -0,0 +1,173 @@
1+import { describe, expect, it } from "@effect/vitest";
2+import {
3+ SignJWT,
4+ exportJWK,
5+ generateKeyPair,
6+ jwtVerify,
7+ type JSONWebKeySet,
8+ type JWK,
9+ type KeyLike,
10+} from "jose";
11+
12+import { createCachedRemoteJWKSet } from "./jwks-cache";
13+
14+const issuer = "https://test-authkit.example.com";
15+const audience = "client_test_fixture";
16+const jwksUrl = new URL("https://test-authkit.example.com/oauth2/jwks");
17+
18+interface Keypair {
19+ readonly kid: string;
20+ readonly publicJwk: JWK;
21+ readonly privateKey: KeyLike;
22+}
23+
24+const generateRotatableKeypair = async (kid: string): Promise<Keypair> => {
25+ const { publicKey, privateKey } = await generateKeyPair("RS256");
26+ const jwk = await exportJWK(publicKey);
27+ return { kid, publicJwk: { ...jwk, kid, alg: "RS256" }, privateKey };
28+};
29+
30+const sign = (keypair: Keypair) =>
31+ new SignJWT({})
32+ .setProtectedHeader({ alg: "RS256", kid: keypair.kid })
33+ .setIssuer(issuer)
34+ .setSubject("user_test")
35+ .setAudience(audience)
36+ .setIssuedAt()
37+ .setExpirationTime("5m")
38+ .sign(keypair.privateKey);
39+
40+interface FetchHarness {
41+ readonly fetch: typeof globalThis.fetch;
42+ readonly callCount: () => number;
43+ readonly setKeys: (keys: ReadonlyArray<JWK>) => void;
44+}
45+
46+const makeFetchHarness = (initialKeys: ReadonlyArray<JWK>): FetchHarness => {
47+ let keys: ReadonlyArray<JWK> = initialKeys;
48+ let calls = 0;
49+
50+ const fetch: typeof globalThis.fetch = async () => {
51+ calls++;
52+ const body: JSONWebKeySet = { keys: keys.map((k) => ({ ...k })) };
53+ return new Response(JSON.stringify(body), {
54+ status: 200,
55+ headers: { "content-type": "application/json" },
56+ });
57+ };
58+
59+ return {
60+ fetch,
61+ callCount: () => calls,
62+ setKeys: (next) => {
63+ keys = next;
64+ },
65+ };
66+};
67+
68+describe("createCachedRemoteJWKSet", () => {
69+ it("FAILING-WITHOUT-CACHE: N verifications hit JWKS endpoint only once within TTL", async () => {
70+ const kp = await generateRotatableKeypair("k1");
71+ const harness = makeFetchHarness([kp.publicJwk]);
72+ const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch: harness.fetch });
73+
74+ for (let i = 0; i < 5; i++) {
75+ const token = await sign(kp);
76+ const { payload } = await jwtVerify(token, jwks, { issuer, audience });
77+ expect(payload.sub).toBe("user_test");
78+ }
79+
80+ expect(harness.callCount()).toBe(1);
81+ });
82+
83+ it("single-flights concurrent cache misses into one fetch", async () => {
84+ const kp = await generateRotatableKeypair("k1");
85+
86+ let resolveFetch!: () => void;
87+ const gate = new Promise<void>((r) => {
88+ resolveFetch = r;
89+ });
90+
91+ let calls = 0;
92+ const fetch: typeof globalThis.fetch = async () => {
93+ calls++;
94+ await gate;
95+ const body: JSONWebKeySet = { keys: [{ ...kp.publicJwk }] };
96+ return new Response(JSON.stringify(body), {
97+ status: 200,
98+ headers: { "content-type": "application/json" },
99+ });
100+ };
101+
102+ const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch });
103+ const token = await sign(kp);
104+
105+ const verifies = Array.from({ length: 10 }, () =>
106+ jwtVerify(token, jwks, { issuer, audience }),
107+ );
108+ // Let microtasks settle so all 10 calls hit the cache miss path.
109+ await new Promise((r) => setTimeout(r, 10));
110+ resolveFetch();
111+ const results = await Promise.all(verifies);
112+ expect(results).toHaveLength(10);
113+ expect(calls).toBe(1);
114+ });
115+
116+ it("forces a refresh when verification fails with a no-matching-key error (key rotation)", async () => {
117+ const oldKey = await generateRotatableKeypair("k_old");
118+ const newKey = await generateRotatableKeypair("k_new");
119+
120+ const harness = makeFetchHarness([oldKey.publicJwk]);
121+ const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch: harness.fetch });
122+
123+ // Warm the cache with the old key.
124+ const t1 = await sign(oldKey);
125+ const ok = await jwtVerify(t1, jwks, { issuer, audience });
126+ expect(ok.payload.sub).toBe("user_test");
127+ expect(harness.callCount()).toBe(1);
128+
129+ // Upstream rotates: only the new key remains in the JWKS endpoint.
130+ harness.setKeys([newKey.publicJwk]);
131+
132+ // A token signed with the new key must verify even though our cache
133+ // still has the old one — the resolver must refetch on miss.
134+ const t2 = await sign(newKey);
135+ const ok2 = await jwtVerify(t2, jwks, { issuer, audience });
136+ expect(ok2.payload.sub).toBe("user_test");
137+ expect(harness.callCount()).toBe(2);
138+ });
139+
140+ it("re-fetches after the TTL window elapses", async () => {
141+ const kp = await generateRotatableKeypair("k1");
142+ const harness = makeFetchHarness([kp.publicJwk]);
143+ const jwks = createCachedRemoteJWKSet(jwksUrl, {
144+ fetch: harness.fetch,
145+ ttlMs: 10,
146+ });
147+
148+ const t1 = await sign(kp);
149+ await jwtVerify(t1, jwks, { issuer, audience });
150+ expect(harness.callCount()).toBe(1);
151+
152+ await new Promise((r) => setTimeout(r, 20));
153+
154+ const t2 = await sign(kp);
155+ await jwtVerify(t2, jwks, { issuer, audience });
156+ expect(harness.callCount()).toBe(2);
157+ });
158+
159+ it("forceRefresh() invalidates the cache so the next call refetches", async () => {
160+ const kp = await generateRotatableKeypair("k1");
161+ const harness = makeFetchHarness([kp.publicJwk]);
162+ const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch: harness.fetch });
163+
164+ const t1 = await sign(kp);
165+ await jwtVerify(t1, jwks, { issuer, audience });
166+ expect(harness.callCount()).toBe(1);
167+
168+ jwks.forceRefresh();
169+
170+ await jwtVerify(t1, jwks, { issuer, audience });
171+ expect(harness.callCount()).toBe(2);
172+ });
173+});
apps/cloud/src/jwks-cache.ts+163−0
@@ -0,0 +1,163 @@
1+// ---------------------------------------------------------------------------
2+// In-memory JWKS cache for MCP JWT verification.
3+// ---------------------------------------------------------------------------
4+//
5+// Cloudflare Workers boot many short-lived isolates. `createRemoteJWKSet`'s
6+// default cooldown (30s) and cache max-age (10m) still results in many JWKS
7+// fetches per hour because each new isolate starts cold. Production p99 for
8+// `mcp.auth.jwt_verify` was 1.7s — almost entirely the JWKS fetch.
9+//
10+// This module offers a drop-in `createCachedRemoteJWKSet` that:
11+//
12+// * Caches the JSON Web Key Set in module-scope memory for a configurable
13+// TTL (default 1 hour).
14+// * Single-flights concurrent fetches so a stampede of verifies during a
15+// cache miss only fires one upstream request.
16+// * Force-refreshes once when verification fails with a cached key, so
17+// genuine key rotation isn't blocked by the TTL.
18+//
19+// The returned function is a `JWTVerifyGetKey` and slots directly into
20+// `jose.jwtVerify`. It also exposes `forceRefresh()` so the verify path can
21+// invalidate the cache and retry on a bad signature.
22+// ---------------------------------------------------------------------------
23+
24+import {
25+ createLocalJWKSet,
26+ type FlattenedJWSInput,
27+ type JSONWebKeySet,
28+ type JWTHeaderParameters,
29+ type JWTVerifyGetKey,
30+ type KeyLike,
31+} from "jose";
32+import { JWKSNoMatchingKey } from "jose/errors";
33+
34+export interface CachedRemoteJWKSetOptions {
35+ /**
36+ * How long a successful fetch is considered fresh. Defaults to 1 hour —
37+ * AuthKit rotates JWKS roughly daily, and a forced refresh on verify
38+ * failure handles unscheduled rotations.
39+ */
40+ readonly ttlMs?: number;
41+ /** Override the fetch implementation for tests. */
42+ readonly fetch?: typeof globalThis.fetch;
43+ /** HTTP request timeout. Defaults to 5s, matching jose. */
44+ readonly timeoutMs?: number;
45+}
46+
47+export interface CachedRemoteJWKSet extends JWTVerifyGetKey {
48+ /** Drop the cached JWKS so the next call refetches. */
49+ readonly forceRefresh: () => void;
50+ /** Inspect the current cache state (testing/diagnostics). */
51+ readonly inspect: () => { fetchedAt: number | null; hasJwks: boolean };
52+}
53+
54+const DEFAULT_TTL_MS = 60 * 60 * 1000;
55+const DEFAULT_TIMEOUT_MS = 5000;
56+
57+interface CacheEntry {
58+ jwks: JSONWebKeySet;
59+ fetchedAt: number;
60+ resolver: (
61+ protectedHeader: JWTHeaderParameters,
62+ token?: FlattenedJWSInput,
63+ ) => Promise<KeyLike>;
64+}
65+
66+const fetchJwksOnce = async (
67+ url: URL,
68+ fetchImpl: typeof globalThis.fetch,
69+ timeoutMs: number,
70+): Promise<JSONWebKeySet> => {
71+ const controller = new AbortController();
72+ const timer = setTimeout(() => controller.abort(), timeoutMs);
73+ try {
74+ const response = await fetchImpl(url.toString(), {
75+ method: "GET",
76+ headers: { accept: "application/json" },
77+ signal: controller.signal,
78+ });
79+ if (!response.ok) {
80+ throw new Error(`JWKS fetch failed: ${response.status} ${response.statusText}`);
81+ }
82+ const body = (await response.json()) as JSONWebKeySet;
83+ if (!body || !Array.isArray((body as JSONWebKeySet).keys)) {
84+ throw new Error("JWKS fetch returned malformed payload");
85+ }
86+ return body;
87+ } finally {
88+ clearTimeout(timer);
89+ }
90+};
91+
92+/**
93+ * Creates a cached, single-flight, force-refreshable JWKS resolver compatible
94+ * with `jose.jwtVerify`. Drop-in replacement for `createRemoteJWKSet` for the
95+ * MCP auth path — see module header for why we don't just use jose's built-in.
96+ */
97+export const createCachedRemoteJWKSet = (
98+ url: URL,
99+ options: CachedRemoteJWKSetOptions = {},
100+): CachedRemoteJWKSet => {
101+ const ttlMs = options.ttlMs ?? DEFAULT_TTL_MS;
102+ const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
103+ // Capture the fetch impl lazily so consumers can swap globalThis.fetch
104+ // (tests do this) without us snapshotting a stale reference.
105+ const fetchImpl = (): typeof globalThis.fetch =>
106+ options.fetch ?? globalThis.fetch.bind(globalThis);
107+
108+ let entry: CacheEntry | null = null;
109+ let inflight: Promise<CacheEntry> | null = null;
110+
111+ const refresh = (): Promise<CacheEntry> => {
112+ if (inflight) return inflight;
113+ inflight = (async () => {
114+ try {
115+ const jwks = await fetchJwksOnce(url, fetchImpl(), timeoutMs);
116+ const next: CacheEntry = {
117+ jwks,
118+ fetchedAt: Date.now(),
119+ resolver: createLocalJWKSet(jwks),
120+ };
121+ entry = next;
122+ return next;
123+ } finally {
124+ inflight = null;
125+ }
126+ })();
127+ return inflight;
128+ };
129+
130+ const ensureFresh = async (forceRefresh: boolean): Promise<CacheEntry> => {
131+ if (forceRefresh) return refresh();
132+ if (entry && Date.now() - entry.fetchedAt < ttlMs) return entry;
133+ return refresh();
134+ };
135+
136+ const get: JWTVerifyGetKey = async (protectedHeader, token) => {
137+ const current = await ensureFresh(false);
138+ try {
139+ return await current.resolver(protectedHeader, token);
140+ } catch (error) {
141+ // Likely cause: keys rotated upstream after our TTL window started.
142+ // Refetch once and try again. Anything still failing bubbles up so
143+ // jose can classify it (we do not silently swallow real failures).
144+ if (!(error instanceof JWKSNoMatchingKey)) throw error;
145+ const refreshed = await ensureFresh(true);
146+ return await refreshed.resolver(protectedHeader, token);
147+ }
148+ };
149+
150+ const result = get as CachedRemoteJWKSet;
151+ Object.defineProperty(result, "forceRefresh", {
152+ value: () => {
153+ entry = null;
154+ },
155+ });
156+ Object.defineProperty(result, "inspect", {
157+ value: () => ({
158+ fetchedAt: entry?.fetchedAt ?? null,
159+ hasJwks: entry !== null,
160+ }),
161+ });
162+ return result;
163+};
apps/cloud/src/mcp.ts+9−2
@@ -18,8 +18,8 @@ import { env } from "cloudflare:workers"; · 17 unmodified lines
1818 import { HttpApp, HttpServerRequest, HttpServerResponse } from "@effect/platform";
1919 import * as Sentry from "@sentry/cloudflare";
2020 import { Context, Effect, Layer, Option, Schema } from "effect";
21−import { createRemoteJWKSet } from "jose";
2221
22+import { createCachedRemoteJWKSet } from "./jwks-cache";
2323 import { TelemetryLive } from "./services/telemetry";
2424 import {
2525 McpJwtVerificationError,
@@ -47,7 +47,14 @@ const AUTHKIT_DOMAIN = env.MCP_AUTHKIT_DOMAIN ?? "https://signin.executor.sh"; · 21 unmodified lines
4747 const RESOURCE_ORIGIN = env.MCP_RESOURCE_ORIGIN ?? "https://executor.sh";
4848 const WORKOS_CLIENT_ID = env.WORKOS_CLIENT_ID;
4949
50−const jwks = createRemoteJWKSet(new URL(`${AUTHKIT_DOMAIN}/oauth2/jwks`));
50+// Module-scope cache survives across MCP requests within the same worker
51+// isolate. AuthKit's JWKS rotates on the order of hours/days, so a 1h TTL
52+// dominates the upstream cooldown without sacrificing rotation safety —
53+// `createCachedRemoteJWKSet` force-refreshes on key-not-found inside its
54+// resolver. Production telemetry showed ~222 fetches/8h with p99 1.7s on
55+// the previous default-cooldown setup; this collapses that to ~1 per
56+// isolate-hour.
57+const jwks = createCachedRemoteJWKSet(new URL(`${AUTHKIT_DOMAIN}/oauth2/jwks`));
5158
5259 const BEARER_PREFIX = "Bearer ";
5360 const INTERNAL_ACCOUNT_ID_HEADER = "x-executor-mcp-account-id";