feat(middleware):为登录与互动接口增加限流中间件

This commit is contained in:
tangerineyu
2026-03-27 16:09:01 +08:00
parent 145bf3fa86
commit e89d5aca5d
3 changed files with 109 additions and 9 deletions

View File

@@ -0,0 +1,71 @@
package ratelimit
import (
rediscache "feedsystem_video_go/internal/middleware/redis"
jwt "feedsystem_video_go/internal/middleware/jwt"
"fmt"
"net/http"
"strings"
"time"
"strconv"
"github.com/gin-gonic/gin"
)
type KeyFunc func(*gin.Context) (string, bool)
func Limit(
cache *rediscache.Client,
keyPrefix string,
maxRequests int64,
window time.Duration,
keyFunc KeyFunc,
) gin.HandlerFunc {
return func(c *gin.Context) {
if cache == nil || keyFunc == nil || maxRequests <= 0 || window <= 0 {
c.Next()
return
}
subject, ok := keyFunc(c)
if !ok {
c.Next()
return
}
key := buildKey(keyPrefix, subject)
count, err := cache.IncrementWithExpire(c.Request.Context(), key, window)
if err != nil {
c.Next()
return
}
if count > maxRequests {
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
"error": "too many requests",
})
return
}
c.Next()
}
}
func buildKey(keyPrefix, subject string) string {
keyPrefix = strings.TrimSpace(keyPrefix)
if keyPrefix == "" {
keyPrefix = "default"
}
return fmt.Sprintf("feedsystem:ratelimit:%s:%s", keyPrefix, strings.TrimSpace(subject))
}
func KeyByIP(c *gin.Context) (string, bool) {
ip := strings.TrimSpace(c.ClientIP())
if ip == "" {
return "", false
}
return ip, true
}
func KeyByAccount(c *gin.Context) (string, bool) {
accountID, err := jwt.GetAccountID(c)
if err != nil || accountID == 0 {
return "", false
}
return strconv.FormatUint(uint64(accountID), 10), true
}