feat: 增加GLOBAL_API_RATE_LIMIT_ENABLEGLOBAL_WEB_RATE_LIMIT_ENABLE环境变量,支持是否开启访问速率控制

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

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

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) {