12 files changed, 451 insertions, 75 deletions

diff --git a/packages/plugins/package.json b/packages/plugins/package.json
index 8b3f73f7..5931d1a7 100644
--- a/packages/plugins/package.json
+++ b/packages/plugins/package.json
@@ -7,6 +7,7 @@
   "exports": {
     "./queue-liteque": "./queue-liteque/index.ts",
     "./queue-restate": "./queue-restate/index.ts",
+    "./ratelimit-memory": "./ratelimit-memory/index.ts",
     "./search-meilisearch": "./search-meilisearch/index.ts"
   },
   "scripts": {
diff --git a/packages/plugins/ratelimit-memory/index.ts b/packages/plugins/ratelimit-memory/index.ts
new file mode 100644
index 00000000..e47a2341
--- /dev/null
+++ b/packages/plugins/ratelimit-memory/index.ts
@@ -0,0 +1,13 @@
+// Auto-register the RateLimit plugin when this package is imported
+import { PluginManager, PluginType } from "@karakeep/shared/plugins";
+
+import { RateLimitProvider } from "./src";
+
+PluginManager.register({
+  type: PluginType.RateLimit,
+  name: "In-Memory Rate Limiter",
+  provider: new RateLimitProvider(),
+});
+
+// Export the provider and rate limiter class for advanced usage
+export { RateLimiter, RateLimitProvider } from "./src";
diff --git a/packages/plugins/ratelimit-memory/src/index.test.ts b/packages/plugins/ratelimit-memory/src/index.test.ts
new file mode 100644
index 00000000..5bbed769
--- /dev/null
+++ b/packages/plugins/ratelimit-memory/src/index.test.ts
@@ -0,0 +1,253 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+import { RateLimiter } from "./index";
+
+describe("RateLimiter", () => {
+  let rateLimiter: RateLimiter;
+
+  beforeEach(() => {
+    rateLimiter = new RateLimiter();
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+    rateLimiter.clear();
+  });
+
+  describe("checkRateLimit", () => {
+    it("should allow requests within rate limit", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 3,
+      };
+
+      const result1 = rateLimiter.checkRateLimit(config, "user1");
+      const result2 = rateLimiter.checkRateLimit(config, "user1");
+      const result3 = rateLimiter.checkRateLimit(config, "user1");
+
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(true);
+      expect(result3.allowed).toBe(true);
+    });
+
+    it("should block requests exceeding rate limit", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 2,
+      };
+
+      const result1 = rateLimiter.checkRateLimit(config, "user1");
+      const result2 = rateLimiter.checkRateLimit(config, "user1");
+      const result3 = rateLimiter.checkRateLimit(config, "user1");
+
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(true);
+      expect(result3.allowed).toBe(false);
+      expect(result3.resetInSeconds).toBeDefined();
+      expect(result3.resetInSeconds).toBeGreaterThan(0);
+    });
+
+    it("should reset after window expires", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 2,
+      };
+
+      // First two requests allowed
+      const result1 = rateLimiter.checkRateLimit(config, "user1");
+      const result2 = rateLimiter.checkRateLimit(config, "user1");
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(true);
+
+      // Third request blocked
+      const result3 = rateLimiter.checkRateLimit(config, "user1");
+      expect(result3.allowed).toBe(false);
+
+      // Advance time past the window
+      vi.advanceTimersByTime(61000);
+
+      // Should allow request after window reset
+      const result4 = rateLimiter.checkRateLimit(config, "user1");
+      expect(result4.allowed).toBe(true);
+    });
+
+    it("should isolate rate limits by identifier", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      const result1 = rateLimiter.checkRateLimit(config, "user1");
+      const result2 = rateLimiter.checkRateLimit(config, "user2");
+
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(true);
+    });
+
+    it("should isolate rate limits by key", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      const result1 = rateLimiter.checkRateLimit(config, "user1:/api/v1");
+      const result2 = rateLimiter.checkRateLimit(config, "user1:/api/v2");
+
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(true);
+    });
+
+    it("should isolate rate limits by config name", () => {
+      const config1 = {
+        name: "api",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+      const config2 = {
+        name: "auth",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      const result1 = rateLimiter.checkRateLimit(config1, "user1");
+      const result2 = rateLimiter.checkRateLimit(config2, "user1");
+
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(true);
+    });
+
+    it("should calculate correct resetInSeconds", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      // First request allowed
+      rateLimiter.checkRateLimit(config, "user1");
+
+      // Advance time by 30 seconds
+      vi.advanceTimersByTime(30000);
+
+      // Second request blocked
+      const result = rateLimiter.checkRateLimit(config, "user1");
+      expect(result.allowed).toBe(false);
+      // Should have ~30 seconds remaining
+      expect(result.resetInSeconds).toBeGreaterThan(29);
+      expect(result.resetInSeconds).toBeLessThanOrEqual(30);
+    });
+  });
+
+  describe("reset", () => {
+    it("should reset rate limit for specific identifier", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      // Use up the limit
+      rateLimiter.checkRateLimit(config, "user1");
+      const result1 = rateLimiter.checkRateLimit(config, "user1");
+      expect(result1.allowed).toBe(false);
+
+      // Reset the limit
+      rateLimiter.reset(config, "user1");
+
+      // Should allow request again
+      const result2 = rateLimiter.checkRateLimit(config, "user1");
+      expect(result2.allowed).toBe(true);
+    });
+
+    it("should reset rate limit for specific key", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      // Use up the limit for key1
+      rateLimiter.checkRateLimit(config, "user1:/path1");
+      const result1 = rateLimiter.checkRateLimit(config, "user1:/path1");
+      expect(result1.allowed).toBe(false);
+
+      // Reset only key1
+      rateLimiter.reset(config, "user1:/path1");
+
+      // key1 should be allowed
+      const result2 = rateLimiter.checkRateLimit(config, "user1:/path1");
+      expect(result2.allowed).toBe(true);
+    });
+
+    it("should not affect other identifiers", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      // Use up limits for both users
+      rateLimiter.checkRateLimit(config, "user1");
+      rateLimiter.checkRateLimit(config, "user2");
+
+      // Reset only user1
+      rateLimiter.reset(config, "user1");
+
+      const result1 = rateLimiter.checkRateLimit(config, "user1");
+      const result2 = rateLimiter.checkRateLimit(config, "user2");
+
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(false);
+    });
+  });
+
+  describe("clear", () => {
+    it("should clear all rate limits", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      // Use up limits for multiple users
+      rateLimiter.checkRateLimit(config, "user1");
+      rateLimiter.checkRateLimit(config, "user2");
+
+      // Clear all limits
+      rateLimiter.clear();
+
+      // All should be allowed
+      const result1 = rateLimiter.checkRateLimit(config, "user1");
+      const result2 = rateLimiter.checkRateLimit(config, "user2");
+
+      expect(result1.allowed).toBe(true);
+      expect(result2.allowed).toBe(true);
+    });
+  });
+
+  describe("cleanup", () => {
+    it("should cleanup expired entries", () => {
+      const config = {
+        name: "test",
+        windowMs: 60000,
+        maxRequests: 1,
+      };
+
+      // Create an entry
+      rateLimiter.checkRateLimit(config, "user1");
+
+      // Advance time past window + cleanup interval
+      vi.advanceTimersByTime(61000 + 60000);
+
+      // Entry should be cleaned up and new request allowed
+      const result = rateLimiter.checkRateLimit(config, "user1");
+      expect(result.allowed).toBe(true);
+    });
+  });
+});
diff --git a/packages/plugins/ratelimit-memory/src/index.ts b/packages/plugins/ratelimit-memory/src/index.ts
new file mode 100644
index 00000000..c47889d3
--- /dev/null
+++ b/packages/plugins/ratelimit-memory/src/index.ts
@@ -0,0 +1,86 @@
+import type {
+  RateLimitClient,
+  RateLimitConfig,
+  RateLimitResult,
+} from "@karakeep/shared/ratelimiting";
+import { PluginProvider } from "@karakeep/shared/plugins";
+
+interface RateLimitEntry {
+  count: number;
+  resetTime: number;
+}
+
+export class RateLimiter implements RateLimitClient {
+  private store = new Map<string, RateLimitEntry>();
+  private cleanupProbability: number;
+
+  constructor(cleanupProbability = 0.01) {
+    // Probability of cleanup on each check (default 1%)
+    this.cleanupProbability = cleanupProbability;
+  }
+
+  private cleanupExpiredEntries() {
+    const now = Date.now();
+    for (const [key, entry] of this.store.entries()) {
+      if (now > entry.resetTime) {
+        this.store.delete(key);
+      }
+    }
+  }
+
+  checkRateLimit(config: RateLimitConfig, key: string): RateLimitResult {
+    if (!key) {
+      return { allowed: true };
+    }
+
+    // Probabilistic cleanup
+    if (Math.random() < this.cleanupProbability) {
+      this.cleanupExpiredEntries();
+    }
+
+    const rateLimitKey = `${config.name}:${key}`;
+    const now = Date.now();
+
+    let entry = this.store.get(rateLimitKey);
+
+    if (!entry || now > entry.resetTime) {
+      entry = {
+        count: 1,
+        resetTime: now + config.windowMs,
+      };
+      this.store.set(rateLimitKey, entry);
+      return { allowed: true };
+    }
+
+    if (entry.count >= config.maxRequests) {
+      const resetInSeconds = Math.ceil((entry.resetTime - now) / 1000);
+      return {
+        allowed: false,
+        resetInSeconds,
+      };
+    }
+
+    entry.count++;
+    return { allowed: true };
+  }
+
+  reset(config: RateLimitConfig, key: string) {
+    const rateLimitKey = `${config.name}:${key}`;
+    this.store.delete(rateLimitKey);
+  }
+
+  clear() {
+    this.store.clear();
+  }
+}
+
+export class RateLimitProvider implements PluginProvider<RateLimitClient> {
+  private client: RateLimiter | null = null;
+
+  async getClient(): Promise<RateLimitClient | null> {
+    if (!this.client) {
+      this.client = new RateLimiter();
+    }
+    return this.client;
+  }
+}
diff --git a/packages/shared-server/src/plugins.ts b/packages/shared-server/src/plugins.ts
index 97503403..9e78bedb 100644
--- a/packages/shared-server/src/plugins.ts
+++ b/packages/shared-server/src/plugins.ts
@@ -10,6 +10,7 @@ export async function loadAllPlugins() {
   await import("@karakeep/plugins/queue-liteque");
   await import("@karakeep/plugins/queue-restate");
   await import("@karakeep/plugins/search-meilisearch");
+  await import("@karakeep/plugins/ratelimit-memory");
   PluginManager.logAllPlugins();
   pluginsLoaded = true;
 }
diff --git a/packages/shared/plugins.ts b/packages/shared/plugins.ts
index e04fd91e..2d03ee39 100644
--- a/packages/shared/plugins.ts
+++ b/packages/shared/plugins.ts
@@ -1,17 +1,20 @@
 // Implementation inspired from Outline
 
 import type { QueueClient } from "./queueing";
+import type { RateLimitClient } from "./ratelimiting";
 import logger from "./logger";
 import { SearchIndexClient } from "./search";
 
 export enum PluginType {
   Search = "search",
   Queue = "queue",
+  RateLimit = "ratelimit",
 }
 
 interface PluginTypeMap {
   [PluginType.Search]: SearchIndexClient;
   [PluginType.Queue]: QueueClient;
+  [PluginType.RateLimit]: RateLimitClient;
 }
 
 export interface TPlugin<T extends PluginType> {
@@ -31,6 +34,7 @@ export class PluginManager {
   private static providers: ProviderMap = {
     [PluginType.Search]: [],
     [PluginType.Queue]: [],
+    [PluginType.RateLimit]: [],
   };
 
   static register<T extends PluginType>(plugin: TPlugin<T>): void {
diff --git a/packages/shared/ratelimiting.ts b/packages/shared/ratelimiting.ts
new file mode 100644
index 00000000..3b22310b
--- /dev/null
+++ b/packages/shared/ratelimiting.ts
@@ -0,0 +1,38 @@
+import { PluginManager, PluginType } from "./plugins";
+
+export interface RateLimitConfig {
+  name: string;
+  windowMs: number;
+  maxRequests: number;
+}
+
+export interface RateLimitResult {
+  allowed: boolean;
+  resetInSeconds?: number;
+}
+
+export interface RateLimitClient {
+  /**
+   * Check if a request should be allowed based on rate limiting rules
+   * @param config Rate limit configuration
+   * @param key Unique rate limiting key (e.g., "ip:127.0.0.1:path:/api/v1")
+   * @returns Result indicating if the request is allowed and reset time if not
+   */
+  checkRateLimit(config: RateLimitConfig, key: string): RateLimitResult;
+
+  /**
+   * Reset rate limit for a specific key
+   * @param config Rate limit configuration
+   * @param key Unique rate limiting key
+   */
+  reset(config: RateLimitConfig, key: string): void;
+
+  /**
+   * Clear all rate limit entries
+   */
+  clear(): void;
+}
+
+export async function getRateLimitClient(): Promise<RateLimitClient | null> {
+  return PluginManager.getClient(PluginType.RateLimit);
+}
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();
-  };
-}
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 4d89fc27..47387568 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1317,6 +1317,9 @@ importers:
       '@karakeep/db':
         specifier: workspace:*
         version: link:../db
+      '@karakeep/plugins':
+        specifier: workspace:*
+        version: link:../plugins
       '@karakeep/shared':
         specifier: workspace:*
         version: link:../shared