Add rate limiting to generate description API

Use atomic increments for rate limit counters

Use server reads for rate limit counts

Shard rate limit counters to reduce contention

Add TTL expiration to rate limit shards

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,
+  };
+}