perms-system-server/internal/middleware/ratelimitMiddleware_test.go

package middleware

import (
	"encoding/json"
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/zeromicro/go-zero/core/stores/redis"
	"github.com/zeromicro/go-zero/rest/httpx"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"perms-system-server/internal/response"
	"perms-system-server/internal/testutil"
	"testing"
	"time"
)

func init() {
	response.Setup()
}

func uniqueIP() string {
	return fmt.Sprintf("10.%d.%d.%d", rand.Intn(256), rand.Intn(256), rand.Intn(256))
}

func newTestRedis() *redis.Redis {
	cfg := testutil.GetTestConfig()
	return redis.MustNewRedis(cfg.CacheRedis.Nodes[0].RedisConf)
}

func newTestMiddleware(rds *redis.Redis, quota int) *RateLimitMiddleware {
	prefix := fmt.Sprintf("test_rl_%d_%d", time.Now().UnixNano(), rand.Intn(100000))
	return NewRateLimitMiddleware(rds, 60, quota, prefix, false)
}

// TC-0546: 正常请求(未超限)
func TestRateLimit_NormalRequest(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddleware(rds, 10)

	nextCalled := false
	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		nextCalled = true
		w.WriteHeader(http.StatusOK)
	})

	req := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req.Header.Set("X-Forwarded-For", uniqueIP())
	w := httptest.NewRecorder()

	handler(w, req)
	assert.True(t, nextCalled, "next handler should be called")
	assert.Equal(t, http.StatusOK, w.Code)
}

// TC-0547: 超限请求被拒绝
func TestRateLimit_OverQuotaRejected(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddleware(rds, 2)
	ip := uniqueIP()

	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		httpx.OkJson(w, nil)
	})

	for i := 0; i < 2; i++ {
		req := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
		req.Header.Set("X-Forwarded-For", ip)
		w := httptest.NewRecorder()
		handler(w, req)
	}

	req := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req.Header.Set("X-Forwarded-For", ip)
	w := httptest.NewRecorder()
	handler(w, req)

	var body response.Body
	err := json.Unmarshal(w.Body.Bytes(), &body)
	require.NoError(t, err)
	assert.False(t, body.Success)
	assert.Equal(t, 429, body.ErrorCode)
	assert.Equal(t, "请求过于频繁，请稍后再试", body.ErrorMessage)
}

// TC-0548: behindProxy=false时XFF被忽略
func TestRateLimit_XForwardedForIgnored(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddleware(rds, 1)

	var nextCount int
	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		nextCount++
		w.WriteHeader(http.StatusOK)
	})

	remoteAddr := uniqueIP() + ":12345"

	req := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req.RemoteAddr = remoteAddr
	req.Header.Set("X-Forwarded-For", uniqueIP())
	handler(httptest.NewRecorder(), req)
	assert.Equal(t, 1, nextCount)

	req2 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req2.RemoteAddr = remoteAddr
	req2.Header.Set("X-Forwarded-For", uniqueIP())
	handler(httptest.NewRecorder(), req2)
	assert.Equal(t, 1, nextCount, "different X-Forwarded-For should NOT bypass rate limit; RemoteAddr is used")
}

// TC-0549: behindProxy=false时X-Real-IP被忽略
func TestRateLimit_XRealIPIgnored(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddleware(rds, 1)

	var nextCount int
	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		nextCount++
		w.WriteHeader(http.StatusOK)
	})

	remoteAddr := uniqueIP() + ":12345"

	req := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req.RemoteAddr = remoteAddr
	req.Header.Set("X-Real-IP", uniqueIP())
	handler(httptest.NewRecorder(), req)
	assert.Equal(t, 1, nextCount)

	req2 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req2.RemoteAddr = remoteAddr
	req2.Header.Set("X-Real-IP", uniqueIP())
	handler(httptest.NewRecorder(), req2)
	assert.Equal(t, 1, nextCount, "different X-Real-IP should NOT bypass rate limit; RemoteAddr is used")
}

// TC-0550: IP从RemoteAddr解析
func TestRateLimit_IPFromRemoteAddr(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddleware(rds, 1)
	ip := uniqueIP()
	remoteAddr := ip + ":12345"

	var gotNext bool
	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		gotNext = true
		w.WriteHeader(http.StatusOK)
	})

	req := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req.RemoteAddr = remoteAddr
	w := httptest.NewRecorder()
	handler(w, req)
	assert.True(t, gotNext)

	gotNext = false
	req2 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req2.RemoteAddr = remoteAddr
	w2 := httptest.NewRecorder()
	handler(w2, req2)
	assert.False(t, gotNext, "should be rate limited by RemoteAddr")
}

// TC-0551: 不同RemoteAddr独立限流
func TestRateLimit_DifferentIPsIndependent(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddleware(rds, 1)
	addr1 := uniqueIP() + ":12345"
	addr2 := uniqueIP() + ":12345"

	var nextCount int
	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		nextCount++
		w.WriteHeader(http.StatusOK)
	})

	req1 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req1.RemoteAddr = addr1
	handler(httptest.NewRecorder(), req1)
	assert.Equal(t, 1, nextCount)

	req2 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req2.RemoteAddr = addr2
	handler(httptest.NewRecorder(), req2)
	assert.Equal(t, 2, nextCount, "different RemoteAddr should have independent quotas")

	req3 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req3.RemoteAddr = addr1
	handler(httptest.NewRecorder(), req3)
	assert.Equal(t, 2, nextCount, "addr1 should be over quota")

	req4 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req4.RemoteAddr = addr2
	handler(httptest.NewRecorder(), req4)
	assert.Equal(t, 2, nextCount, "addr2 should be over quota")
}

func newTestMiddlewareProxy(rds *redis.Redis, quota int) *RateLimitMiddleware {
	prefix := fmt.Sprintf("test_rl_%d_%d", time.Now().UnixNano(), rand.Intn(100000))
	return NewRateLimitMiddleware(rds, 60, quota, prefix, true)
}

// TC-0552: behindProxy=true时信任X-Real-IP
func TestRateLimit_BehindProxy_TrustsXRealIP(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddlewareProxy(rds, 1)
	remoteAddr := uniqueIP() + ":12345"

	var nextCount int
	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		nextCount++
		w.WriteHeader(http.StatusOK)
	})

	req1 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req1.RemoteAddr = remoteAddr
	req1.Header.Set("X-Real-IP", uniqueIP())
	handler(httptest.NewRecorder(), req1)
	assert.Equal(t, 1, nextCount)

	req2 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req2.RemoteAddr = remoteAddr
	req2.Header.Set("X-Real-IP", uniqueIP())
	handler(httptest.NewRecorder(), req2)
	assert.Equal(t, 2, nextCount, "different X-Real-IP should have independent quotas when behindProxy=true")
}

// TC-0553: behindProxy=true时无X-Real-IP回退RemoteAddr
func TestRateLimit_BehindProxy_FallbackToRemoteAddr(t *testing.T) {
	rds := newTestRedis()
	m := newTestMiddlewareProxy(rds, 1)
	remoteAddr := uniqueIP() + ":12345"

	var nextCount int
	handler := m.Handle(func(w http.ResponseWriter, r *http.Request) {
		nextCount++
		w.WriteHeader(http.StatusOK)
	})

	req1 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req1.RemoteAddr = remoteAddr
	handler(httptest.NewRecorder(), req1)
	assert.Equal(t, 1, nextCount)

	req2 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req2.RemoteAddr = remoteAddr
	handler(httptest.NewRecorder(), req2)
	assert.Equal(t, 1, nextCount, "should fall back to RemoteAddr when X-Real-IP is absent")
}

// 原 TC-0554 "behindProxy=true 时 XFF 仍被忽略" 已按  修复反转：behindProxy=true
// 时 XFF 首段优先，契约由 ratelimitMiddlewareXff_audit_test.go (TC-0862~0866) 取代。

// =============================================================================
// audit  回归：产品登录与管后登录必须使用独立的限流桶
// 修复前：两个入口共享同一个 keyPrefix，导致攻击者对产品登录的爆破会消耗管后登录的配额（或反之）
// 修复后：ProductLoginRateLimit 使用 "...:rl:login:product"，AdminLoginRateLimit 使用 "...:rl:login:admin"
// =============================================================================

// TC-0710: 两个不同 keyPrefix 的限流中间件在同一 IP 上互不影响
func TestRateLimit_ProductAndAdminBucketsAreIndependent(t *testing.T) {
	rds := newTestRedis()

	// 模拟 servicecontext.go 里的两个独立桶
	prefixBase := fmt.Sprintf("test_rl_l2_%d_%d", time.Now().UnixNano(), rand.Intn(100000))
	productM := NewRateLimitMiddleware(rds, 60, 1, prefixBase+":rl:login:product", false)
	adminM := NewRateLimitMiddleware(rds, 60, 1, prefixBase+":rl:login:admin", false)

	ip := uniqueIP()
	remoteAddr := ip + ":12345"

	var productNext, adminNext int
	productHandler := productM.Handle(func(w http.ResponseWriter, r *http.Request) {
		productNext++
		w.WriteHeader(http.StatusOK)
	})
	adminHandler := adminM.Handle(func(w http.ResponseWriter, r *http.Request) {
		adminNext++
		w.WriteHeader(http.StatusOK)
	})

	// 对产品登录打一枪（配额=1，刚好用完）
	req1 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req1.RemoteAddr = remoteAddr
	productHandler(httptest.NewRecorder(), req1)
	require.Equal(t, 1, productNext)

	// 再对产品登录打一枪 → 被限流
	req2 := httptest.NewRequest(http.MethodPost, "/api/auth/login", nil)
	req2.RemoteAddr = remoteAddr
	productHandler(httptest.NewRecorder(), req2)
	require.Equal(t, 1, productNext, "产品登录桶已耗尽")

	// 关键：同 IP 对管后登录仍应放行（独立桶）
	req3 := httptest.NewRequest(http.MethodPost, "/api/auth/adminLogin", nil)
	req3.RemoteAddr = remoteAddr
	adminHandler(httptest.NewRecorder(), req3)
	assert.Equal(t, 1, adminNext,
		"产品登录限流不应影响管后登录（不同 keyPrefix）")

	// 再打管后一枪 → 管后桶也应耗尽，但产品桶已经耗尽在先
	req4 := httptest.NewRequest(http.MethodPost, "/api/auth/adminLogin", nil)
	req4.RemoteAddr = remoteAddr
	adminHandler(httptest.NewRecorder(), req4)
	assert.Equal(t, 1, adminNext, "管后桶配额=1，第二次应被限流")
}

// TC-0555: RemoteAddr无端口格式
func TestExtractClientIP_RemoteAddrNoPort(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/api/test", nil)
	req.RemoteAddr = "1.2.3.4"

	ip := ExtractClientIP(req, false)
	assert.Equal(t, "1.2.3.4", ip, "should return raw RemoteAddr when SplitHostPort fails")

	ip2 := ExtractClientIP(req, true)
	assert.Equal(t, "1.2.3.4", ip2, "behindProxy=true without X-Real-IP should also fallback")
}

func TestExtractClientIP_XFFFirstValid(t *testing.T) {
	r := httptest.NewRequest("POST", "/x", nil)
	r.Header.Set("X-Forwarded-For", "1.1.1.1, 2.2.2.2, 3.3.3.3")
	r.Header.Set("X-Real-IP", "9.9.9.9") // 不应被用
	r.RemoteAddr = "5.5.5.5:8080"        // 不应被用

	assert.Equal(t, "1.1.1.1", ExtractClientIP(r, true),
		"XFF 首段合法时优先返回，高于 XRI / RemoteAddr")
}

// TC-0863: behindProxy=true + XFF 全非法 + XRI 合法 → fallthrough 到 XRI。
func TestExtractClientIP_XFFAllInvalid_FallbackXRI(t *testing.T) {
	r := httptest.NewRequest("POST", "/x", nil)
	r.Header.Set("X-Forwarded-For", "garbage, not-an-ip")
	r.Header.Set("X-Real-IP", "10.0.0.1")
	r.RemoteAddr = "5.5.5.5:8080"

	assert.Equal(t, "10.0.0.1", ExtractClientIP(r, true),
		"XFF 全不合法应当 fallthrough 到 X-Real-IP，不得返回 garbage 或 RemoteAddr")
}

// TC-0864: behindProxy=true + 两头均空 → 回落到 RemoteAddr 剥端口后的 host。
func TestExtractClientIP_NoHeaders_FallbackRemoteAddr(t *testing.T) {
	r := httptest.NewRequest("POST", "/x", nil)
	r.RemoteAddr = "198.51.100.9:13579"

	assert.Equal(t, "198.51.100.9", ExtractClientIP(r, true),
		"所有代理头缺失时最终仍能回落到 RemoteAddr 剥端口")
}

// TC-0865: behindProxy=true + XFF 首段带两端空白 → trim 后仍解析合法，返回 trimmed 结果。
func TestExtractClientIP_XFFWhitespaceTrimmed(t *testing.T) {
	r := httptest.NewRequest("POST", "/x", nil)
	r.Header.Set("X-Forwarded-For", "  3.3.3.3  ,  4.4.4.4")

	assert.Equal(t, "3.3.3.3", ExtractClientIP(r, true),
		"XFF 首段 trim 后合法应当被采用；严禁保留首尾空白而误判")
}

// TC-0866: behindProxy=false —— 完全忽略 XFF / XRI，防止客户端伪造头。
func TestExtractClientIP_BehindProxyFalse_IgnoreHeaders(t *testing.T) {
	r := httptest.NewRequest("POST", "/x", nil)
	r.Header.Set("X-Forwarded-For", "1.1.1.1") // 应被忽略
	r.Header.Set("X-Real-IP", "2.2.2.2")       // 应被忽略
	r.RemoteAddr = "5.5.5.5:8080"

	assert.Equal(t, "5.5.5.5", ExtractClientIP(r, false),
		"behindProxy=false 时应完全忽略客户端注入的代理头")
}

// 补充：XFF 包含空段（"1.1.1.1,,2.2.2.2"）不应 panic，空段跳过后首段合法。
func TestExtractClientIP_XFFEmptySegmentsSkipped(t *testing.T) {
	r := httptest.NewRequest("POST", "/x", nil)
	r.Header.Set("X-Forwarded-For", ",,,1.1.1.1,2.2.2.2")

	assert.Equal(t, "1.1.1.1", ExtractClientIP(r, true),
		"XFF 中空段必须跳过，不得 panic 或返回空串")
}

// 补充：XFF 全为合法 IPv6 地址也应能返回首段。
func TestExtractClientIP_XFFIPv6FirstValid(t *testing.T) {
	r := httptest.NewRequest("POST", "/x", nil)
	r.Header.Set("X-Forwarded-For", "2001:db8::1, 2001:db8::2")

	assert.Equal(t, "2001:db8::1", ExtractClientIP(r, true),
		"IPv6 也是 net.ParseIP 合法值，XFF 首段应返回 IPv6")
}