misty.cafe/mistykey
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

Merge branch 'develop' into merge/2025-03-24

Browse source 
# Conflicts:
#	packages/backend/src/server/api/StreamingApiServerService.ts
#	packages/backend/src/server/api/stream/Connection.ts
This commit is contained in:
Hazelnoot 2025-03-31 15:02:52 -04:00
commit e7632c83dc
7 changed files with 284 additions and 114 deletions
Download patch file Download diff file Expand all files Collapse all files

98
packages/backend/src/server/api/StreamingApiServerService.ts
View file

@ -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 { NotificationService } from '@/core/NotificationService.js';
import { bindThis } from '@/decorators.js';
import { CacheService } from '@/core/CacheService.js';
@ -24,13 +26,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(
@ -56,17 +61,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;
@ -86,21 +83,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;
@ -138,15 +120,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)
});
};
@ -163,6 +178,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,
});

104
packages/backend/src/server/api/stream/Connection.ts
View file

@ -22,6 +22,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
@ -30,12 +32,12 @@ 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>();
// TODO see if we should remove this, now that it has no more reads
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();
@ -44,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;
@ -58,11 +59,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');
}
@ -119,25 +119,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 {
@ -170,15 +158,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);
}
};
@ -200,9 +186,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);
@ -217,12 +213,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);
}
}
@ -289,7 +285,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;
}
@ -305,12 +305,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) {
@ -326,11 +330,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);
}
}
@ -345,7 +349,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);
}
@ -357,8 +361,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();
}
}