From ffc2737478c6f9efd5de9fbaf526b13164727f87 Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Sat, 7 Dec 2024 10:22:45 -0500 Subject: implement SkRateLimiterService with Leaky Bucket rate limiting --- .../unit/server/api/SkRateLimiterServiceTests.ts | 703 +++++++++++++++++++++ 1 file changed, 703 insertions(+) create mode 100644 packages/backend/test/unit/server/api/SkRateLimiterServiceTests.ts (limited to 'packages/backend/test/unit/server/api') diff --git a/packages/backend/test/unit/server/api/SkRateLimiterServiceTests.ts b/packages/backend/test/unit/server/api/SkRateLimiterServiceTests.ts new file mode 100644 index 0000000000..8554aa39ef --- /dev/null +++ b/packages/backend/test/unit/server/api/SkRateLimiterServiceTests.ts @@ -0,0 +1,703 @@ +/* + * SPDX-FileCopyrightText: hazelnoot and other Sharkey contributors + * SPDX-License-Identifier: AGPL-3.0-only + */ + +import { KEYWORD } from 'color-convert/conversions.js'; +import type Redis from 'ioredis'; +import { LegacyRateLimit, LimitCounter, RateLimit, SkRateLimiterService } from '@/server/api/SkRateLimiterService.js'; +import { LoggerService } from '@/core/LoggerService.js'; + +/* eslint-disable @typescript-eslint/no-non-null-assertion */ +/* eslint-disable @typescript-eslint/no-unnecessary-condition */ + +describe(SkRateLimiterService, () => { + let mockTimeService: { now: number, date: Date } = null!; + let mockRedisGet: ((key: string) => string | null) | undefined = undefined; + let mockRedisSet: ((args: unknown[]) => void) | undefined = undefined; + let mockEnvironment: Record = null!; + let serviceUnderTest: () => SkRateLimiterService = null!; + + let loggedMessages: { level: string, data: unknown[] }[] = []; + + beforeEach(() => { + mockTimeService = { + now: 0, + get date() { + return new Date(mockTimeService.now); + }, + }; + + mockRedisGet = undefined; + mockRedisSet = undefined; + const mockRedisClient = { + get(key: string) { + if (mockRedisGet) return Promise.resolve(mockRedisGet(key)); + else return Promise.resolve(null); + }, + set(...args: unknown[]): Promise { + if (mockRedisSet) mockRedisSet(args); + return Promise.resolve(); + }, + } as unknown as Redis.Redis; + + mockEnvironment = Object.create(process.env); + mockEnvironment.NODE_ENV = 'production'; + const mockEnvService = { + env: mockEnvironment, + }; + + loggedMessages = []; + const mockLogService = { + getLogger() { + return { + createSubLogger(context: string, color?: KEYWORD) { + return mockLogService.getLogger(context, color); + }, + error(...data: unknown[]) { + loggedMessages.push({ level: 'error', data }); + }, + warn(...data: unknown[]) { + loggedMessages.push({ level: 'warn', data }); + }, + succ(...data: unknown[]) { + loggedMessages.push({ level: 'succ', data }); + }, + debug(...data: unknown[]) { + loggedMessages.push({ level: 'debug', data }); + }, + info(...data: unknown[]) { + loggedMessages.push({ level: 'info', data }); + }, + }; + }, + } as unknown as LoggerService; + + let service: SkRateLimiterService | undefined = undefined; + serviceUnderTest = () => { + return service ??= new SkRateLimiterService(mockTimeService, mockRedisClient, mockLogService, mockEnvService); + }; + }); + + function expectNoUnhandledErrors() { + const unhandledErrors = loggedMessages.filter(m => m.level === 'error'); + if (unhandledErrors.length > 0) { + throw new Error(`Test failed: got unhandled errors ${unhandledErrors.join('\n')}`); + } + } + + describe('limit', () => { + const actor = 'actor'; + const key = 'test'; + + let counter: LimitCounter | undefined = undefined; + let minCounter: LimitCounter | undefined = undefined; + + beforeEach(() => { + counter = undefined; + minCounter = undefined; + + mockRedisGet = (key: string) => { + if (key === 'rl_actor_test' && counter) { + return JSON.stringify(counter); + } + + if (key === 'rl_actor_test_min' && minCounter) { + return JSON.stringify(minCounter); + } + + return null; + }; + + mockRedisSet = (args: unknown[]) => { + const [key, value] = args; + + if (key === 'rl_actor_test') { + if (value == null) counter = undefined; + else if (typeof(value) === 'string') counter = JSON.parse(value); + else throw new Error('invalid redis call'); + } + + if (key === 'rl_actor_test_min') { + if (value == null) minCounter = undefined; + else if (typeof(value) === 'string') minCounter = JSON.parse(value); + else throw new Error('invalid redis call'); + } + }; + }); + + it('should bypass in non-production', async () => { + mockEnvironment.NODE_ENV = 'test'; + + const info = await serviceUnderTest().limit({ key: 'l', type: undefined, max: 0 }, 'actor'); + + expect(info.blocked).toBeFalsy(); + expect(info.remaining).toBe(Number.MAX_SAFE_INTEGER); + expect(info.resetSec).toBe(0); + expect(info.resetMs).toBe(0); + expect(info.fullResetSec).toBe(0); + expect(info.fullResetMs).toBe(0); + }); + + describe('with bucket limit', () => { + let limit: RateLimit = null!; + + beforeEach(() => { + limit = { + type: 'bucket', + key: 'test', + size: 1, + }; + }); + + it('should allow when limit is not reached', async () => { + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should not error when allowed', async () => { + await serviceUnderTest().limit(limit, actor); + + expectNoUnhandledErrors(); + }); + + it('should return correct info when allowed', async () => { + limit.size = 2; + counter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.remaining).toBe(0); + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(2); + expect(info.fullResetMs).toBe(2000); + }); + + it('should increment counter when called', async () => { + await serviceUnderTest().limit(limit, actor); + + expect(counter).not.toBeUndefined(); + expect(counter?.c).toBe(1); + }); + + it('should set timestamp when called', async () => { + mockTimeService.now = 1000; + + await serviceUnderTest().limit(limit, actor); + + expect(counter).not.toBeUndefined(); + expect(counter?.t).toBe(1000); + }); + + it('should decrement counter when dripRate has passed', async () => { + counter = { c: 2, t: 0 }; + mockTimeService.now = 2000; + + await serviceUnderTest().limit(limit, actor); + + expect(counter).not.toBeUndefined(); + expect(counter?.c).toBe(1); // 2 (starting) - 2 (2x1 drip) + 1 (call) = 1 + }); + + it('should decrement counter by dripSize', async () => { + counter = { c: 2, t: 0 }; + limit.dripSize = 2; + mockTimeService.now = 1000; + + await serviceUnderTest().limit(limit, actor); + + expect(counter).not.toBeUndefined(); + expect(counter?.c).toBe(1); // 2 (starting) - 2 (1x2 drip) + 1 (call) = 1 + }); + + it('should maintain counter between calls over time', async () => { + limit.size = 5; + + await serviceUnderTest().limit(limit, actor); // 0 + 1 = 1 + mockTimeService.now += 1000; // 1 - 1 = 0 + await serviceUnderTest().limit(limit, actor); // 0 + 1 = 1 + await serviceUnderTest().limit(limit, actor); // 1 + 1 = 2 + await serviceUnderTest().limit(limit, actor); // 2 + 1 = 3 + mockTimeService.now += 1000; // 3 - 1 = 2 + mockTimeService.now += 1000; // 2 - 1 = 1 + await serviceUnderTest().limit(limit, actor); // 1 + 1 = 2 + + expect(counter?.c).toBe(2); + expect(counter?.t).toBe(3000); + }); + + it('should log error and continue when update fails', async () => { + mockRedisSet = () => { + throw new Error('test error'); + }; + + await serviceUnderTest().limit(limit, actor); + + const matchingError = loggedMessages + .find(m => m.level === 'error' && m.data + .some(d => typeof(d) === 'string' && d.includes('Failed to update limit'))); + expect(matchingError).toBeTruthy(); + }); + + it('should block when bucket is filled', async () => { + counter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeTruthy(); + }); + + it('should calculate correct info when blocked', async () => { + counter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(1); + expect(info.fullResetMs).toBe(1000); + }); + + it('should allow when bucket is filled but should drip', async () => { + counter = { c: 1, t: 0 }; + mockTimeService.now = 1000; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should scale limit by factor', async () => { + counter = { c: 1, t: 0 }; + + const i1 = await serviceUnderTest().limit(limit, actor, 2); // 1 + 1 = 2 + const i2 = await serviceUnderTest().limit(limit, actor, 2); // 2 + 1 = 3 + + expect(i1.blocked).toBeFalsy(); + expect(i2.blocked).toBeTruthy(); + }); + + it('should set key expiration', async () => { + mockRedisSet = args => { + expect(args[2]).toBe('PX'); + expect(args[3]).toBe(1000); + }; + + await serviceUnderTest().limit(limit, actor); + }); + + it('should not increment when already blocked', async () => { + counter = { c: 1, t: 0 }; + mockTimeService.now += 100; + + await serviceUnderTest().limit(limit, actor); + + expect(counter?.c).toBe(1); + expect(counter?.t).toBe(0); + }); + }); + + describe('with min interval', () => { + let limit: MutableLegacyRateLimit = null!; + + beforeEach(() => { + limit = { + type: undefined, + key, + minInterval: 1000, + }; + }); + + it('should allow when limit is not reached', async () => { + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should not error when allowed', async () => { + await serviceUnderTest().limit(limit, actor); + + expectNoUnhandledErrors(); + }); + + it('should calculate correct info when allowed', async () => { + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.remaining).toBe(0); + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(1); + expect(info.fullResetMs).toBe(1000); + }); + + it('should increment counter when called', async () => { + await serviceUnderTest().limit(limit, actor); + + expect(minCounter).not.toBeUndefined(); + expect(minCounter?.c).toBe(1); + }); + + it('should set timestamp when called', async () => { + mockTimeService.now = 1000; + + await serviceUnderTest().limit(limit, actor); + + expect(minCounter).not.toBeUndefined(); + expect(minCounter?.t).toBe(1000); + }); + + it('should decrement counter when minInterval has passed', async () => { + minCounter = { c: 1, t: 0 }; + mockTimeService.now = 1000; + + await serviceUnderTest().limit(limit, actor); + + expect(minCounter).not.toBeUndefined(); + expect(minCounter?.c).toBe(1); // 1 (starting) - 1 (interval) + 1 (call) = 1 + }); + + it('should reset counter entirely', async () => { + minCounter = { c: 2, t: 0 }; + mockTimeService.now = 1000; + + await serviceUnderTest().limit(limit, actor); + + expect(minCounter).not.toBeUndefined(); + expect(minCounter?.c).toBe(1); // 2 (starting) - 2 (interval) + 1 (call) = 1 + }); + + it('should maintain counter between calls over time', async () => { + await serviceUnderTest().limit(limit, actor); // 0 + 1 = 1 + mockTimeService.now += 1000; // 1 - 1 = 0 + await serviceUnderTest().limit(limit, actor); // 0 + 1 = 1 + await serviceUnderTest().limit(limit, actor); // blocked + await serviceUnderTest().limit(limit, actor); // blocked + mockTimeService.now += 1000; // 1 - 1 = 0 + mockTimeService.now += 1000; // 0 - 1 = 0 + await serviceUnderTest().limit(limit, actor); // 0 + 1 = 1 + + expect(minCounter?.c).toBe(1); + expect(minCounter?.t).toBe(3000); + }); + + it('should log error and continue when update fails', async () => { + mockRedisSet = () => { + throw new Error('test error'); + }; + + await serviceUnderTest().limit(limit, actor); + + const matchingError = loggedMessages + .find(m => m.level === 'error' && m.data + .some(d => typeof(d) === 'string' && d.includes('Failed to update limit'))); + expect(matchingError).toBeTruthy(); + }); + + it('should block when interval exceeded', async () => { + minCounter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeTruthy(); + }); + + it('should calculate correct info when blocked', async () => { + minCounter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(1); + expect(info.fullResetMs).toBe(1000); + }); + + it('should allow when bucket is filled but interval has passed', async () => { + minCounter = { c: 1, t: 0 }; + mockTimeService.now = 1000; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should scale limit by factor', async () => { + minCounter = { c: 1, t: 0 }; + + const i1 = await serviceUnderTest().limit(limit, actor, 2); // 1 + 1 = 2 + const i2 = await serviceUnderTest().limit(limit, actor, 2); // 2 + 1 = 3 + + expect(i1.blocked).toBeFalsy(); + expect(i2.blocked).toBeTruthy(); + }); + + it('should set key expiration', async () => { + mockRedisSet = args => { + expect(args[2]).toBe('PX'); + expect(args[3]).toBe(1000); + }; + + await serviceUnderTest().limit(limit, actor); + }); + + it('should not increment when already blocked', async () => { + minCounter = { c: 1, t: 0 }; + mockTimeService.now += 100; + + await serviceUnderTest().limit(limit, actor); + + expect(minCounter?.c).toBe(1); + expect(minCounter?.t).toBe(0); + }); + }); + + describe('with legacy limit', () => { + let limit: MutableLegacyRateLimit = null!; + + beforeEach(() => { + limit = { + type: undefined, + key, + max: 1, + duration: 1000, + }; + }); + + it('should allow when limit is not reached', async () => { + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should not error when allowed', async () => { + await serviceUnderTest().limit(limit, actor); + + expectNoUnhandledErrors(); + }); + + it('should infer dripRate from duration', async () => { + limit.max = 10; + limit.duration = 10000; + counter = { c: 10, t: 0 }; + + const i1 = await serviceUnderTest().limit(limit, actor); + mockTimeService.now += 1000; + const i2 = await serviceUnderTest().limit(limit, actor); + mockTimeService.now += 2000; + const i3 = await serviceUnderTest().limit(limit, actor); + const i4 = await serviceUnderTest().limit(limit, actor); + const i5 = await serviceUnderTest().limit(limit, actor); + mockTimeService.now += 2000; + const i6 = await serviceUnderTest().limit(limit, actor); + + expect(i1.blocked).toBeTruthy(); + expect(i2.blocked).toBeFalsy(); + expect(i3.blocked).toBeFalsy(); + expect(i4.blocked).toBeFalsy(); + expect(i5.blocked).toBeTruthy(); + expect(i6.blocked).toBeFalsy(); + }); + + it('should calculate correct info when allowed', async () => { + limit.max = 10; + limit.duration = 10000; + counter = { c: 10, t: 0 }; + mockTimeService.now += 2000; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.remaining).toBe(1); + expect(info.resetSec).toBe(0); + expect(info.resetMs).toBe(0); + expect(info.fullResetSec).toBe(9); + expect(info.fullResetMs).toBe(9000); + }); + + it('should calculate correct info when blocked', async () => { + limit.max = 10; + limit.duration = 10000; + counter = { c: 10, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.remaining).toBe(0); + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(10); + expect(info.fullResetMs).toBe(10000); + }); + + it('should allow when bucket is filled but interval has passed', async () => { + counter = { c: 10, t: 0 }; + mockTimeService.now = 1000; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeTruthy(); + }); + + it('should scale limit by factor', async () => { + counter = { c: 10, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor, 2); // 10 + 1 = 11 + + expect(info.blocked).toBeTruthy(); + }); + + it('should set key expiration', async () => { + mockRedisSet = args => { + expect(args[2]).toBe('PX'); + expect(args[3]).toBe(1000); + }; + + await serviceUnderTest().limit(limit, actor); + }); + + it('should not increment when already blocked', async () => { + counter = { c: 1, t: 0 }; + mockTimeService.now += 100; + + await serviceUnderTest().limit(limit, actor); + + expect(counter?.c).toBe(1); + expect(counter?.t).toBe(0); + }); + }); + + describe('with legacy limit and min interval', () => { + let limit: MutableLegacyRateLimit = null!; + + beforeEach(() => { + limit = { + type: undefined, + key, + max: 5, + duration: 5000, + minInterval: 1000, + }; + }); + + it('should allow when limit is not reached', async () => { + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should not error when allowed', async () => { + await serviceUnderTest().limit(limit, actor); + + expectNoUnhandledErrors(); + }); + + it('should block when limit exceeded', async () => { + counter = { c: 5, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeTruthy(); + }); + + it('should block when minInterval exceeded', async () => { + minCounter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeTruthy(); + }); + + it('should calculate correct info when allowed', async () => { + counter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.remaining).toBe(0); + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(2); + expect(info.fullResetMs).toBe(2000); + }); + + it('should calculate correct info when blocked by limit', async () => { + counter = { c: 5, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.remaining).toBe(0); + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(5); + expect(info.fullResetMs).toBe(5000); + }); + + it('should calculate correct info when blocked by minInterval', async () => { + minCounter = { c: 1, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.remaining).toBe(0); + expect(info.resetSec).toBe(1); + expect(info.resetMs).toBe(1000); + expect(info.fullResetSec).toBe(1); + expect(info.fullResetMs).toBe(1000); + }); + + it('should allow when counter is filled but interval has passed', async () => { + counter = { c: 5, t: 0 }; + mockTimeService.now = 1000; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should allow when minCounter is filled but interval has passed', async () => { + minCounter = { c: 1, t: 0 }; + mockTimeService.now = 1000; + + const info = await serviceUnderTest().limit(limit, actor); + + expect(info.blocked).toBeFalsy(); + }); + + it('should scale limit by factor', async () => { + minCounter = { c: 5, t: 0 }; + + const info = await serviceUnderTest().limit(limit, actor, 2); + + expect(info.blocked).toBeTruthy(); + }); + + it('should set key expiration', async () => { + mockRedisSet = args => { + expect(args[2]).toBe('PX'); + expect(args[3]).toBe(1000); + }; + + await serviceUnderTest().limit(limit, actor); + }); + + it('should not increment when already blocked', async () => { + counter = { c: 5, t: 0 }; + minCounter = { c: 1, t: 0 }; + mockTimeService.now += 100; + + await serviceUnderTest().limit(limit, actor); + + expect(counter?.c).toBe(5); + expect(counter?.t).toBe(0); + expect(minCounter?.c).toBe(1); + expect(minCounter?.t).toBe(0); + }); + }); + }); +}); + +// The same thing, but mutable +interface MutableLegacyRateLimit extends LegacyRateLimit { + key: string; + duration?: number; + max?: number; + minInterval?: number; +} -- cgit v1.2.3-freya