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.
 *
 * CONCURRENCY CONTROL:
 *
 * Uses optimistic locking with versioned metadata for race condition safety.
 * Each KV entry includes a version number that is incremented on every write.
 * Accepts eventual consistency - no post-write verification needed for
 * abuse prevention use case.
 *
 * KNOWN LIMITATIONS:
 *
 * 1. EVENTUAL CONSISTENCY: KV writes are eventually consistent, which may
 *    cause slight inaccuracies in rate counting across edge locations.
 *    This is acceptable for abuse prevention.
 *
 * 2. RACE CONDITIONS: Multiple concurrent requests may all succeed if they
 *    read before any writes complete. This is acceptable - rate limiting
 *    provides statistical protection, not strict guarantees.
 *
 * For strict rate limiting requirements (billing, quotas), consider:
 * - Cloudflare Durable Objects for atomic counters with strong consistency
 * - Cloudflare's built-in Rate Limiting rules for global enforcement
 *
 * This implementation is suitable for abuse prevention with single KV read
 * per request (30-50ms overhead).
 */

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

const logger = createLogger('[RateLimit]');

/**
 * 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;

  // For non-Cloudflare requests, fail-closed: use a fixed identifier
  // This applies a shared rate limit to all non-Cloudflare traffic
  logger.warn('CF-Connecting-IP missing - applying shared rate limit');
  return 'non-cloudflare-traffic';
}

/**
 * Check if request is rate limited using Cloudflare KV
 *
 * Simplified flow accepting eventual consistency:
 * 1. Read entry with version from KV metadata
 * 2. Check if limit exceeded
 * 3. Increment counter and version
 * 4. Write with new version
 * 5. Return decision based on pre-increment count
 *
 * No post-write verification - accepts eventual consistency for rate limiting.
 * Version metadata is still incremented for race condition safety.
 *
 * @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 key = `ratelimit:${clientIP}:${path}`;
    const now = Date.now();

    // 1. Read current entry with version metadata (single read)
    const kvResult = await env.RATE_LIMIT_KV.getWithMetadata<{ version?: number }>(key);
    const currentVersion = kvResult.metadata?.version ?? 0;
    let entry: RateLimitEntry | null = null;

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

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

      const ttlSeconds = Math.ceil(config.windowMs / 1000);

      // Write with version 0 (reset on new window) - no verification
      await env.RATE_LIMIT_KV.put(key, JSON.stringify(newEntry), {
        expirationTtl: ttlSeconds,
        metadata: { version: 0 },
      });

      // Allow request - first in new window
      return { allowed: true };
    }

    // 3. Check current count against limit
    const currentCount = entry.count;
    const isOverLimit = currentCount >= config.maxRequests;

    // 4. Increment count and version for next request
    const newVersion = currentVersion + 1;
    const updatedEntry: RateLimitEntry = {
      count: currentCount + 1,
      windowStart: entry.windowStart,
    };

    const windowEnd = entry.windowStart + config.windowMs;
    const ttlSeconds = Math.ceil((windowEnd - now) / 1000);

    // 5. Write updated count (only if TTL valid) - no verification
    if (ttlSeconds > 0) {
      await env.RATE_LIMIT_KV.put(key, JSON.stringify(updatedEntry), {
        expirationTtl: ttlSeconds,
        metadata: { version: newVersion },
      });
    }

    // 6. Return decision based on pre-increment count
    if (isOverLimit) {
      const retryAfter = Math.ceil((windowEnd - now) / 1000);
      return { allowed: false, retryAfter };
    }

    return { allowed: true };
  } catch (error) {
    // Fail-closed on KV errors for security
    logger.error('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) {
    logger.error('Status check error', { error });
    return null;
  }
}