Files
toolbox/server/internal/tools/dual_check/ws_handler.go
T
2026-08-30 21:25:22 +08:00

224 lines
6.8 KiB
Go

package dual_check
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"unicode/utf8"
pkgws "cancerai.top/toolbox-server/internal/pkg/websocket"
"cancerai.top/toolbox-server/internal/response"
"github.com/gin-gonic/gin"
gorillaws "github.com/gorilla/websocket"
"go.uber.org/zap"
)
var upgrader = gorillaws.Upgrader{
CheckOrigin: func(r *http.Request) bool { return true },
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
type WSHandler struct {
service *Service
hub *pkgws.Hub
logger *zap.Logger
}
func NewWSHandler(service *Service, hub *pkgws.Hub, logger *zap.Logger) *WSHandler {
return &WSHandler{service: service, hub: hub, logger: logger}
}
// Handle GET /api/v1/dual-check/ws?session_id=xxx&role=initiator|participant&nickname=xxx
func (h *WSHandler) Handle(c *gin.Context) {
sessionID := c.Query("session_id")
role := c.Query("role")
nickname := c.Query("nickname")
browserID := c.Query("browser_id")
if sessionID == "" {
c.JSON(http.StatusBadRequest, gin.H{"code": 4000, "message": "session_id required"})
return
}
if role != "initiator" && role != "participant" {
c.JSON(http.StatusBadRequest, gin.H{"code": 1004, "message": "role must be initiator or participant"})
return
}
if nickname == "" || utf8.RuneCountInString(nickname) > 32 {
c.JSON(http.StatusBadRequest, gin.H{"code": 4000, "message": "nickname required (1-32)"})
return
}
// 先校验 session 存在
info, err := h.service.CheckSession(c.Request.Context(), sessionID)
if err != nil || !info.Exists {
c.JSON(http.StatusNotFound, gin.H{"code": 1001, "message": "session not found"})
return
}
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
h.logger.Warn("ws upgrade failed", zap.Error(err))
return
}
wsConn := &pkgws.Conn{
Hub: h.hub,
SessionID: sessionID,
Role: role,
WS: conn,
Send: make(chan []byte, 256),
}
h.hub.Register(wsConn)
// 参与者首次连接:自动登记昵称(冲突时自动改名)
actualNick := nickname
if role == "participant" {
_, actual, err := h.service.JoinParticipant(c.Request.Context(), sessionID, nickname, browserID)
if err == nil {
actualNick = actual
if actualNick != nickname {
// 自动改名:通知客户端 + 用 actualNick 广播
wsConn.Send <- marshalMsg(response.WSNicknameAutoChanged(actualNick))
h.logger.Info("ws participant auto-renamed",
zap.String("session_id", truncateSessionID(sessionID)),
zap.String("original", nickname),
zap.String("actual", actualNick))
}
h.hub.Broadcast(sessionID, response.WSParticipantJoined(actualNick))
} else {
h.logger.Warn("ws JoinParticipant failed",
zap.String("session_id", truncateSessionID(sessionID)),
zap.String("nickname", nickname),
zap.Error(err))
}
}
// 推一个 init 消息给新连接者,带上当前全量快照让它自己初始化 UI
if sess, err := h.service.GetSession(c.Request.Context(), sessionID); err == nil {
wsConn.Send <- marshalMsg(response.WSFullState(sess))
}
h.logger.Info("ws connected",
zap.String("session_id", truncateSessionID(sessionID)),
zap.String("role", role),
zap.String("nickname", nickname),
)
go pkgws.WritePump(wsConn)
// ReadPump 里的 Unregister 会触发 participant_left 广播
go h.readLoop(wsConn, actualNick)
}
func (h *WSHandler) readLoop(c *pkgws.Conn, initialNickname string) {
// currentNick 是可变的,rename 成功后更新,保证 defer 中的 participant_left 广播用最新昵称
currentNick := initialNickname
defer func() {
c.WS.Close() // 先关,让 WritePump 也能退
if c.Role == "initiator" {
// 发起方断开:立即结束会话 → 广播给所有参与者 → 删除 Redis
h.hub.Broadcast(c.SessionID, response.WSSessionEnded())
if err := h.service.EndSession(context.Background(), c.SessionID); err != nil {
h.logger.Warn("initiator disconnect end session err",
zap.String("session_id", truncateSessionID(c.SessionID)),
zap.Error(err))
} else {
h.logger.Info("initiator disconnected, session ended",
zap.String("session_id", truncateSessionID(c.SessionID)),
zap.String("nickname", currentNick))
}
} else {
// 参与者断开:先从 Redis 删除字段,再广播 participant_left
if err := h.service.RemoveParticipant(context.Background(), c.SessionID, currentNick); err != nil {
h.logger.Warn("participant RemoveParticipant err",
zap.String("session_id", truncateSessionID(c.SessionID)),
zap.String("nickname", currentNick),
zap.Error(err))
}
h.logger.Info("participant disconnected",
zap.String("session_id", truncateSessionID(c.SessionID)),
zap.String("nickname", currentNick))
h.hub.Broadcast(c.SessionID, response.WSParticipantLeft(currentNick))
}
c.Hub.Unregister(c) // 最后从房间移除
}()
c.WS.SetPongHandler(func(string) error { return nil })
for {
_, raw, err := c.WS.ReadMessage()
if err != nil {
return
}
// 解析客户端发来的消息
var msg struct {
Type string `json:"type"`
Data json.RawMessage `json:"data,omitempty"`
}
if err := json.Unmarshal(raw, &msg); err != nil {
continue // 不是合法 JSON,忽略
}
switch msg.Type {
case "rename":
// 只有参与者允许改昵称
if c.Role != "participant" {
h.sendError(c, "rename not allowed for role "+c.Role)
continue
}
var data struct {
NewNickname string `json:"new_nickname"`
}
if err := json.Unmarshal(msg.Data, &data); err != nil || strings.TrimSpace(data.NewNickname) == "" || utf8.RuneCountInString(data.NewNickname) > 32 {
h.sendError(c, "invalid new_nickname")
continue
}
newNick := strings.TrimSpace(data.NewNickname)
if newNick == currentNick {
continue // 无变化
}
err := h.service.UpdateParticipantNickname(context.Background(), c.SessionID, currentNick, newNick)
if err != nil {
var be *BusinessError
if errors.As(err, &be) {
h.sendError(c, be.Error()) // 优先用自定义 Message,否则回退到 ErrMessage(code)
} else {
h.sendError(c, "rename failed: "+err.Error())
}
continue
}
// 成功:更新本地昵称 + 广播给房间其他人
oldNick := currentNick
currentNick = newNick
h.hub.Broadcast(c.SessionID, response.WSParticipantRenamed(oldNick, newNick))
h.logger.Info("ws rename success",
zap.String("session_id", truncateSessionID(c.SessionID)),
zap.String("old", oldNick),
zap.String("new", newNick),
)
}
}
}
// sendError 通过 WS 给客户端发送一个 error 消息
func (h *WSHandler) sendError(c *pkgws.Conn, errMsg string) {
msg := response.WebSocketMessage{Type: "error", Error: errMsg}
select {
case c.Send <- marshalMsg(msg):
default:
}
}
func marshalMsg(m response.WebSocketMessage) []byte {
b, _ := json.Marshal(m)
return b
}
// 确保 fmt/json 被引用
var _ = fmt.Sprintf