diff --git a/bench_test.go b/bench_test.go new file mode 100644 index 0000000..77b19be --- /dev/null +++ b/bench_test.go @@ -0,0 +1,416 @@ +package main + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +// =========================================================================== +// 基准口径说明(改代码时请同步本段,并由 TestBenchmarkPrimeStates 守护) +// +// 【为什么不能宣称「两边记录数相同」】 +// 两个实现的「饱和态」(此后每次调用都被拒绝的状态)记录条数天然不同: +// 旧:拒绝条件 len(entries[ip]) > limit → 饱和记录数 = limit+1 +// 新:拒绝条件 len(times) >= limit → 饱和记录数 = limit +// 在 limit=100 时即 101 条 vs 100 条。这 1 条之差**正是本 PR 要修正的边界偏差** +// (旧实现窗口内会多放行一次),不是可以「统一」掉的实现细节: +// - 新实现**无法**经公开入口(allow / allowAt)达到 limit+1 条记录:第 limit+1 次 +// 调用已被拒绝,不会追加。若要用直接改 map 的方式硬凑 limit+1,那是运行期 +// 不可达的状态,不能代表新实现。 +// - 旧实现若只预填 limit 条,则第 1 次调用会先放行并追加,计时区间内混入一次 +// 与其余调用路径不同的样本。 +// (此前版本的本文件曾把新实现写成「预填 limit+1 次 ⇒ 保留 limit+1 条」,并把 +// 两边记录数写成「一致(101)」——经实测为 100 条,注释与代码不符,已更正。) +// +// 【本文件采用的做法】 +// 不做「记录数相同」的伪统一,而是让**每一侧各自准备在自己的饱和拒绝态**: +// 计时区间内每次调用都是同一条「已饱和 → 拒绝」路径,不含放行与追加; +// 两侧各自的记录条数由 primeNewRejecting / primeOldRejecting 返回, +// 并由 TestBenchmarkPrimeStates 断言,避免注释再次与代码脱节。 +// +// 【计时口径】 +// - BenchmarkLimiterNew_* / BenchmarkLimiterOld_*:注入固定时间 → 纯算法成本,两者可互比 +// - BenchmarkLimiterNewRealClock_*:走生产入口 allow()(锁内采样 time.Now)→ 含真实 +// 时钟成本,只能与 NewRealClock 系列比,不可与前两者混比 +// +// 【这不是放行路径的对比】 +// 同一记录序列下新旧放行次数不同(limit+1 vs limit),不存在可直接对比的同一 +// 放行状态;如需放行路径数据,必须单独说明状态与记录数,不能与本文件的数字混用。 +// =========================================================================== + +const ( + // benchLimit 是基准统一使用的限流额度(生产 config.json 为 500;此处取 100 以 + // 控制单次运行时长,两侧一致即可比)。 + benchLimit int64 = 100 + // benchIPsMany 是「多 IP」场景的 IP 数量。 + benchIPsMany = 1000 +) + +// workerStride 是并行基准中相邻 worker 的起始 IP 偏移步长。 +// 每个 worker 都从 i := 0 开始会让它们初始全部争抢同一个 IP,使多 IP 场景实际 +// 退化成单 IP;给每个 worker 一个互不相同的起始偏移后,访问分布才真正覆盖多个 IP。 +// 取一个与常见 IP 数量互质的较大质数,使起始位置在 IP 空间上分散。 +const workerStride = 1000003 + +// parallelStart 返回当前 worker 的起始索引,仅在每个 worker 启动时调用一次, +// 不进入 pb.Next() 循环,因此不引入每次请求的额外开销。 +func parallelStart(seq *atomic.Uint64, ipCount int) int { + return int(seq.Add(1)-1) * workerStride % ipCount +} + +// 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 +} + +// expectedNewRejectingRecords 返回新实现饱和拒绝态的每 IP 记录条数。 +func expectedNewRejectingRecords(limit int64) int { return int(limit) } + +// expectedOldRejectingRecords 返回旧实现饱和拒绝态的每 IP 记录条数。 +func expectedOldRejectingRecords(limit int64) int { return int(limit + 1) } + +// primeNewRejecting 把新实现预置到饱和拒绝态:每个 IP 经公开入口放行 limit 次 +// (第 limit 次之后 len == limit,后续调用走拒绝路径)。返回实测记录条数。 +// 预置期间 l.sweepEvery 应已设为不会触发回收的值。 +func primeNewRejecting(l *rateLimiter, ips []string, limit int64, now time.Time) int { + for _, ip := range ips { + for j := int64(0); j < limit; j++ { + if !l.allowAt(ip, limit, now) { + panic("基准预置失败:limit 次以内不应被拒绝") + } + } + } + return len(l.entries[ips[0]]) +} + +// primeNewRejectingRealClock 同 primeNewRejecting,但走生产入口 allow()(真实时钟)。 +func primeNewRejectingRealClock(l *rateLimiter, ips []string, limit int64) int { + for _, ip := range ips { + for j := int64(0); j < limit; j++ { + if !l.allow(ip, limit) { + panic("基准预置失败:limit 次以内不应被拒绝") + } + } + } + return len(l.entries[ips[0]]) +} + +// primeOldRejecting 把旧实现预置到饱和拒绝态:每个 IP 直接写入 limit+1 条记录 +// (旧实现拒绝条件为 len > limit,故饱和态为 limit+1 条)。返回实测记录条数。 +func primeOldRejecting(entries map[string][]time.Time, ips []string, limit int64, now time.Time) int { + for _, ip := range ips { + for j := int64(0); j < limit+1; j++ { + entries[ip] = append(entries[ip], now) + } + } + return len(entries[ips[0]]) +} + +func benchNew(b *testing.B, ipCount int, limit int64) { + l := newRateLimiter() + l.sweepEvery = time.Hour // 排除回收开销,只衡量单请求处理成本 + ips := makeIPs(ipCount) + now := time.Now() + + b.ReportAllocs() + b.StopTimer() + if got, want := primeNewRejecting(l, ips, limit, now), expectedNewRejectingRecords(limit); got != want { + b.Fatalf("新实现饱和态记录数 = %d,期望 %d", got, want) + } + // 注意:ResetTimer 只归零累计耗时,**不会启动**计时器;StopTimer 之后必须再 + // 显式 StartTimer,否则计时区间恒为 0,框架会把 b.N 一路放大到上限后才停。 + b.ResetTimer() + b.StartTimer() + + var seq atomic.Uint64 + b.RunParallel(func(pb *testing.PB) { + i := parallelStart(&seq, ipCount) + for pb.Next() { + l.allowAt(ips[i%ipCount], limit, now) + i++ + } + }) +} + +// benchNewRealClock 衡量生产入口(锁内采样 time.Now)的单请求成本。 +// 时钟口径与 New/Old 两组不同,结果只能在这两者之间比较。 +// 预置与循环都用真实时钟:只要单次 benchmark 的循环耗时远小于 window(1 分钟), +// 记录不会被裁剪,循环内始终保持「已饱和 → 拒绝」。 +func benchNewRealClock(b *testing.B, ipCount int, limit int64) { + l := newRateLimiter() + l.sweepEvery = time.Hour + ips := makeIPs(ipCount) + + b.ReportAllocs() + b.StopTimer() + if got, want := primeNewRejectingRealClock(l, ips, limit), expectedNewRejectingRecords(limit); got != want { + b.Fatalf("新实现(真实时钟)饱和态记录数 = %d,期望 %d", got, want) + } + b.ResetTimer() + b.StartTimer() + + var seq atomic.Uint64 + b.RunParallel(func(pb *testing.PB) { + i := parallelStart(&seq, ipCount) + for pb.Next() { + l.allow(ips[i%ipCount], limit) + 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() + + b.ReportAllocs() + b.StopTimer() + if got, want := primeOldRejecting(entries, ips, limit, now), expectedOldRejectingRecords(limit); got != want { + b.Fatalf("旧实现饱和态记录数 = %d,期望 %d", got, want) + } + b.ResetTimer() + b.StartTimer() + + var seq atomic.Uint64 + b.RunParallel(func(pb *testing.PB) { + i := parallelStart(&seq, ipCount) + for pb.Next() { + allowOld(entries, &mu, ips[i%ipCount], limit, now) + i++ + } + }) +} + +func BenchmarkLimiterNew_SingleIP(b *testing.B) { benchNew(b, 1, benchLimit) } +func BenchmarkLimiterNew_ManyIPs(b *testing.B) { benchNew(b, benchIPsMany, benchLimit) } +func BenchmarkLimiterOld_SingleIP(b *testing.B) { benchOld(b, 1, benchLimit) } +func BenchmarkLimiterOld_ManyIPs(b *testing.B) { benchOld(b, benchIPsMany, benchLimit) } + +func BenchmarkLimiterNewRealClock_SingleIP(b *testing.B) { benchNewRealClock(b, 1, benchLimit) } +func BenchmarkLimiterNewRealClock_ManyIPs(b *testing.B) { + benchNewRealClock(b, benchIPsMany, benchLimit) +} + +// sweepRecordsPerIP 是回收基准中每个 IP 的记录条数:经公开入口写入 10 条 +// (sweepLimit 恰好为 10,故 10 次调用全部放行、全部被记录)。 +const ( + sweepLimit int64 = 10 + sweepRecordsPerIP = 10 + sweepIPCount = 5000 +) + +// fixSweepLimiter 构造一个「已填充 n 个 IP、每个 IP sweepRecordsPerIP 条记录、 +// 且下一次 allow 必然触发回收的限流器。 +// 填充走限流器正常入口(allowAt),每个 IP 调用 sweepLimit 次,因此每个 IP 恰好 +// 记录 10 条(t == base),并在 [base+2min] 触发回收时全部过期(window=1min, +// cutoff = base+1min,10 条均早于 cutoff)。 +// +// 返回值中 l.sweepEvery 已置 0,只要传入的触发时刻**不早于** l.lastSweep +// (= 本函数内部 newRateLimiter 时的 time.Now()),下一次 allow 就会触发 sweepLocked。 +// 因此调用方必须把触发时刻算成本轮填充的时刻(base)再加偏移,不能拿到循环外 +// 复用同一个绝对时刻:那会随着真实时间前进使触发时刻反而早于 l.lastSweep, +// 此时 now.Sub(lastSweep) 为负、sweepEvery=0 也满足不了条件,回收会**静默不触发**, +// 基准就变成了在测一条空的拒绝路径。 +// +// 填充成本是否计入计时区间**由调用方决定**: +// - BenchmarkLimiterSweep 在 StopTimer 状态下调用以排除填充成本; +// - BenchmarkLimiterSweepWithFill 则在计时区间内调用,有意把填充计入。 +func fixSweepLimiter(n int, base time.Time) *rateLimiter { + ips := makeIPs(n) + l := newRateLimiter() + l.sweepEvery = time.Hour // 填充期间不触发回收 + for _, ip := range ips { + for j := 0; j < sweepRecordsPerIP; j++ { + if !l.allowAt(ip, sweepLimit, base) { + panic("回收基准填充失败:sweepLimit 次以内不应被拒绝") + } + } + } + // 防注释与代码脱节:每个 IP 必须恰好 sweepRecordsPerIP 条记录。 + for _, ip := range ips { + if got := len(l.entries[ip]); got != sweepRecordsPerIP { + panic(fmt.Sprintf("回收基准填充记录数 = %d,期望 %d", got, sweepRecordsPerIP)) + } + } + l.sweepEvery = 0 // 下一次 allow 立即触发 sweepLocked + return l +} + +// BenchmarkLimiterSweep 只衡量**一次回收操作**本身。 +// +// 工作负载:sweepIPCount(5000) 个 IP × sweepRecordsPerIP(10) 条记录 = 50,000 条, +// 全部已在窗口外过期 → 一次 sweepLocked 应把 5000 个 IP 键全部删除。 +// +// **计时范围**:构造限流器与填充 5000 个 IP 属于初始化成本,通过 +// StopTimer/StartTimer 排除在计时区间之外;报告中的 ns/op、B/op、allocs/op +// 只反映 sweepLocked 及其触发它的那一次 allow 调用本身,可解释为「单次 sweep +// 的成本」,**不含**填充成本。 +// +// **触发时刻每轮重算**:base 与 next 都在循环内、相对于本轮填充时刻计算。 +// 若把 base 提到循环外,长时间运行后 next 会早于新一轮限流器的 lastSweep +// (= newRateLimiter 时的真实 now),now.Sub(lastSweep) 为负 → 回收不触发。 +func BenchmarkLimiterSweep(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + b.StopTimer() + base := time.Now() + l := fixSweepLimiter(sweepIPCount, base) + next := base.Add(2 * time.Minute) // 相对本轮填充时刻,保证晚于 l.lastSweep + b.StartTimer() + + l.allowAt("192.168.0.1", sweepLimit, next) // 触发 sweepLocked + } +} + +// BenchmarkLimiterSweepWithFill 衡量**完整周期**:新建限流器 + 填充 5000 个 IP +// (每个 10 条记录)+ 触发一次回收。 +// +// **计时范围**:与 BenchmarkLimiterSweep 相反,它**有意把填充成本计入**被测区间, +// 因此其结果**不能**当作纯回收耗时解释;保留它是为了观察「填充 + 回收」整体开销 +// 随规模的变化。 +// +// 触发时刻同样每轮重算(见 BenchmarkLimiterSweep 的说明),避免长跑后回收不触发。 +func BenchmarkLimiterSweepWithFill(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + base := time.Now() + l := fixSweepLimiter(sweepIPCount, base) + next := base.Add(2 * time.Minute) + l.allowAt("192.168.0.1", sweepLimit, next) + } +} + +// TestBenchmarkPrimeStates 断言基准预置状态与文件头部注释所述一致: +// 新实现饱和态 = limit 条,旧实现饱和态 = limit+1 条,且饱和后调用被拒绝、 +// 记录数不再增长。两者相差 1 是边界语义差异的必然结果,不可宣称「两边相同」。 +func TestBenchmarkPrimeStates(t *testing.T) { + const limit = int64(10) + now := time.Now() + ips := []string{"10.0.0.1"} + + l := newTestLimiter() + if got, want := primeNewRejecting(l, ips, limit, now), expectedNewRejectingRecords(limit); got != want { + t.Fatalf("新实现饱和态记录数 = %d,期望 %d", got, want) + } + if l.allowAt(ips[0], limit, now) { + t.Fatal("新实现饱和态应拒绝") + } + if got, want := len(l.entries[ips[0]]), expectedNewRejectingRecords(limit); got != want { + t.Fatalf("被拒绝后新实现记录数 = %d,期望保持不变 %d", got, want) + } + + entries := map[string][]time.Time{} + var mu sync.Mutex + if got, want := primeOldRejecting(entries, ips, limit, now), expectedOldRejectingRecords(limit); got != want { + t.Fatalf("旧实现饱和态记录数 = %d,期望 %d", got, want) + } + if allowOld(entries, &mu, ips[0], limit, now) { + t.Fatal("旧实现饱和态应拒绝") + } + if got, want := len(entries[ips[0]]), expectedOldRejectingRecords(limit); got != want { + t.Fatalf("被拒绝后旧实现记录数 = %d,期望保持不变 %d", got, want) + } + + t.Logf("饱和态记录数:新实现 %d / 旧实现 %d(limit=%d,相差 1 即本 PR 修正的边界偏差)", + expectedNewRejectingRecords(limit), expectedOldRejectingRecords(limit), limit) +} + +// TestSweepFixtureRecordsAllExpired 断言回收基准的填充夹具与注释一致: +// 每个 IP 恰好 sweepRecordsPerIP 条记录,且在触发回收的时刻全部已过期 +// (这才是「一次 sweep 删除 5000 个 IP 键」的前提)。 +func TestSweepFixtureRecordsAllExpired(t *testing.T) { + base := time.Now() + l := fixSweepLimiter(64, base) // 用小规模验证夹具本身 + + if got := len(l.entries); got != 64 { + t.Fatalf("填充后 IP 数 = %d,期望 64", got) + } + for ip, times := range l.entries { + if len(times) != sweepRecordsPerIP { + t.Fatalf("IP %s 记录数 = %d,期望 %d", ip, len(times), sweepRecordsPerIP) + } + } + + next := base.Add(2 * time.Minute) + cutoff := next.Add(-l.window) + for ip, times := range l.entries { + for i, ts := range times { + if !ts.Before(cutoff) { + t.Fatalf("IP %s 第 %d 条记录 %v 未早于 cutoff %v,回收时不会过期", ip, i, ts, cutoff) + } + } + } + + if !l.allowAt("192.168.0.1", sweepLimit, next) { + t.Fatal("新 IP 应放行") + } + if got := len(l.entries); got != 1 { + t.Fatalf("回收后应只剩新 IP(1 个),实际 %d", got) + } + t.Logf("回收夹具:64 个 IP × %d 条记录全部过期,一次 sweep 后仅剩新 IP", sweepRecordsPerIP) +} + +// TestSweepTriggersWithStaleBase 守护触发时刻的算法:base 取一个**远早于现在**的 +// 时刻(模拟基准长跑很久之后,循环外的固定 base 已经过期),触发时刻相对 base +// 计算就会早于新一轮限流器的 lastSweep(= newRateLimiter 时的真实 now), +// now.Sub(lastSweep) 为负 → 即便 sweepEvery=0 也**不会**触发回收。 +// 基准必须按此每轮重算触发时刻;本测试把该陷阱固定下来。 +func TestSweepTriggersWithStaleBase(t *testing.T) { + staleBase := time.Now().Add(-10 * time.Minute) + l := fixSweepLimiter(64, staleBase) + + staleNext := staleBase.Add(2 * time.Minute) // 已早于真实 now + if !staleNext.Before(l.lastSweep) { + t.Fatal("前置条件不成立:staleNext 应早于 l.lastSweep") + } + + if !l.allowAt("192.168.0.1", sweepLimit, staleNext) { + t.Fatal("新 IP 应放行") + } + if got := len(l.entries); got != 65 { + t.Fatalf("陈旧触发时刻下回收不应发生(应保留 64 个旧 IP + 1 个新 IP = 65),实际 %d", got) + } + + // 相对本轮填充时刻重算后,回收必须发生。 + l2 := fixSweepLimiter(64, staleBase) + freshNext := time.Now().Add(2 * time.Minute) + if !l2.allowAt("192.168.0.1", sweepLimit, freshNext) { + t.Fatal("新 IP 应放行") + } + if got := len(l2.entries); got != 1 { + t.Fatalf("相对当前时刻的触发应完成回收(仅剩 1 个新 IP),实际 %d", got) + } + t.Log("已确认:触发时刻必须每轮相对 newRateLimiter 的真实时刻重算,否则回收静默不触发") +} diff --git a/config_test.go b/config_test.go new file mode 100644 index 0000000..c2368b0 --- /dev/null +++ b/config_test.go @@ -0,0 +1,108 @@ +package main + +import ( + "fmt" + "io" + "os" + "path/filepath" + "sync" + "testing" +) + +// 配置重载与请求读取并发时不应发生数据竞争,且读取方始终看到完整一致的快照。 +func TestConfigReloadConcurrentWithRead(t *testing.T) { + // 保存并恢复所有被修改的全局状态,避免污染其他测试。 + origConfig := config.Load() + oldOut := log.Out + oldLevel := log.GetLevel() + log.SetOutput(io.Discard) + t.Cleanup(func() { + config.Store(origConfig) + log.SetOutput(oldOut) + log.SetLevel(oldLevel) + }) + + 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..52e772e 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,96 @@ 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 + // now 可注入,默认 time.Now;测试用 allowAt 注入确定性时间。 + now func() time.Time +} + +func newRateLimiter() *rateLimiter { + return &rateLimiter{ + entries: make(map[string][]time.Time), + lastSweep: time.Now(), + window: time.Minute, + sweepEvery: time.Minute, + now: time.Now, + } +} + +// allow 是生产入口:先取得锁、再采样当前时间,保证时间戳的采样顺序与写入顺序 +// 一致,从而维持 entries[ip] 按时间升序的不变量(否则乱序会导致过期记录裁剪 +// 与 sweep 回收判断出错)。 +func (l *rateLimiter) allow(ip string, limit int64) bool { + l.mu.Lock() + defer l.mu.Unlock() + return l.allowLocked(ip, limit, l.now()) +} + +// allowAt 是测试辅助入口,允许注入时间;与 allow 复用同一套加锁及限流逻辑。 +func (l *rateLimiter) allowAt(ip string, limit int64, now time.Time) bool { + l.mu.Lock() + defer l.mu.Unlock() + return l.allowLocked(ip, limit, now) +} + +// allowLocked 执行限流判断与记录,必须在持有 l.mu 时调用。 +// limit <= 0 表示关闭限流:直接放行且不记录、不回收。 +// window 内允许恰好 limit 次请求,第 limit+1 次起返回 false。 +func (l *rateLimiter) allowLocked(ip string, limit int64, now time.Time) bool { + if limit <= 0 { + return true + } + + // 低频回收:最多每 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 +253,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 +275,9 @@ func main() { } func handler(c *gin.Context) { + // 取一次不可变配置快照,保证整个请求内看到一致的配置,且无锁、无数据竞争。 + cfg := config.Load() + rawPath := strings.TrimPrefix(c.Request.URL.RequestURI(), "/") for strings.HasPrefix(rawPath, "/") { @@ -198,16 +290,16 @@ func handler(c *gin.Context) { matches := checkURL(rawPath) if matches != nil { - if len(config.WhiteList) > 0 && !checkList(matches, config.WhiteList) { - render404(c, tmpl404) + if len(cfg.WhiteList) > 0 && !checkList(matches, cfg.WhiteList) { + render404(c, tmpl404, cfg) return } - if len(config.BlackList) > 0 && checkList(matches, config.BlackList) { - render404(c, tmpl404) + if len(cfg.BlackList) > 0 && checkList(matches, cfg.BlackList) { + render404(c, tmpl404, cfg) return } } else { - render404(c, tmpl404) + render404(c, tmpl404, cfg) return } @@ -215,36 +307,18 @@ func handler(c *gin.Context) { rawPath = strings.Replace(rawPath, "/blob/", "/raw/", 1) } - // 新增:基于时间的速率限制 + // 基于时间的速率限制:只处理当前 IP,窗口内允许恰好 limitRate 次。 + // 时间在 limiter.allow 内部持锁后采样,避免并发下时间戳乱序。 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) { + 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 +329,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 { @@ -267,10 +339,10 @@ func handler(c *gin.Context) { } } - proxy(c, rawPath) + proxy(c, rawPath, cfg) } -func proxy(c *gin.Context, u string) { +func proxy(c *gin.Context, u string, cfg *Config) { req, err := http.NewRequest(c.Request.Method, u, c.Request.Body) if err != nil { c.String(http.StatusInternalServerError, fmt.Sprintf("server error %v", err)) @@ -303,9 +375,7 @@ func proxy(c *gin.Context, u string) { return } - configLock.RLock() - limitSize := config.RequestLimit.LimitSize * 1024 * 1024 // Convert MB to Bytes - configLock.RUnlock() + limitSize := cfg.RequestLimit.LimitSize * 1024 * 1024 // Convert MB to Bytes // 如果 limitSize <= 0 则不限制文件大小 if limitSize > 0 && size > limitSize { @@ -320,7 +390,7 @@ func proxy(c *gin.Context, u string) { // 如果 GitHub 返回 404,显示我们的 404 页面 if resp.StatusCode == http.StatusNotFound { - render404(c, tmpl404) + render404(c, tmpl404, cfg) return } @@ -336,7 +406,7 @@ func proxy(c *gin.Context, u string) { if isEntryURL(location) { c.Header("Location", "/"+location) } else { - proxy(c, location) + proxy(c, location, cfg) return } } @@ -348,11 +418,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 +440,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 { @@ -407,11 +482,9 @@ func checkList(matches, list []string) bool { return false } -func render404(c *gin.Context, tmpl *template.Template) { - configLock.RLock() - analyticsURL := config.AnalyticsURL - supportImageURL := config.SupportImageURL - configLock.RUnlock() +func render404(c *gin.Context, tmpl *template.Template, cfg *Config) { + 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..d4b16f8 --- /dev/null +++ b/proxy_test.go @@ -0,0 +1,159 @@ +package main + +import ( + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + + "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 +} + +// setTestGinMode 设置测试模式并在测试结束后恢复原模式,避免污染全局状态。 +func setTestGinMode(t *testing.T) { + t.Helper() + old := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(old) }) +} + +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) { + setTestGinMode(t) + 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", config.Load()) + + 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) { + setTestGinMode(t) + oldCfg := config.Load() + t.Cleanup(func() { config.Store(oldCfg) }) + cfg := &Config{RequestLimit: RequestLimitConfig{LimitSize: 1}} // 1 MB + config.Store(cfg) + + 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", cfg) + + if w.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("状态码应为 413,实际 %d", w.Code) + } + if !body.closed { + t.Fatal("上游响应体未被关闭") + } +} + +// 入口链接的 302 应重写为代理地址,而不是内部跟随。 +func TestProxyRewritesEntryLocation(t *testing.T) { + setTestGinMode(t) + 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", config.Load()) + + 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) { + setTestGinMode(t) + 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) { + 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..033db98 --- /dev/null +++ b/ratelimit_test.go @@ -0,0 +1,214 @@ +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.allowAt("1.1.1.1", 0, now) { + t.Fatalf("关闭限流时应始终放行,第 %d 次被拒绝", i+1) + } + if !l.allowAt("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.allowAt("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.allowAt("1.1.1.1", limit, now) || !l.allowAt("1.1.1.1", limit, now) { + t.Fatal("前两次请求应放行") + } + if l.allowAt("1.1.1.1", limit, now) { + t.Fatal("第 3 次请求应被拒绝") + } + if l.allowAt("1.1.1.1", limit, now) { + t.Fatal("超限后应持续拒绝") + } +} + +// 4. 时间窗口过期后应恢复放行,并验证窗口边界不出现 off-by-one。 +func TestLimiterWindowExpiry(t *testing.T) { + l := newTestLimiter() + t0 := time.Now() + const limit = 2 + l.allowAt("1.1.1.1", limit, t0) + l.allowAt("1.1.1.1", limit, t0) + if l.allowAt("1.1.1.1", limit, t0.Add(30*time.Second)) { + t.Fatal("窗口内应继续拒绝") + } + // 恰好在窗口边界(age == window)时记录仍有效,应拒绝。 + if l.allowAt("1.1.1.1", limit, t0.Add(time.Minute)) { + t.Fatal("窗口边界上的记录仍应计入,必须拒绝") + } + // 刚超过一个窗口后,旧记录过期,应恢复放行。 + if !l.allowAt("1.1.1.1", limit, t0.Add(time.Minute+time.Nanosecond)) { + t.Fatal("窗口过期后应恢复放行") + } +} + +// 5. 多个 IP 之间互不影响。 +func TestLimiterMultipleIPsIndependent(t *testing.T) { + l := newTestLimiter() + now := time.Now() + const limit = 1 + if !l.allowAt("a", limit, now) { + t.Fatal("a 第一次应放行") + } + if l.allowAt("a", limit, now) { + t.Fatal("a 第二次应拒绝") + } + if !l.allowAt("b", limit, now) { + t.Fatal("b 不应受 a 的限制影响") + } + if l.allowAt("b", limit, now) { + t.Fatal("b 第二次应拒绝") + } +} + +// 6a. 同一 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.allowAt("same", limit, now) { + atomic.AddInt64(&allowed, 1) + } + }() + } + wg.Wait() + + if allowed != limit { + t.Fatalf("并发放行数应为 %d,实际为 %d", limit, allowed) + } +} + +// 6b. 使用真实生产入口(锁内采样 time.Now)并发请求: +// 放行数恰当,且 entries[ip] 必须保持按时间升序,验证时间戳顺序修复。 +func TestLimiterConcurrentRealClockOrdering(t *testing.T) { + l := newTestLimiter() + const limit = 100 + const goroutines = 2000 + + 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) { + atomic.AddInt64(&allowed, 1) + } + }() + } + wg.Wait() + + if allowed != limit { + t.Fatalf("并发放行数应为 %d,实际为 %d", limit, allowed) + } + times := l.entries["same"] + if len(times) != limit { + t.Fatalf("窗口内记录数应为 %d,实际 %d", limit, len(times)) + } + for i := 1; i < len(times); i++ { + if times[i].Before(times[i-1]) { + t.Fatalf("时间戳乱序:索引 %d (%v) 早于 %d (%v)", i, times[i], i-1, times[i-1]) + } + } +} + +// 7. 大量不同 IP 请求后,过期条目应被回收,避免无界增长。 +func TestLimiterReclaimsInactiveIPs(t *testing.T) { + l := newTestLimiter() + 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.allowAt(ip, 10, base) + } + if got := len(l.entries); got != n { + t.Fatalf("应记录 %d 个 IP,实际 %d", n, got) + } + + // 强制在窗口之后触发一次回收,清掉全部不活跃 IP。 + l.sweepEvery = 0 + after := base.Add(2 * time.Minute) + if !l.allowAt("192.168.0.1", 10, after) { + t.Fatal("新 IP 应放行") + } + if got := len(l.entries); got != 1 { + t.Fatalf("过期 IP 应被回收,仅剩新 IP(1 个),实际 %d", got) + } +} + +// 8. 直接验证 sweep 的目标代码路径:裁剪活跃 IP 的过期记录、回收不活跃 IP, +// 且回收不破坏仍然有效的窗口内记录。 +func TestLimiterSweepPrunesActiveAndReclaimsInactive(t *testing.T) { + l := newTestLimiter() + base := time.Now() + + // 直接构造 entries,使旧记录在 sweep 被调用时确实仍然存在, + // 不经过 allowAt 触发的 prune,从而真正覆盖 sweepLocked 的裁剪逻辑。 + l.entries["active"] = []time.Time{base, base.Add(80 * time.Second)} + l.entries["stale"] = []time.Time{base} + + now := base.Add(90 * time.Second) // cutoff = now - window = base + 30s + l.mu.Lock() + l.sweepLocked(now) + l.mu.Unlock() + + // 活跃 IP:旧记录被裁剪,窗口内记录保留。 + active := l.entries["active"] + if len(active) != 1 || !active[0].Equal(base.Add(80*time.Second)) { + t.Fatalf("active 应仅保留窗口内的 1 条记录,实际 %v", active) + } + // 不活跃 IP:被正确回收。 + if _, ok := l.entries["stale"]; ok { + t.Fatal("不活跃 IP 应被回收") + } + // 回收不破坏仍有效的限流:active 在窗口内已有 1 条记录,limit=1 时下一次应被拒绝。 + if l.allowAt("active", 1, now) { + t.Fatal("活跃 IP 的窗口内记录应仍然生效(limit=1 时应拒绝)") + } +}