Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6ab7886865 | ||
|
|
d18983e23d | ||
|
|
bc23f5039a | ||
|
|
2a647ad3e3 |
No files matched your search
@@ -16,7 +16,7 @@ jobs:
|
|||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v4
|
- uses: actions/checkout@v4
|
||||||
with:
|
with:
|
||||||
token: ${{ secrets.GITHUB_TOKEN }}
|
token: ${{ secrets.PAT_TOKEN }}
|
||||||
fetch-depth: 0
|
fetch-depth: 0
|
||||||
|
|
||||||
- name: Configure Git
|
- name: Configure Git
|
||||||
|
|||||||
+416
@@ -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 的真实时刻重算,否则回收静默不触发")
|
||||||
|
}
|
||||||
+108
@@ -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("加载失败时不应替换已有配置")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -11,6 +11,7 @@ import (
|
|||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
"text/template"
|
"text/template"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -25,7 +26,8 @@ const (
|
|||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
exps = []*regexp.Regexp{
|
// 入口链接:用户可直接粘贴的 GitHub 链接。命中这些的 302 会被重写为代理地址。
|
||||||
|
entryExps = []*regexp.Regexp{
|
||||||
regexp.MustCompile(`^(?:https?://)?github\.com/([^/]+)/([^/]+)/(?:releases|archive)/.*$`),
|
regexp.MustCompile(`^(?:https?://)?github\.com/([^/]+)/([^/]+)/(?:releases|archive)/.*$`),
|
||||||
regexp.MustCompile(`^(?:https?://)?github\.com/([^/]+)/([^/]+)/(?:blob|raw)/.*$`),
|
regexp.MustCompile(`^(?:https?://)?github\.com/([^/]+)/([^/]+)/(?:blob|raw)/.*$`),
|
||||||
regexp.MustCompile(`^(?:https?://)?github\.com/([^/]+)/([^/]+)/(?:info|git-).*$`),
|
regexp.MustCompile(`^(?:https?://)?github\.com/([^/]+)/([^/]+)/(?:info|git-).*$`),
|
||||||
@@ -33,17 +35,31 @@ var (
|
|||||||
regexp.MustCompile(`^(?:https?://)?gist\.github(?:usercontent|)\.com/([^/]+)/.+?/.+$`),
|
regexp.MustCompile(`^(?:https?://)?gist\.github(?:usercontent|)\.com/([^/]+)/.+?/.+$`),
|
||||||
regexp.MustCompile(`^(?:https?://)?api\.github\.com/.+?/([^/]+)(?:/.*)?$`),
|
regexp.MustCompile(`^(?:https?://)?api\.github\.com/.+?/([^/]+)(?:/.*)?$`),
|
||||||
}
|
}
|
||||||
|
// 新增:GitHub 资源 CDN 主机(release 资产签名直链 / 归档分流 / LFS 媒体)。
|
||||||
|
// 这些链接可以直接代理,但**不参与 Location 重写**——保证 github.com 的 302
|
||||||
|
// 仍在服务端内部跟随,客户端拿到的始终是文件内容而不是真实文件地址。
|
||||||
|
assetExps = []*regexp.Regexp{
|
||||||
|
// release-assets:新版 release 资产签名直链(302 目标)
|
||||||
|
// objects / github-releases:旧版资产 CDN 主机(兼容历史链接)
|
||||||
|
regexp.MustCompile(`^(?:https?://)?(?:release-assets|objects|github-releases)\.githubusercontent\.com/([^/]+).*$`),
|
||||||
|
// codeload:源码归档(archive / tar.gz / zip)分流主机
|
||||||
|
regexp.MustCompile(`^(?:https?://)?codeload\.github\.com/([^/]+)/([^/]+)/.*$`),
|
||||||
|
// media:Git LFS 媒体文件
|
||||||
|
regexp.MustCompile(`^(?:https?://)?media\.githubusercontent\.com/(?:media/)?([^/]+)/([^/]+)/.*$`),
|
||||||
|
}
|
||||||
|
// checkURL 使用:入口 + 资源 CDN
|
||||||
|
exps = append(append([]*regexp.Regexp{}, entryExps...), assetExps...)
|
||||||
httpClient *http.Client
|
httpClient *http.Client
|
||||||
config *Config
|
// config 以不可变快照方式发布:loadConfig 构造新对象后原子替换,
|
||||||
configLock sync.RWMutex
|
// 请求路径只做 atomic.Load,避免配置重载与请求处理之间的数据竞争和锁开销。
|
||||||
log = logrus.New()
|
config atomic.Pointer[Config]
|
||||||
|
log = logrus.New()
|
||||||
|
|
||||||
tmpl *template.Template
|
tmpl *template.Template
|
||||||
tmpl404 *template.Template
|
tmpl404 *template.Template
|
||||||
|
|
||||||
// 新增:用于记录每个 IP 的请求时间
|
// 每个 IP 独立的滑动窗口限流器
|
||||||
ipRequests = make(map[string][]time.Time)
|
limiter = newRateLimiter()
|
||||||
ipLock sync.Mutex
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type Config struct {
|
type Config struct {
|
||||||
@@ -70,6 +86,96 @@ func init() {
|
|||||||
FullTimestamp: true,
|
FullTimestamp: true,
|
||||||
}
|
}
|
||||||
log.SetFormatter(&formatter)
|
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 {
|
func customLogger() gin.HandlerFunc {
|
||||||
@@ -147,9 +253,7 @@ func main() {
|
|||||||
log.Info("start HTTP server @ ", httpBase)
|
log.Info("start HTTP server @ ", httpBase)
|
||||||
|
|
||||||
router.GET("/", func(c *gin.Context) {
|
router.GET("/", func(c *gin.Context) {
|
||||||
configLock.RLock()
|
analyticsURL := config.Load().AnalyticsURL
|
||||||
analyticsURL := config.AnalyticsURL
|
|
||||||
configLock.RUnlock()
|
|
||||||
|
|
||||||
data := struct {
|
data := struct {
|
||||||
AnalyticsURL string
|
AnalyticsURL string
|
||||||
@@ -171,6 +275,9 @@ func main() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func handler(c *gin.Context) {
|
func handler(c *gin.Context) {
|
||||||
|
// 取一次不可变配置快照,保证整个请求内看到一致的配置,且无锁、无数据竞争。
|
||||||
|
cfg := config.Load()
|
||||||
|
|
||||||
rawPath := strings.TrimPrefix(c.Request.URL.RequestURI(), "/")
|
rawPath := strings.TrimPrefix(c.Request.URL.RequestURI(), "/")
|
||||||
|
|
||||||
for strings.HasPrefix(rawPath, "/") {
|
for strings.HasPrefix(rawPath, "/") {
|
||||||
@@ -183,53 +290,35 @@ func handler(c *gin.Context) {
|
|||||||
|
|
||||||
matches := checkURL(rawPath)
|
matches := checkURL(rawPath)
|
||||||
if matches != nil {
|
if matches != nil {
|
||||||
if len(config.WhiteList) > 0 && !checkList(matches, config.WhiteList) {
|
if len(cfg.WhiteList) > 0 && !checkList(matches, cfg.WhiteList) {
|
||||||
render404(c, tmpl404)
|
render404(c, tmpl404, cfg)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if len(config.BlackList) > 0 && checkList(matches, config.BlackList) {
|
if len(cfg.BlackList) > 0 && checkList(matches, cfg.BlackList) {
|
||||||
render404(c, tmpl404)
|
render404(c, tmpl404, cfg)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
render404(c, tmpl404)
|
render404(c, tmpl404, cfg)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if exps[1].MatchString(rawPath) {
|
if entryExps[1].MatchString(rawPath) {
|
||||||
rawPath = strings.Replace(rawPath, "/blob/", "/raw/", 1)
|
rawPath = strings.Replace(rawPath, "/blob/", "/raw/", 1)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 新增:基于时间的速率限制
|
// 基于时间的速率限制:只处理当前 IP,窗口内允许恰好 limitRate 次。
|
||||||
|
// 时间在 limiter.allow 内部持锁后采样,避免并发下时间戳乱序。
|
||||||
clientIP := c.ClientIP()
|
clientIP := c.ClientIP()
|
||||||
ipLock.Lock()
|
limitRate := cfg.RequestLimit.LimitRate
|
||||||
limitRate := config.RequestLimit.LimitRate
|
if !limiter.allow(clientIP, limitRate) {
|
||||||
now := time.Now()
|
log.Debugf("clientIP: %s rate limited (limit=%d)", clientIP, limitRate)
|
||||||
// 移除一分钟前的请求记录
|
|
||||||
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 的请求次数是否超过限制
|
|
||||||
if len(ipRequests[clientIP]) > int(limitRate) {
|
|
||||||
ipLock.Unlock()
|
|
||||||
c.String(http.StatusTooManyRequests, "Too Many Requests.")
|
c.String(http.StatusTooManyRequests, "Too Many Requests.")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
// 记录当前请求时间
|
|
||||||
ipRequests[clientIP] = append(ipRequests[clientIP], now)
|
|
||||||
ipLock.Unlock()
|
|
||||||
|
|
||||||
// 限制访问 IP
|
// 限制访问 IP
|
||||||
configLock.RLock()
|
limitAddr := cfg.RequestLimit.LimitAddr
|
||||||
limitAddr := config.RequestLimit.LimitAddr
|
|
||||||
configLock.RUnlock()
|
|
||||||
if len(limitAddr) > 0 {
|
if len(limitAddr) > 0 {
|
||||||
for _, ip := range limitAddr {
|
for _, ip := range limitAddr {
|
||||||
if clientIP == ip {
|
if clientIP == ip {
|
||||||
@@ -240,9 +329,7 @@ func handler(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// 限制请求参数
|
// 限制请求参数
|
||||||
configLock.RLock()
|
limitParm := cfg.RequestLimit.LimitParm
|
||||||
limitParm := config.RequestLimit.LimitParm
|
|
||||||
configLock.RUnlock()
|
|
||||||
if len(limitParm) > 0 {
|
if len(limitParm) > 0 {
|
||||||
for key, value := range limitParm {
|
for key, value := range limitParm {
|
||||||
if c.Request.Header.Get(key) == value {
|
if c.Request.Header.Get(key) == value {
|
||||||
@@ -252,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)
|
req, err := http.NewRequest(c.Request.Method, u, c.Request.Body)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
c.String(http.StatusInternalServerError, fmt.Sprintf("server error %v", err))
|
c.String(http.StatusInternalServerError, fmt.Sprintf("server error %v", err))
|
||||||
@@ -288,11 +375,10 @@ func proxy(c *gin.Context, u string) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
configLock.RLock()
|
limitSize := cfg.RequestLimit.LimitSize * 1024 * 1024 // Convert MB to Bytes
|
||||||
limitSize := config.RequestLimit.LimitSize * 1024 * 1024 // Convert MB to Bytes
|
|
||||||
configLock.RUnlock()
|
|
||||||
|
|
||||||
if size > limitSize {
|
// 如果 limitSize <= 0 则不限制文件大小
|
||||||
|
if limitSize > 0 && size > limitSize {
|
||||||
c.String(http.StatusRequestEntityTooLarge, "File too large.")
|
c.String(http.StatusRequestEntityTooLarge, "File too large.")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -302,6 +388,12 @@ func proxy(c *gin.Context, u string) {
|
|||||||
resp.Header.Del("Referrer-Policy")
|
resp.Header.Del("Referrer-Policy")
|
||||||
resp.Header.Del("Strict-Transport-Security")
|
resp.Header.Del("Strict-Transport-Security")
|
||||||
|
|
||||||
|
// 如果 GitHub 返回 404,显示我们的 404 页面
|
||||||
|
if resp.StatusCode == http.StatusNotFound {
|
||||||
|
render404(c, tmpl404, cfg)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
for key, values := range resp.Header {
|
for key, values := range resp.Header {
|
||||||
for _, value := range values {
|
for _, value := range values {
|
||||||
c.Header(key, value)
|
c.Header(key, value)
|
||||||
@@ -309,10 +401,12 @@ func proxy(c *gin.Context, u string) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if location := resp.Header.Get("Location"); location != "" {
|
if location := resp.Header.Get("Location"); location != "" {
|
||||||
if checkURL(location) != nil {
|
// 仅入口链接重写为代理地址;资源 CDN 等跳转目标在服务端内部跟随,
|
||||||
|
// 直接返回文件内容(客户端看不到真实文件地址)。
|
||||||
|
if isEntryURL(location) {
|
||||||
c.Header("Location", "/"+location)
|
c.Header("Location", "/"+location)
|
||||||
} else {
|
} else {
|
||||||
proxy(c, location)
|
proxy(c, location, cfg)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -324,11 +418,16 @@ func proxy(c *gin.Context, u string) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func loadConfig() {
|
func loadConfig() {
|
||||||
|
_ = loadConfigFrom("config.json")
|
||||||
|
}
|
||||||
|
|
||||||
|
// loadConfigFrom 从指定路径加载配置并原子发布。返回错误便于测试验证。
|
||||||
|
func loadConfigFrom(path string) error {
|
||||||
log.Info("loading config...")
|
log.Info("loading config...")
|
||||||
file, err := os.Open("config.json")
|
file, err := os.Open(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Errorf("load config error: %v\n", err)
|
log.Errorf("load config error: %v\n", err)
|
||||||
return
|
return err
|
||||||
}
|
}
|
||||||
defer func(file *os.File) {
|
defer func(file *os.File) {
|
||||||
err := file.Close()
|
err := file.Close()
|
||||||
@@ -341,17 +440,17 @@ func loadConfig() {
|
|||||||
decoder := json.NewDecoder(file)
|
decoder := json.NewDecoder(file)
|
||||||
if err := decoder.Decode(&newConfig); err != nil {
|
if err := decoder.Decode(&newConfig); err != nil {
|
||||||
log.Errorf("decod config error: %v\n", err)
|
log.Errorf("decod config error: %v\n", err)
|
||||||
return
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
configLock.Lock()
|
// 构造完成后整体原子替换,读方无需加锁,也不会看到半更新的配置。
|
||||||
config = &newConfig
|
config.Store(&newConfig)
|
||||||
if newConfig.Debug {
|
if newConfig.Debug {
|
||||||
log.SetLevel(logrus.DebugLevel)
|
log.SetLevel(logrus.DebugLevel)
|
||||||
} else {
|
} else {
|
||||||
log.SetLevel(logrus.InfoLevel)
|
log.SetLevel(logrus.InfoLevel)
|
||||||
}
|
}
|
||||||
configLock.Unlock()
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func checkURL(u string) []string {
|
func checkURL(u string) []string {
|
||||||
@@ -363,6 +462,17 @@ func checkURL(u string) []string {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// isEntryURL 判断是否为「入口」链接(github.com / raw / gist / api)。
|
||||||
|
// 用于决定 302 是否重写为代理地址;资源 CDN 链接不在其中。
|
||||||
|
func isEntryURL(u string) bool {
|
||||||
|
for _, exp := range entryExps {
|
||||||
|
if exp.MatchString(u) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
func checkList(matches, list []string) bool {
|
func checkList(matches, list []string) bool {
|
||||||
for _, item := range list {
|
for _, item := range list {
|
||||||
if strings.HasPrefix(matches[0], item) {
|
if strings.HasPrefix(matches[0], item) {
|
||||||
@@ -372,11 +482,9 @@ func checkList(matches, list []string) bool {
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
func render404(c *gin.Context, tmpl *template.Template) {
|
func render404(c *gin.Context, tmpl *template.Template, cfg *Config) {
|
||||||
configLock.RLock()
|
analyticsURL := cfg.AnalyticsURL
|
||||||
analyticsURL := config.AnalyticsURL
|
supportImageURL := cfg.SupportImageURL
|
||||||
supportImageURL := config.SupportImageURL
|
|
||||||
configLock.RUnlock()
|
|
||||||
|
|
||||||
data := struct {
|
data := struct {
|
||||||
AnalyticsURL string
|
AnalyticsURL string
|
||||||
|
|||||||
+159
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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 时应拒绝)")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,61 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import "testing"
|
||||||
|
|
||||||
|
func TestCheckURL(t *testing.T) {
|
||||||
|
shouldMatch := []string{
|
||||||
|
"https://release-assets.githubusercontent.com/github-production-release-asset/1/abc?sig=x",
|
||||||
|
"https://objects.githubusercontent.com/foo/bar",
|
||||||
|
"https://github-releases.githubusercontent.com/foo/bar",
|
||||||
|
"https://codeload.github.com/cli/cli/tar.gz/refs/tags/v1",
|
||||||
|
"https://media.githubusercontent.com/media/octocat/Hello-World/master/README",
|
||||||
|
"https://media.githubusercontent.com/octocat/Hello-World/master/README",
|
||||||
|
"https://github.com/a/b/releases/download/v1/x.zip",
|
||||||
|
"https://raw.githubusercontent.com/a/b/main/c.txt",
|
||||||
|
"https://api.github.com/repos/a/b",
|
||||||
|
}
|
||||||
|
for _, u := range shouldMatch {
|
||||||
|
if checkURL(u) == nil {
|
||||||
|
t.Errorf("checkURL 应命中但未命中: %s", u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
shouldNotMatch := []string{
|
||||||
|
"https://evil.example.com/a.zip",
|
||||||
|
"https://example.com/",
|
||||||
|
"https://github.com.evil.com/a/b",
|
||||||
|
"https://notcodeload.github.com/a/b",
|
||||||
|
}
|
||||||
|
for _, u := range shouldNotMatch {
|
||||||
|
if checkURL(u) != nil {
|
||||||
|
t.Errorf("checkURL 不应命中但命中了: %s", u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsEntryURL(t *testing.T) {
|
||||||
|
// 入口:只有这些才允许重写 Location 返回客户端
|
||||||
|
entry := []string{
|
||||||
|
"https://github.com/a/b/releases/download/v1/x.zip",
|
||||||
|
"https://raw.githubusercontent.com/a/b/main/c.txt",
|
||||||
|
"https://gist.githubusercontent.com/u/1/raw/x.py",
|
||||||
|
"https://api.github.com/repos/a/b",
|
||||||
|
}
|
||||||
|
for _, u := range entry {
|
||||||
|
if !isEntryURL(u) {
|
||||||
|
t.Errorf("isEntryURL 应命中: %s", u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// 非入口:资源 CDN 必须在服务端内部跟随,不能重写给客户端
|
||||||
|
notEntry := []string{
|
||||||
|
"https://release-assets.githubusercontent.com/github-production-release-asset/1/abc?sig=x",
|
||||||
|
"https://objects.githubusercontent.com/foo/bar",
|
||||||
|
"https://github-releases.githubusercontent.com/foo/bar",
|
||||||
|
"https://codeload.github.com/cli/cli/tar.gz/refs/tags/v1",
|
||||||
|
"https://media.githubusercontent.com/media/octocat/Hello-World/master/README",
|
||||||
|
}
|
||||||
|
for _, u := range notEntry {
|
||||||
|
if isEntryURL(u) {
|
||||||
|
t.Errorf("isEntryURL 不应命中(否则会泄露真实文件地址): %s", u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user