152 lines
5.3 KiB
Go
152 lines
5.3 KiB
Go
package middleware
|
||
|
||
import (
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// newRateLimitRouter 构造一个使用 RateLimit 中间件、所有请求都返回 200 的测试 router。
|
||
// 返回 router 与下游被调用次数计数器，便于断言请求是否被放行。
|
||
func newRateLimitRouter(rps float64, burst int) (*gin.Engine, *int32) {
|
||
gin.SetMode(gin.TestMode)
|
||
var count int32
|
||
r := gin.New()
|
||
r.Use(RateLimit(rps, burst))
|
||
r.GET("/test", func(c *gin.Context) {
|
||
atomic.AddInt32(&count, 1)
|
||
c.Status(http.StatusOK)
|
||
})
|
||
return r, &count
|
||
}
|
||
|
||
// doRequest 发起一次测试请求，返回响应状态码。
|
||
func doRequest(r *gin.Engine, remoteAddr string) int {
|
||
w := httptest.NewRecorder()
|
||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||
req.RemoteAddr = remoteAddr
|
||
r.ServeHTTP(w, req)
|
||
return w.Code
|
||
}
|
||
|
||
// TestRateLimit_AllowUnderBurst 验证突发请求 ≤ burst 时全部放行。
|
||
func TestRateLimit_AllowUnderBurst(t *testing.T) {
|
||
r, count := newRateLimitRouter(0.1 /*rps*/, 5 /*burst*/)
|
||
|
||
// 5 个请求应全部放行（初始满桶 5 令牌）
|
||
for i := 0; i < 5; i++ {
|
||
if code := doRequest(r, "1.1.1.1:1234"); code != http.StatusOK {
|
||
t.Fatalf("请求 %d 应放行（200），实际 %d", i+1, code)
|
||
}
|
||
}
|
||
if got := atomic.LoadInt32(count); got != 5 {
|
||
t.Fatalf("下游应被调用 5 次，实际 %d", got)
|
||
}
|
||
}
|
||
|
||
// TestRateLimit_RejectOverBurst 验证突发请求 > burst 时返回 429。
|
||
func TestRateLimit_RejectOverBurst(t *testing.T) {
|
||
r, count := newRateLimitRouter(0.1 /*rps*/, 3 /*burst*/)
|
||
|
||
// 前 3 个放行
|
||
for i := 0; i < 3; i++ {
|
||
if code := doRequest(r, "2.2.2.2:1234"); code != http.StatusOK {
|
||
t.Fatalf("请求 %d 应放行（200），实际 %d", i+1, code)
|
||
}
|
||
}
|
||
// 第 4 个应被限流（429）
|
||
if code := doRequest(r, "2.2.2.2:1234"); code != http.StatusTooManyRequests {
|
||
t.Fatalf("第 4 个请求应被限流（429），实际 %d", code)
|
||
}
|
||
if got := atomic.LoadInt32(count); got != 3 {
|
||
t.Fatalf("下游应被调用 3 次，实际 %d", got)
|
||
}
|
||
}
|
||
|
||
// TestRateLimit_RefillTokens 验证令牌按 rps 速率补充。
|
||
// rps=10 表示每秒补 10 个令牌，即每 100ms 补 1 个；消耗 1 个后等待 ~150ms 应补回 1+ 个。
|
||
func TestRateLimit_RefillTokens(t *testing.T) {
|
||
r, count := newRateLimitRouter(10 /*rps*/, 1 /*burst*/)
|
||
|
||
// 第 1 个放行（消耗唯一令牌）
|
||
if code := doRequest(r, "3.3.3.3:1234"); code != http.StatusOK {
|
||
t.Fatalf("第 1 个请求应放行（200），实际 %d", code)
|
||
}
|
||
// 立即第 2 个应被限流（无令牌）
|
||
if code := doRequest(r, "3.3.3.3:1234"); code != http.StatusTooManyRequests {
|
||
t.Fatalf("第 2 个请求应被限流（429），实际 %d", code)
|
||
}
|
||
// 等待 150ms，应补充 1.5 个令牌（capped 至 burst=1）
|
||
time.Sleep(150 * time.Millisecond)
|
||
// 第 3 个应再次放行
|
||
if code := doRequest(r, "3.3.3.3:1234"); code != http.StatusOK {
|
||
t.Fatalf("等待补充后请求应放行（200），实际 %d", code)
|
||
}
|
||
if got := atomic.LoadInt32(count); got != 2 {
|
||
t.Fatalf("下游应被调用 2 次，实际 %d", got)
|
||
}
|
||
}
|
||
|
||
// TestRateLimit_PerIPIsolation 验证不同 IP 的桶相互独立。
|
||
func TestRateLimit_PerIPIsolation(t *testing.T) {
|
||
r, count := newRateLimitRouter(0.1 /*rps*/, 2 /*burst*/)
|
||
|
||
// IP A 消耗 2 个令牌（满桶）
|
||
for i := 0; i < 2; i++ {
|
||
if code := doRequest(r, "10.0.0.1:1234"); code != http.StatusOK {
|
||
t.Fatalf("IP A 请求 %d 应放行，实际 %d", i+1, code)
|
||
}
|
||
}
|
||
// IP A 第 3 个应被限流
|
||
if code := doRequest(r, "10.0.0.1:1234"); code != http.StatusTooManyRequests {
|
||
t.Fatalf("IP A 第 3 个请求应被限流，实际 %d", code)
|
||
}
|
||
// IP B 仍有独立满桶，应放行
|
||
if code := doRequest(r, "10.0.0.2:1234"); code != http.StatusOK {
|
||
t.Fatalf("IP B 第 1 个请求应放行（独立桶），实际 %d", code)
|
||
}
|
||
if got := atomic.LoadInt32(count); got != 3 {
|
||
t.Fatalf("下游应被调用 3 次（A×2 + B×1），实际 %d", got)
|
||
}
|
||
}
|
||
|
||
// TestRateLimit_CleanupExpiredBuckets 验证 10 分钟无访问的桶被清理。
|
||
// 使用短周期参数（50ms 间隔、500ms 过期）加速测试；直接构造 rateLimiter 而不经过 RateLimit
|
||
// （后者会启动 5min/10min 的清理 goroutine，无法在测试中观察）。
|
||
// 时序：创建旧桶（1h 前）→ 启动 cleanup → 等 100ms（旧桶被清，无新桶）→ 创建新桶 → 等 100ms（新桶未过 500ms，保留）。
|
||
func TestRateLimit_CleanupExpiredBuckets(t *testing.T) {
|
||
rl := &rateLimiter{rps: 1, burst: 1}
|
||
|
||
// 创建旧桶并人为把 lastTime 调到 1 小时前
|
||
b := rl.getBucket("9.9.9.9:1234")
|
||
b.mu.Lock()
|
||
b.lastTime = time.Now().Add(-1 * time.Hour)
|
||
b.mu.Unlock()
|
||
|
||
// 启动 cleanup，50ms 间隔、500ms 过期
|
||
go rl.cleanup(50*time.Millisecond, 500*time.Millisecond)
|
||
|
||
// 等 100ms 让 cleanup 执行（旧桶 lastTime=1h 前，远超 500ms，必被清）
|
||
time.Sleep(100 * time.Millisecond)
|
||
|
||
// 旧桶应已清理
|
||
if _, ok := rl.buckets.Load("9.9.9.9:1234"); ok {
|
||
t.Fatal("1 小时无访问的旧桶应被清理")
|
||
}
|
||
|
||
// 此时创建新鲜桶（lastTime = now）
|
||
rl.getBucket("8.8.8.8:1234")
|
||
|
||
// 再等 100ms 让 cleanup 再执行
|
||
time.Sleep(100 * time.Millisecond)
|
||
|
||
// 新鲜桶 lastTime 仅 ~100ms 前，未超 500ms maxIdle，应保留
|
||
if _, ok := rl.buckets.Load("8.8.8.8:1234"); !ok {
|
||
t.Fatal("新鲜桶（<500ms）不应被清理")
|
||
}
|
||
}
|