Use atomic increments for rate limit counters

Author: bc-rmalyavc

src/lib/rate-limit.ts

@@ -1,4 +1,4 @@
-import { doc, runTransaction } from 'firebase/firestore';
+import { doc, getDoc, increment, setDoc } from 'firebase/firestore';
 import { getDb } from './db';
 
 interface RateLimitOptions {
@@ -26,31 +26,26 @@ export async function enforceRateLimit(
   { 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,
+  const windowStart = Math.floor(now / windowMs) * windowMs;
+  const ref = doc(db, 'rateLimits', `${key}:${windowStart}`);
+
+  await setDoc(
+    ref,
+    { count: increment(1), windowStart },
+    { merge: true },
+  );
+
+  const snapshot = await getDoc(ref);
+  const data = snapshot.data() as RateLimitDocument | undefined;
+  const count = data?.count ?? 0;
+  const remaining = Math.max(0, maxRequests - count);
+  const reset = windowStart + windowMs;
+
+  return {
+    allowed: count <= maxRequests,
+    remaining,
+    reset,
     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;
 }