feat(websocket): integrate WebSocket support and refactor SSE handling

- Added WebSocket support by introducing a new WSHandler and related infrastructure. - Replaced SSE implementations with WebSocket equivalents across message and QR code handlers. - Updated service and repository layers to utilize WebSocket for real-time messaging. - Removed obsolete SSE hub and related code, streamlining the application for WebSocket usage. - Enhanced router to include WebSocket endpoints for real-time communication.
2026-03-26 21:17:49 +08:00
parent c6848aba06
commit 6d335e393d
19 changed files with 1277 additions and 315 deletions
--- a/internal/handler/message_handler.go
+++ b/internal/handler/message_handler.go
@@ -2,10 +2,7 @@ package handler

 import (
 	"context"
-	"fmt"
-	"net/http"
 	"strconv"
-	"time"

 	"github.com/gin-gonic/gin"

@@ -13,7 +10,7 @@ import (
 	"carrot_bbs/internal/model"
 	"carrot_bbs/internal/pkg/cursor"
 	"carrot_bbs/internal/pkg/response"
-	"carrot_bbs/internal/pkg/sse"
+	"carrot_bbs/internal/pkg/ws"
 	"carrot_bbs/internal/service"
 )

@@ -23,95 +20,17 @@ type MessageHandler struct {
 	messageService *service.MessageService
 	userService    service.UserService
 	groupService   service.GroupService
-	sseHub         *sse.Hub
+	wsHub          *ws.Hub
 }

 // NewMessageHandler 创建消息处理器
-func NewMessageHandler(chatService service.ChatService, messageService *service.MessageService, userService service.UserService, groupService service.GroupService, sseHub *sse.Hub) *MessageHandler {
+func NewMessageHandler(chatService service.ChatService, messageService *service.MessageService, userService service.UserService, groupService service.GroupService, wsHub *ws.Hub) *MessageHandler {
 	return &MessageHandler{
 		chatService:    chatService,
 		messageService: messageService,
 		userService:    userService,
 		groupService:   groupService,
-		sseHub:         sseHub,
-	}
-}
-
-// HandleSSE 实时消息订阅（SSE）
-// GET /api/v1/realtime/sse
-func (h *MessageHandler) HandleSSE(c *gin.Context) {
-	userID := c.GetString("user_id")
-	if userID == "" {
-		response.Unauthorized(c, "")
-		return
-	}
-	if h.sseHub == nil {
-		response.InternalServerError(c, "sse hub not available")
-		return
-	}
-
-	lastID := sse.ParseEventID(c.GetHeader("Last-Event-ID"))
-	if lastID == 0 {
-		lastID = sse.ParseEventID(c.Query("last_event_id"))
-	}
-	ch, cancel, replay := h.sseHub.Subscribe(userID, lastID)
-	defer cancel()
-
-	w := c.Writer
-	flusher, ok := w.(http.Flusher)
-	if !ok {
-		response.InternalServerError(c, "streaming unsupported")
-		return
-	}
-
-	w.Header().Set("Content-Type", "text/event-stream")
-	w.Header().Set("Cache-Control", "no-cache, no-transform")
-	w.Header().Set("Connection", "keep-alive")
-	w.Header().Set("X-Accel-Buffering", "no")
-	c.Status(http.StatusOK)
-	flusher.Flush()
-
-	writeEvent := func(ev sse.Event) bool {
-		data, err := sse.EncodeData(ev)
-		if err != nil {
-			return false
-		}
-		if _, err := fmt.Fprintf(w, "id: %d\nevent: %s\ndata: %s\n\n", ev.ID, ev.Event, data); err != nil {
-			return false
-		}
-		flusher.Flush()
-		return true
-	}
-
-	for _, ev := range replay {
-		if !writeEvent(ev) {
-			return
-		}
-	}
-
-	// Use a shorter heartbeat interval to survive common 50-60s proxy idle timeout.
-	heartbeat := time.NewTicker(10 * time.Second)
-	defer heartbeat.Stop()
-
-	for {
-		select {
-		case <-c.Request.Context().Done():
-			return
-		case ev, ok := <-ch:
-			if !ok || !writeEvent(ev) {
-				return
-			}
-		case <-heartbeat.C:
-			// Comment keepalive line for SSE intermediaries/proxies.
-			if _, err := fmt.Fprint(w, ": keepalive\n\n"); err != nil {
-				return
-			}
-			flusher.Flush()
-			if _, err := fmt.Fprint(w, "event: heartbeat\ndata: {}\n\n"); err != nil {
-				return
-			}
-			flusher.Flush()
-		}
+		wsHub:          wsHub,
 	}
 }

--- a/internal/handler/qrcode_handler.go
+++ b/internal/handler/qrcode_handler.go
@@ -8,7 +8,7 @@ import (
 	"github.com/gin-gonic/gin"

 	"carrot_bbs/internal/pkg/response"
-	"carrot_bbs/internal/pkg/sse"
+	"carrot_bbs/internal/pkg/ws"
 	"carrot_bbs/internal/service"
 )

@@ -44,16 +44,16 @@ func (h *QRCodeHandler) GetQRCode(c *gin.Context) {
 	})
 }

-// SSEEvents SSE事件流
+// WSEvents WebSocket事件流
 // GET /api/v1/auth/qrcode/events?session_id=xxx
-func (h *QRCodeHandler) SSEEvents(c *gin.Context) {
+func (h *QRCodeHandler) WSEvents(c *gin.Context) {
 	sessionID := c.Query("session_id")
 	if sessionID == "" {
 		response.BadRequest(c, "session_id is required")
 		return
 	}

-	ch, cancel, replay := h.qrcodeService.GetSSEHub().Subscribe(sessionID, 0)
+	ch, cancel, replay := h.qrcodeService.GetWSHub().Subscribe(sessionID, 0)
 	defer cancel()

 	w := c.Writer
@@ -70,8 +70,8 @@ func (h *QRCodeHandler) SSEEvents(c *gin.Context) {
 	c.Status(http.StatusOK)
 	flusher.Flush()

-	writeEvent := func(ev sse.Event) bool {
-		data, err := sse.EncodeData(ev)
+	writeEvent := func(ev ws.Event) bool {
+		data, err := ws.EncodeData(ev)
 		if err != nil {
 			return false
 		}
--- a/internal/handler/ws_handler.go
+++ b/internal/handler/ws_handler.go
@@ -0,0 +1,389 @@
+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()
+		close(client.Quit)
+	}()
+
+	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