xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-11-04 05:13:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: 增加GLOBAL_API_RATE_LIMIT_ENABLE与GLOBAL_WEB_RATE_LIMIT_ENABLE环境变量，支持是否开启访问速率控制

Browse Source
This commit is contained in:
licoy
2024-11-12 20:02:33 +08:00
parent fad29a8cc2
commit 08023f6d96
2 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
common/constants.go
View File

@@ -144,9 +144,11 @@ var (
// All duration's unit is seconds
// Shouldn't larger then RateLimitKeyExpirationDuration
var (
GlobalApiRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_API_RATE_LIMIT_ENABLE", true)
GlobalApiRateLimitNum = GetEnvOrDefault("GLOBAL_API_RATE_LIMIT", 180)
GlobalApiRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_API_RATE_LIMIT_DURATION", 180))
GlobalWebRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_WEB_RATE_LIMIT_ENABLE", true)
GlobalWebRateLimitNum = GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT", 60)
GlobalWebRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT_DURATION", 180))

14
middleware/rate-limit.go
View File

@@ -13,6 +13,10 @@ var timeFormat = "2006-01-02T15:04:05.000Z"
var inMemoryRateLimiter common.InMemoryRateLimiter
var defNext = func(c *gin.Context) {
c.Next()
}
func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) {
ctx := context.Background()
rdb := common.RDB
@@ -83,11 +87,17 @@ func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gi
}
func GlobalWebRateLimit() func(c *gin.Context) {
return rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW")
if common.GlobalWebRateLimitEnable {
return rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW")
}
return defNext
}
func GlobalAPIRateLimit() func(c *gin.Context) {
return rateLimitFactory(common.GlobalApiRateLimitNum, common.GlobalApiRateLimitDuration, "GA")
if common.GlobalApiRateLimitEnable {
return rateLimitFactory(common.GlobalApiRateLimitNum, common.GlobalApiRateLimitDuration, "GA")
}
return defNext
}
func CriticalRateLimit() func(c *gin.Context) {