package httpapi_test import ( "context" "encoding/json" "net/http" "net/url" "strconv" "strings" "testing" "time" "github.com/coder/websocket" "git.yeij.top/AskaEth/Scoreboard/server/internal/auth" "git.yeij.top/AskaEth/Scoreboard/server/internal/httpapi" "git.yeij.top/AskaEth/Scoreboard/server/internal/realtime" ) // 会话撤销后 /ws 长连接必须断开,且新的握手不能再按原身份进来。 // // 这条链路有四个接头:Cookie → Authenticate 把会话写进 context → 握手时固化角色与会话 id → // 心跳里按会话 id 复查。少任何一环,被停用的人只要不关页面就还能一直收推送, // 所以这里用真实的 auth.Service 接真实 Hub,撤销动作也走真实接口(管理员停用用户)。 func TestWebSocketClosesWhenSessionRevoked(t *testing.T) { var hub *realtime.Hub env := newTestEnvWith(t, func(o *httpapi.Options) { h := realtime.NewHub(o.Event, o.Logger) h.SetHeartbeatInterval(50 * time.Millisecond) h.SetSessionChecker(o.Auth) t.Cleanup(func() { _ = h.Close() }) hub = h o.WS = h.Handler() }) env.createUser("boss", "boss-pass", "admin") victimID := env.createUser("victim", "victim-pass", "admin") // 被停用的人:登录后打开管理端长连接。 victim, _ := env.client() env.login(victim, "victim", "victim-pass") victimCookie := sessionCookie(t, victim, env.server.URL) c := dialWS(t, env.server.URL, victimCookie) c.send("hello", map[string]any{"role": "admin"}) c.recv("snapshot") // 管理员在后台停用该用户:HTTP 侧会连带作废他的全部会话。 boss, _ := env.client() env.login(boss, "boss", "boss-pass") status, body := env.do(boss, http.MethodPut, "/api/users/"+strconv.FormatInt(victimID, 10), map[string]any{"active": false}, true) if status != http.StatusOK || !body.OK { t.Fatalf("停用用户应 200,得到 %d %+v", status, body.Error) } // 最迟一个心跳周期内断开,关闭码 1008「登录失效」:前端据此停止重连并引导重新登录。 if code := c.waitClose(3 * time.Second); code != websocket.StatusPolicyViolation { t.Fatalf("会话撤销后 /ws 应以 1008 关闭,得到 %d", code) } // 握手侧同样要校验会话:仍带着这条失效 Cookie 再连一次,只能降级成 guest—— // 连得上(大屏是公开的),但申请管理端房间会被按策略拒绝。 re := dialWS(t, env.server.URL, victimCookie) re.send("hello", map[string]any{"role": "admin"}) if code := re.waitClose(3 * time.Second); code != websocket.StatusPolicyViolation { t.Fatalf("失效会话的握手不该拿到管理端身份,关闭码得到 %d", code) } if n := len(hub.Instances()); n != 0 { t.Fatalf("被拒的连接不该登记成大屏实例,得到 %d 条", n) } } // 会话仍有效时连接不受影响:复查不能把在线用户误判成失效。 func TestWebSocketStaysOpenWhileSessionActive(t *testing.T) { env := newTestEnvWith(t, func(o *httpapi.Options) { h := realtime.NewHub(o.Event, o.Logger) h.SetHeartbeatInterval(50 * time.Millisecond) h.SetSessionChecker(o.Auth) t.Cleanup(func() { _ = h.Close() }) o.WS = h.Handler() }) env.createUser("admin", "admin-pass", "admin") client, _ := env.client() env.login(client, "admin", "admin-pass") c := dialWS(t, env.server.URL, sessionCookie(t, client, env.server.URL)) c.send("hello", map[string]any{"role": "admin"}) c.recv("snapshot") // 连收三个心跳仍活着,说明会话复查没有误伤。 for i := 0; i < 3; i++ { c.recv("heartbeat") } } // ---------- 测试用 WebSocket 客户端 ---------- type wsFrame struct { event string err error code websocket.StatusCode } type wsTestConn struct { t *testing.T ctx context.Context cancel context.CancelFunc conn *websocket.Conn frames chan wsFrame } // dialWS 用浏览器会话 Cookie 打开 /ws,等价于管理端页面里的连接。 func dialWS(t *testing.T, serverURL, cookie string) *wsTestConn { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) header := http.Header{} if cookie != "" { header.Set("Cookie", cookie) } conn, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(serverURL, "http")+"/ws", &websocket.DialOptions{HTTPHeader: header}) if err != nil { cancel() t.Fatalf("连接 /ws: %v", err) } c := &wsTestConn{t: t, ctx: ctx, cancel: cancel, conn: conn, frames: make(chan wsFrame, 64)} go c.readLoop() t.Cleanup(c.close) return c } func (c *wsTestConn) readLoop() { defer close(c.frames) for { _, data, err := c.conn.Read(c.ctx) if err != nil { c.frames <- wsFrame{err: err, code: websocket.CloseStatus(err)} return } var env struct { Event string `json:"event"` } if err := json.Unmarshal(data, &env); err != nil { c.frames <- wsFrame{err: err} return } c.frames <- wsFrame{event: env.Event} } } func (c *wsTestConn) close() { _ = c.conn.CloseNow() c.cancel() } func (c *wsTestConn) send(event string, data any) { c.t.Helper() msg, err := json.Marshal(struct { Event string `json:"event"` Data any `json:"data"` }{Event: event, Data: data}) if err != nil { c.t.Fatalf("序列化 %s 帧: %v", event, err) } ctx, cancel := context.WithTimeout(c.ctx, 5*time.Second) defer cancel() if err := c.conn.Write(ctx, websocket.MessageText, msg); err != nil { c.t.Fatalf("发送 %s: %v", event, err) } } // recv 等待指定事件,跳过心跳等无关消息。 func (c *wsTestConn) recv(want string) { c.t.Helper() deadline := time.Now().Add(5 * time.Second) for time.Now().Before(deadline) { select { case f, ok := <-c.frames: if !ok || f.err != nil { c.t.Fatalf("等待 %s 时连接已断开: %v", want, f.err) } if f.event == want { return } case <-time.After(time.Until(deadline)): } } c.t.Fatalf("等待事件 %s 超时", want) } // waitClose 等待服务端关闭连接并返回关闭码;超时返回 0。 func (c *wsTestConn) waitClose(d time.Duration) websocket.StatusCode { c.t.Helper() timer := time.NewTimer(d) defer timer.Stop() for { select { case f, ok := <-c.frames: if !ok { return 0 } if f.err != nil { return f.code } case <-timer.C: return 0 } } } // sessionCookie 从登录后的 Cookie jar 里取会话 Cookie,拼成请求头。 func sessionCookie(t *testing.T, client *http.Client, serverURL string) string { t.Helper() u, err := url.Parse(serverURL) if err != nil { t.Fatalf("解析测试服务地址: %v", err) } for _, c := range client.Jar.Cookies(u) { if c.Name == auth.CookieSession { return c.Name + "=" + c.Value } } t.Fatal("登录后应拿到会话 Cookie") return "" }