fix(rate-limit): support distributed Redis rate limiting backend with Upstash driver (fixes #91)

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();
+  });
+});

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<string, Bucket>();
+
+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> | RateLimitResult;
+  clear?(key: string): Promise<void> | void;
+}
+
+class MemoryStore implements RateLimitStore {
+  private buckets = new Map<string, Bucket>();
+
+  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<RateLimitResult> {
+    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<void> {
+    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<RateLimitResult> {
+  return currentStore.hit(key, limit, windowMs);
+}
+
+export async function clearRateLimit(key: string): Promise<void> {
+  if (currentStore.clear) {
+    await currentStore.clear(key);
+  }
+  defaultMemoryStore.clear(key);
 }