Security: - Enforce client_max_body_size in HTTP/2 and HTTP/3 adapters before buffering - Use trusted-proxy-aware client IP extraction for access control and rate limiting - Include Host and Vary headers in proxy cache key to prevent cache poisoning - Harden static file path traversal check with filepath.Clean + prefix validation - Honor proxy_ssl config for WebSocket upstream TLS Proxy/handler bugs: - Decrement WebSocket connection count immediately on return (not deferred in retry loop) - Remove ineffective headersPool - Coalesce concurrent background cache refreshes with singleflight Dev/build: - Fix Makefile run and test-config targets - Remove broken wget healthcheck from docker-compose - Add missing integration build tags - Fix go vet warnings in tests CI: - Run golangci-lint and integration tests in Gitea Actions Refs: docs/superpowers/plans/2026-06-17-fix-review-issues.md
966 lines
22 KiB
Go
966 lines
22 KiB
Go
// Package security 提供速率限制功能的测试。
|
||
//
|
||
// 该文件测试速率限制模块的各项功能,包括:
|
||
// - 速率限制器创建
|
||
// - 令牌桶算法
|
||
// - 令牌补充机制
|
||
// - 计数器重置
|
||
// - 连接数限制
|
||
// - 统计信息获取
|
||
//
|
||
// 作者:xfy
|
||
package security
|
||
|
||
import (
|
||
"net"
|
||
"runtime"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/valyala/fasthttp"
|
||
"rua.plus/lolly/internal/config"
|
||
"rua.plus/lolly/internal/testutil"
|
||
)
|
||
|
||
func TestNewRateLimiter(t *testing.T) {
|
||
tests := []struct {
|
||
cfg *config.RateLimitConfig
|
||
name string
|
||
wantErr bool
|
||
}{
|
||
{
|
||
name: "nil config",
|
||
cfg: nil,
|
||
wantErr: true,
|
||
},
|
||
{
|
||
name: "valid config",
|
||
cfg: &config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
},
|
||
},
|
||
{
|
||
name: "zero rate",
|
||
cfg: &config.RateLimitConfig{
|
||
RequestRate: 0,
|
||
Burst: 100,
|
||
},
|
||
wantErr: true,
|
||
},
|
||
{
|
||
name: "burst less than rate",
|
||
cfg: &config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 50,
|
||
},
|
||
wantErr: true,
|
||
},
|
||
{
|
||
name: "key by IP",
|
||
cfg: &config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
Key: "ip",
|
||
},
|
||
},
|
||
{
|
||
name: "key by header",
|
||
cfg: &config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
Key: "header",
|
||
},
|
||
},
|
||
{
|
||
name: "unknown key type",
|
||
cfg: &config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
Key: "unknown",
|
||
},
|
||
wantErr: true,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
rl, err := NewRateLimiter(tt.cfg)
|
||
if (err != nil) != tt.wantErr {
|
||
t.Errorf("NewRateLimiter() error = %v, wantErr %v", err, tt.wantErr)
|
||
}
|
||
if !tt.wantErr && rl == nil {
|
||
t.Error("Expected non-nil RateLimiter")
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestRateLimiterAllow(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 10,
|
||
Burst: 10,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
// Test burst allowance
|
||
key := "test-key"
|
||
|
||
// Should allow burst requests
|
||
for i := range 10 {
|
||
if !rl.Allow(key) {
|
||
t.Errorf("Expected request %d to be allowed", i+1)
|
||
}
|
||
}
|
||
|
||
// Next request should be denied (burst exhausted)
|
||
if rl.Allow(key) {
|
||
t.Error("Expected request to be denied after burst exhausted")
|
||
}
|
||
}
|
||
|
||
func TestRateLimiterTokenRefill(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100, // 100 tokens per second
|
||
Burst: 100,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
key := "refill-test"
|
||
|
||
// Exhaust the burst
|
||
for range 100 {
|
||
rl.Allow(key)
|
||
}
|
||
|
||
// Should be denied
|
||
if rl.Allow(key) {
|
||
t.Error("Expected request to be denied")
|
||
}
|
||
|
||
// Wait for token refill (10ms should give us 1 token at 100/s)
|
||
time.Sleep(15 * time.Millisecond)
|
||
|
||
// Should be allowed now
|
||
if !rl.Allow(key) {
|
||
t.Error("Expected request to be allowed after refill")
|
||
}
|
||
}
|
||
|
||
func TestRateLimiterReset(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 1,
|
||
Burst: 1,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
key := "reset-test"
|
||
|
||
// Exhaust
|
||
rl.Allow(key)
|
||
if rl.Allow(key) {
|
||
t.Error("Expected denial")
|
||
}
|
||
|
||
// Reset
|
||
rl.Reset(key)
|
||
|
||
// Should be allowed again
|
||
if !rl.Allow(key) {
|
||
t.Error("Expected request to be allowed after reset")
|
||
}
|
||
}
|
||
|
||
func TestRateLimiterResetAll(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 1,
|
||
Burst: 1,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
// Create multiple buckets
|
||
rl.Allow("key1")
|
||
rl.Allow("key2")
|
||
|
||
// Reset all
|
||
rl.ResetAll()
|
||
|
||
stats := rl.GetStats()
|
||
if stats.BucketCount != 0 {
|
||
t.Errorf("Expected 0 buckets after reset, got %d", stats.BucketCount)
|
||
}
|
||
}
|
||
|
||
func TestRateLimiterCleanup(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 100,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
// Create some buckets
|
||
rl.Allow("key1")
|
||
rl.Allow("key2")
|
||
|
||
// Cleanup with very short max age
|
||
rl.Cleanup(1 * time.Nanosecond)
|
||
|
||
stats := rl.GetStats()
|
||
if stats.BucketCount != 0 {
|
||
t.Errorf("Expected 0 buckets after cleanup, got %d", stats.BucketCount)
|
||
}
|
||
}
|
||
|
||
func TestRateLimiterProcess(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 100,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||
_, _ = ctx.WriteString("OK")
|
||
}
|
||
|
||
handler := mw.Process(nextHandler)
|
||
if handler == nil {
|
||
t.Error("Process() returned nil handler")
|
||
}
|
||
}
|
||
|
||
func TestRateLimiterGetStats(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
rl.Allow("key1")
|
||
rl.Allow("key2")
|
||
|
||
stats := rl.GetStats()
|
||
if stats.BucketCount != 2 {
|
||
t.Errorf("Expected BucketCount 2, got %d", stats.BucketCount)
|
||
}
|
||
if stats.Rate != 100 {
|
||
t.Errorf("Expected Rate 100, got %f", stats.Rate)
|
||
}
|
||
if stats.Burst != 200 {
|
||
t.Errorf("Expected Burst 200, got %f", stats.Burst)
|
||
}
|
||
|
||
// 测试优雅关闭
|
||
rl.StopCleanup()
|
||
}
|
||
|
||
func TestRateLimiterAutoCleanup(t *testing.T) {
|
||
// 使用自定义的创建方式,方便测试
|
||
cfg := &config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
Key: "ip",
|
||
}
|
||
|
||
mw, err := NewRateLimiter(cfg)
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
// 创建一些桶
|
||
rl.Allow("key1")
|
||
rl.Allow("key2")
|
||
rl.Allow("key3")
|
||
|
||
// 验证桶已创建
|
||
stats := rl.GetStats()
|
||
if stats.BucketCount != 3 {
|
||
t.Errorf("Expected 3 buckets, got %d", stats.BucketCount)
|
||
}
|
||
|
||
// 手动调用 Cleanup 模拟过期清理(使用很短的过期时间)
|
||
rl.Cleanup(1 * time.Nanosecond)
|
||
|
||
// 验证所有桶已被清理
|
||
stats = rl.GetStats()
|
||
if stats.BucketCount != 0 {
|
||
t.Errorf("Expected 0 buckets after cleanup, got %d", stats.BucketCount)
|
||
}
|
||
|
||
// 测试优雅关闭
|
||
rl.StopCleanup()
|
||
}
|
||
|
||
func TestRateLimiterStopCleanup(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
|
||
// 验证可以正常关闭
|
||
rl.StopCleanup()
|
||
|
||
// 再次调用不应 panic
|
||
rl.StopCleanup()
|
||
}
|
||
|
||
func TestNewConnLimiter(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
keyType string
|
||
max int
|
||
perKey bool
|
||
wantErr bool
|
||
}{
|
||
{
|
||
name: "global limit",
|
||
max: 100,
|
||
perKey: false,
|
||
},
|
||
{
|
||
name: "per-key by IP",
|
||
max: 10,
|
||
perKey: true,
|
||
keyType: "ip",
|
||
},
|
||
{
|
||
name: "per-key by header",
|
||
max: 10,
|
||
perKey: true,
|
||
keyType: "header",
|
||
},
|
||
{
|
||
name: "zero max",
|
||
max: 0,
|
||
wantErr: true,
|
||
},
|
||
{
|
||
name: "negative max",
|
||
max: -1,
|
||
wantErr: true,
|
||
},
|
||
{
|
||
name: "invalid key type",
|
||
max: 10,
|
||
perKey: true,
|
||
keyType: "invalid",
|
||
wantErr: true,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
cl, err := NewConnLimiter(tt.max, tt.perKey, tt.keyType)
|
||
if (err != nil) != tt.wantErr {
|
||
t.Errorf("NewConnLimiter() error = %v, wantErr %v", err, tt.wantErr)
|
||
}
|
||
if !tt.wantErr && cl == nil {
|
||
t.Error("Expected non-nil ConnLimiter")
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestConnLimiterGlobal(t *testing.T) {
|
||
cl, err := NewConnLimiter(2, false, "")
|
||
if err != nil {
|
||
t.Fatalf("NewConnLimiter() error: %v", err)
|
||
}
|
||
|
||
ctx := testutil.NewRequestCtx("GET", "/")
|
||
|
||
// First two should succeed
|
||
if !cl.Acquire(ctx) {
|
||
t.Error("Expected first acquire to succeed")
|
||
}
|
||
if !cl.Acquire(ctx) {
|
||
t.Error("Expected second acquire to succeed")
|
||
}
|
||
|
||
// Third should fail
|
||
if cl.Acquire(ctx) {
|
||
t.Error("Expected third acquire to fail")
|
||
}
|
||
|
||
// Release one
|
||
cl.Release(ctx)
|
||
|
||
// Should succeed now
|
||
if !cl.Acquire(ctx) {
|
||
t.Error("Expected acquire after release to succeed")
|
||
}
|
||
}
|
||
|
||
func TestConnLimiterMiddleware(t *testing.T) {
|
||
cl, err := NewConnLimiter(1, false, "")
|
||
if err != nil {
|
||
t.Fatalf("NewConnLimiter() error: %v", err)
|
||
}
|
||
|
||
middleware := cl.Middleware()
|
||
if middleware == nil {
|
||
t.Error("Expected non-nil middleware")
|
||
}
|
||
if middleware.Name() != "conn_limiter" {
|
||
t.Errorf("Expected name 'conn_limiter', got %s", middleware.Name())
|
||
}
|
||
}
|
||
|
||
// TestNewRateLimiter_SlidingWindow 测试滑动窗口算法
|
||
func TestNewRateLimiter_SlidingWindow(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 100,
|
||
Algorithm: "sliding_window",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
if mw == nil {
|
||
t.Error("Expected non-nil middleware for sliding_window")
|
||
}
|
||
if mw.Name() != "sliding_window_limiter" {
|
||
t.Errorf("Expected name 'sliding_window_limiter', got %s", mw.Name())
|
||
}
|
||
}
|
||
|
||
// TestNewRateLimiter_SlidingWindowDefault 测试滑动窗口默认配置
|
||
func TestNewRateLimiter_SlidingWindowDefault(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 100,
|
||
Algorithm: "sliding_window",
|
||
SlidingWindow: 0, // 使用默认值
|
||
SlidingWindowMode: "", // 使用默认值
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
if mw == nil {
|
||
t.Error("Expected non-nil middleware")
|
||
}
|
||
}
|
||
|
||
// TestNewRateLimiter_SlidingWindowPrecise 测试滑动窗口精确模式
|
||
func TestNewRateLimiter_SlidingWindowPrecise(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 100,
|
||
Algorithm: "sliding_window",
|
||
SlidingWindow: 1,
|
||
SlidingWindowMode: "precise",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
if mw == nil {
|
||
t.Error("Expected non-nil middleware")
|
||
}
|
||
}
|
||
|
||
// TestNewRateLimiter_UnknownAlgorithm 测试未知算法
|
||
func TestNewRateLimiter_UnknownAlgorithm(t *testing.T) {
|
||
_, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 100,
|
||
Algorithm: "unknown_algo",
|
||
})
|
||
if err == nil {
|
||
t.Error("NewRateLimiter() should return error for unknown algorithm")
|
||
}
|
||
}
|
||
|
||
// TestSlidingWindowLimiterWrapper_Process 测试滑动窗口包装器的 Process 方法
|
||
func TestSlidingWindowLimiterWrapper_Process(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 100,
|
||
Algorithm: "sliding_window",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
called := false
|
||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||
called = true
|
||
}
|
||
|
||
handler := mw.Process(nextHandler)
|
||
if handler == nil {
|
||
t.Fatal("Process() returned nil handler")
|
||
}
|
||
|
||
ctx := testutil.NewRequestCtx("GET", "/test")
|
||
handler(ctx)
|
||
|
||
if !called {
|
||
t.Error("Next handler should be called when rate limit allows")
|
||
}
|
||
}
|
||
|
||
// TestSlidingWindowLimiterWrapper_ProcessDenied 测试滑动窗口拒绝请求
|
||
func TestSlidingWindowLimiterWrapper_ProcessDenied(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 1,
|
||
Burst: 1,
|
||
Algorithm: "sliding_window",
|
||
SlidingWindowMode: "precise",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
callCount := 0
|
||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||
callCount++
|
||
}
|
||
|
||
handler := mw.Process(nextHandler)
|
||
ctx := testutil.NewRequestCtx("GET", "/test")
|
||
|
||
// 第一个请求应该被允许
|
||
handler(ctx)
|
||
// 第二个请求应该被拒绝
|
||
handler(ctx)
|
||
|
||
if callCount != 1 {
|
||
t.Errorf("Expected next handler to be called once, got %d", callCount)
|
||
}
|
||
}
|
||
|
||
// TestRateLimiter_GetRetryAfter 测试 getRetryAfter 方法
|
||
func TestRateLimiter_GetRetryAfter(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 10,
|
||
Burst: 10,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
defer rl.StopCleanup()
|
||
|
||
// 测试不存在的键
|
||
retryAfter := rl.getRetryAfter("nonexistent")
|
||
if retryAfter != 1 {
|
||
t.Errorf("getRetryAfter(nonexistent) = %d, want 1", retryAfter)
|
||
}
|
||
|
||
// 创建一个桶并消耗令牌
|
||
key := "test-key"
|
||
for range 10 {
|
||
rl.Allow(key)
|
||
}
|
||
|
||
// 获取重试时间
|
||
retryAfter = rl.getRetryAfter(key)
|
||
if retryAfter < 1 {
|
||
t.Errorf("getRetryAfter() = %d, want at least 1", retryAfter)
|
||
}
|
||
}
|
||
|
||
// TestKeyByIP 测试 keyByIP 函数
|
||
func TestKeyByIP(t *testing.T) {
|
||
ctx := testutil.NewRequestCtx("GET", "/test")
|
||
|
||
key := keyByIP(ctx)
|
||
if key == "" {
|
||
t.Error("keyByIP() should return non-empty string")
|
||
}
|
||
if key == "unknown" {
|
||
t.Error("keyByIP() should return IP address, not 'unknown'")
|
||
}
|
||
}
|
||
|
||
// TestKeyByHeader 测试 keyByHeader 函数
|
||
func TestKeyByHeader(t *testing.T) {
|
||
// 测试有头部的情况
|
||
ctx := testutil.NewRequestCtx("GET", "/test")
|
||
ctx.Request.Header.Set("X-RateLimit-Key", "custom-key")
|
||
|
||
key := keyByHeader(ctx)
|
||
if key != "custom-key" {
|
||
t.Errorf("keyByHeader() = %q, want 'custom-key'", key)
|
||
}
|
||
|
||
// 测试没有头部的情况(应该回退到 IP)
|
||
ctx2 := testutil.NewRequestCtx("GET", "/test")
|
||
key2 := keyByHeader(ctx2)
|
||
if key2 == "" {
|
||
t.Error("keyByHeader() should fallback to IP when header not set")
|
||
}
|
||
}
|
||
|
||
// TestConnLimiter_PerKey 测试按键限制
|
||
func TestConnLimiter_PerKey(t *testing.T) {
|
||
cl, err := NewConnLimiter(2, true, "ip")
|
||
if err != nil {
|
||
t.Fatalf("NewConnLimiter() error: %v", err)
|
||
}
|
||
|
||
ctx := testutil.NewRequestCtx("GET", "/")
|
||
|
||
// 同一个键的前两个应该成功
|
||
if !cl.Acquire(ctx) {
|
||
t.Error("Expected first acquire to succeed")
|
||
}
|
||
if !cl.Acquire(ctx) {
|
||
t.Error("Expected second acquire to succeed")
|
||
}
|
||
|
||
// 第三个应该失败
|
||
if cl.Acquire(ctx) {
|
||
t.Error("Expected third acquire to fail")
|
||
}
|
||
|
||
// 释放一个
|
||
cl.Release(ctx)
|
||
|
||
// 现在应该成功
|
||
if !cl.Acquire(ctx) {
|
||
t.Error("Expected acquire after release to succeed")
|
||
}
|
||
}
|
||
|
||
// TestConnLimiter_ReleaseUnderflow 测试 Release 下溢保护
|
||
func TestConnLimiter_ReleaseUnderflow(t *testing.T) {
|
||
cl, err := NewConnLimiter(2, true, "ip")
|
||
if err != nil {
|
||
t.Fatalf("NewConnLimiter() error: %v", err)
|
||
}
|
||
|
||
ctx := testutil.NewRequestCtx("GET", "/")
|
||
|
||
// 在没有 Acquire 的情况下 Release(测试下溢保护)
|
||
cl.Release(ctx) // 不应该 panic
|
||
|
||
// 验证计数不会变成负数
|
||
cl.Acquire(ctx)
|
||
cl.Acquire(ctx)
|
||
if cl.Acquire(ctx) {
|
||
t.Error("Expected third acquire to fail")
|
||
}
|
||
}
|
||
|
||
// TestConnLimiterMiddleware_Process 测试连接限制中间件 Process
|
||
func TestConnLimiterMiddleware_Process(t *testing.T) {
|
||
cl, err := NewConnLimiter(1, false, "")
|
||
if err != nil {
|
||
t.Fatalf("NewConnLimiter() error: %v", err)
|
||
}
|
||
|
||
mw := cl.Middleware()
|
||
|
||
called := false
|
||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||
called = true
|
||
ctx.SetStatusCode(200)
|
||
}
|
||
|
||
handler := mw.Process(nextHandler)
|
||
ctx := testutil.NewRequestCtx("GET", "/test")
|
||
|
||
// 第一个请求应该成功
|
||
handler(ctx)
|
||
if !called {
|
||
t.Error("Next handler should be called")
|
||
}
|
||
if ctx.Response.StatusCode() != 200 {
|
||
t.Errorf("Status = %d, want 200", ctx.Response.StatusCode())
|
||
}
|
||
}
|
||
|
||
// TestConnLimiterMiddleware_ProcessLimitExceeded 测试连接限制超出
|
||
func TestConnLimiterMiddleware_ProcessLimitExceeded(t *testing.T) {
|
||
cl, err := NewConnLimiter(1, false, "")
|
||
if err != nil {
|
||
t.Fatalf("NewConnLimiter() error: %v", err)
|
||
}
|
||
|
||
mw := cl.Middleware()
|
||
|
||
callCount := 0
|
||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||
callCount++
|
||
ctx.SetStatusCode(200)
|
||
}
|
||
|
||
handler := mw.Process(nextHandler)
|
||
|
||
// 用尽连接限制
|
||
ctx1 := testutil.NewRequestCtx("GET", "/test1")
|
||
cl.Acquire(ctx1) // 手动占用一个槽位
|
||
|
||
// 现在应该无法获取新的连接
|
||
ctx2 := testutil.NewRequestCtx("GET", "/test2")
|
||
handler(ctx2)
|
||
|
||
if callCount != 0 {
|
||
t.Error("Next handler should NOT be called when limit exceeded")
|
||
}
|
||
if ctx2.Response.StatusCode() != 503 {
|
||
t.Errorf("Status = %d, want 503", ctx2.Response.StatusCode())
|
||
}
|
||
}
|
||
|
||
// TestNewSlidingWindowLimiterWrapper_Error 测试滑动窗口包装器错误情况
|
||
func TestNewSlidingWindowLimiterWrapper_Error(t *testing.T) {
|
||
_, err := NewSlidingWindowLimiterWrapper(&config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Key: "invalid_key_type",
|
||
}, time.Second, false)
|
||
if err == nil {
|
||
t.Error("NewSlidingWindowLimiterWrapper should return error for invalid key type")
|
||
}
|
||
}
|
||
|
||
// TestRateLimiter_Name 测试 RateLimiter Name 方法
|
||
func TestRateLimiter_Name(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 10,
|
||
Burst: 10,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
defer rl.StopCleanup()
|
||
|
||
if rl.Name() != "rate_limiter" {
|
||
t.Errorf("Name() = %q, want 'rate_limiter'", rl.Name())
|
||
}
|
||
}
|
||
|
||
// TestRateLimiter_ProcessDenied 测试限流拒绝
|
||
func TestRateLimiter_ProcessDenied(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 1,
|
||
Burst: 1,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
defer rl.StopCleanup()
|
||
|
||
callCount := 0
|
||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||
callCount++
|
||
}
|
||
|
||
handler := rl.Process(nextHandler)
|
||
|
||
// 第一个请求应该成功
|
||
ctx1 := testutil.NewRequestCtx("GET", "/test")
|
||
handler(ctx1)
|
||
|
||
// 第二个请求应该被限流(使用不同的 context)
|
||
ctx2 := testutil.NewRequestCtx("GET", "/test")
|
||
handler(ctx2)
|
||
|
||
if callCount != 1 {
|
||
t.Errorf("Expected next handler to be called once, got %d", callCount)
|
||
}
|
||
|
||
// 检查第二个请求的状态码
|
||
if ctx2.Response.StatusCode() != 429 {
|
||
t.Errorf("Status = %d, want 429", ctx2.Response.StatusCode())
|
||
}
|
||
}
|
||
|
||
// TestKeyByIP_Unknown 测试无法获取 IP 的情况
|
||
func TestKeyByIP_Unknown(t *testing.T) {
|
||
// 创建一个没有设置远程地址的上下文
|
||
ctx := &fasthttp.RequestCtx{}
|
||
ctx.Request.SetRequestURI("/test")
|
||
|
||
key := keyByIP(ctx)
|
||
// netutil.ExtractClientIPNet 会返回默认值 0.0.0.0 而不是 nil
|
||
// 所以这里验证返回值不是空的即可
|
||
if key == "" {
|
||
t.Error("keyByIP() should return non-empty string")
|
||
}
|
||
}
|
||
|
||
// TestConnLimiter_MiddlewareIdentity 验证 Middleware() 返回相同实例
|
||
func TestConnLimiter_MiddlewareIdentity(t *testing.T) {
|
||
cl, err := NewConnLimiter(100, false, "")
|
||
if err != nil {
|
||
t.Fatalf("NewConnLimiter() error: %v", err)
|
||
}
|
||
|
||
// Middleware() 应该返回自身
|
||
middleware := cl.Middleware()
|
||
if middleware != cl {
|
||
t.Error("Middleware() should return the same ConnLimiter instance")
|
||
}
|
||
|
||
// 验证返回的实例实现了 Middleware 接口
|
||
if middleware.Name() != "conn_limiter" {
|
||
t.Errorf("Name() = %s, want 'conn_limiter'", middleware.Name())
|
||
}
|
||
}
|
||
|
||
// TestRateLimiter_MultipleCreateDestroy 验证多次创建/销毁限流器不会泄漏 goroutine 也不会 panic。
|
||
func TestRateLimiter_MultipleCreateDestroy(t *testing.T) {
|
||
cfg := &config.RateLimitConfig{
|
||
RequestRate: 100,
|
||
Burst: 200,
|
||
}
|
||
|
||
// 预热:先创建并销毁一次,让 runtime 稳定
|
||
mw, err := NewRateLimiter(cfg)
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
rl := mw.(*RateLimiter)
|
||
rl.StopCleanup()
|
||
|
||
before := runtime.NumGoroutine()
|
||
|
||
for range 10 {
|
||
mw, err = NewRateLimiter(cfg)
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
mw.(*RateLimiter).StopCleanup()
|
||
// 重复调用不应 panic
|
||
mw.(*RateLimiter).StopCleanup()
|
||
}
|
||
|
||
after := runtime.NumGoroutine()
|
||
// 允许少量波动,但不应有明显泄漏
|
||
if after-before > 5 {
|
||
t.Fatalf("goroutine leak detected: before=%d, after=%d", before, after)
|
||
}
|
||
}
|
||
|
||
// TestRateLimiter_TrustedProxies 验证令牌桶限流器使用可信代理安全提取客户端 IP。
|
||
func TestRateLimiter_TrustedProxies(t *testing.T) {
|
||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||
RequestRate: 1,
|
||
Burst: 1,
|
||
}, "10.0.0.0/8")
|
||
if err != nil {
|
||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||
}
|
||
|
||
rl, ok := mw.(*RateLimiter)
|
||
if !ok {
|
||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||
}
|
||
defer rl.StopCleanup()
|
||
|
||
callCount := 0
|
||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||
callCount++
|
||
ctx.SetStatusCode(200)
|
||
}
|
||
handler := rl.Process(nextHandler)
|
||
|
||
// 来自可信代理的请求,使用 X-Forwarded-For 中的真实客户端 IP
|
||
ctx1 := testutil.NewRequestCtx("GET", "/test")
|
||
ctx1.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")})
|
||
ctx1.Request.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||
handler(ctx1)
|
||
if ctx1.Response.StatusCode() != 200 {
|
||
t.Errorf("first trusted request should be allowed, got %d", ctx1.Response.StatusCode())
|
||
}
|
||
|
||
// 同一个真实客户端 IP 的第二个请求应被拒绝
|
||
ctx2 := testutil.NewRequestCtx("GET", "/test")
|
||
ctx2.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.2")})
|
||
ctx2.Request.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||
handler(ctx2)
|
||
if ctx2.Response.StatusCode() != 429 {
|
||
t.Errorf("second request from same client IP should be rate limited, got %d", ctx2.Response.StatusCode())
|
||
}
|
||
|
||
// 不同真实客户端 IP 的请求应被允许
|
||
ctx3 := testutil.NewRequestCtx("GET", "/test")
|
||
ctx3.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")})
|
||
ctx3.Request.Header.Set("X-Forwarded-For", "5.6.7.8")
|
||
handler(ctx3)
|
||
if ctx3.Response.StatusCode() != 200 {
|
||
t.Errorf("request from different client IP should be allowed, got %d", ctx3.Response.StatusCode())
|
||
}
|
||
|
||
// 不可信代理:忽略 XFF,使用 remoteAddr 作为限流键
|
||
ctx4 := testutil.NewRequestCtx("GET", "/test")
|
||
ctx4.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("192.168.1.1")})
|
||
ctx4.Request.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||
handler(ctx4)
|
||
if ctx4.Response.StatusCode() != 200 {
|
||
t.Errorf("request from untrusted proxy should use remote IP and be allowed, got %d", ctx4.Response.StatusCode())
|
||
}
|
||
}
|