diff --git a/src/gateway/server/preauth-connection-budget.test.ts b/src/gateway/server/preauth-connection-budget.test.ts new file mode 100644 index 000000000000..699221758f9e --- /dev/null +++ b/src/gateway/server/preauth-connection-budget.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, it } from "vitest"; +import { createPreauthConnectionBudget } from "./preauth-connection-budget.js"; + +describe("createPreauthConnectionBudget", () => { + it("caps connections with a finite configured limit", () => { + const budget = createPreauthConnectionBudget(2); + + expect(budget.acquire("127.0.0.1")).toBe(true); + expect(budget.acquire("127.0.0.1")).toBe(true); + expect(budget.acquire("127.0.0.1")).toBe(false); + + budget.release("127.0.0.1"); + expect(budget.acquire("127.0.0.1")).toBe(true); + }); + + it("uses the default cap for non-finite direct limits", () => { + const budget = createPreauthConnectionBudget(Number.NaN); + + for (let i = 0; i < 32; i += 1) { + expect(budget.acquire("127.0.0.1")).toBe(true); + } + expect(budget.acquire("127.0.0.1")).toBe(false); + }); + + it("shares one capped bucket for missing client IPs", () => { + const budget = createPreauthConnectionBudget(Number.POSITIVE_INFINITY); + + for (let i = 0; i < 32; i += 1) { + expect(budget.acquire(i % 2 === 0 ? undefined : " ")).toBe(true); + } + expect(budget.acquire(undefined)).toBe(false); + }); +}); diff --git a/src/gateway/server/preauth-connection-budget.ts b/src/gateway/server/preauth-connection-budget.ts index dccbd1e1868a..86c1ef106f9b 100644 --- a/src/gateway/server/preauth-connection-budget.ts +++ b/src/gateway/server/preauth-connection-budget.ts @@ -1,3 +1,5 @@ +import { resolveIntegerOption } from "../../shared/number-coercion.js"; + const DEFAULT_MAX_PREAUTH_CONNECTIONS_PER_IP = 32; const UNKNOWN_CLIENT_IP_BUDGET_KEY = "__openclaw_unknown_client_ip__"; @@ -23,6 +25,9 @@ export type PreauthConnectionBudget = { export function createPreauthConnectionBudget( limit = getMaxPreauthConnectionsPerIpFromEnv(), ): PreauthConnectionBudget { + const maxConnectionsPerIp = resolveIntegerOption(limit, getMaxPreauthConnectionsPerIpFromEnv(), { + min: 1, + }); const counts = new Map(); const normalizeBudgetKey = (clientIp: string | undefined) => { const ip = clientIp?.trim(); @@ -36,7 +41,7 @@ export function createPreauthConnectionBudget( acquire(clientIp) { const ip = normalizeBudgetKey(clientIp); const next = (counts.get(ip) ?? 0) + 1; - if (next > limit) { + if (next > maxConnectionsPerIp) { return false; } counts.set(ip, next);