Diffinery

codician/diffinerypublic

Open-source, agent-first GitHub and GitHub Actions alternative that runs on your own Cloudflare account: git hosting, pull requests, code review and drop-in GitHub Actions YAML on Workers, Durable Objects, R2 and Containers.

b19a33e

Branches

claude/community-replies claude/connector-bad-gateway-doc claude/early-access-cta claude/github-discoverability claude/readme-reorganize claude/readme-screenshots-up claude/showcase-agents-plugins community-replies lists-archived main marketplace-builds mcp-registry-auth mcp-registry-docs plugins-any-instance privacy-page

Tags

v0.1.0
diffinery/apps/diffinery/test/unit/rate-limit.test.ts
142 lines · 8.1 KBRaw History
1import { readFileSync } from "node:fs";
2import { describe, expect, it } from "vitest";
3import type { Env } from "../../src/env.js";
4import type { Principal } from "../../src/auth/principal.js";
5import { bucketsFor, MODEL_RATE_LIMIT, modelRateLimited, overLimit, RATE_LIMITS, RATE_PERIOD_S, type RateLimiter } from "../../src/lib/rate-limit.js";
6
7const anon: Principal = { type: "anonymous" };
8const person = { type: "account", account: { id: "u1" }, via: "token", scopes: ["write"], repoIds: null } as unknown as Principal;
9const job = { type: "job", repoId: "r1" } as unknown as Principal;
10
11/** A limiter that allows `n` calls per key. */
12function limiter(n: number): RateLimiter & { seen: string[] } {
13 const counts = new Map<string, number>();
14 const seen: string[] = [];
15 return {
16 seen,
17 async limit({ key }) {
18 seen.push(key);
19 const c = (counts.get(key) ?? 0) + 1;
20 counts.set(key, c);
21 return { success: c <= n };
22 },
23 };
24}
25
26describe("rate limit buckets", () => {
27 it("counts sign-in attempts per IP address, and other anonymous requests separately", () => {
28 expect(bucketsFor("POST", "/auth/passkey/login/options", anon, "203.0.113.9")).toEqual([["RL_AUTH", "ip:203.0.113.9"]]);
29 expect(bucketsFor("POST", "/invite/abc/email", anon, "203.0.113.9")).toEqual([["RL_AUTH", "ip:203.0.113.9"]]);
30 expect(bucketsFor("POST", "/signup", anon, "203.0.113.9")).toEqual([["RL_AUTH", "ip:203.0.113.9"]]);
31 expect(bucketsFor("POST", "/signup/verify", anon, "203.0.113.9")).toEqual([["RL_AUTH", "ip:203.0.113.9"]]);
32 // Security reports, signed in or not, count per IP address like sign-up.
33 expect(bucketsFor("POST", "/security/report", anon, "203.0.113.9")).toEqual([["RL_AUTH", "ip:203.0.113.9"]]);
34 for (const path of ["/early-access", "/early-access/confirm", "/early-access/remove", "/early-access/remove/confirm"]) expect(bucketsFor("POST", path, anon, "203.0.113.9"), path).toEqual([["RL_AUTH", "ip:203.0.113.9"]]);
35 expect(bucketsFor("GET", "/early-access", anon, "203.0.113.9")).toEqual([["RL_ANON", "ip:203.0.113.9"]]);
36 expect(bucketsFor("POST", "/security/report", person, "203.0.113.9")).toEqual([["RL_AUTH", "ip:203.0.113.9"]]);
37 expect(bucketsFor("GET", "/security/report", anon, "203.0.113.9")).toEqual([["RL_ANON", "ip:203.0.113.9"]]);
38 expect(bucketsFor("GET", "/signup", anon, "203.0.113.9")).toEqual([["RL_ANON", "ip:203.0.113.9"]]);
39 expect(bucketsFor("GET", "/login", anon, "203.0.113.9")).toEqual([["RL_ANON", "ip:203.0.113.9"]]);
40 expect(bucketsFor("GET", "/api/v3/repos/acme/app", anon, "203.0.113.9")).toEqual([["RL_ANON", "ip:203.0.113.9"]]);
41 expect(bucketsFor("GET", "/acme/app.git/info/refs", anon, "198.51.100.1")).toEqual([["RL_ANON", "ip:198.51.100.1"]]);
42 });
43
44 it("counts a signed-in account's requests, and its new issues, comments and repositories again", () => {
45 expect(bucketsFor("GET", "/api/v3/user/repos", person, "x")).toEqual([["RL_ACCOUNT", "acct:u1"]]);
46 for (const [m, path] of [
47 ["POST", "/api/v3/user/repos"],
48 ["POST", "/api/v3/orgs/acme/repos"],
49 ["POST", "/api/v3/user/orgs"],
50 ["POST", "/api/v3/repos/acme/app/issues"],
51 ["POST", "/api/v3/repos/acme/app/issues/3/comments"],
52 ["POST", "/api/v3/repos/acme/app/pulls/4/reviews"],
53 ["POST", "/api/v3/repos/acme/app/pulls/4/comments/9/replies"],
54 ["POST", "/api/v3/repos/acme/app/commits/abc123/comments"],
55 ["POST", "/new"],
56 ["POST", "/acme/app/issues"],
57 ["POST", "/acme/app/issues/3/comment"],
58 ["POST", "/acme/app/pull/4/review"],
59 ["POST", "/api/v3/diffinery/reports"],
60 ["POST", "/-/report"],
61 ]) expect(bucketsFor(m!, path!, person, "x"), path).toEqual([["RL_ACCOUNT", "acct:u1"], ["RL_CREATE", "acct:u1"]]);
62 expect(bucketsFor("PATCH", "/api/v3/repos/acme/app/issues/3", person, "x")).toEqual([["RL_ACCOUNT", "acct:u1"]]);
63 });
64
65 it("leaves runners, the Actions runtime, health checks and avatars alone", () => {
66 for (const path of ["/healthz", "/api/avatars/u1", "/twirp/github.actions.results.api.v1.ArtifactService/CreateArtifact", "/_runtime/x", "/runner/jobs/claim", "/-/model/anthropic/v1/messages"]) expect(bucketsFor("POST", path, anon, "x"), path).toEqual([]);
67 expect(bucketsFor("GET", "/api/v3/repos/acme/app", job, "x")).toEqual([]);
68 });
69});
70
71describe("rate limit checks", () => {
72 it("refuses once a bucket is used up, per key", async () => {
73 const rl = limiter(2);
74 const env = { RL_ANON: rl } as unknown as Env;
75 const a = bucketsFor("GET", "/", anon, "203.0.113.9");
76 expect(await overLimit(env, a)).toBeNull();
77 expect(await overLimit(env, a)).toBeNull();
78 expect(await overLimit(env, a)).toBe("RL_ANON");
79 expect(await overLimit(env, bucketsFor("GET", "/", anon, "203.0.113.10"))).toBeNull();
80 });
81
82 it("names the create bucket when an account's general budget still has room", async () => {
83 const env = { RL_ACCOUNT: limiter(100), RL_CREATE: limiter(1) } as unknown as Env;
84 const b = bucketsFor("POST", "/api/v3/repos/acme/app/issues", person, "x");
85 expect(await overLimit(env, b)).toBeNull();
86 expect(await overLimit(env, b)).toBe("RL_CREATE");
87 });
88
89 it("lets requests through when no limiter is configured or the limiter fails", async () => {
90 expect(await overLimit({} as Env, [["RL_ANON", "ip:x"]])).toBeNull();
91 const broken = { RL_ANON: { limit: async () => Promise.reject(new Error("down")) } } as unknown as Env;
92 expect(await overLimit(broken, [["RL_ANON", "ip:x"]])).toBeNull();
93 });
94
95 it("matches the limits configured for production", () => {
96 const text = readFileSync(new URL("../../wrangler.jsonc", import.meta.url), "utf8");
97 for (const [name, limit] of Object.entries(RATE_LIMITS)) {
98 const line = text.split("\n").find((l) => l.includes(`"name": "${name}"`)) ?? "";
99 const m = /"limit": (\d+), "period": (\d+)/.exec(line);
100 expect(m, name).not.toBeNull();
101 expect([Number(m![1]), Number(m![2])], name).toEqual([limit, RATE_PERIOD_S]);
102 }
103 });
104});
105
106describe("rate limits in the Worker", () => {
107 // Importing the whole Worker takes a few seconds when the suite runs in parallel.
108 it("answers 429 with retry-after once an anonymous address is over its limit", { timeout: 30_000 }, async () => {
109 const { app } = await import("../../src/app.js");
110 const env = { DB: {}, PUBLIC_URL: "https://d.example", RL_ANON: limiter(1) } as unknown as Env;
111 const ctx = { waitUntil: () => {}, passThroughOnException: () => {}, props: {} } as unknown as ExecutionContext;
112 const get = () => app.request("/healthz-not-exempt", { headers: { "cf-connecting-ip": "203.0.113.9" } }, env, ctx);
113 expect((await get()).status).toBe(404);
114 const r = await get();
115 expect(r.status).toBe(429);
116 expect(r.headers.get("retry-after")).toBe("60");
117 expect(r.headers.get("strict-transport-security")).toBeTruthy();
118 expect(((await r.json()) as { message: string }).message).toMatch(/at most 300 a minute/);
119 });
120});
121
122describe("model requests per paying person (phase 5a)", () => {
123 it("counts per person, and lets requests through without a binding or when it fails", async () => {
124 const rl = limiter(2);
125 const env = { RL_MODEL: rl } as unknown as Env;
126 expect(await modelRateLimited(env, "p1")).toBe(false);
127 expect(await modelRateLimited(env, "p1")).toBe(false);
128 expect(await modelRateLimited(env, "p1")).toBe(true);
129 expect(await modelRateLimited(env, "p2")).toBe(false);
130 expect(rl.seen).toEqual(["model:p1", "model:p1", "model:p1", "model:p2"]);
131 expect(await modelRateLimited({} as Env, "p1")).toBe(false);
132 const broken = { RL_MODEL: { limit: async () => { throw new Error("down"); } } } as unknown as Env;
133 expect(await modelRateLimited(broken, "p1")).toBe(false);
134 });
135
136 it("matches wrangler.jsonc, and the model proxy is not in the per-account buckets", () => {
137 const text = readFileSync(new URL("../../wrangler.jsonc", import.meta.url), "utf8");
138 const line = text.split("\n").find((l) => l.includes('"name": "RL_MODEL"')) ?? "";
139 expect(/"limit": (\d+), "period": (\d+)/.exec(line)?.slice(1).map(Number)).toEqual([MODEL_RATE_LIMIT, RATE_PERIOD_S]);
140 expect(bucketsFor("POST", "/-/model/anthropic/v1/messages", person, "203.0.113.9")).toEqual([]);
141 });
142});