diff --git a/bench_test.go b/bench_test.go new file mode 100644 index 0000000..3e13310 --- /dev/null +++ b/bench_test.go @@ -0,0 +1,106 @@ +package main + +import ( + "fmt" + "sync" + "testing" + "time" +) + +// allowOld 复刻优化前的限流实现，仅用于基准对比： +// 每个请求都加全局锁并遍历所有 IP，且从不删除过期 IP 键。 +func allowOld(entries map[string][]time.Time, mu *sync.Mutex, ip string, limit int64, now time.Time) bool { + mu.Lock() + for k, times := range entries { + var recent []time.Time + for _, ts := range times { + if now.Sub(ts) <= time.Minute { + recent = append(recent, ts) + } + } + entries[k] = recent + } + if limit > 0 && len(entries[ip]) > int(limit) { + mu.Unlock() + return false + } + entries[ip] = append(entries[ip], now) + mu.Unlock() + return true +} + +func makeIPs(n int) []string { + ips := make([]string, n) + for i := range ips { + ips[i] = fmt.Sprintf("10.%d.%d.%d", (i>>16)&0xff, (i>>8)&0xff, i&0xff) + } + return ips +} + +func benchNew(b *testing.B, ipCount int, limit int64) { + l := newRateLimiter() + l.sweepEvery = time.Hour // 排除回收开销，只衡量单请求处理成本 + ips := makeIPs(ipCount) + now := time.Now() + // 预填满各 IP 的窗口，使热路径处于「已满」状态。 + for _, ip := range ips { + for j := int64(0); j < limit; j++ { + l.allow(ip, limit, now) + } + } + + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + i := 0 + for pb.Next() { + l.allow(ips[i%ipCount], limit, now) + i++ + } + }) +} + +func benchOld(b *testing.B, ipCount int, limit int64) { + entries := make(map[string][]time.Time, ipCount) + var mu sync.Mutex + ips := makeIPs(ipCount) + now := time.Now() + for _, ip := range ips { + for j := int64(0); j < limit; j++ { + entries[ip] = append(entries[ip], now) + } + } + + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + i := 0 + for pb.Next() { + allowOld(entries, &mu, ips[i%ipCount], limit, now) + i++ + } + }) +} + +func BenchmarkLimiterNew_SingleIP(b *testing.B) { benchNew(b, 1, 100) } +func BenchmarkLimiterNew_ManyIPs(b *testing.B) { benchNew(b, 1000, 100) } +func BenchmarkLimiterOld_SingleIP(b *testing.B) { benchOld(b, 1, 100) } +func BenchmarkLimiterOld_ManyIPs(b *testing.B) { benchOld(b, 1000, 100) } + +// BenchmarkLimiterSweep 衡量一次回收（清理大量过期 IP）的开销。 +func BenchmarkLimiterSweep(b *testing.B) { + const n = 5000 + ips := makeIPs(n) + base := time.Now() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + l := newRateLimiter() + l.sweepEvery = time.Hour + for _, ip := range ips { + l.allow(ip, 10, base) + } + l.sweepEvery = 0 + l.allow("192.168.0.1", 10, base.Add(2*time.Minute)) + } +} diff --git a/config_test.go b/config_test.go new file mode 100644 index 0000000..8897c95 --- /dev/null +++ b/config_test.go @@ -0,0 +1,102 @@ +package main + +import ( + "fmt" + "io" + "os" + "path/filepath" + "sync" + "testing" +) + +// 配置重载与请求读取并发时不应发生数据竞争，且读取方始终看到完整一致的快照。 +func TestConfigReloadConcurrentWithRead(t *testing.T) { + // 覆盖日志输出，避免大量重载日志污染测试输出。 + oldOut := log.Out + log.SetOutput(io.Discard) + t.Cleanup(func() { log.SetOutput(oldOut) }) + + dir := t.TempDir() + path := filepath.Join(dir, "config.json") + + write := func(n int) error { + data := fmt.Sprintf( + `{"debug":false,"whiteList":["v%d"],"requestLimit":{"limitRate":%d}}`, + n, n, + ) + return os.WriteFile(path, []byte(data), 0o644) + } + + if err := write(0); err != nil { + t.Fatal(err) + } + if err := loadConfigFrom(path); err != nil { + t.Fatalf("初始加载失败: %v", err) + } + + const rounds = 200 + var wg sync.WaitGroup + stop := make(chan struct{}) + + // 写方：持续重载配置。 + wg.Add(1) + go func() { + defer wg.Done() + defer close(stop) + for n := 1; n <= rounds; n++ { + if err := write(n); err != nil { + t.Errorf("写入配置失败: %v", err) + return + } + if err := loadConfigFrom(path); err != nil { + t.Errorf("重载配置失败: %v", err) + return + } + } + }() + + // 读方：持续读取，验证同一快照内字段一致（whiteList 与 limitRate 对应）。 + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + cfg := config.Load() + if cfg == nil { + t.Error("config 不应为 nil") + return + } + if len(cfg.WhiteList) == 1 { + want := cfg.WhiteList[0] + got := fmt.Sprintf("v%d", cfg.RequestLimit.LimitRate) + if want != got { + t.Errorf("快照不一致: whiteList=%q, limitRate=%d", want, cfg.RequestLimit.LimitRate) + return + } + } + } + } + }() + + wg.Wait() + + // 最终配置应等于最后一次写入。 + got := config.Load() + if got.RequestLimit.LimitRate != rounds { + t.Fatalf("最终 limitRate 应为 %d，实际 %d", rounds, got.RequestLimit.LimitRate) + } +} + +// 加载不存在的文件应返回错误且不破坏已有配置。 +func TestLoadConfigFromMissingFileKeepsOldConfig(t *testing.T) { + before := config.Load() + if err := loadConfigFrom(filepath.Join(t.TempDir(), "nope.json")); err == nil { + t.Fatal("加载不存在的文件应返回错误") + } + if config.Load() != before { + t.Fatal("加载失败时不应替换已有配置") + } +} diff --git a/main.go b/main.go index e6288c3..3df014c 100644 --- a/main.go +++ b/main.go @@ -11,6 +11,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "text/template" "time" @@ -49,16 +50,16 @@ var ( // checkURL 使用：入口 + 资源 CDN exps = append(append([]*regexp.Regexp{}, entryExps...), assetExps...) httpClient *http.Client - config *Config - configLock sync.RWMutex - log = logrus.New() + // config 以不可变快照方式发布：loadConfig 构造新对象后原子替换， + // 请求路径只做 atomic.Load，避免配置重载与请求处理之间的数据竞争和锁开销。 + config atomic.Pointer[Config] + log = logrus.New() tmpl *template.Template tmpl404 *template.Template - // 新增：用于记录每个 IP 的请求时间 - ipRequests = make(map[string][]time.Time) - ipLock sync.Mutex + // 每个 IP 独立的滑动窗口限流器 + limiter = newRateLimiter() ) type Config struct { @@ -85,6 +86,80 @@ func init() { FullTimestamp: true, } log.SetFormatter(&formatter) + // 保证在任何请求到来前 config 非 nil。 + config.Store(&Config{}) +} + +// rateLimiter 按 IP 记录滑动窗口内的请求时间。 +// 相比原来「每个请求都遍历全部 IP」的实现，这里只处理当前 IP， +// 并以固定间隔（sweepEvery）回收长期不活跃的 IP，避免数据结构无界增长。 +type rateLimiter struct { + mu sync.Mutex + entries map[string][]time.Time + lastSweep time.Time + window time.Duration + sweepEvery time.Duration +} + +func newRateLimiter() *rateLimiter { + return &rateLimiter{ + entries: make(map[string][]time.Time), + lastSweep: time.Now(), + window: time.Minute, + sweepEvery: time.Minute, + } +} + +// allow 判断并记录 ip 的这一次请求，返回是否放行。 +// limit <= 0 表示关闭限流：直接放行且不记录、不回收。 +// window 内允许恰好 limit 次请求，第 limit+1 次起返回 false。 +func (l *rateLimiter) allow(ip string, limit int64, now time.Time) bool { + if limit <= 0 { + return true + } + + l.mu.Lock() + defer l.mu.Unlock() + + // 低频回收：最多每 sweepEvery 触发一次，避免每请求全表扫描。 + if now.Sub(l.lastSweep) >= l.sweepEvery { + l.sweepLocked(now) + } + + cutoff := now.Add(-l.window) + times := pruneBefore(l.entries[ip], cutoff) + if len(times) >= int(limit) { + l.entries[ip] = times + return false + } + l.entries[ip] = append(times, now) + return true +} + +// sweepLocked 删除窗口内已无请求的 IP 键，并裁剪仍活跃 IP 的过期记录。 +// 必须在持有 l.mu 时调用。 +func (l *rateLimiter) sweepLocked(now time.Time) { + l.lastSweep = now + cutoff := now.Add(-l.window) + for ip, times := range l.entries { + if len(times) == 0 || times[len(times)-1].Before(cutoff) { + delete(l.entries, ip) + continue + } + if pruned := pruneBefore(times, cutoff); len(pruned) != len(times) { + l.entries[ip] = pruned + } + } +} + +// pruneBefore 丢弃时间戳早于 cutoff 的前缀记录。记录按时间升序追加， +// 因此可安全地从头裁剪；保留 t == cutoff 的记录以对齐原有时间窗口语义。 +func pruneBefore(times []time.Time, cutoff time.Time) []time.Time { + i := 0 + for i < len(times) && times[i].Before(cutoff) { + i++ + } + return times[i:] } func customLogger() gin.HandlerFunc { @@ -162,9 +237,7 @@ func main() { log.Info("start HTTP server @ ", httpBase) router.GET("/", func(c *gin.Context) { - configLock.RLock() - analyticsURL := config.AnalyticsURL - configLock.RUnlock() + analyticsURL := config.Load().AnalyticsURL data := struct { AnalyticsURL string @@ -186,6 +259,9 @@ func main() { } func handler(c *gin.Context) { + // 取一次不可变配置快照，保证整个请求内看到一致的配置，且无锁、无数据竞争。 + cfg := config.Load() + rawPath := strings.TrimPrefix(c.Request.URL.RequestURI(), "/") for strings.HasPrefix(rawPath, "/") { @@ -198,11 +274,11 @@ func handler(c *gin.Context) { matches := checkURL(rawPath) if matches != nil { - if len(config.WhiteList) > 0 && !checkList(matches, config.WhiteList) { + if len(cfg.WhiteList) > 0 && !checkList(matches, cfg.WhiteList) { render404(c, tmpl404) return } - if len(config.BlackList) > 0 && checkList(matches, config.BlackList) { + if len(cfg.BlackList) > 0 && checkList(matches, cfg.BlackList) { render404(c, tmpl404) return } @@ -215,36 +291,17 @@ func handler(c *gin.Context) { rawPath = strings.Replace(rawPath, "/blob/", "/raw/", 1) } - // 新增：基于时间的速率限制 + // 基于时间的速率限制：只处理当前 IP，窗口内允许恰好 limitRate 次。 clientIP := c.ClientIP() - ipLock.Lock() - limitRate := config.RequestLimit.LimitRate - now := time.Now() - // 移除一分钟前的请求记录 - for ip, times := range ipRequests { - var recentTimes []time.Time - for _, t := range times { - if now.Sub(t) <= time.Minute { - recentTimes = append(recentTimes, t) - } - } - ipRequests[ip] = recentTimes - } - log.Debugf("clientIP: %s Rate: %d\n", clientIP, len(ipRequests[clientIP])) - // 检查当前 IP 的请求次数是否超过限制（如果 limitRate <= 0 则不限制） - if limitRate > 0 && len(ipRequests[clientIP]) > int(limitRate) { - ipLock.Unlock() + limitRate := cfg.RequestLimit.LimitRate + if !limiter.allow(clientIP, limitRate, time.Now()) { + log.Debugf("clientIP: %s rate limited (limit=%d)", clientIP, limitRate) c.String(http.StatusTooManyRequests, "Too Many Requests.") return } - // 记录当前请求时间 - ipRequests[clientIP] = append(ipRequests[clientIP], now) - ipLock.Unlock() // 限制访问 IP - configLock.RLock() - limitAddr := config.RequestLimit.LimitAddr - configLock.RUnlock() + limitAddr := cfg.RequestLimit.LimitAddr if len(limitAddr) > 0 { for _, ip := range limitAddr { if clientIP == ip { @@ -255,9 +312,7 @@ func handler(c *gin.Context) { } // 限制请求参数 - configLock.RLock() - limitParm := config.RequestLimit.LimitParm - configLock.RUnlock() + limitParm := cfg.RequestLimit.LimitParm if len(limitParm) > 0 { for key, value := range limitParm { if c.Request.Header.Get(key) == value { @@ -303,9 +358,7 @@ func proxy(c *gin.Context, u string) { return } - configLock.RLock() - limitSize := config.RequestLimit.LimitSize * 1024 * 1024 // Convert MB to Bytes - configLock.RUnlock() + limitSize := config.Load().RequestLimit.LimitSize * 1024 * 1024 // Convert MB to Bytes // 如果 limitSize <= 0 则不限制文件大小 if limitSize > 0 && size > limitSize { @@ -348,11 +401,16 @@ func proxy(c *gin.Context, u string) { } func loadConfig() { + _ = loadConfigFrom("config.json") +} + +// loadConfigFrom 从指定路径加载配置并原子发布。返回错误便于测试验证。 +func loadConfigFrom(path string) error { log.Info("loading config...") - file, err := os.Open("config.json") + file, err := os.Open(path) if err != nil { log.Errorf("load config error: %v\n", err) - return + return err } defer func(file *os.File) { err := file.Close() @@ -365,17 +423,17 @@ func loadConfig() { decoder := json.NewDecoder(file) if err := decoder.Decode(&newConfig); err != nil { log.Errorf("decod config error: %v\n", err) - return + return err } - configLock.Lock() - config = &newConfig + // 构造完成后整体原子替换，读方无需加锁，也不会看到半更新的配置。 + config.Store(&newConfig) if newConfig.Debug { log.SetLevel(logrus.DebugLevel) } else { log.SetLevel(logrus.InfoLevel) } - configLock.Unlock() + return nil } func checkURL(u string) []string { @@ -408,10 +466,9 @@ func checkList(matches, list []string) bool { } func render404(c *gin.Context, tmpl *template.Template) { - configLock.RLock() - analyticsURL := config.AnalyticsURL - supportImageURL := config.SupportImageURL - configLock.RUnlock() + cfg := config.Load() + analyticsURL := cfg.AnalyticsURL + supportImageURL := cfg.SupportImageURL data := struct { AnalyticsURL string diff --git a/proxy_test.go b/proxy_test.go new file mode 100644 index 0000000..3d3d0fb --- /dev/null +++ b/proxy_test.go @@ -0,0 +1,151 @@ +package main + +import ( + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" +) + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +// trackingBody 记录 Close 是否被调用，用于验证响应体被正确关闭。 +type trackingBody struct { + r io.Reader + closed bool +} + +func (b *trackingBody) Read(p []byte) (int, error) { return b.r.Read(p) } +func (b *trackingBody) Close() error { + b.closed = true + return nil +} + +func withFakeTransport(t *testing.T, resp *http.Response, body *trackingBody) { + t.Helper() + old := httpClient + httpClient = &http.Client{ + Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) { + resp.Request = r + return resp, nil + }), + // 阻止自动跟随 302，以便直接观察到 proxy 对 Location 的处理。 + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + }, + } + t.Cleanup(func() { httpClient = old }) +} + +func newTestContext(method, target string) (*gin.Context, *httptest.ResponseRecorder) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(method, target, nil) + return c, w +} + +// 9. 代理应正确转发响应内容，并关闭上游响应体。 +func TestProxyForwardsAndClosesBody(t *testing.T) { + gin.SetMode(gin.TestMode) + const payload = "hello ghproxy streaming body" + body := &trackingBody{r: strings.NewReader(payload)} + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Length": []string{strconv.Itoa(len(payload))}}, + Body: body, + } + withFakeTransport(t, resp, body) + + c, w := newTestContext(http.MethodGet, "/https://raw.githubusercontent.com/a/b/main/x.txt") + proxy(c, "https://raw.githubusercontent.com/a/b/main/x.txt") + + if w.Code != http.StatusOK { + t.Fatalf("状态码应为 200，实际 %d", w.Code) + } + if w.Body.String() != payload { + t.Fatalf("转发内容错误: %q", w.Body.String()) + } + if !body.closed { + t.Fatal("上游响应体未被关闭") + } +} + +// 超过文件大小限制时应返回 413 并关闭响应体（不读取整个文件）。 +func TestProxyRejectsOversize(t *testing.T) { + gin.SetMode(gin.TestMode) + oldCfg := config.Load() + t.Cleanup(func() { config.Store(oldCfg) }) + config.Store(&Config{RequestLimit: RequestLimitConfig{LimitSize: 1}}) // 1 MB + + body := &trackingBody{r: strings.NewReader("tiny")} + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Length": []string{strconv.Itoa(2 * 1024 * 1024)}}, + Body: body, + } + withFakeTransport(t, resp, body) + + c, w := newTestContext(http.MethodGet, "/https://raw.githubusercontent.com/a/b/main/big.bin") + proxy(c, "https://raw.githubusercontent.com/a/b/main/big.bin") + + if w.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("状态码应为 413，实际 %d", w.Code) + } + if !body.closed { + t.Fatal("上游响应体未被关闭") + } +} + +// 入口链接的 302 应重写为代理地址，而不是内部跟随。 +func TestProxyRewritesEntryLocation(t *testing.T) { + gin.SetMode(gin.TestMode) + body := &trackingBody{r: strings.NewReader("")} + resp := &http.Response{ + StatusCode: http.StatusFound, + Header: http.Header{"Location": []string{"https://github.com/a/b/releases/download/v1/x.zip"}}, + Body: body, + } + withFakeTransport(t, resp, body) + + c, w := newTestContext(http.MethodGet, "/https://github.com/a/b/releases/download/v1/x.zip") + proxy(c, "https://github.com/a/b/releases/download/v1/x.zip") + + loc := w.Header().Get("Location") + if loc != "/https://github.com/a/b/releases/download/v1/x.zip" { + t.Fatalf("Location 重写错误: %q", loc) + } +} + +// 达到限流上限后，handler 直接返回 429，不进入代理流程。 +func TestHandlerReturns429WhenLimited(t *testing.T) { + gin.SetMode(gin.TestMode) + oldLimiter := limiter + oldCfg := config.Load() + t.Cleanup(func() { + limiter = oldLimiter + config.Store(oldCfg) + }) + + limiter = newRateLimiter() + config.Store(&Config{RequestLimit: RequestLimitConfig{LimitRate: 1}}) + + // 预先占满该 IP 的唯一额度。 + if !limiter.allow("1.2.3.4", 1, time.Now()) { + t.Fatal("预填充失败") + } + + c, w := newTestContext(http.MethodGet, "/https://raw.githubusercontent.com/a/b/main/x.txt") + c.Request.RemoteAddr = "1.2.3.4:12345" + handler(c) + + if w.Code != http.StatusTooManyRequests { + t.Fatalf("应返回 429，实际 %d", w.Code) + } +} diff --git a/ratelimit_test.go b/ratelimit_test.go new file mode 100644 index 0000000..f45c326 --- /dev/null +++ b/ratelimit_test.go @@ -0,0 +1,166 @@ +package main + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +// newTestLimiter 关闭自动回收，便于测试精确控制时间。 +func newTestLimiter() *rateLimiter { + l := newRateLimiter() + l.sweepEvery = time.Hour + return l +} + +// 1. 关闭限流（limit <= 0）时应始终放行且不记录任何状态。 +func TestLimiterDisabled(t *testing.T) { + l := newTestLimiter() + now := time.Now() + for i := 0; i < 1000; i++ { + if !l.allow("1.1.1.1", 0, now) { + t.Fatalf("关闭限流时应始终放行，第 %d 次被拒绝", i+1) + } + if !l.allow("2.2.2.2", -5, now) { + t.Fatalf("limitRate 为负时应始终放行") + } + } + if len(l.entries) != 0 { + t.Fatalf("关闭限流时不应记录任何 IP，实际记录 %d 个", len(l.entries)) + } +} + +// 2. 请求数恰好达到限额时应全部放行。 +func TestLimiterExactlyAtLimit(t *testing.T) { + l := newTestLimiter() + now := time.Now() + const limit = 3 + for i := 0; i < limit; i++ { + if !l.allow("1.1.1.1", limit, now) { + t.Fatalf("第 %d 次请求应放行（恰好达到限额）", i+1) + } + } +} + +// 3. 超过限额时应被拒绝。 +func TestLimiterOverLimitRejected(t *testing.T) { + l := newTestLimiter() + now := time.Now() + const limit = 2 + if !l.allow("1.1.1.1", limit, now) || !l.allow("1.1.1.1", limit, now) { + t.Fatal("前两次请求应放行") + } + if l.allow("1.1.1.1", limit, now) { + t.Fatal("第 3 次请求应被拒绝") + } + if l.allow("1.1.1.1", limit, now) { + t.Fatal("超限后应持续拒绝") + } +} + +// 4. 时间窗口过期后应恢复放行。 +func TestLimiterWindowExpiry(t *testing.T) { + l := newTestLimiter() + t0 := time.Now() + const limit = 2 + l.allow("1.1.1.1", limit, t0) + l.allow("1.1.1.1", limit, t0) + if l.allow("1.1.1.1", limit, t0.Add(30*time.Second)) { + t.Fatal("窗口内应继续拒绝") + } + // 超过一个窗口后，旧的记录过期，应重新放行。 + if !l.allow("1.1.1.1", limit, t0.Add(time.Minute+time.Second)) { + t.Fatal("窗口过期后应恢复放行") + } +} + +// 5. 多个 IP 之间互不影响。 +func TestLimiterMultipleIPsIndependent(t *testing.T) { + l := newTestLimiter() + now := time.Now() + const limit = 1 + if !l.allow("a", limit, now) { + t.Fatal("a 第一次应放行") + } + if l.allow("a", limit, now) { + t.Fatal("a 第二次应拒绝") + } + if !l.allow("b", limit, now) { + t.Fatal("b 不应受 a 的限制影响") + } + if l.allow("b", limit, now) { + t.Fatal("b 第二次应拒绝") + } +} + +// 6. 同一 IP 的并发请求：放行总数必须恰好等于限额。 +func TestLimiterConcurrentSameIP(t *testing.T) { + l := newTestLimiter() + now := time.Now() + const limit = 50 + const goroutines = 500 + + var allowed int64 + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if l.allow("same", limit, now) { + atomic.AddInt64(&allowed, 1) + } + }() + } + wg.Wait() + + if allowed != limit { + t.Fatalf("并发放行数应为 %d，实际为 %d", limit, allowed) + } +} + +// 7. 大量不同 IP 请求后，过期条目应被回收，避免无界增长。 +func TestLimiterReclaimsInactiveIPs(t *testing.T) { + l := newRateLimiter() // 使用默认 window / sweepEvery + base := time.Now() + const n = 5000 + for i := 0; i < n; i++ { + ip := fmt.Sprintf("10.%d.%d.%d", (i>>16)&0xff, (i>>8)&0xff, i&0xff) + l.allow(ip, 10, base) + } + if got := len(l.entries); got != n { + t.Fatalf("应记录 %d 个 IP，实际 %d", n, got) + } + + // 一个窗口之后发起的请求会触发回收，清掉全部不活跃 IP。 + after := base.Add(2 * time.Minute) + if !l.allow("192.168.0.1", 10, after) { + t.Fatal("新 IP 应放行") + } + if got := len(l.entries); got != 1 { + t.Fatalf("过期 IP 应被回收，仅剩新 IP（1 个），实际 %d", got) + } +} + +// 活跃 IP 的过期记录应在回收时被裁剪，但保留窗口内记录。 +func TestLimiterSweepPrunesActiveIP(t *testing.T) { + l := newTestLimiter() + base := time.Now() + // active 有一条旧记录和一条窗口内新记录。 + l.allow("active", 10, base) + l.allow("active", 10, base.Add(80*time.Second)) + + // 强制在 base+90s 触发一次回收：cutoff = base+30s， + // active 的旧记录被裁剪，窗口内的新记录保留。 + l.sweepEvery = 0 + l.allow("other", 10, base.Add(90*time.Second)) + + got := l.entries["active"] + if len(got) != 1 { + t.Fatalf("active 应保留 1 条窗口内记录，实际 %d 条", len(got)) + } + if !got[0].Equal(base.Add(80 * time.Second)) { + t.Fatalf("保留下来的应是窗口内的新记录") + } +}