first commit
This commit is contained in:
huangzhenpc
161
backend/internal/middleware/rate_limiter.go
Normal file
161
backend/internal/middleware/rate_limiter.go
Normal file
@@ -0,0 +1,161 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// RateLimitFailureMode Redis 故障策略
|
||||
type RateLimitFailureMode int
|
||||
|
||||
const (
|
||||
RateLimitFailOpen RateLimitFailureMode = iota
|
||||
RateLimitFailClose
|
||||
)
|
||||
|
||||
// RateLimitOptions 限流可选配置
|
||||
type RateLimitOptions struct {
|
||||
FailureMode RateLimitFailureMode
|
||||
}
|
||||
|
||||
var rateLimitScript = redis.NewScript(`
|
||||
local current = redis.call('INCR', KEYS[1])
|
||||
local ttl = redis.call('PTTL', KEYS[1])
|
||||
local repaired = 0
|
||||
if current == 1 then
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[1])
|
||||
elseif ttl == -1 then
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[1])
|
||||
repaired = 1
|
||||
end
|
||||
return {current, repaired}
|
||||
`)
|
||||
|
||||
// rateLimitRun 允许测试覆写脚本执行逻辑
|
||||
var rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
|
||||
values, err := rateLimitScript.Run(ctx, client, []string{key}, windowMillis).Slice()
|
||||
if err != nil {
|
||||
return 0, false, err
|
||||
}
|
||||
if len(values) < 2 {
|
||||
return 0, false, fmt.Errorf("rate limit script returned %d values", len(values))
|
||||
}
|
||||
count, err := parseInt64(values[0])
|
||||
if err != nil {
|
||||
return 0, false, err
|
||||
}
|
||||
repaired, err := parseInt64(values[1])
|
||||
if err != nil {
|
||||
return 0, false, err
|
||||
}
|
||||
return count, repaired == 1, nil
|
||||
}
|
||||
|
||||
// RateLimiter Redis 速率限制器
|
||||
type RateLimiter struct {
|
||||
redis *redis.Client
|
||||
prefix string
|
||||
}
|
||||
|
||||
// NewRateLimiter 创建速率限制器实例
|
||||
func NewRateLimiter(redisClient *redis.Client) *RateLimiter {
|
||||
return &RateLimiter{
|
||||
redis: redisClient,
|
||||
prefix: "rate_limit:",
|
||||
}
|
||||
}
|
||||
|
||||
// Limit 返回速率限制中间件
|
||||
// key: 限制类型标识
|
||||
// limit: 时间窗口内最大请求数
|
||||
// window: 时间窗口
|
||||
func (r *RateLimiter) Limit(key string, limit int, window time.Duration) gin.HandlerFunc {
|
||||
return r.LimitWithOptions(key, limit, window, RateLimitOptions{})
|
||||
}
|
||||
|
||||
// LimitWithOptions 返回速率限制中间件（带可选配置）
|
||||
func (r *RateLimiter) LimitWithOptions(key string, limit int, window time.Duration, opts RateLimitOptions) gin.HandlerFunc {
|
||||
failureMode := opts.FailureMode
|
||||
if failureMode != RateLimitFailClose {
|
||||
failureMode = RateLimitFailOpen
|
||||
}
|
||||
|
||||
return func(c *gin.Context) {
|
||||
ip := c.ClientIP()
|
||||
redisKey := r.prefix + key + ":" + ip
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
windowMillis := windowTTLMillis(window)
|
||||
|
||||
// 使用 Lua 脚本原子操作增加计数并设置过期
|
||||
count, repaired, err := rateLimitRun(ctx, r.redis, redisKey, windowMillis)
|
||||
if err != nil {
|
||||
log.Printf("[RateLimit] redis error: key=%s mode=%s err=%v", redisKey, failureModeLabel(failureMode), err)
|
||||
if failureMode == RateLimitFailClose {
|
||||
abortRateLimit(c)
|
||||
return
|
||||
}
|
||||
// Redis 错误时放行，避免影响正常服务
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if repaired {
|
||||
log.Printf("[RateLimit] ttl repaired: key=%s window_ms=%d", redisKey, windowMillis)
|
||||
}
|
||||
|
||||
// 超过限制
|
||||
if count > int64(limit) {
|
||||
abortRateLimit(c)
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func windowTTLMillis(window time.Duration) int64 {
|
||||
ttl := window.Milliseconds()
|
||||
if ttl < 1 {
|
||||
return 1
|
||||
}
|
||||
return ttl
|
||||
}
|
||||
|
||||
func abortRateLimit(c *gin.Context) {
|
||||
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
||||
"error": "rate limit exceeded",
|
||||
"message": "Too many requests, please try again later",
|
||||
})
|
||||
}
|
||||
|
||||
func failureModeLabel(mode RateLimitFailureMode) string {
|
||||
if mode == RateLimitFailClose {
|
||||
return "fail-close"
|
||||
}
|
||||
return "fail-open"
|
||||
}
|
||||
|
||||
func parseInt64(value any) (int64, error) {
|
||||
switch v := value.(type) {
|
||||
case int64:
|
||||
return v, nil
|
||||
case int:
|
||||
return int64(v), nil
|
||||
case string:
|
||||
parsed, err := strconv.ParseInt(v, 10, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return parsed, nil
|
||||
default:
|
||||
return 0, fmt.Errorf("unexpected value type %T", value)
|
||||
}
|
||||
}
|
||||
114
backend/internal/middleware/rate_limiter_integration_test.go
Normal file
114
backend/internal/middleware/rate_limiter_integration_test.go
Normal file
@@ -0,0 +1,114 @@
|
||||
//go:build integration
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
tcredis "github.com/testcontainers/testcontainers-go/modules/redis"
|
||||
)
|
||||
|
||||
const redisImageTag = "redis:8.4-alpine"
|
||||
|
||||
func TestRateLimiterSetsTTLAndDoesNotRefresh(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
ctx := context.Background()
|
||||
rdb := startRedis(t, ctx)
|
||||
limiter := NewRateLimiter(rdb)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(limiter.Limit("ttl-test", 10, 2*time.Second))
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
recorder := performRequest(router)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
|
||||
redisKey := limiter.prefix + "ttl-test:127.0.0.1"
|
||||
ttlBefore, err := rdb.PTTL(ctx, redisKey).Result()
|
||||
require.NoError(t, err)
|
||||
require.Greater(t, ttlBefore, time.Duration(0))
|
||||
require.LessOrEqual(t, ttlBefore, 2*time.Second)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
recorder = performRequest(router)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
|
||||
ttlAfter, err := rdb.PTTL(ctx, redisKey).Result()
|
||||
require.NoError(t, err)
|
||||
require.Less(t, ttlAfter, ttlBefore)
|
||||
}
|
||||
|
||||
func TestRateLimiterFixesMissingTTL(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
ctx := context.Background()
|
||||
rdb := startRedis(t, ctx)
|
||||
limiter := NewRateLimiter(rdb)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(limiter.Limit("ttl-missing", 10, 2*time.Second))
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
redisKey := limiter.prefix + "ttl-missing:127.0.0.1"
|
||||
require.NoError(t, rdb.Set(ctx, redisKey, 5, 0).Err())
|
||||
|
||||
ttlBefore, err := rdb.PTTL(ctx, redisKey).Result()
|
||||
require.NoError(t, err)
|
||||
require.Less(t, ttlBefore, time.Duration(0))
|
||||
|
||||
recorder := performRequest(router)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
|
||||
ttlAfter, err := rdb.PTTL(ctx, redisKey).Result()
|
||||
require.NoError(t, err)
|
||||
require.Greater(t, ttlAfter, time.Duration(0))
|
||||
}
|
||||
|
||||
func performRequest(router *gin.Engine) *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "127.0.0.1:1234"
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, req)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func startRedis(t *testing.T, ctx context.Context) *redis.Client {
|
||||
t.Helper()
|
||||
|
||||
redisContainer, err := tcredis.Run(ctx, redisImageTag)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
_ = redisContainer.Terminate(ctx)
|
||||
})
|
||||
|
||||
redisHost, err := redisContainer.Host(ctx)
|
||||
require.NoError(t, err)
|
||||
redisPort, err := redisContainer.MappedPort(ctx, "6379/tcp")
|
||||
require.NoError(t, err)
|
||||
|
||||
rdb := redis.NewClient(&redis.Options{
|
||||
Addr: fmt.Sprintf("%s:%d", redisHost, redisPort.Int()),
|
||||
DB: 0,
|
||||
})
|
||||
require.NoError(t, rdb.Ping(ctx).Err())
|
||||
|
||||
t.Cleanup(func() {
|
||||
_ = rdb.Close()
|
||||
})
|
||||
|
||||
return rdb
|
||||
}
|
||||
100
backend/internal/middleware/rate_limiter_test.go
Normal file
100
backend/internal/middleware/rate_limiter_test.go
Normal file
@@ -0,0 +1,100 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestWindowTTLMillis(t *testing.T) {
|
||||
require.Equal(t, int64(1), windowTTLMillis(500*time.Microsecond))
|
||||
require.Equal(t, int64(1), windowTTLMillis(1500*time.Microsecond))
|
||||
require.Equal(t, int64(2), windowTTLMillis(2500*time.Microsecond))
|
||||
}
|
||||
|
||||
func TestRateLimiterFailureModes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rdb := redis.NewClient(&redis.Options{
|
||||
Addr: "127.0.0.1:1",
|
||||
DialTimeout: 50 * time.Millisecond,
|
||||
ReadTimeout: 50 * time.Millisecond,
|
||||
WriteTimeout: 50 * time.Millisecond,
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
_ = rdb.Close()
|
||||
})
|
||||
|
||||
limiter := NewRateLimiter(rdb)
|
||||
|
||||
failOpenRouter := gin.New()
|
||||
failOpenRouter.Use(limiter.Limit("test", 1, time.Second))
|
||||
failOpenRouter.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "127.0.0.1:1234"
|
||||
recorder := httptest.NewRecorder()
|
||||
failOpenRouter.ServeHTTP(recorder, req)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
|
||||
failCloseRouter := gin.New()
|
||||
failCloseRouter.Use(limiter.LimitWithOptions("test", 1, time.Second, RateLimitOptions{
|
||||
FailureMode: RateLimitFailClose,
|
||||
}))
|
||||
failCloseRouter.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
req = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "127.0.0.1:1234"
|
||||
recorder = httptest.NewRecorder()
|
||||
failCloseRouter.ServeHTTP(recorder, req)
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code)
|
||||
}
|
||||
|
||||
func TestRateLimiterSuccessAndLimit(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
originalRun := rateLimitRun
|
||||
counts := []int64{1, 2}
|
||||
callIndex := 0
|
||||
rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
|
||||
if callIndex >= len(counts) {
|
||||
return counts[len(counts)-1], false, nil
|
||||
}
|
||||
value := counts[callIndex]
|
||||
callIndex++
|
||||
return value, false, nil
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
rateLimitRun = originalRun
|
||||
})
|
||||
|
||||
limiter := NewRateLimiter(redis.NewClient(&redis.Options{Addr: "127.0.0.1:1"}))
|
||||
|
||||
router := gin.New()
|
||||
router.Use(limiter.Limit("test", 1, time.Second))
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "127.0.0.1:1234"
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, req)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
|
||||
req = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "127.0.0.1:1234"
|
||||
recorder = httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, req)
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code)
|
||||
}
|
||||
Reference in New Issue
Block a user