feat: QPS Limit middleware (#2956)

* feat: QPS Limit middleware

* chore: use request-ip to get client ip

* feat: frequencyLimit schema

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,

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);
+    }
+  };
+}

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
+);

packages/service/common/system/frequencyLimit/type.d.ts

@@ -0,0 +1,6 @@
+export type FrequencyLimitSchemaType = {
+  _id: string;
+  eventId: string; // 事件ID
+  amount: number; // 当前数量
+  expiredTime: Date; // 什么时候过期，过期则重置
+};

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) {}
+};

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"
   }

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": "服务器暂时过载或正在维护",