Skip to content

Commit

Permalink
feat: 新增限流功能 (Chanzhaoyu#718)
Browse files Browse the repository at this point in the history
* 请求速率限制

* perf: 优化代码

---------

Co-authored-by: ChenZhaoYu <[email protected]>
  • Loading branch information
2 people authored and Equim-chan committed Mar 21, 2023
1 parent c327e18 commit 27448d3
Show file tree
Hide file tree
Showing 7 changed files with 42 additions and 1 deletion.
4 changes: 4 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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` 一起时生效,可选
Expand Down Expand Up @@ -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 一起时生效
Expand All @@ -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) |
Expand Down
2 changes: 2 additions & 0 deletions docker-compose/docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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 一起时生效
Expand Down
3 changes: 3 additions & 0 deletions service/.env.example
Original file line number Diff line number Diff line change
Expand Up @@ -27,3 +27,6 @@ SOCKS_PROXY_PORT=

# HTTPS PROXY
HTTPS_PROXY=

# Rate Limit
MAX_REQUEST_PER_HOUR=
1 change: 1 addition & 0 deletions service/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
11 changes: 11 additions & 0 deletions service/pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion service/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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', 'text/event-stream; charset=utf-8')

try {
Expand Down
19 changes: 19 additions & 0 deletions service/src/middleware/limiter.ts
Original file line number Diff line number Diff line change
@@ -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 }

0 comments on commit 27448d3

Please sign in to comment.