加固后台管理安全

backend/internal/middleware/rate_limit.go

@@ -1,12 +1,14 @@
 package middleware
 
 import (
+	"context"
 	"net/http"
 	"strconv"
 	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/redis/go-redis/v9"
 )
 
 type rateLimitBucket struct {
@@ -22,7 +24,14 @@ type rateLimiter struct {
 	buckets map[string]rateLimitBucket
 }
 
-func RateLimitPerMinute(limit int) gin.HandlerFunc {
+type redisRateLimiter struct {
+	redis    *redis.Client
+	fallback *rateLimiter
+	limit    int
+	window   time.Duration
+}
+
+func RateLimitPerMinute(limit int, rdb *redis.Client) gin.HandlerFunc {
 	if limit <= 0 {
 		return func(c *gin.Context) {
 			c.Next()
@@ -33,24 +42,57 @@ func RateLimitPerMinute(limit int) gin.HandlerFunc {
 		window:  time.Minute,
 		buckets: make(map[string]rateLimitBucket),
 	}
+	if rdb != nil {
+		return (&redisRateLimiter{
+			redis:    rdb,
+			fallback: limiter,
+			limit:    limit,
+			window:   time.Minute,
+		}).handle
+	}
 	return limiter.handle
 }
 
+func (l *redisRateLimiter) handle(c *gin.Context) {
+	now := time.Now()
+	key := c.ClientIP()
+	allowed, resetAt, err := l.allow(c.Request.Context(), key, now)
+	if err != nil {
+		allowed, resetAt = l.fallback.allow(key, now)
+	}
+	if !allowed {
+		writeRateLimited(c, now, resetAt)
+		return
+	}
+	c.Next()
+}
+
 func (l *rateLimiter) handle(c *gin.Context) {
 	now := time.Now()
 	key := c.ClientIP()
 	allowed, resetAt := l.allow(key, now)
 	if !allowed {
-		c.Header("Retry-After", retryAfterSeconds(now, resetAt))
-		c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
-			"code":    "rate_limited",
-			"message": "请求过于频繁，请稍后再试",
-		})
+		writeRateLimited(c, now, resetAt)
 		return
 	}
 	c.Next()
 }
 
+func (l *redisRateLimiter) allow(ctx context.Context, key string, now time.Time) (bool, time.Time, error) {
+	windowSeconds := int64(l.window / time.Second)
+	windowID := now.Unix() / windowSeconds
+	redisKey := "rate_limit:" + key + ":" + strconv.FormatInt(windowID, 10)
+	count, err := l.redis.Incr(ctx, redisKey).Result()
+	if err != nil {
+		return false, time.Time{}, err
+	}
+	if count == 1 {
+		_ = l.redis.Expire(ctx, redisKey, 2*l.window).Err()
+	}
+	resetAt := time.Unix((windowID+1)*windowSeconds, 0)
+	return count <= int64(l.limit), resetAt, nil
+}
+
 func (l *rateLimiter) allow(key string, now time.Time) (bool, time.Time) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
@@ -81,3 +123,11 @@ func retryAfterSeconds(now time.Time, resetAt time.Time) string {
 	}
 	return strconv.Itoa(seconds)
 }
+
+func writeRateLimited(c *gin.Context, now time.Time, resetAt time.Time) {
+	c.Header("Retry-After", retryAfterSeconds(now, resetAt))
+	c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
+		"code":    "rate_limited",
+		"message": "请求过于频繁，请稍后再试",
+	})
+}