@ordojs/security/src/validation/rate-limiter.ts

Security package for OrdoJS with XSS, CSRF, and injection protection

/**
 * Rate Limiter
 * Provides request throttling and rate limiting functionality
 */

import type { RateLimitOptions, RateLimitStore } from './types';

/**
 * In-memory rate limit store
 */
export class MemoryRateLimitStore implements RateLimitStore {
  private store = new Map<string, { count: number; resetTime: number }>();
  private cleanupInterval: NodeJS.Timeout;

  constructor() {
    // Clean up expired entries every minute
    this.cleanupInterval = setInterval(() => {
      this.cleanup();
    }, 60000);
  }

  async get(key: string): Promise<number | null> {
    const entry = this.store.get(key);
    if (!entry || Date.now() > entry.resetTime) {
      return null;
    }
    return entry.count;
  }

  async set(key: string, value: number, ttl: number): Promise<void> {
    this.store.set(key, {
      count: value,
      resetTime: Date.now() + ttl
    });
  }

  async increment(key: string, ttl: number): Promise<number> {
    const entry = this.store.get(key);
    const now = Date.now();

    if (!entry || now > entry.resetTime) {
      // Create new entry
      this.store.set(key, {
        count: 1,
        resetTime: now + ttl
      });
      return 1;
    } else {
      // Increment existing entry
      entry.count++;
      return entry.count;
    }
  }

  async reset(key: string): Promise<void> {
    this.store.delete(key);
  }

  private cleanup(): void {
    const now = Date.now();
    for (const [key, entry] of this.store.entries()) {
      if (now > entry.resetTime) {
        this.store.delete(key);
      }
    }
  }

  destroy(): void {
    if (this.cleanupInterval) {
      clearInterval(this.cleanupInterval);
    }
    this.store.clear();
  }
}

/**
 * Redis-based rate limit store
 */
export class RedisRateLimitStore implements RateLimitStore {
  constructor(private redisClient: any) {}

  async get(key: string): Promise<number | null> {
    const value = await this.redisClient.get(key);
    return value ? parseInt(value, 10) : null;
  }

  async set(key: string, value: number, ttl: number): Promise<void> {
    await this.redisClient.setex(key, Math.ceil(ttl / 1000), value);
  }

  async increment(key: string, ttl: number): Promise<number> {
    const multi = this.redisClient.multi();
    multi.incr(key);
    multi.expire(key, Math.ceil(ttl / 1000));
    const results = await multi.exec();
    return results[0][1]; // Return the incremented value
  }

  async reset(key: string): Promise<void> {
    await this.redisClient.del(key);
  }
}

/**
 * Rate limiter implementation
 */
export class RateLimiter {
  private store: RateLimitStore;
  private options: Required<RateLimitOptions>;

  constructor(options: RateLimitOptions, store?: RateLimitStore) {
    this.options = {
      windowMs: options.windowMs,
      maxRequests: options.maxRequests,
      keyGenerator: options.keyGenerator || this.defaultKeyGenerator,
      skipSuccessfulRequests: options.skipSuccessfulRequests || false,
      skipFailedRequests: options.skipFailedRequests || false,
      onLimitReached: options.onLimitReached || (() => {})
    };

    this.store = store || new MemoryRateLimitStore();
  }

  /**
   * Check if request should be rate limited
   */
  async checkLimit(request: any): Promise<{
    allowed: boolean;
    limit: number;
    remaining: number;
    resetTime: number;
    retryAfter?: number;
  }> {
    const key = this.options.keyGenerator(request);
    const now = Date.now();
    const windowStart = now - this.options.windowMs;

    // Get current count
    const currentCount = await this.store.get(key) || 0;

    // Check if limit exceeded
    if (currentCount >= this.options.maxRequests) {
      const resetTime = now + this.options.windowMs;
      const retryAfter = Math.ceil(this.options.windowMs / 1000);

      // Call limit reached callback
      this.options.onLimitReached(request);

      return {
        allowed: false,
        limit: this.options.maxRequests,
        remaining: 0,
        resetTime,
        retryAfter
      };
    }

    // Increment counter
    const newCount = await this.store.increment(key, this.options.windowMs);

    return {
      allowed: true,
      limit: this.options.maxRequests,
      remaining: Math.max(0, this.options.maxRequests - newCount),
      resetTime: now + this.options.windowMs
    };
  }

  /**
   * Reset rate limit for a specific key
   */
  async resetLimit(request: any): Promise<void> {
    const key = this.options.keyGenerator(request);
    await this.store.reset(key);
  }

  /**
   * Default key generator using IP address
   */
  private defaultKeyGenerator(request: any): string {
    // Try to get IP from various sources
    const ip = request.ip ||
               request.connection?.remoteAddress ||
               request.socket?.remoteAddress ||
               request.headers?.['x-forwarded-for']?.split(',')[0]?.trim() ||
               request.headers?.['x-real-ip'] ||
               'unknown';

    return `rate_limit:${ip}`;
  }

  /**
   * Create middleware for Express-like frameworks
   */
  middleware() {
    return async (req: any, res: any, next: any) => {
      try {
        const result = await this.checkLimit(req);

        // Set rate limit headers
        res.setHeader('X-RateLimit-Limit', result.limit);
        res.setHeader('X-RateLimit-Remaining', result.remaining);
        res.setHeader('X-RateLimit-Reset', new Date(result.resetTime).toISOString());

        if (!result.allowed) {
          res.setHeader('Retry-After', result.retryAfter);
          res.status(429).json({
            error: 'Too Many Requests',
            message: 'Rate limit exceeded',
            retryAfter: result.retryAfter
          });
          return;
        }

        next();
      } catch (error) {
        console.error('Rate limiter error:', error);
        next(); // Continue on error to avoid breaking the application
      }
    };
  }
}

/**
 * Sliding window rate limiter for more precise control
 */
export class SlidingWindowRateLimiter {
  private store: RateLimitStore;
  private options: Required<RateLimitOptions>;

  constructor(options: RateLimitOptions, store?: RateLimitStore) {
    this.options = {
      windowMs: options.windowMs,
      maxRequests: options.maxRequests,
      keyGenerator: options.keyGenerator || this.defaultKeyGenerator,
      skipSuccessfulRequests: options.skipSuccessfulRequests || false,
      skipFailedRequests: options.skipFailedRequests || false,
      onLimitReached: options.onLimitReached || (() => {})
    };

    this.store = store || new MemoryRateLimitStore();
  }

  async checkLimit(request: any): Promise<{
    allowed: boolean;
    limit: number;
    remaining: number;
    resetTime: number;
    retryAfter?: number;
  }> {
    const key = this.options.keyGenerator(request);
    const now = Date.now();
    const windowStart = now - this.options.windowMs;

    // For sliding window, we need to track individual request timestamps
    // This is a simplified implementation - in production, you'd want to use
    // a more sophisticated data structure like a sorted set in Redis

    const currentCount = await this.store.get(key) || 0;

    if (currentCount >= this.options.maxRequests) {
      this.options.onLimitReached(request);

      return {
        allowed: false,
        limit: this.options.maxRequests,
        remaining: 0,
        resetTime: now + this.options.windowMs,
        retryAfter: Math.ceil(this.options.windowMs / 1000)
      };
    }

    const newCount = await this.store.increment(key, this.options.windowMs);

    return {
      allowed: true,
      limit: this.options.maxRequests,
      remaining: Math.max(0, this.options.maxRequests - newCount),
      resetTime: now + this.options.windowMs
    };
  }

  private defaultKeyGenerator(request: any): string {
    const ip = request.ip ||
               request.connection?.remoteAddress ||
               request.socket?.remoteAddress ||
               'unknown';

    return `sliding_rate_limit:${ip}`;
  }

  middleware() {
    return async (req: any, res: any, next: any) => {
      try {
        const result = await this.checkLimit(req);

        res.setHeader('X-RateLimit-Limit', result.limit);
        res.setHeader('X-RateLimit-Remaining', result.remaining);
        res.setHeader('X-RateLimit-Reset', new Date(result.resetTime).toISOString());

        if (!result.allowed) {
          res.setHeader('Retry-After', result.retryAfter);
          res.status(429).json({
            error: 'Too Many Requests',
            message: 'Rate limit exceeded',
            retryAfter: result.retryAfter
          });
          return;
        }

        next();
      } catch (error) {
        console.error('Sliding window rate limiter error:', error);
        next();
      }
    };
  }
}