Merge pull request #1301 from Wiseman52/fix/issue-1168-persist-ratelimit-counters

fix(ratelimit): persist rate-limit counters across process restarts #1168

Author: RUKAYAT-CODER

src/app/api/tutorials/__tests__/ratelimit.test.ts

@@ -9,6 +9,11 @@ import {
   resetIdentityRateLimit,
 } from '@/lib/ratelimit';
 
+// Mock the DB pool to prevent real database calls during tests
+vi.mock('@/lib/db/pool', () => ({
+  query: vi.fn().mockResolvedValue({ rows: [], rowCount: 0 }),
+}));
+
 // ---------------------------------------------------------------------------
 // Helpers
 // ---------------------------------------------------------------------------

src/lib/db/migrations/006_create_rate_limits_table.sql

@@ -0,0 +1,14 @@
+-- Migration: Create rate_limits table
+-- Description: Persists rate-limit counters across process restarts so that
+--              in-memory counters no longer reset on every deploy.
+
+CREATE TABLE IF NOT EXISTS rate_limits (
+  identifier VARCHAR(512) PRIMARY KEY,
+  count INTEGER NOT NULL DEFAULT 0,
+  reset_at BIGINT NOT NULL,
+  created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+  updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+);
+
+-- Index for cleaning up expired entries
+CREATE INDEX IF NOT EXISTS idx_rate_limits_reset_at ON rate_limits(reset_at);

src/lib/ratelimit.test.ts

@@ -12,6 +12,11 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { parseTrustedProxyIPs } from '@/config/environment';
 import { slidingWindowRateLimit } from './ratelimit';
 
+// Mock the DB pool to prevent real database calls during tests
+vi.mock('@/lib/db/pool', () => ({
+  query: vi.fn().mockResolvedValue({ rows: [], rowCount: 0 }),
+}));
+
 // ---------------------------------------------------------------------------
 // Helpers
 // ---------------------------------------------------------------------------

src/lib/ratelimit.ts

@@ -1,10 +1,38 @@
 import { NextResponse } from 'next/server';
 import { getTrustedProxyConfig } from '@/config/environment';
 
+type DbQueryResult = { rows: Array<Record<string, unknown>> };
+type DbQueryFn = (text: string, params?: unknown[]) => Promise<DbQueryResult>;
+
+let dbQueryPromise: Promise<DbQueryFn | null> | null = null;
+
+/**
+ * Lazily resolves the pg-backed `query` helper.
+ *
+ * `pg` can only run in a Node.js runtime and importing it from a module shared
+ * with Edge routes would break their bundling, so the database module is
+ * loaded dynamically and only outside the Edge runtime. When the DB is
+ * unavailable the in-memory cache remains authoritative.
+ */
+async function loadDbQuery(): Promise<DbQueryFn | null> {
+  if (process.env.NEXT_RUNTIME === 'edge') return null;
+  if (!dbQueryPromise) {
+    dbQueryPromise = import(/* webpackIgnore: true */ '@/lib/db/pool')
+      .then((mod) => mod.query as DbQueryFn)
+      .catch(() => null);
+  }
+  return dbQueryPromise;
+}
+
 /**
- * In-memory sliding window rate limiter for API routes.
+ * Database-backed sliding window rate limiter for API routes.
  * Provides IP-based rate limiting with configurable limits and windows.
  *
+ * Counters are persisted to PostgreSQL so they survive process restarts.
+ * An in-memory Map serves as a fast synchronous cache; every write is also
+ * persisted to the database asynchronously (fire-and-forget) so that
+ * subsequent processes can pick up the state after a deploy.
+ *
  * Security: getClientIP() only trusts x-forwarded-for / x-real-ip headers when
  * the request arrives from a proxy listed in TRUSTED_PROXY_IPS (see
  * src/config/environment.ts). When no proxies are configured the headers are
@@ -32,8 +60,85 @@ interface RateLimitEntry {
   resetAt: number;
 }
 
+/**
+ * Fast in-memory cache used as the synchronous hot path.
+ * Writes are also persisted to the database asynchronously.
+ */
 const stores = new Map<string, RateLimitEntry>();
 
+/**
+ * Persists a single rate-limit entry to the database.
+ * Errors are silently swallowed so that a DB outage never blocks request
+ * processing — the in-memory cache still provides best-effort limiting.
+ */
+async function persistToDb(identifier: string, entry: RateLimitEntry): Promise<void> {
+  const query = await loadDbQuery();
+  if (!query) return;
+  try {
+    await query(
+      `INSERT INTO rate_limits (identifier, count, reset_at, updated_at)
+       VALUES ($1, $2, $3, NOW())
+       ON CONFLICT (identifier) DO UPDATE
+         SET count = EXCLUDED.count,
+             reset_at = EXCLUDED.reset_at,
+             updated_at = NOW()`,
+      [identifier, entry.count, entry.resetAt],
+    );
+  } catch {
+    // Silently ignore — the in-memory store is still authoritative for
+    // the current process, and DB unavailability should not break requests.
+  }
+}
+
+/**
+ * Removes an expired entry from the database.
+ */
+async function removeFromDb(identifier: string): Promise<void> {
+  const query = await loadDbQuery();
+  if (!query) return;
+  try {
+    await query('DELETE FROM rate_limits WHERE identifier = $1', [identifier]);
+  } catch {
+    // Silently ignore.
+  }
+}
+
+/**
+ * Loads all non-expired rate-limit entries from the database into the
+ * in-memory cache on process startup.  Called once at module load time.
+ */
+async function loadFromDb(): Promise<void> {
+  const query = await loadDbQuery();
+  if (!query) return;
+  try {
+    const now = Date.now();
+    const result = await query(
+      'SELECT identifier, count, reset_at FROM rate_limits WHERE reset_at > $1',
+      [now],
+    );
+    for (const row of result.rows) {
+      stores.set(row.identifier as string, {
+        count: row.count as number,
+        resetAt: row.reset_at as number,
+      });
+    }
+  } catch {
+    // DB may not be available yet (e.g. during build or early startup).
+    // The in-memory store will be used as a fallback.
+  }
+}
+
+// Kick off the load-on-startup — non-blocking.
+void loadFromDb();
+
+/**
+ * Synchronous sliding-window rate limiter backed by an in-memory cache
+ * with asynchronous database persistence.
+ *
+ * The function signature is intentionally synchronous so that the 30+
+ * existing call-sites (withRateLimit, certificate routes, etc.) do not
+ * need to become async.
+ */
 export function slidingWindowRateLimit(
   identifier: string,
   config: RateLimitConfig,
@@ -44,9 +149,12 @@ export function slidingWindowRateLimit(
   if (!entry || entry.resetAt <= now) {
     if (entry) {
       stores.delete(identifier);
+      void removeFromDb(identifier);
     }
     const resetAt = now + config.windowMs;
-    stores.set(identifier, { count: 1, resetAt });
+    const newEntry: RateLimitEntry = { count: 1, resetAt };
+    stores.set(identifier, newEntry);
+    void persistToDb(identifier, newEntry);
     return {
       success: true,
       remaining: config.limit - 1,
@@ -68,6 +176,7 @@ export function slidingWindowRateLimit(
 
   entry.count += 1;
   stores.set(identifier, entry);
+  void persistToDb(identifier, entry);
 
   return {
     success: true,
@@ -232,4 +341,4 @@ export function withRateLimit<T extends Request>(
     addHeaders,
     rateLimitResponse: createRateLimitResponse(result),
   };
-}
+}