feat(ws): WebSocket 核心模块(Hub + Client + PubSub + Handler)

- pkg/ws/message.go: 统一消息协议(Message/Response/PushMessage)
- pkg/ws/hub.go: Hub 连接管理(注册/注销/按 userID 查找/在线计数)
- pkg/ws/client.go: 客户端连接封装(readPump/writePump/心跳 30s)
- pkg/ws/pubsub.go: Redis Pub/Sub 消息路由(按用户频道发布/订阅)
- app/ws/handler.go: WebSocket 升级处理(JWT 认证 + 消息分发)
- app/ws/router.go: GET /ws 路由注册
- app/ws/provider.go: Wire Provider Set
- 更新 provider/router 集成 WebSocket 模块

Made-with: Cursor
This commit is contained in:
bujinyuan
2026-03-02 16:46:17 +08:00
parent d392653446
commit bfb4fb0ccb
13 changed files with 582 additions and 8 deletions

View File

@@ -6,8 +6,10 @@ import (
adminController "github.com/echochat/backend/app/admin/controller"
authController "github.com/echochat/backend/app/auth/controller"
"github.com/echochat/backend/app/auth/service"
wsApp "github.com/echochat/backend/app/ws"
"github.com/echochat/backend/config"
"github.com/echochat/backend/pkg/db"
"github.com/echochat/backend/pkg/ws"
"github.com/google/wire"
"github.com/redis/go-redis/v9"
"gorm.io/gorm"
@@ -22,6 +24,9 @@ type App struct {
AuthController *authController.AuthController // 前台认证控制器
AdminAuthController *authController.AdminAuthController // 后台认证控制器
UserManageController *adminController.UserManageController // 管理端用户管理控制器
WSHandler *wsApp.Handler // WebSocket 连接处理器
Hub *ws.Hub // WebSocket Hub 连接管理
PubSub *ws.PubSub // Redis Pub/Sub 消息路由
}
// NewApp 创建应用实例
@@ -33,6 +38,9 @@ func NewApp(
authCtrl *authController.AuthController,
adminAuthCtrl *authController.AdminAuthController,
userManageCtrl *adminController.UserManageController,
wsHandler *wsApp.Handler,
hub *ws.Hub,
pubsub *ws.PubSub,
) *App {
return &App{
Config: cfg,
@@ -42,6 +50,9 @@ func NewApp(
AuthController: authCtrl,
AdminAuthController: adminAuthCtrl,
UserManageController: userManageCtrl,
WSHandler: wsHandler,
Hub: hub,
PubSub: pubsub,
}
}

View File

@@ -6,6 +6,7 @@ package provider
import (
"github.com/echochat/backend/app/admin"
"github.com/echochat/backend/app/auth"
wsApp "github.com/echochat/backend/app/ws"
"github.com/echochat/backend/config"
"github.com/google/wire"
)
@@ -16,6 +17,7 @@ func InitializeApp(cfg *config.Config) (*App, error) {
InfraSet,
auth.AuthSet,
admin.AdminSet,
wsApp.WSSet,
)
return nil, nil
}

View File

@@ -7,12 +7,13 @@
package provider
import (
adminController "github.com/echochat/backend/app/admin/controller"
adminDAO "github.com/echochat/backend/app/admin/dao"
adminService "github.com/echochat/backend/app/admin/service"
controller2 "github.com/echochat/backend/app/admin/controller"
dao2 "github.com/echochat/backend/app/admin/dao"
service2 "github.com/echochat/backend/app/admin/service"
"github.com/echochat/backend/app/auth/controller"
"github.com/echochat/backend/app/auth/dao"
"github.com/echochat/backend/app/auth/service"
"github.com/echochat/backend/app/ws"
"github.com/echochat/backend/config"
"github.com/echochat/backend/pkg/db"
)
@@ -38,9 +39,12 @@ func InitializeApp(cfg *config.Config) (*App, error) {
authService := service.NewAuthService(userDAO, roleDAO, jwtConfig, tokenStore)
authController := controller.NewAuthController(authService)
adminAuthController := controller.NewAdminAuthController(authService)
userManageDAO := adminDAO.NewUserManageDAO(gormDB)
userManageService := adminService.NewUserManageService(userManageDAO, userDAO, roleDAO)
userManageController := adminController.NewUserManageController(userManageService)
app := NewApp(cfg, gormDB, client, authService, authController, adminAuthController, userManageController)
userManageDAO := dao2.NewUserManageDAO(gormDB)
userManageService := service2.NewUserManageService(userManageDAO, userDAO, roleDAO)
userManageController := controller2.NewUserManageController(userManageService)
hub := ws.ProvideHub()
pubSub := ws.ProvidePubSub(client, hub)
handler := ws.ProvideWSHandler(hub, pubSub, jwtConfig)
app := NewApp(cfg, gormDB, client, authService, authController, adminAuthController, userManageController, handler, hub, pubSub)
return app, nil
}

View File

@@ -0,0 +1,106 @@
// Package ws 提供 WebSocket 连接处理
// 负责 HTTP → WebSocket 升级、JWT 认证、消息路由分发
package ws
import (
"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
}
// NewHandler 创建 WebSocket Handler 实例
func NewHandler(hub *ws.Hub, pubsub *ws.PubSub, jwtCfg *config.JWTConfig) *Handler {
return &Handler{
hub: hub,
pubsub: pubsub,
jwtCfg: jwtCfg,
}
}
// 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)
h.hub.Register(client)
h.pubsub.Subscribe(claims.UserID)
logs.Info(nil, funcName, "WebSocket 连接建立",
zap.Int64("user_id", claims.UserID),
zap.String("ip", c.ClientIP()))
go client.WritePump()
go client.ReadPump(h.onMessage)
}
// onMessage 处理客户端发来的 WebSocket 消息
// 根据 event 类型分发到不同的处理逻辑
func (h *Handler) onMessage(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))
// Phase 2a 阶段暂无需要客户端主动发送的事件
// Phase 2b 将在此处添加 im.message.send 等事件路由
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
}

View File

@@ -0,0 +1,32 @@
package ws
import (
"github.com/echochat/backend/config"
"github.com/echochat/backend/pkg/ws"
"github.com/google/wire"
"github.com/redis/go-redis/v9"
)
// ProvideHub 创建并启动 Hub 实例
func ProvideHub() *ws.Hub {
hub := ws.NewHub()
go hub.Run()
return hub
}
// ProvidePubSub 创建 PubSub 实例
func ProvidePubSub(rdb *redis.Client, hub *ws.Hub) *ws.PubSub {
return ws.NewPubSub(rdb, hub)
}
// ProvideWSHandler 创建 WebSocket Handler
func ProvideWSHandler(hub *ws.Hub, pubsub *ws.PubSub, cfg *config.JWTConfig) *Handler {
return NewHandler(hub, pubsub, cfg)
}
// WSSet WebSocket 模块 Wire Provider Set
var WSSet = wire.NewSet(
ProvideHub,
ProvidePubSub,
ProvideWSHandler,
)

View File

@@ -0,0 +1,8 @@
package ws
import "github.com/gin-gonic/gin"
// RegisterRoutes 注册 WebSocket 相关路由
func RegisterRoutes(engine *gin.Engine, handler *Handler) {
engine.GET("/ws", handler.Upgrade)
}

View File

@@ -7,6 +7,7 @@ require (
github.com/golang-jwt/jwt/v5 v5.2.1
github.com/google/uuid v1.6.0
github.com/google/wire v0.7.0
github.com/gorilla/websocket v1.5.3
github.com/redis/go-redis/v9 v9.18.0
github.com/spf13/viper v1.21.0
go.uber.org/zap v1.27.1

View File

@@ -48,6 +48,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4=
github.com/google/wire v0.7.0/go.mod h1:n6YbUQD9cPKTnHXEBN2DXlOp/mVADhVErcMFb0v3J18=
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=

View File

@@ -0,0 +1,123 @@
package ws
import (
"encoding/json"
"time"
"github.com/echochat/backend/pkg/logs"
"github.com/gorilla/websocket"
"go.uber.org/zap"
)
const (
writeWait = 10 * time.Second // 写超时
pongWait = 60 * time.Second // 等待 pong 的超时
pingPeriod = 30 * time.Second // 心跳发送间隔(必须小于 pongWait
maxMessageSize = 4096 // 单条消息最大字节数
sendBufSize = 256 // 发送缓冲区大小
)
// Client 封装单个 WebSocket 客户端连接
// 每个连接持有两个 goroutinereadPump读取客户端消息和 writePump写入消息到客户端
type Client struct {
hub *Hub
conn *websocket.Conn
send chan []byte // 待发送消息缓冲队列
UserID int64 // 关联的用户 ID
}
// NewClient 创建客户端实例
func NewClient(hub *Hub, conn *websocket.Conn, userID int64) *Client {
return &Client{
hub: hub,
conn: conn,
send: make(chan []byte, sendBufSize),
UserID: userID,
}
}
// MessageHandler 消息处理回调函数类型
type MessageHandler func(client *Client, msg *Message)
// ReadPump 读取客户端消息的循环
// 当连接断开或出错时退出,并触发注销流程
func (c *Client) ReadPump(onMessage MessageHandler) {
defer func() {
c.hub.Unregister(c)
c.conn.Close()
}()
c.conn.SetReadLimit(maxMessageSize)
c.conn.SetReadDeadline(time.Now().Add(pongWait))
c.conn.SetPongHandler(func(string) error {
c.conn.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
for {
_, rawMsg, err := c.conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
logs.Warn(nil, "ws.client.ReadPump", "WebSocket 异常关闭",
zap.Int64("user_id", c.UserID), zap.Error(err))
}
return
}
var msg Message
if err := json.Unmarshal(rawMsg, &msg); err != nil {
logs.Warn(nil, "ws.client.ReadPump", "消息格式解析失败",
zap.Int64("user_id", c.UserID), zap.Error(err))
continue
}
if onMessage != nil {
onMessage(c, &msg)
}
}
}
// WritePump 向客户端写入消息的循环
// 从 send channel 读取消息并写入 WebSocket同时负责心跳
func (c *Client) WritePump() {
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
c.conn.Close()
}()
for {
select {
case message, ok := <-c.send:
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
if !ok {
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
logs.Warn(nil, "ws.client.WritePump", "写入消息失败",
zap.Int64("user_id", c.UserID), zap.Error(err))
return
}
case <-ticker.C:
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
// Send 向客户端发送消息(非阻塞,缓冲区满时丢弃)
func (c *Client) Send(data []byte) bool {
select {
case c.send <- data:
return true
default:
logs.Warn(nil, "ws.client.Send", "发送缓冲区已满,丢弃消息",
zap.Int64("user_id", c.UserID))
return false
}
}

View File

@@ -0,0 +1,123 @@
package ws
import (
"sync"
"github.com/echochat/backend/pkg/logs"
"go.uber.org/zap"
)
// Hub 管理所有活跃的 WebSocket 客户端连接
// 提供按 userID 注册/注销/查找连接的能力
type Hub struct {
clients map[int64]*Client // userID -> Client 映射
register chan *Client // 注册通道
unregister chan *Client // 注销通道
mu sync.RWMutex // 保护 clients map 的读写锁
}
// NewHub 创建 Hub 实例
func NewHub() *Hub {
return &Hub{
clients: make(map[int64]*Client),
register: make(chan *Client, 256),
unregister: make(chan *Client, 256),
}
}
// Run 启动 Hub 主循环,处理连接注册和注销
// 应在单独的 goroutine 中运行
func (h *Hub) Run() {
for {
select {
case client := <-h.register:
h.mu.Lock()
if old, ok := h.clients[client.UserID]; ok {
logs.Info(nil, "ws.hub.Run", "用户重复连接,关闭旧连接",
zap.Int64("user_id", client.UserID))
close(old.send)
}
h.clients[client.UserID] = client
h.mu.Unlock()
logs.Info(nil, "ws.hub.Run", "客户端已注册",
zap.Int64("user_id", client.UserID),
zap.Int("online_count", h.OnlineCount()))
case client := <-h.unregister:
h.mu.Lock()
if existing, ok := h.clients[client.UserID]; ok && existing == client {
delete(h.clients, client.UserID)
close(client.send)
}
h.mu.Unlock()
logs.Info(nil, "ws.hub.Run", "客户端已注销",
zap.Int64("user_id", client.UserID),
zap.Int("online_count", h.OnlineCount()))
}
}
}
// Register 注册客户端连接
func (h *Hub) Register(client *Client) {
h.register <- client
}
// Unregister 注销客户端连接
func (h *Hub) Unregister(client *Client) {
h.unregister <- client
}
// GetClient 根据 userID 获取客户端连接(线程安全)
func (h *Hub) GetClient(userID int64) (*Client, bool) {
h.mu.RLock()
defer h.mu.RUnlock()
client, ok := h.clients[userID]
return client, ok
}
// SendToUser 向指定用户发送消息(仅本地 Hub
// 如果用户不在本实例,返回 false
func (h *Hub) SendToUser(userID int64, data []byte) bool {
h.mu.RLock()
client, ok := h.clients[userID]
h.mu.RUnlock()
if !ok {
return false
}
select {
case client.send <- data:
return true
default:
return false
}
}
// OnlineCount 返回当前在线连接数
func (h *Hub) OnlineCount() int {
h.mu.RLock()
defer h.mu.RUnlock()
return len(h.clients)
}
// OnlineUserIDs 返回所有在线用户 ID 列表
func (h *Hub) OnlineUserIDs() []int64 {
h.mu.RLock()
defer h.mu.RUnlock()
ids := make([]int64, 0, len(h.clients))
for id := range h.clients {
ids = append(ids, id)
}
return ids
}
// IsOnline 检查指定用户是否在线(本地 Hub
func (h *Hub) IsOnline(userID int64) bool {
h.mu.RLock()
defer h.mu.RUnlock()
_, ok := h.clients[userID]
return ok
}

View File

@@ -0,0 +1,63 @@
// Package ws 提供 WebSocket 通讯基础设施
// 包含消息协议定义、连接管理、Redis Pub/Sub 消息路由
package ws
import (
"encoding/json"
"time"
)
// Message 客户端 → 服务端的 WebSocket 消息
// 事件命名规范:{模块}.{对象}.{动作},如 contact.request.send
type Message struct {
Event string `json:"event"` // 事件类型
Seq int64 `json:"seq"` // 客户端消息序号,用于 ACK 匹配
Data json.RawMessage `json:"data,omitempty"` // 业务数据(延迟解析)
Time string `json:"time"` // 发送时间
}
// Response 服务端 → 客户端的 ACK 响应
type Response struct {
Event string `json:"event"` // 事件类型(原事件 + ".ack" 后缀)
Seq int64 `json:"seq"` // 对应请求的序号
Code int `json:"code"` // 状态码0=成功
Message string `json:"message"` // 状态描述
Data interface{} `json:"data,omitempty"` // 响应数据
}
// PushMessage 服务端主动推送给客户端的消息
type PushMessage struct {
Event string `json:"event"` // 事件类型
Data interface{} `json:"data,omitempty"` // 推送数据
Time string `json:"time"` // 推送时间
}
// NewPushMessage 创建一条推送消息
func NewPushMessage(event string, data interface{}) *PushMessage {
return &PushMessage{
Event: event,
Data: data,
Time: time.Now().Format("2006-01-02 15:04:05"),
}
}
// NewResponse 创建一条 ACK 响应
func NewResponse(event string, seq int64, code int, message string, data interface{}) *Response {
return &Response{
Event: event + ".ack",
Seq: seq,
Code: code,
Message: message,
Data: data,
}
}
// MarshalPush 将推送消息序列化为 JSON 字节
func MarshalPush(msg *PushMessage) ([]byte, error) {
return json.Marshal(msg)
}
// MarshalResponse 将响应消息序列化为 JSON 字节
func MarshalResponse(resp *Response) ([]byte, error) {
return json.Marshal(resp)
}

View File

@@ -0,0 +1,98 @@
package ws
import (
"context"
"fmt"
"sync"
"github.com/echochat/backend/pkg/logs"
"github.com/redis/go-redis/v9"
"go.uber.org/zap"
)
const channelPrefix = "echo:ws:channel:"
// PubSub 封装 Redis Pub/Sub提供按用户频道的消息发布和订阅
type PubSub struct {
rdb *redis.Client
hub *Hub
subs map[int64]context.CancelFunc // userID -> 订阅取消函数
mu sync.Mutex
}
// NewPubSub 创建 PubSub 实例
func NewPubSub(rdb *redis.Client, hub *Hub) *PubSub {
return &PubSub{
rdb: rdb,
hub: hub,
subs: make(map[int64]context.CancelFunc),
}
}
// channelName 生成用户专属的 Redis 频道名
func channelName(userID int64) string {
return fmt.Sprintf("%s%d", channelPrefix, userID)
}
// Publish 向指定用户的频道发布消息
func (ps *PubSub) Publish(ctx context.Context, userID int64, data []byte) error {
channel := channelName(userID)
return ps.rdb.Publish(ctx, channel, data).Err()
}
// Subscribe 订阅指定用户的频道
// 收到消息后自动转发给本地 Hub 中对应的 Client
func (ps *PubSub) Subscribe(userID int64) {
ps.mu.Lock()
if cancel, ok := ps.subs[userID]; ok {
cancel()
}
ctx, cancel := context.WithCancel(context.Background())
ps.subs[userID] = cancel
ps.mu.Unlock()
channel := channelName(userID)
sub := ps.rdb.Subscribe(ctx, channel)
go func() {
defer sub.Close()
ch := sub.Channel()
logs.Info(nil, "ws.pubsub.Subscribe", "开始订阅用户频道",
zap.Int64("user_id", userID), zap.String("channel", channel))
for {
select {
case <-ctx.Done():
logs.Info(nil, "ws.pubsub.Subscribe", "取消订阅用户频道",
zap.Int64("user_id", userID))
return
case msg, ok := <-ch:
if !ok {
return
}
ps.hub.SendToUser(userID, []byte(msg.Payload))
}
}
}()
}
// Unsubscribe 取消订阅指定用户的频道
func (ps *PubSub) Unsubscribe(userID int64) {
ps.mu.Lock()
defer ps.mu.Unlock()
if cancel, ok := ps.subs[userID]; ok {
cancel()
delete(ps.subs, userID)
}
}
// PublishToUser 便捷方法:序列化推送消息并发布到用户频道
func (ps *PubSub) PublishToUser(ctx context.Context, userID int64, msg *PushMessage) error {
data, err := MarshalPush(msg)
if err != nil {
return fmt.Errorf("序列化推送消息失败: %w", err)
}
return ps.Publish(ctx, userID, data)
}

View File

@@ -9,6 +9,7 @@ import (
"github.com/echochat/backend/app/admin"
"github.com/echochat/backend/app/auth"
"github.com/echochat/backend/app/provider"
wsApp "github.com/echochat/backend/app/ws"
"github.com/echochat/backend/pkg/middleware"
"github.com/echochat/backend/pkg/utils"
"github.com/gin-gonic/gin"
@@ -32,9 +33,9 @@ func Setup(engine *gin.Engine, app *provider.App) {
// --- 各模块路由注册 ---
auth.RegisterRoutes(engine, app.AuthController, app.AdminAuthController, jwtAuth)
admin.RegisterRoutes(engine, app.UserManageController, jwtAuth)
wsApp.RegisterRoutes(engine, app.WSHandler)
// [未来] im.RegisterRoutes(engine, app.ImController, jwtAuth)
// [未来] meeting.RegisterRoutes(engine, app.MeetingController, jwtAuth)
// [未来] contact.RegisterRoutes(engine, app.ContactController, jwtAuth)
// [未来] notify.RegisterRoutes(engine, app.NotifyController, jwtAuth)
}