implement SkRateLimiterService with Leaky Bucket rate limiting

1 files changed, 47 insertions, 17 deletions

diff --git a/packages/backend/src/server/api/ApiCallService.ts b/packages/backend/src/server/api/ApiCallService.ts
index 6f51825494..14367e02bb 100644
--- a/packages/backend/src/server/api/ApiCallService.ts
+++ b/packages/backend/src/server/api/ApiCallService.ts
@@ -8,6 +8,7 @@ import * as fs from 'node:fs';
 import * as stream from 'node:stream/promises';
 import { Inject, Injectable } from '@nestjs/common';
 import * as Sentry from '@sentry/node';
+import { LimiterInfo } from 'ratelimiter';
 import { DI } from '@/di-symbols.js';
 import { getIpHash } from '@/misc/get-ip-hash.js';
 import type { MiLocalUser, MiUser } from '@/models/User.js';
@@ -18,6 +19,7 @@ import { createTemp } from '@/misc/create-temp.js';
 import { bindThis } from '@/decorators.js';
 import { RoleService } from '@/core/RoleService.js';
 import type { Config } from '@/config.js';
+import { isLimitInfo } from '@/server/api/SkRateLimiterService.js';
 import { ApiError } from './error.js';
 import { RateLimiterService } from './RateLimiterService.js';
 import { ApiLoggerService } from './ApiLoggerService.js';
@@ -68,12 +70,17 @@ export class ApiCallService implements OnApplicationShutdown {
 		} else if (err.code === 'RATE_LIMIT_EXCEEDED') {
 			const info: unknown = err.info;
 			const unixEpochInSeconds = Date.now();
-			if (typeof(info) === 'object' && info && 'resetMs' in info && typeof(info.resetMs) === 'number') {
+			if (isLimitInfo(info)) {
+				// Number of seconds to wait before trying again. Left for backwards compatibility.
+				reply.header('Retry-After', info.resetSec.toString());
+				// Number of milliseconds to wait before trying again.
+				reply.header('X-RateLimit-Reset', info.resetMs.toString());
+			} else if (typeof(info) === 'object' && info && 'resetMs' in info && typeof(info.resetMs) === 'number') {
 				const cooldownInSeconds = Math.ceil((info.resetMs - unixEpochInSeconds) / 1000);
 				// もしかするとマイナスになる可能性がなくはないのでマイナスだったら0にしておく
 				reply.header('Retry-After', Math.max(cooldownInSeconds, 0).toString(10));
 			} else {
-				this.logger.warn(`rate limit information has unexpected type ${typeof(err.info?.reset)}`);
+				this.logger.warn(`rate limit information has unexpected type: ${JSON.stringify(info)}`);
 			}
 		} else if (err.kind === 'client') {
 			reply.header('WWW-Authenticate', `Bearer realm="Misskey", error="invalid_request", error_description="${err.message}"`);
@@ -168,7 +175,7 @@ export class ApiCallService implements OnApplicationShutdown {
 			return;
 		}
 		this.authenticateService.authenticate(token).then(([user, app]) => {
-			this.call(endpoint, user, app, body, null, request).then((res) => {
+			this.call(endpoint, user, app, body, null, request, reply).then((res) => {
 				if (request.method === 'GET' && endpoint.meta.cacheSec && !token && !user) {
 					reply.header('Cache-Control', `public, max-age=${endpoint.meta.cacheSec}`);
 				}
@@ -229,7 +236,7 @@ export class ApiCallService implements OnApplicationShutdown {
 			this.call(endpoint, user, app, fields, {
 				name: multipartData.filename,
 				path: path,
-			}, request).then((res) => {
+			}, request, reply).then((res) => {
 				this.send(reply, res);
 			}).catch((err: ApiError) => {
 				this.#sendApiError(reply, err);
@@ -304,6 +311,7 @@ export class ApiCallService implements OnApplicationShutdown {
 			path: string;
 		} | null,
 		request: FastifyRequest<{ Body: Record<string, unknown> | undefined, Querystring: Record<string, unknown> }>,
+		reply: FastifyReply,
 	) {
 		const isSecure = user != null && token == null;
 
@@ -339,19 +347,41 @@ export class ApiCallService implements OnApplicationShutdown {
 
 			if (factor > 0) {
 				// Rate limit
-				await this.rateLimiterService.limit(limit as IEndpointMeta['limit'] & { key: NonNullable<string> }, limitActor, factor).catch(err => {
-					if ('info' in err) {
-						// errはLimiter.LimiterInfoであることが期待される
-						throw new ApiError({
-							message: 'Rate limit exceeded. Please try again later.',
-							code: 'RATE_LIMIT_EXCEEDED',
-							id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
-							httpStatusCode: 429,
-						}, err.info);
-					} else {
-						throw new TypeError('information must be a rate-limiter information.');
-					}
-				});
+				const info = await this.rateLimiterService.limit(limit as IEndpointMeta['limit'] & { key: NonNullable<string> }, limitActor, factor)
+					.then(info => {
+						// We always want these headers, because clients need them for pacing.
+						// Conditional check in case we somehow revert to the old limiter, which does not return info.
+						if (info) {
+							// Number of seconds until the limit has fully reset.
+							reply.header('X-RateLimit-Clear', info.fullResetSec.toString());
+							// Number of calls that can be made before being limited.
+							reply.header('X-RateLimit-Remaining', info.remaining.toString());
+
+							// Only forward the info object if it's blocked, otherwise we'll reject *all* requests
+							if (info.blocked) {
+								return info;
+							}
+						}
+
+						return undefined;
+					})
+					.catch(err => {
+						// The old limiter throws info instead of returning it.
+						if ('info' in err) {
+							return err.info as LimiterInfo;
+						} else {
+							throw err;
+						}
+					});
+
+				if (info) {
+					throw new ApiError({
+						message: 'Rate limit exceeded. Please try again later.',
+						code: 'RATE_LIMIT_EXCEEDED',
+						id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
+						httpStatusCode: 429,
+					}, info);
+				}
 			}
 		}