Expose rateLimitSlidingWindow strategy for low-level usage (#6815)

Author: jnsdls

.changeset/tall-llamas-dress.md

@@ -0,0 +1,5 @@
+---
+"@thirdweb-dev/service-utils": patch
+---
+
+expose `rateLimitSlidingWindow` strategy directly for more low-level usage

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<number>;
-  mget(keys: string[]): Promise<(string | null)[]>;
-  expire(key: string, seconds: number): Promise<number>;
-};
-
 /**
  * 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}`;
-}

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);
 

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<number>;
+  mget(keys: string[]): Promise<(string | null)[]>;
+  expire(key: string, seconds: number): Promise<number>;
+};
+
+export type CoreRateLimitResult = {
+  rateLimited: boolean;
+  requestCount: number;
+  rateLimit: number;
+};

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<CoreRateLimitResult> {
+  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}`;
+}

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";