pay/internal/middleware/ratelimit_test.go

package middleware_test

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gin-gonic/gin"

	"github.com/wangjia/pay/internal/middleware"
)

func newTestEngine(rl *middleware.IPRateLimiter) *gin.Engine {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.GET("/x", rl.Gin(), func(c *gin.Context) { c.Status(http.StatusOK) })
	return r
}

// burst(=RequestsPerMin)个请求应放行,第 burst+1 个应 429。
func TestIPRateLimiterBurstThenBlocks(t *testing.T) {
	rl := middleware.NewIPRateLimiter(3)
	r := newTestEngine(rl)

	for i := 0; i < 3; i++ {
		w := httptest.NewRecorder()
		r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/x", nil))
		if w.Code != http.StatusOK {
			t.Fatalf("第 %d 次(burst 内)应 200, got %d", i+1, w.Code)
		}
	}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/x", nil))
	if w.Code != http.StatusTooManyRequests {
		t.Fatalf("第 4 次(超 burst)应 429, got %d", w.Code)
	}
}

// 不同 IP 各自独立计桶,互不影响。
func TestIPRateLimiterPerIPIsolated(t *testing.T) {
	rl := middleware.NewIPRateLimiter(1)
	r := newTestEngine(rl)

	req1 := httptest.NewRequest(http.MethodGet, "/x", nil)
	req1.RemoteAddr = "10.0.0.1:1234"
	w1 := httptest.NewRecorder()
	r.ServeHTTP(w1, req1)
	if w1.Code != http.StatusOK {
		t.Fatalf("IP1 第 1 次应 200, got %d", w1.Code)
	}

	req2 := httptest.NewRequest(http.MethodGet, "/x", nil)
	req2.RemoteAddr = "10.0.0.2:1234"
	w2 := httptest.NewRecorder()
	r.ServeHTTP(w2, req2)
	if w2.Code != http.StatusOK {
		t.Fatalf("IP2(独立桶)第 1 次应 200, got %d", w2.Code)
	}

	// IP1 burst=1 已耗尽,第 2 次应 429。
	req3 := httptest.NewRequest(http.MethodGet, "/x", nil)
	req3.RemoteAddr = "10.0.0.1:1234"
	w3 := httptest.NewRecorder()
	r.ServeHTTP(w3, req3)
	if w3.Code != http.StatusTooManyRequests {
		t.Fatalf("IP1 第 2 次(超 burst)应 429, got %d", w3.Code)
	}
}

// jiu 反馈波安全 review:gin 默认 trustedProxies=0.0.0.0/0,不显式 SetTrustedProxies(nil)
// 的话 c.ClientIP() 会信任客户端可控的 X-Forwarded-For——攻击者每次伪造不同 XFF 即可
// 让限流器把每次请求都当"新 IP"的新桶,burst 永不触发(绕过限流)。
// 这条测试用同一个 RemoteAddr(模拟同一 socket 源)、每次带不同的伪造 XFF 头发请求:
//   - 修复前(未调用 SetTrustedProxies(nil),沿用 gin 默认信任所有代理):ClientIP() 读
//     伪造的 XFF,每次都是"新 IP"→ 新桶 → 全部放行,本测试会失败(RED)。
//   - 修复后(main.go 起 engine 时调用 r.SetTrustedProxies(nil)):ClientIP() 退回
//     RemoteIP()(socket 对端地址,请求头伪造不了)→ 同一来源共用一个桶 → burst 耗尽后
//     触发 429,本测试通过(GREEN)。
func TestIPRateLimiterSpoofedXFFCannotBypass(t *testing.T) {
	gin.SetMode(gin.TestMode)
	rl := middleware.NewIPRateLimiter(3)
	r := gin.New()
	if err := r.SetTrustedProxies(nil); err != nil {
		t.Fatalf("SetTrustedProxies(nil) 失败: %v", err)
	}
	r.GET("/x", rl.Gin(), func(c *gin.Context) { c.Status(http.StatusOK) })

	const sameSocketSource = "198.51.100.7:54321"
	var codes []int
	for i := 0; i < 4; i++ {
		req := httptest.NewRequest(http.MethodGet, "/x", nil)
		req.RemoteAddr = sameSocketSource
		// 每次伪造一个不同的、格式合法的公网 IP——修复前的攻击手法。
		req.Header.Set("X-Forwarded-For", fmt.Sprintf("203.0.113.%d", i+1))
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)
		codes = append(codes, w.Code)
	}
	if codes[0] != http.StatusOK || codes[1] != http.StatusOK || codes[2] != http.StatusOK {
		t.Fatalf("同一 socket 源前 3 次(=burst)应放行, got codes=%v", codes)
	}
	if codes[3] != http.StatusTooManyRequests {
		t.Fatalf("同一 socket 源第 4 次应共桶触发 429(伪造 XFF 不能绕过限流), got codes=%v", codes)
	}
}

// 不同 socket 源(不同 RemoteAddr)哪怕带相同的伪造 XFF,也应各自独立计桶——
// 佐证限流键确实取自 c.ClientIP()(修复后=RemoteIP()),不是请求头。
func TestIPRateLimiterDifferentSocketSourceIsolatedEvenWithSameXFF(t *testing.T) {
	gin.SetMode(gin.TestMode)
	rl := middleware.NewIPRateLimiter(1)
	r := gin.New()
	if err := r.SetTrustedProxies(nil); err != nil {
		t.Fatalf("SetTrustedProxies(nil) 失败: %v", err)
	}
	r.GET("/x", rl.Gin(), func(c *gin.Context) { c.Status(http.StatusOK) })

	for _, addr := range []string{"198.51.100.10:1111", "198.51.100.11:2222"} {
		req := httptest.NewRequest(http.MethodGet, "/x", nil)
		req.RemoteAddr = addr
		req.Header.Set("X-Forwarded-For", "203.0.113.99") // 两次都带同一个伪造 XFF
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)
		if w.Code != http.StatusOK {
			t.Fatalf("socket 源 %s 首次应 200(独立桶,不受相同 XFF 影响), got %d", addr, w.Code)
		}
	}
}

// perMinute<=0 兜底成默认 30/min(config.RateLimitConfig 零值语义:"零值=默认开")。
func TestNewIPRateLimiterZeroDefaultsTo30(t *testing.T) {
	rl := middleware.NewIPRateLimiter(0)
	r := newTestEngine(rl)

	for i := 0; i < 30; i++ {
		w := httptest.NewRecorder()
		r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/x", nil))
		if w.Code != http.StatusOK {
			t.Fatalf("零值兜底 30/min:第 %d 次应 200, got %d", i+1, w.Code)
		}
	}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/x", nil))
	if w.Code != http.StatusTooManyRequests {
		t.Fatalf("第 31 次(超默认 30/min)应 429, got %d", w.Code)
	}
}