mirrors/barkey
1
0
Fork 0
mirror of https://codeberg.org/yeentown/barkey.git synced 2025-07-07 12:36:57 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

close websocket when rate limit exceeded

Browse source
This commit is contained in:
Hazelnoot 2025-03-27 10:55:46 -04:00
parent 18655386f3
commit bf1c9b67d6
1 changed files with 7 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -31,7 +31,6 @@ const MAX_CHANNELS_PER_CONNECTION = 32;
export default class Connection { export default class Connection {
public user?: MiUser; public user?: MiUser;
public token?: MiAccessToken; public token?: MiAccessToken;
private rateLimiter?: () => Promise<boolean>;
private wsConnection: WebSocket.WebSocket; private wsConnection: WebSocket.WebSocket;
public subscriber: StreamEventEmitter; public subscriber: StreamEventEmitter;
private channels: Channel[] = []; private channels: Channel[] = [];
@ -45,7 +44,6 @@ export default class Connection {
public userIdsWhoMeMutingRenotes: Set<string> = new Set(); public userIdsWhoMeMutingRenotes: Set<string> = new Set();
public userMutedInstances: Set<string> = new Set(); public userMutedInstances: Set<string> = new Set();
private fetchIntervalId: NodeJS.Timeout | null = null; private fetchIntervalId: NodeJS.Timeout | null = null;
private activeRateLimitRequests = 0;
private closingConnection = false; private closingConnection = false;
private logger: Logger; private logger: Logger;
@ -60,11 +58,10 @@ export default class Connection {
user: MiUser | null | undefined, user: MiUser | null | undefined,
token: MiAccessToken | null | undefined, token: MiAccessToken | null | undefined,
private ip: string, private ip: string,
rateLimiter: () => Promise<boolean>, private readonly rateLimiter: () => Promise<boolean>,
) { ) {
if (user) this.user = user; if (user) this.user = user;
if (token) this.token = token; if (token) this.token = token;
if (rateLimiter) this.rateLimiter = rateLimiter;
this.logger = loggerService.getLogger('streaming', 'coral'); this.logger = loggerService.getLogger('streaming', 'coral');
} }
@ -121,26 +118,14 @@ export default class Connection {
if (this.closingConnection) return; if (this.closingConnection) return;
if (this.rateLimiter) { // The rate limit is very high, so we can safely disconnect any client that hits it.
// this 4096 should match the `max` of the `rateLimiter`, see if (await this.rateLimiter()) {
// StreamingApiServerService this.logger.warn(`Closing a connection from ${this.ip} (user=${this.user?.id}}) due to an excessive influx of messages.`);
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.closingConnection = true;
this.wsConnection.close(1008, 'Please stop spamming the streaming API.'); this.wsConnection.close(1008, 'Disconnected - too many requests');
return; return;
} }
}
try { try {
obj = JSON.parse(data.toString()); obj = JSON.parse(data.toString());