History

Invalidate cached api key validations on revoke

Rhys Sullivan committed 2 months ago

Commit 50ee2fd1 parent 8e2515c4 files · +184 −41Browse files

Unified · Split

4 files
  1. .changeset/apikey-revoke-cache-invalidation.md +7 −0
  2. apps/cloud/src/auth/api-key-validation-cache.node.test.ts +101 −8
  3. apps/cloud/src/auth/api-keys.ts +62 −20
  4. apps/cloud/src/mcp/agent-handler.ts +14 −13
.changeset/apikey-revoke-cache-invalidation.md+7−0
@@ -0,0 +1,7 @@
1+---
2+"@executor-js/cloud": patch
3+---
4+
5+**Revoking an API key now invalidates its cached validation**
6+
7+The per-isolate validation cache made a revoked key keep authenticating for up to 60 seconds. Two fixes close that: the cache map now lives at module scope, shared by every build of the key service (the account middleware rebuilds it per request, so a per-build map left the revoke invalidating a fresh empty cache), and both revoke paths drop the revoked key's entries from it. A revoked key is refused on the next request served by the isolate that processed the revoke; other isolates still age the entry out within the TTL.
apps/cloud/src/auth/api-key-validation-cache.node.test.ts+101−8
@@ -171,10 +171,46 @@ describe("makeCachedApiKeyValidate", () => { · 170 unmodified lines
171171 expect(cached.cacheKeys()).not.toContain(rawKey);
172172 }),
173173 );
174+
175+ it.effect("invalidating a key id drops its entries and forces revalidation", () =>
176+ Effect.gen(function* () {
177+ let calls = 0;
178+ const cached = makeCachedApiKeyValidate(() =>
179+ Effect.sync(() => {
180+ calls += 1;
181+ return userOwner("api_key_1");
182+ }),
183+ );
184+
185+ yield* cached.validate("sk_test_abc");
186+ cached.invalidateKeyId("api_key_1");
187+ expect(cached.cacheKeys()).toHaveLength(0);
188+
189+ yield* cached.validate("sk_test_abc");
190+ expect(calls).toBe(2);
191+ }),
192+ );
193+
194+ it.effect("invalidating one key id leaves other keys' entries cached", () =>
195+ Effect.gen(function* () {
196+ const cached = makeCachedApiKeyValidate((value) =>
197+ Effect.succeed(userOwner(`key_for_${value}`)),
198+ );
199+
200+ yield* cached.validate("sk_1");
201+ yield* cached.validate("sk_2");
202+ cached.invalidateKeyId("key_for_sk_1");
203+
204+ expect(cached.cacheKeys()).toHaveLength(1);
205+ }),
206+ );
174207 });
175208
209+// These tests build the REAL WorkOS layer, whose cache map is module-scope —
210+// entries written by one test are visible to the next within this file. Each
211+// test therefore uses key values of its own; none may reuse another's.
176212 describe("ApiKeyService.WorkOS validation cache", () => {
177− it.effect("validates a repeated key through one upstream call per layer build", () =>
213+ it.effect("serves a key cached by one layer build from a second build", () =>
178214 Effect.gen(function* () {
179215 let calls = 0;
180216 const layer = ApiKeyService.WorkOS.pipe(
@@ -185,7 +221,7 @@ describe("ApiKeyService.WorkOS validation cache", () => { · 4 unmodified lines
185221 calls += 1;
186222 return {
187223 apiKey: {
188− id: "api_key_123",
224+ id: "api_key_shared",
189225 owner: { type: "user", id: "user_123", organizationId: "org_123" },
190226 },
191227 };
@@ -194,17 +230,74 @@ describe("ApiKeyService.WorkOS validation cache", () => { · 2 unmodified lines
194230 ),
195231 );
196232
197− const { first, second } = yield* Effect.gen(function* () {
233+ // Two separate `Effect.provide(layer)` runs are two separate layer
234+ // builds — the account middleware rebuilds this layer per request, so
235+ // cross-build reuse is exactly what production needs the cache to do.
236+ const first = yield* Effect.gen(function* () {
198237 const apiKeys = yield* ApiKeyService;
199− return {
200− first: yield* apiKeys.validate("sk_test_abc"),
201− second: yield* apiKeys.validate("sk_test_abc"),
202− };
238+ return yield* apiKeys.validate("sk_test_cross_build");
239+ }).pipe(Effect.provide(layer));
240+ const second = yield* Effect.gen(function* () {
241+ const apiKeys = yield* ApiKeyService;
242+ return yield* apiKeys.validate("sk_test_cross_build");
203243 }).pipe(Effect.provide(layer));
204244
205245 expect(calls).toBe(1);
206− expect(first?.keyId).toBe("api_key_123");
246+ expect(first?.keyId).toBe("api_key_shared");
207247 expect(second).toEqual(first);
208248 }),
209249 );
250+
251+ it.effect("a revoke in one layer build refuses the cached value in every build", () =>
252+ Effect.gen(function* () {
253+ // Live until revoked: the stub answers with the key while `revoked` is
254+ // false and with a null apiKey afterwards, the same flip the WorkOS
255+ // backend performs. The revoke runs in its OWN layer build — the shape
256+ // of the real console flow, where the account middleware rebuilds the
257+ // layer per request — and must still drop the entry the validating
258+ // plane cached.
259+ let revoked = false;
260+ let calls = 0;
261+ const layer = ApiKeyService.WorkOS.pipe(
262+ Layer.provide(
263+ stubWorkOS({
264+ validateApiKey: () =>
265+ Effect.sync(() => {
266+ calls += 1;
267+ return {
268+ apiKey: revoked
269+ ? null
270+ : {
271+ id: "api_key_revoked",
272+ owner: { type: "user", id: "user_123", organizationId: "org_123" },
273+ },
274+ };
275+ }),
276+ deleteApiKey: () =>
277+ Effect.sync(() => {
278+ revoked = true;
279+ return {};
280+ }),
281+ }),
282+ ),
283+ );
284+
285+ const before = yield* Effect.gen(function* () {
286+ const apiKeys = yield* ApiKeyService;
287+ return yield* apiKeys.validate("sk_test_revocable");
288+ }).pipe(Effect.provide(layer));
289+ yield* Effect.gen(function* () {
290+ const apiKeys = yield* ApiKeyService;
291+ return yield* apiKeys.revokeUserKey({ keyId: "api_key_revoked" });
292+ }).pipe(Effect.provide(layer));
293+ const after = yield* Effect.gen(function* () {
294+ const apiKeys = yield* ApiKeyService;
295+ return yield* apiKeys.validate("sk_test_revocable");
296+ }).pipe(Effect.provide(layer));
297+
298+ expect(before?.keyId).toBe("api_key_revoked");
299+ expect(after, "the revoke dropped the cached entry, not just the upstream key").toBeNull();
300+ expect(calls).toBe(2);
301+ }),
302+ );
210303 });
apps/cloud/src/auth/api-keys.ts+62−20
@@ -296,10 +296,14 @@ const createdFromResponse = (value: unknown): CreatedApiKey | null => · 295 unmodified lines
296296 // failure always miss: an attacker probing bad keys cannot poison the map
297297 // (the size bound alone handles memory), and a just-created key works
298298 // immediately instead of waiting out a stale negative entry.
299−// - REVOCATION WINDOW: a revoked key keeps working for up to the TTL (60s)
300−// in any isolate that validated it before revocation. That is the price
301−// of the cache, and it is far tighter than the 1h JWKS rotation window
302−// the JWT path already accepts.
299+// - REVOCATION WINDOW: the isolate that serves the revoke drops its own
300+// entries for that key id immediately (`invalidateKeyId`, called by both
301+// revoke paths below), so revoking a key through the console refuses it
302+// on the next request wherever the same isolate validates. OTHER isolates
303+// that validated the key before revocation keep serving it for up to the
304+// TTL (60s). That residual window is the price of the cache, and it is
305+// far tighter than the 1h JWKS rotation window the JWT path already
306+// accepts.
303307 // - Concurrent misses for the same key each call WorkOS; there is no
304308 // in-flight dedupe. The duplicate window is one round trip per key per
305309 // TTL per isolate — exactly the cost EVERY request paid before this
@@ -312,18 +316,34 @@ const API_KEY_VALIDATION_CACHE_TTL_MS = 60_000; · 6 unmodified lines
312316 // map unbounded (mirrors BALANCE_CACHE_MAX_ENTRIES in execution-gate.ts).
313317 const API_KEY_VALIDATION_CACHE_MAX_ENTRIES = 10_000;
314318
319+type ApiKeyValidationCacheEntry = {
320+ readonly owner: ApiKeyOwner;
321+ readonly expiresAtMs: number;
322+};
323+
324+// The map lives at MODULE scope — one per isolate — not inside
325+// `makeCachedApiKeyValidate`, because the `ApiKeyService.WorkOS` layer is NOT
326+// built once per isolate on every plane: the account middleware rebuilds it
327+// per request. A map owned by each build would give the revoke paths a fresh
328+// empty cache to invalidate while the long-lived identity plane kept serving
329+// the revoked key from its own. Every build shares this map instead (entries
330+// are pure validation RESULTS, so which build's WorkOS client wrote them does
331+// not matter), which is also what makes the entries survive between requests
332+// on the per-request plane at all.
333+const isolateApiKeyValidationCache = new Map<string, ApiKeyValidationCacheEntry>();
334+
315335 export type ApiKeyValidate = (
316336 value: string,
317337 ) => Effect.Effect<ApiKeyOwner | null, ApiKeyValidationError>;
318338
319339 /**
320− * Wrap a `validate` function with the per-isolate success cache described
321− * above. Exported for tests only — production use is the single instance the
322− * `ApiKeyService.WorkOS` layer builds, which lives in `boot` and is therefore
323− * constructed once per isolate.
340+ * Wrap a `validate` function with the success cache described above. Exported
341+ * for tests only — production use is the `ApiKeyService.WorkOS` layer below,
342+ * every build of which shares the module-scope map.
324343 *
325− * The `ttlMs` / `maxEntries` / `now` knobs exist so tests can exercise expiry
326− * and the size bound without waiting on wall time; production passes nothing.
344+ * The `ttlMs` / `maxEntries` / `now` / `cache` knobs exist so tests can
345+ * exercise expiry, the size bound, and map sharing without wall time or
346+ * module state; production passes only `cache`.
327347 */
328348 export const makeCachedApiKeyValidate = (
329349 validate: ApiKeyValidate,
@@ -331,9 +351,19 @@ export const makeCachedApiKeyValidate = ( · 1 unmodified line
331351 readonly ttlMs?: number;
332352 readonly maxEntries?: number;
333353 readonly now?: () => number;
354+ /** The map to cache into. Production passes the module-scope
355+ * `isolateApiKeyValidationCache`; tests default to a private map. */
356+ readonly cache?: Map<string, ApiKeyValidationCacheEntry>;
334357 },
335358 ): {
336359 readonly validate: ApiKeyValidate;
360+ /**
361+ * Drop every cached entry that resolved to this key id. The cache is keyed
362+ * by the digest of the presented VALUE and a revoke only knows the key ID,
363+ * so this walks the map — bounded by `maxEntries`, and revocation is a rare
364+ * human action, not a hot path.
365+ */
366+ readonly invalidateKeyId: (keyId: string) => void;
337367 /** Test seam: the digests currently cached, so tests can assert the raw
338368 * credential never appears as a map key. */
339369 readonly cacheKeys: () => ReadonlyArray<string>;
@@ -341,7 +371,7 @@ export const makeCachedApiKeyValidate = ( · 1 unmodified line
341371 const ttlMs = options?.ttlMs ?? API_KEY_VALIDATION_CACHE_TTL_MS;
342372 const maxEntries = options?.maxEntries ?? API_KEY_VALIDATION_CACHE_MAX_ENTRIES;
343373 const now = options?.now ?? Date.now;
344− const cache = new Map<string, { readonly owner: ApiKeyOwner; readonly expiresAtMs: number }>();
374+ const cache = options?.cache ?? new Map<string, ApiKeyValidationCacheEntry>();
345375
346376 const writeCache = (digest: string, owner: ApiKeyOwner, nowMs: number): void => {
347377 if (cache.size >= maxEntries) {
@@ -356,6 +386,11 @@ export const makeCachedApiKeyValidate = ( · 8 unmodified lines
356386
357387 return {
358388 cacheKeys: () => [...cache.keys()],
389+ invalidateKeyId: (keyId) => {
390+ for (const [digest, entry] of cache) {
391+ if (entry.owner.keyId === keyId) cache.delete(digest);
392+ }
393+ },
359394 validate: (value) =>
360395 Effect.gen(function* () {
361396 const digest = yield* sha256Hex(value);
@@ -414,12 +449,17 @@ export class ApiKeyService extends Context.Service< · 52 unmodified lines
414449 static WorkOS = Layer.effect(this)(
415450 Effect.gen(function* () {
416451 const workos = yield* WorkOSClient;
417− // Built once with the boot layer, so the success cache is per-isolate.
418− const cachedValidate = makeCachedApiKeyValidate((value) =>
419− workos.validateApiKey(value).pipe(
420− Effect.map(ownerFromResponse),
421− Effect.mapError((cause) => new ApiKeyValidationError({ cause })),
422− ),
452+ // The cache MAP is the module-scope per-isolate one (see its comment);
453+ // this build only contributes the validate function that fills it, so
454+ // however many times a plane rebuilds this layer, hits, misses, and
455+ // invalidations all land in the same map.
456+ const cachedValidate = makeCachedApiKeyValidate(
457+ (value) =>
458+ workos.validateApiKey(value).pipe(
459+ Effect.map(ownerFromResponse),
460+ Effect.mapError((cause) => new ApiKeyValidationError({ cause })),
461+ ),
462+ { cache: isolateApiKeyValidationCache },
423463 );
424464 return {
425465 validate: cachedValidate.validate,
@@ -439,9 +479,10 @@ export class ApiKeyService extends Context.Service< · 13 unmodified lines
439479 }),
440480 ),
441481 revokeUserKey: ({ keyId }) =>
442− workos
443− .deleteApiKey(keyId)
444− .pipe(Effect.mapError((cause) => new ApiKeyManagementError({ cause }))),
482+ workos.deleteApiKey(keyId).pipe(
483+ Effect.mapError((cause) => new ApiKeyManagementError({ cause })),
484+ Effect.tap(() => Effect.sync(() => cachedValidate.invalidateKeyId(keyId))),
485+ ),
445486 listOrgKeys: ({ organizationId }) =>
446487 workos.listOrgApiKeys(organizationId).pipe(
447488 Effect.map((response) => orgListFromResponse(response, organizationId)),
@@ -475,6 +516,7 @@ export class ApiKeyService extends Context.Service< · 27 unmodified lines
475516 yield* workos
476517 .deleteApiKey(keyId)
477518 .pipe(Effect.mapError((cause) => new ApiKeyManagementError({ cause })));
519+ cachedValidate.invalidateKeyId(keyId);
478520 }),
479521 };
480522 }),
apps/cloud/src/mcp/agent-handler.ts+14−13
@@ -133,25 +133,26 @@ const authenticate = (request: Request) => · 132 unmodified lines
133133
134134 // ONE auth layer per isolate. `authenticate` used to end in
135135 // `Effect.provide(cloudMcpAuth)`, which rebuilt the entire auth stack — the
136−// WorkOS client, `ApiKeyService` and with it the api-key validation success
137−// cache — on every MCP request, so that cache never survived a request on this
138−// plane (the /api/* plane resolves the same service from the app's `boot`
139−// layer, built once per isolate). A lazily created ManagedRuntime builds
140−// `cloudMcpAuth` on the first request and memoizes it; `runTraced` runs every
141−// program on it. Nothing inside the layer is request-scoped: the WorkOS client
142−// holds no sockets (just config), the JWKS cache is already module-scope, and
143−// `McpOrganizationAuthLive` builds its postgres socket FRESH inside every
144−// `authorize` call precisely so the service itself can outlive a request
145−// (Cloudflare Workers' I/O isolation — see makeMcpOrganizationAuthServices).
136+// WorkOS client and `ApiKeyService` — on every MCP request. A lazily created
137+// ManagedRuntime builds `cloudMcpAuth` on the first request and memoizes it;
138+// `runTraced` runs every program on it. Nothing inside the layer is
139+// request-scoped: the WorkOS client holds no sockets (just config), the JWKS
140+// cache is already module-scope, and `McpOrganizationAuthLive` builds its
141+// postgres socket FRESH inside every `authorize` call precisely so the
142+// service itself can outlive a request (Cloudflare Workers' I/O isolation —
143+// see makeMcpOrganizationAuthServices).
146144 //
147145 // The memo is keyed on the WorkOS binding values the layer captures at build,
148146 // NOT held forever: Cloudflare reuses warm isolates across binding-only
149147 // deployments, so a bare `??=` would keep authenticating with a rotated-out
150148 // WORKOS_API_KEY until isolate eviction. `runTraced` fingerprints the
151149 // request's own `env` (a string compare per request) and the cache swaps in a
152−// fresh runtime — fresh ApiKeyService validation cache included, as a rotated
153−// key must not serve cached validations — when the bindings changed. See
154−// auth-runtime.ts for the drop-vs-dispose reasoning.
150+// fresh runtime when the bindings changed. The api-key validation success
151+// cache is NOT swapped with it — it is module-scope in api-keys.ts (shared
152+// with the /api/* plane so revocation can invalidate it everywhere), so after
153+// a WORKOS_API_KEY rotation its remaining entries age out within the 60s TTL
154+// rather than dropping instantly. See auth-runtime.ts for the drop-vs-dispose
155+// reasoning.
155156 const mcpAuthRuntimeFor: (
156157 fingerprint: string,
157158 ) => ManagedRuntime.ManagedRuntime<McpAuthProvider, never> = makeBindingKeyedRuntime(() =>