Gin 框架如何实现 WebSocket？Hub 模式与连接管理详解

Gin 本身不自带 WebSocket

先说清楚一件事：Gin 是 HTTP 框架，WebSocket 是另一个协议。Gin 能做的只是把 HTTP 请求接住，剩下的升级握手和长连接管理得交给专门的库。Go 生态里最成熟的选择是 gorilla/websocket，虽然 gorilla 组织已归档，但这个库仍被广泛使用且稳定。

bash
go get github.com/gorilla/websocket

最小可用的 WebSocket 端点

三步走：创建 Upgrader → 升级连接 → 读写消息。

go
var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    // 开发阶段允许所有来源，生产环境必须限制
    CheckOrigin: func(r *http.Request) bool { return true },
}

func handleWS(c *gin.Context) {
    conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
    if err != nil {
        log.Printf("升级失败: %v", err)
        return
    }
    defer conn.Close()

    for {
        msgType, msg, err := conn.ReadMessage()
        if err != nil {
            // 客户端断开或连接异常，退出循环即可
            break
        }
        log.Printf("收到: %s", msg)
        conn.WriteMessage(msgType, msg) // echo 回去
    }
}

func main() {
    r := gin.Default()
    r.GET("/ws", handleWS)
    r.Run(":8080")
}

这段代码能跑，但不能上生产——没有连接管理、没有并发控制、客户端断了你都不知道。

生产级架构：Hub + Client 模式

单连接玩玩可以，真正的 WebSocket 服务需要管理成百上千个连接。经典的模式是用 Hub 集中管理所有 Client，Client 各自负责自己的读写：

go
type Client struct {
    conn *websocket.Conn
    send chan []byte
    hub  *Hub
}

type Hub struct {
    clients    map[*Client]bool
    broadcast  chan []byte
    register   chan *Client
    unregister chan *Client
}

func newHub() *Hub {
    return &Hub{
        clients:    make(map[*Client]bool),
        broadcast:  make(chan []byte),
        register:   make(chan *Client),
        unregister: make(chan *Client),
    }
}

func (h *Hub) run() {
    for {
        select {
        case c := <-h.register:
            h.clients[c] = true
        case c := <-h.unregister:
            if _, ok := h.clients[c]; ok {
                delete(h.clients, c)
                close(c.send)
            }
        case msg := <-h.broadcast:
            for c := range h.clients {
                select {
                case c.send <- msg:
                default:
                    // send 满了说明客户端卡住了，直接踢掉
                    close(c.send)
                    delete(h.clients, c)
                }
            }
        }
    }
}

Hub 用 channel 而不是 mutex 来管理状态，是因为 register/unregister/broadcast 三个操作天然是事件驱动的，select 比加锁更清晰，也避免了死锁风险。

读写分离：readPump 和 writePump

每个 Client 需要两个 goroutine：一个专门读，一个专门写。为什么分开？因为 conn.ReadMessage() 是阻塞调用，和 conn.WriteMessage() 放在同一个 goroutine 里会互相卡。

go
func (c *Client) readPump() {
    defer func() {
        c.hub.unregister <- c
        c.conn.Close()
    }()

    c.conn.SetReadLimit(512)              // 限制单条消息大小
    c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
    c.conn.SetPongHandler(func(string) error {
        c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
        return nil
    })

    for {
        _, msg, err := c.conn.ReadMessage()
        if err != nil {
            break
        }
        c.hub.broadcast <- msg
    }
}

func (c *Client) writePump() {
    ticker := time.NewTicker(30 * time.Second)
    defer func() {
        ticker.Stop()
        c.conn.Close()
    }()

    for {
        select {
        case msg, ok := <-c.send:
            c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
            if !ok {
                c.conn.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }
            if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
                return
            }
        case <-ticker.C:
            c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
            if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                return
            }
        }
    }
}

几个关键细节：

• SetReadLimit 防止恶意客户端发超大消息撑爆内存
• SetReadDeadline + PongHandler 实现：60 秒内没收到任何消息就断开
• writePump 里的 ticker 每 30 秒发一次 Ping，客户端不回 Pong 就会被 readPump 的超时机制踢掉
• send channel 满了（default 分支），说明客户端消费不过来，直接断开

在 Gin 里串起来

go
func serveWS(hub *Hub) gin.HandlerFunc {
    return func(c *gin.Context) {
        conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
        if err != nil {
            return
        }
        client := &Client{
            conn: conn,
            send: make(chan []byte, 256),
            hub:  hub,
        }
        hub.register <- client

        go client.writePump()
        go client.readPump()
    }
}

func main() {
    hub := newHub()
    go hub.run()

    r := gin.Default()
    r.GET("/ws", serveWS(hub))
    r.Run(":8080")
}

注意 serveWS 返回 gin.HandlerFunc，这样 Hub 作为闭包变量传入，不用全局变量。

认证怎么做

WebSocket 握手是 HTTP GET 请求，所以在升级之前做认证就行：

go
func authMiddleware() gin.HandlerFunc {
    return func(c *gin.Context) {
        token := c.Query("token")
        if token == "" {
            // WebSocket 不能返回 JSON，用 HTTP 状态码拒绝
            c.AbortWithStatus(http.StatusUnauthorized)
            return
        }
        claims, err := parseToken(token)
        if err != nil {
            c.AbortWithStatus(http.StatusUnauthorized)
            return
        }
        c.Set("userID", claims.UserID)
        c.Next()
    }
}

// 路由
r.GET("/ws", authMiddleware(), serveWS(hub))

客户端连接时带 token：new WebSocket('ws://localhost:8080/ws?token=xxx')。

不要把 token 放在 URL path 里（如 /ws/:token），URL 会被记录到访问日志和浏览器历史里，有泄露风险。Query parameter 稍好，但最安全的方案是先通过 HTTP 接口换取一次性 ticket，再用 ticket 连 WebSocket。

gorilla/websocket 已归档，怎么办

gorilla 组织在 2022 年底归档了所有项目。gorilla/websocket 目前还能用，但不再有新功能更新。替代方案：

• nhooyr.io/websocket：更现代的 API，支持 context 取消，API 更简洁
• gobwas/ws：零拷贝升级，性能更好，但 API 更底层
• codenotary/websocket：gorilla/websocket 的社区 fork，持续维护

如果你的项目是新开始的，建议直接用 nhooyr.io/websocket，API 更干净。已有的项目不用急着迁移，gorilla/websocket 稳定且经过了大量生产验证。