// web UI 相关 package main import ( "context" "fmt" "net/http" "net/http/httptest" "net/url" "testing" "time" ) // 测试期间换一个管理口令 func withAuthToken(t *testing.T, token string) { t.Helper() streamers.Lock() old := config.AuthToken config.AuthToken = token streamers.Unlock() t.Cleanup(func() { streamers.Lock() config.AuthToken = old streamers.Unlock() }) } // 页面上发给服务端的 X-Auth-Token 是 encodeURIComponent 过的口令, // 这里用 url.QueryEscape 造出同样的形状(对中文而言两者产出一致) func pageHeader(r *http.Request, token string) *http.Request { r.Header.Set("X-Auth-Token", url.QueryEscape(token)) return r } func TestCheckAuthTokenPlainASCII(t *testing.T) { withAuthToken(t, "a1b2c3d4e5f6") r := httptest.NewRequest(http.MethodGet, "/api/status", nil) r.Header.Set("X-Auth-Token", "a1b2c3d4e5f6") if !checkAuthToken(r) { t.Fatal("ASCII 口令应该通过") } // 带 % 的口令:页面会编码成 %25,服务端解码后要能对上 withAuthToken(t, "ab%20cd") r = pageHeader(httptest.NewRequest(http.MethodGet, "/api/status", nil), "ab%20cd") if r.Header.Get("X-Auth-Token") != "ab%2520cd" { t.Fatalf("编码结果不对:%q", r.Header.Get("X-Auth-Token")) } if !checkAuthToken(r) { t.Fatal("含 % 的口令应该通过") } } func TestCheckAuthTokenNonASCII(t *testing.T) { // 中文口令是最容易踩的:原先直接把中文塞进请求头,浏览器会抛 // TypeError(Failed to read the 'headers' property…),请求根本发不出去, // 页面上只显示「无法连接到服务器」,用户被彻底锁在门外 withAuthToken(t, "我的口令abc123") r := pageHeader(httptest.NewRequest(http.MethodGet, "/api/status", nil), "我的口令abc123") if r.Header.Get("X-Auth-Token") == "我的口令abc123" { t.Fatal("测试前提不成立:头里不该是原文") } if !checkAuthToken(r) { t.Fatal("编码后的中文口令应该通过") } // 手写 curl 按文档原样发中文口令,也要能通过 r = httptest.NewRequest(http.MethodGet, "/api/status", nil) r.Header.Set("X-Auth-Token", "我的口令abc123") if !checkAuthToken(r) { t.Fatal("原样的中文口令也应该通过") } } func TestCheckAuthTokenAllSources(t *testing.T) { withAuthToken(t, "口令abc123") // Authorization: Bearer r := httptest.NewRequest(http.MethodGet, "/api/status", nil) r.Header.Set("Authorization", "Bearer "+url.QueryEscape("口令abc123")) if !checkAuthToken(r) { t.Fatal("Bearer 写法应该通过") } // ?token=(net/http 已经把查询串解码好了) r = httptest.NewRequest(http.MethodGet, "/api/status?token="+url.QueryEscape("口令abc123"), nil) if !checkAuthToken(r) { t.Fatal("查询参数写法应该通过") } // 下载链接里的 ?token= 也是这个路径,混着 % 的口令同样要对上 withAuthToken(t, "口令%20abc") r = httptest.NewRequest(http.MethodGet, "/api/file?token="+url.QueryEscape("口令%20abc"), nil) if !checkAuthToken(r) { t.Fatal("查询参数里的混合口令应该通过") } } func TestCheckAuthTokenRejects(t *testing.T) { withAuthToken(t, "正确的口令") cases := map[string]*http.Request{ "口令错了": func() *http.Request { return pageHeader(httptest.NewRequest(http.MethodGet, "/api/status", nil), "错误的口令") }(), "没带口令": httptest.NewRequest(http.MethodGet, "/api/status", nil), "空口令头": func() *http.Request { r := httptest.NewRequest(http.MethodGet, "/api/status", nil) r.Header.Set("X-Auth-Token", "") return r }(), "编码了但不是口令": func() *http.Request { r := httptest.NewRequest(http.MethodGet, "/api/status", nil) r.Header.Set("X-Auth-Token", "%E5%8F%A3%E4%BB%A4") return r }(), } for name, r := range cases { if checkAuthToken(r) { t.Fatalf("%s:不该通过", name) } } // 口令没配时一律拒绝,不能因为「都是空」就放行 withAuthToken(t, "") r := httptest.NewRequest(http.MethodGet, "/api/status", nil) r.Header.Set("X-Auth-Token", "") if checkAuthToken(r) { t.Fatal("口令为空时不该通过") } } func TestAuthTokenCandidatesOrder(t *testing.T) { // 三种写法同时出现时按 Bearer → X-Auth-Token → ?token= 的顺序给候选, // 顺序本身不影响结果,但别写出重复项(重复比较纯属浪费) r := httptest.NewRequest(http.MethodGet, "/api/status?token=qqq", nil) r.Header.Set("Authorization", "Bearer aaa") r.Header.Set("X-Auth-Token", "bbb") got := authTokenCandidates(r) want := []string{"aaa", "bbb", "qqq"} if len(got) != len(want) { t.Fatalf("候选数量不对:%q", got) } for i := range want { if got[i] != want[i] { t.Fatalf("候选顺序不对:%q,想要 %q", got, want) } } // Authorization 不是 Bearer 前缀时不算数 r = httptest.NewRequest(http.MethodGet, "/api/status", nil) r.Header.Set("Authorization", "Basic zzz") r.Header.Set("X-Auth-Token", "bbb") got = authTokenCandidates(r) if len(got) != 1 || got[0] != "bbb" { 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") } }