1c87a466a1
- app/ws/online_service.go: 上线/下线/心跳续期/批量查询在线状态
- 集成到 WebSocket handler: 连接时写入 Redis, 断线时清除
- client.go: 添加 DisconnectHandler 回调
- Redis 键: echo:user:online (SET) + echo:user:status:{uid} (TTL 60s)
Made-with: Cursor
123 lines
3.4 KiB
Go
123 lines
3.4 KiB
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"
|
|
)
|
|
|
|
var upgrader = websocket.Upgrader{
|
|
ReadBufferSize: 1024,
|
|
WriteBufferSize: 1024,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
return true // 开发阶段允许所有来源，生产环境需限制
|
|
},
|
|
}
|
|
|
|
// Handler WebSocket 连接处理器
|
|
type Handler struct {
|
|
hub *ws.Hub
|
|
pubsub *ws.PubSub
|
|
jwtCfg *config.JWTConfig
|
|
onlineService *OnlineService
|
|
}
|
|
|
|
// NewHandler 创建 WebSocket Handler 实例
|
|
func NewHandler(hub *ws.Hub, pubsub *ws.PubSub, jwtCfg *config.JWTConfig, onlineService *OnlineService) *Handler {
|
|
return &Handler{
|
|
hub: hub,
|
|
pubsub: pubsub,
|
|
jwtCfg: jwtCfg,
|
|
onlineService: onlineService,
|
|
}
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
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) {
|
|
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, _ := ws.MarshalResponse(resp)
|
|
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
|
|
}
|