Merge pull request #72 from bigcommerce/codex/implement-rate-limits-for-generatedescription

Add rate limiting to generate description API

Author: bc-rmalyavc

src/app/api/generateDescription/route.ts

@@ -2,6 +2,11 @@ import { type NextRequest, NextResponse } from 'next/server';
 import generateDescription from '~/server/google-ai';
 import { aiSchema } from './schema';
 import { authorize } from '~/lib/authorize';
+import { enforceRateLimit } from '~/lib/rate-limit';
+
+const RATE_LIMIT_KEY = 'generateDescription';
+const RATE_LIMIT_MAX_REQUESTS = 60;
+const RATE_LIMIT_WINDOW_MS = 60_000;
 
 export async function POST(req: NextRequest) {
   const authToken = req.headers.get('X-Auth-Token') || 'missing';
@@ -10,6 +15,27 @@ export async function POST(req: NextRequest) {
     return new NextResponse('Unauthorized', { status: 401 });
   }
 
+  const clientIdentifier =
+    req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || req.ip || 'unknown';
+  const rateLimit = await enforceRateLimit(
+    `${RATE_LIMIT_KEY}:${clientIdentifier}`,
+    { windowMs: RATE_LIMIT_WINDOW_MS, maxRequests: RATE_LIMIT_MAX_REQUESTS }
+  );
+
+  if (!rateLimit.allowed) {
+    const retryAfter = Math.max(0, Math.ceil((rateLimit.reset - Date.now()) / 1000));
+
+    return new NextResponse('Too Many Requests', {
+      status: 429,
+      headers: {
+        'Retry-After': retryAfter.toString(),
+        'X-RateLimit-Limit': rateLimit.limit.toString(),
+        'X-RateLimit-Remaining': rateLimit.remaining.toString(),
+        'X-RateLimit-Reset': rateLimit.reset.toString(),
+      },
+    });
+  }
+
   const data: unknown = await req.json();
   const parsedParams = aiSchema.safeParse(data);
 
@@ -19,5 +45,11 @@ export async function POST(req: NextRequest) {
 
   const description = await generateDescription(parsedParams.data);
 
-  return NextResponse.json({ description });
+  const response = NextResponse.json({ description });
+
+  response.headers.set('X-RateLimit-Limit', rateLimit.limit.toString());
+  response.headers.set('X-RateLimit-Remaining', rateLimit.remaining.toString());
+  response.headers.set('X-RateLimit-Reset', rateLimit.reset.toString());
+
+  return response;
 }

src/lib/db.ts

@@ -35,7 +35,7 @@ export interface ClientTokenData {
 let app: FirebaseApp;
 let db: Firestore;
 
-function getDb() {
+export function getDb() {
   if (!db) {
     const { FIRE_API_KEY, FIRE_DOMAIN, FIRE_PROJECT_ID } = env;
 

src/lib/rate-limit.ts

@@ -0,0 +1,69 @@
+import {
+  collection,
+  doc,
+  getDocsFromServer,
+  increment,
+  setDoc,
+  Timestamp,
+} from 'firebase/firestore';
+import { getDb } from './db';
+
+interface RateLimitOptions {
+  windowMs?: number;
+  maxRequests?: number;
+}
+
+interface RateLimitDocument {
+  count: number;
+  windowStart: number;
+  expiresAt?: Timestamp;
+}
+
+export interface RateLimitState {
+  allowed: boolean;
+  remaining: number;
+  reset: number;
+  limit: number;
+}
+
+const DEFAULT_WINDOW_MS = 60_000;
+const DEFAULT_MAX_REQUESTS = 30;
+const SHARD_COUNT = 20;
+// Documents are cleaned up via Firestore TTL configured on the `expiresAt` field.
+// Keeping a modest retention allows slow TTL sweeps without impacting active windows.
+const RETENTION_WINDOWS = 24; // number of windows to keep for TTL cleanup
+
+export async function enforceRateLimit(
+  key: string,
+  { windowMs = DEFAULT_WINDOW_MS, maxRequests = DEFAULT_MAX_REQUESTS }: RateLimitOptions = {}
+): Promise<RateLimitState> {
+  const db = getDb();
+  const now = Date.now();
+  const windowStart = Math.floor(now / windowMs) * windowMs;
+  const windowRef = collection(db, 'rateLimits', `${key}:${windowStart}`, 'shards');
+  const shardId = Math.floor(Math.random() * SHARD_COUNT).toString();
+  const shardRef = doc(windowRef, shardId);
+  const expiresAt = Timestamp.fromMillis(windowStart + windowMs * RETENTION_WINDOWS);
+
+  await setDoc(
+    shardRef,
+    { count: increment(1), windowStart, expiresAt },
+    { merge: true },
+  );
+
+  const shardsSnapshot = await getDocsFromServer(windowRef);
+  const count = shardsSnapshot.docs.reduce((sum, snapshot) => {
+    const data = snapshot.data() as RateLimitDocument | undefined;
+
+    return sum + (data?.count ?? 0);
+  }, 0);
+  const remaining = Math.max(0, maxRequests - count);
+  const reset = windowStart + windowMs;
+
+  return {
+    allowed: count <= maxRequests,
+    remaining,
+    reset,
+    limit: maxRequests,
+  };
+}