Shade/packages/shade-server/src/rate-limit.ts

import { RateLimitError } from '@shade/core';

/**
 * Simple token-bucket rate limiter with pluggable storage.
 *
 * Default storage is in-memory (Map). For distributed deployments,
 * swap in a Redis-backed RateLimitStore implementation.
 */

export interface RateLimitStore {
  /** Get the current token count and last refill time for a key */
  get(key: string): Promise<{ tokens: number; lastRefill: number } | null>;
  /** Store the current token count and last refill time */
  set(key: string, tokens: number, lastRefill: number): Promise<void>;
}

export class MemoryRateLimitStore implements RateLimitStore {
  private entries = new Map<string, { tokens: number; lastRefill: number }>();

  async get(key: string) {
    return this.entries.get(key) ?? null;
  }

  async set(key: string, tokens: number, lastRefill: number) {
    this.entries.set(key, { tokens, lastRefill });
  }

  /** Periodic cleanup: drop entries older than `maxAge` ms */
  cleanup(maxAge: number) {
    const now = Date.now();
    for (const [key, entry] of this.entries) {
      if (now - entry.lastRefill > maxAge) {
        this.entries.delete(key);
      }
    }
  }
}

export interface RateLimitConfig {
  /** Maximum tokens in the bucket (burst capacity) */
  capacity: number;
  /** Tokens added per second */
  refillPerSecond: number;
}

export class RateLimiter {
  constructor(
    private readonly store: RateLimitStore,
    private readonly config: RateLimitConfig,
  ) {}

  /**
   * Attempt to consume one token for the given key.
   * Throws RateLimitError if no tokens available.
   */
  async consume(key: string, tokens = 1): Promise<void> {
    const now = Date.now();
    const entry = await this.store.get(key);

    let currentTokens: number;
    if (!entry) {
      currentTokens = this.config.capacity;
    } else {
      const elapsed = (now - entry.lastRefill) / 1000;
      const refilled = elapsed * this.config.refillPerSecond;
      currentTokens = Math.min(this.config.capacity, entry.tokens + refilled);
    }

    if (currentTokens < tokens) {
      const needed = tokens - currentTokens;
      const retryAfter = Math.ceil(needed / this.config.refillPerSecond);
      throw new RateLimitError(`Rate limit exceeded for ${key}`, retryAfter);
    }

    await this.store.set(key, currentTokens - tokens, now);
  }
}

// ─── Preset configurations ──────────────────────────────────

/** Register: 5 per hour (prevent identity flooding) */
export const REGISTER_LIMIT: RateLimitConfig = {
  capacity: 5,
  refillPerSecond: 5 / 3600, // 5 per hour
};

/** Fetch bundle: 60 per minute (anonymous, high limit) */
export const FETCH_LIMIT: RateLimitConfig = {
  capacity: 60,
  refillPerSecond: 1,
};

/** Replenish: 10 per minute per identity */
export const REPLENISH_LIMIT: RateLimitConfig = {
  capacity: 10,
  refillPerSecond: 10 / 60,
};

/** Delete: 5 per hour per identity */
export const DELETE_LIMIT: RateLimitConfig = {
  capacity: 5,
  refillPerSecond: 5 / 3600,
};