Extract a pluggable, multi-instance-safe RateLimitStore abstraction

app/api/integration/membership/route.ts

@@ -37,7 +37,7 @@ export async function GET(req: NextRequest) {
     }
 
     // Enforce per-IP / per-wallet rate limit before any downstream call.
-    const rl = rateLimitRequest(req, address)
+    const rl = await rateLimitRequest(req, address)
     rateLimit = rl.limited ? 'limited' : 'allowed'
     if (rl.limited) {
       status = 429

app/api/integration/verify/route.ts

@@ -37,7 +37,7 @@ export async function GET(req: NextRequest) {
     }
 
     // Enforce per-IP / per-wallet rate limit before any downstream call.
-    const rl = rateLimitRequest(req, address)
+    const rl = await rateLimitRequest(req, address)
     rateLimit = rl.limited ? 'limited' : 'allowed'
     if (rl.limited) {
       status = 429

docs/deployment.md

@@ -25,17 +25,76 @@ key.
 If you run **more than one instance** (horizontal scaling, containers behind a
 load balancer) or a **serverless / edge** runtime, each process keeps its own
 counters, so the effective limit is multiplied by the number of instances and
-state is lost on cold starts. To keep a true global limit, replace the in-memory
-`Map` store in `lib/rate-limit.ts` with a shared backend:
+state is lost on cold starts. To keep a true global limit, implement the
+`RateLimitStore` interface exported from `lib/rate-limit.ts` against a shared
+backend and inject it into `rateLimitRequest()`:
+
+```ts
+export interface RateLimitStore {
+  take(key: string, now: number): Promise<{ tokens: number; consumed: boolean }>
+}
+```
 
 - **Redis (recommended):** `@upstash/ratelimit` + `@upstash/redis`, or
   `ioredis` with a Lua token-bucket script for atomic decrement.
 - **Database counter:** a small row per key with `UPDATE … SET tokens = …`
   guarded by a transaction.
 - **Edge KV:** `@vercel/kv` or Cloudflare Workers KV with a TTL-backed counter.
 
-The `rateLimitRequest()` signature stays the same — only the bucket store
-behind `getBucket()` / `take()` needs to be swapped for a shared implementation.
+**`take()` must be atomic.** Two separate remote calls — a `GET` to read the
+current bucket followed by a client-side computation and a `SET` to write it
+back — are **not sufficient**, even though each call is individually async.
+Two concurrent callers (two instances, or two in-flight requests racing
+across an `await` on one instance) can both `GET` the same starting state,
+compute independently, and both `SET` — silently losing one of the two token
+consumptions. That race defeats the entire purpose of a shared store. A real
+implementation must push the read-refill-consume-write into a single atomic
+operation on the backend: a Redis Lua script (`EVAL`) or `MULTI`/`EXEC`
+transaction, a database transaction, or an equivalent atomic primitive.
+
+Example wiring for the gateway routes, once a `RedisRateLimitStore` exists:
+
+```ts
+// lib/rate-limit-redis.ts (illustrative — not part of this repo)
+import { RateLimitStore } from '@/lib/rate-limit'
+import { redis } from '@/lib/redis-client'
+
+const TAKE_SCRIPT = `
+  -- KEYS[1] = bucket key, ARGV[1] = now (ms),
+  -- ARGV[2] = maxTokens, ARGV[3] = refillPerMs
+  -- Reads, refills, conditionally decrements, and writes back atomically.
+  ...
+`
+
+export class RedisRateLimitStore implements RateLimitStore {
+  constructor(
+    private readonly maxTokens: number,
+    private readonly refillPerMs: number,
+  ) {}
+
+  async take(key: string, now: number) {
+    const [tokens, consumed] = await redis.eval(
+      TAKE_SCRIPT,
+      [key],
+      [now, this.maxTokens, this.refillPerMs],
+    )
+    return { tokens, consumed: consumed === 1 }
+  }
+}
+```
+
+```ts
+// app/api/integration/membership/route.ts
+import { rateLimitRequest } from '@/lib/rate-limit'
+import { sharedRateLimitStore } from '@/lib/rate-limit-redis-instance'
+
+const rl = await rateLimitRequest(req, address, sharedRateLimitStore)
+```
+
+`rateLimitRequest()` is `async` and accepts the store as an optional third
+argument, defaulting to the built-in `InMemoryRateLimitStore` — existing
+callers that omit the argument keep today's single-instance behavior
+unchanged.
 
 ## Troubleshooting
 

lib/rate-limit.ts

@@ -1,14 +1,15 @@
 /**
- * In-memory token-bucket rate limiter for the integration gateway routes.
+ * Pluggable token-bucket rate limiter for the integration gateway routes.
  *
- * SINGLE-INSTANCE CAVEAT:
- * The bucket state lives in process memory. This is sufficient for a single
- * Next.js instance (one node/server process). If you run more than one
- * instance behind a load balancer, each instance keeps its own counters and
- * the effective limit is multiplied by the instance count. For multi-instance
- * or serverless (edge) deployments, replace the `Map` store with a shared
- * backend (e.g. Redis via @upstash/ratelimit, or a DB-backed counter) — see
- * docs/deployment.md "Production rate-limiting" for the upgrade path.
+ * The token-bucket refill/consume algorithm is decoupled from storage via
+ * the RateLimitStore interface. `InMemoryRateLimitStore` (the default)
+ * keeps bucket state in a per-process Map — sufficient for a single
+ * Next.js instance. If you run more than one instance behind a load
+ * balancer, each instance keeps its own counters and the effective limit is
+ * multiplied by the instance count. For multi-instance or serverless (edge)
+ * deployments, inject a shared RateLimitStore backed by Redis or another
+ * shared backend — see docs/deployment.md "Production rate-limiting" for
+ * the upgrade path and why the store's take() must be atomic.
  */
 
 export interface RateLimitResult {
@@ -19,44 +20,80 @@ export interface RateLimitResult {
   remaining: number
 }
 
+/**
+ * Storage contract for the token-bucket algorithm. A single atomic method
+ * on purpose: it must create-if-absent, refill for elapsed time, and
+ * consume one token as one indivisible operation.
+ *
+ * Implementations MUST perform this atomically against the backing store
+ * (e.g. a Redis Lua script or a MULTI/EXEC transaction). A plain "GET the
+ * current state, compute the next state in the client, SET it back" pair of
+ * remote calls is NOT sufficient: two concurrent callers (two instances, or
+ * two in-flight requests on one instance racing across an await) can both
+ * read the same starting state, both compute independently, and both write
+ * back — silently losing one of the two token consumptions. That failure
+ * mode defeats the entire purpose of a shared store, so a store must never
+ * be built that way.
+ */
+export interface RateLimitStore {
+  take(key: string, now: number): Promise<{ tokens: number; consumed: boolean }>
+}
+
 interface Bucket {
   tokens: number
   /** epoch ms of the last refill */
   last: number
 }
 
+/**
+ * Default RateLimitStore: per-process Map, sufficient for a single
+ * Next.js instance. See docs/deployment.md for the multi-instance upgrade
+ * path.
+ */
+export class InMemoryRateLimitStore implements RateLimitStore {
+  private readonly buckets = new Map<string, Bucket>()
+
+  constructor(
+    private readonly maxTokens: number,
+    private readonly refillPerMs: number,
+  ) {}
+
+  async take(key: string, now: number): Promise<{ tokens: number; consumed: boolean }> {
+    let bucket = this.buckets.get(key)
+    if (!bucket) {
+      bucket = { tokens: this.maxTokens, last: now }
+      this.buckets.set(key, bucket)
+    } else {
+      // refill based on elapsed time
+      const elapsed = now - bucket.last
+      if (elapsed > 0) {
+        bucket.tokens = Math.min(this.maxTokens, bucket.tokens + elapsed * this.refillPerMs)
+        bucket.last = now
+      }
+    }
+
+    if (bucket.tokens >= 1) {
+      bucket.tokens -= 1
+      return { tokens: bucket.tokens, consumed: true }
+    }
+    // not enough tokens — caller must wait for a full token to refill
+    return { tokens: bucket.tokens, consumed: false }
+  }
+}
+
 const WINDOW_MS = 60_000 // 1 minute
 const MAX_TOKENS = 30 // 30 requests / minute per key
 const REFILL_PER_MS = MAX_TOKENS / WINDOW_MS
 
 // keyed by `${scope}:${id}` — e.g. "ip:1.2.3.4" or "wallet:GA…"
-const buckets = new Map<string, Bucket>()
-
-function getBucket(key: string): Bucket {
-  let bucket = buckets.get(key)
-  const now = Date.now()
-  if (!bucket) {
-    bucket = { tokens: MAX_TOKENS, last: now }
-    buckets.set(key, bucket)
-    return bucket
-  }
-  // refill based on elapsed time
-  const elapsed = now - bucket.last
-  if (elapsed > 0) {
-    bucket.tokens = Math.min(MAX_TOKENS, bucket.tokens + elapsed * REFILL_PER_MS)
-    bucket.last = now
-  }
-  return bucket
-}
+const defaultStore = new InMemoryRateLimitStore(MAX_TOKENS, REFILL_PER_MS)
 
-function take(key: string): RateLimitResult {
-  const bucket = getBucket(key)
-  if (bucket.tokens >= 1) {
-    bucket.tokens -= 1
-    return { limited: false, retryAfter: 0, remaining: Math.floor(bucket.tokens) }
+async function take(store: RateLimitStore, key: string): Promise<RateLimitResult> {
+  const { tokens, consumed } = await store.take(key, Date.now())
+  if (consumed) {
+    return { limited: false, retryAfter: 0, remaining: Math.floor(tokens) }
   }
-  // not enough tokens — caller must wait for a full token to refill
-  const deficit = 1 - bucket.tokens
+  const deficit = 1 - tokens
   const retryAfter = Math.ceil(deficit / REFILL_PER_MS / 1000)
   return { limited: true, retryAfter, remaining: 0 }
 }
@@ -65,20 +102,21 @@ function take(key: string): RateLimitResult {
  * Rate-limit a request by IP and (when present) wallet address.
  * Either key exceeding the limit triggers a 429.
  */
-export function rateLimitRequest(
+export async function rateLimitRequest(
   req: Request,
   address?: string | null,
-): RateLimitResult {
+  store: RateLimitStore = defaultStore,
+): Promise<RateLimitResult> {
   const ip =
     req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() ||
     req.headers.get('x-real-ip') ||
     'unknown'
 
-  const ipResult = take(`ip:${ip}`)
+  const ipResult = await take(store, `ip:${ip}`)
   if (ipResult.limited) return ipResult
 
   if (address) {
-    const walletResult = take(`wallet:${address}`)
+    const walletResult = await take(store, `wallet:${address}`)
     if (walletResult.limited) return walletResult
   }