70 lines
1.5 KiB
Go
70 lines
1.5 KiB
Go
package middleware
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"golang.org/x/time/rate"
|
|
|
|
"filestoragesystem/internal/repository"
|
|
"filestoragesystem/internal/utils"
|
|
)
|
|
|
|
// RateLimit 令牌桶限流中间件
|
|
// 限额从系统设置api_rate_limit读取(次/分钟)，按用户ID(已认证)或IP进行限制
|
|
func RateLimit(settingRepo *repository.SettingRepo) gin.HandlerFunc {
|
|
var mu sync.Mutex
|
|
buckets := make(map[string]*rate.Limiter)
|
|
lastClean := time.Now()
|
|
|
|
getLimit := func() int {
|
|
if v, err := settingRepo.GetValue("api_rate_limit"); err == nil {
|
|
var n int
|
|
if _, e := fmt.Sscanf(v, "%d", &n); e == nil && n > 0 {
|
|
return n
|
|
}
|
|
}
|
|
return 100
|
|
}
|
|
|
|
return func(c *gin.Context) {
|
|
// 管理员后台不限流
|
|
if IsAdmin(c) {
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
key := utils.ClientIP(c)
|
|
if uid := CurrentUserID(c); uid > 0 {
|
|
key = fmt.Sprintf("user:%d", uid)
|
|
}
|
|
|
|
limit := getLimit()
|
|
mu.Lock()
|
|
// 定期清理空闲桶
|
|
if time.Since(lastClean) > 10*time.Minute {
|
|
buckets = make(map[string]*rate.Limiter)
|
|
lastClean = time.Now()
|
|
}
|
|
limiter, ok := buckets[key]
|
|
if !ok {
|
|
// 每分钟limit次：速率=limit/60每秒，突发=limit
|
|
limiter = rate.NewLimiter(rate.Limit(float64(limit)/60.0), limit)
|
|
buckets[key] = limiter
|
|
}
|
|
mu.Unlock()
|
|
|
|
if !limiter.Allow() {
|
|
c.Header("Retry-After", "60")
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, utils.Response{
|
|
Code: 429, Message: "请求过于频繁，请稍后再试",
|
|
})
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|