feat: 新增限流功能 (Chanzhaoyu#718)

* 请求速率限制

* perf: 优化代码

---------

Co-authored-by: ChenZhaoYu <[email protected]>

README.md

@@ -170,6 +170,7 @@ pnpm dev
 通用：
 
 - `AUTH_SECRET_KEY` 访问权限密钥，可选
+- `MAX_REQUEST_PER_HOUR` 每小时最大请求次数，可选，默认无限
 - `TIMEOUT_MS` 超时，单位毫秒，可选
 - `SOCKS_PROXY_HOST` 和 `SOCKS_PROXY_PORT` 一起时生效，可选
 - `SOCKS_PROXY_PORT` 和 `SOCKS_PROXY_HOST` 一起时生效，可选
@@ -224,6 +225,8 @@ services:
       API_REVERSE_PROXY: xxx
       # 访问权限密钥，可选
       AUTH_SECRET_KEY: xxx
+      # 每小时最大请求次数，可选，默认无限
+      MAX_REQUEST_PER_HOUR: 0
       # 超时，单位毫秒，可选
       TIMEOUT_MS: 60000
       # Socks代理，可选，和 SOCKS_PROXY_PORT 一起时生效
@@ -245,6 +248,7 @@ services:
 | --------------------- | ---------------------- | -------------------------------------------------------------------------------------------------- |
 | `PORT`                | 必填                   | 默认 `3002`
 | `AUTH_SECRET_KEY`          | 可选                   | 访问权限密钥                                        |
+| `MAX_REQUEST_PER_HOUR`          | 可选                   | 每小时最大请求次数，可选，默认无限                                        |
 | `TIMEOUT_MS`          | 可选                   | 超时时间，单位毫秒                                                                             |
 | `OPENAI_API_KEY`      | `OpenAI API` 二选一    | 使用 `OpenAI API` 所需的 `apiKey` [(获取 apiKey)](https://platform.openai.com/overview)            |
 | `OPENAI_ACCESS_TOKEN` | `Web API` 二选一       | 使用 `Web API` 所需的 `accessToken` [(获取 accessToken)](https://chat.openai.com/api/auth/session) |

docker-compose/docker-compose.yml

@@ -20,6 +20,8 @@ services:
       API_REVERSE_PROXY: xxx
       # 访问jwt加密参数，可选 不为空则允许登录 同时需要设置 MONGODB_URL
       AUTH_SECRET_KEY: xxx
+      # 每小时最大请求次数，可选，默认无限
+      MAX_REQUEST_PER_HOUR: 0
       # 超时，单位毫秒，可选
       TIMEOUT_MS: 60000
       # Socks代理，可选，和 SOCKS_PROXY_PORT 一起时生效

service/.env.example

@@ -59,4 +59,7 @@ SMTP_TSL=true
 SMTP_USERNAME=[email protected]
 SMTP_PASSWORD=yourpassword
 
-# ----- Only valid after setting AUTH_SECRET_KEY end ----
+# ----- Only valid after setting AUTH_SECRET_KEY end ----
+# Rate Limit
+MAX_REQUEST_PER_HOUR=
+

service/package.json

@@ -29,6 +29,7 @@
     "dotenv": "^16.0.3",
     "esno": "^0.16.3",
     "express": "^4.18.2",
+    "express-rate-limit": "^6.7.0",
     "https-proxy-agent": "^5.0.1",
     "isomorphic-fetch": "^3.0.0",
     "mongodb": "^5.1.0",

service/src/index.ts

@@ -8,6 +8,9 @@ import { Status } from './storage/model'
 import { clearChat, createChatRoom, createUser, deleteChat, deleteChatRoom, existsChatRoom, getChat, getChatRooms, getChats, getUser, insertChat, renameChatRoom, updateChat, verifyUser } from './storage/mongo'
 import { sendMail } from './utils/mail'
 import { checkUserVerify, getUserVerifyUrl, md5 } from './utils/security'
+import { limiter } from './middleware/limiter'
+import { isNotEmptyString } from './utils/is'
+
 
 const app = express()
 const router = express.Router()
@@ -147,7 +150,9 @@ router.post('/chat', auth, async (req, res) => {
   }
 })
 
-router.post('/chat-process', auth, async (req, res) => {
+
+router.post('/chat-process', [auth, limiter], async (req, res) => {
+
   res.setHeader('Content-type', 'application/octet-stream')
 
   try {

service/src/middleware/limiter.ts

@@ -0,0 +1,19 @@
+import { rateLimit } from 'express-rate-limit'
+import { isNotEmptyString } from '../utils/is'
+
+const MAX_REQUEST_PER_HOUR = process.env.MAX_REQUEST_PER_HOUR
+
+const maxCount = (isNotEmptyString(MAX_REQUEST_PER_HOUR) && !isNaN(Number(MAX_REQUEST_PER_HOUR)))
+  ? parseInt(MAX_REQUEST_PER_HOUR)
+  : 0 // 0 means unlimited
+
+const limiter = rateLimit({
+  windowMs: 60 * 60 * 1000, // Maximum number of accesses within an hour
+  max: maxCount,
+  statusCode: 200, // 200 means success，but the message is 'Too many request from this IP in 1 hour'
+  message: async (req, res) => {
+    res.send({ status: 'Fail', message: 'Too many request from this IP in 1 hour', data: null })
+  },
+})
+
+export { limiter }