websocket(gorilla)

简介

Gorilla WebSocket是WebSocket协议的Go实现。可以收&发websocket数据包

如何使用

import "github.com/gorilla/websocket"

func main() {
    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        serverWs(hub, w, r)
    })
    http.ListenAndServe("127.0.0.1:8080", nil)
}

func serverWs(w http.ResponseWriter, r *http.Request) {
    upgrader := websocket.Upgrader{
        HandshakeTimeout: 0,    // 握手超时时间
        ReadBufferSize:   1024, // 读取缓冲区大小
        WriteBufferSize:  1024, // 写入缓冲区大小
        WriteBufferPool:  nil,  // 写入缓冲池
        Error:            nil,  // http错误响应函数
        CheckOrigin: func(r *http.Request) bool {
            return true
        }, // 同源检测函数
        EnableCompression: false, // 是否开启消息压缩
    }
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Println(err.Error())
        return
    }
    // 新的协程读取消息
    go readPump(conn)
    go writePump(conn)
}

// 消息读取泵
func readPump(conn *websocket.Conn) {
    defer func() {
        _ = conn.Close()
    }()
    for {
        // 读取消息
        _, message, err := conn.ReadMessage()
        if err != nil {
            // 判断ws是否断开
            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
                fmt.Println(err.Error())
            }
            break
        }
        fmt.Printf("message: %s\n", message)
    }
}

func writePump(c *websocket.Conn) {
    ticker := time.NewTicker(1 * time.Second)
    defer func() {
        _ = c.Close()
    }()
    for {
        select {
        case <-ticker.C:
            // 写入字符串数据
            _ = c.WriteMessage(websocket.TextMessage, []byte("你好"))
            // 写入json数据
            _ = c.WriteJSON(map[string]interface{}{
                "code": 200,
            })
            // 尝试向客户端发送消息
            if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
                // 发送失败后关闭客户端连接
                return
            }
        }
    }
}