352 lines · 11.9 KB · TypeScript
RawDownload
import { describe, expect, it } from "@effect/vitest";import {  SignJWT,  exportJWK,  generateKeyPair,  jwtVerify,  type JSONWebKeySet,  type JWK,  type KeyLike,} from "jose"; import { createCachedRemoteJWKSet, type JwksStore, type StoredJwks } from "./jwks-cache"; const issuer = "https://test-authkit.example.com";const audience = "client_test_fixture";const jwksUrl = new URL("https://test-authkit.example.com/oauth2/jwks"); interface Keypair {  readonly kid: string;  readonly publicJwk: JWK;  readonly privateKey: KeyLike;} const generateRotatableKeypair = async (kid: string): Promise<Keypair> => {  const { publicKey, privateKey } = await generateKeyPair("RS256");  const jwk = await exportJWK(publicKey);  return { kid, publicJwk: { ...jwk, kid, alg: "RS256" }, privateKey };}; const sign = (keypair: Keypair) =>  new SignJWT({})    .setProtectedHeader({ alg: "RS256", kid: keypair.kid })    .setIssuer(issuer)    .setSubject("user_test")    .setAudience(audience)    .setIssuedAt()    .setExpirationTime("5m")    .sign(keypair.privateKey); interface FetchHarness {  readonly fetch: typeof globalThis.fetch;  readonly callCount: () => number;  readonly setKeys: (keys: ReadonlyArray<JWK>) => void;} const makeFetchHarness = (initialKeys: ReadonlyArray<JWK>): FetchHarness => {  let keys: ReadonlyArray<JWK> = initialKeys;  let calls = 0;   const fetch: typeof globalThis.fetch = async () => {    calls++;    const body: JSONWebKeySet = { keys: keys.map((k) => ({ ...k })) };    return new Response(JSON.stringify(body), {      status: 200,      headers: { "content-type": "application/json" },    });  };   return {    fetch,    callCount: () => calls,    setKeys: (next) => {      keys = next;    },  };}; interface StoreHarness extends JwksStore {  readonly seed: (stored: StoredJwks) => void;  readonly reads: () => number;  readonly writes: () => number;} /** Stands in for the Workers Cache API: shared across "isolates", in memory. */const makeStoreHarness = (): StoreHarness => {  const entries = new Map<string, StoredJwks>();  let reads = 0;  let writes = 0;  return {    get: async (url) => {      reads++;      return entries.get(url.toString()) ?? null;    },    put: async (url, stored) => {      writes++;      entries.set(url.toString(), stored);    },    seed: (stored) => {      entries.set(jwksUrl.toString(), stored);    },    reads: () => reads,    writes: () => writes,  };}; /** A fetch that always fails, standing in for a slow/down key server. */const failingFetch: typeof globalThis.fetch = async () => {  // oxlint-disable-next-line executor/no-try-catch-or-throw, executor/no-error-constructor -- test double: `fetch` signals an unreachable key server by rejecting  throw new Error("JWKS endpoint unreachable");}; describe("createCachedRemoteJWKSet", () => {  it("FAILING-WITHOUT-CACHE: N verifications hit JWKS endpoint only once within TTL", async () => {    const kp = await generateRotatableKeypair("k1");    const harness = makeFetchHarness([kp.publicJwk]);    const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch: harness.fetch });     for (let i = 0; i < 5; i++) {      const token = await sign(kp);      const { payload } = await jwtVerify(token, jwks, { issuer, audience });      expect(payload.sub).toBe("user_test");    }     expect(harness.callCount()).toBe(1);  });   it("single-flights concurrent cache misses into one fetch", async () => {    const kp = await generateRotatableKeypair("k1");     let resolveFetch!: () => void;    const gate = new Promise<void>((r) => {      resolveFetch = r;    });     let calls = 0;    const fetch: typeof globalThis.fetch = async () => {      calls++;      await gate;      const body: JSONWebKeySet = { keys: [{ ...kp.publicJwk }] };      return new Response(JSON.stringify(body), {        status: 200,        headers: { "content-type": "application/json" },      });    };     const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch });    const token = await sign(kp);     const verifies = Array.from({ length: 10 }, () => jwtVerify(token, jwks, { issuer, audience }));    // Let microtasks settle so all 10 calls hit the cache miss path.    await new Promise((r) => setTimeout(r, 10));    resolveFetch();    const results = await Promise.all(verifies);    expect(results).toHaveLength(10);    expect(calls).toBe(1);  });   it("forces a refresh when verification fails with a no-matching-key error (key rotation)", async () => {    const oldKey = await generateRotatableKeypair("k_old");    const newKey = await generateRotatableKeypair("k_new");     const harness = makeFetchHarness([oldKey.publicJwk]);    const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch: harness.fetch });     // Warm the cache with the old key.    const t1 = await sign(oldKey);    const ok = await jwtVerify(t1, jwks, { issuer, audience });    expect(ok.payload.sub).toBe("user_test");    expect(harness.callCount()).toBe(1);     // Upstream rotates: only the new key remains in the JWKS endpoint.    harness.setKeys([newKey.publicJwk]);     // A token signed with the new key must verify even though our cache    // still has the old one — the resolver must refetch on miss.    const t2 = await sign(newKey);    const ok2 = await jwtVerify(t2, jwks, { issuer, audience });    expect(ok2.payload.sub).toBe("user_test");    expect(harness.callCount()).toBe(2);  });   it("re-fetches after the TTL window elapses", async () => {    const kp = await generateRotatableKeypair("k1");    const harness = makeFetchHarness([kp.publicJwk]);    const jwks = createCachedRemoteJWKSet(jwksUrl, {      fetch: harness.fetch,      ttlMs: 10,    });     const t1 = await sign(kp);    await jwtVerify(t1, jwks, { issuer, audience });    expect(harness.callCount()).toBe(1);     await new Promise((r) => setTimeout(r, 20));     const t2 = await sign(kp);    await jwtVerify(t2, jwks, { issuer, audience });    expect(harness.callCount()).toBe(2);  });   it("forceRefresh() invalidates the cache so the next call refetches", async () => {    const kp = await generateRotatableKeypair("k1");    const harness = makeFetchHarness([kp.publicJwk]);    const jwks = createCachedRemoteJWKSet(jwksUrl, { fetch: harness.fetch });     const t1 = await sign(kp);    await jwtVerify(t1, jwks, { issuer, audience });    expect(harness.callCount()).toBe(1);     jwks.forceRefresh();     await jwtVerify(t1, jwks, { issuer, audience });    expect(harness.callCount()).toBe(2);  });   // -------------------------------------------------------------------------  // Cross-isolate store — the cold-isolate path that caused the 2026-08-18  // 500s: no module-scope entry, so every verify paid an upstream fetch.  // -------------------------------------------------------------------------   it("a cold isolate serves from the store instead of going upstream", async () => {    const kp = await generateRotatableKeypair("k1");    const store = makeStoreHarness();     // First isolate: cold everywhere, so it fetches and populates the store.    const warm = makeFetchHarness([kp.publicJwk]);    const first = createCachedRemoteJWKSet(jwksUrl, { fetch: warm.fetch, store });    const t1 = await sign(kp);    await jwtVerify(t1, first, { issuer, audience });    expect(warm.callCount()).toBe(1);    expect(store.writes()).toBe(1);     // A brand-new isolate (fresh module scope). If it reaches upstream at all    // the fetch fails, so verifying proves the store answered.    const second = createCachedRemoteJWKSet(jwksUrl, { fetch: failingFetch, store });    const { payload } = await jwtVerify(t1, second, { issuer, audience });    expect(payload.sub).toBe("user_test");    expect(store.reads()).toBeGreaterThan(0);  });   it("keeps serving the last good keys when the key server is down", async () => {    const kp = await generateRotatableKeypair("k1");    const harness = makeFetchHarness([kp.publicJwk]);    let upstreamUp = true;    const flaky: typeof globalThis.fetch = (...args) =>      upstreamUp ? harness.fetch(...args) : failingFetch(...args);     // ttl short enough that the next call is past it, stale window generous.    const jwks = createCachedRemoteJWKSet(jwksUrl, {      fetch: flaky,      store: null,      ttlMs: 10,      staleMaxMs: 60_000,    });     const token = await sign(kp);    await jwtVerify(token, jwks, { issuer, audience });     upstreamUp = false;    await new Promise((r) => setTimeout(r, 20));     // Past the TTL with a dead upstream: the cached keys still verify.    const { payload } = await jwtVerify(token, jwks, { issuer, audience });    expect(payload.sub).toBe("user_test");  });   it("stops serving stale keys once the stale window closes", async () => {    const kp = await generateRotatableKeypair("k1");    const harness = makeFetchHarness([kp.publicJwk]);    let upstreamUp = true;    const flaky: typeof globalThis.fetch = (...args) =>      upstreamUp ? harness.fetch(...args) : failingFetch(...args);     const jwks = createCachedRemoteJWKSet(jwksUrl, {      fetch: flaky,      store: null,      ttlMs: 5,      staleMaxMs: 10,    });     const token = await sign(kp);    await jwtVerify(token, jwks, { issuer, audience });     upstreamUp = false;    await new Promise((r) => setTimeout(r, 30));     // Beyond staleMaxMs the keys are no longer trustworthy — fail, don't    // silently honour a key set we can no longer confirm.    await expect(jwtVerify(token, jwks, { issuer, audience })).rejects.toThrow();  });   it("attributes background revalidation to fetchCount but not blockingFetchCount", async () => {    const kp = await generateRotatableKeypair("k1");    const harness = makeFetchHarness([kp.publicJwk]);    const jwks = createCachedRemoteJWKSet(jwksUrl, {      fetch: harness.fetch,      store: null,      ttlMs: 10,      staleMaxMs: 60_000,    });     const token = await sign(kp);    await jwtVerify(token, jwks, { issuer, audience });    expect(jwks.inspect().blockingFetchCount).toBe(1);     await new Promise((r) => setTimeout(r, 20));     // Past the TTL: served from the stale entry, revalidated behind it. The    // caller waited on nothing, so latency must not be attributed to it.    const before = jwks.inspect();    await jwtVerify(token, jwks, { issuer, audience });    const after = jwks.inspect();     expect(after.blockingFetchCount).toBe(before.blockingFetchCount);    expect(after.fetchCount).toBeGreaterThan(before.fetchCount);  });   it("records a cold-isolate store read as a store hit, not a fetch", async () => {    const kp = await generateRotatableKeypair("k1");    const store = makeStoreHarness();    const warm = makeFetchHarness([kp.publicJwk]);     const first = createCachedRemoteJWKSet(jwksUrl, { fetch: warm.fetch, store });    const token = await sign(kp);    await jwtVerify(token, first, { issuer, audience });     const second = createCachedRemoteJWKSet(jwksUrl, { fetch: failingFetch, store });    await jwtVerify(token, second, { issuer, audience });     const stats = second.inspect();    expect(stats.storeHitCount).toBe(1);    expect(stats.blockingFetchCount).toBe(0);  });   it("does not block a request on revalidation once keys are cached", async () => {    const kp = await generateRotatableKeypair("k1");    const harness = makeFetchHarness([kp.publicJwk]);    let hang = false;    const slowAfterFirst: typeof globalThis.fetch = async (...args) => {      if (hang) await new Promise((r) => setTimeout(r, 5_000));      return harness.fetch(...args);    };     const jwks = createCachedRemoteJWKSet(jwksUrl, {      fetch: slowAfterFirst,      store: null,      ttlMs: 10,      staleMaxMs: 60_000,    });     const token = await sign(kp);    await jwtVerify(token, jwks, { issuer, audience });     hang = true;    await new Promise((r) => setTimeout(r, 20));     // The revalidation behind this call hangs for 5s; the verify must not.    const startedAt = Date.now();    await jwtVerify(token, jwks, { issuer, audience });    expect(Date.now() - startedAt).toBeLessThan(1_000);  });});