UNPKG

@limitrate/core

Version:

Core rate limiting and cost control engine for LimitRate

1,726 lines (1,674 loc) 72 kB
// src/stores/memory.ts function isCacheEntry(entry) { return "count" in entry; } var MemoryStore = class { cache; maxKeys; cleanupInterval; constructor(options = {}) { this.cache = /* @__PURE__ */ new Map(); this.maxKeys = options.maxKeys ?? 1e4; this.cleanupInterval = null; const intervalMs = options.cleanupIntervalMs ?? 6e4; this.cleanupInterval = setInterval(() => this.cleanup(), intervalMs); } async checkRate(key, limit, windowSeconds, burst) { const now = Date.now(); const entry = this.cache.get(key); if (entry && !isCacheEntry(entry)) { const expiresAt = now + windowSeconds * 1e3; const burstTokens = burst; const cacheEntry = { count: 1, expiresAt, burstTokens }; this.cache.set(key, cacheEntry); this.evictIfNeeded(); return { allowed: true, current: 1, remaining: limit - 1, resetInSeconds: windowSeconds, limit, burstTokens }; } if (!entry || entry.expiresAt <= now) { const expiresAt = now + windowSeconds * 1e3; const burstTokens = burst; this.cache.set(key, { count: 1, expiresAt, burstTokens }); this.evictIfNeeded(); return { allowed: true, current: 1, remaining: limit - 1, resetInSeconds: windowSeconds, limit, burstTokens }; } const current = entry.count; const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3); if (current < limit) { entry.count += 1; return { allowed: true, current: entry.count, remaining: limit - entry.count, resetInSeconds, limit, burstTokens: entry.burstTokens }; } if (burst !== void 0 && entry.burstTokens !== void 0 && entry.burstTokens > 0) { entry.burstTokens -= 1; entry.count += 1; return { allowed: true, current: entry.count, remaining: 0, resetInSeconds, limit, burstTokens: entry.burstTokens }; } return { allowed: false, current, remaining: 0, resetInSeconds, limit, burstTokens: entry.burstTokens ?? 0 }; } async peekRate(key, limit, windowSeconds) { const now = Date.now(); const entry = this.cache.get(key); if (entry && !isCacheEntry(entry)) { return { allowed: true, current: 0, remaining: limit, resetInSeconds: windowSeconds, limit }; } if (!entry || entry.expiresAt <= now) { return { allowed: true, current: 0, remaining: limit, resetInSeconds: windowSeconds, limit }; } const current = entry.count; const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3); const remaining = Math.max(0, limit - current); return { allowed: current < limit, current, remaining, resetInSeconds, limit, burstTokens: entry.burstTokens }; } async incrementCost(key, cost, windowSeconds, cap) { const now = Date.now(); const entry = this.cache.get(key); if (entry && !isCacheEntry(entry)) { const expiresAt = now + windowSeconds * 1e3; const cacheEntry = { count: cost, expiresAt }; this.cache.set(key, cacheEntry); this.evictIfNeeded(); return { allowed: cost <= cap, current: cost, remaining: cap - cost, resetInSeconds: windowSeconds, cap }; } if (!entry || entry.expiresAt <= now) { const expiresAt = now + windowSeconds * 1e3; this.cache.set(key, { count: cost, expiresAt }); this.evictIfNeeded(); return { allowed: cost <= cap, current: cost, remaining: cap - cost, resetInSeconds: windowSeconds, cap }; } const newCost = entry.count + cost; const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3); if (newCost > cap) { return { allowed: false, current: entry.count, remaining: cap - entry.count, resetInSeconds, cap }; } entry.count = newCost; return { allowed: true, current: newCost, remaining: cap - newCost, resetInSeconds, cap }; } async incrementTokens(key, tokens, windowSeconds, limit) { const now = Date.now(); const entry = this.cache.get(key); if (entry && !isCacheEntry(entry)) { const expiresAt = now + windowSeconds * 1e3; const cacheEntry = { count: tokens, expiresAt }; this.cache.set(key, cacheEntry); this.evictIfNeeded(); return { allowed: tokens <= limit, current: tokens, remaining: Math.max(0, limit - tokens), resetInSeconds: windowSeconds, limit }; } if (!entry || entry.expiresAt <= now) { const expiresAt = now + windowSeconds * 1e3; this.cache.set(key, { count: tokens, expiresAt }); this.evictIfNeeded(); return { allowed: tokens <= limit, current: tokens, remaining: Math.max(0, limit - tokens), resetInSeconds: windowSeconds, limit }; } const newTokens = entry.count + tokens; const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3); if (newTokens > limit) { return { allowed: false, current: entry.count, remaining: Math.max(0, limit - entry.count), resetInSeconds, limit }; } entry.count = newTokens; return { allowed: true, current: newTokens, remaining: limit - newTokens, resetInSeconds, limit }; } async ping() { return true; } async close() { if (this.cleanupInterval) { clearInterval(this.cleanupInterval); this.cleanupInterval = null; } this.cache.clear(); } /** * Remove expired entries */ cleanup() { const now = Date.now(); for (const [key, entry] of this.cache.entries()) { if (entry.expiresAt <= now) { this.cache.delete(key); } } } /** * Evict oldest entry if cache is full (simple LRU) */ evictIfNeeded() { if (this.cache.size > this.maxKeys) { const firstKey = this.cache.keys().next().value; if (firstKey) { this.cache.delete(firstKey); } } } /** * Generic get method for arbitrary data (v2.0.0 - D4) */ async get(key) { const now = Date.now(); const entry = this.cache.get(key); if (!entry) { return null; } if ("value" in entry) { const genericEntry = entry; if (genericEntry.expiresAt <= now) { this.cache.delete(key); return null; } return genericEntry.value; } return null; } /** * Generic set method for arbitrary data (v2.0.0 - D4) */ async set(key, value, ttl) { const now = Date.now(); const expiresAt = ttl ? now + ttl * 1e3 : now + 86400 * 1e3; const entry = { value, expiresAt }; this.cache.set(key, entry); this.evictIfNeeded(); } /** * Generic delete method (v2.0.0 - D4) */ async delete(key) { this.cache.delete(key); } /** * Get current cache size (for testing/debugging) */ getCacheSize() { return this.cache.size; } }; // src/stores/redis.ts import Redis from "ioredis"; var RATE_PEEK_SCRIPT = ` local key = KEYS[1] local limit = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local current = redis.call('GET', key) if not current then return {0, limit, window, limit} end current = tonumber(current) local ttl = redis.call('TTL', key) local remaining = math.max(0, limit - current) return {current, remaining, ttl, limit} `; var RATE_CHECK_SCRIPT = ` local key = KEYS[1] local limit = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local now = tonumber(ARGV[3]) local current = redis.call('GET', key) if not current then redis.call('SETEX', key, window, 1) return {1, limit - 1, window, limit} end current = tonumber(current) if current >= limit then local ttl = redis.call('TTL', key) return {current, 0, ttl, limit} end redis.call('INCR', key) local ttl = redis.call('TTL', key) return {current + 1, limit - current - 1, ttl, limit} `; var RATE_CHECK_BURST_SCRIPT = ` local rateKey = KEYS[1] local burstKey = KEYS[2] local limit = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local now = tonumber(ARGV[3]) local burst = tonumber(ARGV[4]) local current = redis.call('GET', rateKey) -- No entry: create new with initial burst tokens if not current then redis.call('SETEX', rateKey, window, 1) redis.call('SETEX', burstKey, window, burst) return {1, limit - 1, window, limit, burst} end current = tonumber(current) -- Within limit: increment and return if current < limit then redis.call('INCR', rateKey) local ttl = redis.call('TTL', rateKey) local burstTokens = tonumber(redis.call('GET', burstKey) or 0) return {current + 1, limit - current - 1, ttl, limit, burstTokens} end -- Limit reached: try burst local burstTokens = tonumber(redis.call('GET', burstKey) or 0) if burstTokens > 0 then redis.call('DECR', burstKey) redis.call('INCR', rateKey) local ttl = redis.call('TTL', rateKey) return {current + 1, 0, ttl, limit, burstTokens - 1} end -- Both exhausted local ttl = redis.call('TTL', rateKey) return {current, 0, ttl, limit, 0} `; var COST_INCREMENT_SCRIPT = ` local key = KEYS[1] local cost = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local cap = tonumber(ARGV[3]) local current = redis.call('GET', key) if not current then if cost > cap then redis.call('SETEX', key, window, 0) return {0, false, cap, window, cap} end redis.call('SETEX', key, window, cost) return {cost, true, cap - cost, window, cap} end current = tonumber(current) local newCost = current + cost if newCost > cap then local ttl = redis.call('TTL', key) return {current, false, cap - current, ttl, cap} end redis.call('SET', key, newCost, 'KEEPTTL') local ttl = redis.call('TTL', key) return {newCost, true, cap - newCost, ttl, cap} `; var TOKEN_INCREMENT_SCRIPT = ` local key = KEYS[1] local tokens = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local limit = tonumber(ARGV[3]) local current = redis.call('GET', key) if not current then if tokens > limit then redis.call('SETEX', key, window, 0) return {0, false, limit, window, limit} end redis.call('SETEX', key, window, tokens) return {tokens, true, limit - tokens, window, limit} end current = tonumber(current) local newTokens = current + tokens if newTokens > limit then local ttl = redis.call('TTL', key) return {current, false, math.max(0, limit - current), ttl, limit} end redis.call('SET', key, newTokens, 'KEEPTTL') local ttl = redis.call('TTL', key) return {newTokens, true, limit - newTokens, ttl, limit} `; var RedisStore = class { client; ownClient; keyPrefix; constructor(options = {}) { this.keyPrefix = options.keyPrefix ?? "limitrate:"; this.ownClient = false; if (options.client instanceof Redis) { this.client = options.client; } else if (typeof options.client === "string") { this.client = new Redis(options.client, options.redisOptions); this.ownClient = true; } else { throw new Error("RedisStore requires either a Redis URL or ioredis instance"); } } async checkRate(key, limit, windowSeconds, burst) { const now = Math.floor(Date.now() / 1e3); try { if (burst !== void 0) { const rateKey = `${this.keyPrefix}rate:${key}`; const burstKey = `${this.keyPrefix}burst:${key}`; const result2 = await this.client.eval( RATE_CHECK_BURST_SCRIPT, 2, rateKey, burstKey, limit.toString(), windowSeconds.toString(), now.toString(), burst.toString() ); const [current2, remaining2, resetInSeconds2, returnedLimit2, burstTokens] = result2; return { allowed: current2 <= limit + burst, current: current2, remaining: Math.max(0, remaining2), resetInSeconds: Math.max(1, resetInSeconds2), limit: returnedLimit2, burstTokens }; } const prefixedKey = `${this.keyPrefix}rate:${key}`; const result = await this.client.eval( RATE_CHECK_SCRIPT, 1, prefixedKey, limit.toString(), windowSeconds.toString(), now.toString() ); const [current, remaining, resetInSeconds, returnedLimit] = result; return { allowed: current <= limit, current, remaining: Math.max(0, remaining), resetInSeconds: Math.max(1, resetInSeconds), limit: returnedLimit }; } catch (error) { console.error("[LimitRate] Redis rate check error:", error); throw error; } } async peekRate(key, limit, windowSeconds) { const prefixedKey = `${this.keyPrefix}rate:${key}`; try { const result = await this.client.eval( RATE_PEEK_SCRIPT, 1, prefixedKey, limit.toString(), windowSeconds.toString() ); const [current, remaining, resetInSeconds, returnedLimit] = result; return { allowed: current < limit, current, remaining: Math.max(0, remaining), resetInSeconds: Math.max(1, resetInSeconds), limit: returnedLimit }; } catch (error) { console.error("[LimitRate] Redis peek rate error:", error); throw error; } } async incrementCost(key, cost, windowSeconds, cap) { const prefixedKey = `${this.keyPrefix}cost:${key}`; try { const result = await this.client.eval( COST_INCREMENT_SCRIPT, 1, prefixedKey, cost.toString(), windowSeconds.toString(), cap.toString() ); const [current, allowed, remaining, resetInSeconds, returnedCap] = result; return { allowed: allowed === 1, current, remaining: Math.max(0, remaining), resetInSeconds: Math.max(1, resetInSeconds), cap: returnedCap }; } catch (error) { console.error("[LimitRate] Redis cost increment error:", error); throw error; } } async incrementTokens(key, tokens, windowSeconds, limit) { const prefixedKey = `${this.keyPrefix}tokens:${key}`; try { const result = await this.client.eval( TOKEN_INCREMENT_SCRIPT, 1, prefixedKey, tokens.toString(), windowSeconds.toString(), limit.toString() ); const [current, allowed, remaining, resetInSeconds, returnedLimit] = result; return { allowed: allowed === 1, current, remaining: Math.max(0, remaining), resetInSeconds: Math.max(1, resetInSeconds), limit: returnedLimit }; } catch (error) { console.error("[LimitRate] Redis token increment error:", error); throw error; } } async ping() { try { const result = await this.client.ping(); return result === "PONG"; } catch { return false; } } async close() { if (this.ownClient) { await this.client.quit(); } } /** * Generic get method for arbitrary data (v2.0.0 - D4) */ async get(key) { try { const prefixedKey = `${this.keyPrefix}generic:${key}`; const value = await this.client.get(prefixedKey); if (!value) { return null; } return JSON.parse(value); } catch (error) { console.error("[LimitRate] Redis get error:", error); throw error; } } /** * Generic set method for arbitrary data (v2.0.0 - D4) */ async set(key, value, ttl) { try { const prefixedKey = `${this.keyPrefix}generic:${key}`; const serialized = JSON.stringify(value); if (ttl) { await this.client.setex(prefixedKey, ttl, serialized); } else { await this.client.setex(prefixedKey, 86400, serialized); } } catch (error) { console.error("[LimitRate] Redis set error:", error); throw error; } } /** * Generic delete method (v2.0.0 - D4) */ async delete(key) { try { const prefixedKey = `${this.keyPrefix}generic:${key}`; await this.client.del(prefixedKey); } catch (error) { console.error("[LimitRate] Redis delete error:", error); throw error; } } /** * Get underlying Redis client (for advanced use cases) */ getClient() { return this.client; } }; // src/stores/upstash.ts import { Redis as Redis2 } from "@upstash/redis"; var RATE_CHECK_SCRIPT2 = ` local key = KEYS[1] local limit = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local now = tonumber(ARGV[3]) local current = redis.call('GET', key) if not current then redis.call('SETEX', key, window, 1) return {1, limit - 1, window, limit} end current = tonumber(current) if current >= limit then local ttl = redis.call('TTL', key) return {current, 0, ttl, limit} end redis.call('INCR', key) local ttl = redis.call('TTL', key) return {current + 1, limit - current - 1, ttl, limit} `; var RATE_CHECK_BURST_SCRIPT2 = ` local rateKey = KEYS[1] local burstKey = KEYS[2] local limit = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local now = tonumber(ARGV[3]) local burst = tonumber(ARGV[4]) local current = redis.call('GET', rateKey) if not current then redis.call('SETEX', rateKey, window, 1) redis.call('SETEX', burstKey, window, burst) return {1, limit - 1, window, limit, burst} end current = tonumber(current) if current < limit then redis.call('INCR', rateKey) local ttl = redis.call('TTL', rateKey) local burstTokens = tonumber(redis.call('GET', burstKey) or 0) return {current + 1, limit - current - 1, ttl, limit, burstTokens} end local burstTokens = tonumber(redis.call('GET', burstKey) or 0) if burstTokens > 0 then redis.call('DECR', burstKey) redis.call('INCR', rateKey) local ttl = redis.call('TTL', rateKey) return {current + 1, 0, ttl, limit, burstTokens - 1} end local ttl = redis.call('TTL', rateKey) return {current, 0, ttl, limit, 0} `; var COST_INCREMENT_SCRIPT2 = ` local key = KEYS[1] local cost = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local cap = tonumber(ARGV[3]) local current = redis.call('GET', key) if not current then if cost > cap then redis.call('SETEX', key, window, 0) return {0, false, cap, window, cap} end redis.call('SETEX', key, window, cost) return {cost, true, cap - cost, window, cap} end current = tonumber(current) local newCost = current + cost if newCost > cap then local ttl = redis.call('TTL', key) return {current, false, cap - current, ttl, cap} end redis.call('SET', key, newCost, 'KEEPTTL') local ttl = redis.call('TTL', key) return {newCost, true, cap - newCost, ttl, cap} `; var TOKEN_INCREMENT_SCRIPT2 = ` local key = KEYS[1] local tokens = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local limit = tonumber(ARGV[3]) local current = redis.call('GET', key) if not current then if tokens > limit then redis.call('SETEX', key, window, 0) return {0, false, limit, window, limit} end redis.call('SETEX', key, window, tokens) return {tokens, true, limit - tokens, window, limit} end current = tonumber(current) local newTokens = current + tokens if newTokens > limit then local ttl = redis.call('TTL', key) return {current, false, math.max(0, limit - current), ttl, limit} end redis.call('SET', key, newTokens, 'KEEPTTL') local ttl = redis.call('TTL', key) return {newTokens, true, limit - newTokens, ttl, limit} `; var UpstashStore = class { client; keyPrefix; url; token; constructor(options) { if (!options.url || !options.token) { throw new Error("UpstashStore requires both url and token"); } this.keyPrefix = options.keyPrefix ?? "limitrate:"; this.url = options.url; this.token = options.token; this.client = new Redis2({ url: options.url, token: options.token }); } async checkRate(key, limit, windowSeconds, burst) { const now = Math.floor(Date.now() / 1e3); try { if (burst !== void 0) { const rateKey = `${this.keyPrefix}rate:${key}`; const burstKey = `${this.keyPrefix}burst:${key}`; const result2 = await this.client.eval( RATE_CHECK_BURST_SCRIPT2, [rateKey, burstKey], [limit.toString(), windowSeconds.toString(), now.toString(), burst.toString()] ); const [current2, remaining2, resetInSeconds2, returnedLimit2, burstTokens] = result2; return { allowed: current2 <= limit + burst, current: current2, remaining: Math.max(0, remaining2), resetInSeconds: Math.max(1, resetInSeconds2), limit: returnedLimit2, burstTokens }; } const prefixedKey = `${this.keyPrefix}rate:${key}`; const result = await this.client.eval( RATE_CHECK_SCRIPT2, [prefixedKey], [limit.toString(), windowSeconds.toString(), now.toString()] ); const [current, remaining, resetInSeconds, returnedLimit] = result; return { allowed: current <= limit, current, remaining: Math.max(0, remaining), resetInSeconds: Math.max(1, resetInSeconds), limit: returnedLimit }; } catch (error) { console.error("[LimitRate] Upstash rate check error:", error); throw error; } } async peekRate(key, limit, windowSeconds) { const prefixedKey = `${this.keyPrefix}rate:${key}`; try { const response = await fetch(`${this.url}/get/${encodeURIComponent(prefixedKey)}`, { headers: { "Authorization": `Bearer ${this.token}` } }); if (!response.ok && response.status !== 404) { throw new Error(`Upstash API error: ${response.status}`); } const data = await response.json(); const current = data.result ? parseInt(String(data.result), 10) : 0; const ttlResponse = await fetch(`${this.url}/ttl/${encodeURIComponent(prefixedKey)}`, { headers: { "Authorization": `Bearer ${this.token}` } }); const ttlData = await ttlResponse.json(); const resetInSeconds = ttlData.result && ttlData.result > 0 ? ttlData.result : windowSeconds; const remaining = Math.max(0, limit - current); return { allowed: current < limit, current, remaining, resetInSeconds, limit }; } catch (error) { console.error("[LimitRate] Upstash peek rate error:", error); throw error; } } async incrementCost(key, cost, windowSeconds, cap) { const prefixedKey = `${this.keyPrefix}cost:${key}`; try { const result = await this.client.eval( COST_INCREMENT_SCRIPT2, [prefixedKey], [cost.toString(), windowSeconds.toString(), cap.toString()] ); const [current, allowed, remaining, resetInSeconds, returnedCap] = result; return { allowed: allowed === 1, current, remaining: Math.max(0, remaining), resetInSeconds: Math.max(1, resetInSeconds), cap: returnedCap }; } catch (error) { console.error("[LimitRate] Upstash cost increment error:", error); throw error; } } async incrementTokens(key, tokens, windowSeconds, limit) { const prefixedKey = `${this.keyPrefix}tokens:${key}`; try { const result = await this.client.eval( TOKEN_INCREMENT_SCRIPT2, [prefixedKey], [tokens.toString(), windowSeconds.toString(), limit.toString()] ); const [current, allowed, remaining, resetInSeconds, returnedLimit] = result; return { allowed: allowed === 1, current, remaining: Math.max(0, remaining), resetInSeconds: Math.max(1, resetInSeconds), limit: returnedLimit }; } catch (error) { console.error("[LimitRate] Upstash token increment error:", error); throw error; } } async ping() { try { const result = await this.client.ping(); return result === "PONG"; } catch { return false; } } async close() { return Promise.resolve(); } /** * Generic get method for arbitrary data (v2.0.0 - D4) */ async get(key) { try { const prefixedKey = `${this.keyPrefix}generic:${key}`; const value = await this.client.get(prefixedKey); if (!value) { return null; } return value; } catch (error) { console.error("[LimitRate] Upstash get error:", error); throw error; } } /** * Generic set method for arbitrary data (v2.0.0 - D4) */ async set(key, value, ttl) { try { const prefixedKey = `${this.keyPrefix}generic:${key}`; const expirySeconds = ttl || 86400; await this.client.setex(prefixedKey, expirySeconds, value); } catch (error) { console.error("[LimitRate] Upstash set error:", error); throw error; } } /** * Generic delete method (v2.0.0 - D4) */ async delete(key) { try { const prefixedKey = `${this.keyPrefix}generic:${key}`; await this.client.del(prefixedKey); } catch (error) { console.error("[LimitRate] Upstash delete error:", error); throw error; } } /** * Get underlying Upstash client (for advanced use cases) */ getClient() { return this.client; } }; // src/stores/index.ts function createStore(config) { switch (config.type) { case "memory": return new MemoryStore(); case "redis": if (!config.url) { throw new Error("Redis store requires url"); } return new RedisStore({ client: config.url, redisOptions: config.options }); case "upstash": if (!config.url || !config.token) { throw new Error("Upstash store requires url and token"); } return new UpstashStore({ url: config.url, token: config.token }); default: throw new Error(`Unknown store type: ${config.type}`); } } function createSharedMemoryStore(options) { return new MemoryStore(options); } function createSharedRedisStore(options) { if ("url" in options) { return new RedisStore({ client: options.url, keyPrefix: options.keyPrefix, redisOptions: options.redisOptions }); } return new RedisStore({ client: options.client, keyPrefix: options.keyPrefix }); } function createSharedUpstashStore(options) { return new UpstashStore(options); } // src/utils/events.ts var EventEmitter = class { handlers; constructor() { this.handlers = /* @__PURE__ */ new Set(); } /** * Register an event handler */ on(handler) { this.handlers.add(handler); } /** * Unregister an event handler */ off(handler) { this.handlers.delete(handler); } /** * Emit an event to all handlers */ async emit(event) { const promises = []; for (const handler of this.handlers) { try { const result = handler(event); if (result instanceof Promise) { promises.push(result); } } catch (error) { console.error("[LimitRate] Event handler error:", error); } } if (promises.length > 0) { await Promise.allSettled(promises); } } /** * Get number of registered handlers */ getHandlerCount() { return this.handlers.size; } /** * Clear all handlers */ clear() { this.handlers.clear(); } }; // src/penalty/manager.ts var PenaltyManager = class { store; constructor(store) { this.store = store; } /** * Get penalty state key for a user/endpoint */ getPenaltyKey(user, endpoint) { return `penalty:${user}:${endpoint}`; } /** * Get current penalty/reward multiplier for a user/endpoint * Returns 1.0 if no penalty/reward is active */ async getMultiplier(user, endpoint) { try { const key = this.getPenaltyKey(user, endpoint); const state = await this.store.get(key); if (!state) { return 1; } if (Date.now() > state.expiresAt) { await this.store.delete(key); return 1; } return state.multiplier; } catch (error) { console.warn("[LimitRate] Failed to get penalty multiplier:", error); return 1; } } /** * Apply a penalty to a user/endpoint */ async applyPenalty(user, endpoint, config) { try { const key = this.getPenaltyKey(user, endpoint); const state = { multiplier: config.multiplier, expiresAt: Date.now() + config.duration * 1e3, reason: "violation" }; await this.store.set(key, state, config.duration); } catch (error) { console.error("[LimitRate] Failed to apply penalty:", error); } } /** * Apply a reward to a user/endpoint */ async applyReward(user, endpoint, config) { try { const key = this.getPenaltyKey(user, endpoint); const state = { multiplier: config.multiplier, expiresAt: Date.now() + config.duration * 1e3, reason: "reward" }; await this.store.set(key, state, config.duration); } catch (error) { console.error("[LimitRate] Failed to apply reward:", error); } } /** * Check if a reward should be granted based on usage */ shouldGrantReward(currentUsage, limit, rewardConfig) { const usagePercent = currentUsage / limit * 100; switch (rewardConfig.trigger) { case "below_10_percent": return usagePercent < 10; case "below_25_percent": return usagePercent < 25; case "below_50_percent": return usagePercent < 50; default: return false; } } /** * Clear penalty/reward for a user/endpoint */ async clear(user, endpoint) { try { const key = this.getPenaltyKey(user, endpoint); await this.store.delete(key); } catch (error) { console.error("[LimitRate] Failed to clear penalty:", error); } } }; // src/engine.ts var PolicyEngine = class { store; policies; events; penaltyManager; constructor(store, policies) { this.store = store; this.policies = policies; this.events = new EventEmitter(); this.penaltyManager = new PenaltyManager(store); } /** * Register event handler */ onEvent(handler) { this.events.on(handler); } /** * Check if request should be allowed */ async check(context) { const policy = context.policyOverride || this.resolvePolicy(context.plan, context.endpoint); if (!policy) { await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "allowed" }); return { allowed: true, action: "allow", details: { used: 0, limit: Infinity, remaining: Infinity, resetInSeconds: 0 } }; } let finalDetails = { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 }; if (policy.rate) { const rateResult = await this.checkRate(context, policy); if (!rateResult.allowed || rateResult.action !== "allow") { return rateResult; } finalDetails = rateResult.details; } if (policy.rate && context.tokens !== void 0 && context.tokens > 0) { const tokenResult = await this.checkTokens(context, policy); if (!tokenResult.allowed || tokenResult.action !== "allow") { return tokenResult; } } if (policy.cost) { const costResult = await this.checkCost(context, policy); if (!costResult.allowed || costResult.action !== "allow") { return costResult; } } await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "allowed" }); return { allowed: true, action: "allow", details: finalDetails }; } /** * Check rate limit */ async checkRate(context, policy) { if (!policy.rate) { return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } }; } let { maxPerSecond, maxPerMinute, maxPerHour, maxPerDay, burst, actionOnExceed, slowdownMs } = policy.rate; if (context.userOverride) { const override = context.userOverride; const isValidLimit = (val) => { return typeof val === "number" && val > 0 && !isNaN(val) && isFinite(val); }; const endpointOverride = override.endpoints?.[context.endpoint]; if (endpointOverride) { if (isValidLimit(endpointOverride.maxPerSecond)) maxPerSecond = endpointOverride.maxPerSecond; if (isValidLimit(endpointOverride.maxPerMinute)) maxPerMinute = endpointOverride.maxPerMinute; if (isValidLimit(endpointOverride.maxPerHour)) maxPerHour = endpointOverride.maxPerHour; if (isValidLimit(endpointOverride.maxPerDay)) maxPerDay = endpointOverride.maxPerDay; } else { if (isValidLimit(override.maxPerSecond)) maxPerSecond = override.maxPerSecond; if (isValidLimit(override.maxPerMinute)) maxPerMinute = override.maxPerMinute; if (isValidLimit(override.maxPerHour)) maxPerHour = override.maxPerHour; if (isValidLimit(override.maxPerDay)) maxPerDay = override.maxPerDay; } } let limit; let windowSeconds; let windowLabel; if (maxPerSecond !== void 0) { limit = maxPerSecond; windowSeconds = 1; windowLabel = "1s"; } else if (maxPerMinute !== void 0) { limit = maxPerMinute; windowSeconds = 60; windowLabel = "1m"; } else if (maxPerHour !== void 0) { limit = maxPerHour; windowSeconds = 3600; windowLabel = "1h"; } else if (maxPerDay !== void 0) { limit = maxPerDay; windowSeconds = 86400; windowLabel = "1d"; } else { throw new Error("No time window specified in rate rule"); } const originalLimit = limit; if (policy.penalty?.enabled) { const multiplier = await this.penaltyManager.getMultiplier(context.user, context.endpoint); limit = Math.floor(originalLimit * multiplier); if (limit < 1) limit = 1; } const rateKey = `${context.user}:${context.endpoint}`; const result = await this.store.checkRate(rateKey, limit, windowSeconds, burst); if (result.allowed) { if (policy.penalty?.enabled && policy.penalty.rewards) { const shouldReward = this.penaltyManager.shouldGrantReward( result.current, originalLimit, policy.penalty.rewards ); if (shouldReward) { await this.penaltyManager.applyReward( context.user, context.endpoint, policy.penalty.rewards ); } } return { allowed: true, action: "allow", details: { used: result.current, limit: result.limit, remaining: result.remaining, resetInSeconds: result.resetInSeconds, burstTokens: result.burstTokens } }; } if (policy.penalty?.enabled && policy.penalty.onViolation) { await this.penaltyManager.applyPenalty( context.user, context.endpoint, policy.penalty.onViolation ); } await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "rate_exceeded", window: windowLabel, value: result.current, threshold: limit }); if (actionOnExceed === "block") { return { allowed: false, action: "block", reason: "rate_exceeded", retryAfterSeconds: result.resetInSeconds, details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds, burstTokens: result.burstTokens } }; } if (actionOnExceed === "slowdown") { await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "slowdown_applied", value: slowdownMs }); return { allowed: true, action: "slowdown", slowdownMs, details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds, burstTokens: result.burstTokens } }; } if (actionOnExceed === "allow-and-log") { return { allowed: true, action: "allow-and-log", details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds, burstTokens: result.burstTokens } }; } return { allowed: true, action: "allow", details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds, burstTokens: result.burstTokens } }; } /** * Check cost limit */ async checkCost(context, policy) { if (!policy.cost) { return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } }; } const { estimateCost, hourlyCap, dailyCap, actionOnExceed } = policy.cost; const cost = estimateCost(context.costContext); const cap = dailyCap ?? hourlyCap; const windowSeconds = dailyCap ? 86400 : 3600; const costKey = `${context.user}:${context.endpoint}:cost`; const result = await this.store.incrementCost(costKey, cost, windowSeconds, cap); if (result.allowed) { return { allowed: true, action: "allow", details: { used: result.current, limit: cap, remaining: result.remaining, resetInSeconds: result.resetInSeconds } }; } await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "cost_exceeded", window: dailyCap ? "1d" : "1h", value: result.current + cost, threshold: cap }); if (actionOnExceed === "block") { return { allowed: false, action: "block", reason: "cost_exceeded", retryAfterSeconds: result.resetInSeconds, details: { used: result.current, limit: cap, remaining: 0, resetInSeconds: result.resetInSeconds } }; } if (actionOnExceed === "slowdown") { return { allowed: false, action: "block", reason: "cost_exceeded", retryAfterSeconds: result.resetInSeconds, details: { used: result.current, limit: cap, remaining: 0, resetInSeconds: result.resetInSeconds } }; } if (actionOnExceed === "allow-and-log") { return { allowed: true, action: "allow-and-log", details: { used: result.current, limit: cap, remaining: 0, resetInSeconds: result.resetInSeconds } }; } return { allowed: true, action: "allow", details: { used: result.current, limit: cap, remaining: 0, resetInSeconds: result.resetInSeconds } }; } /** * Check token limits (v1.4.0 - AI feature) */ async checkTokens(context, policy) { if (!policy.rate || !context.tokens) { return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } }; } const { maxTokensPerMinute, maxTokensPerHour, maxTokensPerDay, actionOnExceed, slowdownMs } = policy.rate; if (!maxTokensPerMinute && !maxTokensPerHour && !maxTokensPerDay) { return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } }; } const checks = []; if (maxTokensPerMinute) { checks.push({ limit: maxTokensPerMinute, windowSeconds: 60, windowLabel: "1m" }); } if (maxTokensPerHour) { checks.push({ limit: maxTokensPerHour, windowSeconds: 3600, windowLabel: "1h" }); } if (maxTokensPerDay) { checks.push({ limit: maxTokensPerDay, windowSeconds: 86400, windowLabel: "1d" }); } for (const check of checks) { const tokenKey = `${context.user}:${context.endpoint}:tokens:${check.windowLabel}`; const result = await this.store.incrementTokens( tokenKey, context.tokens, check.windowSeconds, check.limit ); if (!result.allowed) { await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "token_limit_exceeded", window: check.windowLabel, value: result.current, threshold: check.limit, tokens: context.tokens }); if (actionOnExceed === "block") { return { allowed: false, action: "block", reason: "token_limit_exceeded", retryAfterSeconds: result.resetInSeconds, details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds } }; } if (actionOnExceed === "slowdown") { await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "slowdown_applied", value: slowdownMs }); return { allowed: true, action: "slowdown", slowdownMs, details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds } }; } if (actionOnExceed === "allow-and-log") { return { allowed: true, action: "allow-and-log", details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds } }; } return { allowed: true, action: "allow", details: { used: result.current, limit: result.limit, remaining: 0, resetInSeconds: result.resetInSeconds } }; } } await this.emitEvent({ timestamp: Date.now(), user: context.user, plan: context.plan, endpoint: context.endpoint, type: "token_usage_tracked", tokens: context.tokens }); return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } }; } /** * Resolve policy for plan and endpoint */ resolvePolicy(plan, endpoint) { const planConfig = this.policies[plan]; if (!planConfig) { return null; } if (planConfig.endpoints?.[endpoint]) { return planConfig.endpoints[endpoint]; } return planConfig.defaults ?? null; } /** * Emit event */ async emitEvent(event) { await this.events.emit(event); } /** * Get event emitter (for external handlers) */ getEventEmitter() { return this.events; } }; // src/validation.ts var ValidationError = class extends Error { constructor(message) { super(`[LimitRate] Invalid configuration: ${message}`); this.name = "ValidationError"; } }; function isPositiveNumber(value) { return typeof value === "number" && value > 0 && !isNaN(value); } function isFunction(value) { return typeof value === "function"; } function validateRateRule(rule, path) { if (rule.maxPerSecond !== void 0 && !isPositiveNumber(rule.maxPerSecond)) { throw new ValidationError(`${path}.maxPerSecond must be a positive number, got: ${rule.maxPerSecond}`); } if (rule.maxPerMinute !== void 0 && !isPositiveNumber(rule.maxPerMinute)) { throw new ValidationError(`${path}.maxPerMinute must be a positive number, got: ${rule.maxPerMinute}`); } if (rule.maxPerHour !== void 0 && !isPositiveNumber(rule.maxPerHour)) { throw new ValidationError(`${path}.maxPerHour must be a positive number, got: ${rule.maxPerHour}`); } if (rule.maxPerDay !== void 0 && !isPositiveNumber(rule.maxPerDay)) { throw new ValidationError(`${path}.maxPerDay must be a positive number, got: ${rule.maxPerDay}`); } const timeWindows = [ rule.maxPerSecond, rule.maxPerMinute, rule.maxPerHour, rule.maxPerDay ].filter((w) => w !== void 0); if (timeWindows.length === 0) { throw new ValidationError( `${path} must specify exactly one time window (maxPerSecond, maxPerMinute, maxPerHour, or maxPerDay)` ); } if (timeWindows.length > 1) { throw new ValidationError( `${path} can only specify one time window, but found ${timeWindows.length}. Use separate endpoint policies for multiple limits.` ); } if (rule.burst !== void 0 && !isPositiveNumber(rule.burst)) { throw new ValidationError(`${path}.burst must be a positive number, got: ${rule.burst}`); } if (rule.slowdownMs !== void 0) { if (!isPositiveNumber(rule.slowdownMs)) { throw new ValidationError(`${path}.slowdownMs must be a positive number, got: ${rule.slowdownMs}`); } if (rule.slowdownMs > 6e4) { throw new ValidationError(`${path}.slowdownMs must be <= 60000 (60 seconds), got: ${rule.slowdownMs}`); } } const validActions = ["allow", "block", "slowdown", "allow-and-log"]; if (!validActions.includes(rule.actionOnExceed)) { throw new ValidationError( `${path}.actionOnExceed must be one of: ${validActions.join(", ")}, got: ${rule.actionOnExceed}` ); } if (rule.actionOnExceed === "slowdown" && !rule.slowdownMs) { throw new ValidationError(`${path}.slowdownMs is required when actionOnExceed is 'slowdown'`); } } function validateCostRule(rule, path) { if (!isFunction(rule.estimateCost)) { throw new ValidationError(`${path}.estimateCost must be a function`); } if (rule.hourlyCap !== void 0 && !isPositiveNumber(rule.hourlyCap)) { throw new ValidationError(`${path}.hourlyCap must be a positive number, got: ${rule.hourlyCap}`); } if (rule.dailyCap !== void 0 && !isPositiveNumber(rule.dailyCap)) { throw new ValidationError(`${path}.dailyCap must be a positive number, got: ${rule.dailyCap}`); } const validActions = ["allow", "block", "slowdown", "allow-and-log"]; if (!validActions.includes(rule.actionOnExceed)) { throw new ValidationError( `${path}.actionOnExceed must be one of: ${validActions.join(", ")}, got: ${rule.actionOnExceed}` ); } } function validatePolicyConfig(config) { if (!config || typeof config !== "object") { throw new ValidationError("PolicyConfig must be an object"); } for (const [plan, planConfig] of Object.entries(config)) { if (!planConfig || typeof planConfig !== "object") { throw new ValidationError(`Plan '${plan}' config must be an object`); } const hasEndpoints = planConfig.endpoints && typeof planConfig.endpoints === "object"; const hasDefaults = planConfig.defaults && typeof planConfig.defaults === "object"; if (!hasEndpoints && !hasDefaults) { throw new ValidationError(`Plan '${plan}' must have 'endpoints' or 'defaults' object`); } if (hasEndpoints) { for (const [endpoint, policy] of Object.entries(planConfig.endpoints)) { const basePath = `policies.${plan}.endpoints['${endpoint}']`; if (policy.rate) { validateRateRule(policy.rate, `${basePath}.rate`); } if (policy.cost) { validateCostRule(policy.cost, `${basePath}.cost`); } if (!policy.rate && !policy.cost && !policy.concurrency) { throw new ValidationError(`${basePath} must have at least one of: rate, cost, concurrency`); } } } if (planConfig.defaults) { const basePath = `policies.${plan}.defaults`; if (planConfig.defaults.rate) { validateRateRule(planConfig.defaults.rate, `${basePath}.rate`); } if (planConfig.defaults.cost) { validateCostRule(planConfig.defaults.cost, `${basePath}.cost`); } } } } function validateStoreConfig(config) { if (!config || typeof config !== "object") { throw new ValidationError("StoreConfig must be an object"); } const validTypes = ["memory", "redis", "upstash"]; if (!validTypes.includes(config.type)) { throw new ValidationError(`store.type must be one of: ${validTypes.join(", ")}, got: ${config.type}`); } if (config.type === "redis" && !config.url) { throw new ValidationError('store.url is required when type is "redis"'); } if (config.type === "upstash") { if (!config.url) { throw new ValidationError('store.url is required when type is "upstash"'); } if (!config.token) { throw new ValidationError('store.token is required when type is "upstash"'); } } } function validateIPAddress(ip) { const ipv4Regex = /^(\d{1,3}\.){3}\d{1,3}$/; if (ipv4Regex.test(ip)) { const parts = ip.split("."); return parts.every((part) => { const num = parseInt(part, 10); return num >= 0 && num <= 255; }); } const ipv4MappedRegex = /^::ffff:(\d{1,3}\.){3}\d{1,3}$/i; if (ipv4MappedRegex.test(ip)) { const ipv4Part = ip.substring(7); const parts = ipv4Part.split("."); return parts.every((part) => { const num = parseInt(part, 10);