// Package ws 提供 WebSocket 连接处理 // 负责 HTTP → WebSocket 升级、JWT 认证、消息路由分发 package ws import ( "context" "net/http" "net/url" "strings" "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, username string) bool } // OfflineMessagePusher 离线消息推送接口 // 由 im.handler.OfflinePusher 实现,WebSocket 连接建立后触发推送 type OfflineMessagePusher interface { PushOfflineMessages(ctx context.Context, userID int64) } // NotifyConnectHook 通知未读补偿钩子接口 // 由 notify.service.NotifyService 隐式实现 // WebSocket 连接建立后触发,向客户端推送 notify.unread.total 事件 // 用于断线重连场景下的徽标状态同步 type NotifyConnectHook interface { PushUnreadTotalOnConnect(ctx context.Context, userID int64) } // MeetingDisconnectHook 会议 WS 断线钩子接口(Phase 2e-2 Task 8) // 由 meeting.service.MeetingSignalService 隐式实现 // 触发时机:用户最后一条 WS 连接被移除(最终下线) // 职责:清理该用户在会议中的媒体资源;若为 host 则启动 host 宽限期 // 抽象为接口避免 ws 包反向依赖 meeting 包引发循环引用 type MeetingDisconnectHook interface { OnWSDisconnect(ctx context.Context, userID int64) } // Handler WebSocket 连接处理器 type Handler struct { hub *ws.Hub pubsub *ws.PubSub jwtCfg *config.JWTConfig serverCfg *config.ServerConfig // Task 16 Nit:CheckOrigin 白名单需要 onlineService *OnlineService tokenValidator TokenValidator offlinePusher OfflineMessagePusher notifyConnectHook NotifyConnectHook meetingDisconnectHook MeetingDisconnectHook // Task 8 注入 upgrader websocket.Upgrader } // NewHandler 创建 WebSocket Handler 实例 // Task 16 Nit:新增 serverCfg 参数,按 server.ws_allowed_origins + server.mode 收敛 CheckOrigin func NewHandler(hub *ws.Hub, pubsub *ws.PubSub, jwtCfg *config.JWTConfig, serverCfg *config.ServerConfig, onlineService *OnlineService, tokenValidator TokenValidator) *Handler { h := &Handler{ hub: hub, pubsub: pubsub, jwtCfg: jwtCfg, serverCfg: serverCfg, onlineService: onlineService, tokenValidator: tokenValidator, } h.upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: h.checkOrigin, } return h } // checkOrigin 按配置收敛 WebSocket 握手 Origin: // - 同源(Origin 为空 或 Origin.Host == Request.Host)→ 放行 // - dev 模式(server.mode != release)+ 未配置白名单 → 放行全部(便于本地开发调试) // - release 模式 / 配置了白名单 → 仅放行白名单匹配的 Origin // // 精确匹配 scheme+host+port,不做前缀/通配,避免误匹配 func (h *Handler) checkOrigin(r *http.Request) bool { origin := strings.TrimSpace(r.Header.Get("Origin")) // 空 Origin:非浏览器场景(如 curl / server-to-server),默认放行 if origin == "" { return true } u, err := url.Parse(origin) if err != nil || u.Host == "" { logs.Warn(r.Context(), "ws.handler.checkOrigin", "Origin 非法,拒绝", zap.String("origin", origin)) return false } // 同源(Origin.Host == Request.Host)直接放行 if strings.EqualFold(u.Host, r.Host) { return true } allowed := h.serverCfg.AllowedOrigins() // 未配置白名单:dev 放行,release 拒绝 if len(allowed) == 0 { if h.serverCfg.IsRelease() { logs.Warn(r.Context(), "ws.handler.checkOrigin", "release 模式未配置 WSAllowedOrigins,拒绝跨源", zap.String("origin", origin)) return false } return true } // 精确匹配白名单 for _, ao := range allowed { if strings.EqualFold(ao, origin) { return true } } logs.Warn(r.Context(), "ws.handler.checkOrigin", "Origin 不在白名单,拒绝", zap.String("origin", origin)) return false } // SetOfflinePusher 设置离线消息推送器(由 IM 模块在初始化时注入) func (h *Handler) SetOfflinePusher(pusher OfflineMessagePusher) { h.offlinePusher = pusher } // SetNotifyConnectHook 设置通知未读补偿钩子(由 notify 模块在初始化时注入) func (h *Handler) SetNotifyConnectHook(hook NotifyConnectHook) { h.notifyConnectHook = hook } // SetMeetingDisconnectHook 设置会议 WS 断线钩子(由 meeting 模块在初始化时注入) // Phase 2e-2 Task 8:WS 最终下线时触发 host 宽限期 / 媒体资源清理 func (h *Handler) SetMeetingDisconnectHook(hook MeetingDisconnectHook) { h.meetingDisconnectHook = hook } // 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, claims.Username) { logs.Warn(nil, funcName, "WebSocket Token 已失效(Redis 校验)", zap.Int64("user_id", claims.UserID)) utils.ResponseUnauthorized(c, "认证已失效,请重新登录") return } conn, err := h.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() && h.hub.IsOnline(userID) { logs.Info(nil, "ws.handler.onDisconnect", "连接被新连接替换,跳过下线清理", zap.Int64("user_id", userID)) return } h.pubsub.Unsubscribe(userID) h.onlineService.UserOffline(context.Background(), userID) // Task 8:通知会议模块处理 host 宽限期 / 媒体资源清理 // 钩子为可选(测试或 meeting 模块未加载场景安全跳过) if h.meetingDisconnectHook != nil { go h.meetingDisconnectHook.OnWSDisconnect(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)) if h.offlinePusher != nil { go h.offlinePusher.PushOfflineMessages(context.Background(), claims.UserID) } if h.notifyConnectHook != nil { go h.notifyConnectHook.PushUnreadTotalOnConnect(context.Background(), claims.UserID) } } // createReadHandler 创建带生命周期管理的消息处理函数 // 优先查 Hub 事件路由表(业务模块注册的处理器),未命中再走内置 fallback 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)) // 优先查事件路由表(IM、Meeting 等模块注册的处理器) if h.hub.DispatchEvent(client, msg) { return } // 内置事件 fallback 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: logs.Warn(nil, funcName, "未知事件类型", zap.String("event", msg.Event), zap.Int64("user_id", client.UserID)) resp := ws.NewResponse(msg.Event, msg.Seq, -1, "未知事件", 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 }