Skip to content

Commit 48abe17

Browse files
evanmcneelyevanmcneely
committed
Merge branch 'em/deploy2' of https://github.com/oslabs-beta/GraphQL-Gate into em/deploy2
2 parents d7cfa46 + ffa0296 commit 48abe17

File tree

3 files changed

+309
-2
lines changed

3 files changed

+309
-2
lines changed

‎src/@types/rateLimit.d.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,13 @@ export interface RedisBucket {
2424
timestamp: number;
2525
}
2626

27-
export interface RedisWindow {
27+
export interface FixedWindow {
2828
currentTokens: number;
29-
previousTokens: number;
3029
fixedWindowStart: number;
3130
}
31+
export interface RedisWindow extends FixedWindow {
32+
previousTokens: number;
33+
}
3234

3335
export type RedisLog = RedisBucket[];
3436

‎src/rateLimiters/fixedWindow.ts‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import Redis from 'ioredis';
2+
import { RateLimiter, RateLimiterResponse, FixedWindow as Window } from '../@types/rateLimit';
3+
4+
/**
5+
* The FixedWindow instance of a RateLimiter limits requests based on a unique user ID and a fixed time window.
6+
* Whenever a user makes a request the following steps are performed:
7+
* 1. Define the time window with fixed amount of queries.
8+
* 2. Update the timestamp of the last request.
9+
* 3. Allow the request and decrease the allowed amount of requests if the user has enough at this time window.
10+
* 4. Otherwise, disallow the request until the next time window opens.
11+
*/
12+
13+
class FixedWindow implements RateLimiter {
14+
private capacity: number;
15+
16+
private windowSize: number;
17+
18+
private client: Redis;
19+
20+
/**
21+
* Create a new instance of a FixedWindow rate limiter that can be connected to any database store
22+
* @param capacity max requests capacity in one time window
23+
* @param windowSize rate at which the token bucket is refilled
24+
* @param client redis client where rate limiter will cache information
25+
*/
26+
27+
constructor(capacity: number, windowSize: number, client: Redis) {
28+
this.capacity = capacity;
29+
this.windowSize = windowSize;
30+
this.client = client;
31+
if (windowSize <= 0 || capacity <= 0)
32+
throw Error('FixedWindow windowSize and capacity must be positive');
33+
}
34+
35+
/**
36+
* @function processRequest - Fixed Window algorithm to allow or block
37+
* based on the depth/complexity (in amount of tokens) of incoming requests.
38+
* Fixed Window
39+
* _________________________________
40+
* | *full capacity |
41+
* | | move to next time window
42+
* | token adds up until full | ---------->
43+
*____._________________________________.____
44+
* |<-- window size -->|
45+
*current timestamp next timestamp
46+
*
47+
* First, checks if a window exists in the redis cache.
48+
* If not, then `fixedWindowStart` is set as the current timestamp, and `currentTokens` is checked against `capacity`.
49+
* If enough room exists for the request, returns success as true and tokens as how many tokens remain in the current fixed window.
50+
*
51+
* If a window does exist in the cache, we first check if the timestamp is greater than the fixedWindowStart + windowSize.
52+
* If it isn't, we update currentToken with the incoming token until reach the capcity
53+
*
54+
* @param {string} uuid - unique identifer used to throttle requests
55+
* @param {number} timestamp - time the request was recieved
56+
* @param {number} [tokens=1] - complexity of the query for throttling requests
57+
* @return {*} {Promise<RateLimiterResponse>}
58+
* @memberof FixedWindow
59+
*/
60+
async processRequest(
61+
uuid: string,
62+
timestamp: number,
63+
tokens = 1
64+
): Promise<RateLimiterResponse> {
65+
// set the expiry of key-value pairs in the cache to 24 hours
66+
const keyExpiry = 86400000;
67+
68+
// attempt to get the value for the uuid from the redis cache
69+
const windowJSON = await this.client.get(uuid);
70+
71+
if (windowJSON === null) {
72+
const newUserWindow: Window = {
73+
currentTokens: tokens > this.capacity ? 0 : tokens,
74+
fixedWindowStart: timestamp,
75+
};
76+
77+
if (tokens > this.capacity) {
78+
await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
79+
return { success: false, tokens: this.capacity };
80+
}
81+
await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
82+
return { success: true, tokens: this.capacity - newUserWindow.currentTokens };
83+
}
84+
const window: Window = await JSON.parse(windowJSON);
85+
86+
const updatedUserWindow = this.updateTimeWindow(window, timestamp);
87+
updatedUserWindow.currentTokens += tokens;
88+
// update the currentToken until reaches its capacity
89+
if (updatedUserWindow.currentTokens > this.capacity) {
90+
updatedUserWindow.currentTokens -= tokens;
91+
return {
92+
success: false,
93+
tokens: this.capacity - updatedUserWindow.currentTokens,
94+
};
95+
}
96+
97+
// update a new time window, check the current capacity situation
98+
// if (tokens > this.capacity) {
99+
// await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
100+
// return { success: false, tokens: this.capacity };
101+
// }
102+
await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
103+
return {
104+
success: true,
105+
tokens: this.capacity - updatedUserWindow.currentTokens,
106+
};
107+
}
108+
109+
/**
110+
* Resets the rate limiter to the intial state by clearing the redis store.
111+
*/
112+
public reset(): void {
113+
this.client.flushall();
114+
}
115+
116+
private updateTimeWindow = (window: Window, timestamp: number): Window => {
117+
const updatedUserWindow: Window = {
118+
currentTokens: window.currentTokens,
119+
fixedWindowStart: window.fixedWindowStart,
120+
};
121+
if (timestamp >= window.fixedWindowStart + this.windowSize) {
122+
if (timestamp >= window.fixedWindowStart + this.windowSize * 2) {
123+
updatedUserWindow.fixedWindowStart = timestamp;
124+
updatedUserWindow.currentTokens = 0;
125+
} else {
126+
updatedUserWindow.fixedWindowStart = window.fixedWindowStart + this.windowSize;
127+
updatedUserWindow.currentTokens = 0;
128+
}
129+
}
130+
return updatedUserWindow;
131+
};
132+
}
133+
134+
export default FixedWindow;

‎test/rateLimiters/fixedWindow.test.ts‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
import * as ioredis from 'ioredis';
2+
import { FixedWindow as Window } from '../../src/@types/rateLimit';
3+
import FixedWindow from '../../src/rateLimiters/fixedWindow';
4+
5+
// eslint-disable-next-line @typescript-eslint/no-var-requires
6+
const RedisMock = require('ioredis-mock');
7+
8+
const CAPACITY = 10;
9+
const WINDOW_SIZE = 6000;
10+
11+
let limiter: FixedWindow;
12+
let client: ioredis.Redis;
13+
let timestamp: number;
14+
const user1 = '1';
15+
const user2 = '2';
16+
const user3 = '3';
17+
18+
async function getWindowFromClient(redisClient: ioredis.Redis, uuid: string): Promise<Window> {
19+
const res = await redisClient.get(uuid);
20+
// if no uuid is found, return -1 for tokens and timestamp, which are both impossible
21+
if (res === null) return { currentTokens: -1, fixedWindowStart: -1 };
22+
return JSON.parse(res);
23+
}
24+
25+
async function setTokenCountInClient(
26+
redisClient: ioredis.Redis,
27+
uuid: string,
28+
currentTokens: number,
29+
fixedWindowStart: number
30+
) {
31+
const value: Window = { currentTokens, fixedWindowStart };
32+
await redisClient.set(uuid, JSON.stringify(value));
33+
}
34+
describe('Test FixedWindow Rate Limiter', () => {
35+
beforeEach(async () => {
36+
client = new RedisMock();
37+
limiter = new FixedWindow(CAPACITY, WINDOW_SIZE, client);
38+
timestamp = new Date().valueOf();
39+
});
40+
describe('FixedWindow returns correct number of tokens and updates redis store as expected', () => {
41+
describe('after an ALLOWED request...', () => {
42+
afterEach(() => {
43+
client.flushall();
44+
});
45+
test('current time window has no token initially', async () => {
46+
// zero token used in this time window
47+
const withdraw5 = 5;
48+
expect((await limiter.processRequest(user1, timestamp, withdraw5)).tokens).toBe(
49+
CAPACITY - withdraw5
50+
);
51+
const tokenCountFull = await getWindowFromClient(client, user1);
52+
expect(tokenCountFull.currentTokens).toBe(5);
53+
});
54+
test('reached 40% capacity in current time window and still can pass request', async () => {
55+
const initial = 5;
56+
await setTokenCountInClient(client, user2, initial, timestamp);
57+
const partialWithdraw = 2;
58+
expect(
59+
(
60+
await limiter.processRequest(
61+
user2,
62+
timestamp + WINDOW_SIZE * 0.4,
63+
partialWithdraw
64+
)
65+
).tokens
66+
).toBe(CAPACITY - initial - partialWithdraw);
67+
68+
const tokenCountPartial = await getWindowFromClient(client, user2);
69+
expect(tokenCountPartial.currentTokens).toBe(initial + partialWithdraw);
70+
});
71+
72+
test('window is partially full and request has no leftover tokens', async () => {
73+
const initial = 6;
74+
const partialWithdraw = 4;
75+
await setTokenCountInClient(client, user2, initial, timestamp);
76+
expect(
77+
(await limiter.processRequest(user2, timestamp, partialWithdraw)).success
78+
).toBe(true);
79+
expect(
80+
(await limiter.processRequest(user2, timestamp, partialWithdraw)).tokens
81+
).toBe(0);
82+
});
83+
84+
test('window is partially full and request exceeds tokens in availability', async () => {
85+
const initial = 6;
86+
const partialWithdraw = 5;
87+
await setTokenCountInClient(client, user2, initial, timestamp);
88+
expect(
89+
(await limiter.processRequest(user2, timestamp, partialWithdraw)).success
90+
).toBe(false);
91+
expect(
92+
(await limiter.processRequest(user2, timestamp, partialWithdraw)).tokens
93+
).toBe(4);
94+
});
95+
});
96+
describe('after a BLOCKED request...', () => {
97+
afterEach(() => {
98+
client.flushall();
99+
});
100+
test('initial request is greater than capacity', async () => {
101+
// expect remaining tokens to be 10, b/c the 11 token request should be blocked
102+
expect((await limiter.processRequest(user1, timestamp, 11)).success).toBe(false);
103+
// expect current tokens in the window to still be 0
104+
expect((await getWindowFromClient(client, user1)).currentTokens).toBe(0);
105+
});
106+
test('window is partially full but not enough time elapsed to reach new window', async () => {
107+
const requestedTokens = 9;
108+
109+
await setTokenCountInClient(client, user2, requestedTokens, timestamp);
110+
// expect remaining tokens to be 1, b/c the 2-token-request should be blocked
111+
const result = await limiter.processRequest(user2, timestamp + WINDOW_SIZE - 1, 2);
112+
113+
expect(result.success).toBe(false);
114+
expect(result.tokens).toBe(1);
115+
116+
// expect current tokens in the window to still be 9
117+
expect((await getWindowFromClient(client, user2)).currentTokens).toBe(9);
118+
});
119+
});
120+
describe('updateTimeWindow function works as expect', () => {
121+
afterEach(() => {
122+
client.flushall();
123+
});
124+
test('New window is initialized after reaching the window size', async () => {
125+
const fullRequest = 10;
126+
await setTokenCountInClient(client, user3, fullRequest, timestamp);
127+
const noAccess = await limiter.processRequest(
128+
user3,
129+
timestamp + WINDOW_SIZE - 1,
130+
2
131+
);
132+
133+
// expect not passing any request
134+
expect(noAccess.tokens).toBe(0);
135+
expect(noAccess.success).toBe(false);
136+
137+
const newRequest = 1;
138+
expect(
139+
(await limiter.processRequest(user3, timestamp + WINDOW_SIZE, newRequest))
140+
.success
141+
).toBe(true);
142+
const count = await getWindowFromClient(client, user3);
143+
expect(count.currentTokens).toBe(1);
144+
});
145+
test('Request will be passed after two window sizes', async () => {
146+
const fullRequest = 10;
147+
await setTokenCountInClient(client, user3, fullRequest, timestamp);
148+
const noAccess = await limiter.processRequest(
149+
user3,
150+
timestamp + WINDOW_SIZE - 1,
151+
2
152+
);
153+
154+
// expect not passing any request
155+
expect(noAccess.tokens).toBe(0);
156+
expect(noAccess.success).toBe(false);
157+
158+
const newRequest = 6;
159+
// check if current time is over one window size
160+
const newAccess = await limiter.processRequest(
161+
user3,
162+
timestamp + WINDOW_SIZE * 2,
163+
newRequest
164+
);
165+
166+
expect(newAccess.tokens).toBe(4);
167+
expect(newAccess.success).toBe(true);
168+
});
169+
});
170+
});
171+
});

0 commit comments

Comments
 (0)