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

* 请求速率限制 * perf: 优化代码 --------- Co-authored-by: ChenZhaoYu <[email protected]>
AstraSurge · Mar 21, 2023 · 9baf59a · 9baf59a
1 parent 2cf7a52
commit 9baf59a
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 1 deletion.
diff --git a/README.md b/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) |

diff --git a/docker-compose/docker-compose.yml b/docker-compose/docker-compose.yml
@@ -18,6 +18,8 @@ services:
       API_REVERSE_PROXY: xxx
       # 访问权限密钥，可选
       AUTH_SECRET_KEY: xxx
+      # 每小时最大请求次数，可选，默认无限
+      MAX_REQUEST_PER_HOUR: 0
       # 超时，单位毫秒，可选
       TIMEOUT_MS: 60000
       # Socks代理，可选，和 SOCKS_PROXY_PORT 一起时生效

diff --git a/service/.env.example b/service/.env.example
@@ -27,3 +27,6 @@ SOCKS_PROXY_PORT=
 
 # HTTPS PROXY
 HTTPS_PROXY=
+
+# Rate Limit
+MAX_REQUEST_PER_HOUR=
diff --git a/service/package.json b/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",
     "node-fetch": "^3.3.0",

diff --git a/service/pnpm-lock.yaml b/service/pnpm-lock.yaml
diff --git a/service/src/index.ts b/service/src/index.ts
@@ -2,6 +2,7 @@ import express from 'express'
 import type { ChatContext, ChatMessage } from './chatgpt'
 import { chatConfig, chatReplyProcess, currentModel } from './chatgpt'
 import { auth } from './middleware/auth'
+import { limiter } from './middleware/limiter'
 import { isNotEmptyString } from './utils/is'
 
 const app = express()
@@ -17,7 +18,7 @@ app.all('*', (_, res, next) => {
   next()
 })
 
-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 {

diff --git a/service/src/middleware/limiter.ts b/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 }