spec(aiscript): Mk:apiの呼び出しにエンドポイントごとのレートリミットを設定 (MisskeyIO#522)

packages/frontend/package.json

@@ -19,6 +19,7 @@
 	"dependencies": {
 		"@discordapp/twemoji": "15.0.2",
 		"@github/webauthn-json": "2.1.1",
+		"@isaacs/ttlcache": "1.4.1",
 		"@mcaptcha/vanilla-glue": "0.1.0-alpha-3",
 		"@misskey-dev/browser-image-resizer": "2024.1.0",
 		"@rollup/plugin-json": "6.1.0",

packages/frontend/src/scripts/aiscript/api.ts

@@ -11,6 +11,7 @@ import { miLocalStorage } from '@/local-storage.js';
 import { customEmojis } from '@/custom-emojis.js';
 import { url, lang } from '@/config.js';
 import { nyaize } from '@/scripts/nyaize.js';
+import { RateLimiter } from '@/scripts/rate-limiter.js';
 
 export function aiScriptReadline(q: string): Promise<string> {
 	return new Promise(ok => {
@@ -23,6 +24,8 @@ export function aiScriptReadline(q: string): Promise<string> {
 }
 
 export function createAiScriptEnv(opts) {
+	const rateLimiter = new RateLimiter<string>({ duration: 1000 * 15, max: 30 });
+
 	return {
 		USER_ID: $i ? values.STR($i.id) : values.NULL,
 		USER_NAME: $i ? values.STR($i.name) : values.NULL,
@@ -55,6 +58,7 @@ export function createAiScriptEnv(opts) {
 				if (typeof token.value !== 'string') throw new Error('invalid token');
 			}
 			const actualToken: string|null = token?.value ?? opts.token ?? null;
+			if (!rateLimiter.hit(ep.value)) return values.ERROR('rate_limited', values.NULL);
 			return misskeyApi(ep.value, utils.valToJs(param), actualToken).then(res => {
 				return utils.jsToVal(res);
 			}, err => {

packages/frontend/src/scripts/rate-limiter.ts

@@ -0,0 +1,71 @@
+/*
+ * SPDX-FileCopyrightText: Isaac Z. Schlueter and Contributors of https://github.com/isaacs/ttlcache
+ * SPDX-License-Identifier: ISC
+ *
+ * This file is derived from the project that has licensed under the ISC license.
+ * This file SHOULD NOT be considered as a part of the this project that has licensed under AGPL-3.0-only
+ * Adapted from https://github.com/isaacs/ttlcache/blob/b6002f971e122e3b35e23d00ac6a8365d505c14d/examples/rate-limiter-window.ts
+ */
+
+import TTLCache from '@isaacs/ttlcache';
+import type { Options as TTLCacheOptions } from '@isaacs/ttlcache';
+
+export interface Options {
+	duration: number;
+	max: number;
+}
+
+interface RLEntryOptions extends TTLCacheOptions<number, boolean> {
+	onEmpty: () => void;
+}
+
+class RLEntry extends TTLCache<number, boolean> {
+	onEmpty: () => void;
+
+	constructor(options: RLEntryOptions) {
+		super(options);
+		this.onEmpty = options.onEmpty;
+	}
+
+	purgeStale() {
+		const ret = super.purgeStale();
+		if (this.size === 0 && ret) {
+			this.onEmpty();
+		}
+		return ret;
+	}
+}
+
+export class RateLimiter<K> extends Map<K, TTLCache<number, boolean>> {
+	duration: number;
+	max: number;
+
+	constructor(options: Options) {
+		super();
+		this.duration = options.duration;
+		this.max = options.max;
+	}
+
+	hit(key: K) {
+		const c =
+			super.get(key) ??
+			new RLEntry({
+				ttl: this.duration,
+				onEmpty: () => this.delete(key),
+			});
+
+		this.set(key, c);
+
+		if (c.size > this.max) {
+			// rejected, too many hits within window
+			return false;
+		}
+		c.set(performance.now(), true);
+		return true;
+	}
+
+	count(key: K) {
+		const c = super.get(key);
+		return c ? c.size : 0;
+	}
+}