Add rate limits

Co-authored-by: mrjvs <mistrjvs@gmail.com>

src/config/schema.ts

@@ -59,4 +59,11 @@ export const configSchema = z.object({
       clientKey: z.string().min(1).optional(),
     })
     .default({}),
+  ratelimits: z
+    .object({
+      // enabled captchas on register
+      enabled: z.coerce.boolean().default(false),
+      redisUrl: z.string().optional(),
+    })
+    .default({}),
 });

src/main.ts

@@ -2,6 +2,7 @@ import { setupFastify, startFastify } from '@/modules/fastify';
 import { setupJobs } from '@/modules/jobs';
 import { setupMetrics } from '@/modules/metrics';
 import { setupMikroORM } from '@/modules/mikro';
+import { setupRatelimits } from '@/modules/ratelimits';
 import { scopedLogger } from '@/services/logger';
 
 const log = scopedLogger('mw-backend');
@@ -11,6 +12,7 @@ async function bootstrap(): Promise<void> {
     evt: 'setup',
   });
 
+  await setupRatelimits();
   const app = await setupFastify();
   await setupMikroORM();
   await setupMetrics(app);

src/modules/ratelimits/index.ts

@@ -0,0 +1,24 @@
+import { conf } from '@/config';
+import { Limiter } from '@/modules/ratelimits/limiter';
+import { connectRedis } from '@/modules/ratelimits/redis';
+import { scopedLogger } from '@/services/logger';
+
+const log = scopedLogger('ratelimits');
+
+let limiter: null | Limiter = null;
+
+export function getLimiter() {
+  return limiter;
+}
+
+export async function setupRatelimits() {
+  if (!conf.ratelimits.enabled) {
+    log.warn('Ratelimits disabled!');
+    return;
+  }
+  const redis = await connectRedis();
+  limiter = new Limiter({
+    redis,
+  });
+  log.info('Ratelimits have been setup!');
+}

src/modules/ratelimits/limiter.ts

@@ -0,0 +1,63 @@
+import Redis from 'ioredis';
+import RateLimiter from 'async-ratelimiter';
+import ms from 'ms';
+import { StatusError } from '@/services/error';
+
+export interface LimiterOptions {
+  redis: Redis;
+}
+
+interface LimitBucket {
+  limiter: RateLimiter;
+}
+
+interface BucketOptions {
+  id: string;
+  window: string;
+  max: number;
+  inc?: number;
+}
+
+export class Limiter {
+  private redis: Redis;
+  private buckets: Record<string, LimitBucket> = {};
+
+  constructor(ops: LimiterOptions) {
+    this.redis = ops.redis;
+  }
+
+  async bump(req: { ip: string }, ops: BucketOptions) {
+    const ip = req.ip;
+    if (!this.buckets[ops.id]) {
+      this.buckets[ops.id] = {
+        limiter: new RateLimiter({
+          db: this.redis,
+          namespace: `RATELIMIT_${ops.id}`,
+          duration: ms(ops.window),
+          max: ops.max,
+        }),
+      };
+    }
+
+    for (let i = 1; i < (ops.inc ?? 0); i++) {
+      await this.buckets[ops.id].limiter.get({
+        id: ip,
+      });
+    }
+    const currentLimit = await this.buckets[ops.id].limiter.get({
+      id: ip,
+    });
+
+    return {
+      hasBeenLimited: currentLimit.remaining <= 0,
+      limit: currentLimit,
+    };
+  }
+
+  async assertAndBump(req: { ip: string }, ops: BucketOptions) {
+    const { hasBeenLimited } = await this.bump(req, ops);
+    if (hasBeenLimited) {
+      throw new StatusError('Ratelimited', 429);
+    }
+  }
+}

src/modules/ratelimits/redis.ts

@@ -0,0 +1,7 @@
+import { conf } from '@/config';
+import Redis from 'ioredis';
+
+export function connectRedis() {
+  if (!conf.ratelimits.redisUrl) throw new Error('missing redis URL');
+  return new Redis(conf.ratelimits.redisUrl);
+}

src/routes/auth/login.ts

@@ -25,7 +25,13 @@ export const loginAuthRouter = makeRouter((app) => {
   app.post(
     '/auth/login/start',
     { schema: { body: startSchema } },
-    handle(async ({ em, body }) => {
+    handle(async ({ em, body, limiter, req }) => {
+      await limiter?.assertAndBump(req, {
+        id: 'login_challenge_tokens',
+        max: 20,
+        window: '10m',
+      });
+
       const user = await em.findOne(User, { publicKey: body.publicKey });
 
       if (user == null) {
@@ -46,7 +52,13 @@ export const loginAuthRouter = makeRouter((app) => {
     app.post(
       '/auth/login/complete',
       { schema: { body: completeSchema } },
-      handle(async ({ em, body, req }) => {
+      handle(async ({ em, body, req, limiter }) => {
+        await limiter?.assertAndBump(req, {
+          id: 'login_complete',
+          max: 20,
+          window: '10m',
+        });
+
         await assertChallengeCode(
           em,
           body.challenge.code,

src/routes/auth/manage.ts

@@ -32,7 +32,12 @@ export const manageAuthRouter = makeRouter((app) => {
   app.post(
     '/auth/register/start',
     { schema: { body: startSchema } },
-    handle(async ({ em, body }) => {
+    handle(async ({ em, body, limiter, req }) => {
+      await limiter?.assertAndBump(req, {
+        id: 'register_challenge_tokens',
+        max: 10,
+        window: '10m',
+      });
       await assertCaptcha(body.captchaToken);
 
       const challenge = new ChallengeCode();
@@ -50,7 +55,13 @@ export const manageAuthRouter = makeRouter((app) => {
   app.post(
     '/auth/register/complete',
     { schema: { body: completeSchema } },
-    handle(async ({ em, body, req }) => {
+    handle(async ({ em, body, req, limiter }) => {
+      await limiter?.assertAndBump(req, {
+        id: 'register_complete',
+        max: 10,
+        window: '10m',
+      });
+
       await assertChallengeCode(
         em,
         body.challenge.code,

src/routes/metrics.ts

@@ -29,7 +29,14 @@ export const metricsRouter = makeRouter((app) => {
         body: metricsProviderInputSchema,
       },
     },
-    handle(async ({ em, body }) => {
+    handle(async ({ em, body, req, limiter }) => {
+      await limiter?.assertAndBump(req, {
+        id: 'provider_metrics',
+        max: 300,
+        inc: body.items.length,
+        window: '30m',
+      });
+
       const entities = body.items.map((v) => {
         const metric = new ProviderMetric();
         em.assign(metric, {

src/services/handler.ts

@@ -1,4 +1,6 @@
 import { getORM } from '@/modules/mikro';
+import { getLimiter } from '@/modules/ratelimits';
+import { Limiter } from '@/modules/ratelimits/limiter';
 import { makeAuthContext } from '@/services/auth';
 import { EntityManager } from '@mikro-orm/postgresql';
 import {
@@ -74,6 +76,7 @@ export type RequestContext<
     Logger
   >['query'];
   em: EntityManager;
+  limiter: Limiter | null;
   auth: ReturnType<typeof makeAuthContext>;
 };
 
@@ -124,6 +127,7 @@ export function handle<
         query: req.query,
         em,
         auth: makeAuthContext(em, req),
+        limiter: getLimiter(),
       }),
     );
   };