diff --git a/lib/api/apiUtils/rateLimit/cache.js b/lib/api/apiUtils/rateLimit/cache.js index a9bef0c1cd..a73147b580 100644 --- a/lib/api/apiUtils/rateLimit/cache.js +++ b/lib/api/apiUtils/rateLimit/cache.js @@ -1,63 +1,66 @@ const configCache = new Map(); -// Load tracking for adaptive burst capacity -// Map> - rolling 1-second window -const requestTimestamps = new Map(); +const namespace = { + bucket: 'bkt', +}; -function setCachedConfig(key, limitConfig, ttl) { +function cacheSet(cache, key, value, ttl) { const expiry = Date.now() + ttl; - configCache.set(key, { expiry, config: limitConfig }); + cache.set(key, { expiry, value }); } -function getCachedConfig(key) { - const value = configCache.get(key); - if (value === undefined) { +function cacheGet(cache, key) { + const cachedValue = cache.get(key); + if (cachedValue === undefined) { return undefined; } - const { expiry, config } = value; + const { expiry, value } = cachedValue; if (expiry <= Date.now()) { - configCache.delete(key); + cache.delete(key); return undefined; } - return config; + return value; +} + +function cacheDelete(cache, key) { + cache.delete(key); } -function expireCachedConfigs(now) { +function cacheExpire(cache) { + const now = Date.now(); + const toRemove = []; - for (const [key, { expiry }] of configCache.entries()) { + for (const [key, { expiry }] of cache.entries()) { if (expiry <= now) { toRemove.push(key); } } for (const key of toRemove) { - configCache.delete(key); + cache.delete(key); } return toRemove.length; } -/** - * Invalidate cached config for a specific bucket - * - * @param {string} bucketName - Bucket name - * @returns {boolean} True if entry was found and removed - */ -function invalidateCachedConfig(bucketName) { - const cacheKey = `bucket:${bucketName}`; - return configCache.delete(cacheKey); +function formatKeyDecorator(fn) { + return (resourceClass, resourceId, ...args) => fn(`${resourceClass}:${resourceId}`, ...args); } +const getCachedConfig = formatKeyDecorator(cacheGet.bind(null, configCache)); +const setCachedConfig = formatKeyDecorator(cacheSet.bind(null, configCache)); +const deleteCachedConfig = formatKeyDecorator(cacheDelete.bind(null, configCache)); +const expireCachedConfigs = cacheExpire.bind(null, configCache); + module.exports = { + namespace, setCachedConfig, getCachedConfig, expireCachedConfigs, - invalidateCachedConfig, - + deleteCachedConfig, // Do not access directly // Used only for tests configCache, - requestTimestamps, }; diff --git a/lib/api/apiUtils/rateLimit/helpers.js b/lib/api/apiUtils/rateLimit/helpers.js index 80c039b097..8bc8d7b83c 100644 --- a/lib/api/apiUtils/rateLimit/helpers.js +++ b/lib/api/apiUtils/rateLimit/helpers.js @@ -1,24 +1,12 @@ const { config } = require('../../../Config'); const cache = require('./cache'); -const constants = require('../../../../constants'); const { getTokenBucket } = require('./tokenBucket'); const { policies: { actionMaps: { actionMapBucketRateLimit } } } = require('arsenal'); const rateLimitApiActions = Object.keys(actionMapBucketRateLimit); /** - * Get rate limit configuration from cache only (no metadata fetch) - * - * @param {string} bucketName - Bucket name - * @returns {object|null|undefined} Rate limit config, null (no limit), or undefined (not cached) - */ -function getRateLimitFromCache(bucketName) { - const cacheKey = `bucket:${bucketName}`; - return cache.getCachedConfig(cacheKey); -} - -/** - * Extract rate limit configuration from bucket metadata and cache it + * Extract rate limit configuration from bucket metadata or global rate limit configuration. * * Resolves in priority order: * 1. Per-bucket configuration (from bucket metadata) @@ -30,24 +18,21 @@ function getRateLimitFromCache(bucketName) { * @param {object} log - Logger instance * @returns {object|null} Rate limit config or null if no limit */ -function extractAndCacheRateLimitConfig(bucket, bucketName, log) { - const cacheKey = `bucket:${bucketName}`; - const cacheTTL = config.rateLimiting.bucket?.configCacheTTL || - constants.rateLimitDefaultConfigCacheTTL; - +function extractBucketRateLimitConfig(bucket, bucketName, log) { // Try per-bucket config first const bucketConfig = bucket.getRateLimitConfiguration(); if (bucketConfig) { const data = bucketConfig.getData(); const limitConfig = { limit: data.RequestsPerSecond.Limit, + burstCapacity: config.rateLimiting.bucket.defaultBurstCapacity * 1000, source: 'bucket', }; - cache.setCachedConfig(cacheKey, limitConfig, cacheTTL); log.debug('Extracted per-bucket rate limit config', { bucketName, limit: limitConfig.limit, + burstCapacity: config.rateLimiting.bucket.defaultBurstCapacity * 1000, }); return limitConfig; @@ -58,10 +43,10 @@ function extractAndCacheRateLimitConfig(bucket, bucketName, log) { if (globalLimit !== undefined && globalLimit > 0) { const limitConfig = { limit: globalLimit, + burstCapacity: config.rateLimiting.bucket.defaultBurstCapacity * 1000, source: 'global', }; - cache.setCachedConfig(cacheKey, limitConfig, cacheTTL); log.debug('Using global default rate limit config', { bucketName, limit: limitConfig.limit, @@ -71,58 +56,86 @@ function extractAndCacheRateLimitConfig(bucket, bucketName, log) { } // No rate limiting configured - cache null to avoid repeated lookups - cache.setCachedConfig(cacheKey, null, cacheTTL); log.trace('No rate limit configured for bucket', { bucketName }); - return null; } -/** - * Check rate limit with pre-resolved configuration using token reservation - * - * Uses token bucket: Workers maintain local tokens granted by Redis. - * Token consumption is pure in-memory (fast). Refills happen async in background. - * - * @param {string} bucketName - Bucket name - * @param {object|null} limitConfig - Pre-resolved rate limit config - * @param {object} log - Logger instance - * @param {function} callback - Callback(err, rateLimited boolean) - * @returns {undefined} - */ -function checkRateLimitWithConfig(bucketName, limitConfig, log, callback) { - // No rate limiting configured - if (!limitConfig || limitConfig.limit === 0) { - return callback(null, false); +function extractRateLimitConfigFromRequest(request) { + const applyRateLimit = config.rateLimiting?.enabled + && !rateLimitApiActions.includes(request.apiMethod) // Don't limit any rate limit admin actions + && !request.isInternalServiceRequest; // Don't limit any calls from internal services + + if (!applyRateLimit) { + return { needsCheck: false }; + } + + const limitConfigs = {}; + let needsCheck = false; + + if (request.accountLimits) { + limitConfigs.account = { + ...request.accountLimits, + source: 'account', + }; + needsCheck = true; + } + + if (request.bucketName) { + const cachedConfig = cache.getCachedConfig(cache.namespace.bucket, request.bucketName); + if (cachedConfig) { + limitConfigs.bucket = cachedConfig; + needsCheck = true; + } + + if (!request.accountLimits) { + const cachedOwner = cache.getCachedBucketOwner(request.bucketName); + if (cachedOwner !== undefined) { + const cachedConfig = cache.getCachedConfig(cache.namespace.account, cachedOwner); + if (cachedConfig) { + limitConfigs.account = cachedConfig; + limitConfigs.bucketOwner = cachedOwner; + needsCheck = true; + } + } + } } - // Get or create token bucket for this bucket - const tokenBucket = getTokenBucket(bucketName, limitConfig, log); + return { needsCheck, limitConfigs }; +} + +function checkRateLimitsForRequest(checks, log) { + const buckets = []; + for (const check of checks) { + const bucket = getTokenBucket(check.resourceClass, check.resourceId, check.measure, check.config, log); + if (!bucket.hasCapacity()) { + log.debug('Rate limit check: denied (no tokens available)', { + resourceClass: bucket.resourceClass, + resourceId: bucket.resourceId, + measure: bucket.measure, + limit: bucket.limitConfig.limit, + source: bucket.limitConfig.source, + }); - // Try to consume a token (in-memory, no Redis) - const allowed = tokenBucket.tryConsume(); + return { allowed: false, rateLimitSource: `${check.resourceClass}:${check.source}`}; + } + + buckets.push(bucket); - if (allowed) { log.trace('Rate limit check: allowed (token consumed)', { - bucketName, - tokensRemaining: tokenBucket.tokens, - }); - } else { - log.debug('Rate limit check: denied (no tokens available)', { - bucketName, - limit: limitConfig.limit, - source: limitConfig.source, + resourceClass: bucket.resourceClass, + resourceId: bucket.resourceId, + measure: bucket.measure, + source: bucket.limitConfig.source, }); } - // Return inverse: callback expects "rateLimited" boolean - // allowed=true → rateLimited=false - // allowed=false → rateLimited=true - return callback(null, !allowed); + buckets.forEach(bucket => bucket.tryConsume()); + return { allowed: true }; } module.exports = { - getRateLimitFromCache, - extractAndCacheRateLimitConfig, - checkRateLimitWithConfig, rateLimitApiActions, + extractBucketRateLimitConfig, + extractRateLimitConfigFromRequest, + checkRateLimitsForRequest, }; diff --git a/lib/api/bucketDeleteRateLimit.js b/lib/api/bucketDeleteRateLimit.js index 81bfde32eb..f500463746 100644 --- a/lib/api/bucketDeleteRateLimit.js +++ b/lib/api/bucketDeleteRateLimit.js @@ -4,7 +4,7 @@ const metadata = require('../metadata/wrapper'); const { standardMetadataValidateBucket } = require('../metadata/metadataUtils'); const collectCorsHeaders = require('../utilities/collectCorsHeaders'); const { isRateLimitServiceUser } = require('./apiUtils/authorization/serviceUser'); -const { invalidateCachedConfig } = require('./apiUtils/rateLimit/cache'); +const cache = require('./apiUtils/rateLimit/cache'); const { removeTokenBucket } = require('./apiUtils/rateLimit/tokenBucket'); /** @@ -52,7 +52,7 @@ function bucketDeleteRateLimit(authInfo, request, log, callback) { return callback(err, corsHeaders); } // Invalidate cache and remove token bucket - invalidateCachedConfig(bucketName); + cache.deleteCachedConfig(cache.namespace.bucket, bucketName); removeTokenBucket(bucketName); log.debug('invalidated rate limit cache and token bucket for bucket', { bucketName }); // TODO: implement Utapi metric support diff --git a/lib/api/bucketPutRateLimit.js b/lib/api/bucketPutRateLimit.js index 116c8940d9..d895678775 100644 --- a/lib/api/bucketPutRateLimit.js +++ b/lib/api/bucketPutRateLimit.js @@ -6,7 +6,7 @@ const { config } = require('../Config'); const metadata = require('../metadata/wrapper'); const { standardMetadataValidateBucket } = require('../metadata/metadataUtils'); const { isRateLimitServiceUser } = require('./apiUtils/authorization/serviceUser'); -const { invalidateCachedConfig } = require('./apiUtils/rateLimit/cache'); +const cache = require('./apiUtils/rateLimit/cache'); function parseRequestBody(requestBody, callback) { try { @@ -94,7 +94,7 @@ function bucketPutRateLimit(authInfo, request, log, callback) { return callback(err, corsHeaders); } // Invalidate cache so new limit takes effect immediately - invalidateCachedConfig(bucketName); + cache.deleteCachedConfig(cache.namespace.bucket, bucketName); log.debug('invalidated rate limit cache for bucket', { bucketName }); // TODO: implement Utapi metric support return callback(null, corsHeaders); diff --git a/tests/unit/api/apiUtils/rateLimit/cache.js b/tests/unit/api/apiUtils/rateLimit/cache.js index dc8b15552e..cb3e705fa9 100644 --- a/tests/unit/api/apiUtils/rateLimit/cache.js +++ b/tests/unit/api/apiUtils/rateLimit/cache.js @@ -3,11 +3,12 @@ const sinon = require('sinon'); const constants = require('../../../../../constants'); const { + namespace, configCache, getCachedConfig, setCachedConfig, expireCachedConfigs, - invalidateCachedConfig, + deleteCachedConfig, } = require('../../../../../lib/api/apiUtils/rateLimit/cache'); describe('test limit config cache storage', () => { @@ -22,66 +23,70 @@ describe('test limit config cache storage', () => { clock.restore(); }); + beforeEach(() => { + configCache.clear(); + }); + it('should add config to cache', () => { - setCachedConfig('foo', 10, constants.rateLimitDefaultConfigCacheTTL); + setCachedConfig(namespace.bucket, 'foo', 10, constants.rateLimitDefaultConfigCacheTTL); assert.deepStrictEqual( - configCache.get('foo'), + configCache.get(`${namespace.bucket}:foo`), { expiry: now + constants.rateLimitDefaultConfigCacheTTL, - config: 10, + value: 10, } ); }); it('should get a non expired config', () => { - setCachedConfig('foo', 10, constants.rateLimitDefaultConfigCacheTTL); - assert.strictEqual(getCachedConfig('foo'), 10); + setCachedConfig(namespace.bucket, 'foo', 10, constants.rateLimitDefaultConfigCacheTTL); + assert.strictEqual(getCachedConfig(namespace.bucket, 'foo'), 10); }); it('should return undefined and delete the key for an expired config', () => { - configCache.set('foo', { + configCache.set(`${namespace.bucket}:foo`, { expiry: now - 10000, - config: 10, + value: 10, }); - assert.strictEqual(getCachedConfig('foo'), undefined); + assert.strictEqual(getCachedConfig(namespace.bucket, 'foo'), undefined); }); - it('should expire configs less than or equal to the given timestamp', () => { + it('should expire configs less than or equal to current time', () => { configCache.set('past', { expiry: now - 10000, - config: 10, + value: 10, }); configCache.set('present', { expiry: now, - config: 10, + value: 10, }); configCache.set('future', { expiry: now + 10000, - config: 10, + value: 10, }); - expireCachedConfigs(now); + // expireCachedConfigs uses Date.now() internally; fake clock is set to `now` + expireCachedConfigs(); assert.strictEqual(configCache.get('past'), undefined); assert.strictEqual(configCache.get('present'), undefined); assert.deepStrictEqual(configCache.get('future'), { expiry: now + 10000, - config: 10, + value: 10, }); }); - it('should invalidate cached config for a specific bucket', () => { - setCachedConfig('bucket:my-bucket', { limit: 100 }, constants.rateLimitDefaultConfigCacheTTL); - setCachedConfig('bucket:other-bucket', { limit: 200 }, constants.rateLimitDefaultConfigCacheTTL); + it('should delete cached config for a specific resource', () => { + setCachedConfig(namespace.bucket, 'my-bucket', { limit: 100 }, constants.rateLimitDefaultConfigCacheTTL); + setCachedConfig(namespace.bucket, 'other-bucket', { limit: 200 }, constants.rateLimitDefaultConfigCacheTTL); - const result = invalidateCachedConfig('my-bucket'); + deleteCachedConfig(namespace.bucket, 'my-bucket'); - assert.strictEqual(result, true); - assert.strictEqual(getCachedConfig('bucket:my-bucket'), undefined); - assert.deepStrictEqual(getCachedConfig('bucket:other-bucket'), { limit: 200 }); + assert.strictEqual(getCachedConfig(namespace.bucket, 'my-bucket'), undefined); + assert.deepStrictEqual(getCachedConfig(namespace.bucket, 'other-bucket'), { limit: 200 }); }); - it('should return false when invalidating non-existent bucket', () => { - const result = invalidateCachedConfig('non-existent-bucket'); - - assert.strictEqual(result, false); + it('should be a no-op when deleting a non-existent key', () => { + assert.doesNotThrow(() => { + deleteCachedConfig(namespace.bucket, 'non-existent-bucket'); + }); }); }); diff --git a/tests/unit/api/apiUtils/rateLimit/helpers.js b/tests/unit/api/apiUtils/rateLimit/helpers.js index b1b14d4e24..39c49060b3 100644 --- a/tests/unit/api/apiUtils/rateLimit/helpers.js +++ b/tests/unit/api/apiUtils/rateLimit/helpers.js @@ -28,36 +28,7 @@ describe('Rate limit helpers', () => { sandbox.restore(); }); - describe('getRateLimitFromCache', () => { - it('should return cached config on cache hit', () => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; - cache.setCachedConfig(`bucket:${bucketName}`, limitConfig, 60000); - - const result = helpers.getRateLimitFromCache(bucketName); - - assert.strictEqual(result, limitConfig); - }); - - it('should return undefined on cache miss', () => { - const bucketName = 'test-bucket'; - - const result = helpers.getRateLimitFromCache(bucketName); - - assert.strictEqual(result, undefined); - }); - - it('should return null when null is cached', () => { - const bucketName = 'test-bucket'; - cache.setCachedConfig(`bucket:${bucketName}`, null, 60000); - - const result = helpers.getRateLimitFromCache(bucketName); - - assert.strictEqual(result, null); - }); - }); - - describe('extractAndCacheRateLimitConfig', () => { + describe('extractBucketRateLimitConfig', () => { let configStub; beforeEach(() => { @@ -65,11 +36,12 @@ describe('Rate limit helpers', () => { enabled: true, bucket: { configCacheTTL: 30000, + defaultBurstCapacity: 1, }, }); }); - it('should extract and cache per-bucket config', () => { + it('should extract per-bucket config', () => { const bucketName = 'test-bucket'; const mockBucket = { getRateLimitConfiguration: () => ({ @@ -79,15 +51,12 @@ describe('Rate limit helpers', () => { }), }; - const result = helpers.extractAndCacheRateLimitConfig(mockBucket, bucketName, mockLog); + const result = helpers.extractBucketRateLimitConfig(mockBucket, bucketName, mockLog); - assert.deepStrictEqual(result, { limit: 200, source: 'bucket' }); - // Verify it was cached - const cached = cache.getCachedConfig(`bucket:${bucketName}`); - assert.deepStrictEqual(cached, { limit: 200, source: 'bucket' }); + assert.deepStrictEqual(result, { limit: 200, burstCapacity: 1000, source: 'bucket' }); }); - it('should fall back to global config when no bucket config', () => { + it('should fall back to global default config when no bucket config', () => { const bucketName = 'test-bucket'; const mockBucket = { getRateLimitConfiguration: () => null, @@ -97,16 +66,14 @@ describe('Rate limit helpers', () => { enabled: true, bucket: { defaultConfig: { limit: 100 }, + defaultBurstCapacity: 1, configCacheTTL: 30000, }, }); - const result = helpers.extractAndCacheRateLimitConfig(mockBucket, bucketName, mockLog); + const result = helpers.extractBucketRateLimitConfig(mockBucket, bucketName, mockLog); - assert.deepStrictEqual(result, { limit: 100, source: 'global' }); - // Verify it was cached - const cached = cache.getCachedConfig(`bucket:${bucketName}`); - assert.deepStrictEqual(cached, { limit: 100, source: 'global' }); + assert.deepStrictEqual(result, { limit: 100, burstCapacity: 1000, source: 'global' }); }); it('should return null when no config exists', () => { @@ -118,19 +85,17 @@ describe('Rate limit helpers', () => { configStub.value({ enabled: true, bucket: { + defaultBurstCapacity: 1, configCacheTTL: 30000, }, }); - const result = helpers.extractAndCacheRateLimitConfig(mockBucket, bucketName, mockLog); + const result = helpers.extractBucketRateLimitConfig(mockBucket, bucketName, mockLog); assert.strictEqual(result, null); - // Verify null was cached - const cached = cache.getCachedConfig(`bucket:${bucketName}`); - assert.strictEqual(cached, null); }); - it('should skip global default if limit is 0', () => { + it('should return null when global default limit is 0', () => { const bucketName = 'test-bucket'; const mockBucket = { getRateLimitConfiguration: () => null, @@ -140,16 +105,17 @@ describe('Rate limit helpers', () => { enabled: true, bucket: { defaultConfig: { limit: 0 }, + defaultBurstCapacity: 1, configCacheTTL: 30000, }, }); - const result = helpers.extractAndCacheRateLimitConfig(mockBucket, bucketName, mockLog); + const result = helpers.extractBucketRateLimitConfig(mockBucket, bucketName, mockLog); assert.strictEqual(result, null); }); - it('should use default TTL when not configured', () => { + it('should use default TTL when configCacheTTL is not set', () => { const bucketName = 'test-bucket'; const mockBucket = { getRateLimitConfiguration: () => ({ @@ -161,240 +127,194 @@ describe('Rate limit helpers', () => { configStub.value({ enabled: true, - bucket: {}, + bucket: { + defaultBurstCapacity: 1, + }, }); sandbox.stub(constants, 'rateLimitDefaultConfigCacheTTL').value(60000); - helpers.extractAndCacheRateLimitConfig(mockBucket, bucketName, mockLog); + const result = helpers.extractBucketRateLimitConfig(mockBucket, bucketName, mockLog); - // Verify it was cached with default TTL - const cached = cache.getCachedConfig(`bucket:${bucketName}`); - assert.deepStrictEqual(cached, { limit: 200, source: 'bucket' }); + assert.deepStrictEqual(result, { limit: 200, burstCapacity: 1000, source: 'bucket' }); }); }); - describe('checkRateLimitWithConfig', () => { - let configStub; - - beforeEach(() => { - configStub = sandbox.stub(config, 'rateLimiting').value({ + describe('checkRateLimitsForRequest', () => { + beforeEach(() => + sandbox.stub(config, 'rateLimiting').value({ enabled: true, nodes: 1, + tokenBucketBufferSize: 50, + tokenBucketRefillThreshold: 20, bucket: { defaultBurstCapacity: 1, }, - }); - sandbox.stub(config, 'clusters').value(1); - }); + }) + ); - it('should allow request when no limit configured', done => { - const bucketName = 'test-bucket'; - const limitConfig = null; - - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(err, null); - assert.strictEqual(rateLimited, false); - done(); - }); - }); + afterEach(() => sinon.restore()); - it('should allow request when limit is 0', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 0, source: 'global' }; + it('should allow request when checks array is empty', () => { + const result = helpers.checkRateLimitsForRequest([], mockLog); - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(err, null); - assert.strictEqual(rateLimited, false); - done(); - }); + assert.deepStrictEqual(result, { allowed: true }); }); - it('should allow request when bucket has capacity', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; + it('should allow request when bucket has capacity', () => { + const check = { + resourceClass: 'bkt', resourceId: 'test-bucket', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; // Pre-populate token bucket with tokens - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); + const bucket = tokenBucket.getTokenBucket('bkt', 'test-bucket', 'rps', check.config, mockLog); bucket.tokens = 50; - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(err, null); - assert.strictEqual(rateLimited, false); - // Verify token was consumed - assert.strictEqual(bucket.tokens, 49); - done(); - }); + const result = helpers.checkRateLimitsForRequest([check], mockLog); + + assert.deepStrictEqual(result, { allowed: true }); + // Verify token was consumed + assert.strictEqual(bucket.tokens, 49); }); - it('should deny request when bucket is full', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; + it('should deny request when bucket has no tokens', () => { + const check = { + resourceClass: 'bkt', resourceId: 'test-bucket', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); + const bucket = tokenBucket.getTokenBucket('bkt', 'test-bucket', 'rps', check.config, mockLog); // Explicitly set tokens to 0 to simulate exhausted quota bucket.tokens = 0; - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(err, null); - assert.strictEqual(rateLimited, true); - done(); - }); + const result = helpers.checkRateLimitsForRequest([check], mockLog); + + assert.strictEqual(result.allowed, false); + assert.strictEqual(result.rateLimitSource, 'bkt:bucket'); }); - it('should use configured burst capacity', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; + it('should not consume tokens when denied', () => { + const check = { + resourceClass: 'bkt', resourceId: 'test-bucket', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; - configStub.value({ - enabled: true, - nodes: 1, - bucket: { - defaultBurstCapacity: 2, - }, - }); + const bucket = tokenBucket.getTokenBucket('bkt', 'test-bucket', 'rps', check.config, mockLog); + bucket.tokens = 0; - // Pre-populate token bucket - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); - bucket.tokens = 50; + helpers.checkRateLimitsForRequest([check], mockLog); - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(err, null); - assert.strictEqual(rateLimited, false); - done(); - }); + assert.strictEqual(bucket.tokens, 0); }); - it('should use default burst capacity when not configured', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; - - configStub.value({ - enabled: true, - nodes: 1, - bucket: {}, - }); + it('should consume tokens from all buckets when all have capacity', () => { + const check1 = { + resourceClass: 'bkt', resourceId: 'bucket-1', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; + const check2 = { + resourceClass: 'acc', resourceId: 'account-1', measure: 'rps', + config: { limit: 200, burstCapacity: 1000, source: 'account' }, source: 'account', + }; - sandbox.stub(constants, 'rateLimitDefaultBurstCapacity').value(1); + const bucket1 = tokenBucket.getTokenBucket('bkt', 'bucket-1', 'rps', check1.config, mockLog); + const bucket2 = tokenBucket.getTokenBucket('acc', 'account-1', 'rps', check2.config, mockLog); + bucket1.tokens = 50; + bucket2.tokens = 50; - // Pre-populate token bucket - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); - bucket.tokens = 50; + const result = helpers.checkRateLimitsForRequest([check1, check2], mockLog); - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(err, null); - assert.strictEqual(rateLimited, false); - done(); - }); + assert.deepStrictEqual(result, { allowed: true }); + assert.strictEqual(bucket1.tokens, 49); + assert.strictEqual(bucket2.tokens, 49); }); - it('should calculate interval for distributed setup', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 1000, source: 'global' }; + it('should deny on first exhausted bucket and not consume other buckets', () => { + const check1 = { + resourceClass: 'bkt', resourceId: 'bucket-1', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; + const check2 = { + resourceClass: 'acc', resourceId: 'account-1', measure: 'rps', + config: { limit: 200, burstCapacity: 1000, source: 'account' }, source: 'account', + }; - configStub.value({ - enabled: true, - nodes: 10, - bucket: { - defaultBurstCapacity: 1, - }, - }); - sandbox.stub(config, 'clusters').value(5); + const bucket1 = tokenBucket.getTokenBucket('bkt', 'bucket-1', 'rps', check1.config, mockLog); + const bucket2 = tokenBucket.getTokenBucket('acc', 'account-1', 'rps', check2.config, mockLog); + bucket1.tokens = 0; // exhausted + bucket2.tokens = 50; - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, err => { - assert.strictEqual(err, null); - // Should work correctly with distributed calculation - done(); - }); + const result = helpers.checkRateLimitsForRequest([check1, check2], mockLog); + + assert.strictEqual(result.allowed, false); + // bucket2 tokens should be unchanged (not consumed when an earlier check fails) + assert.strictEqual(bucket2.tokens, 50); }); - it('should log debug info when using token bucket', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; + it('should log debug info when request is denied', () => { + const check = { + resourceClass: 'bkt', resourceId: 'test-bucket', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); + const bucket = tokenBucket.getTokenBucket('bkt', 'test-bucket', 'rps', check.config, mockLog); // Explicitly set tokens to 0 to trigger denial log bucket.tokens = 0; - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, () => { - // Check for token bucket denial log - const deniedCall = mockLog.debug.getCalls().find( - call => call.args[0] === 'Rate limit check: denied (no tokens available)' - ); - assert(deniedCall, 'Should have logged denied message'); - const logArgs = deniedCall.args[1]; - assert.strictEqual(logArgs.bucketName, bucketName); - assert.strictEqual(logArgs.limit, 100); - assert.strictEqual(logArgs.source, 'bucket'); - done(); - }); + helpers.checkRateLimitsForRequest([check], mockLog); + + const deniedCall = mockLog.debug.getCalls().find( + call => call.args[0] === 'Rate limit check: denied (no tokens available)' + ); + assert(deniedCall, 'Should have logged denied message'); + const logArgs = deniedCall.args[1]; + assert.strictEqual(logArgs.resourceClass, 'bkt'); + assert.strictEqual(logArgs.resourceId, 'test-bucket'); + assert.strictEqual(logArgs.limit, 100); + assert.strictEqual(logArgs.source, 'bucket'); }); - it('should log trace info when request allowed', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; + it('should log trace info when request is allowed', () => { + const check = { + resourceClass: 'bkt', resourceId: 'test-bucket', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; // Pre-populate token bucket - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); + const bucket = tokenBucket.getTokenBucket('bkt', 'test-bucket', 'rps', check.config, mockLog); bucket.tokens = 50; - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(rateLimited, false); - // Check for trace log (token consumed) - const allowedCall = mockLog.trace.getCalls().find( - call => call.args[0] === 'Rate limit check: allowed (token consumed)' - ); - assert(allowedCall, 'Should have logged allowed message'); - assert.strictEqual(allowedCall.args[1].bucketName, bucketName); - assert.strictEqual(allowedCall.args[1].tokensRemaining, 49); - done(); - }); - }); - - it('should log debug info when request denied (no tokens)', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; + const result = helpers.checkRateLimitsForRequest([check], mockLog); - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); - // Explicitly set tokens to 0 to simulate exhausted quota - bucket.tokens = 0; - - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(rateLimited, true); - // Find the "denied" log call - const deniedCall = mockLog.debug.getCalls().find( - call => call.args[0] === 'Rate limit check: denied (no tokens available)' - ); - assert(deniedCall, 'Should have logged denied message'); - assert.strictEqual(deniedCall.args[1].bucketName, bucketName); - assert.strictEqual(deniedCall.args[1].limit, 100); - assert.strictEqual(deniedCall.args[1].source, 'bucket'); - done(); - }); + assert.strictEqual(result.allowed, true); + const allowedCall = mockLog.trace.getCalls().find( + call => call.args[0] === 'Rate limit check: allowed (token consumed)' + ); + assert(allowedCall, 'Should have logged allowed message'); + assert.strictEqual(allowedCall.args[1].resourceClass, 'bkt'); + assert.strictEqual(allowedCall.args[1].resourceId, 'test-bucket'); }); - it('should handle multiple sequential requests correctly', done => { - const bucketName = 'test-bucket'; - const limitConfig = { limit: 100, source: 'bucket' }; + it('should handle multiple sequential requests correctly', () => { + const check = { + resourceClass: 'bkt', resourceId: 'test-bucket', measure: 'rps', + config: { limit: 100, burstCapacity: 1000, source: 'bucket' }, source: 'bucket', + }; // Pre-populate token bucket with multiple tokens - const bucket = tokenBucket.getTokenBucket(bucketName, limitConfig, mockLog); + const bucket = tokenBucket.getTokenBucket('bkt', 'test-bucket', 'rps', check.config, mockLog); bucket.tokens = 50; // First request should be allowed - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err, rateLimited) => { - assert.strictEqual(err, null); - assert.strictEqual(rateLimited, false); - assert.strictEqual(bucket.tokens, 49); - - // Second request should also be allowed (still has tokens) - helpers.checkRateLimitWithConfig(bucketName, limitConfig, mockLog, (err2, rateLimited2) => { - assert.strictEqual(err2, null); - assert.strictEqual(rateLimited2, false); - assert.strictEqual(bucket.tokens, 48); - done(); - }); - }); + const result1 = helpers.checkRateLimitsForRequest([check], mockLog); + assert.strictEqual(result1.allowed, true); + assert.strictEqual(bucket.tokens, 49); + + // Second request should also be allowed (still has tokens) + const result2 = helpers.checkRateLimitsForRequest([check], mockLog); + assert.strictEqual(result2.allowed, true); + assert.strictEqual(bucket.tokens, 48); }); }); });