cloud-server/src/middleware/rateLimit.ts

/**
 * Rate Limiting Middleware - Cloudflare KV Based
 *
 * Distributed rate limiting using Cloudflare KV for multi-worker support.
 * Different limits for different endpoints.
 */

import { Env } from '../types';
import { RATE_LIMIT_DEFAULTS, HTTP_STATUS } from '../constants';

/**
 * Rate limit configuration per endpoint
 */
interface RateLimitConfig {
  /** Maximum requests allowed in the time window */
  maxRequests: number;
  /** Time window in milliseconds */
  windowMs: number;
}

/**
 * Rate limit tracking entry stored in KV
 */
interface RateLimitEntry {
  /** Request count in current window */
  count: number;
  /** Window start timestamp */
  windowStart: number;
}

/**
 * Rate limit configurations by endpoint
 */
const RATE_LIMITS: Record<string, RateLimitConfig> = {
  '/instances': {
    maxRequests: RATE_LIMIT_DEFAULTS.MAX_REQUESTS_INSTANCES,
    windowMs: RATE_LIMIT_DEFAULTS.WINDOW_MS,
  },
  '/sync': {
    maxRequests: RATE_LIMIT_DEFAULTS.MAX_REQUESTS_SYNC,
    windowMs: RATE_LIMIT_DEFAULTS.WINDOW_MS,
  },
};

/**
 * Get client IP from request
 *
 * Security: Only trust CF-Connecting-IP header from Cloudflare.
 * X-Forwarded-For can be spoofed by clients and should NOT be trusted.
 */
function getClientIP(request: Request): string {
  // Only trust Cloudflare-provided IP header
  const cfIP = request.headers.get('CF-Connecting-IP');
  if (cfIP) return cfIP;

  // If CF-Connecting-IP is missing, request may not be from Cloudflare
  // Use unique identifier to still apply rate limit
  console.warn('[RateLimit] CF-Connecting-IP missing, possible direct access');
  return `unknown-${Date.now()}-${Math.random().toString(36).substring(7)}`;
}

/**
 * Check if request is rate limited using Cloudflare KV
 *
 * @param request - HTTP request to check
 * @param path - Request path for rate limit lookup
 * @param env - Cloudflare Worker environment with KV binding
 * @returns Object with allowed status and optional retry-after seconds
 */
export async function checkRateLimit(
  request: Request,
  path: string,
  env: Env
): Promise<{ allowed: boolean; retryAfter?: number }> {
  const config = RATE_LIMITS[path];
  if (!config) {
    // No rate limit configured for this path
    return { allowed: true };
  }

  try {
    const clientIP = getClientIP(request);
    const now = Date.now();
    const key = `ratelimit:${clientIP}:${path}`;

    // Get current entry from KV
    const entryJson = await env.RATE_LIMIT_KV.get(key);
    let entry: RateLimitEntry | null = null;

    if (entryJson) {
      try {
        entry = JSON.parse(entryJson);
      } catch {
        // Invalid JSON, treat as no entry
        entry = null;
      }
    }

    // Check if window has expired
    if (!entry || entry.windowStart + config.windowMs <= now) {
      // New window - allow and create new entry
      const newEntry: RateLimitEntry = {
        count: 1,
        windowStart: now,
      };

      // Store with TTL (convert ms to seconds, round up)
      const ttlSeconds = Math.ceil(config.windowMs / 1000);
      await env.RATE_LIMIT_KV.put(key, JSON.stringify(newEntry), {
        expirationTtl: ttlSeconds,
      });

      return { allowed: true };
    }

    // Increment count
    entry.count++;

    // Check if over limit
    if (entry.count > config.maxRequests) {
      const windowEnd = entry.windowStart + config.windowMs;
      const retryAfter = Math.ceil((windowEnd - now) / 1000);

      // Still update KV to persist the attempt
      const ttlSeconds = Math.ceil((windowEnd - now) / 1000);
      if (ttlSeconds > 0) {
        await env.RATE_LIMIT_KV.put(key, JSON.stringify(entry), {
          expirationTtl: ttlSeconds,
        });
      }

      return { allowed: false, retryAfter };
    }

    // Update entry in KV
    const windowEnd = entry.windowStart + config.windowMs;
    const ttlSeconds = Math.ceil((windowEnd - now) / 1000);
    if (ttlSeconds > 0) {
      await env.RATE_LIMIT_KV.put(key, JSON.stringify(entry), {
        expirationTtl: ttlSeconds,
      });
    }

    return { allowed: true };
  } catch (error) {
    // Fail-closed on KV errors for security
    console.error('[RateLimit] KV error, blocking request for safety:', error);
    return { allowed: false, retryAfter: 60 };
  }
}

/**
 * Create 429 Too Many Requests response
 */
export function createRateLimitResponse(retryAfter: number): Response {
  return Response.json(
    {
      error: 'Too Many Requests',
      message: 'Rate limit exceeded. Please try again later.',
      retry_after_seconds: retryAfter,
      timestamp: new Date().toISOString(),
    },
    {
      status: HTTP_STATUS.TOO_MANY_REQUESTS,
      headers: {
        'Retry-After': retryAfter.toString(),
        'X-RateLimit-Retry-After': retryAfter.toString(),
      },
    }
  );
}

/**
 * Get current rate limit status for a client
 * (useful for debugging or adding X-RateLimit-* headers)
 */
export async function getRateLimitStatus(
  request: Request,
  path: string,
  env: Env
): Promise<{ limit: number; remaining: number; resetAt: number } | null> {
  const config = RATE_LIMITS[path];
  if (!config) return null;

  try {
    const clientIP = getClientIP(request);
    const now = Date.now();
    const key = `ratelimit:${clientIP}:${path}`;

    const entryJson = await env.RATE_LIMIT_KV.get(key);

    if (!entryJson) {
      return {
        limit: config.maxRequests,
        remaining: config.maxRequests,
        resetAt: now + config.windowMs,
      };
    }

    const entry: RateLimitEntry = JSON.parse(entryJson);

    // Check if window expired
    if (entry.windowStart + config.windowMs <= now) {
      return {
        limit: config.maxRequests,
        remaining: config.maxRequests,
        resetAt: now + config.windowMs,
      };
    }

    return {
      limit: config.maxRequests,
      remaining: Math.max(0, config.maxRequests - entry.count),
      resetAt: entry.windowStart + config.windowMs,
    };
  } catch (error) {
    console.error('[RateLimit] Status check error:', error);
    return null;
  }
}