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()
+	}
+}

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"))
+}