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
}

View File

@@ -77,3 +77,20 @@ func (c *Client) Unlock(ctx context.Context, key string, token string) error {
_, err := unlockScript.Run(ctx, c.rdb, []string{key}, token).Result()
return err
}
func (c *Client) IncrementWithExpire(ctx context.Context, key string, expire time.Duration) (int64, error) {
if c == nil || c.rdb == nil {
return 0, nil
}
count, err := c.rdb.Incr(ctx, key).Result()
if err != nil {
return 0, err
}
if count == 1 {
err = c.rdb.Expire(ctx, key, expire).Err()
if err != nil {
return 0, err
}
}
return count, nil
}