92 lines
3.1 KiB
Go
92 lines
3.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"hfb_sys/backend/internal/modules/auth"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func runRateLimitedRequest(router *gin.Engine, bearerToken string) int {
|
|
req := httptest.NewRequest(http.MethodGet, "/ping", nil)
|
|
if bearerToken != "" {
|
|
req.Header.Set("Authorization", "Bearer "+bearerToken)
|
|
}
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
return recorder.Code
|
|
}
|
|
|
|
func newRateLimitRouter(limit int, rdb interface{}, handler gin.HandlerFunc) *gin.Engine {
|
|
gin.SetMode(gin.TestMode)
|
|
router := gin.New()
|
|
router.Use(handler)
|
|
router.GET("/ping", func(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
})
|
|
return router
|
|
}
|
|
|
|
func adminAccessToken(t *testing.T, jwtManager *auth.JWTManager, adminID uint64, phone string) string {
|
|
t.Helper()
|
|
pair, err := jwtManager.GenerateSubjectPairWithVersion(adminID, phone, "admin", 0)
|
|
if err != nil {
|
|
t.Fatalf("生成后台令牌失败: %v", err)
|
|
}
|
|
return pair.AccessToken
|
|
}
|
|
|
|
func TestAdminAwareRateLimitIsolatesAdminsBehindSharedIP(t *testing.T) {
|
|
jwtManager := auth.NewJWTManager("test-secret")
|
|
adminA := adminAccessToken(t, jwtManager, 101, "13800000001")
|
|
adminB := adminAccessToken(t, jwtManager, 102, "13800000002")
|
|
router := newRateLimitRouter(2, nil, AdminAwareRateLimitPerMinute(2, nil, jwtManager))
|
|
|
|
for i := 0; i < 2; i++ {
|
|
if code := runRateLimitedRequest(router, adminA); code != http.StatusOK {
|
|
t.Fatalf("客服A 第 %d 次请求 = %d, want 200", i+1, code)
|
|
}
|
|
}
|
|
if code := runRateLimitedRequest(router, adminA); code != http.StatusTooManyRequests {
|
|
t.Fatalf("客服A 超额请求 = %d, want 429", code)
|
|
}
|
|
if code := runRateLimitedRequest(router, adminB); code != http.StatusOK {
|
|
t.Fatalf("同出口IP的客服B 应不受客服A影响, got %d, want 200", code)
|
|
}
|
|
}
|
|
|
|
func TestAdminAwareRateLimitFallsBackToIP(t *testing.T) {
|
|
jwtManager := auth.NewJWTManager("test-secret")
|
|
userPair, err := jwtManager.GenerateSubjectPairWithVersion(201, "13900000001", "user", 0)
|
|
if err != nil {
|
|
t.Fatalf("生成用户令牌失败: %v", err)
|
|
}
|
|
router := newRateLimitRouter(2, nil, AdminAwareRateLimitPerMinute(2, nil, jwtManager))
|
|
|
|
for i := 0; i < 2; i++ {
|
|
if code := runRateLimitedRequest(router, ""); code != http.StatusOK {
|
|
t.Fatalf("匿名请求 第 %d 次 = %d, want 200", i+1, code)
|
|
}
|
|
}
|
|
if code := runRateLimitedRequest(router, ""); code != http.StatusTooManyRequests {
|
|
t.Fatalf("匿名请求超额 = %d, want 429", code)
|
|
}
|
|
// 用户令牌无法按 admin 身份解析，应与匿名共享同一个 IP 桶。
|
|
if code := runRateLimitedRequest(router, userPair.AccessToken); code != http.StatusTooManyRequests {
|
|
t.Fatalf("用户令牌应回退到IP维度, got %d, want 429", code)
|
|
}
|
|
}
|
|
|
|
func TestLegacyRateLimitPerMinuteKeepsIPKey(t *testing.T) {
|
|
router := newRateLimitRouter(1, nil, RateLimitPerMinute(1, nil))
|
|
if code := runRateLimitedRequest(router, ""); code != http.StatusOK {
|
|
t.Fatalf("首次匿名请求 = %d, want 200", code)
|
|
}
|
|
if code := runRateLimitedRequest(router, ""); code != http.StatusTooManyRequests {
|
|
t.Fatalf("第二次匿名请求 = %d, want 429", code)
|
|
}
|
|
}
|