ref: move from sliding window to token bucket in rate limiter

bun.lock

@@ -19,6 +19,7 @@
         "hono": "^4.7.8",
         "hono-rate-limiter": "^0.4.2",
         "ioredis": "^5.6.1",
+        "limiter": "^3.0.0",
         "postgres": "^3.4.5",
         "rate-limit-redis": "^4.2.0",
         "yup": "^1.6.1",
@@ -845,6 +846,8 @@
 
     "lightningcss-win32-x64-msvc": ["lightningcss-win32-x64-msvc@1.29.2", "", { "os": "win32", "cpu": "x64" }, "sha512-EdIUW3B2vLuHmv7urfzMI/h2fmlnOQBk1xlsDxkN1tCWKjNFjfLhGxYk8C8mzpSfr+A6jFFIi8fU6LbQGsRWjA=="],
 
+    "limiter": ["limiter@3.0.0", "", {}, "sha512-hev7DuXojsTFl2YwyzUJMDnZ/qBDd3yZQLSH3aD4tdL1cqfc3TMnoecEJtWFaQFdErZsKoFMBTxF/FBSkgDbEg=="],
+
     "lines-and-columns": ["lines-and-columns@1.2.4", "", {}, "sha512-7ylylesZQ/PV29jhEDl3Ufjo6ZX7gCqJr5F7PKrqc93v7fzSymt1BpwEU8nAUXs8qzzvqhbjhK5QZg6Mt/HkBg=="],
 
     "locate-character": ["locate-character@3.0.0", "", {}, "sha512-SW13ws7BjaeJ6p7Q6CO2nchbYEc3X3J6WrmTTDto7yMPqVSZTUyY5Tjbid+Ab8gLnATtygYtiDIJGQRRn2ZOiA=="],

packages/backend/middleware/index.ts

@@ -1,7 +1,7 @@
 import { Hono } from "hono";
 import { Variables } from "hono/types";
 import { bodyLimitForPing } from "./bodyLimits.ts";
-import { pingHandler } from "../routes/ping.ts";
+import { pingHandler } from "routes/ping";
 import { registerRateLimiter } from "./rateLimiters.ts";
 import { preetifyResponse } from "./preetifyResponse.ts";
 import { logger } from "./logger.ts";

packages/backend/middleware/rateLimiters.ts

@@ -11,14 +11,18 @@ export const registerRateLimiter = rateLimiter<BlankEnv, "/user", {}>({
 	limit: 10,
 	standardHeaders: "draft-6",
 	keyGenerator: (c) => {
+		let ipAddr = crypto.randomUUID() as string;
 		const info = getConnInfo(c as unknown as Context<BlankEnv, "/user", {}>);
-		if (!info.remote || !info.remote.address) {
+		if (info.remote && info.remote.address) {
-			return crypto.randomUUID();
+			ipAddr = info.remote.address;
+		}
+		const forwardedFor = c.req.header("X-Forwarded-For");
+		if (forwardedFor) {
+			ipAddr = forwardedFor.split(",")[0];
 		}
-		const addr = info.remote.address;
 		const path = new URL(c.req.url).pathname;
 		const method = c.req.method;
-		return `${method}-${path}@${addr}`;
+		return `${method}-${path}@${ipAddr}`;
 	},
 	store: new RedisStore({
 		// @ts-expect-error - Known issue: the `c`all` function is not present in @types/ioredis

packages/backend/package.json

@@ -10,6 +10,7 @@
 		"hono": "^4.7.8",
 		"hono-rate-limiter": "^0.4.2",
 		"ioredis": "^5.6.1",
+		"limiter": "^3.0.0",
 		"postgres": "^3.4.5",
 		"rate-limit-redis": "^4.2.0",
 		"yup": "^1.6.1",

packages/backend/routes/validation/session/POST.ts

@@ -0,0 +1,25 @@
+import { createHandlers } from "src/utils.ts";
+
+const DIFFICULTY = 200000;
+
+const createNewChallenge = async (difficulty: number) => {
+    const baseURL = process.env["UCAPTCHA_URL"];
+    const url = new URL(baseURL);
+    url.pathname = "/challenge";
+    const res = await fetch(url.toString(), {
+        method: "POST",
+        headers: {
+            "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+            difficulty: difficulty,
+        })
+    });
+    const data = await res.json();
+    return data;
+}
+
+export const createValidationSessionHandler = createHandlers(async (c) => {
+    const challenge = await createNewChallenge(DIFFICULTY);
+    return c.json(challenge);
+});

packages/backend/routes/validation/session/index.ts

@@ -0,0 +1 @@
+export * from "./POST.ts";

packages/backend/src/routing.ts

@@ -4,6 +4,7 @@ import { registerHandler } from "routes/user";
 import { videoInfoHandler, getSnapshotsHanlder } from "routes/video";
 import { Hono } from "hono";
 import { Variables } from "hono/types";
+import { createValidationSessionHandler } from "routes/validation/session";
 
 export function configureRoutes(app: Hono<{ Variables: Variables }>) {
 	app.get("/", ...rootHandler);
@@ -13,4 +14,6 @@ export function configureRoutes(app: Hono<{ Variables: Variables }>) {
 	app.post("/user", ...registerHandler);
 
 	app.get("/video/:id/info", ...videoInfoHandler);
+
+	app.post("/validation/session", ...createValidationSessionHandler)
 }

packages/core/mq/rateLimiter.ts

@@ -1,56 +1,68 @@
-import type { SlidingWindow } from "./slidingWindow.ts";
+import { TokenBucket } from "./tokenBucket.ts";
 
 export interface RateLimiterConfig {
-	window: SlidingWindow;
+    duration: number;
-	max: number;
+    max: number;
 }
 
 export class RateLimiter {
-	private readonly configs: RateLimiterConfig[];
+    private configs: RateLimiterConfig[] = [];
-	private readonly configEventNames: string[];
+    private buckets: TokenBucket[] = [];
+    private identifierFn: (configIndex: number) => string;
 
-	/*
+    /*
-	 * @param name The name of the rate limiter
+     * @param name The name of the rate limiter
-	 * @param configs The configuration of the rate limiter, containing:
+     * @param configs The configuration of the rate limiter, containing:
-	 * - window: The sliding window to use
+     * - tokenBucket: The token bucket instance
-	 * - max: The maximum number of events allowed in the window
+     * - max: The maximum number of tokens allowed per operation
-	 */
+     */
-	constructor(name: string, configs: RateLimiterConfig[]) {
+    constructor(
-		this.configs = configs;
+        name: string,
-		this.configEventNames = configs.map((_, index) => `${name}_config_${index}`);
+        configs: RateLimiterConfig[],
-	}
+        identifierFn?: (configIndex: number) => string
+    ) {
+        this.configs = configs;
+        this.identifierFn = identifierFn || ((index) => `${name}_config_${index}`);
+        for (let i = 0; i < configs.length; i++) {
+            const config = configs[i];
+            const bucket = new TokenBucket({
+                capacity: config.max,
+                rate: config.max / config.duration,
+                identifier: this.identifierFn(i),
+            })
+            this.buckets.push(bucket);
+        }
+    }
 
-	/*
+    /*
-	 * Check if the event has reached the rate limit
+     * Check if the event has reached the rate limit
-	 */
+     */
-	async getAvailability(): Promise<boolean> {
+    async getAvailability(): Promise<boolean> {
-		for (let i = 0; i < this.configs.length; i++) {
+        for (let i = 0; i < this.configs.length; i++) {
-			const config = this.configs[i];
+            const remaining = await this.buckets[i].getRemainingTokens();
-			const eventName = this.configEventNames[i];
-			const count = await config.window.count(eventName);
-			if (count >= config.max) {
-				return false;
-			}
-		}
-		return true;
-	}
 
-	/*
+            if (remaining === null) {
-	 * Trigger an event in the rate limiter
+                return false; // Rate limit exceeded
-	 */
+            }
-	async trigger(): Promise<void> {
+        }
-		for (let i = 0; i < this.configs.length; i++) {
+        return true;
-			const config = this.configs[i];
+    }
-			const eventName = this.configEventNames[i];
-			await config.window.event(eventName);
-		}
-	}
 
-	async clear(): Promise<void> {
+    /*
-		for (let i = 0; i < this.configs.length; i++) {
+     * Trigger an event in the rate limiter
-			const config = this.configs[i];
+     */
-			const eventName = this.configEventNames[i];
+    async trigger(): Promise<void> {
-			await config.window.clear(eventName);
+        for (let i = 0; i < this.configs.length; i++) {
-		}
+            await this.buckets[i].consume(1);
-	}
+        }
+    }
+
+    /*
+     * Clear all buckets for all configurations
+     */
+    async clear(): Promise<void> {
+        for (let i = 0; i < this.configs.length; i++) {
+            await this.buckets[i].reset();
+        }
+    }
 }

packages/core/mq/slidingWindow.ts

@@ -1,51 +0,0 @@
-import type { Redis } from "ioredis";
-
-export class SlidingWindow {
-	private redis: Redis;
-	private readonly windowSize: number;
-
-	/*
-	 * Create a new sliding window
-	 * @param redisClient The Redis client used to store the data
-	 * @param windowSize The size of the window in seconds
-	 */
-	constructor(redisClient: Redis, windowSize: number) {
-		this.redis = redisClient;
-		this.windowSize = windowSize * 1000;
-	}
-
-	/*
-	 * Trigger an event in the sliding window
-	 * @param eventName The name of the event
-	 */
-	async event(eventName: string): Promise<void> {
-		const now = Date.now();
-		const key = `cvsa:sliding_window:${eventName}`;
-
-		const uniqueMember = `${now}-${Math.random()}`;
-		// Add current timestamp to an ordered set
-		await this.redis.zadd(key, now, uniqueMember);
-
-		// Remove timestamps outside the window
-		await this.redis.zremrangebyscore(key, 0, now - this.windowSize);
-	}
-
-	/*
-	 * Count the number of events in the sliding window
-	 * @param eventName The name of the event
-	 */
-	async count(eventName: string): Promise<number> {
-		const key = `cvsa:sliding_window:${eventName}`;
-		const now = Date.now();
-
-		// Remove timestamps outside the window
-		await this.redis.zremrangebyscore(key, 0, now - this.windowSize);
-		// Get the number of timestamps in the window
-		return this.redis.zcard(key);
-	}
-
-	clear(eventName: string): Promise<number> {
-		const key = `cvsa:sliding_window:${eventName}`;
-		return this.redis.del(key);
-	}
-}

packages/core/mq/tokenBucket.ts

@@ -0,0 +1,84 @@
+import { redis } from "@core/db/redis";
+import { SECOND } from "@core/const/time";
+
+interface TokenBucketOptions {
+	capacity: number;
+	rate: number;
+	identifier: string;
+	keyPrefix?: string; 
+}
+
+export class TokenBucket {
+	private readonly capacity: number;
+	private readonly rate: number;
+	private readonly keyPrefix: string;
+	private readonly identifier: string;
+
+	constructor(options: TokenBucketOptions) {
+		if (options.capacity <= 0 || options.rate <= 0) {
+			throw new Error("Capacity and rate must be greater than zero.");
+		}
+
+		this.capacity = options.capacity;
+		this.rate = options.rate;
+		this.identifier = options.identifier;
+		this.keyPrefix = options.keyPrefix || "cvsa:token_bucket:";
+	}
+
+	getKey(): string {
+		return `${this.keyPrefix}${this.identifier}`;
+	}
+
+	/**
+	 * Try to consume a specified number of tokens
+	 * @param count The number of tokens to be consumed
+	 * @returns If consumption is successful, returns the number of remaining tokens; otherwise returns null
+	 */
+	public async consume(count: number): Promise<number | null> {
+		const key = this.getKey();
+		const now = Math.floor(Date.now() / SECOND);
+
+		const script = `
+            local tokens_key = KEYS[1]
+            local last_refilled_key = KEYS[2]
+            local now = tonumber(ARGV[1])
+            local count = tonumber(ARGV[2])
+            local capacity = tonumber(ARGV[3])
+            local rate = tonumber(ARGV[4])
+
+            local last_refilled = tonumber(redis.call('GET', last_refilled_key)) or now
+            local current_tokens = tonumber(redis.call('GET', tokens_key)) or capacity
+
+            local elapsed = now - last_refilled
+            local new_tokens = elapsed * rate
+            current_tokens = math.min(capacity, current_tokens + new_tokens)
+
+            if current_tokens >= count then
+                current_tokens = current_tokens - count
+                redis.call('SET', tokens_key, current_tokens)
+                redis.call('SET', last_refilled_key, now)
+                return current_tokens
+            else
+                return nil
+            end
+        `;
+
+		const keys = [`${key}:tokens`, `${key}:last_refilled`];
+		const args = [now, count, this.capacity, this.rate];
+
+		const result = await redis.eval(script, keys.length, ...keys, ...args);
+
+		return result as number | null;
+	}
+
+	public async getRemainingTokens(): Promise<number> {
+		const key = this.getKey();
+		const tokens = await redis.get(`${key}:tokens`);
+		return Number(tokens) || this.capacity;
+	}
+
+	public async reset(): Promise<void> {
+		const key = this.getKey();
+		await redis.del(`${key}:tokens`, `${key}:last_refilled`);
+	}
+}

packages/core/net/delegate.ts

@@ -1,6 +1,5 @@
 import logger from "@core/log/logger.ts";
 import { RateLimiter, type RateLimiterConfig } from "mq/rateLimiter.ts";
-import { SlidingWindow } from "mq/slidingWindow.ts";
 import { redis } from "db/redis.ts";
 import { ReplyError } from "ioredis";
 import { SECOND } from "../const/time.ts";
@@ -316,37 +315,37 @@ class NetworkDelegate {
 const networkDelegate = new NetworkDelegate();
 const videoInfoRateLimiterConfig: RateLimiterConfig[] = [
 	{
-		window: new SlidingWindow(redis, 0.3),
+		duration: 0.3,
 		max: 1,
 	},
 	{
-		window: new SlidingWindow(redis, 3),
+		duration: 3,
 		max: 5,
 	},
 	{
-		window: new SlidingWindow(redis, 30),
+		duration: 30,
 		max: 30,
 	},
 	{
-		window: new SlidingWindow(redis, 2 * 60),
+		duration: 2 * 60,
 		max: 50,
 	},
 ];
 const biliLimiterConfig: RateLimiterConfig[] = [
 	{
-		window: new SlidingWindow(redis, 1),
+		duration: 1,
 		max: 6,
 	},
 	{
-		window: new SlidingWindow(redis, 5),
+		duration: 5,
 		max: 20,
 	},
 	{
-		window: new SlidingWindow(redis, 30),
+		duration: 30,
 		max: 100,
 	},
 	{
-		window: new SlidingWindow(redis, 5 * 60),
+		duration: 5 * 60,
 		max: 200,
 	},
 ];