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 @@
1import { describe, expect, it } from "@effect/vitest";
2import {
3 SignJWT,
4 exportJWK,
5 generateKeyPair,
6 jwtVerify,
7 type JSONWebKeySet,
8 type JWK,
9 type KeyLike,
10} from "jose";
11
12import { createCachedRemoteJWKSet } from "./jwks-cache";
13
14const issuer = "https://test-authkit.example.com";
15const audience = "client_test_fixture";
16const jwksUrl = new URL("https://test-authkit.example.com/oauth2/jwks");
17
18interface Keypair {
19 readonly kid: string;
20 readonly publicJwk: JWK;
21 readonly privateKey: KeyLike;
22}
23
24const 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
30const 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
40interface FetchHarness {
41 readonly fetch: typeof globalThis.fetch;
42 readonly callCount: () => number;
43 readonly setKeys: (keys: ReadonlyArray<JWK>) => void;
44}
45
46const 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
68describe("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
24import {
25 createLocalJWKSet,
26 type FlattenedJWSInput,
27 type JSONWebKeySet,
28 type JWTHeaderParameters,
29 type JWTVerifyGetKey,
30 type KeyLike,
31} from "jose";
32import { JWKSNoMatchingKey } from "jose/errors";
33
34export 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
47export 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
54const DEFAULT_TTL_MS = 60 * 60 * 1000;
55const DEFAULT_TIMEOUT_MS = 5000;
56
57interface CacheEntry {
58 jwks: JSONWebKeySet;
59 fetchedAt: number;
60 resolver: (
61 protectedHeader: JWTHeaderParameters,
62 token?: FlattenedJWSInput,
63 ) => Promise<KeyLike>;
64}
65
66const 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 */
97export 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
18import { HttpApp, HttpServerRequest, HttpServerResponse } from "@effect/platform";18import { HttpApp, HttpServerRequest, HttpServerResponse } from "@effect/platform";
19import * as Sentry from "@sentry/cloudflare";19import * as Sentry from "@sentry/cloudflare";
20import { Context, Effect, Layer, Option, Schema } from "effect";20import { Context, Effect, Layer, Option, Schema } from "effect";
21import { createRemoteJWKSet } from "jose";
22 21
22import { createCachedRemoteJWKSet } from "./jwks-cache";
23import { TelemetryLive } from "./services/telemetry";23import { TelemetryLive } from "./services/telemetry";
24import {24import {
25 McpJwtVerificationError,25 McpJwtVerificationError,
@@ -47,7 +47,14 @@ const AUTHKIT_DOMAIN = env.MCP_AUTHKIT_DOMAIN ?? "https://signin.executor.sh"; · 21 unmodified lines
47const RESOURCE_ORIGIN = env.MCP_RESOURCE_ORIGIN ?? "https://executor.sh";47const RESOURCE_ORIGIN = env.MCP_RESOURCE_ORIGIN ?? "https://executor.sh";
48const WORKOS_CLIENT_ID = env.WORKOS_CLIENT_ID;48const WORKOS_CLIENT_ID = env.WORKOS_CLIENT_ID;
49 49
50const 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.
57const jwks = createCachedRemoteJWKSet(new URL(`${AUTHKIT_DOMAIN}/oauth2/jwks`));
51 58
52const BEARER_PREFIX = "Bearer ";59const BEARER_PREFIX = "Bearer ";
53const INTERNAL_ACCOUNT_ID_HEADER = "x-executor-mcp-account-id";60const INTERNAL_ACCOUNT_ID_HEADER = "x-executor-mcp-account-id";