EchoChat/backend/go-service/app/ws/handler.go

// Package ws 提供 WebSocket 连接处理
// 负责 HTTP → WebSocket 升级、JWT 认证、消息路由分发
package ws

import (
	"context"
	"net/http"

	"github.com/echochat/backend/config"
	"github.com/echochat/backend/pkg/logs"
	"github.com/echochat/backend/pkg/utils"
	"github.com/echochat/backend/pkg/ws"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
)

// TokenValidator 有状态 JWT 验证接口（检查 Token 是否在 Redis 中有效）
// 由 auth.AuthService 实现，用于防止已登出用户建立 WebSocket 连接
type TokenValidator interface {
	ValidateAccessToken(ctx context.Context, userID int64, clientType, token string) bool
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // TODO: 生产环境通过配置限制 allowed origins
	},
}

// Handler WebSocket 连接处理器
type Handler struct {
	hub            *ws.Hub
	pubsub         *ws.PubSub
	jwtCfg         *config.JWTConfig
	onlineService  *OnlineService
	tokenValidator TokenValidator
}

// NewHandler 创建 WebSocket Handler 实例
func NewHandler(hub *ws.Hub, pubsub *ws.PubSub, jwtCfg *config.JWTConfig, onlineService *OnlineService, tokenValidator TokenValidator) *Handler {
	return &Handler{
		hub:            hub,
		pubsub:         pubsub,
		jwtCfg:         jwtCfg,
		onlineService:  onlineService,
		tokenValidator: tokenValidator,
	}
}

// Upgrade 处理 WebSocket 升级请求
// GET /ws?token=xxx → JWT 认证 → 升级连接 → 注册 Hub → 订阅 Redis 频道
func (h *Handler) Upgrade(c *gin.Context) {
	funcName := "ws.handler.Upgrade"

	token := c.Query("token")
	if token == "" {
		utils.ResponseUnauthorized(c, "缺少认证 Token")
		return
	}

	claims, err := utils.ParseToken(h.jwtCfg, token)
	if err != nil {
		logs.Warn(nil, funcName, "WebSocket Token 验证失败", zap.Error(err))
		utils.ResponseUnauthorized(c, "Token 无效或已过期")
		return
	}

	clientType := claims.ClientType
	if clientType == "" {
		clientType = "frontend"
	}
	if h.tokenValidator != nil && !h.tokenValidator.ValidateAccessToken(c.Request.Context(), claims.UserID, clientType, token) {
		logs.Warn(nil, funcName, "WebSocket Token 已失效（Redis 校验）",
			zap.Int64("user_id", claims.UserID))
		utils.ResponseUnauthorized(c, "认证已失效，请重新登录")
		return
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logs.Error(nil, funcName, "WebSocket 升级失败",
			zap.Int64("user_id", claims.UserID), zap.Error(err))
		return
	}

	client := ws.NewClient(h.hub, conn, claims.UserID)
	client.SetOnDisconnect(func(userID int64) {
		if client.IsClosedByHub() {
			logs.Info(nil, "ws.handler.onDisconnect", "连接被 Hub 踢出（重复连接），跳过下线清理",
				zap.Int64("user_id", userID))
			return
		}
		h.pubsub.Unsubscribe(userID)
		h.onlineService.UserOffline(context.Background(), userID)
	})
	h.hub.Register(client)
	h.pubsub.Subscribe(claims.UserID)
	h.onlineService.UserOnline(c.Request.Context(), claims.UserID, c.ClientIP())

	logs.Info(nil, funcName, "WebSocket 连接建立",
		zap.Int64("user_id", claims.UserID),
		zap.String("ip", c.ClientIP()))

	go client.WritePump()
	go client.ReadPump(h.createReadHandler(claims.UserID))
}

// createReadHandler 创建带生命周期管理的消息处理函数
// 客户端断开时自动执行下线清理
func (h *Handler) createReadHandler(userID int64) ws.MessageHandler {
	return func(client *ws.Client, msg *ws.Message) {
		funcName := "ws.handler.onMessage"
		logs.Debug(nil, funcName, "收到 WebSocket 消息",
			zap.Int64("user_id", client.UserID),
			zap.String("event", msg.Event),
			zap.Int64("seq", msg.Seq))

		switch msg.Event {
		case "heartbeat":
			h.onlineService.HeartbeatRenew(context.Background(), userID)
			resp := ws.NewResponse(msg.Event, msg.Seq, 0, "pong", nil)
			data, err := ws.MarshalResponse(resp)
			if err != nil {
				logs.Error(nil, funcName, "序列化心跳响应失败", zap.Error(err))
				return
			}
			client.Send(data)
		default:
			resp := ws.NewResponse(msg.Event, msg.Seq, 0, "ok", nil)
			data, err := ws.MarshalResponse(resp)
			if err != nil {
				logs.Error(nil, funcName, "序列化响应失败", zap.Error(err))
				return
			}
			client.Send(data)
		}
	}
}

// GetHub 返回 Hub 实例（供在线状态等模块访问）
func (h *Handler) GetHub() *ws.Hub {
	return h.hub
}

// GetPubSub 返回 PubSub 实例（供业务模块发送推送）
func (h *Handler) GetPubSub() *ws.PubSub {
	return h.pubsub
}