OpenFrame/lib/rate-limit.ts

import { createHash } from 'crypto';
import { db } from '@/lib/db';
import { NextResponse } from 'next/server';
import { logError, logWarn } from '@/lib/logger';

const RATE_LIMIT_CLEANUP_INTERVAL_MS = 5 * 60 * 1000;

const globalForRateLimitCleanup = globalThis as unknown as {
  rateLimitCleanupIntervalStarted?: boolean;
};

interface RateLimitConfig {
  windowMs: number; // Time window in milliseconds
  maxRequests: number; // Max requests per window
}

interface RateLimitResult {
  allowed: boolean;
  remaining: number;
  resetAt: Date;
}

const TRUTHY_ENV_VALUES = new Set(['1', 'true', 'yes', 'on']);

function isRateLimitDisabled(): boolean {
  const rawValue = process.env.DISABLE_RATE_LIMIT?.trim().toLowerCase();
  return rawValue !== undefined && TRUTHY_ENV_VALUES.has(rawValue);
}

if (process.env.NODE_ENV === 'production' && isRateLimitDisabled()) {
  throw new Error(
    'DISABLE_RATE_LIMIT must not be set in production. ' +
      'Remove or unset the environment variable before deploying.'
  );
}

// Without a proxy mode every caller resolves to 127.0.0.1, so the limiter counts the whole
// world in one bucket. That is the deliberate trade-off (trusting a spoofable header is
// worse), but a deployment behind a proxy should know it is running with global rather
// than per-client limits rather than discover it under load.
if (process.env.NODE_ENV === 'production' && !process.env.TRUSTED_PROXY_MODE?.trim()) {
  logWarn(
    'TRUSTED_PROXY_MODE is not set. Every request resolves to 127.0.0.1, so rate limits ' +
      'apply per process rather than per client. Set TRUSTED_PROXY_MODE=cloudflare or ' +
      'TRUSTED_PROXY_MODE=nginx once you have confirmed your proxy overwrites the ' +
      'corresponding header on every inbound request.'
  );
}

// Industry-standard rate limit defaults per action
export const RATE_LIMIT_CONFIGS: Record<string, RateLimitConfig> = {
  // Auth — strict to prevent brute force / credential stuffing
  register: { windowMs: 60 * 60 * 1000, maxRequests: 5 }, // 5 per hour
  login: { windowMs: 15 * 60 * 1000, maxRequests: 10 }, // 10 per 15 min
  'share-unlock': { windowMs: 15 * 60 * 1000, maxRequests: 20 }, // 20 per 15 min per IP
  'share-unlock-token': { windowMs: 15 * 60 * 1000, maxRequests: 8 }, // 8 per 15 min per IP+token

  // Content creation — moderate limits
  comment: { windowMs: 60 * 1000, maxRequests: 15 }, // 15 per minute
  'image-upload': { windowMs: 60 * 1000, maxRequests: 20 }, // 20 per minute
  'voice-upload': { windowMs: 60 * 1000, maxRequests: 10 }, // 10 per minute
  'feedback-submit': { windowMs: 60 * 1000, maxRequests: 8 }, // 8 per minute
  'feedback-upload': { windowMs: 60 * 1000, maxRequests: 20 }, // 20 per minute
  'create-project': { windowMs: 60 * 60 * 1000, maxRequests: 20 }, // 20 per hour
  'create-video': { windowMs: 60 * 1000, maxRequests: 10 }, // 10 per minute
  'create-version': { windowMs: 60 * 1000, maxRequests: 10 }, // 10 per minute
  'create-workspace': { windowMs: 60 * 60 * 1000, maxRequests: 10 }, // 10 per hour
  'asset-list': { windowMs: 60 * 1000, maxRequests: 120 }, // 120 per minute
  'asset-create': { windowMs: 60 * 1000, maxRequests: 20 }, // 20 per minute
  'asset-delete': { windowMs: 60 * 1000, maxRequests: 20 }, // 20 per minute
  'asset-download': { windowMs: 60 * 1000, maxRequests: 10 }, // 10 per minute
  'asset-bunny-init': { windowMs: 60 * 1000, maxRequests: 10 }, // 10 per minute
  'asset-r2-init': { windowMs: 60 * 1000, maxRequests: 10 }, // 10 per minute

  // Search — debounced on client but protect against scripted callers
  search: { windowMs: 60 * 1000, maxRequests: 60 }, // 60 per minute

  // Watch progress — allow frequent updates but prevent abuse
  'watch-progress': { windowMs: 60 * 1000, maxRequests: 30 }, // 30 per minute (pausing + periodic + visibility changes)

  // Exports
  'comment-export': { windowMs: 60 * 1000, maxRequests: 10 }, // 10 per minute

  // Downloads — strict enough to limit upstream probing/cost abuse
  'video-download': { windowMs: 60 * 1000, maxRequests: 8 }, // 8 per minute
  'video-download-prepare': { windowMs: 60 * 1000, maxRequests: 5 }, // 5 per minute
  'project-download': { windowMs: 60 * 1000, maxRequests: 3 }, // 3 per minute

  // Email verification
  'verify-email': { windowMs: 15 * 60 * 1000, maxRequests: 20 }, // 20 per 15 min (clicked link)
  'resend-verification': { windowMs: 60 * 60 * 1000, maxRequests: 5 }, // 5 per hour

  // Onboarding — one-time action, very strict
  'onboarding-complete': { windowMs: 60 * 60 * 1000, maxRequests: 5 }, // 5 per hour

  // Member management
  'invite-member': { windowMs: 60 * 60 * 1000, maxRequests: 30 }, // 30 per hour
  // Unauthenticated invitation preview (accept landing + invited sign-up). Two buckets,
  // as with share-unlock: a generous per-IP one that bounds token enumeration, and a
  // tight per-IP+token one that stops repeated probing of a single invitation.
  'invitation-preview': { windowMs: 15 * 60 * 1000, maxRequests: 120 }, // 120 per 15 min per IP
  'invitation-preview-token': { windowMs: 15 * 60 * 1000, maxRequests: 12 }, // 12 per 15 min per IP+token
  'manage-member': { windowMs: 60 * 1000, maxRequests: 20 }, // 20 per minute

  // Mutations (update/delete) — moderate
  mutate: { windowMs: 60 * 1000, maxRequests: 30 }, // 30 per minute

  // Analytics beacon — anonymous and public, so bound it per IP
  'analytics-beacon': { windowMs: 60 * 60 * 1000, maxRequests: 30 }, // 30 per hour

  // General reads — generous
  api: { windowMs: 60 * 1000, maxRequests: 100 }, // 100 per minute
};

// Column widths of rate_limits.key and rate_limits.action in prisma/schema.prisma. A value
// wider than its column would fail the INSERT with SQLSTATE 22001.
const RATE_LIMIT_KEY_MAX_LENGTH = 255;
const RATE_LIMIT_ACTION_MAX_LENGTH = 50;

/**
 * Fits a value to its column without ever giving up on counting it. A SHA-256 hex digest
 * is 64 characters, so it is truncated for the narrower action column; 50 hex characters
 * is 200 bits, far past any collision that matters for a rate limit bucket.
 */
function fitToColumn(value: string, maxLength: number): string {
  if (value.length <= maxLength) return value;
  return createHash('sha256').update(value).digest('hex').slice(0, maxLength);
}

/**
 * Check and update rate limit for a given key and action
 * Uses PostgreSQL UNLOGGED table for performance
 */
export async function checkRateLimit(
  key: string,
  action: string,
  config?: RateLimitConfig
): Promise<RateLimitResult> {
  const { windowMs, maxRequests } = config || RATE_LIMIT_CONFIGS[action] || RATE_LIMIT_CONFIGS.api;

  if (isRateLimitDisabled()) {
    return {
      allowed: true,
      remaining: maxRequests,
      resetAt: new Date(Date.now() + windowMs),
    };
  }

  const windowSeconds = Math.floor(windowMs / 1000);

  // Anything wider than its column is replaced by a digest rather than skipped. Skipping
  // meant the limit stopped applying altogether, and letting the value through meant the
  // INSERT failed with SQLSTATE 22001 and the catch below allowed the request anyway.
  // Both were fail-open. A digest is stable, so the same caller keeps the same bucket.
  const storedKey = fitToColumn(key, RATE_LIMIT_KEY_MAX_LENGTH);
  const storedAction = fitToColumn(action, RATE_LIMIT_ACTION_MAX_LENGTH);

  try {
    // Atomic upsert with window check
    // If window expired, reset count; otherwise increment
    const result = await db.$queryRaw<
      Array<{
        count: number;
        window_start: Date;
        is_new_window: boolean;
      }>
    >`
            INSERT INTO rate_limits (key, action, count, window_start)
            VALUES (${storedKey}, ${storedAction}, 1, NOW())
            ON CONFLICT (key, action) DO UPDATE SET
                count = CASE 
                    WHEN rate_limits.window_start < NOW() - (${windowSeconds} || ' seconds')::INTERVAL 
                    THEN 1 
                    ELSE rate_limits.count + 1 
                END,
                window_start = CASE 
                    WHEN rate_limits.window_start < NOW() - (${windowSeconds} || ' seconds')::INTERVAL 
                    THEN NOW() 
                    ELSE rate_limits.window_start 
                END
            RETURNING count, window_start, 
                (window_start = NOW()) as is_new_window
        `;

    const record = result[0];
    const resetAt = new Date(record.window_start.getTime() + windowMs);
    const remaining = Math.max(0, maxRequests - record.count);
    const allowed = record.count <= maxRequests;

    return { allowed, remaining, resetAt };
  } catch (error) {
    // If table doesn't exist, allow the request but log warning
    logError('Rate limit check failed (table may not exist):', error);
    return {
      allowed: true,
      remaining: maxRequests,
      resetAt: new Date(Date.now() + windowMs),
    };
  }
}

// Basic IP format validation — IPv4 or IPv6 (loose check, rejects obvious garbage)
const IP_PATTERN = /^[\da-fA-F.:]+$/;

function isPlausibleIp(value: string): boolean {
  return value.length <= 45 && IP_PATTERN.test(value);
}

/**
 * Get client IP from request headers.
 *
 * Trusting proxy-injected headers is only safe when a known trusted proxy sits in front
 * of this server and strips or overwrites those headers before forwarding requests.
 * Set TRUSTED_PROXY_MODE to opt in:
 *
 *   TRUSTED_PROXY_MODE=cloudflare  — trust cf-connecting-ip (Cloudflare edge)
 *   TRUSTED_PROXY_MODE=nginx       — trust x-real-ip / x-forwarded-for (Nginx real_ip_header)
 *
 * Without TRUSTED_PROXY_MODE set, no proxy headers are trusted: all requests appear
 * as 127.0.0.1, which means rate limits apply per-process rather than per-client IP.
 * In that configuration, prefer session/user-keyed rate limits for authenticated endpoints.
 *
 * WARNING: Do not set TRUSTED_PROXY_MODE unless you have confirmed that your proxy
 * strips or overwrites the corresponding headers on every inbound request. Failing to
 * do so allows clients to spoof their IP and bypass rate limits.
 */
export function getClientIp(request: Request): string {
  return getClientIpFromHeaders(request.headers);
}

/**
 * Same resolution as {@link getClientIp}, for callers that only have headers rather than a
 * Request — server components reading `await headers()`.
 */
export function getClientIpFromHeaders(headers: Headers): string {
  const mode = process.env.TRUSTED_PROXY_MODE?.trim().toLowerCase();

  if (mode === 'cloudflare') {
    // cf-connecting-ip is injected by Cloudflare and cannot be set by clients
    // when origin access is restricted to Cloudflare's IP ranges.
    const cfIp = headers.get('cf-connecting-ip');
    if (cfIp && isPlausibleIp(cfIp)) {
      return cfIp;
    }
  }

  if (mode === 'nginx') {
    // x-real-ip is set by Nginx's real_ip_header directive (connection-level, not spoofable
    // by clients when set_real_ip_from is configured for the upstream proxy).
    const realIp = headers.get('x-real-ip');
    if (realIp && isPlausibleIp(realIp)) return realIp;

    // x-forwarded-for last entry added by Nginx when proxy_add_x_forwarded_for is used.
    const forwardedFor = headers.get('x-forwarded-for');
    if (forwardedFor) {
      const entries = forwardedFor.split(',');
      const last = entries[entries.length - 1].trim();
      if (isPlausibleIp(last)) return last;
    }
  }

  // No trusted proxy configured — fall back to a constant value.
  // Rate limiting will apply per-process; use userId-keyed limits for authenticated endpoints.
  return '127.0.0.1';
}

/**
 * Create rate limit headers for response
 */
export function rateLimitHeaders(result: RateLimitResult, maxRequests: number): HeadersInit {
  return {
    'X-RateLimit-Limit': maxRequests.toString(),
    'X-RateLimit-Remaining': result.remaining.toString(),
    'X-RateLimit-Reset': Math.floor(result.resetAt.getTime() / 1000).toString(),
  };
}

/**
 * Cleanup old rate limit entries (call periodically)
 */
export async function cleanupRateLimits(): Promise<void> {
  try {
    await db.$executeRaw`SELECT cleanup_rate_limits()`;
  } catch (error) {
    logError('Rate limit cleanup failed:', error);
  }
}

// Start cleanup interval once per process to avoid duplicate scheduling on module reload.
if (
  !globalForRateLimitCleanup.rateLimitCleanupIntervalStarted &&
  typeof setInterval !== 'undefined'
) {
  const interval = setInterval(() => {
    cleanupRateLimits().catch((err) => logError('Unexpected error:', err));
  }, RATE_LIMIT_CLEANUP_INTERVAL_MS);

  // Avoid keeping Node.js process alive because of housekeeping timers.
  interval.unref?.();
  globalForRateLimitCleanup.rateLimitCleanupIntervalStarted = true;
}

/**
 * One-call rate limit check that returns a 429 NextResponse if blocked, or null if allowed.
 * Use at the top of any API handler:
 *   const limited = await rateLimit(request, 'comment');
 *   if (limited) return limited;
 */
export async function rateLimit(
  request: Request,
  action: string,
  config?: RateLimitConfig
): Promise<NextResponse | null> {
  const ip = getClientIp(request);
  const cfg = config || RATE_LIMIT_CONFIGS[action] || RATE_LIMIT_CONFIGS.api;
  const result = await checkRateLimit(ip, action, cfg);

  if (!result.allowed) {
    return NextResponse.json(
      { error: 'Too many requests. Please try again later.' },
      {
        status: 429,
        headers: rateLimitHeaders(result, cfg.maxRequests),
      }
    );
  }

  return null;
}