ebebbbc165
- Added comments to clarify that the client.Quit channel should not be closed within the readPump function, as it is managed by Hub.Unregister. This change improves code safety and maintainability.
389 lines
9.1 KiB
Go
389 lines
9.1 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
|
||
|
||
// 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)
|
||
|
||
// 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)
|
||
}
|
||
|
||
// readPump 读取客户端消息
|
||
func (h *WSHandler) readPump(conn *websocket.Conn, client *ws.Client) {
|
||
defer func() {
|
||
conn.Close()
|
||
}()
|
||
|
||
conn.SetReadLimit(64 * 1024) // 64KB
|
||
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||
conn.SetPongHandler(func(string) error {
|
||
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||
return nil
|
||
})
|
||
|
||
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),
|
||
)
|
||
}
|
||
break
|
||
}
|
||
|
||
// 解析消息
|
||
var msg ws.Message
|
||
if err := json.Unmarshal(message, &msg); err != nil {
|
||
h.wsHub.SendError(client, "parse_error", "invalid message format")
|
||
continue
|
||
}
|
||
|
||
// 处理消息
|
||
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
|