joke-app/apps/server/src/ws/rate-limit.ts

/** Token bucket (docs/03 §8). */
export class TokenBucket {
  private tokens: number;
  private last: number;
  constructor(
    private ratePerSec: number,
    private burst: number,
    now = Date.now(),
  ) {
    this.tokens = burst;
    this.last = now;
  }

  take(now = Date.now()): boolean {
    this.tokens = Math.min(this.burst, this.tokens + ((now - this.last) / 1000) * this.ratePerSec);
    this.last = now;
    if (this.tokens < 1) return false;
    this.tokens -= 1;
    return true;
  }
}

/** Fixed-window counter per key (HTTP endpoints). */
export class WindowLimiter {
  private hits = new Map<string, { start: number; n: number }>();
  constructor(
    private limit: number,
    private windowMs: number,
  ) {}

  allow(key: string, now = Date.now()): boolean {
    const h = this.hits.get(key);
    if (!h || now - h.start >= this.windowMs) {
      this.hits.set(key, { start: now, n: 1 });
      if (this.hits.size > 50_000) this.prune(now);
      return true;
    }
    h.n++;
    return h.n <= this.limit;
  }

  private prune(now: number): void {
    for (const [k, v] of this.hits) if (now - v.start >= this.windowMs) this.hits.delete(k);
  }
}