From bf1c9b67d63e1d9263be3f6e6d7184606c992696 Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Thu, 27 Mar 2025 10:55:46 -0400 Subject: close websocket when rate limit exceeded --- .../backend/src/server/api/stream/Connection.ts | 31 ++++++---------------- 1 file changed, 8 insertions(+), 23 deletions(-) (limited to 'packages/backend/src/server/api/stream') diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts index e98e2a2f3f..9ca209f08b 100644 --- a/packages/backend/src/server/api/stream/Connection.ts +++ b/packages/backend/src/server/api/stream/Connection.ts @@ -31,7 +31,6 @@ const MAX_CHANNELS_PER_CONNECTION = 32; export default class Connection { public user?: MiUser; public token?: MiAccessToken; - private rateLimiter?: () => Promise; private wsConnection: WebSocket.WebSocket; public subscriber: StreamEventEmitter; private channels: Channel[] = []; @@ -45,7 +44,6 @@ export default class Connection { public userIdsWhoMeMutingRenotes: Set = new Set(); public userMutedInstances: Set = new Set(); private fetchIntervalId: NodeJS.Timeout | null = null; - private activeRateLimitRequests = 0; private closingConnection = false; private logger: Logger; @@ -60,11 +58,10 @@ export default class Connection { user: MiUser | null | undefined, token: MiAccessToken | null | undefined, private ip: string, - rateLimiter: () => Promise, + private readonly rateLimiter: () => Promise, ) { if (user) this.user = user; if (token) this.token = token; - if (rateLimiter) this.rateLimiter = rateLimiter; this.logger = loggerService.getLogger('streaming', 'coral'); } @@ -121,25 +118,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--; - - 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; - } + // 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.`); + + this.closingConnection = true; + this.wsConnection.close(1008, 'Disconnected - too many requests'); + return; } try { -- cgit v1.2.3-freya From 831329499dd6b29638a6390a860cded30c448215 Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Thu, 27 Mar 2025 11:07:26 -0400 Subject: limit the number of note subscriptions per connection --- .../backend/src/server/api/stream/Connection.ts | 25 ++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) (limited to 'packages/backend/src/server/api/stream') diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts index 9ca209f08b..36e5769216 100644 --- a/packages/backend/src/server/api/stream/Connection.ts +++ b/packages/backend/src/server/api/stream/Connection.ts @@ -23,6 +23,7 @@ import type { EventEmitter } from 'events'; import type Channel from './channel.js'; const MAX_CHANNELS_PER_CONNECTION = 32; +const MAX_SUBSCRIPTIONS_PER_CONNECTION = 256; /** * Main stream connection @@ -34,7 +35,7 @@ export default class Connection { private wsConnection: WebSocket.WebSocket; public subscriber: StreamEventEmitter; private channels: Channel[] = []; - private subscribingNotes: Partial> = {}; + private subscribingNotes = new Map(); private cachedNotes: Packed<'Note'>[] = []; public userProfile: MiUserProfile | null = null; public following: Record | undefined> = {}; @@ -200,9 +201,21 @@ 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; + + // Limit the number of distinct notes that can be subscribed to. + // If current is-zero, then this is a new note and we need to check the limit + if (current === 0 && 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); + } + const updated = current + 1; - this.subscribingNotes[payload.id] = updated; + this.subscribingNotes.set(payload.id, updated); if (updated === 1) { this.subscriber.on(`noteStream:${payload.id}`, this.onNoteStreamMessage); @@ -217,12 +230,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); } } -- cgit v1.2.3-freya From b8fd9d0bc04656d30628e5f7a3c7c74632014154 Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Thu, 27 Mar 2025 12:17:23 -0400 Subject: clear subscriptions when connection closes --- packages/backend/src/server/api/stream/Connection.ts | 7 +++++++ 1 file changed, 7 insertions(+) (limited to 'packages/backend/src/server/api/stream') diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts index 36e5769216..691ce54feb 100644 --- a/packages/backend/src/server/api/stream/Connection.ts +++ b/packages/backend/src/server/api/stream/Connection.ts @@ -373,5 +373,12 @@ export default class Connection { for (const c of this.channels.filter(c => c.dispose)) { if (c.dispose) c.dispose(); } + for (const k of this.subscribingNotes.keys()) { + this.subscriber.off(`noteStream:${k}`, this.onNoteStreamMessage); + } + + this.fetchIntervalId = null; + this.channels = []; + this.subscribingNotes.clear(); } } -- cgit v1.2.3-freya From 045ff5d2c0037fe0770d619b9f7a21d79a9b668d Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Thu, 27 Mar 2025 12:23:14 -0400 Subject: make sure that note subscriptions can't stay above limit --- packages/backend/src/server/api/stream/Connection.ts | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) (limited to 'packages/backend/src/server/api/stream') diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts index 691ce54feb..96b968d890 100644 --- a/packages/backend/src/server/api/stream/Connection.ts +++ b/packages/backend/src/server/api/stream/Connection.ts @@ -202,10 +202,11 @@ export default class Connection { if (!payload.id || typeof payload.id !== 'string') return; const current = this.subscribingNotes.get(payload.id) ?? 0; + const updated = current + 1; + this.subscribingNotes.set(payload.id, updated); // Limit the number of distinct notes that can be subscribed to. - // If current is-zero, then this is a new note and we need to check the limit - if (current === 0 && this.subscribingNotes.size >= MAX_SUBSCRIPTIONS_PER_CONNECTION) { + 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!; @@ -214,9 +215,6 @@ export default class Connection { this.subscriber.off(`noteStream:${oldestKey}`, this.onNoteStreamMessage); } - const updated = current + 1; - this.subscribingNotes.set(payload.id, updated); - if (updated === 1) { this.subscriber.on(`noteStream:${payload.id}`, this.onNoteStreamMessage); } -- cgit v1.2.3-freya From 14a7309cfbc35cb8e448ef8ee6e8e3cd13b62013 Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Thu, 27 Mar 2025 12:28:42 -0400 Subject: avoid leaking cached notes in WS connection --- .../backend/src/server/api/stream/Connection.ts | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) (limited to 'packages/backend/src/server/api/stream') diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts index 96b968d890..7cc6157999 100644 --- a/packages/backend/src/server/api/stream/Connection.ts +++ b/packages/backend/src/server/api/stream/Connection.ts @@ -24,6 +24,7 @@ import type Channel from './channel.js'; const MAX_CHANNELS_PER_CONNECTION = 32; const MAX_SUBSCRIPTIONS_PER_CONNECTION = 256; +const MAX_CACHED_NOTES_PER_CONNECTION = 32; /** * Main stream connection @@ -36,7 +37,7 @@ export default class Connection { public subscriber: StreamEventEmitter; private channels: Channel[] = []; private subscribingNotes = new Map(); - private cachedNotes: Packed<'Note'>[] = []; + private cachedNotes = new Map>(); public userProfile: MiUserProfile | null = null; public following: Record | undefined> = {}; public followingChannels: Set = new Set(); @@ -158,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); } }; @@ -178,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)) { @@ -378,5 +377,6 @@ export default class Connection { this.fetchIntervalId = null; this.channels = []; this.subscribingNotes.clear(); + this.cachedNotes.clear(); } } -- cgit v1.2.3-freya From eff73218604089148a3382c1c8530afc652880a4 Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Thu, 27 Mar 2025 12:33:38 -0400 Subject: avoid duplicate channels in WS connection --- .../backend/src/server/api/stream/Connection.ts | 28 ++++++++++++++-------- 1 file changed, 18 insertions(+), 10 deletions(-) (limited to 'packages/backend/src/server/api/stream') diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts index 7cc6157999..4617fba9d1 100644 --- a/packages/backend/src/server/api/stream/Connection.ts +++ b/packages/backend/src/server/api/stream/Connection.ts @@ -35,7 +35,7 @@ export default class Connection { public token?: MiAccessToken; private wsConnection: WebSocket.WebSocket; public subscriber: StreamEventEmitter; - private channels: Channel[] = []; + private channels = new Map(); private subscribingNotes = new Map(); private cachedNotes = new Map>(); public userProfile: MiUserProfile | null = null; @@ -299,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; } @@ -315,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) { @@ -336,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); } } @@ -355,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); } @@ -367,7 +375,7 @@ 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()) { @@ -375,7 +383,7 @@ export default class Connection { } this.fetchIntervalId = null; - this.channels = []; + this.channels.clear(); this.subscribingNotes.clear(); this.cachedNotes.clear(); } -- cgit v1.2.3-freya From fafb811333e1bb375517810c05f95d048edb106b Mon Sep 17 00:00:00 2001 From: Hazelnoot Date: Fri, 28 Mar 2025 11:44:29 -0400 Subject: increase limits on WS note subscriptions and cached notes --- packages/backend/src/server/api/stream/Connection.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'packages/backend/src/server/api/stream') diff --git a/packages/backend/src/server/api/stream/Connection.ts b/packages/backend/src/server/api/stream/Connection.ts index 4617fba9d1..a69c1b8b7f 100644 --- a/packages/backend/src/server/api/stream/Connection.ts +++ b/packages/backend/src/server/api/stream/Connection.ts @@ -23,8 +23,8 @@ import type { EventEmitter } from 'events'; import type Channel from './channel.js'; const MAX_CHANNELS_PER_CONNECTION = 32; -const MAX_SUBSCRIPTIONS_PER_CONNECTION = 256; -const MAX_CACHED_NOTES_PER_CONNECTION = 32; +const MAX_SUBSCRIPTIONS_PER_CONNECTION = 512; +const MAX_CACHED_NOTES_PER_CONNECTION = 64; /** * Main stream connection -- cgit v1.2.3-freya