From 07194e6da437efa395d0d8daf47566c453209624 Mon Sep 17 00:00:00 2001 From: lex00 <121451605+lex00@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:57:38 -0600 Subject: [PATCH] feat: exact in-flight concurrency cap via a Durable Object (#4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The KV breaker bounds rate (audits per minute); nothing bounded concurrency, so a burst that fit the per-minute budgets could still pile simultaneous tree-walks onto the shared git token. One ConcurrencyGate DO (idFromName("global")) now counts in-flight audits exactly and sheds with 429 past MAX_IN_FLIGHT (6). Advisory by design, per the issue's fallback criterion: a DO error or missing binding never blocks an audit — only an explicit at-capacity answer sheds. The slot ledger is in-memory (a DO is single-threaded, so it is exact while the object lives; eviction resets to all-free, the harmless direction), leaked slots self-heal after STALE_MS, and release is idempotent. SQLite-backed class in the migration — the only kind new migrations may create on the free plan; it uses no storage. Verified on workerd via wrangler dev --local in fixture mode: binding resolves, audits flow through acquire/release. Closes #4 Co-Authored-By: Claude Fable 5 --- README.md | 6 ++++ src/gate.test.ts | 78 ++++++++++++++++++++++++++++++++++++++++++ src/gate.ts | 89 ++++++++++++++++++++++++++++++++++++++++++++++++ src/handler.ts | 38 +++++++++++++++++++++ wrangler.toml | 12 +++++++ 5 files changed, 223 insertions(+) create mode 100644 src/gate.test.ts create mode 100644 src/gate.ts diff --git a/README.md b/README.md index 8130213..7b694ea 100644 --- a/README.md +++ b/README.md @@ -223,6 +223,12 @@ All edge-side; the audit engine adds the SSRF base (chant `fetch.ts`). - **Bot gate** (`src/turnstile.ts`): when `TURNSTILE_SECRET` is set, every audit requires a valid Cloudflare Turnstile token (verified server-side, fail-closed). The SPA renders the widget when `VITE_TURNSTILE_SITEKEY` is set. +- **Concurrency cap** (`src/gate.ts`): one Durable Object counts audits in + flight exactly and sheds with `429` past `MAX_IN_FLIGHT` — a burst that fits + the per-minute budgets still can't pile simultaneous tree-walks onto the + shared git token. Advisory: if the DO is unavailable the audit proceeds (the + KV breaker still stands). Leaked slots self-heal after `STALE_MS`. Tune both + in `src/gate.ts`. - **CORS**: same-origin by default (the SPA is served by this Worker), so other sites' browsers can't call `/audit`. Set `ALLOWED_ORIGIN` only if the SPA is ever hosted on a different origin. diff --git a/src/gate.test.ts b/src/gate.test.ts new file mode 100644 index 0000000..fb0e2c6 --- /dev/null +++ b/src/gate.test.ts @@ -0,0 +1,78 @@ +import { describe, test, expect } from "vitest"; +import { ConcurrencySlots, ConcurrencyGate, MAX_IN_FLIGHT, STALE_MS } from "./gate"; + +describe("ConcurrencySlots (#4)", () => { + test("hands out slots up to the cap, then refuses", () => { + const slots = new ConcurrencySlots(2); + const a = slots.acquire(0); + const b = slots.acquire(0); + expect(a).not.toBeNull(); + expect(b).not.toBeNull(); + expect(a).not.toBe(b); + expect(slots.acquire(0)).toBeNull(); + expect(slots.inFlight).toBe(2); + }); + + test("release frees a slot for the next acquire", () => { + const slots = new ConcurrencySlots(1); + const a = slots.acquire(0)!; + expect(slots.acquire(0)).toBeNull(); + slots.release(a); + expect(slots.acquire(0)).not.toBeNull(); + }); + + test("release is idempotent and ignores unknown tokens", () => { + const slots = new ConcurrencySlots(1); + const a = slots.acquire(0)!; + slots.release(a); + slots.release(a); + slots.release("never-issued"); + expect(slots.inFlight).toBe(0); + }); + + test("a leaked slot (worker died before releasing) is reclaimed after STALE_MS", () => { + const slots = new ConcurrencySlots(1, 1000); + slots.acquire(0); + // Still held inside the window… + expect(slots.acquire(500)).toBeNull(); + // …reclaimed past it, so the gate never wedges shut. + expect(slots.acquire(1501)).not.toBeNull(); + }); + + test("defaults are the documented knobs", () => { + expect(MAX_IN_FLIGHT).toBeGreaterThan(0); + expect(STALE_MS).toBeGreaterThan(60_000); + }); +}); + +describe("ConcurrencyGate DO surface", () => { + const acquire = (gate: ConcurrencyGate) => gate.fetch(new Request("https://gate/acquire", { method: "POST" })); + const release = (gate: ConcurrencyGate, token: string) => + gate.fetch(new Request("https://gate/release", { method: "POST", body: JSON.stringify({ token }) })); + + test("acquire returns a token; at capacity it sheds with 429; release reopens", async () => { + const gate = new ConcurrencyGate(); + const tokens: string[] = []; + for (let i = 0; i < MAX_IN_FLIGHT; i++) { + const res = await acquire(gate); + expect(res.status).toBe(200); + tokens.push(((await res.json()) as { token: string }).token); + } + expect((await acquire(gate)).status).toBe(429); + + expect((await release(gate, tokens[0])).status).toBe(200); + expect((await acquire(gate)).status).toBe(200); + }); + + test("release with a garbage body is a tolerated no-op", async () => { + const gate = new ConcurrencyGate(); + const res = await gate.fetch(new Request("https://gate/release", { method: "POST", body: "not json" })); + expect(res.status).toBe(200); + }); + + test("unknown routes 404", async () => { + const gate = new ConcurrencyGate(); + expect((await gate.fetch(new Request("https://gate/acquire", { method: "GET" }))).status).toBe(404); + expect((await gate.fetch(new Request("https://gate/other", { method: "POST" }))).status).toBe(404); + }); +}); diff --git a/src/gate.ts b/src/gate.ts new file mode 100644 index 0000000..56f423d --- /dev/null +++ b/src/gate.ts @@ -0,0 +1,89 @@ +/** + * Exact in-flight concurrency cap (#4) — the Durable Object upgrade the rate + * limiter's doc block promised. KV windows bound *rate* (audits per minute); + * this bounds *concurrency* (audits running right now), so a burst that fits + * the per-minute budgets still can't pile N simultaneous tree-walks onto the + * shared git token. One DO instance (`idFromName("global")`) is the single + * point of truth; the worker acquires a slot before auditing and releases it + * after, and treats the gate as advisory — if the DO errors or the binding is + * absent, the audit proceeds (defense-in-depth, not a single point of failure; + * the KV breaker still stands). + */ + +/** Concurrent audits allowed before shedding with 429. Tune here. */ +export const MAX_IN_FLIGHT = 6; + +/** + * A slot held longer than this is presumed leaked (the worker died between + * acquire and release) and is reclaimed. Generously above the worst honest + * audit: the fetch layer caps each request at 10s and the whole walk is + * bounded, so a two-minute-old slot is not a running audit. + */ +export const STALE_MS = 120_000; + +/** + * The slot ledger, pure and clock-injected so it unit-tests without a DO + * runtime. In-memory only: a Durable Object is single-threaded, so this is + * exact while the object lives, and an eviction resets to "all slots free" — + * the harmless direction. + */ +export class ConcurrencySlots { + private readonly held = new Map(); + private seq = 0; + + constructor( + private readonly maxInFlight: number = MAX_IN_FLIGHT, + private readonly staleMs: number = STALE_MS, + ) {} + + /** Reclaim slots whose holder never released (crashed worker). */ + private sweep(now: number): void { + for (const [token, at] of this.held) { + if (now - at > this.staleMs) this.held.delete(token); + } + } + + /** A slot token when one is free, `null` when the cap is reached. */ + acquire(now: number): string | null { + this.sweep(now); + if (this.held.size >= this.maxInFlight) return null; + const token = `s${++this.seq}`; + this.held.set(token, now); + return token; + } + + /** Idempotent — releasing an unknown or already-released token is a no-op. */ + release(token: string): void { + this.held.delete(token); + } + + get inFlight(): number { + return this.held.size; + } +} + +/** + * The Durable Object wrapper: POST /acquire → `{ token }` (200) or 429 when + * full; POST /release `{ token }` → 200. Anything else 404. State is the + * in-memory ledger above — no storage API, nothing persisted. + */ +export class ConcurrencyGate { + private readonly slots = new ConcurrencySlots(); + + async fetch(req: Request): Promise { + const path = new URL(req.url).pathname; + if (req.method === "POST" && path === "/acquire") { + const token = this.slots.acquire(Date.now()); + if (token === null) { + return new Response(JSON.stringify({ error: "At capacity" }), { status: 429 }); + } + return new Response(JSON.stringify({ token }), { status: 200 }); + } + if (req.method === "POST" && path === "/release") { + const body = (await req.json().catch(() => ({}))) as { token?: string }; + if (body.token) this.slots.release(body.token); + return new Response(JSON.stringify({ ok: true }), { status: 200 }); + } + return new Response(null, { status: 404 }); + } +} diff --git a/src/handler.ts b/src/handler.ts index f861d5e..07846c0 100644 --- a/src/handler.ts +++ b/src/handler.ts @@ -21,6 +21,11 @@ import { checkLimits } from "./limit"; import { verifyTurnstile } from "./turnstile"; import { bumpStats, readStats } from "./stats"; import { uaFetch } from "./ua-fetch"; +import { ConcurrencyGate } from "./gate"; + +// The Durable Object class must be exported from the worker entry for the +// binding/migration to resolve (#4). +export { ConcurrencyGate }; import { detectTemplate as detectK8s } from "@intentius/chant-lexicon-k8s/detect"; import { detectTemplate as detectDocker } from "@intentius/chant-lexicon-docker/detect"; @@ -111,6 +116,8 @@ interface Env { GIT_TOKEN?: string; /** Anonymous audit counter + rate-limit windows (KV). Rate limiting is on when bound. */ STATS?: KVNamespace; + /** Exact in-flight concurrency cap (#4). Advisory: unavailable ⇒ audits proceed. */ + GATE?: DurableObjectNamespace; /** Turnstile secret — when set, every audit requires a valid challenge token. */ TURNSTILE_SECRET?: string; /** Comma-separated allowed origins for cross-origin calls. Unset = same-origin @@ -186,6 +193,27 @@ export default { } } + // Exact in-flight cap (#4) — a burst that fits the per-minute budgets + // still can't pile simultaneous tree-walks onto the shared git token. + // Advisory by design: a DO error or missing binding never blocks an audit + // (the KV breaker above still stands); only an explicit "at capacity" + // sheds, with 429. + let gate: DurableObjectStub | undefined; + let gateToken: string | undefined; + if (env.GATE) { + try { + gate = env.GATE.get(env.GATE.idFromName("global")); + const res = await gate.fetch("https://gate/acquire", { method: "POST" }); + if (res.status === 429) { + console.warn(`reject reason=concurrency ip=${ip}`); + return json({ error: "Busy — too many audits running. Try again shortly." }, 429, { "retry-after": "10" }); + } + if (res.ok) gateToken = ((await res.json()) as { token?: string }).token; + } catch { + gate = undefined; // gate unavailable — proceed uncapped + } + } + try { const fetchImpl = env.BLACKLIGHT_FIXTURE === "1" ? (await import("./fixture")).fixtureFetch() : uaFetch; const files = await fetchRepoFiles(target, { token: env.GIT_TOKEN, fetchImpl }); @@ -213,6 +241,16 @@ export default { const msg = err instanceof Error ? err.message : String(err); const status = err instanceof FetchError ? 400 : 502; return json({ error: msg }, status); + } finally { + // Free the slot on every path; a lost release self-heals via the DO's + // stale-slot sweep, so failure here is tolerable and never surfaced. + if (gate && gateToken) { + try { + await gate.fetch("https://gate/release", { method: "POST", body: JSON.stringify({ token: gateToken }) }); + } catch { + // reclaimed by the sweep + } + } } }, }; diff --git a/wrangler.toml b/wrangler.toml index 8e052f6..6743374 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -36,3 +36,15 @@ directory = "./web/dist" [[kv_namespaces]] binding = "STATS" id = "a9d537a35aab42b19acd1f22787535a0" + +# Exact in-flight concurrency cap (#4): one Durable Object instance sheds with +# 429 when MAX_IN_FLIGHT audits are already running. Advisory — the handler +# proceeds uncapped if the DO is unavailable. SQLite-backed class (the only +# kind new migrations may create on the free plan); it uses no storage. +[[durable_objects.bindings]] +name = "GATE" +class_name = "ConcurrencyGate" + +[[migrations]] +tag = "v1" +new_sqlite_classes = ["ConcurrencyGate"]