ubiquitous-invention/apps/web/server/lib/rate-limit.ts

import {
  consumeTokenBucket,
  type TokenBucketState,
} from "@tasks/shared";

/**
 * In-process, in-memory rate limiter built on top of the pure token bucket
 * in `@tasks/shared`. Used as a first-line guardrail against brute-force
 * (credentials sign-in) and casual abuse (invite spam) — see
 * `plans/Plan-multitenant-saas-hardening/Epic-tenant-lifecycle/Task-rate-limit-and-abuse-guardrails.md`.
 *
 * Storage is a module-scoped `Map<string, TokenBucketState>`. Entries are
 * lazily expired on lookup (no background timer); a horizontally-scaled
 * deployment would defeat the in-process map because every Next.js worker
 * has its own copy. That's an explicit v1 trade-off — the limits in this
 * file are sized so a single worker is enough to cap brute-force at the
 * "annoy a human" threshold even if N workers each get N attempts. When
 * the app moves to multi-region or multi-replica deploy, swap the body of
 * `rateLimit` for a Redis-backed implementation behind the same exported
 * type. Tracked as a follow-up in the task spec (the "Storage" section).
 *
 * NOT async on purpose: the in-memory implementation has no I/O and we
 * want call sites to stay synchronous. The Redis follow-up will need its
 * own async surface (`rateLimitAsync` or similar) rather than retrofitting
 * `Promise<RateLimitResult>` here, so existing callers don't have to be
 * audited for `await` correctness when the storage backend changes.
 */

export type RateLimitResult =
  | { allowed: true; remaining: number; resetAt: Date }
  | {
      allowed: false;
      remaining: 0;
      resetAt: Date;
      retryAfterMs: number;
    };

export interface RateLimitOptions {
  /**
   * Globally unique bucket identifier. Convention: `${scope}:${actor}`,
   * e.g. `signin:192.0.2.1`, `invite:create:<userId>`. The scope prefix
   * is required — without it, two unrelated features sharing a key like
   * a bare userId would collide.
   */
  key: string;
  /** Max allowed actions in `windowMs`. */
  limit: number;
  /** Window length in milliseconds. */
  windowMs: number;
}

// Eviction threshold: once the map grows past this, we opportunistically
// sweep expired entries on the *next* lookup. Picked to be large enough
// that we don't sweep on every cold start, small enough that a single
// rogue IP iterating keys can't OOM the process before we notice.
const EVICTION_THRESHOLD = 10_000;

const buckets = new Map<string, TokenBucketState>();

function maybeSweep(now: number): void {
  if (buckets.size <= EVICTION_THRESHOLD) return;
  for (const [k, v] of buckets) {
    if (now >= v.resetAt) buckets.delete(k);
  }
}

export function rateLimit(opts: RateLimitOptions): RateLimitResult {
  const now = Date.now();
  maybeSweep(now);

  const prior = buckets.get(opts.key);
  const result = consumeTokenBucket(prior, now, {
    limit: opts.limit,
    windowMs: opts.windowMs,
  });
  buckets.set(opts.key, result.state);

  if (result.allowed) {
    return {
      allowed: true,
      remaining: result.remaining,
      resetAt: new Date(result.resetAt),
    };
  }
  return {
    allowed: false,
    remaining: 0,
    resetAt: new Date(result.resetAt),
    retryAfterMs: result.retryAfterMs,
  };
}

/**
 * Test-only escape hatch so colocated unit tests (when we add a vitest
 * harness for `apps/web`) can reset shared module state between cases.
 * Not exported through any barrel — call sites should never need this.
 */
export function __resetRateLimitForTests(): void {
  buckets.clear();
}