jiu/backend/internal/middleware/ratelimit_test.go

package middleware

import (
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"

	"github.com/wangjia/jiu/backend/config"
)

func init() {
	gin.SetMode(gin.TestMode)
}

// enableRateLimit 临时打开限流开关，返回恢复函数。
func enableRateLimit(t *testing.T) {
	t.Helper()
	prev := config.C.RateLimit.Enabled
	config.C.RateLimit.Enabled = true
	t.Cleanup(func() { config.C.RateLimit.Enabled = prev })
}

// hit 发一个带指定客户端 IP 的请求，返回状态码。SetTrustedProxies 让 X-Forwarded-For 被采信，
// 便于在测试里模拟不同来源 IP。
func hit(r *gin.Engine, ip string) *httptest.ResponseRecorder {
	w := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = ip + ":12345"
	r.ServeHTTP(w, req)
	return w
}

func newIPRouter(h gin.HandlerFunc) *gin.Engine {
	r := gin.New()
	r.GET("/", h, func(c *gin.Context) { c.String(http.StatusOK, "ok") })
	return r
}

func TestRateLimitByIP_BurstThen429(t *testing.T) {
	enableRateLimit(t)
	// 每分钟 3 次，burst 3：前 3 个放行，第 4 个 429。
	r := newIPRouter(RateLimitByIP(PerMinute(3), 3))

	for i := 0; i < 3; i++ {
		if w := hit(r, "1.1.1.1"); w.Code != http.StatusOK {
			t.Fatalf("第 %d 个请求应放行，得到 %d", i+1, w.Code)
		}
	}
	w := hit(r, "1.1.1.1")
	if w.Code != http.StatusTooManyRequests {
		t.Fatalf("超出 burst 应 429，得到 %d", w.Code)
	}
	if ra := w.Header().Get("Retry-After"); ra == "" {
		t.Fatal("429 响应应带 Retry-After 头")
	}
}

func TestRateLimitByIP_KeysIndependent(t *testing.T) {
	enableRateLimit(t)
	r := newIPRouter(RateLimitByIP(PerMinute(1), 1))

	// 不同 IP 各有独立令牌桶，互不影响。
	if w := hit(r, "2.2.2.2"); w.Code != http.StatusOK {
		t.Fatalf("IP A 首次应放行，得到 %d", w.Code)
	}
	if w := hit(r, "3.3.3.3"); w.Code != http.StatusOK {
		t.Fatalf("IP B 首次应放行（独立桶），得到 %d", w.Code)
	}
	// IP A 再来一发应 429（桶已空）。
	if w := hit(r, "2.2.2.2"); w.Code != http.StatusTooManyRequests {
		t.Fatalf("IP A 第二发应 429，得到 %d", w.Code)
	}
}

func TestRateLimitByIP_Refill(t *testing.T) {
	enableRateLimit(t)
	// 每秒 50 次：耗尽 burst 后等一小会儿令牌补回，能再次放行。
	r := newIPRouter(RateLimitByIP(rate.Limit(50), 1))
	if w := hit(r, "4.4.4.4"); w.Code != http.StatusOK {
		t.Fatalf("首发应放行，得到 %d", w.Code)
	}
	if w := hit(r, "4.4.4.4"); w.Code != http.StatusTooManyRequests {
		t.Fatalf("紧接第二发应 429，得到 %d", w.Code)
	}
	time.Sleep(40 * time.Millisecond) // 50/s → 20ms 补一个令牌
	if w := hit(r, "4.4.4.4"); w.Code != http.StatusOK {
		t.Fatalf("等待补充后应放行，得到 %d", w.Code)
	}
}

func TestRateLimitDisabledPassthrough(t *testing.T) {
	prev := config.C.RateLimit.Enabled
	config.C.RateLimit.Enabled = false
	defer func() { config.C.RateLimit.Enabled = prev }()

	r := newIPRouter(RateLimitByIP(PerMinute(1), 1))
	// 关闭时远超额度也全放行。
	for i := 0; i < 10; i++ {
		if w := hit(r, "5.5.5.5"); w.Code != http.StatusOK {
			t.Fatalf("限流关闭时第 %d 发应放行，得到 %d", i+1, w.Code)
		}
	}
}

func TestKeyedLimiterSweepEvictsIdle(t *testing.T) {
	kl := newKeyedLimiter(PerMinute(60), 1)
	kl.get("ip:a")
	kl.get("ip:b")
	if len(kl.entries) != 2 {
		t.Fatalf("应有 2 个 entry，得到 %d", len(kl.entries))
	}
	// 把 a 的活动时间推到很久以前，sweep 应只淘汰 a。
	kl.mu.Lock()
	kl.entries["ip:a"].lastSeen = time.Now().Add(-time.Hour)
	kl.mu.Unlock()

	kl.sweep(10 * time.Minute)

	kl.mu.Lock()
	defer kl.mu.Unlock()
	if _, ok := kl.entries["ip:a"]; ok {
		t.Fatal("空闲 key a 应被淘汰")
	}
	if _, ok := kl.entries["ip:b"]; !ok {
		t.Fatal("活跃 key b 不应被淘汰")
	}
}

// TestTrustedProxyRealIP 复刻 main.go 的可信代理配置：只信任本机写的 X-Real-IP，
// 客户端伪造的 X-Forwarded-For 不被采信 → c.ClientIP() 返回真实 IP，限流不可被请求头绕过。
func TestTrustedProxyRealIP(t *testing.T) {
	r := gin.New()
	_ = r.SetTrustedProxies([]string{"127.0.0.1", "::1"})
	r.RemoteIPHeaders = []string{"X-Real-IP"}
	var got string
	r.GET("/", func(c *gin.Context) {
		got = c.ClientIP()
		c.Status(http.StatusOK)
	})

	w := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "127.0.0.1:9999"            // 本机 nginx
	req.Header.Set("X-Real-IP", "9.9.9.9")       // nginx 写入的真实 IP
	req.Header.Set("X-Forwarded-For", "1.2.3.4") // 客户端伪造，应被忽略
	r.ServeHTTP(w, req)

	if got != "9.9.9.9" {
		t.Fatalf("ClientIP 应取可信的 X-Real-IP=9.9.9.9，得到 %q（伪造的 XFF 不应生效）", got)
	}
}

func TestRateLimitByShop_NoShopPassthrough(t *testing.T) {
	enableRateLimit(t)
	// 未设置 shop_id（无 JWT 上下文）时放行，交由 JWT 中间件拦截。
	r := gin.New()
	r.GET("/", RateLimitByShop(PerSecond(1), 1), func(c *gin.Context) { c.String(http.StatusOK, "ok") })
	for i := 0; i < 5; i++ {
		if w := hit(r, "6.6.6.6"); w.Code != http.StatusOK {
			t.Fatalf("无 shop_id 时第 %d 发应放行，得到 %d", i+1, w.Code)
		}
	}
}

func TestRateLimitByShop_PerShop(t *testing.T) {
	enableRateLimit(t)
	r := gin.New()
	// 模拟 JWT 已注入 shop_id（取自 ?shop= 查询参数）。
	r.GET("/", func(c *gin.Context) {
		id, _ := strconv.ParseUint(c.Query("shop"), 10, 64)
		c.Set(CtxShopID, id)
	}, RateLimitByShop(PerSecond(1), 1), func(c *gin.Context) { c.String(http.StatusOK, "ok") })

	shopHit := func(shop string) int {
		w := httptest.NewRecorder()
		req := httptest.NewRequest(http.MethodGet, "/?shop="+shop, nil)
		r.ServeHTTP(w, req)
		return w.Code
	}
	if c := shopHit("1"); c != http.StatusOK {
		t.Fatalf("shop1 首发应放行，得到 %d", c)
	}
	if c := shopHit("2"); c != http.StatusOK {
		t.Fatalf("shop2 首发应放行（独立桶），得到 %d", c)
	}
	if c := shopHit("1"); c != http.StatusTooManyRequests {
		t.Fatalf("shop1 第二发应 429，得到 %d", c)
	}
}