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 包装一个 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()
		}
	})
}

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