working on test

Author: Yufei Wu

src/@types/rateLimit.d.ts

@@ -22,6 +22,12 @@ export interface RedisBucket {
     tokens: number;
     timestamp: number;
 }
+export interface RedisWindow {
+    currentTokens: number;
+    // null if limiter is currently on the initial fixed window
+    previousTokens?: number | null;
+    fixedWindowStart: number;
+}
 
 export type RateLimiterSelection =
     | 'TOKEN_BUCKET'

src/rateLimiters/fixedWindow.ts

@@ -0,0 +1,109 @@
+import Redis from 'ioredis';
+import { RateLimiter, RateLimiterResponse, RedisWindow } from '../@types/rateLimit';
+
+/**
+ * The FixedWindow instance of a RateLimiter limits requests based on a unique user ID and a fixed time window.
+ * Whenever a user makes a request the following steps are performed:
+ *  1. Define the time window with fixed amount of queries.
+ *  2. Update the timestamp of the last request.
+ *  3. Allow the request and decrease the allowed amount of requests if the user has enough at this time window.
+ *  4. Otherwise, disallow the request until the next time window opens.
+ */
+
+class FixedWindow implements RateLimiter {
+    private capacity: number;
+
+    private windowSize: number;
+
+    private client: Redis;
+
+    /**
+     * Create a new instance of a FixedWindow rate limiter that can be connected to any database store
+     * @param capacity max requests capacity in one time window
+     * @param windowSize rate at which the token bucket is refilled
+     * @param client redis client where rate limiter will cache information
+     */
+
+    constructor(capacity: number, windowSize: number, client: Redis) {
+        this.capacity = capacity;
+        this.windowSize = windowSize;
+        this.client = client;
+        if (windowSize <= 0 || capacity <= 0)
+            throw Error('FixedWindow windowSize and capacity must be positive');
+    }
+
+    /**
+     *               Fixed Window
+     *     _________________________________
+     *    |         *full capacity          |
+     *    |                                 |  move to next time window
+     *    |    token adds up until full     |       ---------->
+     *____._________________________________.____
+     *    |<--        window size        -->|
+     *current timestamp              next timestamp
+     *
+     *
+     * @param {string} uuid - unique identifer used to throttle requests
+     * @param {number} timestamp - time the request was recieved
+     * @param {number} [tokens=1]  - complexity of the query for throttling requests
+     * @return {*}  {Promise<RateLimiterResponse>}
+     * @memberof FixedWindow
+     */
+    async processRequest(
+        uuid: string,
+        timestamp: number,
+        tokens = 1
+    ): Promise<RateLimiterResponse> {
+        // set the expiry of key-value pairs in the cache to 24 hours
+        const keyExpiry = 86400000;
+
+        // attempt to get the value for the uuid from the redis cache
+        const windowJSON = await this.client.get(uuid);
+
+        if (windowJSON === null) {
+            const newUserWindow: RedisWindow = {
+                currentTokens: tokens <= this.capacity ? tokens : 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));
+        }
+
+        const window: RedisWindow = await JSON.parse(windowJSON as string);
+
+        const updatedUserWindow = this.updateTimeWindow(window, timestamp);
+        if (window.currentTokens > this.capacity) {
+            await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
+            return { success: false, tokens: window.currentTokens };
+        }
+        await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
+
+        return { success: false, tokens: updatedUserWindow.currentTokens };
+    }
+
+    /**
+     * Resets the rate limiter to the intial state by clearing the redis store.
+     */
+    public reset(): void {
+        this.client.flushall();
+    }
+
+    private updateTimeWindow = (window: RedisWindow, timestamp: number): RedisWindow => {
+        const updatedUserWindow: RedisWindow = {
+            currentTokens: window.currentTokens,
+            fixedWindowStart: window.fixedWindowStart,
+        };
+        if (timestamp > window.fixedWindowStart + this.windowSize) {
+            updatedUserWindow.fixedWindowStart = window.fixedWindowStart + this.windowSize;
+            updatedUserWindow.currentTokens = 0;
+        }
+        return updatedUserWindow;
+    };
+}
+
+export default FixedWindow;

test/rateLimiters/fixedWindow.test.ts

@@ -0,0 +1,57 @@
+import * as ioredis from 'ioredis';
+import { RedisWindow } from '../../src/@types/rateLimit';
+import FixedWindow from '../../src/rateLimiters/fixedWindow';
+
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const RedisMock = require('ioredis-mock');
+
+const CAPACITY = 10;
+const WINDOW_SIZE = 6000;
+
+let limiter: FixedWindow;
+let client: ioredis.Redis;
+let timestamp: number;
+const user1 = '1';
+const user2 = '2';
+const user3 = '3';
+const user4 = '4';
+
+async function getWindowFromClient(redisClient: ioredis.Redis, uuid: string): Promise<RedisWindow> {
+    const res = await redisClient.get(uuid);
+    // if no uuid is found, return -1 for tokens and timestamp, which are both impossible
+    if (res === null) return { currentTokens: -1, fixedWindowStart: -1 };
+    return JSON.parse(res);
+}
+
+async function setTokenCountInClient(
+    redisClient: ioredis.Redis,
+    uuid: string,
+    tokens: number,
+    time: number
+) {
+    const value: RedisWindow = { currentTokens: tokens, fixedWindowStart: time };
+    await redisClient.set(uuid, JSON.stringify(value));
+}
+describe('Test FixedWindow Rate Limiter', () => {
+    beforeEach(async () => {
+        client = new RedisMock();
+        limiter = new FixedWindow(CAPACITY, WINDOW_SIZE, client);
+        timestamp = new Date().valueOf();
+    });
+    describe('FixedWindow returns correct number of tokens and updates redis store as expected', () => {
+        describe('after an ALLOWED request...', () => {
+            afterEach(() => {
+                client.flushall();
+            });
+            test('current time window has no token initially', async () => {
+                // zero token used in this time window
+                const withdraw5 = 5;
+                expect((await limiter.processRequest(user1, timestamp, withdraw5)).tokens).toBe(
+                    CAPACITY - withdraw5
+                );
+                const tokenCountFull = await getWindowFromClient(client, user1);
+                expect(tokenCountFull.currentTokens).toBe(CAPACITY - withdraw5);
+            });
+        });
+    });
+});