jiang13-bbs/backend/realtime/client.go

package realtime

import (
	"encoding/json/v2"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

func uintToString(n uint) string {
	return strconv.FormatUint(uint64(n), 10)
}

const (
	// writeWait 单次写帧超时
	writeWait = 10 * time.Second
	// pongWait 无 Pong 超时：浏览器自动响应协议层 Ping，超时即判定断连
	pongWait = 70 * time.Second
	// pingPeriod 协议层 Ping 周期（须小于 pongWait）
	pingPeriod = 30 * time.Second
	// maxMessageSize 仅接收应用层心跳，限制 4KiB 足够
	maxMessageSize = 4 << 10
	// sendBufferSize 慢消费保护阈值
	sendBufferSize = 32
	// CloseAuthInvalid 鉴权失效（封禁/强制下线）自定义关闭码，前端不自动重连
	CloseAuthInvalid = 4401
)

// Client 一条 WebSocket 连接
type Client struct {
	hub    *Hub
	conn   *websocket.Conn
	userID uint
	// staff 建连时按 DB 实时 Actor 判定，决定是否加入管理团队房间
	staff bool
	// validate 应用层心跳时复查 token_version / 封禁状态，失效则立即断连
	validate func() bool
	// onHeartbeat 心跳顺带刷新 last_seen_at（SQL 60s 限频），
	// 让只挂着 WS、不发 HTTP 请求的标签页也保持在线统计准确
	onHeartbeat func()
	// authorizeRoom 客户端动态订阅房间前的业务鉴权（如校验群成员身份）；
	// 仅 "chat:" 前缀房间允许订阅，user:/staff 等内部房间禁止
	authorizeRoom func(room string) bool

	send    chan []byte
	rooms   map[string]struct{}
	closed  atomic.Bool
	closeCh chan struct{}
}

func NewClient(
	hub *Hub,
	conn *websocket.Conn,
	userID uint,
	staff bool,
	validate func() bool,
	onHeartbeat func(),
	authorizeRoom func(room string) bool,
) *Client {
	return &Client{
		hub:           hub,
		conn:          conn,
		userID:        userID,
		staff:         staff,
		validate:      validate,
		onHeartbeat:   onHeartbeat,
		authorizeRoom: authorizeRoom,
		send:          make(chan []byte, sendBufferSize),
		rooms:         make(map[string]struct{}),
		closeCh:       make(chan struct{}),
	}
}

// Serve 注册到 Hub 后阻塞运行读写泵，断开时注销（调用方在升级成功后调用一次）。
func (c *Client) Serve(sendHello func(*Client)) {
	cameOnline := c.hub.register(c)
	sendHello(c)
	if cameOnline {
		c.hub.BroadcastPresence(c.userID, true)
	}

	go c.writePump()
	c.readPump() // 阻塞至断连

	wentOffline := c.hub.unregister(c)
	c.forceClose()
	if wentOffline {
		c.hub.BroadcastPresence(c.userID, false)
	}
}

// forceClose 幂等关闭连接并触发 closeCh
func (c *Client) forceClose() {
	if c.closed.CompareAndSwap(false, true) {
		_ = c.conn.Close()
		close(c.closeCh)
	}
}

// clientInbound 浏览器→服务端帧：ping 心跳；join/leave 动态订阅/退订房间
type clientInbound struct {
	Type string `json:"type"`
	Room string `json:"room,omitempty"`
}

// readPump 单读泵：Pong 续期 + 应用层心跳鉴权复查 + 房间订阅
func (c *Client) readPump() {
	c.conn.SetReadLimit(maxMessageSize)
	_ = c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error {
		return c.conn.SetReadDeadline(time.Now().Add(pongWait))
	})

	for {
		_, raw, err := c.conn.ReadMessage()
		if err != nil {
			return
		}
		var msg clientInbound
		if err := json.Unmarshal(raw, &msg); err != nil {
			continue // 非法帧忽略，不踢连接
		}
		switch msg.Type {
		case "ping":
			// 心跳鉴权复查：被封禁/改密/降级强制下线后 ~25s 内断开
			if c.validate != nil && !c.validate() {
				_ = c.conn.WriteControl(
					websocket.CloseMessage,
					websocket.FormatCloseMessage(CloseAuthInvalid, "auth invalid"),
					time.Now().Add(writeWait),
				)
				return
			}
			if c.onHeartbeat != nil {
				c.onHeartbeat()
			}
			c.enqueueMustMarshal(Envelope{Type: EventPong, Data: map[string]int64{
				"ts": time.Now().Unix(),
			}})
		case "join", "leave":
			c.handleRoomFrame(msg)
		}
	}
}

// handleRoomFrame 处理房间订阅：只允许 chat:{id}，且实时校验调用方业务身份
func (c *Client) handleRoomFrame(msg clientInbound) {
	room := msg.Room
	if room == "" || len(room) > 32 || !strings.HasPrefix(room, "chat:") {
		return
	}
	idStr := strings.TrimPrefix(room, "chat:")
	roomID, err := strconv.ParseUint(idStr, 10, 64)
	if err != nil || roomID == 0 {
		return
	}
	if c.authorizeRoom != nil && !c.authorizeRoom(room) {
		return
	}
	if msg.Type == "join" {
		c.hub.JoinClientRoom(c, room)
	} else {
		c.hub.LeaveClientRoom(c, room)
	}
}

// writePump 单写泵：唯一持有连接写权限；协议层 Ping + 业务帧都从这里发出
func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	for {
		select {
		case payload := <-c.send:
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.TextMessage, payload); err != nil {
				return
			}
		case <-ticker.C:
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		case <-c.closeCh:
			return
		}
	}
}

// Send 业务帧入队（由唯一 writePump 发出）；连接已关闭时静默丢弃
func (c *Client) Send(env Envelope) {
	c.enqueueMustMarshal(env)
}

func (c *Client) enqueueMustMarshal(env Envelope) {
	payload, err := json.Marshal(env)
	if err != nil {
		return
	}
	select {
	case c.send <- payload:
	case <-c.closeCh:
	}
}