72 lines
2.0 KiB
TypeScript
72 lines
2.0 KiB
TypeScript
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
|
import { consume, _resetMemoryBuckets, clientIp } from "./ratelimit";
|
|
|
|
describe("memory rate limiter", () => {
|
|
beforeEach(() => {
|
|
_resetMemoryBuckets();
|
|
vi.useFakeTimers();
|
|
vi.setSystemTime(new Date(0));
|
|
});
|
|
afterEach(() => {
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it("allows requests under the limit", async () => {
|
|
const r1 = await consume("k", 3, 1000);
|
|
const r2 = await consume("k", 3, 1000);
|
|
const r3 = await consume("k", 3, 1000);
|
|
expect(r1.ok).toBe(true);
|
|
expect(r2.ok).toBe(true);
|
|
expect(r3.ok).toBe(true);
|
|
expect(r3.remaining).toBe(0);
|
|
});
|
|
|
|
it("rejects the request that exceeds the limit", async () => {
|
|
await consume("k", 2, 1000);
|
|
await consume("k", 2, 1000);
|
|
const r3 = await consume("k", 2, 1000);
|
|
expect(r3.ok).toBe(false);
|
|
expect(r3.remaining).toBe(0);
|
|
});
|
|
|
|
it("resets the bucket after the window elapses", async () => {
|
|
await consume("k", 1, 1000);
|
|
const blocked = await consume("k", 1, 1000);
|
|
expect(blocked.ok).toBe(false);
|
|
|
|
vi.advanceTimersByTime(1001);
|
|
|
|
const allowed = await consume("k", 1, 1000);
|
|
expect(allowed.ok).toBe(true);
|
|
});
|
|
|
|
it("isolates buckets by key", async () => {
|
|
await consume("a", 1, 1000);
|
|
const a2 = await consume("a", 1, 1000);
|
|
const b1 = await consume("b", 1, 1000);
|
|
expect(a2.ok).toBe(false);
|
|
expect(b1.ok).toBe(true);
|
|
});
|
|
|
|
it("returns a sensible resetAt timestamp", async () => {
|
|
const r = await consume("k", 5, 1000);
|
|
expect(r.resetAt).toBe(1000);
|
|
});
|
|
});
|
|
|
|
describe("clientIp", () => {
|
|
it("trusts the first hop in x-forwarded-for", () => {
|
|
const h = new Headers({ "x-forwarded-for": "1.2.3.4, 5.6.7.8" });
|
|
expect(clientIp(h)).toBe("1.2.3.4");
|
|
});
|
|
|
|
it("falls back to x-real-ip", () => {
|
|
const h = new Headers({ "x-real-ip": "9.9.9.9" });
|
|
expect(clientIp(h)).toBe("9.9.9.9");
|
|
});
|
|
|
|
it("returns 'unknown' when no header is present", () => {
|
|
expect(clientIp(new Headers())).toBe("unknown");
|
|
});
|
|
});
|