diff options
| author | Mohamed Bassem <me@mbassem.com> | 2025-11-09 17:10:54 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2025-11-09 17:10:54 +0000 |
| commit | 03161482b44bd67f6eafb3e3d51107811b638d4b (patch) | |
| tree | bc32d16ee187fd94f31a0c8b06c5861c7694176e | |
| parent | ec87813a257e63f8a161e7bc04679e9fab6fbcf6 (diff) | |
| download | karakeep-03161482b44bd67f6eafb3e3d51107811b638d4b.tar.zst | |
refactor: Extract ratelimiter into separate plugin (#2112)
* refactor(trpc): extract rate limiter into dedicated plugin
Move the rate limiting middleware from the trpc package to the
centralized plugins package. This improves code organization by
consolidating all plugins in a single location.
Changes:
- Created packages/plugins/trpc-ratelimit/ plugin
- Moved rate limiter from packages/trpc/rateLimit.ts to
packages/plugins/trpc-ratelimit/src/index.ts
- Added trpc-ratelimit export to plugins package.json
- Added @trpc/server dependency to plugins package
- Updated trpc package to import from @karakeep/plugins/trpc-ratelimit
- Added @karakeep/plugins dependency to trpc package
- Removed packages/trpc/plugins/ directory
* refactor(plugins): decouple rate limiter from tRPC
Refactor the rate limiting plugin to be framework-agnostic, allowing
it to be used outside of tRPC contexts. The plugin now has a generic
core with a tRPC-specific adapter.
Changes:
- Renamed trpc-ratelimit to ratelimit plugin
- Created generic RateLimiter class with framework-agnostic API
- Added checkRateLimit() method that returns allow/deny results
- Created separate tRPC adapter (src/trpc.ts) that uses the generic core
- Exported both generic (RateLimiter, globalRateLimiter) and
tRPC-specific (createRateLimitMiddleware) APIs
- Updated trpc package to import from @karakeep/plugins/ratelimit
- Updated plugins package.json exports
Benefits:
- Rate limiter can now be used in any context (HTTP handlers, WebSocket, etc.)
- Cleaner separation of concerns
- Easy to create adapters for other frameworks
- Generic API allows for custom error handling
* refactor(plugins): integrate rate limiter with plugin registry
Refactor the rate limiting plugin to use the centralized plugin
system with PluginManager, making it consistent with other plugins
like queue and search providers.
Changes:
- Added RateLimit plugin type to PluginType enum
- Created RateLimitClient interface in packages/shared/ratelimiting.ts
- Created RateLimitProvider class implementing PluginProvider
- Updated plugin to auto-register with PluginManager on import
- Updated tRPC adapter to use getRateLimitClient() from PluginManager
- Added ratelimit plugin to loadAllPlugins() in shared-server
- Updated shared/plugins.ts with RateLimit type mapping
Benefits:
- Consistent plugin architecture across the codebase
- Rate limiter can be swapped with alternative implementations
- Centralized plugin management and logging
- Better separation of concerns
- Framework-agnostic core with tRPC adapter pattern
* refactor(trpc): move rate limit middleware to trpc package
Move the tRPC-specific rate limiting middleware from the plugins
package to the trpc package, making the plugins package
framework-agnostic.
Changes:
- Moved packages/plugins/ratelimit/src/trpc.ts to
packages/trpc/lib/rateLimit.ts
- Updated packages/trpc/index.ts to import from local lib/rateLimit
- Removed tRPC export from packages/plugins/ratelimit/index.ts
- Removed @trpc/server dependency from packages/plugins/package.json
Benefits:
- plugins package is now framework-agnostic
- tRPC-specific code lives in the trpc package where it belongs
- Cleaner separation of concerns
- Rate limiter plugin can be used in any context without tRPC
* refactor(plugins): rename to ratelimit-memory and add tests
Rename the rate limiting plugin from "ratelimit" to "ratelimit-memory"
to better indicate it's an in-memory implementation. This naming leaves
room for future implementations like ratelimit-redis. Also added
comprehensive test coverage.
Changes:
- Renamed packages/plugins/ratelimit to ratelimit-memory
- Updated package.json export from ./ratelimit to ./ratelimit-memory
- Updated shared-server to import @karakeep/plugins/ratelimit-memory
- Added comprehensive unit tests (index.test.ts):
- Rate limit enforcement tests
- Window expiration tests
- Identifier and path isolation tests
- Reset functionality tests
- Cleanup mechanism tests
- Added provider integration tests (provider.test.ts):
- PluginProvider interface compliance
- Client singleton behavior
- End-to-end rate limiting functionality
Benefits:
- More descriptive plugin name indicating the storage mechanism
- Better test coverage ensuring reliability
- Easier to add alternative implementations (Redis, etc.)
* change the api to only take the key
* move the serverConfig check to the trpc
* fix lockfile
* get rid of the timer
---------
Co-authored-by: Claude <noreply@anthropic.com>
| -rw-r--r-- | packages/plugins/package.json | 1 | ||||
| -rw-r--r-- | packages/plugins/ratelimit-memory/index.ts | 13 | ||||
| -rw-r--r-- | packages/plugins/ratelimit-memory/src/index.test.ts | 253 | ||||
| -rw-r--r-- | packages/plugins/ratelimit-memory/src/index.ts | 86 | ||||
| -rw-r--r-- | packages/shared-server/src/plugins.ts | 1 | ||||
| -rw-r--r-- | packages/shared/plugins.ts | 4 | ||||
| -rw-r--r-- | packages/shared/ratelimiting.ts | 38 | ||||
| -rw-r--r-- | packages/trpc/index.ts | 6 | ||||
| -rw-r--r-- | packages/trpc/lib/rateLimit.ts | 48 | ||||
| -rw-r--r-- | packages/trpc/package.json | 1 | ||||
| -rw-r--r-- | packages/trpc/rateLimit.ts | 72 | ||||
| -rw-r--r-- | pnpm-lock.yaml | 3 |
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 |
