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

});

});

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__";

export function createPreauthConnectionBudget(

limit = getMaxPreauthConnectionsPerIpFromEnv(),

): PreauthConnectionBudget {

const maxConnectionsPerIp = resolveIntegerOption(limit, getMaxPreauthConnectionsPerIpFromEnv(), {

min: 1,

});

const counts = new Map<string, number>();

const normalizeBudgetKey = (clientIp: string | undefined) => {

const ip = clientIp?.trim();

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