@omni-api/plugin-ratelimit 0.0.1

package/README.md

@@ -0,0 +1,25 @@
+# @omni/plugin-ratelimit
+
+OmniAPI 限流中间件。固定窗口计数算法，可插拔后端（默认内存）。
+
+## 用法
+```ts
+import { rateLimit, createMemoryStore } from '@omni/plugin-ratelimit';
+
+defineProcedure({
+  name: 'order.create',
+  middleware: [rateLimit({ rpm: 30 })],     // 每分钟 30 次
+  // ...
+});
+
+// 自定义 key（按租户）
+rateLimit({ limit: 100, windowMs: 60_000, key: (ctx) => ctx.state.tenant as string });
+
+// 自定义 store（如 Redis）
+rateLimit({ rpm: 100, store: myRedisStore });
+```
+
+## 注意
+- 默认共享 module-level 内存 store —— 单进程足够；多实例部署需替换为分布式 store。
+- 算法是固定窗口，不保证严格平滑；需要平滑请用令牌桶/漏桶（社区扩展）。
+- key 自动加上 `procedureName:` 前缀，不同 procedure 互不影响。

package/dist/index.d.ts

@@ -0,0 +1,78 @@
+import { Context, Middleware } from '@omni-api/core';
+
+/**
+ * 限流后端存储接口（可替换：内存 / Redis / Memcached…）
+ *
+ * incr 在原子操作中：
+ *   1. 增加计数器
+ *   2. 若是首次（count==1），设置 TTL = windowMs
+ *   3. 返回最新计数 + 剩余 TTL
+ */
+interface RateLimitStore {
+    /** 增加 1，返回 [新计数, 剩余窗口毫秒数] */
+    incr(key: string, windowMs: number): Promise<{
+        count: number;
+        resetMs: number;
+    }>;
+    /** 重置某 key（可选实现） */
+    reset?(key: string): Promise<void>;
+}
+/**
+ * 单进程内存版 Store —— 默认实现。
+ *
+ * 适合单实例 / 开发期；分布式部署应替换为 Redis 实现。
+ *
+ * 自带懒清理：incr 时若 key 已过期则重置；不需要后台 timer。
+ */
+declare function createMemoryStore(): RateLimitStore;
+
+interface RateLimitOptions {
+    /** 每分钟请求数。与 limit/windowMs 互斥；二选一 */
+    rpm?: number;
+    /** 窗口内最大请求数（与 windowMs 配合）。优先级高于 rpm */
+    limit?: number;
+    /** 窗口时长（毫秒）。默认 60_000 */
+    windowMs?: number;
+    /** key 生成函数。默认 user.id ?? 'anon' */
+    key?: (ctx: Context) => string;
+    /** 限流后端，默认共享 module-level 内存实例 */
+    store?: RateLimitStore;
+    /** 自定义错误消息 */
+    message?: string;
+}
+/**
+ * 限流中间件。
+ *
+ * 算法：固定窗口计数。简单、零开销；不保证严格平滑。
+ * 需要平滑限流（令牌桶/漏桶）请用 Redis 版扩展。
+ */
+declare function rateLimit(options?: RateLimitOptions): Middleware;
+
+/**
+ * Redis 客户端的最小接口。
+ *
+ * ioredis / node-redis@4 都满足该接口（参数顺序略有差异，下面适配）。
+ */
+interface RedisLike {
+    /**
+     * 执行 Lua 脚本。
+     *
+     * 不同客户端的 eval 签名不同，我们这里用最通用的形态：
+     *   - ioredis:  eval(script, numKeys, ...keys, ...args) → result
+     *   - node-redis@4: client.eval(script, { keys, arguments }) → result
+     * 我们要求传入的 `eval` 已经是 ioredis 风格的 wrapper（业务方包一下即可）。
+     */
+    eval(script: string, numKeys: number, ...keysAndArgs: (string | number)[]): Promise<unknown>;
+}
+interface RedisStoreOptions {
+    /** 客户端 */
+    redis: RedisLike;
+    /** key 前缀，避免和其他业务冲突。默认 'omni:rl:' */
+    keyPrefix?: string;
+}
+/**
+ * Redis 后端的限流 Store —— 真正的分布式限流（不像内存版会被 N 个 pod 放大）。
+ */
+declare function createRedisStore(options: RedisStoreOptions): RateLimitStore;
+
+export { type RateLimitOptions, type RateLimitStore, type RedisLike, type RedisStoreOptions, createMemoryStore, createRedisStore, rateLimit };

package/dist/index.js

@@ -0,0 +1,95 @@
+import { RateLimitError } from '@omni-api/core';
+
+// src/middleware.ts
+
+// src/store.ts
+function createMemoryStore() {
+  const map = /* @__PURE__ */ new Map();
+  return {
+    async incr(key, windowMs) {
+      const now = Date.now();
+      const entry = map.get(key);
+      if (!entry || entry.expiresAt <= now) {
+        const fresh = { count: 1, expiresAt: now + windowMs };
+        map.set(key, fresh);
+        return { count: 1, resetMs: windowMs };
+      }
+      entry.count += 1;
+      return { count: entry.count, resetMs: entry.expiresAt - now };
+    },
+    async reset(key) {
+      map.delete(key);
+    }
+  };
+}
+
+// src/middleware.ts
+var defaultStore = createMemoryStore();
+function rateLimit(options = {}) {
+  const windowMs = options.windowMs ?? 6e4;
+  const limit = options.limit ?? options.rpm ?? 60;
+  const store = options.store ?? defaultStore;
+  const keyFn = options.key ?? defaultKeyFn;
+  if (limit <= 0) throw new Error("rateLimit: `limit` must be > 0");
+  if (windowMs <= 0) throw new Error("rateLimit: `windowMs` must be > 0");
+  return async (ctx, next) => {
+    const procedureName = ctx.procedure?.name ?? "unknown";
+    const key = `${procedureName}:${keyFn(ctx)}`;
+    const { count, resetMs } = await store.incr(key, windowMs);
+    if (count > limit) {
+      throw new RateLimitError(options.message ?? "Too many requests", {
+        retryAfterMs: resetMs,
+        details: { limit, windowMs, current: count }
+      });
+    }
+    return next();
+  };
+}
+function defaultKeyFn(ctx) {
+  return ctx.user?.id ?? "anon";
+}
+
+// src/redis-store.ts
+var SCRIPT = `
+local count = redis.call('INCR', KEYS[1])
+if count == 1 then
+  redis.call('PEXPIRE', KEYS[1], ARGV[1])
+end
+local ttl = redis.call('PTTL', KEYS[1])
+if ttl < 0 then ttl = tonumber(ARGV[1]) end
+return {count, ttl}
+`;
+function createRedisStore(options) {
+  const { redis, keyPrefix = "omni:rl:" } = options;
+  return {
+    async incr(key, windowMs) {
+      const fullKey = keyPrefix + key;
+      try {
+        const result = await redis.eval(SCRIPT, 1, fullKey, windowMs);
+        const count = Number(Array.isArray(result) ? result[0] : 0);
+        const ttl = Number(Array.isArray(result) ? result[1] : windowMs);
+        return { count, resetMs: ttl > 0 ? ttl : windowMs };
+      } catch (err) {
+        console.error("[ratelimit] redis error, fail-open:", err);
+        return { count: 1, resetMs: windowMs };
+      }
+    },
+    async reset(key) {
+      const fullKey = keyPrefix + key;
+      try {
+        const r = redis;
+        if (typeof r.del === "function") {
+          await r.del(fullKey);
+        } else {
+          await redis.eval('return redis.call("DEL", KEYS[1])', 1, fullKey);
+        }
+      } catch (err) {
+        console.error("[ratelimit] redis reset error:", err);
+      }
+    }
+  };
+}
+
+export { createMemoryStore, createRedisStore, rateLimit };
+//# sourceMappingURL=index.js.map
+//# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -0,0 +1 @@
+{"version":3,"sources":["../src/store.ts","../src/middleware.ts","../src/redis-store.ts"],"names":[],"mappings":";;;;;AA4BO,SAAS,iBAAA,GAAoC;AAClD,EAAA,MAAM,GAAA,uBAAU,GAAA,EAAsB;AAEtC,EAAA,OAAO;AAAA,IACL,MAAM,IAAA,CAAK,GAAA,EAAK,QAAA,EAAU;AACxB,MAAA,MAAM,GAAA,GAAM,KAAK,GAAA,EAAI;AACrB,MAAA,MAAM,KAAA,GAAQ,GAAA,CAAI,GAAA,CAAI,GAAG,CAAA;AACzB,MAAA,IAAI,CAAC,KAAA,IAAS,KAAA,CAAM,SAAA,IAAa,GAAA,EAAK;AACpC,QAAA,MAAM,QAAkB,EAAE,KAAA,EAAO,CAAA,EAAG,SAAA,EAAW,MAAM,QAAA,EAAS;AAC9D,QAAA,GAAA,CAAI,GAAA,CAAI,KAAK,KAAK,CAAA;AAClB,QAAA,OAAO,EAAE,KAAA,EAAO,CAAA,EAAG,OAAA,EAAS,QAAA,EAAS;AAAA,MACvC;AACA,MAAA,KAAA,CAAM,KAAA,IAAS,CAAA;AACf,MAAA,OAAO,EAAE,KAAA,EAAO,KAAA,CAAM,OAAO,OAAA,EAAS,KAAA,CAAM,YAAY,GAAA,EAAI;AAAA,IAC9D,CAAA;AAAA,IACA,MAAM,MAAM,GAAA,EAAK;AACf,MAAA,GAAA,CAAI,OAAO,GAAG,CAAA;AAAA,IAChB;AAAA,GACF;AACF;;;ACzBA,IAAM,eAAe,iBAAA,EAAkB;AAQhC,SAAS,SAAA,CAAU,OAAA,GAA4B,EAAC,EAAe;AACpE,EAAA,MAAM,QAAA,GAAW,QAAQ,QAAA,IAAY,GAAA;AACrC,EAAA,MAAM,KAAA,GAAQ,OAAA,CAAQ,KAAA,IAAS,OAAA,CAAQ,GAAA,IAAO,EAAA;AAC9C,EAAA,MAAM,KAAA,GAAQ,QAAQ,KAAA,IAAS,YAAA;AAC/B,EAAA,MAAM,KAAA,GAAQ,QAAQ,GAAA,IAAO,YAAA;AAE7B,EAAA,IAAI,KAAA,IAAS,CAAA,EAAG,MAAM,IAAI,MAAM,gCAAgC,CAAA;AAChE,EAAA,IAAI,QAAA,IAAY,CAAA,EAAG,MAAM,IAAI,MAAM,mCAAmC,CAAA;AAEtE,EAAA,OAAO,OAAO,KAAK,IAAA,KAAS;AAC1B,IAAA,MAAM,aAAA,GAAgB,GAAA,CAAI,SAAA,EAAW,IAAA,IAAQ,SAAA;AAC7C,IAAA,MAAM,MAAM,CAAA,EAAG,aAAa,CAAA,CAAA,EAAI,KAAA,CAAM,GAAG,CAAC,CAAA,CAAA;AAC1C,IAAA,MAAM,EAAE,OAAO,OAAA,EAAQ,GAAI,MAAM,KAAA,CAAM,IAAA,CAAK,KAAK,QAAQ,CAAA;AAEzD,IAAA,IAAI,QAAQ,KAAA,EAAO;AACjB,MAAA,MAAM,IAAI,cAAA,CAAe,OAAA,CAAQ,OAAA,IAAW,mBAAA,EAAqB;AAAA,QAC/D,YAAA,EAAc,OAAA;AAAA,QACd,OAAA,EAAS,EAAE,KAAA,EAAO,QAAA,EAAU,SAAS,KAAA;AAAM,OAC5C,CAAA;AAAA,IACH;AAEA,IAAA,OAAO,IAAA,EAAK;AAAA,EACd,CAAA;AACF;AAEA,SAAS,aAAa,GAAA,EAAsB;AAC1C,EAAA,OAAO,GAAA,CAAI,MAAM,EAAA,IAAM,MAAA;AACzB;;;AC5BA,IAAM,MAAA,GAAS;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,CAAA;AAoBR,SAAS,iBAAiB,OAAA,EAA4C;AAC3E,EAAA,MAAM,EAAE,KAAA,EAAO,SAAA,GAAY,UAAA,EAAW,GAAI,OAAA;AAE1C,EAAA,OAAO;AAAA,IACL,MAAM,IAAA,CAAK,GAAA,EAAK,QAAA,EAAU;AACxB,MAAA,MAAM,UAAU,SAAA,GAAY,GAAA;AAC5B,MAAA,IAAI;AACF,QAAA,MAAM,SAAU,MAAM,KAAA,CAAM,KAAK,MAAA,EAAQ,CAAA,EAAG,SAAS,QAAQ,CAAA;AAE7D,QAAA,MAAM,KAAA,GAAQ,OAAO,KAAA,CAAM,OAAA,CAAQ,MAAM,CAAA,GAAI,MAAA,CAAO,CAAC,CAAA,GAAI,CAAC,CAAA;AAC1D,QAAA,MAAM,GAAA,GAAM,OAAO,KAAA,CAAM,OAAA,CAAQ,MAAM,CAAA,GAAI,MAAA,CAAO,CAAC,CAAA,GAAI,QAAQ,CAAA;AAC/D,QAAA,OAAO,EAAE,KAAA,EAAO,OAAA,EAAS,GAAA,GAAM,CAAA,GAAI,MAAM,QAAA,EAAS;AAAA,MACpD,SAAS,GAAA,EAAK;AAGZ,QAAA,OAAA,CAAQ,KAAA,CAAM,uCAAuC,GAAG,CAAA;AACxD,QAAA,OAAO,EAAE,KAAA,EAAO,CAAA,EAAG,OAAA,EAAS,QAAA,EAAS;AAAA,MACvC;AAAA,IACF,CAAA;AAAA,IACA,MAAM,MAAM,GAAA,EAAK;AACf,MAAA,MAAM,UAAU,SAAA,GAAY,GAAA;AAC5B,MAAA,IAAI;AAEF,QAAA,MAAM,CAAA,GAAI,KAAA;AACV,QAAA,IAAI,OAAO,CAAA,CAAE,GAAA,KAAQ,UAAA,EAAY;AAC/B,UAAA,MAAM,CAAA,CAAE,IAAI,OAAO,CAAA;AAAA,QACrB,CAAA,MAAO;AACL,UAAA,MAAM,KAAA,CAAM,IAAA,CAAK,mCAAA,EAAqC,CAAA,EAAG,OAAO,CAAA;AAAA,QAClE;AAAA,MACF,SAAS,GAAA,EAAK;AACZ,QAAA,OAAA,CAAQ,KAAA,CAAM,kCAAkC,GAAG,CAAA;AAAA,MACrD;AAAA,IACF;AAAA,GACF;AACF","file":"index.js","sourcesContent":["/**\n * 限流后端存储接口（可替换：内存 / Redis / Memcached…）\n *\n * incr 在原子操作中：\n *   1. 增加计数器\n *   2. 若是首次（count==1），设置 TTL = windowMs\n *   3. 返回最新计数 + 剩余 TTL\n */\nexport interface RateLimitStore {\n  /** 增加 1，返回 [新计数, 剩余窗口毫秒数] */\n  incr(key: string, windowMs: number): Promise<{ count: number; resetMs: number }>;\n  /** 重置某 key（可选实现） */\n  reset?(key: string): Promise<void>;\n}\n\ninterface MemEntry {\n  count: number;\n  /** 窗口结束时刻（绝对毫秒） */\n  expiresAt: number;\n}\n\n/**\n * 单进程内存版 Store —— 默认实现。\n *\n * 适合单实例 / 开发期；分布式部署应替换为 Redis 实现。\n *\n * 自带懒清理：incr 时若 key 已过期则重置；不需要后台 timer。\n */\nexport function createMemoryStore(): RateLimitStore {\n  const map = new Map<string, MemEntry>();\n\n  return {\n    async incr(key, windowMs) {\n      const now = Date.now();\n      const entry = map.get(key);\n      if (!entry || entry.expiresAt <= now) {\n        const fresh: MemEntry = { count: 1, expiresAt: now + windowMs };\n        map.set(key, fresh);\n        return { count: 1, resetMs: windowMs };\n      }\n      entry.count += 1;\n      return { count: entry.count, resetMs: entry.expiresAt - now };\n    },\n    async reset(key) {\n      map.delete(key);\n    },\n  };\n}\n","import { RateLimitError, type Context, type Middleware } from '@omni-api/core';\nimport { createMemoryStore, type RateLimitStore } from './store.js';\n\nexport interface RateLimitOptions {\n  /** 每分钟请求数。与 limit/windowMs 互斥；二选一 */\n  rpm?: number;\n  /** 窗口内最大请求数（与 windowMs 配合）。优先级高于 rpm */\n  limit?: number;\n  /** 窗口时长（毫秒）。默认 60_000 */\n  windowMs?: number;\n\n  /** key 生成函数。默认 user.id ?? 'anon' */\n  key?: (ctx: Context) => string;\n\n  /** 限流后端，默认共享 module-level 内存实例 */\n  store?: RateLimitStore;\n\n  /** 自定义错误消息 */\n  message?: string;\n}\n\n/** module-level 共享 store，避免每次调用 rateLimit() 都新建一个独立的 map */\nconst defaultStore = createMemoryStore();\n\n/**\n * 限流中间件。\n *\n * 算法：固定窗口计数。简单、零开销；不保证严格平滑。\n * 需要平滑限流（令牌桶/漏桶）请用 Redis 版扩展。\n */\nexport function rateLimit(options: RateLimitOptions = {}): Middleware {\n  const windowMs = options.windowMs ?? 60_000;\n  const limit = options.limit ?? options.rpm ?? 60;\n  const store = options.store ?? defaultStore;\n  const keyFn = options.key ?? defaultKeyFn;\n\n  if (limit <= 0) throw new Error('rateLimit: `limit` must be > 0');\n  if (windowMs <= 0) throw new Error('rateLimit: `windowMs` must be > 0');\n\n  return async (ctx, next) => {\n    const procedureName = ctx.procedure?.name ?? 'unknown';\n    const key = `${procedureName}:${keyFn(ctx)}`;\n    const { count, resetMs } = await store.incr(key, windowMs);\n\n    if (count > limit) {\n      throw new RateLimitError(options.message ?? 'Too many requests', {\n        retryAfterMs: resetMs,\n        details: { limit, windowMs, current: count },\n      });\n    }\n\n    return next();\n  };\n}\n\nfunction defaultKeyFn(ctx: Context): string {\n  return ctx.user?.id ?? 'anon';\n}\n","import type { RateLimitStore } from './store.js';\n\n/**\n * Redis 客户端的最小接口。\n *\n * ioredis / node-redis@4 都满足该接口（参数顺序略有差异，下面适配）。\n */\nexport interface RedisLike {\n  /**\n   * 执行 Lua 脚本。\n   *\n   * 不同客户端的 eval 签名不同，我们这里用最通用的形态：\n   *   - ioredis:  eval(script, numKeys, ...keys, ...args) → result\n   *   - node-redis@4: client.eval(script, { keys, arguments }) → result\n   * 我们要求传入的 `eval` 已经是 ioredis 风格的 wrapper（业务方包一下即可）。\n   */\n  eval(script: string, numKeys: number, ...keysAndArgs: (string | number)[]): Promise<unknown>;\n}\n\n/**\n * 固定窗口限流的 Lua 脚本（原子操作）。\n *\n * 步骤：\n *   1. INCR key\n *   2. 如果是首次（count == 1），PEXPIRE key windowMs\n *   3. PTTL 返回剩余 ms\n *\n * 返回值：[count, remainingMs]\n */\nconst SCRIPT = `\nlocal count = redis.call('INCR', KEYS[1])\nif count == 1 then\n  redis.call('PEXPIRE', KEYS[1], ARGV[1])\nend\nlocal ttl = redis.call('PTTL', KEYS[1])\nif ttl < 0 then ttl = tonumber(ARGV[1]) end\nreturn {count, ttl}\n`;\n\nexport interface RedisStoreOptions {\n  /** 客户端 */\n  redis: RedisLike;\n  /** key 前缀，避免和其他业务冲突。默认 'omni:rl:' */\n  keyPrefix?: string;\n}\n\n/**\n * Redis 后端的限流 Store —— 真正的分布式限流（不像内存版会被 N 个 pod 放大）。\n */\nexport function createRedisStore(options: RedisStoreOptions): RateLimitStore {\n  const { redis, keyPrefix = 'omni:rl:' } = options;\n\n  return {\n    async incr(key, windowMs) {\n      const fullKey = keyPrefix + key;\n      try {\n        const result = (await redis.eval(SCRIPT, 1, fullKey, windowMs)) as [number, number];\n        // ioredis 返回 [Number, Number]；node-redis 同；保险起见 toNumber\n        const count = Number(Array.isArray(result) ? result[0] : 0);\n        const ttl = Number(Array.isArray(result) ? result[1] : windowMs);\n        return { count, resetMs: ttl > 0 ? ttl : windowMs };\n      } catch (err) {\n        // Redis 故障时降级：不阻塞业务（fail-open）\n        // 调用方可能想 fail-close，但默认 open 更安全 —— 否则 redis 挂导致业务全挂\n        console.error('[ratelimit] redis error, fail-open:', err);\n        return { count: 1, resetMs: windowMs };\n      }\n    },\n    async reset(key) {\n      const fullKey = keyPrefix + key;\n      try {\n        // 直接 DEL；用 eval 也行，但 DEL 简单且通常各客户端都暴露\n        const r = redis as unknown as { del?: (k: string) => Promise<number> };\n        if (typeof r.del === 'function') {\n          await r.del(fullKey);\n        } else {\n          await redis.eval('return redis.call(\"DEL\", KEYS[1])', 1, fullKey);\n        }\n      } catch (err) {\n        console.error('[ratelimit] redis reset error:', err);\n      }\n    },\n  };\n}\n"]}

package/package.json

@@ -0,0 +1,36 @@
+{
+  "name": "@omni-api/plugin-ratelimit",
+  "version": "0.0.1",
+  "description": "Rate limit middleware for OmniAPI (in-memory, pluggable store)",
+  "type": "module",
+  "main": "./dist/index.js",
+  "module": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
+  "files": ["dist", "README.md"],
+  "scripts": {
+    "build": "tsup",
+    "test": "vitest run",
+    "test:watch": "vitest",
+    "typecheck": "tsc --noEmit",
+    "clean": "rm -rf dist .turbo coverage"
+  },
+  "dependencies": {
+    "@omni-api/core": "workspace:*"
+  },
+  "devDependencies": {
+    "tsup": "^8.3.0",
+    "typescript": "^5.7.0",
+    "vitest": "^2.1.0",
+    "zod": "^3.23.8",
+    "@types/node": "^22.10.0"
+  },
+  "publishConfig": {
+    "access": "public"
+  }
+}