From 12ab488e3bcca90eb5838bf7cd14a3ddf6257698 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Tue, 22 Apr 2025 18:41:49 +0200 Subject: [PATCH] Expose rateLimitSlidingWindow strategy for low-level usage --- .changeset/tall-llamas-dress.md | 5 ++ .../service-utils/src/core/rateLimit/index.ts | 78 ++++-------------- .../src/core/rateLimit/rateLimit.test.ts | 4 +- .../src/core/rateLimit/strategies/shared.ts | 12 +++ .../rateLimit/strategies/sliding-window.ts | 81 +++++++++++++++++++ packages/service-utils/src/index.ts | 3 + 6 files changed, 120 insertions(+), 63 deletions(-) create mode 100644 .changeset/tall-llamas-dress.md create mode 100644 packages/service-utils/src/core/rateLimit/strategies/shared.ts create mode 100644 packages/service-utils/src/core/rateLimit/strategies/sliding-window.ts diff --git a/.changeset/tall-llamas-dress.md b/.changeset/tall-llamas-dress.md new file mode 100644 index 00000000000..f5a21008f9c --- /dev/null +++ b/.changeset/tall-llamas-dress.md @@ -0,0 +1,5 @@ +--- +"@thirdweb-dev/service-utils": patch +--- + +expose `rateLimitSlidingWindow` strategy directly for more low-level usage diff --git a/packages/service-utils/src/core/rateLimit/index.ts b/packages/service-utils/src/core/rateLimit/index.ts index 7ca2ae43745..dd1b98e7504 100644 --- a/packages/service-utils/src/core/rateLimit/index.ts +++ b/packages/service-utils/src/core/rateLimit/index.ts @@ -1,15 +1,10 @@ import type { CoreServiceConfig, TeamResponse } from "../api.js"; +import type { IRedis } from "./strategies/shared.js"; +import { rateLimitSlidingWindow } from "./strategies/sliding-window.js"; import type { RateLimitResult } from "./types.js"; const SLIDING_WINDOW_SECONDS = 10; -// Redis interface compatible with ioredis (Node) and upstash (Cloudflare Workers). -type IRedis = { - incrby(key: string, value: number): Promise; - mget(keys: string[]): Promise<(string | null)[]>; - expire(key: string, seconds: number): Promise; -}; - /** * Increments the request count for this team and returns whether the team has hit their rate limit. * Uses a sliding 10 second window. @@ -30,68 +25,29 @@ export async function rateLimit(args: { const { team, limitPerSecond, serviceConfig, redis, increment = 1 } = args; const { serviceScope } = serviceConfig; - if (limitPerSecond === 0) { - // No rate limit is provided. Assume the request is not rate limited. - return { - rateLimited: false, - requestCount: 0, - rateLimit: 0, - }; - } - - // Enforce rate limit: sum the total requests in the last `SLIDING_WINDOW_SECONDS` seconds. - const currentSecond = Math.floor(Date.now() / 1000); - const keys = Array.from({ length: SLIDING_WINDOW_SECONDS }, (_, i) => - getRequestCountAtSecondCacheKey(serviceScope, team.id, currentSecond - i), - ); - const counts = await redis.mget(keys); - const totalCount = counts.reduce( - (sum, count) => sum + (count ? Number.parseInt(count) : 0), - 0, - ); + const rateLimitResult = await rateLimitSlidingWindow({ + redis, + limitPerSecond, + key: `rate-limit:${serviceScope}:${team.id}`, + increment, + windowSeconds: SLIDING_WINDOW_SECONDS, + }); - const limitPerWindow = limitPerSecond * SLIDING_WINDOW_SECONDS; - - if (totalCount > limitPerWindow) { + // if the request is rate limited, return the rate limit result. + if (rateLimitResult.rateLimited) { return { rateLimited: true, - requestCount: totalCount, - rateLimit: limitPerWindow, + requestCount: rateLimitResult.requestCount, + rateLimit: rateLimitResult.rateLimit, status: 429, - errorMessage: `You've exceeded your ${serviceScope} rate limit at ${limitPerSecond} reqs/sec. Please upgrade your plan to get higher rate limits.`, + errorMessage: `You've exceeded your ${serviceScope} rate limit at ${limitPerSecond} requests per second. Please upgrade your plan to increase your limits: https://thirdweb.com/team/${team.slug}/~/settings/billing`, errorCode: "RATE_LIMIT_EXCEEDED", }; } - - // Non-blocking: increment the request count for the current second. - (async () => { - try { - const key = getRequestCountAtSecondCacheKey( - serviceScope, - team.id, - currentSecond, - ); - await redis.incrby(key, increment); - // If this is the first time setting this key, expire it after the sliding window is past. - if (counts[0] === null) { - await redis.expire(key, SLIDING_WINDOW_SECONDS + 1); - } - } catch (error) { - console.error("Error updating rate limit key:", error); - } - })(); - + // otherwise, the request is not rate limited. return { rateLimited: false, - requestCount: totalCount + increment, - rateLimit: limitPerWindow, + requestCount: rateLimitResult.requestCount, + rateLimit: rateLimitResult.rateLimit, }; } - -function getRequestCountAtSecondCacheKey( - serviceScope: CoreServiceConfig["serviceScope"], - teamId: string, - second: number, -) { - return `rate-limit:${serviceScope}:${teamId}:${second}`; -} diff --git a/packages/service-utils/src/core/rateLimit/rateLimit.test.ts b/packages/service-utils/src/core/rateLimit/rateLimit.test.ts index 3161bd796d4..2f5bd40de09 100644 --- a/packages/service-utils/src/core/rateLimit/rateLimit.test.ts +++ b/packages/service-utils/src/core/rateLimit/rateLimit.test.ts @@ -31,7 +31,7 @@ describe("rateLimit", () => { expect(result).toEqual({ rateLimited: false, - requestCount: 0, + requestCount: 1, rateLimit: 0, }); expect(mockRedis.mget).not.toHaveBeenCalled(); @@ -62,7 +62,7 @@ describe("rateLimit", () => { // Verify correct keys are checked const expectedKeys = Array.from( { length: SLIDING_WINDOW_SECONDS }, - (_, i) => `rate-limit:storage:1:${currentSecond - i}`, + (_, i) => `rate-limit:storage:1:s_${currentSecond - i}`, ); expect(mockRedis.mget).toHaveBeenCalledWith(expectedKeys); diff --git a/packages/service-utils/src/core/rateLimit/strategies/shared.ts b/packages/service-utils/src/core/rateLimit/strategies/shared.ts new file mode 100644 index 00000000000..d95db31686d --- /dev/null +++ b/packages/service-utils/src/core/rateLimit/strategies/shared.ts @@ -0,0 +1,12 @@ +// Redis interface compatible with ioredis (Node) and upstash (Cloudflare Workers). +export type IRedis = { + incrby(key: string, value: number): Promise; + mget(keys: string[]): Promise<(string | null)[]>; + expire(key: string, seconds: number): Promise; +}; + +export type CoreRateLimitResult = { + rateLimited: boolean; + requestCount: number; + rateLimit: number; +}; diff --git a/packages/service-utils/src/core/rateLimit/strategies/sliding-window.ts b/packages/service-utils/src/core/rateLimit/strategies/sliding-window.ts new file mode 100644 index 00000000000..eba6372e1f7 --- /dev/null +++ b/packages/service-utils/src/core/rateLimit/strategies/sliding-window.ts @@ -0,0 +1,81 @@ +import type { CoreRateLimitResult, IRedis } from "./shared.js"; + +type RateLimitSlidingWindowOptions = { + redis: IRedis; + limitPerSecond: number; + key: string; + /** + * The number of requests to increment by. + * @default 1 + */ + increment?: number; + /** + * The number of seconds to look back for the sliding window. + * @default 10 + */ + windowSeconds?: number; +}; + +export async function rateLimitSlidingWindow( + options: RateLimitSlidingWindowOptions, +): Promise { + const WINDOW_SIZE = options.windowSeconds || 10; + const INCREMENT_BY = options.increment || 1; + + // No rate limit is provided. Assume the request is not rate limited. + if (options.limitPerSecond <= 0) { + return { + rateLimited: false, + requestCount: INCREMENT_BY, + rateLimit: 0, + }; + } + + // Enforce rate limit: sum the total requests in the last `SLIDING_WINDOW_SECONDS` seconds. + const currentSecond = Math.floor(Date.now() / 1000); + const keys = Array.from({ length: WINDOW_SIZE }, (_, i) => + getRequestCountAtSecondCacheKey(options.key, currentSecond - i), + ); + const counts = await options.redis.mget(keys); + const totalCount = counts.reduce( + (sum, count) => sum + (count ? Number.parseInt(count) : 0), + 0, + ); + + const limitPerWindow = options.limitPerSecond * WINDOW_SIZE; + + if (totalCount > limitPerWindow) { + return { + rateLimited: true, + requestCount: totalCount, + rateLimit: limitPerWindow, + }; + } + + // Non-blocking: increment the request count for the current second. + (async () => { + try { + const incrKey = getRequestCountAtSecondCacheKey( + options.key, + currentSecond, + ); + await options.redis.incrby(incrKey, INCREMENT_BY); + // If this is the first time setting this key, expire it after the sliding window is past. + if (counts[0] === null) { + await options.redis.expire(incrKey, WINDOW_SIZE + 1); + } + } catch (error) { + console.error("Error updating rate limit key:", error); + } + })(); + + return { + rateLimited: false, + requestCount: totalCount + INCREMENT_BY, + rateLimit: limitPerWindow, + }; +} + +function getRequestCountAtSecondCacheKey(key: string, second: number) { + return `${key}:s_${second}`; +} diff --git a/packages/service-utils/src/index.ts b/packages/service-utils/src/index.ts index 7c55de31136..07b03d88fcd 100644 --- a/packages/service-utils/src/index.ts +++ b/packages/service-utils/src/index.ts @@ -20,3 +20,6 @@ export { authorizeBundleId, authorizeDomain, } from "./core/authorize/client.js"; + +export { rateLimitSlidingWindow } from "./core/rateLimit/strategies/sliding-window.js"; +export { rateLimit } from "./core/rateLimit/index.js";