Cyrene_v2/internal/services/gateway/ws_hub_conn.go

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