merge: Prevent streaming API denial-of-service (resolves #1019) (!951)

View MR for information: https://activitypub.software/TransFem-org/Sharkey/-/merge_requests/951

Closes #1019

Approved-by: dakkar <dakkar@thenautilus.net>
Approved-by: Marie <github@yuugi.dev>

2 files changed, 122 insertions, 83 deletions

diff --git a/packages/backend/src/server/api/StreamingApiServerService.ts b/packages/backend/src/server/api/StreamingApiServerService.ts
index 0954744f81..d86deef1d7 100644
--- a/packages/backend/src/server/api/StreamingApiServerService.ts
+++ b/packages/backend/src/server/api/StreamingApiServerService.ts
@@ -10,7 +10,9 @@ import * as WebSocket from 'ws';
 import proxyAddr from 'proxy-addr';
 import ms from 'ms';
 import { DI } from '@/di-symbols.js';
-import type { UsersRepository, MiAccessToken } from '@/models/_.js';
+import type { UsersRepository, MiAccessToken, MiUser } from '@/models/_.js';
+import type { Config } from '@/config.js';
+import type { Keyed, RateLimit } from '@/misc/rate-limit-utils.js';
 import { NoteReadService } from '@/core/NoteReadService.js';
 import { NotificationService } from '@/core/NotificationService.js';
 import { bindThis } from '@/decorators.js';
@@ -25,13 +27,16 @@ import { AuthenticateService, AuthenticationError } from './AuthenticateService.
 import MainStreamConnection from './stream/Connection.js';
 import { ChannelsService } from './stream/ChannelsService.js';
 import type * as http from 'node:http';
-import type { IEndpointMeta } from './endpoints.js';
-import type { Config } from "@/config.js";
+
+// Maximum number of simultaneous connections by client (user ID or IP address).
+// Excess connections will be closed automatically.
+const MAX_CONNECTIONS_PER_CLIENT = 32;
 
 @Injectable()
 export class StreamingApiServerService {
 	#wss: WebSocket.WebSocketServer;
 	#connections = new Map<WebSocket.WebSocket, number>();
+	#connectionsByClient = new Map<string, Set<WebSocket.WebSocket>>(); // key: IP / user ID -> value: connection
 	#cleanConnectionsIntervalId: NodeJS.Timeout | null = null;
 
 	constructor(
@@ -58,17 +63,9 @@ export class StreamingApiServerService {
 
 	@bindThis
 	private async rateLimitThis(
-		user: MiLocalUser | null | undefined,
-		requestIp: string,
-		limit: IEndpointMeta['limit'] & { key: NonNullable<string> },
+		limitActor: MiUser | string,
+		limit: Keyed<RateLimit>,
 	) : Promise<boolean> {
-		let limitActor: string | MiLocalUser;
-		if (user) {
-			limitActor = user;
-		} else {
-			limitActor = getIpHash(requestIp);
-		}
-
 		// Rate limit
 		const rateLimit = await this.rateLimiterService.limit(limit, limitActor);
 		return rateLimit.blocked;
@@ -88,21 +85,6 @@ export class StreamingApiServerService {
 				return;
 			}
 
-			// ServerServices sets `trustProxy: true`, which inside
-			// fastify/request.js ends up calling `proxyAddr` in this way,
-			// so we do the same
-			const requestIp = proxyAddr(request, () => { return true; } );
-
-			if (await this.rateLimitThis(null, requestIp, {
-				key: 'wsconnect',
-				duration: ms('5min'),
-				max: 32,
-			})) {
-				socket.write('HTTP/1.1 429 Rate Limit Exceeded\r\n\r\n');
-				socket.destroy();
-				return;
-			}
-
 			const q = new URL(request.url, `http://${request.headers.host}`).searchParams;
 
 			let user: MiLocalUser | null = null;
@@ -140,15 +122,48 @@ export class StreamingApiServerService {
 				return;
 			}
 
+			// ServerServices sets `trustProxy: true`, which inside fastify/request.js ends up calling `proxyAddr` in this way, so we do the same.
+			const requestIp = proxyAddr(request, () => true );
+			const limitActor = user?.id ?? getIpHash(requestIp);
+			if (await this.rateLimitThis(limitActor, {
+				key: 'wsconnect',
+				duration: ms('5min'),
+				max: 32,
+			})) {
+				socket.write('HTTP/1.1 429 Rate Limit Exceeded\r\n\r\n');
+				socket.destroy();
+				return;
+			}
+
+			// For performance and code simplicity, obtain and hold this reference for the lifetime of the connection.
+			// This should be safe because the map entry should only be deleted after *all* connections close.
+			let connectionsForClient = this.#connectionsByClient.get(limitActor);
+			if (!connectionsForClient) {
+				connectionsForClient = new Set();
+				this.#connectionsByClient.set(limitActor, connectionsForClient);
+			}
+
+			// Close excess connections
+			while (connectionsForClient.size >= MAX_CONNECTIONS_PER_CLIENT) {
+				// Set maintains insertion order, so first entry is the oldest.
+				// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+				const oldestConnection = connectionsForClient.values().next().value!;
+
+				// Technically, the close() handler should remove this entry.
+				// But if that ever fails, then we could enter an infinite loop.
+				// We manually remove the connection here just in case.
+				oldestConnection.close(1008, 'Disconnected - too many simultaneous connections');
+				connectionsForClient.delete(oldestConnection);
+			}
+
 			const rateLimiter = () => {
-				// rather high limit, because when catching up at the top of a
-				// timeline, the frontend may render many many notes, each of
-				// which causes a message via `useNoteCapture` to ask for
-				// realtime updates of that note
-				return this.rateLimitThis(user, requestIp, {
+				// Rather high limit because when catching up at the top of a timeline, the frontend may render many many notes.
+				// Each of which causes a message via `useNoteCapture` to ask for realtime updates of that note.
+				return this.rateLimitThis(limitActor, {
+					type: 'bucket',
 					key: 'wsmessage',
-					duration: ms('2sec'),
-					max: 4096,
+					size: 4096, // Allow spikes of up to 4096
+					dripRate: 50, // Then once every 50ms (20/second rate)
 				});
 			};
 
@@ -166,6 +181,19 @@ export class StreamingApiServerService {
 			await stream.init();
 
 			this.#wss.handleUpgrade(request, socket, head, (ws) => {
+				connectionsForClient.add(ws);
+
+				// Call before emit() in case it throws an error.
+				// We don't want to leave dangling references!
+				ws.once('close', () => {
+					connectionsForClient.delete(ws);
+
+					// Make sure we don't leak the Set objects!
+					if (connectionsForClient.size < 1) {
+						this.#connectionsByClient.delete(limitActor);
+					}
+				});
+
 				this.#wss.emit('connection', ws, request, {
 					stream, user, app,
 				});
diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts
index e98e2a2f3f..a69c1b8b7f 100644
--- a/packages/backend/src/server/api/stream/Connection.ts
+++ b/packages/backend/src/server/api/stream/Connection.ts
@@ -23,6 +23,8 @@ import type { EventEmitter } from 'events';
 import type Channel from './channel.js';
 
 const MAX_CHANNELS_PER_CONNECTION = 32;
+const MAX_SUBSCRIPTIONS_PER_CONNECTION = 512;
+const MAX_CACHED_NOTES_PER_CONNECTION = 64;
 
 /**
  * Main stream connection
@@ -31,12 +33,11 @@ const MAX_CHANNELS_PER_CONNECTION = 32;
 export default class Connection {
 	public user?: MiUser;
 	public token?: MiAccessToken;
-	private rateLimiter?: () => Promise<boolean>;
 	private wsConnection: WebSocket.WebSocket;
 	public subscriber: StreamEventEmitter;
-	private channels: Channel[] = [];
-	private subscribingNotes: Partial<Record<string, number>> = {};
-	private cachedNotes: Packed<'Note'>[] = [];
+	private channels = new Map<string, Channel>();
+	private subscribingNotes = new Map<string, number>();
+	private cachedNotes = new Map<string, Packed<'Note'>>();
 	public userProfile: MiUserProfile | null = null;
 	public following: Record<string, Pick<MiFollowing, 'withReplies'> | undefined> = {};
 	public followingChannels: Set<string> = new Set();
@@ -45,7 +46,6 @@ export default class Connection {
 	public userIdsWhoMeMutingRenotes: Set<string> = new Set();
 	public userMutedInstances: Set<string> = new Set();
 	private fetchIntervalId: NodeJS.Timeout | null = null;
-	private activeRateLimitRequests = 0;
 	private closingConnection = false;
 	private logger: Logger;
 
@@ -60,11 +60,10 @@ export default class Connection {
 		user: MiUser | null | undefined,
 		token: MiAccessToken | null | undefined,
 		private ip: string,
-		rateLimiter: () => Promise<boolean>,
+		private readonly rateLimiter: () => Promise<boolean>,
 	) {
 		if (user) this.user = user;
 		if (token) this.token = token;
-		if (rateLimiter) this.rateLimiter = rateLimiter;
 
 		this.logger = loggerService.getLogger('streaming', 'coral');
 	}
@@ -121,25 +120,13 @@ export default class Connection {
 
 		if (this.closingConnection) return;
 
-		if (this.rateLimiter) {
-			// this 4096 should match the `max` of the `rateLimiter`, see
-			// StreamingApiServerService
-			if (this.activeRateLimitRequests <= 4096) {
-				this.activeRateLimitRequests++;
-				const shouldRateLimit = await this.rateLimiter();
-				this.activeRateLimitRequests--;
+		// The rate limit is very high, so we can safely disconnect any client that hits it.
+		if (await this.rateLimiter()) {
+			this.logger.warn(`Closing a connection from ${this.ip} (user=${this.user?.id}}) due to an excessive influx of messages.`);
 
-				if (shouldRateLimit) return;
-				if (this.closingConnection) return;
-			} else {
-				let connectionInfo = `IP ${this.ip}`;
-				if (this.user) connectionInfo += `, user ID ${this.user.id}`;
-
-				this.logger.warn(`Closing a connection (${connectionInfo}) due to an excessive influx of messages.`);
-				this.closingConnection = true;
-				this.wsConnection.close(1008, 'Please stop spamming the streaming API.');
-				return;
-			}
+			this.closingConnection = true;
+			this.wsConnection.close(1008, 'Disconnected - too many requests');
+			return;
 		}
 
 		try {
@@ -172,15 +159,13 @@ export default class Connection {
 	@bindThis
 	public cacheNote(note: Packed<'Note'>) {
 		const add = (note: Packed<'Note'>) => {
-			const existIndex = this.cachedNotes.findIndex(n => n.id === note.id);
-			if (existIndex > -1) {
-				this.cachedNotes[existIndex] = note;
-				return;
-			}
+			this.cachedNotes.set(note.id, note);
 
-			this.cachedNotes.unshift(note);
-			if (this.cachedNotes.length > 32) {
-				this.cachedNotes.splice(32);
+			while (this.cachedNotes.size > MAX_CACHED_NOTES_PER_CONNECTION) {
+				// Map maintains insertion order, so first key is always the oldest
+				// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+				const oldestKey = this.cachedNotes.keys().next().value!;
+				this.cachedNotes.delete(oldestKey);
 			}
 		};
 
@@ -192,9 +177,9 @@ export default class Connection {
 	@bindThis
 	private readNote(body: JsonValue | undefined) {
 		if (!isJsonObject(body)) return;
-		const id = body.id;
+		const id = body.id as string;
 
-		const note = this.cachedNotes.find(n => n.id === id);
+		const note = this.cachedNotes.get(id);
 		if (note == null) return;
 
 		if (this.user && (note.userId !== this.user.id)) {
@@ -215,9 +200,19 @@ export default class Connection {
 		if (!isJsonObject(payload)) return;
 		if (!payload.id || typeof payload.id !== 'string') return;
 
-		const current = this.subscribingNotes[payload.id] ?? 0;
+		const current = this.subscribingNotes.get(payload.id) ?? 0;
 		const updated = current + 1;
-		this.subscribingNotes[payload.id] = updated;
+		this.subscribingNotes.set(payload.id, updated);
+
+		// Limit the number of distinct notes that can be subscribed to.
+		while (this.subscribingNotes.size > MAX_SUBSCRIPTIONS_PER_CONNECTION) {
+			// Map maintains insertion order, so first key is always the oldest
+			// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+			const oldestKey = this.subscribingNotes.keys().next().value!;
+
+			this.subscribingNotes.delete(oldestKey);
+			this.subscriber.off(`noteStream:${oldestKey}`, this.onNoteStreamMessage);
+		}
 
 		if (updated === 1) {
 			this.subscriber.on(`noteStream:${payload.id}`, this.onNoteStreamMessage);
@@ -232,12 +227,12 @@ export default class Connection {
 		if (!isJsonObject(payload)) return;
 		if (!payload.id || typeof payload.id !== 'string') return;
 
-		const current = this.subscribingNotes[payload.id];
+		const current = this.subscribingNotes.get(payload.id);
 		if (current == null) return;
 		const updated = current - 1;
-		this.subscribingNotes[payload.id] = updated;
+		this.subscribingNotes.set(payload.id, updated);
 		if (updated <= 0) {
-			delete this.subscribingNotes[payload.id];
+			this.subscribingNotes.delete(payload.id);
 			this.subscriber.off(`noteStream:${payload.id}`, this.onNoteStreamMessage);
 		}
 	}
@@ -304,7 +299,11 @@ export default class Connection {
 	 */
 	@bindThis
 	public connectChannel(id: string, params: JsonObject | undefined, channel: string, pong = false) {
-		if (this.channels.length >= MAX_CHANNELS_PER_CONNECTION) {
+		if (this.channels.has(id)) {
+			this.disconnectChannel(id);
+		}
+
+		if (this.channels.size >= MAX_CHANNELS_PER_CONNECTION) {
 			return;
 		}
 
@@ -320,12 +319,16 @@ export default class Connection {
 		}
 
 		// 共有可能チャンネルに接続しようとしていて、かつそのチャンネルに既に接続していたら無意味なので無視
-		if (channelService.shouldShare && this.channels.some(c => c.chName === channel)) {
-			return;
+		if (channelService.shouldShare) {
+			for (const c of this.channels.values()) {
+				if (c.chName === channel) {
+					return;
+				}
+			}
 		}
 
 		const ch: Channel = channelService.create(id, this);
-		this.channels.push(ch);
+		this.channels.set(ch.id, ch);
 		ch.init(params ?? {});
 
 		if (pong) {
@@ -341,11 +344,11 @@ export default class Connection {
 	 */
 	@bindThis
 	public disconnectChannel(id: string) {
-		const channel = this.channels.find(c => c.id === id);
+		const channel = this.channels.get(id);
 
 		if (channel) {
 			if (channel.dispose) channel.dispose();
-			this.channels = this.channels.filter(c => c.id !== id);
+			this.channels.delete(id);
 		}
 	}
 
@@ -360,7 +363,7 @@ export default class Connection {
 		if (typeof data.type !== 'string') return;
 		if (typeof data.body === 'undefined') return;
 
-		const channel = this.channels.find(c => c.id === data.id);
+		const channel = this.channels.get(data.id);
 		if (channel != null && channel.onMessage != null) {
 			channel.onMessage(data.type, data.body);
 		}
@@ -372,8 +375,16 @@ export default class Connection {
 	@bindThis
 	public dispose() {
 		if (this.fetchIntervalId) clearInterval(this.fetchIntervalId);
-		for (const c of this.channels.filter(c => c.dispose)) {
+		for (const c of this.channels.values()) {
 			if (c.dispose) c.dispose();
 		}
+		for (const k of this.subscribingNotes.keys()) {
+			this.subscriber.off(`noteStream:${k}`, this.onNoteStreamMessage);
+		}
+
+		this.fetchIntervalId = null;
+		this.channels.clear();
+		this.subscribingNotes.clear();
+		this.cachedNotes.clear();
 	}
 }