From 7c1c8b2749af0a8d8afa7183285af09a7bdba377 Mon Sep 17 00:00:00 2001 From: Hampus Date: Mon, 31 Aug 2026 00:06:52 +0200 Subject: [PATCH] perf(rate-limit): precompute bucket hash and client identifier (#2205) --- .../src/api/middleware/RateLimitMiddleware.ts | 7 +- .../tests/RateLimitMiddleware.test.ts | 148 ++++++++++++++++++ 2 files changed, 151 insertions(+), 4 deletions(-) create mode 100644 fluxer_api/src/api/middleware/tests/RateLimitMiddleware.test.ts diff --git a/fluxer_api/src/api/middleware/RateLimitMiddleware.ts b/fluxer_api/src/api/middleware/RateLimitMiddleware.ts index 86aa2a760..a59b68a09 100644 --- a/fluxer_api/src/api/middleware/RateLimitMiddleware.ts +++ b/fluxer_api/src/api/middleware/RateLimitMiddleware.ts @@ -83,13 +83,12 @@ function getGlobalRateLimit(ctx: Context): number { return 50; } -function resolveBucket(bucket: string, ctx: Context): string { +function resolveBucket(bucket: string, clientId: string, ctx: Context): string { let resolved = bucket; const params = ctx.req.param(); for (const [key, value] of Object.entries(params)) { resolved = resolved.replace(`:${key}`, String(value)); } - const clientId = getClientIdentifier(ctx); return `${clientId}:${resolved}`; } @@ -134,6 +133,7 @@ async function revokeAuthenticatedSessionOnGlobalRateLimit(ctx: Context } export function RateLimitMiddleware(routeConfig: RouteRateLimitConfig): MiddlewareHandler { + const routeBucketHash = getBucketHash(routeConfig.bucket); return createMiddleware(async (ctx, next) => { if (!shouldEnforceRateLimits(ctx)) { await next(); @@ -152,7 +152,6 @@ export function RateLimitMiddleware(routeConfig: RouteRateLimitConfig): Middlewa const accountType = getAccountType(ctx); const showHeaders = shouldShowHeadersOnSuccess(accountType); const clientId = getClientIdentifier(ctx); - const routeBucketHash = getBucketHash(routeConfig.bucket); if (!routeConfig.config.exemptFromGlobal) { const globalLimit = getGlobalRateLimit(ctx); const globalResult = await rateLimitService.checkGlobalLimit(clientId, globalLimit); @@ -169,7 +168,7 @@ export function RateLimitMiddleware(routeConfig: RouteRateLimitConfig): Middlewa }); } } - const bucket = resolveBucket(routeConfig.bucket, ctx); + const bucket = resolveBucket(routeConfig.bucket, clientId, ctx); const bucketConfigWithAlgorithm: BucketConfig = { ...routeConfig.config, algorithm: 'leaky_bucket', diff --git a/fluxer_api/src/api/middleware/tests/RateLimitMiddleware.test.ts b/fluxer_api/src/api/middleware/tests/RateLimitMiddleware.test.ts new file mode 100644 index 000000000..2917ec75f --- /dev/null +++ b/fluxer_api/src/api/middleware/tests/RateLimitMiddleware.test.ts @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +import {createHash} from 'node:crypto'; +import type { + BucketConfig, + IRateLimitService, + RateLimitConfig, + RateLimitResult, +} from '@pkgs/rate_limit/src/IRateLimitService'; +import {type Context, Hono} from 'hono'; +import {describe, expect, test} from 'vitest'; +import type {HonoEnv} from '../../types/HonoEnv'; +import {RateLimitMiddleware, type RouteRateLimitConfig} from '../RateLimitMiddleware'; + +const CLIENT_IP = '203.0.113.10'; +const SWAPPED_CLIENT_IP = '198.51.100.7'; + +const WEBHOOK_READ: RouteRateLimitConfig = { + bucket: 'webhook:read::webhook_id', + config: {limit: 40, windowMs: 10000}, +}; + +const WEBHOOK_UPDATE: RouteRateLimitConfig = { + bucket: 'webhook:update::webhook_id', + config: {limit: 20, windowMs: 10000}, +}; + +function createAllowedResult(limit: number): RateLimitResult { + return { + allowed: true, + limit, + remaining: limit - 1, + resetTime: new Date(Date.now() + 10000), + resetAfterDecimal: 10, + }; +} + +class RecordingRateLimitService implements IRateLimitService { + readonly globalIdentifiers: Array = []; + readonly buckets: Array = []; + onGlobalCheck: () => void = () => undefined; + + async checkLimit(config: RateLimitConfig): Promise { + return createAllowedResult(config.maxAttempts); + } + + async peekLimit(config: RateLimitConfig): Promise { + return createAllowedResult(config.maxAttempts); + } + + async checkBucketLimit(bucket: string, config: BucketConfig): Promise { + this.buckets.push(bucket); + return createAllowedResult(config.limit); + } + + async checkGlobalLimit(identifier: string, limit: number): Promise { + this.globalIdentifiers.push(identifier); + this.onGlobalCheck(); + return createAllowedResult(limit); + } + + async resetLimit(_identifier: string): Promise {} + + async clearLimitsByIdentifierPrefix(_identifierPrefix: string): Promise { + return 0; + } +} + +interface Harness { + app: Hono; + service: RecordingRateLimitService; + getContext(): Context; +} + +function buildHarness(routeConfig: RouteRateLimitConfig): Harness { + const service = new RecordingRateLimitService(); + let context: Context | null = null; + const app = new Hono({strict: true}); + app.use('*', async (ctx, next) => { + context = ctx; + ctx.set('rateLimitService', service); + await next(); + }); + app.get('/webhooks/:webhook_id/:token', RateLimitMiddleware(routeConfig), (ctx) => ctx.text('ok')); + return { + app, + service, + getContext(): Context { + if (!context) { + throw new Error('no request has run yet'); + } + return context; + }, + }; +} + +async function callRoute(harness: Harness, webhookId: string, clientIp = CLIENT_IP): Promise { + return await harness.app.request(`http://localhost/webhooks/${webhookId}/secret`, { + headers: { + 'x-forwarded-for': clientIp, + 'x-fluxer-test-enable-rate-limits': 'true', + }, + }); +} + +function expectedBucketHash(bucket: string): string { + return createHash('sha256').update(bucket).digest('hex').slice(0, 16); +} + +describe('RateLimitMiddleware', () => { + test('reports the same bucket hash for every request to a route', async () => { + const harness = buildHarness(WEBHOOK_READ); + + const first = await callRoute(harness, '111'); + const second = await callRoute(harness, '222'); + + expect(first.status).toBe(200); + expect(second.status).toBe(200); + expect(first.headers.get('X-RateLimit-Bucket')).toBe(expectedBucketHash(WEBHOOK_READ.bucket)); + expect(second.headers.get('X-RateLimit-Bucket')).toBe(first.headers.get('X-RateLimit-Bucket')); + expect(harness.service.buckets).toEqual([`ip:${CLIENT_IP}:webhook:read:111`, `ip:${CLIENT_IP}:webhook:read:222`]); + }); + + test('gives routes with different buckets different bucket hashes', async () => { + const readHarness = buildHarness(WEBHOOK_READ); + const updateHarness = buildHarness(WEBHOOK_UPDATE); + + const read = await callRoute(readHarness, '111'); + const update = await callRoute(updateHarness, '111'); + + expect(read.headers.get('X-RateLimit-Bucket')).toBe(expectedBucketHash(WEBHOOK_READ.bucket)); + expect(update.headers.get('X-RateLimit-Bucket')).toBe(expectedBucketHash(WEBHOOK_UPDATE.bucket)); + expect(read.headers.get('X-RateLimit-Bucket')).not.toBe(update.headers.get('X-RateLimit-Bucket')); + }); + + test('resolves the client identifier once and reuses it for the bucket key', async () => { + const harness = buildHarness(WEBHOOK_READ); + harness.service.onGlobalCheck = () => { + harness.getContext().req.raw.headers.set('x-forwarded-for', SWAPPED_CLIENT_IP); + }; + + const response = await callRoute(harness, '111'); + + expect(response.status).toBe(200); + expect(harness.service.globalIdentifiers).toEqual([`ip:${CLIENT_IP}`]); + expect(harness.service.buckets).toEqual([`ip:${CLIENT_IP}:webhook:read:111`]); + }); +});