feat: QPS Limit middleware (#2956)

* feat: QPS Limit middleware * chore: use request-ip to get client ip * feat: frequencyLimit schema
2025-07-23 13:03:50 +00:00 · 2024-10-25 10:08:59 +08:00
parent bb727b0710
commit 75494f8d01
8 changed files with 144 additions and 28 deletions
--- a/packages/global/common/error/errorCode.ts
+++ b/packages/global/common/error/errorCode.ts
@@ -19,6 +19,7 @@ export const ERROR_CODE: { [key: number]: string } = {
  406: i18nT('common:code_error.error_code.406'),
  410: i18nT('common:code_error.error_code.410'),
  422: i18nT('common:code_error.error_code.422'),
+  429: i18nT('common:code_error.error_code.429'),
  500: i18nT('common:code_error.error_code.500'),
  502: i18nT('common:code_error.error_code.502'),
  503: i18nT('common:code_error.error_code.503'),
@@ -39,7 +40,8 @@ export enum ERROR_ENUM {
  insufficientQuota = 'insufficientQuota',
  unAuthModel = 'unAuthModel',
  unAuthApiKey = 'unAuthApiKey',
-  unAuthFile = 'unAuthFile'
+  unAuthFile = 'unAuthFile',
+  QPSLimitExceed = 'QPSLimitExceed'
 }

 export type ErrType<T> = Record<
@@ -67,6 +69,12 @@ export const ERROR_RESPONSE: Record<
    message: i18nT('common:code_error.error_message.403'),
    data: null
  },
+  [ERROR_ENUM.QPSLimitExceed]: {
+    code: 429,
+    statusText: ERROR_ENUM.QPSLimitExceed,
+    message: i18nT('common:code_error.error_code.429'),
+    data: null
+  },
  [ERROR_ENUM.insufficientQuota]: {
    code: 510,
    statusText: ERROR_ENUM.insufficientQuota,
--- a/packages/service/common/middle/qpsLimit.ts
+++ b/packages/service/common/middle/qpsLimit.ts
@@ -0,0 +1,26 @@
+import { ApiRequestProps } from 'type/next';
+import requestIp from 'request-ip';
+import { ERROR_ENUM } from '@fastgpt/global/common/error/errorCode';
+import { authFrequencyLimit } from 'common/system/frequencyLimit/utils';
+import { addSeconds } from 'date-fns';
+
+// unit: times/s
+// how to use?
+// export default NextAPI(useQPSLimit(10), handler); // limit 10 times per second for a ip
+export function useQPSLimit(limit: number) {
+  return async (req: ApiRequestProps) => {
+    const ip = requestIp.getClientIp(req);
+    if (!ip) {
+      return;
+    }
+    try {
+      await authFrequencyLimit({
+        eventId: 'ip-qps-limit' + ip,
+        maxAmount: limit,
+        expiredTime: addSeconds(new Date(), 1)
+      });
+    } catch (_) {
+      return Promise.reject(ERROR_ENUM.QPSLimitExceed);
+    }
+  };
+}
--- a/packages/service/common/system/frequencyLimit/schema.ts
+++ b/packages/service/common/system/frequencyLimit/schema.ts
@@ -0,0 +1,27 @@
+import { getMongoModel, Schema } from '../../mongo';
+import type { FrequencyLimitSchemaType } from './type';
+
+const FrequencyLimitSchema = new Schema({
+  eventId: {
+    type: String,
+    required: true
+  },
+  amount: {
+    type: Number,
+    default: 0
+  },
+  expiredTime: {
+    type: Date,
+    required: true
+  }
+});
+
+try {
+  FrequencyLimitSchema.index({ eventId: 1 }, { unique: true });
+  FrequencyLimitSchema.index({ expiredTime: 1 }, { expireAfterSeconds: 0 });
+} catch (error) {}
+
+export const MongoFrequencyLimit = getMongoModel<FrequencyLimitSchemaType>(
+  'frequency_limit',
+  FrequencyLimitSchema
+);
--- a/packages/service/common/system/frequencyLimit/type.d.ts
+++ b/packages/service/common/system/frequencyLimit/type.d.ts
@@ -0,0 +1,6 @@
+export type FrequencyLimitSchemaType = {
+  _id: string;
+  eventId: string; // 事件ID
+  amount: number; // 当前数量
+  expiredTime: Date; // 什么时候过期，过期则重置
+};
--- a/packages/service/common/system/frequencyLimit/utils.ts
+++ b/packages/service/common/system/frequencyLimit/utils.ts
@@ -0,0 +1,32 @@
+import { AuthFrequencyLimitProps } from '@fastgpt/global/common/frequenctLimit/type';
+import { MongoFrequencyLimit } from './schema';
+import { readFromSecondary } from '../../mongo/utils';
+
+export const authFrequencyLimit = async ({
+  eventId,
+  maxAmount,
+  expiredTime
+}: AuthFrequencyLimitProps) => {
+  try {
+    // 对应 eventId 的 account+1, 不存在的话，则创建一个
+    const result = await MongoFrequencyLimit.findOneAndUpdate(
+      {
+        eventId
+      },
+      {
+        $inc: { amount: 1 },
+        $setOnInsert: { expiredTime }
+      },
+      {
+        upsert: true,
+        new: true,
+        ...readFromSecondary
+      }
+    );
+
+    // 因为始终会返回+1的结果，所以这里不能直接等，需要多一个。
+    if (result.amount > maxAmount) {
+      return Promise.reject(result);
+    }
+  } catch (error) {}
+};
--- a/packages/service/package.json
+++ b/packages/service/package.json
@@ -33,6 +33,7 @@
    "papaparse": "5.4.1",
    "pdfjs-dist": "4.4.168",
    "pg": "^8.10.0",
+    "request-ip": "^3.3.0",
    "tiktoken": "^1.0.15",
    "tunnel": "^0.0.6",
    "turndown": "^7.1.2"
@@ -46,6 +47,7 @@
    "@types/node-cron": "^3.0.11",
    "@types/papaparse": "5.3.7",
    "@types/pg": "^8.6.6",
+    "@types/request-ip": "^0.0.37",
    "@types/tunnel": "^0.0.4",
    "@types/turndown": "^5.0.4"
  }
--- a/packages/web/i18n/zh/common.json
+++ b/packages/web/i18n/zh/common.json
@@ -49,6 +49,7 @@
  "code_error.error_code.406": "请求格式错误",
  "code_error.error_code.410": "资源已删除",
  "code_error.error_code.422": "验证错误",
+  "code_error.error_code.429": "请求过于频繁",
  "code_error.error_code.500": "服务器发生错误",
  "code_error.error_code.502": "网关错误",
  "code_error.error_code.503": "服务器暂时过载或正在维护",
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -220,6 +220,9 @@ importers:
      pg:
        specifier: ^8.10.0
        version: 8.12.0
+      request-ip:
+        specifier: ^3.3.0
+        version: 3.3.0
      tiktoken:
        specifier: ^1.0.15
        version: 1.0.15
@@ -254,6 +257,9 @@ importers:
      '@types/pg':
        specifier: ^8.6.6
        version: 8.11.6
+      '@types/request-ip':
+        specifier: ^0.0.37
+        version: 0.0.37
      '@types/tunnel':
        specifier: ^0.0.4
        version: 0.0.4
@@ -554,7 +560,7 @@ importers:
        version: 1.77.8
      ts-jest:
        specifier: ^29.1.0
-        version: 29.2.2(@babel/core@7.24.9)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.24.9))(jest@29.7.0(@types/node@20.14.11)(babel-plugin-macros@3.1.0)(ts-node@10.9.2(@types/node@20.14.11)(typescript@5.5.3)))(typescript@5.5.3)
+        version: 29.2.2(@babel/core@7.24.9)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.24.9))(jest@29.7.0(@types/node@20.14.11)(babel-plugin-macros@3.1.0))(typescript@5.5.3)
      use-context-selector:
        specifier: ^1.4.4
        version: 1.4.4(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(scheduler@0.23.2)
@@ -688,7 +694,7 @@ importers:
        version: 6.3.4
      ts-jest:
        specifier: ^29.1.0
-        version: 29.2.2(@babel/core@7.24.9)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.24.9))(jest@29.7.0(@types/node@20.14.11)(babel-plugin-macros@3.1.0)(ts-node@10.9.2(@types/node@20.14.11)(typescript@5.5.3)))(typescript@5.5.3)
+        version: 29.2.2(@babel/core@7.24.9)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.24.9))(jest@29.7.0(@types/node@20.14.11)(babel-plugin-macros@3.1.0))(typescript@5.5.3)
      ts-loader:
        specifier: ^9.4.3
        version: 9.5.1(typescript@5.5.3)(webpack@5.92.1)
@@ -1985,7 +1991,7 @@ packages:
  '@emotion/use-insertion-effect-with-fallbacks@1.0.1':
    resolution: {integrity: sha512-jT/qyKZ9rzLErtrjGgdkMBn2OP8wl0G3sQlBb3YPryvKHsjvINUhVaPFfP+fpBcOkmrVOVEEHQFJ7nbj2TH2gw==}
    peerDependencies:
-      react: '>=16.8.0'
+      react: 18.3.1

  '@emotion/utils@1.2.1':
    resolution: {integrity: sha512-Y2tGf3I+XVnajdItskUCn6LX+VUDmP6lTL4fcqsXAv43dnlbZiuW4MWQW38rW/BVWSE7Q/7+XQocmpnRYILUmg==}
@@ -2604,8 +2610,8 @@ packages:
    resolution: {integrity: sha512-RFkU9/i7cN2bsq/iTkurMWOEErmYcY6JiQI3Jn+WeR/FGISH8JbHERjpS9oRuSOPvDMJI0Z8nJeKkbOs9sBYQw==}
    peerDependencies:
      monaco-editor: '>= 0.25.0 < 1'
-      react: ^16.8.0 || ^17.0.0 || ^18.0.0
-      react-dom: ^16.8.0 || ^17.0.0 || ^18.0.0
+      react: 18.3.1
+      react-dom: 18.3.1

  '@mongodb-js/saslprep@1.1.7':
    resolution: {integrity: sha512-dCHW/oEX0KJ4NjDULBo3JiOaK5+6axtpBbS+ao2ZInoAL9/YRQLhXzSNAFz7hP4nzLkIqsfYAK/PDE3+XHny0Q==}
@@ -2956,8 +2962,8 @@ packages:
  '@reactflow/node-resizer@2.2.14':
    resolution: {integrity: sha512-fwqnks83jUlYr6OHcdFEedumWKChTHRGw/kbCxj0oqBd+ekfs+SIp4ddyNU0pdx96JIm5iNFS0oNrmEiJbbSaA==}
    peerDependencies:
-      react: '>=17'
-      react-dom: '>=17'
+      react: 18.3.1
+      react-dom: 18.3.1

  '@reactflow/node-toolbar@1.3.14':
    resolution: {integrity: sha512-rbynXQnH/xFNu4P9H+hVqlEUafDCkEoCy0Dg9mG22Sg+rY/0ck6KkrAQrYrTgXusd+cEJOMK0uOOFCK2/5rSGQ==}
@@ -3449,6 +3455,9 @@ packages:
  '@types/react-syntax-highlighter@15.5.13':
    resolution: {integrity: sha512-uLGJ87j6Sz8UaBAooU0T6lWJ0dBmjZgN1PZTrj05TNql2/XpC6+4HhMT5syIdFUUt+FASfCeLLv4kBygNU+8qA==}

+  '@types/react@18.3.0':
+    resolution: {integrity: sha512-DiUcKjzE6soLyln8NNZmyhcQjVv+WsUIFSqetMN0p8927OztKT4VTfFTqsbAi5oAGIcgOmOajlfBqyptDDjZRw==}
+
  '@types/react@18.3.1':
    resolution: {integrity: sha512-V0kuGBX3+prX+DQ/7r2qsv1NsdfnCLnTgnRJ1pYnxykBhGMz+qj+box5lq7XsO5mtZsBqpjwwTu/7wszPfMBcw==}

@@ -7045,8 +7054,8 @@ packages:
    peerDependencies:
      '@opentelemetry/api': ^1.1.0
      '@playwright/test': ^1.41.2
-      react: ^18.2.0
-      react-dom: ^18.2.0
+      react: 18.3.1
+      react-dom: 18.3.1
      sass: ^1.3.0
    peerDependenciesMeta:
      '@opentelemetry/api':
@@ -7698,8 +7707,8 @@ packages:
  react-photo-view@1.2.6:
    resolution: {integrity: sha512-Fq17yxkMIv0oFp7HOJr39HgCZRP6A9K5T5rixJ4flSUYT2OO3V8vNxEExjhIKgIrfmTu+mDnHYEsI9RRWi1JHw==}
    peerDependencies:
-      react: '>=16.8.0'
-      react-dom: '>=16.8.0'
+      react: 18.3.1
+      react-dom: 18.3.1

  react-redux@7.2.9:
    resolution: {integrity: sha512-Gx4L3uM182jEEayZfRbI/G11ZpYdNAnBs70lFVMNdHJI76XYtR+7m0MN+eAs7UHBPhWXcnFPaS+9owSCJQHNpQ==}
@@ -7717,8 +7726,8 @@ packages:
    resolution: {integrity: sha512-DtSYaao4mBmX+HDo5YWYdBWQwYIQQshUV/dVxFxK+KM26Wjwp1gZ6rv6OC3oujI6Bfu6Xyg3TwK533AQutsn/g==}
    engines: {node: '>=10'}
    peerDependencies:
-      '@types/react': ^16.8.0 || ^17.0.0 || ^18.0.0
-      react: ^16.8.0 || ^17.0.0 || ^18.0.0
+      '@types/react': 18.3.1
+      react: 18.3.1
    peerDependenciesMeta:
      '@types/react':
        optional: true
@@ -7737,8 +7746,8 @@ packages:
    resolution: {integrity: sha512-ZWj0fHEMyWkHzKYUr2Bs/4zU6XLmq9HsgBURm7g5pAVfyn49DgUiNgY2d4lXRlYSiCif9YBGpQleewkcqddc7g==}
    engines: {node: '>=10'}
    peerDependencies:
-      '@types/react': ^16.8.0 || ^17.0.0 || ^18.0.0
-      react: ^16.8.0 || ^17.0.0 || ^18.0.0
+      '@types/react': 18.3.1
+      react: 18.3.1
    peerDependenciesMeta:
      '@types/react':
        optional: true
@@ -8753,8 +8762,8 @@ packages:
    resolution: {integrity: sha512-elOQwe6Q8gqZgDA8mrh44qRTQqpIHDcZ3hXTLjBe1i4ph8XpNJnO+aQf3NaG+lriLopI4HMx9VjQLfPQ6vhnoA==}
    engines: {node: '>=10'}
    peerDependencies:
-      '@types/react': ^16.8.0 || ^17.0.0 || ^18.0.0
-      react: ^16.8.0 || ^17.0.0 || ^18.0.0
+      '@types/react': 18.3.1
+      react: 18.3.1
    peerDependenciesMeta:
      '@types/react':
        optional: true
@@ -8804,8 +8813,8 @@ packages:
    resolution: {integrity: sha512-epTbsLuzZ7lPClpz2TyryBfztm7m+28DlEv2ZCQ3MDr5ssiwyOwGH/e5F9CkfWjJ1t4clvI58yF822/GUkjjhw==}
    engines: {node: '>=10'}
    peerDependencies:
-      '@types/react': ^16.9.0 || ^17.0.0 || ^18.0.0
-      react: ^16.8.0 || ^17.0.0 || ^18.0.0
+      '@types/react': 18.3.1
+      react: 18.3.1
    peerDependenciesMeta:
      '@types/react':
        optional: true
@@ -12464,7 +12473,7 @@ snapshots:
  '@types/react-redux@7.1.33':
    dependencies:
      '@types/hoist-non-react-statics': 3.3.5
-      '@types/react': 18.3.1
+      '@types/react': 18.3.0
      hoist-non-react-statics: 3.3.2
      redux: 4.2.1

@@ -12472,6 +12481,11 @@ snapshots:
    dependencies:
      '@types/react': 18.3.1

+  '@types/react@18.3.0':
+    dependencies:
+      '@types/prop-types': 15.7.12
+      csstype: 3.1.3
+
  '@types/react@18.3.1':
    dependencies:
      '@types/prop-types': 15.7.12
@@ -14370,7 +14384,7 @@ snapshots:
      eslint: 8.56.0
      eslint-import-resolver-node: 0.3.9
      eslint-import-resolver-typescript: 3.6.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1(eslint@8.56.0))(eslint@8.56.0)
-      eslint-plugin-import: 2.29.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1(eslint@8.56.0))(eslint@8.56.0))(eslint@8.56.0)
+      eslint-plugin-import: 2.29.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-typescript@3.6.1)(eslint@8.56.0)
      eslint-plugin-jsx-a11y: 6.9.0(eslint@8.56.0)
      eslint-plugin-react: 7.34.4(eslint@8.56.0)
      eslint-plugin-react-hooks: 4.6.2(eslint@8.56.0)
@@ -14393,8 +14407,8 @@ snapshots:
      debug: 4.3.5
      enhanced-resolve: 5.17.0
      eslint: 8.56.0
-      eslint-module-utils: 2.8.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1(eslint@8.56.0))(eslint@8.56.0))(eslint@8.56.0)
-      eslint-plugin-import: 2.29.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1(eslint@8.56.0))(eslint@8.56.0))(eslint@8.56.0)
+      eslint-module-utils: 2.8.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1)(eslint@8.56.0)
+      eslint-plugin-import: 2.29.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-typescript@3.6.1)(eslint@8.56.0)
      fast-glob: 3.3.2
      get-tsconfig: 4.7.5
      is-core-module: 2.14.0
@@ -14405,7 +14419,7 @@ snapshots:
      - eslint-import-resolver-webpack
      - supports-color

-  eslint-module-utils@2.8.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1(eslint@8.56.0))(eslint@8.56.0))(eslint@8.56.0):
+  eslint-module-utils@2.8.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1)(eslint@8.56.0):
    dependencies:
      debug: 3.2.7
    optionalDependencies:
@@ -14416,7 +14430,7 @@ snapshots:
    transitivePeerDependencies:
      - supports-color

-  eslint-plugin-import@2.29.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1(eslint@8.56.0))(eslint@8.56.0))(eslint@8.56.0):
+  eslint-plugin-import@2.29.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-typescript@3.6.1)(eslint@8.56.0):
    dependencies:
      array-includes: 3.1.8
      array.prototype.findlastindex: 1.2.5
@@ -14426,7 +14440,7 @@ snapshots:
      doctrine: 2.1.0
      eslint: 8.56.0
      eslint-import-resolver-node: 0.3.9
-      eslint-module-utils: 2.8.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1(eslint@8.56.0))(eslint@8.56.0))(eslint@8.56.0)
+      eslint-module-utils: 2.8.1(@typescript-eslint/parser@6.21.0(eslint@8.56.0)(typescript@5.5.3))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1)(eslint@8.56.0)
      hasown: 2.0.2
      is-core-module: 2.14.0
      is-glob: 4.0.3
@@ -18818,7 +18832,7 @@ snapshots:

  ts-dedent@2.2.0: {}

-  ts-jest@29.2.2(@babel/core@7.24.9)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.24.9))(jest@29.7.0(@types/node@20.14.11)(babel-plugin-macros@3.1.0)(ts-node@10.9.2(@types/node@20.14.11)(typescript@5.5.3)))(typescript@5.5.3):
+  ts-jest@29.2.2(@babel/core@7.24.9)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.24.9))(jest@29.7.0(@types/node@20.14.11)(babel-plugin-macros@3.1.0))(typescript@5.5.3):
    dependencies:
      bs-logger: 0.2.6
      ejs: 3.1.10