lookup and cache rate limit factors directly within SkRateLimiterService

3 files changed, 53 insertions, 45 deletions

diff --git a/packages/backend/src/server/api/ApiCallService.ts b/packages/backend/src/server/api/ApiCallService.ts
index 03f25a51fe..9c3952d541 100644
--- a/packages/backend/src/server/api/ApiCallService.ts
+++ b/packages/backend/src/server/api/ApiCallService.ts
@@ -313,35 +313,30 @@ export class ApiCallService implements OnApplicationShutdown {
 		// eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
 		if (endpointLimit) {
 			// koa will automatically load the `X-Forwarded-For` header if `proxy: true` is configured in the app.
-			let limitActor: string;
+			let limitActor: string | MiLocalUser;
 			if (user) {
-				limitActor = user.id;
+				limitActor = user;
 			} else {
 				limitActor = getIpHash(request.ip);
 			}
 
-			// TODO: 毎リクエスト計算するのもあれだしキャッシュしたい
-			const factor = user ? (await this.roleService.getUserPolicies(user.id)).rateLimitFactor : 1;
+			const limit = {
+				key: ep.name,
+				...endpointLimit,
+			};
 
-			if (factor > 0) {
-				const limit = {
-					key: ep.name,
-					...endpointLimit,
-				};
+			// Rate limit
+			const info = await this.rateLimiterService.limit(limit, limitActor);
 
-				// Rate limit
-				const info = await this.rateLimiterService.limit(limit, limitActor, factor);
+			sendRateLimitHeaders(reply, info);
 
-				sendRateLimitHeaders(reply, info);
-
-				if (info.blocked) {
-					throw new ApiError({
-						message: 'Rate limit exceeded. Please try again later.',
-						code: 'RATE_LIMIT_EXCEEDED',
-						id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
-						httpStatusCode: 429,
-					}, info);
-				}
+			if (info.blocked) {
+				throw new ApiError({
+					message: 'Rate limit exceeded. Please try again later.',
+					code: 'RATE_LIMIT_EXCEEDED',
+					id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
+					httpStatusCode: 429,
+				}, info);
 			}
 		}
 
diff --git a/packages/backend/src/server/api/SkRateLimiterService.ts b/packages/backend/src/server/api/SkRateLimiterService.ts
index 38c97b63df..70103222f3 100644
--- a/packages/backend/src/server/api/SkRateLimiterService.ts
+++ b/packages/backend/src/server/api/SkRateLimiterService.ts
@@ -5,36 +5,59 @@
 
 import { Inject, Injectable } from '@nestjs/common';
 import Redis from 'ioredis';
-import { TimeService } from '@/core/TimeService.js';
-import { EnvService } from '@/core/EnvService.js';
+import type { TimeService } from '@/core/TimeService.js';
+import type { EnvService } from '@/core/EnvService.js';
 import { BucketRateLimit, LegacyRateLimit, LimitInfo, RateLimit, hasMinLimit, isLegacyRateLimit, Keyed, hasMaxLimit, disabledLimitInfo, MaxLegacyLimit, MinLegacyLimit } from '@/misc/rate-limit-utils.js';
 import { DI } from '@/di-symbols.js';
+import { MemoryKVCache } from '@/misc/cache.js';
+import type { MiUser } from '@/models/_.js';
+import type { RoleService } from '@/core/RoleService.js';
+
+// Sentinel value used for caching the default role template.
+// Required because MemoryKVCache doesn't support null keys.
+const defaultUserKey = '';
 
 @Injectable()
 export class SkRateLimiterService {
+	// 1-minute cache interval
+	private readonly factorCache = new MemoryKVCache<number>(1000 * 60);
 	private readonly disabled: boolean;
 
 	constructor(
-		@Inject(TimeService)
+		@Inject('TimeService')
 		private readonly timeService: TimeService,
 
 		@Inject(DI.redis)
 		private readonly redisClient: Redis.Redis,
 
-		@Inject(EnvService)
+		@Inject('RoleService')
+		private readonly roleService: RoleService,
+
+		@Inject('EnvService')
 		envService: EnvService,
 	) {
 		this.disabled = envService.env.NODE_ENV === 'test';
 	}
 
 	/**
-	 * Check & increment a rate limit
+	 * Check & increment a rate limit for a client
 	 * @param limit The limit definition
-	 * @param actor Client who is calling this limit
-	 * @param factor Scaling factor - smaller = larger limit (less restrictive)
+	 * @param actorOrUser authenticated client user or IP hash
 	 */
-	public async limit(limit: Keyed<RateLimit>, actor: string, factor = 1): Promise<LimitInfo> {
-		if (this.disabled || factor === 0) {
+	public async limit(limit: Keyed<RateLimit>, actorOrUser: string | MiUser): Promise<LimitInfo> {
+		if (this.disabled) {
+			return disabledLimitInfo;
+		}
+
+		const actor = typeof(actorOrUser) === 'object' ? actorOrUser.id : actorOrUser;
+		const userCacheKey = typeof(actorOrUser) === 'object' ? actorOrUser.id : defaultUserKey;
+		const userRoleKey = typeof(actorOrUser) === 'object' ? actorOrUser.id : null;
+		const factor = this.factorCache.get(userCacheKey) ?? await this.factorCache.fetch(userCacheKey, async () => {
+			const role = await this.roleService.getUserPolicies(userRoleKey);
+			return role.rateLimitFactor;
+		});
+
+		if (factor === 0) {
 			return disabledLimitInfo;
 		}
 
@@ -42,10 +65,6 @@ export class SkRateLimiterService {
 			throw new Error(`Rate limit factor is zero or negative: ${factor}`);
 		}
 
-		return await this.tryLimit(limit, actor, factor);
-	}
-
-	private async tryLimit(limit: Keyed<RateLimit>, actor: string, factor: number): Promise<LimitInfo> {
 		if (isLegacyRateLimit(limit)) {
 			return await this.limitLegacy(limit, actor, factor);
 		} else {
diff --git a/packages/backend/src/server/api/StreamingApiServerService.ts b/packages/backend/src/server/api/StreamingApiServerService.ts
index e3fd1312ae..f30bbb928b 100644
--- a/packages/backend/src/server/api/StreamingApiServerService.ts
+++ b/packages/backend/src/server/api/StreamingApiServerService.ts
@@ -18,7 +18,6 @@ import { CacheService } from '@/core/CacheService.js';
 import { MiLocalUser } from '@/models/User.js';
 import { UserService } from '@/core/UserService.js';
 import { ChannelFollowingService } from '@/core/ChannelFollowingService.js';
-import { RoleService } from '@/core/RoleService.js';
 import { getIpHash } from '@/misc/get-ip-hash.js';
 import { LoggerService } from '@/core/LoggerService.js';
 import { SkRateLimiterService } from '@/server/api/SkRateLimiterService.js';
@@ -49,7 +48,6 @@ export class StreamingApiServerService {
 		private usersService: UserService,
 		private channelFollowingService: ChannelFollowingService,
 		private rateLimiterService: SkRateLimiterService,
-		private roleService: RoleService,
 		private loggerService: LoggerService,
 	) {
 	}
@@ -57,22 +55,18 @@ export class StreamingApiServerService {
 	@bindThis
 	private async rateLimitThis(
 		user: MiLocalUser | null | undefined,
-		requestIp: string | undefined,
+		requestIp: string,
 		limit: IEndpointMeta['limit'] & { key: NonNullable<string> },
 	) : Promise<boolean> {
-		let limitActor: string;
+		let limitActor: string | MiLocalUser;
 		if (user) {
-			limitActor = user.id;
+			limitActor = user;
 		} else {
-			limitActor = getIpHash(requestIp || 'wtf');
+			limitActor = getIpHash(requestIp);
 		}
 
-		const factor = user ? (await this.roleService.getUserPolicies(user.id)).rateLimitFactor : 1;
-
-		if (factor <= 0) return false;
-
 		// Rate limit
-		const rateLimit = await this.rateLimiterService.limit(limit, limitActor, factor);
+		const rateLimit = await this.rateLimiterService.limit(limit, limitActor);
 		return rateLimit.blocked;
 	}