acfunlive-remake/webui_test.go

// 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")
	}
}