feat: Add API ratelimits

5 files changed, 166 insertions, 13 deletions

diff --git a/packages/trpc/index.ts b/packages/trpc/index.ts
index 90f37ae4..cc62c534 100644
--- a/packages/trpc/index.ts
+++ b/packages/trpc/index.ts
@@ -5,6 +5,7 @@ import { ZodError } from "zod";
 import type { db } from "@karakeep/db";
 import serverConfig from "@karakeep/shared/config";
 
+import { createRateLimitMiddleware } from "./rateLimit";
 import {
   apiErrorsTotalCounter,
   apiRequestDurationSummary,
@@ -86,21 +87,38 @@ export const procedure = t.procedure
     end();
     return res;
   });
-export const publicProcedure = procedure;
 
-export const authedProcedure = procedure.use(function isAuthed(opts) {
-  const user = opts.ctx.user;
+// Default public procedure rate limiting
+export const publicProcedure = procedure.use(
+  createRateLimitMiddleware({
+    name: "globalPublic",
+    windowMs: 60 * 1000,
+    maxRequests: 1000,
+  }),
+);
 
-  if (!user?.id) {
-    throw new TRPCError({ code: "UNAUTHORIZED" });
-  }
+export const authedProcedure = procedure
+  // Default authed procedure rate limiting
+  .use(
+    createRateLimitMiddleware({
+      name: "globalAuthed",
+      windowMs: 60 * 1000,
+      maxRequests: 3000,
+    }),
+  )
+  .use(function isAuthed(opts) {
+    const user = opts.ctx.user;
+
+    if (!user?.id) {
+      throw new TRPCError({ code: "UNAUTHORIZED" });
+    }
 
-  return opts.next({
-    ctx: {
-      user,
-    },
+    return opts.next({
+      ctx: {
+        user,
+      },
+    });
   });
-});
 
 export const adminProcedure = authedProcedure.use(function isAdmin(opts) {
   const user = opts.ctx.user;
@@ -109,3 +127,6 @@ export const adminProcedure = authedProcedure.use(function isAdmin(opts) {
   }
   return opts.next(opts);
 });
+
+// Export the rate limiting utilities for use in routers
+export { createRateLimitMiddleware };
diff --git a/packages/trpc/rateLimit.ts b/packages/trpc/rateLimit.ts
new file mode 100644
index 00000000..b9aa4aa1
--- /dev/null
+++ b/packages/trpc/rateLimit.ts
@@ -0,0 +1,72 @@
+import { TRPCError } from "@trpc/server";
+
+import serverConfig from "@karakeep/shared/config";
+
+import { Context } from ".";
+
+interface RateLimitConfig {
+  name: string;
+  windowMs: number;
+  maxRequests: number;
+}
+
+interface RateLimitEntry {
+  count: number;
+  resetTime: number;
+}
+
+const rateLimitStore = new Map<string, RateLimitEntry>();
+
+function cleanupExpiredEntries() {
+  const now = Date.now();
+  for (const [key, entry] of rateLimitStore.entries()) {
+    if (now > entry.resetTime) {
+      rateLimitStore.delete(key);
+    }
+  }
+}
+
+setInterval(cleanupExpiredEntries, 60000);
+
+export function createRateLimitMiddleware<T>(config: RateLimitConfig) {
+  return function rateLimitMiddleware(opts: {
+    path: string;
+    ctx: Context;
+    next: () => Promise<T>;
+  }) {
+    if (!serverConfig.rateLimiting.enabled) {
+      return opts.next();
+    }
+    const ip = opts.ctx.req.ip;
+
+    if (!ip) {
+      return opts.next();
+    }
+
+    // TODO: Better fingerprinting
+    const key = `${config.name}:${ip}:${opts.path}`;
+    const now = Date.now();
+
+    let entry = rateLimitStore.get(key);
+
+    if (!entry || now > entry.resetTime) {
+      entry = {
+        count: 1,
+        resetTime: now + config.windowMs,
+      };
+      rateLimitStore.set(key, entry);
+      return opts.next();
+    }
+
+    if (entry.count >= config.maxRequests) {
+      const resetInSeconds = Math.ceil((entry.resetTime - now) / 1000);
+      throw new TRPCError({
+        code: "TOO_MANY_REQUESTS",
+        message: `Rate limit exceeded. Try again in ${resetInSeconds} seconds.`,
+      });
+    }
+
+    entry.count++;
+    return opts.next();
+  };
+}
diff --git a/packages/trpc/routers/apiKeys.ts b/packages/trpc/routers/apiKeys.ts
index eb52189b..d4e01aa5 100644
--- a/packages/trpc/routers/apiKeys.ts
+++ b/packages/trpc/routers/apiKeys.ts
@@ -11,7 +11,12 @@ import {
   logAuthenticationError,
   validatePassword,
 } from "../auth";
-import { authedProcedure, publicProcedure, router } from "../index";
+import {
+  authedProcedure,
+  createRateLimitMiddleware,
+  publicProcedure,
+  router,
+} from "../index";
 
 const zApiKeySchema = z.object({
   id: z.string(),
@@ -70,6 +75,13 @@ export const apiKeysAppRouter = router({
   // Exchange the username and password with an API key.
   // Homemade oAuth. This is used by the extension.
   exchange: publicProcedure
+    .use(
+      createRateLimitMiddleware({
+        name: "apiKey.exchange",
+        windowMs: 15 * 60 * 1000,
+        maxRequests: 10,
+      }),
+    ) // 10 requests per 15 minutes
     .input(
       z.object({
         keyName: z.string(),
@@ -97,6 +109,13 @@ export const apiKeysAppRouter = router({
       return await generateApiKey(input.keyName, user.id);
     }),
   validate: publicProcedure
+    .use(
+      createRateLimitMiddleware({
+        name: "apiKey.validate",
+        windowMs: 60 * 1000,
+        maxRequests: 30,
+      }),
+    ) // 30 requests per minute
     .input(z.object({ apiKey: z.string() }))
     .output(z.object({ success: z.boolean() }))
     .mutation(async ({ input, ctx }) => {
diff --git a/packages/trpc/routers/invites.ts b/packages/trpc/routers/invites.ts
index 5f7897c5..0a98f36a 100644
--- a/packages/trpc/routers/invites.ts
+++ b/packages/trpc/routers/invites.ts
@@ -7,7 +7,12 @@ import { invites, users } from "@karakeep/db/schema";
 
 import { generatePasswordSalt, hashPassword } from "../auth";
 import { sendInviteEmail } from "../email";
-import { adminProcedure, publicProcedure, router } from "../index";
+import {
+  adminProcedure,
+  createRateLimitMiddleware,
+  publicProcedure,
+  router,
+} from "../index";
 import { createUserRaw } from "./users";
 
 export const invitesAppRouter = router({
@@ -113,6 +118,13 @@ export const invitesAppRouter = router({
     }),
 
   get: publicProcedure
+    .use(
+      createRateLimitMiddleware({
+        name: "invites.get",
+        windowMs: 60 * 1000,
+        maxRequests: 10,
+      }),
+    )
     .input(
       z.object({
         token: z.string(),
@@ -153,6 +165,13 @@ export const invitesAppRouter = router({
     }),
 
   accept: publicProcedure
+    .use(
+      createRateLimitMiddleware({
+        name: "invites.accept",
+        windowMs: 60 * 1000,
+        maxRequests: 10,
+      }),
+    )
     .input(
       z.object({
         token: z.string(),
diff --git a/packages/trpc/routers/users.ts b/packages/trpc/routers/users.ts
index 58093b42..ebe7d96f 100644
--- a/packages/trpc/routers/users.ts
+++ b/packages/trpc/routers/users.ts
@@ -31,6 +31,7 @@ import {
   adminProcedure,
   authedProcedure,
   Context,
+  createRateLimitMiddleware,
   publicProcedure,
   router,
 } from "../index";
@@ -124,6 +125,13 @@ export async function createUser(
 
 export const usersAppRouter = router({
   create: publicProcedure
+    .use(
+      createRateLimitMiddleware({
+        name: "users.create",
+        windowMs: 60 * 1000,
+        maxRequests: 3,
+      }),
+    )
     .input(zSignUpSchema)
     .output(
       z.object({
@@ -541,6 +549,13 @@ export const usersAppRouter = router({
         .where(eq(userSettings.userId, ctx.user.id));
     }),
   verifyEmail: publicProcedure
+    .use(
+      createRateLimitMiddleware({
+        name: "users.verifyEmail",
+        windowMs: 5 * 60 * 1000,
+        maxRequests: 10,
+      }),
+    ) // 10 requests per 5 minutes
     .input(
       z.object({
         email: z.string().email(),
@@ -572,6 +587,13 @@ export const usersAppRouter = router({
       return { success: true };
     }),
   resendVerificationEmail: publicProcedure
+    .use(
+      createRateLimitMiddleware({
+        name: "users.resendVerificationEmail",
+        windowMs: 5 * 60 * 1000,
+        maxRequests: 3,
+      }),
+    ) // 3 requests per 5 minutes
     .input(
       z.object({
         email: z.string().email(),