From d5983dc12ecaa2e8a9fb193d9d238b9a6550872f Mon Sep 17 00:00:00 2001
From: sugar <sugar@sylveon.social>
Date: Sat, 11 Jan 2025 19:03:07 +0100
Subject: [PATCH] implement rate limit support for mastodon endpoints

---
 .../backend/src/server/api/ApiCallService.ts  | 71 +++++++++++--------
 .../src/server/api/mastodon-endpoints.ts      |  2 +
 2 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/packages/backend/src/server/api/ApiCallService.ts b/packages/backend/src/server/api/ApiCallService.ts
index 9c0421849..f81531f62 100644
--- a/packages/backend/src/server/api/ApiCallService.ts
+++ b/packages/backend/src/server/api/ApiCallService.ts
@@ -198,6 +198,11 @@ export class ApiCallService implements OnApplicationShutdown {
 		const [user] = await this.authenticateService.authenticate(token);
 		let res;
 		try {
+			try {
+				await this.#checkRateLimit(endpoint, user, request);
+			} catch (e) {
+				throw new MastodonApiError('too many requests', 429);
+			}
 			res = await endpoint.exec(body, user, token, null, request.ip, request.headers);
 		} catch (e) {
 			if (e instanceof MastodonApiError) {
@@ -273,36 +278,15 @@ export class ApiCallService implements OnApplicationShutdown {
 			throw new ApiError(accessDenied);
 		}
 
-		const bypassRateLimit = this.config.bypassRateLimit?.some(({ header, value }) => request.headers[header] === value) ?? false;
-		if (ep.meta.limit && !bypassRateLimit) {
-			// koa will automatically load the `X-Forwarded-For` header if `proxy: true` is configured in the app.
-			let limitActor: string;
-			if (user) {
-				limitActor = user.id;
-			} else {
-				limitActor = getIpHash(request.ip);
-			}
-
-			const limit = Object.assign({}, ep.meta.limit);
-
-			if (limit.key == null) {
-				(limit as any).key = ep.name;
-			}
-
-			// TODO: 毎リクエスト計算するのもあれだしキャッシュしたい
-			const factor = user ? (await this.roleService.getUserPolicies(user.id)).rateLimitFactor : 1;
-
-			if (factor > 0) {
-				// Rate limit
-				await this.rateLimiterService.limit(limit as IEndpointMeta['limit'] & { key: NonNullable<string> }, limitActor, factor).catch(err => {
-					throw new ApiError({
-						message: 'Rate limit exceeded. Please try again later.',
-						code: 'RATE_LIMIT_EXCEEDED',
-						id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
-						httpStatusCode: 429,
-					});
-				});
-			}
+		try {
+			await this.#checkRateLimit(ep, user, request);
+		} catch (e) {
+			throw new ApiError({
+				message: 'Rate limit exceeded. Please try again later.',
+				code: 'RATE_LIMIT_EXCEEDED',
+				id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
+				httpStatusCode: 429,
+			});
 		}
 
 		if (ep.meta.requireCredential || ep.meta.requireModerator || ep.meta.requireAdmin) {
@@ -454,6 +438,33 @@ export class ApiCallService implements OnApplicationShutdown {
 		});
 	}
 
+	async #checkRateLimit(ep: IEndpoint | IMastodonEndpoint, user: MiLocalUser | null | undefined, request: FastifyRequest<{ Body: Record<string, unknown> | undefined, Querystring: Record<string, unknown> }>): Promise<void> {
+		const bypassRateLimit = this.config.bypassRateLimit?.some(({ header, value }) => request.headers[header] === value) ?? false;
+		if (ep.meta.limit && !bypassRateLimit) {
+			// koa will automatically load the `X-Forwarded-For` header if `proxy: true` is configured in the app.
+			let limitActor: string;
+			if (user) {
+				limitActor = user.id;
+			} else {
+				limitActor = getIpHash(request.ip);
+			}
+
+			const limit = Object.assign({}, ep.meta.limit);
+
+			if (limit.key == null) {
+				(limit as any).key = ep.name;
+			}
+
+			// TODO: 毎リクエスト計算するのもあれだしキャッシュしたい
+			const factor = user ? (await this.roleService.getUserPolicies(user.id)).rateLimitFactor : 1;
+
+			if (factor > 0) {
+				// Rate limit
+				await this.rateLimiterService.limit(limit as IEndpointMeta['limit'] & { key: NonNullable<string> }, limitActor, factor);
+			}
+		}
+	}
+
 	@bindThis
 	public dispose(): void {
 		clearInterval(this.userIpHistoriesClearIntervalId);
diff --git a/packages/backend/src/server/api/mastodon-endpoints.ts b/packages/backend/src/server/api/mastodon-endpoints.ts
index 6c138c206..60ab5c932 100644
--- a/packages/backend/src/server/api/mastodon-endpoints.ts
+++ b/packages/backend/src/server/api/mastodon-endpoints.ts
@@ -1,5 +1,6 @@
 import { Schema } from '@/misc/json-schema.js';
 import * as mep___accounts_lookup_v1_get from './mastodon/accounts/lookup/v1/get.js';
+import { IEndpointMeta } from './endpoints.js';
 
 const eps = [
 	['GET', 'v1/accounts/lookup', mep___accounts_lookup_v1_get],
@@ -7,6 +8,7 @@ const eps = [
 
 export interface IMastodonEndpointMeta {
 	readonly requireCredential: boolean;
+	readonly limit?: IEndpointMeta['limit'];
 }
 
 export interface IMastodonEndpoint {
-- 
2.45.2