diff --git a/backend/internal/ratelimit/middleware.go b/backend/internal/ratelimit/middleware.go new file mode 100644 index 0000000..dceeee5 --- /dev/null +++ b/backend/internal/ratelimit/middleware.go @@ -0,0 +1,42 @@ +package ratelimit + +import ( + "fmt" + "net/http" + + "github.com/gin-gonic/gin" +) + +// Middleware 返回 Gin 中间件,按 key 维度限流。 +// keyFunc 从请求中提取限流 key(如 IP、用户 ID)。 +func Middleware(limiter Limiter, keyFunc func(*gin.Context) string) gin.HandlerFunc { + return func(c *gin.Context) { + if limiter == nil { + c.Next() + return + } + + key := keyFunc(c) + if key == "" { + // key 为空时跳过限流 + c.Next() + return + } + + allowed, retryAfter := limiter.Allow(c.Request.Context(), key) + + if !allowed { + // 设置 Retry-After header(秒) + c.Header("Retry-After", fmt.Sprintf("%d", int(retryAfter.Seconds()+0.5))) + + c.JSON(http.StatusTooManyRequests, gin.H{ + "code": "RATE_LIMITED", + "message": fmt.Sprintf("too many requests, retry after %s", retryAfter.Round(1)), + }) + c.Abort() + return + } + + c.Next() + } +} diff --git a/backend/internal/ratelimit/middleware_test.go b/backend/internal/ratelimit/middleware_test.go new file mode 100644 index 0000000..c9fd953 --- /dev/null +++ b/backend/internal/ratelimit/middleware_test.go @@ -0,0 +1,196 @@ +package ratelimit + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// mockLimiter 用于测试的 mock 限流器。 +type mockLimiter struct { + allowFunc func(ctx context.Context, key string) (bool, time.Duration) +} + +func (m *mockLimiter) Allow(ctx context.Context, key string) (bool, time.Duration) { + if m.allowFunc != nil { + return m.allowFunc(ctx, key) + } + return true, 0 +} + +func (m *mockLimiter) Stop() {} + +// 编译期接口检查 +var _ Limiter = (*mockLimiter)(nil) + +func TestMiddleware_Allow(t *testing.T) { + gin.SetMode(gin.TestMode) + + limiter := &mockLimiter{ + allowFunc: func(ctx context.Context, key string) (bool, time.Duration) { + return true, 0 + }, + } + + router := gin.New() + router.Use(Middleware(limiter, func(c *gin.Context) string { + return "user1:test" + })) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + }) + + req := httptest.NewRequest(http.MethodGet, "/test", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + + var resp map[string]interface{} + err := json.Unmarshal(w.Body.Bytes(), &resp) + require.NoError(t, err) + assert.Equal(t, "ok", resp["status"]) +} + +func TestMiddleware_Deny(t *testing.T) { + gin.SetMode(gin.TestMode) + + limiter := &mockLimiter{ + allowFunc: func(ctx context.Context, key string) (bool, time.Duration) { + return false, 5 * time.Second + }, + } + + router := gin.New() + router.Use(Middleware(limiter, func(c *gin.Context) string { + return "user1:test" + })) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + }) + + req := httptest.NewRequest(http.MethodGet, "/test", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + // 验证返回 429 + assert.Equal(t, http.StatusTooManyRequests, w.Code) + + // 验证 Retry-After header + assert.Equal(t, "5", w.Header().Get("Retry-After")) + + // 验证响应体 + var resp map[string]interface{} + err := json.Unmarshal(w.Body.Bytes(), &resp) + require.NoError(t, err) + assert.Equal(t, "RATE_LIMITED", resp["code"]) + assert.Contains(t, resp["message"], "retry after") +} + +func TestMiddleware_NilLimiter(t *testing.T) { + gin.SetMode(gin.TestMode) + + router := gin.New() + router.Use(Middleware(nil, func(c *gin.Context) string { + return "user1:test" + })) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + }) + + req := httptest.NewRequest(http.MethodGet, "/test", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + // nil limiter 应该放行 + assert.Equal(t, http.StatusOK, w.Code) +} + +func TestMiddleware_EmptyKey(t *testing.T) { + gin.SetMode(gin.TestMode) + + limiter := &mockLimiter{ + allowFunc: func(ctx context.Context, key string) (bool, time.Duration) { + // 不应该被调用 + t.Error("Allow should not be called with empty key") + return false, 0 + }, + } + + router := gin.New() + router.Use(Middleware(limiter, func(c *gin.Context) string { + return "" // 返回空 key + })) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + }) + + req := httptest.NewRequest(http.MethodGet, "/test", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + // 空 key 应该放行 + assert.Equal(t, http.StatusOK, w.Code) +} + +func TestMiddleware_KeyFunc(t *testing.T) { + gin.SetMode(gin.TestMode) + + var capturedKey string + limiter := &mockLimiter{ + allowFunc: func(ctx context.Context, key string) (bool, time.Duration) { + capturedKey = key + return true, 0 + }, + } + + router := gin.New() + router.Use(Middleware(limiter, func(c *gin.Context) string { + // 从 query 参数提取 user_id + userID := c.Query("user_id") + return userID + ":test" + })) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + }) + + req := httptest.NewRequest(http.MethodGet, "/test?user_id=user123", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "user123:test", capturedKey) +} + +func TestMiddleware_RetryAfterRounding(t *testing.T) { + gin.SetMode(gin.TestMode) + + limiter := &mockLimiter{ + allowFunc: func(ctx context.Context, key string) (bool, time.Duration) { + return false, 2500 * time.Millisecond // 2.5 秒 + }, + } + + router := gin.New() + router.Use(Middleware(limiter, func(c *gin.Context) string { + return "user1:test" + })) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + }) + + req := httptest.NewRequest(http.MethodGet, "/test", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusTooManyRequests, w.Code) + // 2.5 秒向上取整为 3 秒 + assert.Equal(t, "3", w.Header().Get("Retry-After")) +}