package handler import ( "context" "encoding/json" "fmt" "net/http" "strconv" "strings" "sync/atomic" "time" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "go.uber.org/zap" "carrot_bbs/internal/dto" "carrot_bbs/internal/model" "carrot_bbs/internal/pkg/response" "carrot_bbs/internal/pkg/ws" "carrot_bbs/internal/service" ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 4096, CheckOrigin: func(r *http.Request) bool { return true // 允许所有来源,实际生产环境应该限制 }, HandshakeTimeout: 10 * time.Second, } // WSHandler WebSocket处理器 type WSHandler struct { wsHub *ws.Hub chatService service.ChatService groupService service.GroupService jwtService *service.JWTService clientSeq uint64 } // NewWSHandler 创建WebSocket处理器 func NewWSHandler( wsHub *ws.Hub, chatService service.ChatService, groupService service.GroupService, jwtService *service.JWTService, ) *WSHandler { return &WSHandler{ wsHub: wsHub, chatService: chatService, groupService: groupService, jwtService: jwtService, } } // HandleWebSocket WebSocket连接处理 // GET /api/v1/realtime/ws?token=xxx func (h *WSHandler) HandleWebSocket(c *gin.Context) { // 1. 认证 - 支持Query参数和Header两种方式 token := c.Query("token") if token == "" { authHeader := c.GetHeader("Authorization") if authHeader != "" { parts := strings.SplitN(authHeader, " ", 2) if len(parts) == 2 && parts[0] == "Bearer" { token = parts[1] } } } if token == "" { response.Unauthorized(c, "token is required") return } claims, err := h.jwtService.ParseToken(token) if err != nil { response.Unauthorized(c, "invalid token") return } userID := claims.UserID zap.L().Info("WebSocket connection attempt", zap.String("user_id", userID), ) // 2. 升级HTTP连接为WebSocket conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { zap.L().Error("Failed to upgrade WebSocket connection", zap.String("user_id", userID), zap.Error(err), ) return } // 3. 创建客户端 clientID := atomic.AddUint64(&h.clientSeq, 1) client := &ws.Client{ ID: clientID, UserID: userID, Send: make(chan []byte, defaultUserBufferSize), Quit: make(chan struct{}), } // 4. 注册客户端 replayEvents := h.wsHub.Register(client) defer h.wsHub.Unregister(client) zap.L().Info("WebSocket client connected", zap.String("user_id", userID), zap.Uint64("client_id", clientID), ) // 5. 发送历史回放消息 go func() { for _, ev := range replayEvents { msg := ws.ResponseMessage{ EventID: ev.ID, Type: ev.Type, TS: ev.TS, Payload: ev.Payload, } data, _ := json.Marshal(msg) select { case <-client.Quit: return case client.Send <- data: } } }() // 6. 启动读写goroutine go h.writePump(conn, client) h.readPump(conn, client) zap.L().Info("WebSocket client disconnected", zap.String("user_id", userID), zap.Uint64("client_id", clientID), ) } // readPump 读取客户端消息 func (h *WSHandler) readPump(conn *websocket.Conn, client *ws.Client) { defer func() { zap.L().Info("WebSocket readPump exiting", zap.String("user_id", client.UserID), zap.Uint64("client_id", client.ID), ) conn.Close() }() conn.SetReadLimit(64 * 1024) // 64KB conn.SetReadDeadline(time.Now().Add(60 * time.Second)) conn.SetPongHandler(func(string) error { zap.L().Debug("WebSocket pong received", zap.String("user_id", client.UserID), ) conn.SetReadDeadline(time.Now().Add(60 * time.Second)) return nil }) zap.L().Info("WebSocket readPump started", zap.String("user_id", client.UserID), zap.Uint64("client_id", client.ID), ) for { _, message, err := conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { zap.L().Debug("WebSocket read error", zap.String("user_id", client.UserID), zap.Error(err), ) } else { zap.L().Info("WebSocket connection closed", zap.String("user_id", client.UserID), zap.Error(err), ) } break } // 解析消息 var msg ws.Message if err := json.Unmarshal(message, &msg); err != nil { h.wsHub.SendError(client, "parse_error", "invalid message format") continue } zap.L().Debug("WebSocket message received", zap.String("user_id", client.UserID), zap.String("type", msg.Type), ) // 处理消息 h.handleMessage(client, &msg) } } // writePump 向客户端发送消息 func (h *WSHandler) writePump(conn *websocket.Conn, client *ws.Client) { ticker := time.NewTicker(30 * time.Second) defer func() { ticker.Stop() conn.Close() }() for { select { case <-client.Quit: conn.WriteMessage(websocket.CloseMessage, []byte{}) return case message, ok := <-client.Send: conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if !ok { conn.WriteMessage(websocket.CloseMessage, []byte{}) return } w, err := conn.NextWriter(websocket.TextMessage) if err != nil { return } w.Write(message) // 批量发送缓冲区中的消息 n := len(client.Send) for i := 0; i < n; i++ { w.Write([]byte{'\n'}) w.Write(<-client.Send) } if err := w.Close(); err != nil { return } case <-ticker.C: conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } // handleMessage 处理客户端消息 func (h *WSHandler) handleMessage(client *ws.Client, msg *ws.Message) { ctx := context.Background() switch msg.Type { case "ping": h.handlePing(client) case "chat": h.handleChat(ctx, client, msg.Payload) case "typing": h.handleTyping(ctx, client, msg.Payload) case "read": h.handleRead(ctx, client, msg.Payload) case "recall": h.handleRecall(ctx, client, msg.Payload) default: h.wsHub.SendError(client, "unknown_type", fmt.Sprintf("unknown message type: %s", msg.Type)) } } // handlePing 处理ping消息 func (h *WSHandler) handlePing(client *ws.Client) { pong := ws.ResponseMessage{ EventID: h.wsHub.NextID(), Type: "pong", TS: time.Now().UnixMilli(), } data, _ := json.Marshal(pong) select { case <-client.Quit: case client.Send <- data: } } // handleChat 处理聊天消息 func (h *WSHandler) handleChat(ctx context.Context, client *ws.Client, payload json.RawMessage) { var req struct { ConversationID string `json:"conversation_id"` DetailType string `json:"detail_type"` Segments model.MessageSegments `json:"segments"` ReplyToID *string `json:"reply_to_id,omitempty"` } if err := json.Unmarshal(payload, &req); err != nil { h.wsHub.SendError(client, "parse_error", "invalid chat message format") return } if req.ConversationID == "" { h.wsHub.SendError(client, "invalid_params", "conversation_id is required") return } if len(req.Segments) == 0 { h.wsHub.SendError(client, "invalid_params", "segments is required") return } // 发送消息 message, err := h.chatService.SendMessage(ctx, client.UserID, req.ConversationID, req.Segments, req.ReplyToID) if err != nil { h.wsHub.SendError(client, "send_failed", err.Error()) return } // 发送成功响应 detailType := "private" if conv, _ := h.chatService.GetConversationByID(ctx, req.ConversationID, client.UserID); conv != nil && conv.Type == model.ConversationTypeGroup { detailType = "group" } resp := dto.WSEventResponse{ ID: message.ID, Time: message.CreatedAt.UnixMilli(), Type: "message", DetailType: detailType, Seq: strconv.FormatInt(message.Seq, 10), Segments: req.Segments, SenderID: client.UserID, } msg := ws.ResponseMessage{ EventID: h.wsHub.NextID(), Type: "message_sent", TS: time.Now().UnixMilli(), Payload: resp, } data, _ := json.Marshal(msg) select { case <-client.Quit: case client.Send <- data: } } // handleTyping 处理输入状态 func (h *WSHandler) handleTyping(ctx context.Context, client *ws.Client, payload json.RawMessage) { var req struct { ConversationID string `json:"conversation_id"` IsTyping bool `json:"is_typing"` } if err := json.Unmarshal(payload, &req); err != nil { return } if req.ConversationID == "" { return } // 输入状态通常不需要持久化,直接通过ChatService的SendTyping处理 // 这里简化处理,直接推送 h.chatService.SendTyping(ctx, client.UserID, req.ConversationID) } // handleRead 处理已读标记 func (h *WSHandler) handleRead(ctx context.Context, client *ws.Client, payload json.RawMessage) { var req struct { ConversationID string `json:"conversation_id"` Seq int64 `json:"seq"` } if err := json.Unmarshal(payload, &req); err != nil { return } if req.ConversationID == "" || req.Seq <= 0 { return } err := h.chatService.MarkAsRead(ctx, req.ConversationID, client.UserID, req.Seq) if err != nil { zap.L().Debug("Failed to mark as read via WebSocket", zap.String("user_id", client.UserID), zap.String("conversation_id", req.ConversationID), zap.Error(err), ) } } // handleRecall 处理消息撤回 func (h *WSHandler) handleRecall(ctx context.Context, client *ws.Client, payload json.RawMessage) { var req struct { MessageID string `json:"message_id"` } if err := json.Unmarshal(payload, &req); err != nil { h.wsHub.SendError(client, "parse_error", "invalid recall message format") return } if req.MessageID == "" { h.wsHub.SendError(client, "invalid_params", "message_id is required") return } err := h.chatService.RecallMessage(ctx, req.MessageID, client.UserID) if err != nil { h.wsHub.SendError(client, "recall_failed", err.Error()) return } // 发送撤回成功响应 resp := ws.ResponseMessage{ EventID: h.wsHub.NextID(), Type: "message_recalled", TS: time.Now().UnixMilli(), Payload: map[string]string{"message_id": req.MessageID}, } data, _ := json.Marshal(resp) select { case <-client.Quit: case client.Send <- data: } } const defaultUserBufferSize = 128