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

.changeset/apikey-revoke-cache-invalidation.md +7 −0
apps/cloud/src/auth/api-key-validation-cache.node.test.ts +101 −8
apps/cloud/src/auth/api-keys.ts +62 −20
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
171	expect(cached.cacheKeys()).not.toContain(rawKey);	171	expect(cached.cacheKeys()).not.toContain(rawKey);
172	}),	172	}),
173	);	173	);
		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	);
174	});	207	});
175		208
		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.
176	describe("ApiKeyService.WorkOS validation cache", () => {	212	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", () =>
178	Effect.gen(function* () {	214	Effect.gen(function* () {
179	let calls = 0;	215	let calls = 0;
180	const layer = ApiKeyService.WorkOS.pipe(	216	const layer = ApiKeyService.WorkOS.pipe(
@@ -185,7 +221,7 @@ describe("ApiKeyService.WorkOS validation cache", () => { · 4 unmodified lines
185	calls += 1;	221	calls += 1;
186	return {	222	return {
187	apiKey: {	223	apiKey: {
188	id: "api_key_123",	224	id: "api_key_shared",
189	owner: { type: "user", id: "user_123", organizationId: "org_123" },	225	owner: { type: "user", id: "user_123", organizationId: "org_123" },
190	},	226	},
191	};	227	};
@@ -194,17 +230,74 @@ describe("ApiKeyService.WorkOS validation cache", () => { · 2 unmodified lines
194	),	230	),
195	);	231	);
196		232
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* () {
198	const apiKeys = yield* ApiKeyService;	237	const apiKeys = yield* ApiKeyService;
199	return {	238	return yield* apiKeys.validate("sk_test_cross_build");
200	first: yield* apiKeys.validate("sk_test_abc"),	239	}).pipe(Effect.provide(layer));
201	second: yield* apiKeys.validate("sk_test_abc"),	240	const second = yield* Effect.gen(function* () {
202	};	241	const apiKeys = yield* ApiKeyService;
		242	return yield* apiKeys.validate("sk_test_cross_build");
203	}).pipe(Effect.provide(layer));	243	}).pipe(Effect.provide(layer));
204		244
205	expect(calls).toBe(1);	245	expect(calls).toBe(1);
206	expect(first?.keyId).toBe("api_key_123");	246	expect(first?.keyId).toBe("api_key_shared");
207	expect(second).toEqual(first);	247	expect(second).toEqual(first);
208	}),	248	}),
209	);	249	);
		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	);
210	});	303	});

apps/cloud/src/auth/api-keys.ts+62−20

 //     failure always miss: an attacker probing bad keys cannot poison the map
 //     (the size bound alone handles memory), and a just-created key works
 //     immediately instead of waiting out a stale negative entry.
-//   - REVOCATION WINDOW: a revoked key keeps working for up to the TTL (60s)
+//   - REVOCATION WINDOW: the isolate that serves the revoke drops its own
-//     in any isolate that validated it before revocation. That is the price
+//     entries for that key id immediately (`invalidateKeyId`, called by both
-//     of the cache, and it is far tighter than the 1h JWKS rotation window
+//     revoke paths below), so revoking a key through the console refuses it
-//     the JWT path already accepts.
+//     on the next request wherever the same isolate validates. OTHER isolates
+//     that validated the key before revocation keep serving it for up to the
+//     TTL (60s). That residual window is the price of the cache, and it is
+//     far tighter than the 1h JWKS rotation window the JWT path already
+//     accepts.
 //   - Concurrent misses for the same key each call WorkOS; there is no
 //     in-flight dedupe. The duplicate window is one round trip per key per
 //     TTL per isolate — exactly the cost EVERY request paid before this
 // map unbounded (mirrors BALANCE_CACHE_MAX_ENTRIES in execution-gate.ts).
 const API_KEY_VALIDATION_CACHE_MAX_ENTRIES = 10_000;
+type ApiKeyValidationCacheEntry = {
+  readonly owner: ApiKeyOwner;
+  readonly expiresAtMs: number;
+};
+// The map lives at MODULE scope — one per isolate — not inside
+// `makeCachedApiKeyValidate`, because the `ApiKeyService.WorkOS` layer is NOT
+// built once per isolate on every plane: the account middleware rebuilds it
+// per request. A map owned by each build would give the revoke paths a fresh
+// empty cache to invalidate while the long-lived identity plane kept serving
+// the revoked key from its own. Every build shares this map instead (entries
+// are pure validation RESULTS, so which build's WorkOS client wrote them does
+// not matter), which is also what makes the entries survive between requests
+// on the per-request plane at all.
+const isolateApiKeyValidationCache = new Map<string, ApiKeyValidationCacheEntry>();
 export type ApiKeyValidate = (
   value: string,
 ) => Effect.Effect<ApiKeyOwner | null, ApiKeyValidationError>;
 /**
- * Wrap a `validate` function with the per-isolate success cache described
+ * Wrap a `validate` function with the success cache described above. Exported
- * above. Exported for tests only — production use is the single instance the
+ * for tests only — production use is the `ApiKeyService.WorkOS` layer below,
- * `ApiKeyService.WorkOS` layer builds, which lives in `boot` and is therefore
+ * every build of which shares the module-scope map.
- * constructed once per isolate.
  *
- * The `ttlMs` / `maxEntries` / `now` knobs exist so tests can exercise expiry
+ * The `ttlMs` / `maxEntries` / `now` / `cache` knobs exist so tests can
- * and the size bound without waiting on wall time; production passes nothing.
+ * exercise expiry, the size bound, and map sharing without wall time or
+ * module state; production passes only `cache`.
  */
 export const makeCachedApiKeyValidate = (
   validate: ApiKeyValidate,
     readonly ttlMs?: number;
     readonly maxEntries?: number;
     readonly now?: () => number;
+    /** The map to cache into. Production passes the module-scope
+     *  `isolateApiKeyValidationCache`; tests default to a private map. */
+    readonly cache?: Map<string, ApiKeyValidationCacheEntry>;
   },
 ): {
   readonly validate: ApiKeyValidate;
+  /**
+   * Drop every cached entry that resolved to this key id. The cache is keyed
+   * by the digest of the presented VALUE and a revoke only knows the key ID,
+   * so this walks the map — bounded by `maxEntries`, and revocation is a rare
+   * human action, not a hot path.
+   */
+  readonly invalidateKeyId: (keyId: string) => void;
   /** Test seam: the digests currently cached, so tests can assert the raw
    *  credential never appears as a map key. */
   readonly cacheKeys: () => ReadonlyArray<string>;
   const ttlMs = options?.ttlMs ?? API_KEY_VALIDATION_CACHE_TTL_MS;
   const maxEntries = options?.maxEntries ?? API_KEY_VALIDATION_CACHE_MAX_ENTRIES;
   const now = options?.now ?? Date.now;
-  const cache = new Map<string, { readonly owner: ApiKeyOwner; readonly expiresAtMs: number }>();
+  const cache = options?.cache ?? new Map<string, ApiKeyValidationCacheEntry>();
   const writeCache = (digest: string, owner: ApiKeyOwner, nowMs: number): void => {
     if (cache.size >= maxEntries) {
   return {
     cacheKeys: () => [...cache.keys()],
+    invalidateKeyId: (keyId) => {
+      for (const [digest, entry] of cache) {
+        if (entry.owner.keyId === keyId) cache.delete(digest);
+      }
+    },
     validate: (value) =>
       Effect.gen(function* () {
         const digest = yield* sha256Hex(value);
   static WorkOS = Layer.effect(this)(
     Effect.gen(function* () {
       const workos = yield* WorkOSClient;
-      // Built once with the boot layer, so the success cache is per-isolate.
+      // The cache MAP is the module-scope per-isolate one (see its comment);
-      const cachedValidate = makeCachedApiKeyValidate((value) =>
+      // this build only contributes the validate function that fills it, so
-        workos.validateApiKey(value).pipe(
+      // however many times a plane rebuilds this layer, hits, misses, and
-          Effect.map(ownerFromResponse),
+      // invalidations all land in the same map.
-          Effect.mapError((cause) => new ApiKeyValidationError({ cause })),
+      const cachedValidate = makeCachedApiKeyValidate(
-        ),
+        (value) =>
+          workos.validateApiKey(value).pipe(
+            Effect.map(ownerFromResponse),
+            Effect.mapError((cause) => new ApiKeyValidationError({ cause })),
+          ),
+        { cache: isolateApiKeyValidationCache },
       );
       return {
         validate: cachedValidate.validate,
             }),
           ),
         revokeUserKey: ({ keyId }) =>
-          workos
+          workos.deleteApiKey(keyId).pipe(
-            .deleteApiKey(keyId)
+            Effect.mapError((cause) => new ApiKeyManagementError({ cause })),
-            .pipe(Effect.mapError((cause) => new ApiKeyManagementError({ cause }))),
+            Effect.tap(() => Effect.sync(() => cachedValidate.invalidateKeyId(keyId))),
+          ),
         listOrgKeys: ({ organizationId }) =>
           workos.listOrgApiKeys(organizationId).pipe(
             Effect.map((response) => orgListFromResponse(response, organizationId)),
             yield* workos
               .deleteApiKey(keyId)
               .pipe(Effect.mapError((cause) => new ApiKeyManagementError({ cause })));
+            cachedValidate.invalidateKeyId(keyId);
           }),
       };
     }),

apps/cloud/src/mcp/agent-handler.ts+14−13

@@ -133,25 +133,26 @@ const authenticate = (request: Request) => · 132 unmodified lines
133		133
134	// ONE auth layer per isolate. `authenticate` used to end in	134	// ONE auth layer per isolate. `authenticate` used to end in
135	// `Effect.provide(cloudMcpAuth)`, which rebuilt the entire auth stack — the	135	// `Effect.provide(cloudMcpAuth)`, which rebuilt the entire auth stack — the
136	// WorkOS client, `ApiKeyService` and with it the api-key validation success	136	// WorkOS client and `ApiKeyService` — on every MCP request. A lazily created
137	// cache — on every MCP request, so that cache never survived a request on this	137	// ManagedRuntime builds `cloudMcpAuth` on the first request and memoizes it;
138	// plane (the /api/* plane resolves the same service from the app's `boot`	138	// `runTraced` runs every program on it. Nothing inside the layer is
139	// layer, built once per isolate). A lazily created ManagedRuntime builds	139	// request-scoped: the WorkOS client holds no sockets (just config), the JWKS
140	// `cloudMcpAuth` on the first request and memoizes it; `runTraced` runs every	140	// cache is already module-scope, and `McpOrganizationAuthLive` builds its
141	// program on it. Nothing inside the layer is request-scoped: the WorkOS client	141	// postgres socket FRESH inside every `authorize` call precisely so the
142	// holds no sockets (just config), the JWKS cache is already module-scope, and	142	// service itself can outlive a request (Cloudflare Workers' I/O isolation —
143	// `McpOrganizationAuthLive` builds its postgres socket FRESH inside every	143	// see makeMcpOrganizationAuthServices).
144	// `authorize` call precisely so the service itself can outlive a request
145	// (Cloudflare Workers' I/O isolation — see makeMcpOrganizationAuthServices).
146	//	144	//
147	// The memo is keyed on the WorkOS binding values the layer captures at build,	145	// The memo is keyed on the WorkOS binding values the layer captures at build,
148	// NOT held forever: Cloudflare reuses warm isolates across binding-only	146	// NOT held forever: Cloudflare reuses warm isolates across binding-only
149	// deployments, so a bare `??=` would keep authenticating with a rotated-out	147	// deployments, so a bare `??=` would keep authenticating with a rotated-out
150	// WORKOS_API_KEY until isolate eviction. `runTraced` fingerprints the	148	// WORKOS_API_KEY until isolate eviction. `runTraced` fingerprints the
151	// request's own `env` (a string compare per request) and the cache swaps in a	149	// 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	150	// fresh runtime when the bindings changed. The api-key validation success
153	// key must not serve cached validations — when the bindings changed. See	151	// cache is NOT swapped with it — it is module-scope in api-keys.ts (shared
154	// auth-runtime.ts for the drop-vs-dispose reasoning.	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.
155	const mcpAuthRuntimeFor: (	156	const mcpAuthRuntimeFor: (
156	fingerprint: string,	157	fingerprint: string,
157	) => ManagedRuntime.ManagedRuntime<McpAuthProvider, never> = makeBindingKeyedRuntime(() =>	158	) => ManagedRuntime.ManagedRuntime<McpAuthProvider, never> = makeBindingKeyedRuntime(() =>