package gateway import ( "context" "encoding/json" "net/http" "strings" "sync" "time" "github.com/gorilla/websocket" "git.yeij.top/AskaEth/Cyrene_v2/internal/auth" apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors" "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/history" ) // wsConn 包装一个 WebSocket 连接。写入经缓冲的 send // 通道串行化,由单个 writePump 排空,因此 gorilla conn // 在任一时刻只被一个 goroutine 触碰。 type wsConn struct { conn *websocket.Conn send chan []byte done chan struct{} h *hub claims *auth.Claims once sync.Once } // handleWS 在校验来自查询 token 或 Authorization 头的 JWT 之后 // 升级 /ws/chat。升级失败会在握手前写出 401 {error:...}, // 使浏览器/IM 客户端看到干净的拒绝。 func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "websocket requires GET")) return } claims, err := s.auth.Verify(wsToken(r)) if err != nil { s.log.Debug("ws auth", "err", err) writeError(w, apperr.ErrUnauthorized) return } conn, err := s.hub.upgrader.Upgrade(w, r, nil) if err != nil { s.log.Debug("ws upgrade", "err", err) return } c := &wsConn{ conn: conn, send: make(chan []byte, wsSendBuffer), done: make(chan struct{}), h: s.hub, claims: claims, } s.hub.register(c) s.hub.replayHistory(c, history.MainSessionID) go c.writePump() go c.readPump() } // readPump 读取客户端帧直到连接关闭,按类型路由每个帧。 func (c *wsConn) readPump() { defer func() { c.h.unregister(c) c.shutdown() }() c.conn.SetReadLimit(wsReadLimit) _ = c.conn.SetReadDeadline(time.Now().Add(wsPongWait)) c.conn.SetPongHandler(func(string) error { return c.conn.SetReadDeadline(time.Now().Add(wsPongWait)) }) for { _, msg, err := c.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { c.h.log.Debug("ws read", "err", err) } return } var in wsInbound if err := json.Unmarshal(msg, &in); err != nil { c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "invalid json"}}) continue } c.route(in) } } // writePump 排空 send 通道(并发送 keepalive ping)直到 // 连接关闭。它是 gorilla conn 上唯一的写入者。 func (c *wsConn) writePump() { ticker := time.NewTicker(wsPingPeriod) defer func() { ticker.Stop() _ = c.conn.Close() }() for { select { case <-c.done: return case msg := <-c.send: _ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil { return } case <-ticker.C: _ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } // publishPublic 发布一个 bus 事件,且不做任何按连接的回复。 func (c *wsConn) publishPublic(subject string, payload any) { if c.h.bus == nil { return } _ = c.h.bus.Publish(context.Background(), subject, payload) } // shutdown 恰好关闭连接一次。它容忍 nil 的 *websocket.Conn, // 因此构建假连接的测试可以依赖安全的清理。 func (c *wsConn) shutdown() { c.once.Do(func() { close(c.done) if c.conn != nil { _ = c.conn.Close() } }) } // closed 报告连接是否已被关闭(readPump 退出时 shutdown 关闭 done)。异步收尾 // 路径(如语音回合结算后的定向回帧)用它判断「这条回帧还送得到吗」—— send 通道 // 不会在关闭时被关上,入队后只会静默丢失,所以必须显式问一次,才能给出一条带 // 来源的 WARN 而不是无声无息。nil 字段(零值 wsConn)视为未关闭。 func (c *wsConn) closed() bool { select { case <-c.done: return true default: return false } } // enqueue 把序列化后的帧入队,慢客户端缓冲满时丢弃它, // 使生产者绝不阻塞。 func (c *wsConn) enqueue(b []byte) { select { case c.send <- b: default: } } // enqueueJSON 序列化 ev 并把它入队。 func (c *wsConn) enqueueJSON(ev wsEvent) { b, err := json.Marshal(ev) if err != nil { c.h.log.Warn("marshal ws event", "err", err) return } c.enqueue(b) } // wsToken 从 ?token= 查询或 Authorization 头中提取 JWT。 func wsToken(r *http.Request) string { if t := r.URL.Query().Get("token"); t != "" { return t } const prefix = "Bearer " h := r.Header.Get("Authorization") if strings.HasPrefix(h, prefix) { return strings.TrimSpace(h[len(prefix):]) } return "" }