diff --git a/webui.go b/webui.go index 20b8a5f..9846778 100644 --- a/webui.go +++ b/webui.go @@ -2460,7 +2460,8 @@ func checkAuthToken(r *http.Request) bool { // 口令错误次数的记录，给反复尝试的请求加一点延迟 var authFail struct { sync.Mutex - m map[string]authFailEntry + m map[string]authFailEntry + sinceSweep int // 距上次扫表又记了多少条 } type authFailEntry struct { @@ -2471,6 +2472,62 @@ type authFailEntry struct { // 口令错误时最多延迟这么久，避免拖垮服务器 const authMaxDelay = 2 * time.Second +const ( + // 表里最多留这么多条。IP 是外部给的，必须有个硬上限：只靠过期清理的话， + // 每换一个新来源就多一条，清理追不上写入，最后把内存写满 + authFailMaxEntries = 10000 + // 超过这么久没再失败，这条记录就当没人用了 + authFailExpire = 10 * time.Minute + // 同一个来源最多记到这个次数：再往上涨延迟也是封顶的，没必要让数字一直加 + authFailMaxCount = 1000 + // 每记这么多条扫一次表。扫一次要遍历（必要时还要排序），不必每次都做 + authFailSweepEvery = 256 +) + +// 扫一遍口令错误记录：先删过期项，还多就把「最久没动静」的淘汰掉。 +// 淘汰时留出一成余量，这样满了之后还能再收一千来个新来源才需要重排下一次全表， +// 不至于每来一个就把整张表排一遍（那反倒成了另一条拖慢服务的路子）。 +func sweepAuthFail(now time.Time) { + for k, v := range authFail.m { + if now.Sub(v.last) > authFailExpire { + delete(authFail.m, k) + } + } + target := authFailMaxEntries * 9 / 10 + over := len(authFail.m) - target + if over <= 0 { + return + } + type record struct { + ip string + last time.Time + } + all := make([]record, 0, len(authFail.m)) + for k, v := range authFail.m { + all = append(all, record{k, v.last}) + } + sort.Slice(all, func(i, j int) bool { return all[i].last.Before(all[j].last) }) + for i := 0; i < over && i < len(all); i++ { + delete(authFail.m, all[i].ip) + } +} + +// 等待口令错误的惩罚延迟。客户端一断开就立刻返回， +// 否则并发撞口令时一堆 goroutine 各占着一条连接干等两秒，很容易把服务拖住 +func awaitAuthDelay(ctx context.Context, d time.Duration) bool { + if d <= 0 { + return true + } + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-timer.C: + return true + case <-ctx.Done(): + return false + } +} + // 取请求的来源 IP func clientIP(r *http.Request) string { if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { @@ -2486,17 +2543,23 @@ func authFailDelay(ip string) time.Duration { if authFail.m == nil { authFail.m = make(map[string]authFailEntry) } - // 记录太多时清掉过期的，避免大量不同来源的请求把内存撑满 - if len(authFail.m) > 1000 { - for k, v := range authFail.m { - if time.Since(v.last) > 10*time.Minute { - delete(authFail.m, k) - } - } + now := time.Now() + // 到上限就立刻扫一次。只靠定期扫的话，两个周期之间涌入的新来源 + // 会把表顶破上限（IP 是外部给的，写入速度不受我们控制） + if len(authFail.m) >= authFailMaxEntries { + authFail.sinceSweep = 0 + sweepAuthFail(now) + } + authFail.sinceSweep++ + if authFail.sinceSweep >= authFailSweepEvery { + authFail.sinceSweep = 0 + sweepAuthFail(now) } e := authFail.m[ip] - e.count++ - e.last = time.Now() + if e.count < authFailMaxCount { + e.count++ + } + e.last = now authFail.m[ip] = e d := time.Duration(e.count) * 100 * time.Millisecond if d > authMaxDelay { @@ -2527,7 +2590,10 @@ func webAuth(next http.Handler) http.Handler { return } if d := authFailDelay(ip); d > 0 { - time.Sleep(d) + // 客户端已经断开就不用再等了，直接收工 + if !awaitAuthDelay(r.Context(), d) { + return + } } w.Header().Set("Content-Type", "application/json; charset=utf-8") w.WriteHeader(http.StatusUnauthorized) diff --git a/webui_test.go b/webui_test.go index 3c7ba32..9923ea8 100644 --- a/webui_test.go +++ b/webui_test.go @@ -2,10 +2,13 @@ package main import ( + "context" + "fmt" "net/http" "net/http/httptest" "net/url" "testing" + "time" ) // 测试期间换一个管理口令 @@ -155,3 +158,148 @@ func TestAuthTokenCandidatesOrder(t *testing.T) { t.Fatalf("非 Bearer 的 Authorization 不该当候选：%q", got) } } + +// ==================== 口令错误记录的容量 ==================== +// IP 是外部给的：只靠「过期清理」挡不住换着地址撞口令的情况， +// 必须要有硬上限。这些用例锁住这一点。 + +// 换着来源持续撞口令，表必须被压在上限以内，不能一直涨 +func TestAuthFailTableHasHardCap(t *testing.T) { + saved, savedSweep := authFail.m, authFail.sinceSweep + t.Cleanup(func() { + authFail.Lock() + authFail.m, authFail.sinceSweep = saved, savedSweep + authFail.Unlock() + }) + authFail.Lock() + authFail.m, authFail.sinceSweep = nil, 0 + authFail.Unlock() + + // 每个都是新来源、且都是「刚刚」失败过 —— 过期清理一条都删不掉的情况 + for i := 0; i < authFailMaxEntries*2; i++ { + authFailDelay(fmt.Sprintf("198.51.100.%d", i)) + } + got := 0 + authFail.Lock() + got = len(authFail.m) + authFail.Unlock() + if got > authFailMaxEntries { + t.Errorf("记录数应被压在上限 %d 以内，实际 %d（换 IP 撞口令会把内存写满）", + authFailMaxEntries, got) + } + t.Logf(" %d 个新来源撞完，表内留下 %d 条（上限 %d）", authFailMaxEntries*2, got, authFailMaxEntries) +} + +// 超过过期时间的记录要被扫掉，不是在表里一直躺着 +func TestAuthFailExpiredEntriesGetSwept(t *testing.T) { + saved, savedSweep := authFail.m, authFail.sinceSweep + t.Cleanup(func() { + authFail.Lock() + authFail.m, authFail.sinceSweep = saved, savedSweep + authFail.Unlock() + }) + authFail.Lock() + authFail.m, authFail.sinceSweep = nil, 0 + authFail.Unlock() + + for i := 0; i < authFailSweepEvery+5; i++ { + authFailDelay(fmt.Sprintf("203.0.113.%d", i)) + } + authFail.Lock() + before := len(authFail.m) + // 把记录的最后时间推到过期之前 + for k, v := range authFail.m { + authFail.m[k] = authFailEntry{count: v.count, last: time.Now().Add(-2 * authFailExpire)} + } + authFail.Unlock() + + sweepAuthFail(time.Now()) + authFail.Lock() + after := len(authFail.m) + authFail.Unlock() + if after != 0 { + t.Errorf("已经全部过期的记录应被扫干净，扫之前 %d 条，扫之后 %d 条", before, after) + } +} + +// 扫表必须是自动发生的：不要求调用方记着手动清理， +// 混在里面的过期项应该在若干次记录之后自己消失 +func TestAuthFailSweepsPeriodically(t *testing.T) { + saved, savedSweep := authFail.m, authFail.sinceSweep + t.Cleanup(func() { + authFail.Lock() + authFail.m, authFail.sinceSweep = saved, savedSweep + authFail.Unlock() + }) + authFail.Lock() + authFail.m, authFail.sinceSweep = nil, 0 + authFail.Unlock() + + // 直接往里塞一条早就过期了的记录（绕过 authFailDelay，模拟「以前留下的」） + authFail.Lock() + if authFail.m == nil { + authFail.m = make(map[string]authFailEntry) + } + authFail.m["203.0.113.99"] = authFailEntry{count: 1, last: time.Now().Add(-2 * authFailExpire)} + authFail.Unlock() + + // 之后只要还在有人失败，过一阵子它就应该被清掉 + for i := 0; i < authFailSweepEvery+1; i++ { + authFailDelay(fmt.Sprintf("198.51.100.%d", i)) + } + + authFail.Lock() + _, still := authFail.m["203.0.113.99"] + authFail.Unlock() + if still { + t.Errorf("记了 %d 次之后，过期记录应当已经被自动扫掉", authFailSweepEvery+1) + } +} + +// 同一个来源的计数不能无限涨（临界点是延迟早已封顶，再涨只是让数字溢出） +func TestAuthFailCountCapped(t *testing.T) { + saved, savedSweep := authFail.m, authFail.sinceSweep + t.Cleanup(func() { + authFail.Lock() + authFail.m, authFail.sinceSweep = saved, savedSweep + authFail.Unlock() + }) + authFail.Lock() + authFail.m, authFail.sinceSweep = nil, 0 + authFail.Unlock() + + for i := 0; i < authFailMaxCount+50; i++ { + authFailDelay("192.0.2.9") + } + authFail.Lock() + c := authFail.m["192.0.2.9"].count + authFail.Unlock() + if c > authFailMaxCount { + t.Errorf("同一来源的次数应封顶在 %d，实际 %d", authFailMaxCount, c) + } + if d := authFailDelay("192.0.2.9"); d != authMaxDelay { + t.Errorf("打满之后延迟应封顶在 %v，实际 %v", authMaxDelay, d) + } +} + +// 客户端断开时延迟等待必须立刻返回，不能傻等满 2 秒 +func TestAwaitAuthDelayStopsOnCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(80 * time.Millisecond) + cancel() + }() + start := time.Now() + waited := awaitAuthDelay(ctx, 2*time.Second) + used := time.Since(start) + if waited { + t.Error("上游取消时 awaitAuthDelay 应返回 false") + } + if used > time.Second { + t.Errorf("上游取消后应立刻返回，实际等了 %v", used) + } + // 正常等待不受影响 + if !awaitAuthDelay(context.Background(), 10*time.Millisecond) { + t.Error("没被取消时应当正常等完再返回 true") + } +}