packages/trpc/lib/rateLimit.ts


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

import { TRPCError } from "@trpc/server";

import type { RateLimitConfig } from "@karakeep/shared/ratelimiting";
import serverConfig from "@karakeep/shared/config";
import { getRateLimitClient } from "@karakeep/shared/ratelimiting";

/**
 * Create a tRPC middleware for rate limiting
 * @param config Rate limit configuration
 * @returns tRPC middleware function
 */
export function createRateLimitMiddleware<T>(config: RateLimitConfig) {
  return async function rateLimitMiddleware(opts: {
    path: string;
    ctx: { req: { ip: string | null } };
    next: () => Promise<T>;
  }) {
    if (!serverConfig.rateLimiting.enabled) {
      return opts.next();
    }

    const ip = opts.ctx.req.ip;

    if (!ip) {
      return opts.next();
    }

    const client = await getRateLimitClient();

    if (!client) {
      // If no rate limit client is registered, allow the request
      return opts.next();
    }

    // Build the rate limiting key from IP and path
    const key = `${ip}:${opts.path}`;
    const result = client.checkRateLimit(config, key);

    if (!result.allowed) {
      throw new TRPCError({
        code: "TOO_MANY_REQUESTS",
        message: `Rate limit exceeded. Try again in ${result.resetInSeconds} seconds.`,
      });
    }

    return opts.next();
  };
}