UNPKG

zenin-limiter

Version:

Universal rate & throttle limiter middleware for Express, Fastify, and custom handlers

856 lines (843 loc) 25.8 kB
'use strict'; var common = require('@nestjs/common'); var __getOwnPropDesc = Object.getOwnPropertyDescriptor; var __decorateClass = (decorators, target, key, kind) => { var result = kind > 1 ? void 0 : kind ? __getOwnPropDesc(target, key) : target; for (var i = decorators.length - 1, decorator; i >= 0; i--) if (decorator = decorators[i]) result = (decorator(result)) || result; return result; }; // src/strategies/memoryStore.ts var memoryStore = /* @__PURE__ */ new Map(); var heap = []; var hits = 0; var rejections = 0; var callCount = 0; var lruHead = null; var lruTail = null; var perKeyStats = /* @__PURE__ */ new Map(); var lockPromise = Promise.resolve(); var gcInterval = null; function heapPush(node) { heap.push(node); let i = heap.length - 1; while (i > 0) { const parent = Math.floor((i - 1) / 2); if (heap[parent].expiresAt <= node.expiresAt) break; heap[i] = heap[parent]; i = parent; } heap[i] = node; } function heapPop() { if (heap.length === 0) return void 0; const result = heap[0]; const last = heap.pop(); if (heap.length > 0) { heap[0] = last; let i = 0; while (true) { const left = 2 * i + 1; const right = 2 * i + 2; let smallest = i; if (left < heap.length && heap[left].expiresAt < heap[smallest].expiresAt) { smallest = left; } if (right < heap.length && heap[right].expiresAt < heap[smallest].expiresAt) { smallest = right; } if (smallest === i) break; [heap[i], heap[smallest]] = [heap[smallest], heap[i]]; i = smallest; } } return result; } function addLruNode(key) { const node = { key, prev: null, next: lruHead }; if (lruHead) lruHead.prev = node; lruHead = node; if (!lruTail) lruTail = node; return node; } function moveToFront(node) { if (node === lruHead) return; if (node.prev) node.prev.next = node.next; if (node.next) node.next.prev = node.prev; if (node === lruTail) lruTail = node.prev; node.next = lruHead; node.prev = null; if (lruHead) lruHead.prev = node; lruHead = node; if (!lruTail) lruTail = node; } function removeLruTail() { if (!lruTail) return; memoryStore.delete(lruTail.key); perKeyStats.delete(lruTail.key); if (lruTail.prev) { lruTail.prev.next = null; lruTail = lruTail.prev; } else { lruHead = null; lruTail = null; } } async function resetKey(key) { if (typeof key !== "string" || key.trim() === "") { throw new Error("Invalid key"); } const unlock = await acquireLock(); try { const entry = memoryStore.get(key); if (entry) { if (entry.lruNode.prev) entry.lruNode.prev.next = entry.lruNode.next; if (entry.lruNode.next) entry.lruNode.next.prev = entry.lruNode.prev; if (entry.lruNode === lruHead) lruHead = entry.lruNode.next; if (entry.lruNode === lruTail) lruTail = entry.lruNode.prev; memoryStore.delete(key); perKeyStats.delete(key); } } finally { unlock(); } } async function resetAll() { const unlock = await acquireLock(); try { memoryStore.clear(); heap.length = 0; hits = 0; rejections = 0; callCount = 0; lruHead = null; lruTail = null; perKeyStats.clear(); } finally { unlock(); } } async function getMetrics(key) { const unlock = await acquireLock(); try { if (key && perKeyStats.has(key)) { return { ...perKeyStats.get(key) }; } return { hits, rejections }; } finally { unlock(); } } async function acquireLock() { const currentLock = lockPromise; let resolveLock; lockPromise = new Promise((resolve) => { resolveLock = resolve; }); await currentLock; return () => resolveLock(); } function sweepExpiredKeys(maxBatchCleanup = 1e3) { const now = Date.now(); let cleaned = 0; while (heap.length > 0 && heap[0].expiresAt <= now && cleaned < maxBatchCleanup) { const expired = heapPop(); const current = memoryStore.get(expired.key); if (current && current.expiresAt === expired.expiresAt) { if (current.lruNode.prev) current.lruNode.prev.next = current.lruNode.next; if (current.lruNode.next) current.lruNode.next.prev = current.lruNode.prev; if (current.lruNode === lruHead) lruHead = current.lruNode.next; if (current.lruNode === lruTail) lruTail = current.lruNode.prev; memoryStore.delete(expired.key); perKeyStats.delete(expired.key); } cleaned++; } } async function isAllowedMemory(key, limit, windowInSeconds, config = {}, now = Date.now) { if (typeof key !== "string" || key.trim() === "") throw new Error("Invalid key"); if (!Number.isFinite(limit) || limit <= 0 || !Number.isFinite(windowInSeconds) || windowInSeconds <= 0) throw new Error("Invalid limit or windowInSeconds"); if (windowInSeconds * 1e3 > Number.MAX_SAFE_INTEGER) { throw new Error("Window too large for safe expiration"); } const { maxStoreSize = 1e6, cleanupInterval = 1e3, enablePerKeyStats = false, maxBatchCleanup = 1e3 } = config; const unlock = await acquireLock(); try { const currentTime = now(); callCount++; let cleanupCount = 0; while (heap.length > 0 && heap[0].expiresAt <= currentTime && cleanupCount < maxBatchCleanup) { const expired = heapPop(); const current = memoryStore.get(expired.key); if (current && current.expiresAt === expired.expiresAt) { if (current.lruNode.prev) current.lruNode.prev.next = current.lruNode.next; if (current.lruNode.next) current.lruNode.next.prev = current.lruNode.prev; if (current.lruNode === lruHead) lruHead = current.lruNode.next; if (current.lruNode === lruTail) lruTail = current.lruNode.prev; memoryStore.delete(expired.key); perKeyStats.delete(expired.key); } cleanupCount++; } if (callCount % cleanupInterval === 0 && heap.length > 0) { const threshold = currentTime - windowInSeconds * 1e3; cleanupCount = 0; while (heap.length > 0 && heap[0].expiresAt <= threshold && cleanupCount < maxBatchCleanup) { const expired = heapPop(); const current = memoryStore.get(expired.key); if (current && current.expiresAt === expired.expiresAt) { if (current.lruNode.prev) current.lruNode.prev.next = current.lruNode.next; if (current.lruNode.next) current.lruNode.next.prev = current.lruNode.prev; if (current.lruNode === lruHead) lruHead = current.lruNode.next; if (current.lruNode === lruTail) lruTail = current.lruNode.prev; memoryStore.delete(expired.key); perKeyStats.delete(expired.key); } cleanupCount++; } } while (memoryStore.size >= maxStoreSize) { removeLruTail(); } let entry = memoryStore.get(key); if (!entry) { const expiresAt = currentTime + windowInSeconds * 1e3; const lruNode = addLruNode(key); entry = { count: 1, expiresAt, lruNode }; memoryStore.set(key, entry); heapPush({ key, expiresAt }); hits++; if (enablePerKeyStats) { perKeyStats.set(key, { hits: (perKeyStats.get(key)?.hits || 0) + 1, rejections: perKeyStats.get(key)?.rejections || 0 }); } return true; } moveToFront(entry.lruNode); if (entry.expiresAt <= currentTime) { const expiresAt = currentTime + windowInSeconds * 1e3; entry.count = 1; entry.expiresAt = expiresAt; heapPush({ key, expiresAt }); hits++; if (enablePerKeyStats) { perKeyStats.set(key, { hits: (perKeyStats.get(key)?.hits || 0) + 1, rejections: perKeyStats.get(key)?.rejections || 0 }); } return true; } if (entry.count < limit) { entry.count++; hits++; if (enablePerKeyStats) { perKeyStats.set(key, { hits: (perKeyStats.get(key)?.hits || 0) + 1, rejections: perKeyStats.get(key)?.rejections || 0 }); } return true; } rejections++; if (enablePerKeyStats) { perKeyStats.set(key, { hits: perKeyStats.get(key)?.hits || 0, rejections: (perKeyStats.get(key)?.rejections || 0) + 1 }); } return false; } finally { unlock(); } } var FixedWindowStrategy = class { constructor(config) { this.gcStarted = false; this.config = config; this.maxEntries = config.limiterConfig?.maxStoreSize || 1e6; this.getLimitFn = typeof config.limit === "function" ? config.limit : void 0; this.startGC(); } getLimit(req) { if (this.getLimitFn) { return this.getLimitFn(req); } return this.config.limit; } startGC() { if (this.gcStarted) return; gcInterval = setInterval(() => { sweepExpiredKeys(this.config.limiterConfig?.maxBatchCleanup || 1e3); while (memoryStore.size > this.maxEntries) { removeLruTail(); } }, 3e4); this.gcStarted = true; } stopGC() { if (gcInterval) clearInterval(gcInterval); this.gcStarted = false; } async isAllowed(key, req) { while (memoryStore.size >= this.maxEntries) { removeLruTail(); } return isAllowedMemory( key, this.getLimit(req), this.config.windowInSeconds, this.config.limiterConfig ); } async reset(key) { return resetKey(key); } async getState(key) { return { remaining: 0, resetAt: 0, limit: this.config.limit }; } static async resetAll() { await resetAll(); } }; // src/strategies/slidingWindow.ts var slidingWindowStore = /* @__PURE__ */ new Map(); var gcInterval2 = null; function sweepExpiredKeys2(maxEntries = 1e6) { const now = Date.now(); for (const [key, entry] of slidingWindowStore.entries()) { entry.timestamps = entry.timestamps.filter( (ts) => ts > now - entry.expiresAt ); if (entry.timestamps.length === 0) { slidingWindowStore.delete(key); } } while (slidingWindowStore.size > maxEntries) { const firstKey = slidingWindowStore.keys().next().value; slidingWindowStore.delete(firstKey); } } var SlidingWindowStrategy = class { constructor(config) { this.gcStarted = false; this.config = config; this.limit = config.limit; this.windowMs = config.windowInSeconds * 1e3; this.maxEntries = config.limiterConfig?.maxStoreSize || 1e6; this.getLimitFn = typeof config.limit === "function" ? config.limit : void 0; this.startGC(); } startGC() { if (this.gcStarted) return; gcInterval2 = setInterval(() => { sweepExpiredKeys2(this.maxEntries); }, 3e4); this.gcStarted = true; } stopGC() { if (gcInterval2) clearInterval(gcInterval2); this.gcStarted = false; } getLimit(req) { if (this.getLimitFn) { return this.getLimitFn(req); } return this.limit; } async isAllowed(key, req) { while (slidingWindowStore.size >= this.maxEntries) { const firstKey = slidingWindowStore.keys().next().value; slidingWindowStore.delete(firstKey); } const now = Date.now(); const windowStart = now - this.windowMs; let entry = slidingWindowStore.get(key); if (!entry) { entry = { timestamps: [], expiresAt: now + this.windowMs }; slidingWindowStore.set(key, entry); } entry.timestamps = entry.timestamps.filter( (timestamp) => timestamp > windowStart ); if (entry.timestamps.length < this.getLimit(req)) { entry.timestamps.push(now); return true; } return false; } async getState(key, req) { const now = Date.now(); const windowStart = now - this.windowMs; const entry = slidingWindowStore.get(key); const limit = this.getLimit(req); if (!entry) { return { remaining: limit, resetAt: now + this.windowMs, limit }; } entry.timestamps = entry.timestamps.filter( (timestamp) => timestamp > windowStart ); return { remaining: Math.max(0, limit - entry.timestamps.length), resetAt: entry.timestamps.length > 0 ? entry.timestamps[0] + this.windowMs : now + this.windowMs, limit }; } async reset(key) { slidingWindowStore.delete(key); } static async resetAll() { slidingWindowStore.clear(); } }; // src/strategies/tokenBucket.ts var tokenBucketStore = /* @__PURE__ */ new Map(); var gcInterval3 = null; function sweepExpiredKeys3(maxEntries = 1e6) { const now = Date.now(); for (const [key, entry] of tokenBucketStore.entries()) { if (entry.tokens <= 0 && now - entry.lastRefill > entry.capacity / entry.refillRate) { tokenBucketStore.delete(key); } } while (tokenBucketStore.size > maxEntries) { const firstKey = tokenBucketStore.keys().next().value; tokenBucketStore.delete(firstKey); } } var TokenBucketStrategy = class { constructor(config) { this.gcStarted = false; this.config = config; this.capacity = config.limit; this.refillRate = config.limit / (config.windowInSeconds * 1e3); this.maxEntries = config.limiterConfig?.maxStoreSize || 1e6; this.getLimitFn = typeof config.limit === "function" ? config.limit : void 0; this.startGC(); } startGC() { if (this.gcStarted) return; gcInterval3 = setInterval(() => { sweepExpiredKeys3(this.maxEntries); }, 3e4); this.gcStarted = true; } stopGC() { if (gcInterval3) clearInterval(gcInterval3); this.gcStarted = false; } getLimit(req) { if (this.getLimitFn) { return this.getLimitFn(req); } return this.capacity; } async isAllowed(key, req) { while (tokenBucketStore.size >= this.maxEntries) { const firstKey = tokenBucketStore.keys().next().value; tokenBucketStore.delete(firstKey); } const now = Date.now(); const limit = this.getLimit(req); let entry = tokenBucketStore.get(key); if (!entry) { entry = { tokens: limit, lastRefill: now, capacity: limit, refillRate: limit / (this.config.windowInSeconds * 1e3) }; tokenBucketStore.set(key, entry); } const timeElapsed = now - entry.lastRefill; const tokensToAdd = timeElapsed * entry.refillRate; entry.tokens = Math.min(entry.capacity, entry.tokens + tokensToAdd); entry.lastRefill = now; if (entry.tokens >= 1) { entry.tokens -= 1; return true; } return false; } async getState(key, req) { const now = Date.now(); const limit = this.getLimit(req); const entry = tokenBucketStore.get(key); if (!entry) { return { remaining: limit, resetAt: now + limit / (limit / (this.config.windowInSeconds * 1e3)), limit }; } const timeElapsed = now - entry.lastRefill; const tokensToAdd = timeElapsed * entry.refillRate; const currentTokens = Math.min(entry.capacity, entry.tokens + tokensToAdd); return { remaining: Math.floor(currentTokens), resetAt: now + (limit - currentTokens) / entry.refillRate, limit }; } async reset(key) { tokenBucketStore.delete(key); } static async resetAll() { tokenBucketStore.clear(); } }; // src/utils/configDefaults.ts function applyDefaults(config) { const defaultLimiterConfig = { maxStoreSize: 1e6, cleanupInterval: 1e3, enablePerKeyStats: false, maxBatchCleanup: 1e3 }; return { limit: 100, windowInSeconds: 60, strategy: "fixed", keyType: "ip", debug: false, dryRun: false, silent: false, limiterConfig: { ...defaultLimiterConfig, ...config.limiterConfig }, ...config }; } // src/utils/configValidator.ts function validateConfig(config) { const errors = []; if (typeof config.limit === "function") ; else if (typeof config.limit !== "number" || config.limit <= 0) { errors.push({ field: "limit", message: "Limit must be a positive number or a function" }); } if (typeof config.windowInSeconds !== "number" || config.windowInSeconds <= 0) { errors.push({ field: "windowInSeconds", message: "windowInSeconds must be a positive number" }); } if (config.strategy && !["fixed", "sliding", "tokenBucket"].includes(config.strategy)) { errors.push({ field: "strategy", message: "Strategy must be one of: fixed, sliding, tokenBucket" }); } if (config.keyType && !["ip", "user-agent", "path", "custom"].includes(config.keyType) && !config.keyType.startsWith("header:")) { errors.push({ field: "keyType", message: "keyType must be one of: ip, user-agent, path, custom, or header:HEADER_NAME" }); } if (config.limiterConfig) { if (config.limiterConfig.maxStoreSize && (typeof config.limiterConfig.maxStoreSize !== "number" || config.limiterConfig.maxStoreSize <= 0)) { errors.push({ field: "limiterConfig.maxStoreSize", message: "maxStoreSize must be a positive number" }); } if (config.limiterConfig.cleanupInterval && (typeof config.limiterConfig.cleanupInterval !== "number" || config.limiterConfig.cleanupInterval <= 0)) { errors.push({ field: "limiterConfig.cleanupInterval", message: "cleanupInterval must be a positive number" }); } } return errors; } function throwIfInvalid(config) { const errors = validateConfig(config); if (errors.length > 0) { const errorMessage = errors.map((error) => `${error.field}: ${error.message}`).join(", "); throw new Error(`Invalid rate limiter configuration: ${errorMessage}`); } } // src/core/RateLimiter.ts var RateLimiter = class { constructor(config, strategy) { this.stats = { totalRequests: 0, hits: 0, rejections: 0 }; this.config = applyDefaults(config); throwIfInvalid(this.config); this.strategy = strategy || this.createStrategy(this.config); } createStrategy(config) { const strategyType = config.strategy || "fixed"; switch (strategyType) { case "sliding": return new SlidingWindowStrategy(config); case "tokenBucket": return new TokenBucketStrategy(config); case "fixed": default: return new FixedWindowStrategy(config); } } getLimit(req) { const limit = this.config.limit; if (typeof limit === "function") { return limit(req); } return limit || 100; } callHook(hookName, ...args) { const hook = this.config[hookName]; if (hook) { try { hook(...args); } catch (error) { if (this.config.onError) { this.config.onError(error); } } } } logDebug(message, data) { if (this.config.debug) { console.log(`[RateLimiter] ${message}`, data || ""); } } async isAllowed(key, req) { this.stats.totalRequests++; this.logDebug(`Checking rate limit for key: ${key}`); try { const allowed = await this.strategy.isAllowed(key, req); if (allowed) { this.stats.hits++; this.callHook("onPass", key, req); this.logDebug(`Request allowed for key: ${key}`); } else { this.stats.rejections++; this.callHook("onLimitReached", key, req); this.logDebug(`Request rejected for key: ${key}`); } if (this.config.dryRun) { this.logDebug( `Dry run mode - would ${allowed ? "allow" : "reject"} request for key: ${key}` ); return true; } if (this.config.silent) { this.logDebug( `Silent mode - ${allowed ? "allowing" : "rejecting"} request for key: ${key}` ); return allowed; } return allowed; } catch (error) { this.callHook("onError", error); this.logDebug(`Error checking rate limit for key: ${key}`, error); throw error; } } async reset(key) { this.logDebug(`Resetting rate limit for key: ${key}`); if (this.strategy.reset) { await this.strategy.reset(key); this.callHook("onReset", key); } } async getState(key) { if (this.strategy.getState) { return this.strategy.getState(key); } return null; } getStats() { return { ...this.stats, activeKeys: this.getActiveKeysCount() }; } getActiveKeysCount() { return 0; } }; // src/utils/keyGenerator.ts function createKeyGenerator(options = {}) { const { keyType = "ip", headerName, customKeyGenerator } = options; if (keyType === "custom" && typeof customKeyGenerator === "function") { return customKeyGenerator; } switch (keyType) { case "ip": return (req) => req.ip || req.connection?.remoteAddress || "__unknown_ip__"; case "user-agent": return (req) => req.headers?.["user-agent"] || "__unknown_ua__"; case "path": return (req) => req.path || req.url || "__unknown_path__"; default: if (keyType.startsWith("header:")) { const header = keyType.split(":")[1] || headerName; return (req) => req.headers?.[header?.toLowerCase()] || `__unknown_header_${header}__`; } return (req) => req.ip || "__unknown__"; } } // src/middleware/express.ts function expressLimiter(config) { const limiter = new RateLimiter(config); const keyFn = createKeyGenerator({ keyType: config.customKeyGenerator ? "custom" : config.keyType, headerName: config.headerName, customKeyGenerator: config.customKeyGenerator }); return async function limiterMiddleware(req, res, next) { try { const key = keyFn(req); const allowed = await limiter.isAllowed(key, req); if (!allowed) { return res.status(429).json({ error: "Too many requests. Please try again later." }); } next(); } catch (err) { console.error("Rate limiter error:", err); res.status(500).json({ error: "Internal rate limiter error" }); } }; } // src/middleware/fastify.ts function fastifyLimiter(config) { const limiter = new RateLimiter(config); const keyFn = createKeyGenerator({ keyType: config.customKeyGenerator ? "custom" : config.keyType, headerName: config.headerName, customKeyGenerator: config.customKeyGenerator }); return async function(req, reply) { const key = keyFn(req); const allowed = await limiter.isAllowed(key, req); if (!allowed) { reply.status(429).send({ message: "Too Many Requests" }); } }; } var RATE_LIMIT_METADATA_KEY = "rate_limit_config"; function RateLimit(config) { return function(target, propertyKey, descriptor) { Reflect.defineMetadata(RATE_LIMIT_METADATA_KEY, config, descriptor.value); return descriptor; }; } exports.NestLimiterGuard = class NestLimiterGuard { constructor(config) { this.config = config; this.limiter = new RateLimiter(config); this.keyFn = createKeyGenerator({ keyType: config.customKeyGenerator ? "custom" : config.keyType, headerName: config.headerName, customKeyGenerator: config.customKeyGenerator }); } async canActivate(context) { const ctx = context.switchToHttp(); const request = ctx.getRequest(); const handler = context.getHandler(); const classRef = context.getClass(); const routeConfig = Reflect.getMetadata(RATE_LIMIT_METADATA_KEY, handler) || Reflect.getMetadata(RATE_LIMIT_METADATA_KEY, classRef); let limiter = this.limiter; let keyFn = this.keyFn; if (routeConfig) { limiter = new RateLimiter(routeConfig); keyFn = createKeyGenerator({ keyType: routeConfig.keyType, headerName: routeConfig.headerName, customKeyGenerator: routeConfig.customKeyGenerator }); } const key = keyFn(request); const allowed = await limiter.isAllowed(key, request); if (!allowed) { throw new common.UnauthorizedException("Too Many Requests"); } return true; } }; exports.NestLimiterGuard = __decorateClass([ common.Injectable() ], exports.NestLimiterGuard); exports.RateLimitModule = class RateLimitModule { }; exports.RateLimitModule = __decorateClass([ common.Global(), common.Module({ providers: [exports.NestLimiterGuard], exports: [exports.NestLimiterGuard] }) ], exports.RateLimitModule); // src/middleware/handler.ts function universalLimiter(config) { const limiter = new RateLimiter(config); const keyFn = createKeyGenerator({ keyType: config.customKeyGenerator ? "custom" : config.keyType, headerName: config.headerName, customKeyGenerator: config.customKeyGenerator }); return async function(req, res, next) { const key = keyFn(req); const allowed = await limiter.isAllowed(key, req); if (!allowed) { if (res?.status && res?.send) { return res.status(429).send({ message: "Too Many Requests" }); } else if (res?.code && res?.send) { return res.code(429).send({ message: "Too Many Requests" }); } else { throw new Error("Rate limit exceeded"); } } next(); }; } exports.FixedWindowStrategy = FixedWindowStrategy; exports.RateLimit = RateLimit; exports.RateLimiter = RateLimiter; exports.SlidingWindowStrategy = SlidingWindowStrategy; exports.TokenBucketStrategy = TokenBucketStrategy; exports.applyDefaults = applyDefaults; exports.expressLimiter = expressLimiter; exports.fastifyLimiter = fastifyLimiter; exports.getMetrics = getMetrics; exports.isAllowedMemory = isAllowedMemory; exports.resetAll = resetAll; exports.resetKey = resetKey; exports.throwIfInvalid = throwIfInvalid; exports.universalLimiter = universalLimiter; exports.validateConfig = validateConfig; //# sourceMappingURL=index.js.map //# sourceMappingURL=index.js.map