feat: exact in-flight concurrency cap via a Durable Object (#4)

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.

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);
+  });
+});

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<string, number>();
+  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<Response> {
+    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 });
+  }
+}

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
+        }
+      }
     }
   },
 };

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"]