9ecb29225a
- Added logging for WebSocket connection attempts, client connections, disconnections, and message receptions to improve traceability and debugging. - Implemented detailed logs in readPump to capture user and client IDs during various WebSocket events, enhancing visibility into client interactions.
424 lines
10 KiB
Go
424 lines
10 KiB
Go
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
|