streamingController: rate-limit with ttl-cache

This commit is contained in:
Alex Gleason 2024-06-20 16:15:31 -05:00
parent 766290bd45
commit 44dfd15502
No known key found for this signature in database
GPG key ID: 7211D1F99744FBB7

View file

@ -9,9 +9,10 @@ import { MuteListPolicy } from '@/policies/MuteListPolicy.ts';
import { getFeedPubkeys } from '@/queries.ts';
import { hydrateEvents } from '@/storages/hydrate.ts';
import { Storages } from '@/storages.ts';
import { bech32ToPubkey } from '@/utils.ts';
import { bech32ToPubkey, Time } from '@/utils.ts';
import { renderReblog, renderStatus } from '@/views/mastodon/statuses.ts';
import { renderNotification } from '@/views/mastodon/notifications.ts';
import TTLCache from '@isaacs/ttlcache';
const debug = Debug('ditto:streaming');
@ -36,6 +37,11 @@ const streamSchema = z.enum([
type Stream = z.infer<typeof streamSchema>;
const LIMITER_WINDOW = Time.minutes(5);
const LIMITER_LIMIT = 100;
const limiter = new TTLCache<string, number>();
const streamingController: AppController = async (c) => {
const upgrade = c.req.header('upgrade');
const token = c.req.header('sec-websocket-protocol');
@ -56,6 +62,7 @@ const streamingController: AppController = async (c) => {
const store = await Storages.db();
const pubsub = await Storages.pubsub();
const ip = c.req.header('x-real-ip');
const policy = pubkey ? new MuteListPolicy(pubkey, await Storages.admin()) : undefined;
function send(name: string, payload: object) {
@ -119,6 +126,22 @@ const streamingController: AppController = async (c) => {
}
};
socket.onmessage = (e) => {
if (ip) {
const count = limiter.get(ip) ?? 0;
limiter.set(ip, count + 1, { ttl: LIMITER_WINDOW });
if (count > LIMITER_LIMIT) {
socket.close(1008, 'Rate limit exceeded');
return;
}
}
if (typeof e.data !== 'string') {
socket.close(1003, 'Invalid message');
}
};
socket.onclose = () => {
controller.abort();
};