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