如何使用Golang读取WebSocket消息的前X字节？

我有如下WebSocket代码：

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
}

func handleConnections(w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Fatal(err)
    }
    defer ws.Close()

    for {
        messageType, message, err := ws.ReadMessage()
         // 上面的代码可以正常运行...
       // 我希望仅读取消息的前20字节，而不是完整消息
      
    }
}

目前我可以通过ws.ReadMessage()读取完整消息，但希望仅读取消息的前20字节而非全部，该如何实现？

要实现仅读取WebSocket消息的前20字节，你可以使用websocket.Conn的NextReader()方法替代ReadMessage()。ReadMessage()会一次性读取完整消息并返回字节切片，而NextReader()会返回一个可以逐部分读取消息的io.Reader，这样你就能精确控制读取的字节数。

具体实现步骤：

• 调用ws.NextReader()获取消息类型和对应的读取器
• 创建长度为20的字节切片，使用io.ReadFull()尝试读取满20字节（若消息长度不足，会返回io.ErrUnexpectedEOF，可按需处理）
• 读取完前20字节后，必须丢弃剩余的消息内容，避免残留数据影响后续消息读取

修改后的完整代码：

import (
    "io"
    "log"
    "net/http"

    "github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
}

func handleConnections(w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Fatal(err)
    }
    defer ws.Close()

    for {
        messageType, reader, err := ws.NextReader()
        if err != nil {
            log.Println("获取消息读取器失败:", err)
            break
        }

        // 准备存储前20字节的缓冲区
        first20Bytes := make([]byte, 20)
        readCount, err := io.ReadFull(reader, first20Bytes)
        if err != nil {
            if err == io.ErrUnexpectedEOF {
                // 消息长度不足20字节的情况
                log.Printf("消息长度不足20字节，实际读取%d字节", readCount)
                // 使用first20Bytes[:readCount]获取实际读取的内容
                processedContent := first20Bytes[:readCount]
                log.Println("处理内容:", processedContent)
            } else {
                log.Println("读取消息失败:", err)
                break
            }
        } else {
            log.Printf("成功读取前20字节: %s", first20Bytes)
            // 这里可以处理读取到的20字节内容
        }

        // 丢弃剩余的消息内容，避免影响下一次读取
        if _, err := io.Copy(io.Discard, reader); err != nil {
            log.Println("丢弃剩余消息失败:", err)
            break
        }
    }
}

关键注意事项：

• 必须处理剩余消息：如果不丢弃未读取的部分，这些数据会留在WebSocket的缓冲区中，导致下一次读取的消息内容混乱
• 错误处理：针对消息长度不足的场景，io.ErrUnexpectedEOF是预期错误，可根据业务需求决定是否保留部分读取的内容
• 消息类型兼容：messageType与ReadMessage()返回的类型一致，支持文本和二进制消息的处理

内容的提问来源于stack exchange，提问作者Alexander Mills