diff --git a/src/@types/rateLimit.d.ts b/src/@types/rateLimit.d.ts index e62dd5a..9f5d873 100644 --- a/src/@types/rateLimit.d.ts +++ b/src/@types/rateLimit.d.ts @@ -47,7 +47,19 @@ export interface TokenBucketOptions { refillRate: number; } +/** + * @type {number} windowSize - Size of each fixed window and the rolling window + * @type {number} capacity - Number of tokens a window can hold + */ +export interface SlidingWindowCounterOptions { + windowSize: number; + capacity: number; +} + // TODO: This will be a union type where we can specify Option types for other Rate Limiters // Record represents the empty object for alogorithms that don't require settings // and might be able to be removed in the future. -export type RateLimiterOptions = TokenBucketOptions | Record; +export type RateLimiterOptions = + | TokenBucketOptions + | SlidingWindowCounterOptions + | Record; diff --git a/src/middleware/rateLimiterSetup.ts b/src/middleware/rateLimiterSetup.ts index 734bc5f..d7a3fb3 100644 --- a/src/middleware/rateLimiterSetup.ts +++ b/src/middleware/rateLimiterSetup.ts @@ -1,5 +1,6 @@ import Redis from 'ioredis'; -import { RateLimiterOptions, RateLimiterSelection } from '../@types/rateLimit'; +import { RateLimiterOptions, RateLimiterSelection, TokenBucketOptions } from '../@types/rateLimit'; +import SlidingWindowCounter from '../rateLimiters/slidingWindowCounter'; import TokenBucket from '../rateLimiters/tokenBucket'; /** @@ -19,6 +20,8 @@ export default function setupRateLimiter( switch (selection) { case 'TOKEN_BUCKET': // todo validate options + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore return new TokenBucket(options.bucketSize, options.refillRate, client); break; case 'LEAKY_BUCKET': @@ -31,7 +34,9 @@ export default function setupRateLimiter( throw new Error('Sliding Window Log has not be implemented.'); break; case 'SLIDING_WINDOW_COUNTER': - throw new Error('Sliding Window Counter algonithm has not be implemented.'); + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + return new SlidingWindowCounter(options.windowSize, options.capacity, client); break; default: // typescript should never let us invoke this function with anything other than the options above diff --git a/src/rateLimiters/slidingWindowCounter.ts b/src/rateLimiters/slidingWindowCounter.ts index ff24e81..a5c0d35 100644 --- a/src/rateLimiters/slidingWindowCounter.ts +++ b/src/rateLimiters/slidingWindowCounter.ts @@ -9,12 +9,12 @@ import { RateLimiter, RateLimiterResponse, RedisWindow } from '../@types/rateLim * takeup in each. * * Whenever a user makes a request the following steps are performed: - * 1. Fixed minute windows are defined along with redis caches if previously undefined. - * 2. Rolling minute windows are defined or updated based on the timestamp of the new request. + * 1. Fixed windows are defined along with redis caches if previously undefined. + * 2. Rolling windows are defined or updated based on the timestamp of the new request. * 3. Counter of the current fixed window is updated with the new request's token usage. * 4. If a new minute interval is reached, the averaging formula is run to prevent fixed window's flaw * of flooded requests around window borders - * (ex. 10 token capacity: 1m59s 10 reqs 2m2s 10 reqs) + * (ex. 1m windows, 10 token capacity: 1m59s 10 reqs 2m2s 10 reqs) */ class SlidingWindowCounter implements RateLimiter { private windowSize: number; @@ -24,10 +24,10 @@ class SlidingWindowCounter implements RateLimiter { private client: Redis; /** - * Create a new instance of a TokenBucket rate limiter that can be connected to any database store - * @param windowSize - size of each window in milliseconds (fixed and rolling) - * @param capacity - max capacity of tokens allowed per fixed window - * @param client - redis client where rate limiter will cache information + * Create a new instance of a SlidingWindowCounter rate limiter that can be connected to any database store + * @param windowSize size of each window in milliseconds (fixed and rolling) + * @param capacity max capacity of tokens allowed per fixed window + * @param client redis client where rate limiter will cache information */ constructor(windowSize: number, capacity: number, client: Redis) { this.windowSize = windowSize; @@ -38,12 +38,13 @@ class SlidingWindowCounter implements RateLimiter { } /** - * @function processRequest - current timestamp and number of tokens required for - * the request to go through are passed in. We first check if a window exists in the redis - * cache. + * @function processRequest - Sliding window counter algorithm to allow or block + * based on the depth/complexity (in amount of tokens) of incoming requests. * - * If not, then fixedWindowStart is set as the current timestamp, and currentTokens - * is checked against capacity. If we have enough capacity for the request, we return + * First, checks if a window exists in the redis cache. + * + * If not, then `fixedWindowStart` is set as the current timestamp, and `currentTokens` + * is checked against `capacity`. If enough room exists for the request, returns * success as true and tokens as how many tokens remain in the current fixed window. * * If a window does exist in the cache, we first check if the timestamp is greater than @@ -66,6 +67,8 @@ class SlidingWindowCounter implements RateLimiter { * @param {number} timestamp - time the request was recieved * @param {number} [tokens=1] - complexity of the query for throttling requests * @return {*} {Promise} + * RateLimiterResponse: {success: boolean, tokens: number} + * (tokens represents the remaining available capacity of the window) * @memberof SlidingWindowCounter */ async processRequest( @@ -79,31 +82,90 @@ class SlidingWindowCounter implements RateLimiter { // attempt to get the value for the uuid from the redis cache const windowJSON = await this.client.get(uuid); - // // if the response is null, we need to create a window for the user - // if (windowJSON === null) { - // // rolling window is 1 minute long - // const rollingWindowEnd = timestamp + 60000; - - // // grabs the actual minute from the timestamp to create fixed window - // const fixedWindowStart = timestamp - (timestamp % 10000); - // const fixedWindowEnd = fixedWindowStart + 60000; - - // const newUserWindow: RedisWindow = { - // // conditionally set tokens depending on how many are requested compared to the capacity - // tokens: tokens > this.capacity ? this.capacity : this.capacity - tokens, - // timestamp, - // }; - - // // reject the request, not enough tokens could even be in the bucket - // if (tokens > this.capacity) { - // await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow)); - // return { success: false, tokens: this.capacity }; - // } - // await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow)); - // return { success: true, tokens: newUserWindow.tokens }; - // } - - return { success: true, tokens: 0 }; + // if the response is null, we need to create a window for the user + if (windowJSON === null) { + const newUserWindow: RedisWindow = { + // current and previous tokens represent how many tokens are in each window + currentTokens: tokens <= this.capacity ? tokens : 0, + previousTokens: 0, + fixedWindowStart: timestamp, + }; + + if (tokens <= this.capacity) { + await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow)); + return { success: true, tokens: this.capacity - newUserWindow.currentTokens }; + } + + await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow)); + // tokens property represents how much capacity remains + return { success: false, tokens: this.capacity }; + } + + // if the cache is populated + + const window: RedisWindow = await JSON.parse(windowJSON); + + const updatedUserWindow: RedisWindow = { + currentTokens: window.currentTokens, + previousTokens: window.previousTokens, + fixedWindowStart: window.fixedWindowStart, + }; + + // if request time is in a new window + if (window.fixedWindowStart && timestamp >= window.fixedWindowStart + this.windowSize) { + // if more than one window was skipped + if (timestamp >= window.fixedWindowStart + this.windowSize * 2) { + // if one or more windows was skipped, reset new window to be at current timestamp + updatedUserWindow.previousTokens = 0; + updatedUserWindow.currentTokens = 0; + updatedUserWindow.fixedWindowStart = timestamp; + } else { + updatedUserWindow.previousTokens = updatedUserWindow.currentTokens; + updatedUserWindow.currentTokens = 0; + updatedUserWindow.fixedWindowStart = window.fixedWindowStart + this.windowSize; + } + } + + // assigned to avoid TS error, this var will never be used as 0 + // var is declared here so that below can be inside a conditional for efficiency's sake + let rollingWindowProportion = 0; + let previousRollingTokens = 0; + + if (updatedUserWindow.fixedWindowStart && updatedUserWindow.previousTokens) { + // proportion of rolling window present in previous window + rollingWindowProportion = + (this.windowSize - (timestamp - updatedUserWindow.fixedWindowStart)) / + this.windowSize; + + // remove unecessary decimals, 0.xx is enough + // rollingWindowProportion -= rollingWindowProportion % 0.01; + + // # of tokens present in rolling & previous window + previousRollingTokens = Math.floor( + updatedUserWindow.previousTokens * rollingWindowProportion + ); + } + + // # of tokens present in rolling and/or current window + // if previous tokens is null, previousRollingTokens will be 0 + const rollingTokens = updatedUserWindow.currentTokens + previousRollingTokens; + + // if request is allowed + if (tokens + rollingTokens <= this.capacity) { + updatedUserWindow.currentTokens += tokens; + await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow)); + return { + success: true, + tokens: this.capacity - (updatedUserWindow.currentTokens + previousRollingTokens), + }; + } + + // if request is blocked + await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow)); + return { + success: false, + tokens: this.capacity - (updatedUserWindow.currentTokens + previousRollingTokens), + }; } /** diff --git a/test/rateLimiters/slidingWindowCounter.test.ts b/test/rateLimiters/slidingWindowCounter.test.ts index ceba055..949e31e 100644 --- a/test/rateLimiters/slidingWindowCounter.test.ts +++ b/test/rateLimiters/slidingWindowCounter.test.ts @@ -35,7 +35,7 @@ async function setTokenCountInClient( await redisClient.set(uuid, JSON.stringify(value)); } -xdescribe('Test SlidingWindowCounter Rate Limiter', () => { +describe('Test SlidingWindowCounter Rate Limiter', () => { beforeEach(async () => { // init a mock redis cache client = new RedisMock(); @@ -249,7 +249,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { const result = await limiter.processRequest( user4, timestamp + WINDOW_SIZE * 1.99, - 4 + 10 ); expect(result.tokens).toBe(0); expect(result.success).toBe(true); @@ -257,7 +257,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { // currentTokens (in current fixed window): 4 // previousTokens (in previous fixed window): 8 const count1 = await getWindowFromClient(client, user4); - expect(count1.currentTokens).toBe(4); + expect(count1.currentTokens).toBe(10); expect(count1.previousTokens).toBe(8); }); }); @@ -304,7 +304,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { // 3 + 8 * 1 = 11, above capacity (request should be blocked) const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE, 3); - expect(result.tokens).toBe(10); + expect(result.tokens).toBe(2); expect(result.success).toBe(false); // currentTokens (in current fixed window): 0 @@ -332,7 +332,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { timestamp + WINDOW_SIZE * 1.25, 5 ); - expect(result.tokens).toBe(10); + expect(result.tokens).toBe(4); expect(result.success).toBe(false); // currentTokens (in current fixed window): 0 @@ -358,7 +358,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { // 7 + 8 * .5 = 11, over capacity (request should be blocked) const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE * 1.5, 7); - expect(result.tokens).toBe(10); + expect(result.tokens).toBe(6); expect(result.success).toBe(false); // currentTokens (in current fixed window): 0 @@ -383,7 +383,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { // 9 + 8 * .25 = 11, over capacity (request should be blocked) const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE * 1.75, 9); - expect(result.tokens).toBe(10); + expect(result.tokens).toBe(8); expect(result.success).toBe(false); // currentTokens (in current fixed window): 0 @@ -407,7 +407,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { // 11 + 8 * .01 = 11, above capacity (request should be blocked) const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE, 11); - expect(result.tokens).toBe(10); + expect(result.tokens).toBe(2); expect(result.success).toBe(false); // currentTokens (in current fixed window): 0 @@ -465,7 +465,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => { await ( await limiter.processRequest(user1, timestamp + WINDOW_SIZE, 4) ).tokens - ).toBe(2); + ).toBe(1); // currentTokens (in current fixed window): 0 // previousTokens (in previous fixed window): 8 const count = await getWindowFromClient(client, user1);