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,56 @@
+import { doc, runTransaction } from 'firebase/firestore';
+import { getDb } from './db';
+
+interface RateLimitOptions {
+  windowMs?: number;
+  maxRequests?: number;
+}
+
+interface RateLimitDocument {
+  count: number;
+  windowStart: number;
+}
+
+export interface RateLimitState {
+  allowed: boolean;
+  remaining: number;
+  reset: number;
+  limit: number;
+}
+
+const DEFAULT_WINDOW_MS = 60_000;
+const DEFAULT_MAX_REQUESTS = 30;
+
+export async function enforceRateLimit(
+  key: string,
+  { windowMs = DEFAULT_WINDOW_MS, maxRequests = DEFAULT_MAX_REQUESTS }: RateLimitOptions = {}
+): Promise<RateLimitState> {
+  const db = getDb();
+  const ref = doc(db, 'rateLimits', key);
+  const now = Date.now();
+
+  let state: RateLimitState = {
+    allowed: true,
+    remaining: maxRequests,
+    reset: now + windowMs,
+    limit: maxRequests,
+  };
+
+  await runTransaction(db, async (transaction) => {
+    const snapshot = await transaction.get(ref);
+    const data = snapshot.data() as RateLimitDocument | undefined;
+    const withinWindow = data && now - data.windowStart < windowMs;
+
+    const windowStart = withinWindow ? data.windowStart : now;
+    const newCount = (withinWindow ? data.count : 0) + 1;
+    const remaining = Math.max(0, maxRequests - newCount);
+    const reset = windowStart + windowMs;
+    const allowed = newCount <= maxRequests;
+
+    transaction.set(ref, { count: newCount, windowStart });
+
+    state = { allowed, remaining, reset, limit: maxRequests };
+  });
+
+  return state;
+}