diff --git a/app/api/integration/membership/route.ts b/app/api/integration/membership/route.ts index b643c17..50fbb4a 100644 --- a/app/api/integration/membership/route.ts +++ b/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 diff --git a/app/api/integration/verify/route.ts b/app/api/integration/verify/route.ts index 5aa4cd0..f4aa5a2 100644 --- a/app/api/integration/verify/route.ts +++ b/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 diff --git a/docs/deployment.md b/docs/deployment.md index 62fc1a2..2d56c93 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -25,8 +25,15 @@ 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. @@ -34,8 +41,60 @@ state is lost on cold starts. To keep a true global limit, replace the in-memory 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 diff --git a/lib/rate-limit.ts b/lib/rate-limit.ts index 36bdf32..b1ee664 100644 --- a/lib/rate-limit.ts +++ b/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() + + 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() - -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 { + 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 { 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 } diff --git a/test/rate-limit.test.ts b/test/rate-limit.test.ts index 9dfdfc3..7498543 100644 --- a/test/rate-limit.test.ts +++ b/test/rate-limit.test.ts @@ -1,6 +1,6 @@ import { describe, test, afterEach } from 'node:test' import * as assert from 'node:assert/strict' -import { rateLimitRequest } from '../lib/rate-limit' +import { rateLimitRequest, InMemoryRateLimitStore, type RateLimitStore } from '../lib/rate-limit' function makeReq(opts: { ip?: string; address?: string } = {}): Request { const headers = new Headers() @@ -16,45 +16,161 @@ afterEach(() => { }) describe('rateLimitRequest', () => { - test('allows requests under the limit', () => { + test('allows requests under the limit', async () => { const req = makeReq({ ip: '10.0.0.1', address: 'GA1' }) - const r = rateLimitRequest(req, 'GA1') + const r = await rateLimitRequest(req, 'GA1') assert.equal(r.limited, false) assert.ok(r.remaining >= 0) }) - test('returns limited=true once the per-IP bucket is exhausted', () => { + test('returns limited=true once the per-IP bucket is exhausted', async () => { const ip = '10.0.0.2' let last for (let i = 0; i < 30; i++) { - last = rateLimitRequest(makeReq({ ip }), 'GA2') + last = await rateLimitRequest(makeReq({ ip }), 'GA2') assert.equal(last.limited, false, `request ${i} should pass`) } // 31st request exceeds the 30/min bucket - const over = rateLimitRequest(makeReq({ ip }), 'GA2') + const over = await rateLimitRequest(makeReq({ ip }), 'GA2') assert.equal(over.limited, true) assert.ok(over.retryAfter > 0, 'retryAfter should be a positive number of seconds') }) - test('keys IP and wallet independently', () => { + test('keys IP and wallet independently', async () => { const ip = '10.0.0.3' // exhaust wallet key with one address for (let i = 0; i < 30; i++) { - rateLimitRequest(makeReq({ ip, address: 'GA3' }), 'GA3') + await rateLimitRequest(makeReq({ ip, address: 'GA3' }), 'GA3') } - const walletOver = rateLimitRequest(makeReq({ ip, address: 'GA3' }), 'GA3') + const walletOver = await rateLimitRequest(makeReq({ ip, address: 'GA3' }), 'GA3') assert.equal(walletOver.limited, true) // a different wallet from the SAME ip still has its own bucket state // (ip bucket also exhausted at 30, so this exercises wallet-keying separately) - const otherWallet = rateLimitRequest(makeReq({ ip, address: 'GA4' }), 'GA4') + const otherWallet = await rateLimitRequest(makeReq({ ip, address: 'GA4' }), 'GA4') // ip is exhausted, so still limited — but retryAfter comes from ip bucket assert.equal(otherWallet.limited, true) }) - test('different IPs do not share bucket state', () => { - const a = rateLimitRequest(makeReq({ ip: '10.0.0.4' }), null) - const b = rateLimitRequest(makeReq({ ip: '10.0.0.5' }), null) + test('different IPs do not share bucket state', async () => { + const a = await rateLimitRequest(makeReq({ ip: '10.0.0.4' }), null) + const b = await rateLimitRequest(makeReq({ ip: '10.0.0.5' }), null) assert.equal(a.limited, false) assert.equal(b.limited, false) }) }) + +// =========================================================================== +// InMemoryRateLimitStore — deterministic tests against a directly +// constructed store with a caller-controlled `now`, so refill/consume math +// can be asserted exactly without depending on real wall-clock timing. +// =========================================================================== + +describe('InMemoryRateLimitStore', () => { + test('consumes one token per call and reports the exact residual token count', async () => { + const store = new InMemoryRateLimitStore(5, 5 / 60_000) + const now = 0 + const r1 = await store.take('k', now) + assert.deepEqual(r1, { tokens: 4, consumed: true }) + const r2 = await store.take('k', now) + assert.deepEqual(r2, { tokens: 3, consumed: true }) + }) + + test('reports consumed=false with the exact residual token count once exhausted', async () => { + const store = new InMemoryRateLimitStore(2, 2 / 60_000) + const now = 0 + await store.take('k', now) + await store.take('k', now) + const r = await store.take('k', now) + assert.equal(r.consumed, false) + assert.equal(r.tokens, 0) + }) + + test('refills exactly the expected number of tokens after elapsed time', async () => { + // 60 tokens per 60_000ms => 1 token refilled per 1000ms elapsed + const store = new InMemoryRateLimitStore(60, 60 / 60_000) + await store.take('k', 0) // tokens: 59 + const r = await store.take('k', 1_000) // +1 refilled, then -1 consumed => back to 59 + assert.equal(r.consumed, true) + assert.equal(r.tokens, 59) + }) + + test('does not refill beyond maxTokens', async () => { + const store = new InMemoryRateLimitStore(10, 10 / 60_000) + await store.take('k', 0) // tokens: 9 + const r = await store.take('k', 10_000_000) // huge elapsed time — refill clamps at 10, then -1 + assert.equal(r.consumed, true) + assert.equal(r.tokens, 9) + }) + + test('different keys have independent state within the same store instance', async () => { + const store = new InMemoryRateLimitStore(1, 1 / 60_000) + const a = await store.take('a', 0) + const b = await store.take('b', 0) + assert.equal(a.consumed, true) + assert.equal(b.consumed, true) + }) +}) + +// =========================================================================== +// rateLimitRequest with an injected fake store — pins down key-check order, +// short-circuit behavior, and the exact retryAfter/remaining derivation +// without needing real time or 30 real requests to exhaust a bucket. +// =========================================================================== + +describe('rateLimitRequest with an injected store', () => { + test('checks the IP bucket before the wallet bucket and short-circuits on an IP limit', async () => { + const calls: string[] = [] + const fakeStore: RateLimitStore = { + async take(key) { + calls.push(key) + return { tokens: 0, consumed: false } + }, + } + const req = makeReq({ ip: '10.0.0.9' }) + const result = await rateLimitRequest(req, 'GAX', fakeStore) + assert.equal(result.limited, true) + assert.deepEqual(calls, ['ip:10.0.0.9']) + }) + + test('falls through to the wallet bucket only when the IP bucket allows the request', async () => { + const calls: string[] = [] + const fakeStore: RateLimitStore = { + async take(key) { + calls.push(key) + return { tokens: 5, consumed: true } + }, + } + const req = makeReq({ ip: '10.0.0.10' }) + const result = await rateLimitRequest(req, 'GAY', fakeStore) + assert.equal(result.limited, false) + assert.deepEqual(calls, ['ip:10.0.0.10', 'wallet:GAY']) + }) + + test('returns the IP result (not the wallet result) when both buckets allow the request', async () => { + const fakeStore: RateLimitStore = { + async take(key) { + if (key.startsWith('ip:')) return { tokens: 7, consumed: true } + return { tokens: 2, consumed: true } + }, + } + const req = makeReq({ ip: '10.0.0.11' }) + const result = await rateLimitRequest(req, 'GAZ', fakeStore) + assert.equal(result.limited, false) + assert.equal(result.remaining, 7) // from the IP bucket, not the wallet bucket (2) + }) + + test('computes retryAfter from the exact deficit reported by the store', async () => { + const fakeStore: RateLimitStore = { + async take() { + return { tokens: 0.5, consumed: false } // deficit = 1 - 0.5 = 0.5 + }, + } + const req = makeReq({ ip: '10.0.0.12' }) + const result = await rateLimitRequest(req, null, fakeStore) + assert.equal(result.limited, true) + assert.equal(result.remaining, 0) + // production REFILL_PER_MS = 30 / 60_000 = 0.0005 + // retryAfter = ceil(0.5 / 0.0005 / 1000) = ceil(1) = 1 + assert.equal(result.retryAfter, 1) + }) +})