wangjia
149 lines
5.3 KiB
Go
149 lines
5.3 KiB
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)
|
|
}
|
|
}
|