golang websocket 服务端开启压缩如何写-Go服务器开发

在 Go 中，我们可以通过设置 http.Header 中的 "Sec-WebSocket-Extensions" 字段来开启 WebSocket 数据压缩。具体步骤如下：

导入相关包

import (
    "compress/flate"
    "github.com/gorilla/websocket"
    "net/http"
)

定义一个 http.HandlerFunc 处理函数，并在其中使用 websocket.Upgrader 升级 HTTP 连接为 WebSocket 连接。

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    upgrader := &websocket.Upgrader{
        ReadBufferSize:  1024,
        WriteBufferSize: 1024,
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    }
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("upgrade:", err)
        return
    }
    defer conn.Close()
    // TODO: 处理 WebSocket 消息
}

在升级连接时，设置 http.Header 中的 "Sec-WebSocket-Extensions" 字段，以开启数据压缩。压缩算法可以选择 "permessage-deflate" 或 "x-webkit-deflate-frame"。

extensions := []string{"permessage-deflate"}
responseHeader := http.Header{}
responseHeader.Set("Sec-WebSocket-Extensions", strings.Join(extensions, ", "))
upgrader.Subprotocols = []string{"chat"}
upgrader.EnableCompression = true
conn, err := upgrader.Upgrade(w, r, responseHeader)
if err != nil {
    // ...
}

这样就开启了 WebSocket 数据压缩。在处理 WebSocket 消息时，可以通过 conn.EnableWriteCompression() 和 conn.EnableReadCompression() 分别启用写入和读取的数据压缩。

if err := conn.WriteMessage(websocket.TextMessage, []byte("Hello, World!")); err != nil {
    // ...
}
conn.EnableWriteCompression(true)
if err := conn.WriteMessage(websocket.TextMessage, []byte("Hello, Compressed World!")); err != nil {
    // ...
}
conn.EnableReadCompression(true)
_, message, err := conn.ReadMessage()
if err != nil {
    // ...
}
log.Printf("received message: %s", message)

完整代码示例：

package main

import (
    "compress/flate"
    "github.com/gorilla/websocket"
    "log"
    "net/http"
)

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    upgrader := &websocket.Upgrader{
        ReadBufferSize:  1024,
        WriteBufferSize: 1024,
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    }
    extensions := []string{"permessage-deflate"}
    responseHeader := http.Header{}
    responseHeader.Set("Sec-WebSocket-Extensions", strings.Join(extensions, ", "))
    upgrader.Subprotocols = []string{"chat"}
    upgrader.EnableCompression = true
    conn, err := upgrader.Upgrade(w, r, responseHeader)
    if err != nil {
        log.Println("upgrade:", err)
        return
    }
    defer conn.Close()
    conn.EnableWriteCompression(true)
    conn.EnableReadCompression(true)
    for {
        messageType, p, err := conn.ReadMessage()
        if err != nil {
            log.Println("read:", err)
            break
        }
        log.Printf("recv: %s", p)
        if err := conn.WriteMessage(messageType, p); err != nil {
            log.Println("write:", err)
            break
        }
    }
}

func main() {
    http.HandleFunc("/ws", websocketHandler)
    log.Fatal(http.ListenAndServe(":8080", nil))
}