4 files changed, 52 insertions, 75 deletions
diff --git a/packages/trpc/index.ts b/packages/trpc/index.ts
index cc62c534..555ca3ba 100644
--- a/packages/trpc/index.ts
+++ b/packages/trpc/index.ts
@@ -5,7 +5,7 @@ import { ZodError } from "zod";
 import type { db } from "@karakeep/db";
 import serverConfig from "@karakeep/shared/config";
 
-import { createRateLimitMiddleware } from "./rateLimit";
+import { createRateLimitMiddleware } from "./lib/rateLimit";
 import {
   apiErrorsTotalCounter,
   apiRequestDurationSummary,
@@ -128,5 +128,5 @@ export const adminProcedure = authedProcedure.use(function isAdmin(opts) {
   return opts.next(opts);
 });
 
-// Export the rate limiting utilities for use in routers
-export { createRateLimitMiddleware };
+// Export the rate limiting middleware for use in routers
+export { createRateLimitMiddleware } from "./lib/rateLimit";
diff --git a/packages/trpc/lib/rateLimit.ts b/packages/trpc/lib/rateLimit.ts
new file mode 100644
index 00000000..bf8a8e8b
--- /dev/null
+++ b/packages/trpc/lib/rateLimit.ts
@@ -0,0 +1,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();
+  };
+}
diff --git a/packages/trpc/package.json b/packages/trpc/package.json
index d1896a0b..d9fa12c0 100644
--- a/packages/trpc/package.json
+++ b/packages/trpc/package.json
@@ -14,6 +14,7 @@
   },
   "dependencies": {
     "@karakeep/db": "workspace:*",
+    "@karakeep/plugins": "workspace:*",
     "@karakeep/shared": "workspace:*",
     "@karakeep/shared-server": "workspace:*",
     "@trpc/server": "^11.4.3",
diff --git a/packages/trpc/rateLimit.ts b/packages/trpc/rateLimit.ts
deleted file mode 100644
index b9aa4aa1..00000000
--- a/packages/trpc/rateLimit.ts
+++ /dev/null
@@ -1,72 +0,0 @@
-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();
-  };
-}