Scoreboard/server/internal/httpapi/realtime_session_test.go

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