From 4cd8bb3901058871524c30eb5399ebb6f03d9d00 Mon Sep 17 00:00:00 2001 From: produckq-boop <351437814+produckq-boop@users.noreply.github.com> Date: Tue, 3 Nov 2026 10:45:07 +0700 Subject: [PATCH] feat(rate-limit): make tier quotas configurable and emit standard headers (#50) --- .env.example | 14 ++ README.md | 52 +++++ package.json | 3 +- src/index.ts | 10 + src/middleware/apiKeyAuth.ts | 28 ++- src/middleware/rateLimit.ts | 109 ++++++----- src/middleware/tokenBucket.ts | 198 +++++++++++++++---- tests/rate-limit.test.ts | 92 ++++++++- tests/token-bucket.test.ts | 345 ++++++++++++++++++++++++++++++++++ 9 files changed, 763 insertions(+), 88 deletions(-) create mode 100644 tests/token-bucket.test.ts diff --git a/.env.example b/.env.example index 6e4ff54..d196e3f 100644 --- a/.env.example +++ b/.env.example @@ -63,6 +63,20 @@ INDEXER_BATCH_SIZE=100 INDEXER_CATCHUP_WORKERS=4 # ─── Rate limiting ──────────────────────────────────────────────────────────── +# Token-bucket rate limit tiers (per-tier request quota and window in milliseconds). +# Optional RATE_LIMIT_TIER__BURST overrides burst allowance per tier. +RATE_LIMIT_TIER_UNAUTHENTICATED_LIMIT=60 +RATE_LIMIT_TIER_UNAUTHENTICATED_WINDOW_MS=60000 +RATE_LIMIT_TIER_FREE_LIMIT=1000 +RATE_LIMIT_TIER_FREE_WINDOW_MS=60000 +RATE_LIMIT_TIER_DEVELOPER_LIMIT=1000 +RATE_LIMIT_TIER_DEVELOPER_WINDOW_MS=60000 +RATE_LIMIT_TIER_PRO_LIMIT=10000 +RATE_LIMIT_TIER_PRO_WINDOW_MS=60000 +RATE_LIMIT_TIER_ENTERPRISE_LIMIT=60000 +RATE_LIMIT_TIER_ENTERPRISE_WINDOW_MS=60000 + +# Fallback / adaptive rate limiting RATE_LIMIT_WINDOW_MS=60000 RATE_LIMIT_MAX=100 RATE_LIMIT_PUBLIC_MAX=100 diff --git a/README.md b/README.md index 3abdd62..074b15e 100644 --- a/README.md +++ b/README.md @@ -171,6 +171,55 @@ Because the server caches keys, a typical rotation involves: *Note: Never log or commit the `JWT_PRIVATE_KEY` to version control. The server performs a startup validation check and will crash if the keys are missing or malformed when `NODE_ENV=production`.* +## Rate Limiting + +Every API request passes through `apiKeyAuth` followed by `tieredRateLimit` ([`src/middleware/tokenBucket.ts`](./src/middleware/tokenBucket.ts), [`src/middleware/rateLimit.ts`](./src/middleware/rateLimit.ts), [`src/middleware/apiKeyAuth.ts`](./src/middleware/apiKeyAuth.ts)). Quotas are enforced with an atomic Redis Lua token bucket when `REDIS_URL` is configured, and automatically fall back to an in-process token bucket when Redis is unavailable. + +### Tiers & Configuration + +Requests without an `X-Api-Key` header use the `unauthenticated` tier (keyed by client IP). Validated API keys resolve their tier from `DevApiKey.tier` (defaulting to `free`). Each tier's request limit and window duration are configurable via environment variables, falling back to the defaults below when unset or invalid: + +| Tier | Default Limit | Default Window | Default Burst | Limit Variable | Window Variable | +| ----------------- | ------------- | -------------- | ------------- | --------------------------------------- | ------------------------------------------- | +| `unauthenticated` | `60` | `60000` ms | `10` | `RATE_LIMIT_TIER_UNAUTHENTICATED_LIMIT` | `RATE_LIMIT_TIER_UNAUTHENTICATED_WINDOW_MS` | +| `free` | `1000` | `60000` ms | `50` | `RATE_LIMIT_TIER_FREE_LIMIT` | `RATE_LIMIT_TIER_FREE_WINDOW_MS` | +| `developer` | `1000` | `60000` ms | `50` | `RATE_LIMIT_TIER_DEVELOPER_LIMIT` | `RATE_LIMIT_TIER_DEVELOPER_WINDOW_MS` | +| `pro` | `10000` | `60000` ms | `200` | `RATE_LIMIT_TIER_PRO_LIMIT` | `RATE_LIMIT_TIER_PRO_WINDOW_MS` | +| `enterprise` | `60000` | `60000` ms | `500` | `RATE_LIMIT_TIER_ENTERPRISE_LIMIT` | `RATE_LIMIT_TIER_ENTERPRISE_WINDOW_MS` | + +Optional `RATE_LIMIT_TIER__BURST` variables can also be set to adjust per-tier burst allowances. + +### Endpoint Group Multipliers + +Resource-intensive route groups scale the effective quota for a request by a multiplier (`Math.ceil(limit * multiplier)`): + +| Route Pattern | Group | Multiplier | +| ------------------------------------- | ----------- | ---------- | +| `/api/v1/events*` | `events` | `1.0x` | +| `/api/v1/search*`, `/api/v1/nft/search*` | `search` | `0.5x` | +| `POST /api/v1/contracts*` | `contracts` | `0.17x` | +| `/api/v1/simulate*` | `simulate` | `0.08x` | +| `/ws*` | `websocket` | `0.05x` | +| All other endpoints | default | `1.0x` | + +### Per-Key & Admin Overrides + +- **API key override (`DevApiKey.rateLimitOverride`)** — Setting a positive integer on `DevApiKey.rateLimitOverride` overrides the key's base tier limit for all requests authenticated with that `X-Api-Key` (endpoint multipliers still apply on top of the custom limit). After updating a key record in PostgreSQL, call `invalidateKeyCache(keyHash)` so the updated limit takes effect immediately rather than after `KEY_CACHE_TTL_MS` (default `10000` ms). +- **Admin endpoint override** — `POST /api/v1/admin/rate-limits` accepts `{ identifier, endpoint, max, windowMs }` to persist targeted overrides in `RateLimitOverride` for the fallback limiter (`X-RateLimit-Policy: user-override`). + +### Response Headers + +Standard rate-limit headers are returned on all rate-limited routes so clients can track remaining quota and back off cleanly: + +| Header | Returned On | Description | +| --------------------- | --------------- | --------------------------------------------------------------------------- | +| `RateLimit-Limit` | Every response | Maximum requests permitted in the window for the active tier/override | +| `RateLimit-Remaining` | Every response | Remaining requests available in the current window (`0` when throttled) | +| `RateLimit-Reset` | Every response | Seconds until the rate-limit bucket replenishes to full capacity | +| `Retry-After` | `429` responses | Seconds (`>= 1`) the client should wait before retrying a throttled request | + +Legacy `X-RateLimit-Limit`, `X-RateLimit-Remaining`, `X-RateLimit-Reset` (Unix epoch seconds), and `X-RateLimit-Tier` headers are also emitted for backward compatibility. + ## Environment Variables | Variable | Default | Description | @@ -185,6 +234,9 @@ Because the server caches keys, a typical rotation involves: | `INDEXER_BATCH_SIZE` | `100` | Ledgers per batch | | `ADMIN_SECRET` | — | Bearer token required by every `/api/admin/*` route (indexer). See [`docs/ADMIN_AUTH.md`](./docs/ADMIN_AUTH.md). | | `METRICS_TOKEN` | — | Bearer/`X-Metrics-Token` required to scrape the API service's `GET /metrics` from outside the host. Unset restricts it to loopback callers. See [Metrics](#metrics). | +| `REDIS_URL` | _(unset)_ | Shared Redis URL for distributed token-bucket rate limiting and caching (falls back to in-memory when unset). See [Rate Limiting](#rate-limiting). | +| `RATE_LIMIT_TIER__LIMIT` | Tier default | Per-window request quota for `UNAUTHENTICATED`, `FREE`, `DEVELOPER`, `PRO`, or `ENTERPRISE`. See [Rate Limiting](#rate-limiting). | +| `RATE_LIMIT_TIER__WINDOW_MS` | `60000` | Window duration in milliseconds for `UNAUTHENTICATED`, `FREE`, `DEVELOPER`, `PRO`, or `ENTERPRISE`. See [Rate Limiting](#rate-limiting). | | `MOCK_DATA` | `false` | Gates experimental endpoints that fabricate demo data instead of a real integration. See [Experimental & Mock Data](#experimental--mock-data). `ENABLE_EXPERIMENTAL` is accepted as an alias. | ## Experimental & Mock Data diff --git a/package.json b/package.json index 42e1518..2f44874 100644 --- a/package.json +++ b/package.json @@ -23,8 +23,7 @@ "repair": "ts-node src/indexer/repair-run.ts", "archive": "ts-node src/archival/run.ts", "seed": "ts-node prisma/seed.ts", - "test": "DATABASE_URL=postgresql://test:test@localhost:5432/test TESTNET_DATABASE_URL=postgresql://test:test@localhost:5432/test STELLAR_NETWORK=testnet vitest run tests/reentrancy-fortress.test.ts tests/api/nlq.test.ts tests/arbitrage-engine.test.ts tests/indexer/token-metadata.test.ts tests/error-handling-integration.test.ts tests/build-queue.test.ts tests/indexer/decoder-parity.test.ts tests/tracer.test.ts tests/api/mock-data-gating.test.ts tests/api/reputation-leaderboard.test.ts", - "test": "DATABASE_URL=postgresql://test:test@localhost:5432/test TESTNET_DATABASE_URL=postgresql://test:test@localhost:5432/test STELLAR_NETWORK=testnet vitest run tests/reentrancy-fortress.test.ts tests/api/nlq.test.ts tests/arbitrage-engine.test.ts tests/indexer/token-metadata.test.ts tests/error-handling-integration.test.ts tests/build-queue.test.ts tests/indexer/decoder-parity.test.ts tests/tracer.test.ts tests/webhooks/dispatcher.test.ts", + "test": "DATABASE_URL=postgresql://test:test@localhost:5432/test TESTNET_DATABASE_URL=postgresql://test:test@localhost:5432/test STELLAR_NETWORK=testnet vitest run tests/reentrancy-fortress.test.ts tests/api/nlq.test.ts tests/arbitrage-engine.test.ts tests/indexer/token-metadata.test.ts tests/error-handling-integration.test.ts tests/build-queue.test.ts tests/indexer/decoder-parity.test.ts tests/tracer.test.ts tests/api/mock-data-gating.test.ts tests/api/reputation-leaderboard.test.ts tests/webhooks/dispatcher.test.ts tests/api-key-auth.test.ts tests/rate-limit.test.ts tests/token-bucket.test.ts", "test:watch": "vitest", "test:coverage": "vitest run --coverage", "test:ui": "vitest --ui", diff --git a/src/index.ts b/src/index.ts index c873892..aa276c7 100644 --- a/src/index.ts +++ b/src/index.ts @@ -130,6 +130,16 @@ app.use( origin: corsOrigin, methods: ['GET', 'POST', 'PUT', 'PATCH', 'DELETE', 'OPTIONS'], allowedHeaders: ['Content-Type', 'Authorization', 'X-Api-Key', 'X-Request-Id'], + exposedHeaders: [ + 'RateLimit-Limit', + 'RateLimit-Remaining', + 'RateLimit-Reset', + 'Retry-After', + 'X-RateLimit-Limit', + 'X-RateLimit-Remaining', + 'X-RateLimit-Reset', + 'X-RateLimit-Tier', + ], credentials: true, }), ); diff --git a/src/middleware/apiKeyAuth.ts b/src/middleware/apiKeyAuth.ts index d1117c6..6fe045b 100644 --- a/src/middleware/apiKeyAuth.ts +++ b/src/middleware/apiKeyAuth.ts @@ -49,6 +49,30 @@ declare global { } } +const VALID_RATE_LIMIT_TIERS = new Set([ + 'unauthenticated', + 'free', + 'developer', + 'pro', + 'enterprise', +]); + +export function normalizeRateLimitTier( + rawTier: string | null | undefined, + fallback: RateLimitTier = 'free', +): RateLimitTier { + if (!rawTier) return fallback; + const candidate = rawTier.trim().toLowerCase() as RateLimitTier; + return VALID_RATE_LIMIT_TIERS.has(candidate) ? candidate : fallback; +} + +function normalizeRateLimitOverride(rawOverride: number | null | undefined): number | undefined { + if (typeof rawOverride !== 'number' || !Number.isFinite(rawOverride) || rawOverride <= 0) { + return undefined; + } + return Math.floor(rawOverride); +} + function hashKey(raw: string): string { return crypto.createHash('sha256').update(raw).digest('hex'); } @@ -247,8 +271,8 @@ async function resolveApiKey(raw: string): Promise { id: record.id, keyName: record.name, developerId: record.developerId, - tier: (record.tier as RateLimitTier) ?? 'free', - rateLimitOverride: record.rateLimitOverride ?? undefined, + tier: normalizeRateLimitTier(record.tier, 'free'), + rateLimitOverride: normalizeRateLimitOverride(record.rateLimitOverride), allowedIps: Array.isArray(record.allowedIps) ? (record.allowedIps as string[]) : undefined, allowedEndpoints: Array.isArray(record.allowedEndpoints) ? (record.allowedEndpoints as string[]) diff --git a/src/middleware/rateLimit.ts b/src/middleware/rateLimit.ts index 7502377..e0b9b4c 100644 --- a/src/middleware/rateLimit.ts +++ b/src/middleware/rateLimit.ts @@ -5,9 +5,9 @@ import { prismaRead } from '../db'; import { logger } from '../logger'; import { checkTokenBucket, - RateLimitTier, + type RateLimitTier, setRateLimitRedisClient, - TokenBucketResult, + type TokenBucketResult, } from './tokenBucket'; const developerKeys = new Set((process.env.API_KEYS_DEVELOPER ?? '').split(',').filter(Boolean)); @@ -102,13 +102,15 @@ function applyAdaptiveThrottle( const currentMax = Math.min(throttledMax, tierConfigValue.max); res.setHeader('X-RateLimit-Warn', 'true'); res.setHeader('X-RateLimit-Predicted', `${currentMax}`); - req.app.locals.rateLimitPredictedMax = currentMax; + if (req.app?.locals) { + req.app.locals.rateLimitPredictedMax = currentMax; + } return { ...tierConfigValue, max: currentMax }; } function getRequestBucketKey(req: Request, tier: TierName, userIdentifier?: string): string { const endpoint = req.path || req.originalUrl || '/'; - const keySource = userIdentifier ?? req.ip ?? 'unknown'; + const keySource = userIdentifier ?? req.apiKey?.id ?? req.ip ?? 'unknown'; return `${tier}:${keySource}:${endpoint}`; } @@ -159,7 +161,6 @@ function buildLimiters(store?: Store): Limiters { } let legacyLimiters: Limiters = buildLimiters(); -let useTokenBucket = false; export async function initRateLimitStore(): Promise { const redisUrl = process.env.REDIS_URL; @@ -174,7 +175,6 @@ export async function initRateLimitStore(): Promise { await client.connect(); setRateLimitRedisClient(client as any); - useTokenBucket = true; const store = new RedisStore({ sendCommand: (...args: string[]) => (client as any).sendCommand(args), @@ -188,6 +188,23 @@ export async function initRateLimitStore(): Promise { } } +function writeRateLimitHeaders(res: Response, result: TokenBucketResult): number { + const nowSec = Math.floor(Date.now() / 1000); + const remaining = Math.max(0, result.remaining); + const resetSeconds = Math.max(0, result.resetAt - nowSec); + + res.setHeader('RateLimit-Limit', String(result.limit)); + res.setHeader('RateLimit-Remaining', String(remaining)); + res.setHeader('RateLimit-Reset', String(resetSeconds)); + + res.setHeader('X-RateLimit-Limit', String(result.limit)); + res.setHeader('X-RateLimit-Remaining', String(remaining)); + res.setHeader('X-RateLimit-Reset', String(result.resetAt)); + res.setHeader('X-RateLimit-Tier', result.tier); + + return resetSeconds; +} + function runLocalRateLimit(req: Request, res: Response, next: NextFunction): void { void (async () => { const apiKey = req.headers['x-api-key'] as string | undefined; @@ -195,7 +212,11 @@ function runLocalRateLimit(req: Request, res: Response, next: NextFunction): voi const userIdentifier = req.headers['x-user-id'] as string | undefined; const endpoint = req.path || req.originalUrl || '/'; const configOverride = userIdentifier ? await getUserOverride(userIdentifier, endpoint) : null; - const baseConfig = configOverride ?? normalizeTierConfig()[tier]; + const apiKeyOverride = + req.apiKey?.rateLimitOverride && req.apiKey.rateLimitOverride > 0 + ? { windowMs: DEFAULT_TIERS.public.windowMs, max: req.apiKey.rateLimitOverride } + : null; + const baseConfig = apiKeyOverride ?? configOverride ?? normalizeTierConfig()[tier]; const effectiveConfig = applyAdaptiveThrottle(req, res, baseConfig); const now = Date.now(); const bucketKey = getRequestBucketKey(req, tier, userIdentifier); @@ -212,15 +233,24 @@ function runLocalRateLimit(req: Request, res: Response, next: NextFunction): voi resetAt: now + effectiveConfig.windowMs, }; const remaining = Math.max(0, effectiveConfig.max - bucket.count); + const resetEpochSeconds = Math.ceil(bucket.resetAt / 1000); + const resetSeconds = Math.max(0, resetEpochSeconds - Math.floor(now / 1000)); + + res.setHeader('RateLimit-Limit', `${effectiveConfig.max}`); + res.setHeader('RateLimit-Remaining', `${remaining}`); + res.setHeader('RateLimit-Reset', `${resetSeconds}`); res.setHeader('X-RateLimit-Limit', `${effectiveConfig.max}`); res.setHeader('X-RateLimit-Remaining', `${remaining}`); - res.setHeader('X-RateLimit-Reset', `${Math.ceil(bucket.resetAt / 1000)}`); + res.setHeader('X-RateLimit-Reset', `${resetEpochSeconds}`); - if (configOverride) res.setHeader('X-RateLimit-Policy', 'user-override'); + if (configOverride || apiKeyOverride) res.setHeader('X-RateLimit-Policy', 'user-override'); if (bucket.count > effectiveConfig.max) { + const retryAfter = Math.max(1, resetSeconds); + res.setHeader('RateLimit-Remaining', '0'); res.setHeader('X-RateLimit-Remaining', '0'); - res.status(429).json({ error: 'Too many requests' }); + res.setHeader('Retry-After', `${retryAfter}`); + res.status(429).json({ error: 'Too many requests', retryAfter }); return; } @@ -238,42 +268,37 @@ export async function tieredRateLimit( const tier: RateLimitTier = keyCtx?.tier ?? getTierFromEnvKey(rawKey); const clientKey = keyCtx?.id ?? req.ip ?? 'anonymous'; - if (useTokenBucket) { - try { - const result: TokenBucketResult = await checkTokenBucket( - clientKey, - tier, - req.method, - req.path, - keyCtx?.rateLimitOverride, - ); - - res.setHeader('X-RateLimit-Limit', result.limit); - res.setHeader('X-RateLimit-Remaining', result.remaining); - res.setHeader('X-RateLimit-Reset', result.resetAt); - res.setHeader('X-RateLimit-Tier', result.tier); - (req as Request & { rateLimitResult?: TokenBucketResult }).rateLimitResult = result; - - if (!result.allowed) { - const retryAfter = Math.max(1, result.resetAt - Math.floor(Date.now() / 1000)); - res.setHeader('Retry-After', retryAfter); - res.status(429).json({ - error: 'Rate limit exceeded', - tier: result.tier, - retryAfter, - resetAt: new Date(result.resetAt * 1000).toISOString(), - }); - return; - } - - next(); + try { + const result: TokenBucketResult = await checkTokenBucket( + clientKey, + tier, + req.method, + req.path, + keyCtx?.rateLimitOverride, + ); + + const resetSeconds = writeRateLimitHeaders(res, result); + (req as Request & { rateLimitResult?: TokenBucketResult }).rateLimitResult = result; + + if (!result.allowed) { + const retryAfter = Math.max(1, resetSeconds); + res.setHeader('Retry-After', String(retryAfter)); + res.status(429).json({ + error: 'Rate limit exceeded', + tier: result.tier, + retryAfter, + resetAt: new Date(result.resetAt * 1000).toISOString(), + }); return; - } catch (err) { - logger.warn(`[rate-limit] Token bucket error, falling through: ${String(err)}`); } + + next(); + return; + } catch (err) { + logger.warn(`[rate-limit] Token bucket error, falling through: ${String(err)}`); } - if (keyCtx?.rateLimitOverride) { + if (keyCtx?.rateLimitOverride || !req.app) { runLocalRateLimit(req, res, next); return; } diff --git a/src/middleware/tokenBucket.ts b/src/middleware/tokenBucket.ts index d2081cd..7aa03b2 100644 --- a/src/middleware/tokenBucket.ts +++ b/src/middleware/tokenBucket.ts @@ -5,11 +5,13 @@ * (similar to CL.THROTTLE from redis-cell). Falls back to an in-memory map * when Redis is unavailable. * - * Tiers: - * unauthenticated 60 req/min burst 10 - * free 1000 req/min burst 50 - * pro 10000 req/min burst 200 - * enterprise custom custom + * Default tiers (configurable via RATE_LIMIT_TIER__LIMIT and + * RATE_LIMIT_TIER__WINDOW_MS environment variables): + * unauthenticated 60 req / 60s burst 10 + * free 1000 req / 60s burst 50 + * developer 1000 req / 60s burst 50 + * pro 10000 req / 60s burst 200 + * enterprise 60000 req / 60s burst 500 * * Per-endpoint groups add an additional multiplier on top of the tier limit. */ @@ -21,20 +23,108 @@ import { logger } from '../logger'; export type RateLimitTier = 'unauthenticated' | 'free' | 'developer' | 'pro' | 'enterprise'; -interface TierLimits { +export const RATE_LIMIT_TIERS: readonly RateLimitTier[] = [ + 'unauthenticated', + 'free', + 'developer', + 'pro', + 'enterprise', +] as const; + +export interface TierLimits { + /** Requests allowed per window (alias retained for backward compatibility). */ perMinute: number; + /** Requests allowed per window. */ + limit: number; + /** Window duration in milliseconds. */ + windowMs: number; + /** Burst allowance tokens. */ burst: number; + /** Redis key TTL in seconds. */ ttlSeconds: number; } -export const TIER_LIMITS: Record = { - unauthenticated: { perMinute: 60, burst: 10, ttlSeconds: 3600 }, - free: { perMinute: 1000, burst: 50, ttlSeconds: 86400 }, - developer: { perMinute: 1000, burst: 50, ttlSeconds: 86400 }, - pro: { perMinute: 10000, burst: 200, ttlSeconds: 2592000 }, - enterprise: { perMinute: 60000, burst: 500, ttlSeconds: 2592000 }, +export const DEFAULT_TIER_LIMITS: Readonly>> = { + unauthenticated: { + perMinute: 60, + limit: 60, + windowMs: 60_000, + burst: 10, + ttlSeconds: 3600, + }, + free: { + perMinute: 1000, + limit: 1000, + windowMs: 60_000, + burst: 50, + ttlSeconds: 86400, + }, + developer: { + perMinute: 1000, + limit: 1000, + windowMs: 60_000, + burst: 50, + ttlSeconds: 86400, + }, + pro: { + perMinute: 10000, + limit: 10000, + windowMs: 60_000, + burst: 200, + ttlSeconds: 2592000, + }, + enterprise: { + perMinute: 60000, + limit: 60000, + windowMs: 60_000, + burst: 500, + ttlSeconds: 2592000, + }, }; +function parsePositiveIntEnv(raw: string | undefined, fallback: number): number { + if (!raw || raw.trim() === '') return fallback; + const parsed = Number(raw.trim()); + if (!Number.isInteger(parsed) || parsed <= 0) return fallback; + return parsed; +} + +export function getTierLimits( + tier: RateLimitTier, + env: NodeJS.ProcessEnv = process.env, +): TierLimits { + const safeTier: RateLimitTier = tier in DEFAULT_TIER_LIMITS ? tier : 'unauthenticated'; + const defaults = DEFAULT_TIER_LIMITS[safeTier]; + const prefix = `RATE_LIMIT_TIER_${safeTier.toUpperCase()}`; + + const limit = parsePositiveIntEnv(env[`${prefix}_LIMIT`], defaults.limit); + const windowMs = parsePositiveIntEnv(env[`${prefix}_WINDOW_MS`], defaults.windowMs); + const burst = parsePositiveIntEnv(env[`${prefix}_BURST`], defaults.burst); + const ttlSeconds = Math.max(defaults.ttlSeconds, Math.ceil(windowMs / 1000)); + + return { + perMinute: limit, + limit, + windowMs, + burst, + ttlSeconds, + }; +} + +export function resolveTierLimits( + env: NodeJS.ProcessEnv = process.env, +): Record { + return { + unauthenticated: getTierLimits('unauthenticated', env), + free: getTierLimits('free', env), + developer: getTierLimits('developer', env), + pro: getTierLimits('pro', env), + enterprise: getTierLimits('enterprise', env), + }; +} + +export const TIER_LIMITS: Record = resolveTierLimits(); + // ─── Endpoint group multipliers (fraction of tier limit) ───────────────────── interface EndpointGroup { @@ -129,6 +219,11 @@ function checkMemoryBucket( return [allowed, Math.floor(bucket.tokens), resetAt]; } +/** Exposed for testing only — clears in-process token buckets. */ +export function _resetMemoryBuckets(): void { + memBuckets.clear(); +} + // ─── Public interface ───────────────────────────────────────────────────────── interface TokenBucketRedisClient { @@ -137,7 +232,7 @@ interface TokenBucketRedisClient { let redisClient: TokenBucketRedisClient | null = null; -export function setRateLimitRedisClient(client: TokenBucketRedisClient): void { +export function setRateLimitRedisClient(client: TokenBucketRedisClient | null): void { redisClient = client; } @@ -156,18 +251,24 @@ export async function checkTokenBucket( path: string, overridePerMinute?: number, ): Promise { - const baseLimits = TIER_LIMITS[tier]; + const resolvedTier: RateLimitTier = tier in DEFAULT_TIER_LIMITS ? tier : 'unauthenticated'; + const baseLimits = getTierLimits(resolvedTier); const multiplier = getEndpointMultiplier(method, path); - const capacity = overridePerMinute - ? Math.ceil(overridePerMinute * multiplier) - : Math.ceil(baseLimits.perMinute * multiplier); - const refillRate = capacity / 60; // tokens per second - const burst = Math.min(capacity, overridePerMinute ? capacity : baseLimits.burst); + const hasValidOverride = + typeof overridePerMinute === 'number' && + Number.isFinite(overridePerMinute) && + overridePerMinute > 0; + + const baseQuota = hasValidOverride ? overridePerMinute : baseLimits.limit; + const capacity = Math.max(1, Math.ceil(baseQuota * multiplier)); + const windowSeconds = Math.max(0.001, baseLimits.windowMs / 1000); + const refillRate = capacity / windowSeconds; // tokens per second + const burst = Math.min(capacity, hasValidOverride ? capacity : baseLimits.burst); const effectiveCapacity = Math.max(capacity, burst); const now = Date.now(); - const redisKey = `tb:${tier}:${clientKey}:${path.replace(/\//g, '_')}`; + const redisKey = `tb:${resolvedTier}:${clientKey}:${path.replace(/\//g, '_')}`; if (redisClient) { try { @@ -186,7 +287,7 @@ export async function checkTokenBucket( remaining: result[1], limit: capacity, resetAt: result[2], - tier, + tier: resolvedTier, }; } catch (err) { logger.warn('[token-bucket] Redis eval failed, using memory fallback', { err }); @@ -199,30 +300,63 @@ export async function checkTokenBucket( refillRate, now, ); - return { allowed, remaining, limit: capacity, resetAt, tier }; + return { allowed, remaining, limit: capacity, resetAt, tier: resolvedTier }; +} + +export function applyRateLimitHeaders(res: Response, result: TokenBucketResult): number { + const nowSec = Math.floor(Date.now() / 1000); + const remaining = Math.max(0, result.remaining); + const resetSeconds = Math.max(0, result.resetAt - nowSec); + + res.setHeader('RateLimit-Limit', String(result.limit)); + res.setHeader('RateLimit-Remaining', String(remaining)); + res.setHeader('RateLimit-Reset', String(resetSeconds)); + + res.setHeader('X-RateLimit-Limit', String(result.limit)); + res.setHeader('X-RateLimit-Remaining', String(remaining)); + res.setHeader('X-RateLimit-Reset', String(result.resetAt)); + res.setHeader('X-RateLimit-Tier', result.tier); + + return resetSeconds; +} + +function defaultTierResolver(req: Request): { + tier: RateLimitTier; + key: string; + override?: number; +} { + if (req.apiKey) { + return { + tier: req.apiKey.tier, + key: req.apiKey.id, + override: req.apiKey.rateLimitOverride, + }; + } + return { + tier: 'unauthenticated', + key: req.ip || 'anonymous', + }; } // ─── Express middleware ─────────────────────────────────────────────────────── export function tokenBucketMiddleware( - getTier: (req: Request) => { tier: RateLimitTier; key: string; override?: number }, + getTier: (req: Request) => { + tier: RateLimitTier; + key: string; + override?: number; + } = defaultTierResolver, ) { return async (req: Request, res: Response, next: NextFunction): Promise => { const { tier, key, override } = getTier(req); const result = await checkTokenBucket(key, tier, req.method, req.path, override); - // Set standard rate limit headers - res.setHeader('X-RateLimit-Limit', result.limit); - res.setHeader('X-RateLimit-Remaining', result.remaining); - res.setHeader('X-RateLimit-Reset', result.resetAt); - res.setHeader('X-RateLimit-Tier', result.tier); - - // Attach to request for audit logging + const resetSeconds = applyRateLimitHeaders(res, result); (req as Request & { rateLimitResult?: TokenBucketResult }).rateLimitResult = result; if (!result.allowed) { - const retryAfter = Math.max(1, result.resetAt - Math.floor(Date.now() / 1000)); - res.setHeader('Retry-After', retryAfter); + const retryAfter = Math.max(1, resetSeconds); + res.setHeader('Retry-After', String(retryAfter)); res.status(429).json({ error: 'Rate limit exceeded', tier: result.tier, diff --git a/tests/rate-limit.test.ts b/tests/rate-limit.test.ts index 51c41e1..909e18f 100644 --- a/tests/rate-limit.test.ts +++ b/tests/rate-limit.test.ts @@ -24,11 +24,21 @@ function makeRes(): { status: ReturnType; json: ReturnType; setHeader: ReturnType; + headers: Record; } { + const headers: Record = {}; const json = vi.fn(); const status = vi.fn().mockReturnValue({ json }); - const setHeader = vi.fn(); - return { res: { status, json, setHeader } as unknown as Response, status, json, setHeader }; + const setHeader = vi.fn((name: string, value: unknown) => { + headers[name.toLowerCase()] = String(value); + }); + return { + res: { status, json, setHeader } as unknown as Response, + status, + json, + setHeader, + headers, + }; } describe('rate limit configuration', () => { @@ -60,35 +70,97 @@ describe('tieredRateLimit', () => { beforeEach(() => { next = vi.fn(); + vi.clearAllMocks(); + const nowSec = Math.floor(Date.now() / 1000); mockCheck.mockResolvedValue({ allowed: true, limit: 100, remaining: 99, - resetAt: 9999999999, + resetAt: nowSec + 45, tier: 'free', }); }); - it('calls next() when the fallback limiter allows the request', async () => { - const { res } = makeRes(); + it('calls next() and emits standard RateLimit-* headers when allowed', async () => { + const { res, headers } = makeRes(); await tieredRateLimit(makeReq(), res, next); - expect(next).toHaveBeenCalled(); + + expect(next).toHaveBeenCalledOnce(); + expect(headers['ratelimit-limit']).toBe('100'); + expect(headers['ratelimit-remaining']).toBe('99'); + expect(Number(headers['ratelimit-reset'])).toBeGreaterThanOrEqual(0); + expect(headers['retry-after']).toBeUndefined(); }); - it('uses apiKey tier when present', async () => { + it('uses apiKey tier and forwards rateLimitOverride when present', async () => { const req = makeReq({ - apiKey: { id: 'k', keyName: 'n', developerId: 'd', tier: 'pro' }, + apiKey: { + id: 'key-pro-1', + keyName: 'prod-key', + developerId: 'dev-1', + tier: 'pro', + rateLimitOverride: 2500, + }, } as any); const { res } = makeRes(); await tieredRateLimit(req, res, next); - expect(next).toHaveBeenCalled(); + + expect(mockCheck).toHaveBeenCalledWith('key-pro-1', 'pro', 'GET', '/test', 2500); + expect(next).toHaveBeenCalledOnce(); }); it('falls back to unauthenticated tier when no apiKey header is present', async () => { const req = makeReq({ headers: {} }); const { res } = makeRes(); await tieredRateLimit(req, res, next); - expect(next).toHaveBeenCalled(); + + expect(mockCheck).toHaveBeenCalledWith( + '127.0.0.1', + 'unauthenticated', + 'GET', + '/test', + undefined, + ); + expect(next).toHaveBeenCalledOnce(); + }); + + it('returns 429 with Retry-After and standard RateLimit-* headers when limit is exceeded', async () => { + const nowSec = Math.floor(Date.now() / 1000); + mockCheck.mockResolvedValueOnce({ + allowed: false, + limit: 60, + remaining: 0, + resetAt: nowSec + 20, + tier: 'unauthenticated', + }); + + const { res, status, json, headers } = makeRes(); + await tieredRateLimit(makeReq(), res, next); + + expect(next).not.toHaveBeenCalled(); + expect(status).toHaveBeenCalledWith(429); + expect(headers['ratelimit-limit']).toBe('60'); + expect(headers['ratelimit-remaining']).toBe('0'); + expect(Number(headers['ratelimit-reset'])).toBeGreaterThanOrEqual(1); + expect(Number(headers['retry-after'])).toBeGreaterThanOrEqual(1); + expect(json).toHaveBeenCalledWith( + expect.objectContaining({ + error: 'Rate limit exceeded', + tier: 'unauthenticated', + }), + ); + }); + + it('falls back gracefully to local rate limiter if checkTokenBucket throws', async () => { + mockCheck.mockRejectedValueOnce(new Error('unexpected failure')); + const { res, headers } = makeRes(); + await tieredRateLimit(makeReq(), res, next); + + await vi.waitFor(() => { + expect(next).toHaveBeenCalledOnce(); + }); + expect(headers['ratelimit-limit']).toBeDefined(); + expect(headers['ratelimit-remaining']).toBeDefined(); }); }); diff --git a/tests/token-bucket.test.ts b/tests/token-bucket.test.ts new file mode 100644 index 0000000..670eef3 --- /dev/null +++ b/tests/token-bucket.test.ts @@ -0,0 +1,345 @@ +import express from 'express'; +import request from 'supertest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +vi.mock('../src/db', () => ({ + prismaRead: { + devApiKey: { findFirst: vi.fn() }, + rateLimitOverride: { findUnique: vi.fn().mockResolvedValue(null) }, + }, + prismaWrite: { + devApiKey: { update: vi.fn().mockResolvedValue(null) }, + }, +})); +vi.mock('../src/logger', () => ({ + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }, +})); + +import { prismaRead } from '../src/db'; +import { _clearKeyCache, apiKeyAuth, normalizeRateLimitTier } from '../src/middleware/apiKeyAuth'; +import { tieredRateLimit } from '../src/middleware/rateLimit'; +import { + _resetMemoryBuckets, + checkTokenBucket, + DEFAULT_TIER_LIMITS, + getTierLimits, + resolveTierLimits, + setRateLimitRedisClient, + tokenBucketMiddleware, +} from '../src/middleware/tokenBucket'; + +const mockFindKey = (prismaRead as any).devApiKey.findFirst as ReturnType; + +function buildTestApp(useDirectTokenBucket = false) { + const app = express(); + app.use(express.json()); + app.use(apiKeyAuth); + app.use(useDirectTokenBucket ? tokenBucketMiddleware() : tieredRateLimit); + app.get('/api/v1/events', (_req, res) => { + res.status(200).json({ ok: true }); + }); + app.get('/api/v1/search', (_req, res) => { + res.status(200).json({ results: [] }); + }); + return app; +} + +describe('tokenBucket tier configuration via environment variables', () => { + beforeEach(() => { + _resetMemoryBuckets(); + _clearKeyCache(); + setRateLimitRedisClient(null); + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + vi.useRealTimers(); + setRateLimitRedisClient(null); + }); + + it('uses current hardcoded values as defaults when env vars are unset', () => { + const resolved = resolveTierLimits({}); + + expect(resolved.unauthenticated).toMatchObject({ + limit: 60, + perMinute: 60, + windowMs: 60_000, + burst: 10, + }); + expect(resolved.free).toMatchObject({ + limit: 1000, + perMinute: 1000, + windowMs: 60_000, + burst: 50, + }); + expect(resolved.developer).toMatchObject({ + limit: 1000, + perMinute: 1000, + windowMs: 60_000, + burst: 50, + }); + expect(resolved.pro).toMatchObject({ + limit: 10000, + perMinute: 10000, + windowMs: 60_000, + burst: 200, + }); + expect(resolved.enterprise).toMatchObject({ + limit: 60000, + perMinute: 60000, + windowMs: 60_000, + burst: 500, + }); + }); + + it('configures each tier limit and window from environment variables', () => { + const customEnv: NodeJS.ProcessEnv = { + RATE_LIMIT_TIER_UNAUTHENTICATED_LIMIT: '25', + RATE_LIMIT_TIER_UNAUTHENTICATED_WINDOW_MS: '30000', + RATE_LIMIT_TIER_FREE_LIMIT: '500', + RATE_LIMIT_TIER_FREE_WINDOW_MS: '120000', + RATE_LIMIT_TIER_DEVELOPER_LIMIT: '2000', + RATE_LIMIT_TIER_DEVELOPER_WINDOW_MS: '45000', + RATE_LIMIT_TIER_PRO_LIMIT: '15000', + RATE_LIMIT_TIER_PRO_WINDOW_MS: '90000', + RATE_LIMIT_TIER_ENTERPRISE_LIMIT: '120000', + RATE_LIMIT_TIER_ENTERPRISE_WINDOW_MS: '180000', + }; + + const resolved = resolveTierLimits(customEnv); + + expect(resolved.unauthenticated).toMatchObject({ limit: 25, windowMs: 30_000 }); + expect(resolved.free).toMatchObject({ limit: 500, windowMs: 120_000 }); + expect(resolved.developer).toMatchObject({ limit: 2000, windowMs: 45_000 }); + expect(resolved.pro).toMatchObject({ limit: 15000, windowMs: 90_000 }); + expect(resolved.enterprise).toMatchObject({ limit: 120000, windowMs: 180_000 }); + }); + + it('falls back to defaults when env vars are invalid, zero, negative, or non-integer', () => { + const noisyEnv: NodeJS.ProcessEnv = { + RATE_LIMIT_TIER_UNAUTHENTICATED_LIMIT: '0', + RATE_LIMIT_TIER_UNAUTHENTICATED_WINDOW_MS: '-5000', + RATE_LIMIT_TIER_FREE_LIMIT: 'not-a-number', + RATE_LIMIT_TIER_FREE_WINDOW_MS: '12.5', + RATE_LIMIT_TIER_PRO_LIMIT: ' ', + }; + + expect(getTierLimits('unauthenticated', noisyEnv)).toMatchObject({ + limit: DEFAULT_TIER_LIMITS.unauthenticated.limit, + windowMs: DEFAULT_TIER_LIMITS.unauthenticated.windowMs, + }); + expect(getTierLimits('free', noisyEnv)).toMatchObject({ + limit: DEFAULT_TIER_LIMITS.free.limit, + windowMs: DEFAULT_TIER_LIMITS.free.windowMs, + }); + expect(getTierLimits('pro', noisyEnv)).toMatchObject({ + limit: DEFAULT_TIER_LIMITS.pro.limit, + windowMs: DEFAULT_TIER_LIMITS.pro.windowMs, + }); + }); +}); + +describe('standard rate-limit headers and 429 Retry-After integration', () => { + beforeEach(() => { + _resetMemoryBuckets(); + _clearKeyCache(); + setRateLimitRedisClient(null); + vi.clearAllMocks(); + mockFindKey.mockResolvedValue(null); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + vi.useRealTimers(); + setRateLimitRedisClient(null); + }); + + it('includes RateLimit-Limit, RateLimit-Remaining, and RateLimit-Reset on within-limit requests and 429 with Retry-After when exceeded', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-06-01T00:00:00.000Z')); + vi.stubEnv('RATE_LIMIT_TIER_UNAUTHENTICATED_LIMIT', '3'); + vi.stubEnv('RATE_LIMIT_TIER_UNAUTHENTICATED_WINDOW_MS', '30000'); + + const app = buildTestApp(); + + const r1 = await request(app).get('/api/v1/events'); + expect(r1.status).toBe(200); + expect(r1.headers['ratelimit-limit']).toBe('3'); + expect(r1.headers['ratelimit-remaining']).toBe('2'); + expect(Number(r1.headers['ratelimit-reset'])).toBeGreaterThan(0); + expect(r1.headers['retry-after']).toBeUndefined(); + + const r2 = await request(app).get('/api/v1/events'); + expect(r2.status).toBe(200); + expect(r2.headers['ratelimit-limit']).toBe('3'); + expect(r2.headers['ratelimit-remaining']).toBe('1'); + expect(r2.headers['retry-after']).toBeUndefined(); + + const r3 = await request(app).get('/api/v1/events'); + expect(r3.status).toBe(200); + expect(r3.headers['ratelimit-limit']).toBe('3'); + expect(r3.headers['ratelimit-remaining']).toBe('0'); + expect(r3.headers['retry-after']).toBeUndefined(); + + const r4 = await request(app).get('/api/v1/events'); + expect(r4.status).toBe(429); + expect(r4.headers['ratelimit-limit']).toBe('3'); + expect(r4.headers['ratelimit-remaining']).toBe('0'); + expect(Number(r4.headers['ratelimit-reset'])).toBe(30); + expect(r4.headers['retry-after']).toBe('30'); + expect(r4.body).toMatchObject({ + error: 'Rate limit exceeded', + tier: 'unauthenticated', + retryAfter: 30, + }); + + // Advance past the 30s configured window and verify quota refills + vi.advanceTimersByTime(30_000); + + const r5 = await request(app).get('/api/v1/events'); + expect(r5.status).toBe(200); + expect(r5.headers['ratelimit-limit']).toBe('3'); + expect(r5.headers['ratelimit-remaining']).toBe('2'); + expect(r5.headers['retry-after']).toBeUndefined(); + }); + + it('applies authenticated API key tier limits configured via env vars', async () => { + vi.stubEnv('RATE_LIMIT_TIER_DEVELOPER_LIMIT', '4'); + vi.stubEnv('RATE_LIMIT_TIER_DEVELOPER_WINDOW_MS', '60000'); + + mockFindKey.mockResolvedValue({ + id: 'dev-key-1', + name: 'ci-key', + developerId: 'dev-42', + tier: 'developer', + rateLimitOverride: null, + allowedIps: null, + allowedEndpoints: null, + allowedDomains: null, + expiresAt: null, + revokedAt: null, + }); + + const app = buildTestApp(); + + const res = await request(app).get('/api/v1/events').set('X-Api-Key', 'secret-dev-key'); + expect(res.status).toBe(200); + expect(res.headers['ratelimit-limit']).toBe('4'); + expect(res.headers['ratelimit-remaining']).toBe('3'); + expect(res.headers['x-ratelimit-tier']).toBe('developer'); + }); + + it('enforces per-key rateLimitOverride from apiKeyAuth over the tier default', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-06-01T12:00:00.000Z')); + + mockFindKey.mockResolvedValue({ + id: 'custom-override-key', + name: 'custom-quota', + developerId: 'dev-99', + tier: 'pro', + rateLimitOverride: 2, + allowedIps: null, + allowedEndpoints: null, + allowedDomains: null, + expiresAt: null, + revokedAt: null, + }); + + const app = buildTestApp(); + + const first = await request(app).get('/api/v1/events').set('X-Api-Key', 'override-raw-key'); + expect(first.status).toBe(200); + expect(first.headers['ratelimit-limit']).toBe('2'); + expect(first.headers['ratelimit-remaining']).toBe('1'); + + const second = await request(app).get('/api/v1/events').set('X-Api-Key', 'override-raw-key'); + expect(second.status).toBe(200); + expect(second.headers['ratelimit-limit']).toBe('2'); + expect(second.headers['ratelimit-remaining']).toBe('0'); + + const third = await request(app).get('/api/v1/events').set('X-Api-Key', 'override-raw-key'); + expect(third.status).toBe(429); + expect(third.headers['ratelimit-limit']).toBe('2'); + expect(third.headers['ratelimit-remaining']).toBe('0'); + expect(Number(third.headers['retry-after'])).toBeGreaterThanOrEqual(1); + expect(third.body.tier).toBe('pro'); + }); + + it('normalizes casing and unknown tier values safely in apiKeyAuth', async () => { + expect(normalizeRateLimitTier(' PRO ')).toBe('pro'); + expect(normalizeRateLimitTier('Enterprise')).toBe('enterprise'); + expect(normalizeRateLimitTier('legacy_unknown')).toBe('free'); + expect(normalizeRateLimitTier(null)).toBe('free'); + + mockFindKey.mockResolvedValue({ + id: 'unknown-tier-key', + name: 'legacy', + developerId: 'dev-7', + tier: 'legacy_unknown', + rateLimitOverride: 0, + allowedIps: null, + allowedEndpoints: null, + allowedDomains: null, + expiresAt: null, + revokedAt: null, + }); + + const app = buildTestApp(); + const res = await request(app).get('/api/v1/events').set('X-Api-Key', 'legacy-key'); + + expect(res.status).toBe(200); + expect(res.headers['ratelimit-limit']).toBe('1000'); + expect(res.headers['ratelimit-remaining']).toBe('999'); + expect(res.headers['x-ratelimit-tier']).toBe('free'); + }); + + it('emits standard headers and Retry-After via standalone tokenBucketMiddleware', async () => { + vi.stubEnv('RATE_LIMIT_TIER_UNAUTHENTICATED_LIMIT', '1'); + vi.stubEnv('RATE_LIMIT_TIER_UNAUTHENTICATED_WINDOW_MS', '15000'); + + const app = buildTestApp(true); + + const okRes = await request(app).get('/api/v1/events'); + expect(okRes.status).toBe(200); + expect(okRes.headers['ratelimit-limit']).toBe('1'); + expect(okRes.headers['ratelimit-remaining']).toBe('0'); + + const blockedRes = await request(app).get('/api/v1/events'); + expect(blockedRes.status).toBe(429); + expect(blockedRes.headers['ratelimit-limit']).toBe('1'); + expect(blockedRes.headers['ratelimit-remaining']).toBe('0'); + expect(Number(blockedRes.headers['retry-after'])).toBeGreaterThanOrEqual(1); + }); + + it('scales quota by endpoint multiplier and falls back cleanly if Redis eval fails', async () => { + vi.stubEnv('RATE_LIMIT_TIER_FREE_LIMIT', '100'); + vi.stubEnv('RATE_LIMIT_TIER_FREE_WINDOW_MS', '20000'); + + const redisEval = vi + .fn() + .mockResolvedValueOnce([1, 49, Math.floor(Date.now() / 1000) + 20]) + .mockRejectedValueOnce(new Error('READONLY Redis error')); + + setRateLimitRedisClient({ eval: redisEval }); + + // /api/v1/search has a 0.5x multiplier -> effective capacity 50 over 20s (2.5 tokens/sec) + const first = await checkTokenBucket('client-1', 'free', 'GET', '/api/v1/search'); + expect(first.allowed).toBe(true); + expect(first.limit).toBe(50); + expect(first.remaining).toBe(49); + expect(redisEval).toHaveBeenCalledWith( + expect.any(String), + expect.objectContaining({ + arguments: expect.arrayContaining(['50', '2.5']), + }), + ); + + // Second call triggers Redis error and seamlessly falls back to in-memory bucket + const second = await checkTokenBucket('client-1', 'free', 'GET', '/api/v1/search'); + expect(second.allowed).toBe(true); + expect(second.limit).toBe(50); + expect(second.remaining).toBe(49); + }); +});