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 wraps one WebSocket connection. Writes are serialized through the // buffered send channel drained by a single writePump, so the gorilla conn is // only ever touched by one goroutine at a time. type wsConn struct { conn *websocket.Conn send chan []byte done chan struct{} h *hub claims *auth.Claims once sync.Once } // handleWS upgrades /ws/chat after authenticating the JWT from the query token or // the Authorization header. A failed upgrade writes a 401 {error:...} before the // handshake, so a browser/IM client sees a clean reject. 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 reads client frames until the connection closes, routing each by type. 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 drains the send channel (and sends keepalive pings) until the // connection closes. It is the only writer on the 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 publishes a bus event without any per-connection reply. func (c *wsConn) publishPublic(subject string, payload any) { if c.h.bus == nil { return } _ = c.h.bus.Publish(context.Background(), subject, payload) } // shutdown closes the connection exactly once. It tolerates a nil *websocket.Conn // so tests that build a fake connection can rely on a safe teardown. func (c *wsConn) shutdown() { c.once.Do(func() { close(c.done) if c.conn != nil { _ = c.conn.Close() } }) } // enqueue queues a serialized frame, dropping it when a slow client's buffer is // full so the producer never blocks. func (c *wsConn) enqueue(b []byte) { select { case c.send <- b: default: } } // enqueueJSON serializes ev and queues it. 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 extracts the JWT from the ?token= query or the Authorization header. 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 "" }