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
|