diff --git a/apps/web/lib/rate-limit.test.ts b/apps/web/lib/rate-limit.test.ts new file mode 100644 index 0000000..4d17fc3 --- /dev/null +++ b/apps/web/lib/rate-limit.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, it, vi, beforeEach } from "vitest"; +import { + rateLimit, + rateLimitAsync, + clearRateLimit, + setRateLimitStore, + UpstashStore, + RateLimitStore, +} from "./rate-limit"; + +describe("rateLimit", () => { + beforeEach(async () => { + setRateLimitStore(null); + await clearRateLimit("test-key"); + }); + + it("limits requests within window and resets after expiry", () => { + const res1 = rateLimit("test-key", 2, 1000); + expect(res1.ok).toBe(true); + expect(res1.remaining).toBe(1); + + const res2 = rateLimit("test-key", 2, 1000); + expect(res2.ok).toBe(true); + expect(res2.remaining).toBe(0); + + const res3 = rateLimit("test-key", 2, 1000); + expect(res3.ok).toBe(false); + expect(res3.remaining).toBe(0); + }); + + it("clears rate limit state with clearRateLimit", async () => { + rateLimit("test-key", 1, 1000); + expect(rateLimit("test-key", 1, 1000).ok).toBe(false); + + await clearRateLimit("test-key"); + expect(rateLimit("test-key", 1, 1000).ok).toBe(true); + }); + + it("supports custom store via setRateLimitStore", async () => { + const mockStore: RateLimitStore = { + hit: vi.fn().mockReturnValue({ ok: true, remaining: 99, resetAt: 12345 }), + clear: vi.fn(), + }; + setRateLimitStore(mockStore); + + const res = await rateLimitAsync("test-key", 100, 60000); + expect(res.ok).toBe(true); + expect(res.remaining).toBe(99); + expect(mockStore.hit).toHaveBeenCalledWith("test-key", 100, 60000); + }); + + it("handles UpstashStore REST pipeline responses", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValueOnce({ + ok: true, + json: async () => [{ result: 2 }, { result: "OK" }, { result: 50 }], + } as Response); + + const upstash = new UpstashStore("https://example.upstash.io", "mock-token"); + const res = await upstash.hit("api-ip", 5, 60000); + + expect(res.ok).toBe(true); + expect(res.remaining).toBe(3); + expect(fetchSpy).toHaveBeenCalledWith( + "https://example.upstash.io/pipeline", + expect.objectContaining({ + headers: expect.objectContaining({ Authorization: "Bearer mock-token" }), + }) + ); + fetchSpy.mockRestore(); + }); +}); diff --git a/apps/web/lib/rate-limit.ts b/apps/web/lib/rate-limit.ts index 3397bbf..60cf17a 100644 --- a/apps/web/lib/rate-limit.ts +++ b/apps/web/lib/rate-limit.ts @@ -1,31 +1,137 @@ /** - * Minimal in-memory fixed-window rate limiter. Good enough for the MVP on a - * single Railway instance; swap for Upstash/Redis when horizontally scaled. + * Fixed-window rate limiter supporting both in-memory local fallback and + * distributed multi-instance storage (e.g. Upstash Redis REST API). */ -type Bucket = { count: number; resetAt: number }; -const buckets = new Map(); + +export type Bucket = { count: number; resetAt: number }; +export type RateLimitResult = { ok: boolean; remaining: number; resetAt: number }; + +export interface RateLimitStore { + hit(key: string, limit: number, windowMs: number): Promise | RateLimitResult; + clear?(key: string): Promise | void; +} + +class MemoryStore implements RateLimitStore { + private buckets = new Map(); + + constructor() { + if (typeof setInterval !== "undefined") { + setInterval(() => { + const now = Date.now(); + for (const [k, v] of this.buckets) { + if (v.resetAt < now) this.buckets.delete(k); + } + }, 60_000).unref?.(); + } + } + + hit(key: string, limit: number, windowMs: number): RateLimitResult { + const now = Date.now(); + const b = this.buckets.get(key); + if (!b || b.resetAt < now) { + const resetAt = now + windowMs; + this.buckets.set(key, { count: 1, resetAt }); + return { ok: true, remaining: limit - 1, resetAt }; + } + b.count++; + const ok = b.count <= limit; + return { ok, remaining: Math.max(0, limit - b.count), resetAt: b.resetAt }; + } + + clear(key: string): void { + this.buckets.delete(key); + } +} + +/** Upstash REST API store for multi-instance distributed deployments */ +export class UpstashStore implements RateLimitStore { + private url: string; + private token: string; + + constructor(url: string, token: string) { + this.url = url.replace(/\/$/, ""); + this.token = token; + } + + async hit(key: string, limit: number, windowMs: number): Promise { + const now = Date.now(); + const windowSec = Math.ceil(windowMs / 1000); + const redisKey = `ratelimit:${key}`; + + try { + // Atomic pipeline: INCR and EXPIRE if new + const res = await fetch(`${this.url}/pipeline`, { + method: "POST", + headers: { + Authorization: `Bearer ${this.token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify([ + ["INCR", redisKey], + ["EXPIRE", redisKey, windowSec, "NX"], + ["TTL", redisKey], + ]), + }); + + if (!res.ok) throw new Error(`Upstash HTTP ${res.status}`); + const data = await res.json(); + const count = Number(data[0]?.result ?? 1); + const ttlSec = Number(data[2]?.result ?? windowSec); + + const resetAt = now + (ttlSec > 0 ? ttlSec * 1000 : windowMs); + const ok = count <= limit; + return { ok, remaining: Math.max(0, limit - count), resetAt }; + } catch { + // Fallback to local in-memory store if Redis request fails + return defaultMemoryStore.hit(key, limit, windowMs); + } + } + + async clear(key: string): Promise { + try { + await fetch(`${this.url}/del/ratelimit:${key}`, { + method: "POST", + headers: { Authorization: `Bearer ${this.token}` }, + }); + } catch {} + } +} + +const defaultMemoryStore = new MemoryStore(); +let currentStore: RateLimitStore = defaultMemoryStore; + +export function setRateLimitStore(store: RateLimitStore | null): void { + currentStore = store || defaultMemoryStore; +} + +export function getRateLimitStore(): RateLimitStore { + return currentStore; +} export function rateLimit( key: string, limit: number, windowMs: number, -): { ok: boolean; remaining: number; resetAt: number } { - const now = Date.now(); - const b = buckets.get(key); - if (!b || b.resetAt < now) { - const resetAt = now + windowMs; - buckets.set(key, { count: 1, resetAt }); - return { ok: true, remaining: limit - 1, resetAt }; - } - b.count++; - const ok = b.count <= limit; - return { ok, remaining: Math.max(0, limit - b.count), resetAt: b.resetAt }; -} - -// Occasional cleanup to bound memory. -if (typeof setInterval !== "undefined") { - setInterval(() => { - const now = Date.now(); - for (const [k, v] of buckets) if (v.resetAt < now) buckets.delete(k); - }, 60_000).unref?.(); +): RateLimitResult { + const result = currentStore.hit(key, limit, windowMs); + if (result instanceof Promise) { + // If the configured store is async (e.g. Upstash), fall back to memory store for sync rateLimit callers + return defaultMemoryStore.hit(key, limit, windowMs); + } + return result; +} + +export async function rateLimitAsync( + key: string, + limit: number, + windowMs: number, +): Promise { + return currentStore.hit(key, limit, windowMs); +} + +export async function clearRateLimit(key: string): Promise { + if (currentStore.clear) { + await currentStore.clear(key); + } + defaultMemoryStore.clear(key); }