玩转 Go 生态｜Hertz WebSocket 扩展简析

WebSocket 是一种可以在单个 TCP 连接上进行全双工通信，位于 OSI 模型的应用层。WebSocket 使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。在 WebSocket API 中，浏览器和服务器只需要完成一次握手，两者之间就可以创建持久性的连接，并进行双向数据传输。

Hertz 提供了 WebSocket 的支持，参考 gorilla/websocket 库使用 hijack 的方式在 Hertz 进行了适配，用法和参数基本保持一致。

安装

go get github.com/hertz-contrib/websocket

示例代码

package main
​
import (
    "context"
    "flag"
    "html/template"
    "log"
​
    "github.com/cloudwego/hertz/pkg/app"
    "github.com/cloudwego/hertz/pkg/app/server"
    "github.com/hertz-contrib/websocket"
)
​
var addr = flag.String("addr", "localhost:8080", "http service address")
​
var upgrader = websocket.HertzUpgrader{} // use default options
​
func echo(_ context.Context, c *app.RequestContext) {
    err := upgrader.Upgrade(c, func(conn *websocket.Conn) {
        for {
            mt, message, err := conn.ReadMessage()
            if err != nil {
                log.Println("read:", err)
                break
            }
            log.Printf("recv: %s", message)
            err = conn.WriteMessage(mt, message)
            if err != nil {
                log.Println("write:", err)
                break
            }
        }
    })
    if err != nil {
        log.Print("upgrade:", err)
        return
    }
}
​
func home(_ context.Context, c *app.RequestContext) {
    c.SetContentType("text/html; charset=utf-8")
    homeTemplate.Execute(c, "ws://"+string(c.Host())+"/echo")
}
​
func main() {
    flag.Parse()
    h := server.Default(server.WithHostPorts(*addr))
    // https://github.com/cloudwego/hertz/issues/121
    h.NoHijackConnPool = true
    h.GET("/", home)
    h.GET("/echo", echo)
    h.Spin()
}
​
var homeTemplate = template.Must(template.New("").Parse(`









Click "Open" to create a connection to the server, 
"Send" to send a message to the server and "Close" to close the connection. 
You can change the message and send multiple times.


Open
Close

Send






`))

运行 server：

go run server.go

上述示例代码中，服务器包括一个简单的网络客户端。要使用该客户端，在浏览器中打开 http://127.0.0.1:8080，并按照页面上的指示操作。

Upgrade

websocket.Conn 类型代表一个 WebSocket 连接。服务器应用程序从 HTTP 请求处理程序中调用 HertzUpgrader.Upgrade 方法，将 HTTP 协议的连接请求升级为 WebSocket 协议的连接请求。

这部分逻辑对应着示例代码的 echo() 函数，此处着重介绍 HertzUpgrader.Upgrade。

函数签名：

func (u *HertzUpgrader) Upgrade(ctx *app.RequestContext, handler HertzHandler) error

内部处理逻辑：

func (u *HertzUpgrader) Upgrade(ctx *app.RequestContext, handler HertzHandler) error {
    if !ctx.IsGet() {
        return u.returnError(ctx, consts.StatusMethodNotAllowed, fmt.Sprintf("%s request method is not GET", badHandshake))
    }
    // 校验 requsetHeader 中与 websocket 相关的字段（此处省略部分逻辑代码）
​
    subprotocol := u.selectSubprotocol(ctx)
    compress := u.isCompressionEnable(ctx)
​
    ctx.SetStatusCode(consts.StatusSwitchingProtocols)
    // 构造协议升级后的响应头部信息
    ctx.Response.Header.Set("Upgrade", "websocket")
    ctx.Response.Header.Set("Connection", "Upgrade")
    ctx.Response.Header.Set("Sec-WebSocket-Accept", computeAcceptKeyBytes(challengeKey))
    // “无上下文接管”模式
    if compress {
        ctx.Response.Header.Set("Sec-WebSocket-Extensions", "permessage-deflate; server_no_context_takeover; client_no_context_takeover")
    }
    if subprotocol != nil {
        ctx.Response.Header.SetBytesV("Sec-WebSocket-Protocol", subprotocol)
    }
​
    // 通过 Hijack 的方式，实现 websocket 全双工的通信
    ctx.Hijack(func(netConn network.Conn) {
        writeBuf := poolWriteBuffer.Get().([]byte)
        c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize, u.WriteBufferPool, nil, writeBuf)
        if subprotocol != nil {
            c.subprotocol = b2s(subprotocol)
        }
​
        if compress {
            c.newCompressionWriter = compressNoContextTakeover
            c.newDecompressionReader = decompressNoContextTakeover
        }
​
        netConn.SetDeadline(time.Time{})
​
        handler(c)
​
        writeBuf = writeBuf[0:0]
        poolWriteBuffer.Put(writeBuf)
    })
​
    return nil
}

HertzHandler

HertzHandler 是上述 HertzUpgrader.Upgrade 函数的第二个参数。HertzHandler 在握手完成后接收一个 websocket 连接，通过劫持这个连接，完成全双工的通信。

HertzHandler 必须由用户提供，内部定义了 WebSocket 请求和响应的具体流程。

函数签名：

type HertzHandler func(*Conn)

上述 echo 服务器的 websocket 处理流程：

err := upgrader.Upgrade(c, func(conn *websocket.Conn) {
    for {
        // 读取客户端发送的信息
        mt, message, err := conn.ReadMessage()
        if err != nil {
            log.Println("read:", err)
            break
        }
        log.Printf("recv: %s", message)
        // 向客户端发送信息
        err = conn.WriteMessage(mt, message)
        if err != nil {
            log.Println("write:", err)
            break
        }
    }
})

配置

上述文档已经讲述了Hertz WebSocket 最核心的协议升级与连接劫持的逻辑，下面将罗列 Hertz WebSocket 使用过程中可选的配置参数。

这部分将围绕 websocket.HertzUpgrader 结构展开说明。

参数	介绍
ReadBufferSize	用于设置输入缓冲区的大小，单位为字节。如果缓冲区大小为零，那么就使用 HTTP 服务器分配的大小。输入缓冲区大小并不限制可以接收的信息的大小。
WriteBufferSize	用于设置输出缓冲区的大小，单位为字节。如果缓冲区大小为零，那么就使用 HTTP 服务器分配的大小。输出缓冲区大小并不限制可以发送的信息的大小。
WriteBufferPool	用于设置写操作的缓冲池。
Subprotocols	用于按优先顺序设置服务器支持的协议。如果这个字段不是 nil，那么 Upgrade 方法通过选择这个列表中与客户端请求的协议的第一个匹配来协商一个子协议。如果没有匹配，那么就不协商协议（Sec-Websocket-Protocol 头不包括在握手响应中）。
Error	用于设置生成 HTTP 错误响应的函数。
CheckOrigin	用于设置针对请求的 Origin 头的校验函数， 如果请求的 Origin 头是可接受的，CheckOrigin 返回 true。
EnableCompression	用于设置服务器是否应该尝试协商每个消息的压缩（RFC 7692）。将此值设置为 true 并不能保证压缩会被支持。

WriteBufferPool

如果该值没有被设置，则额外初始化写缓冲区，并在当前生命周期内分配给该连接。当应用程序在大量的连接上有适度的写入量时，缓冲池是最有用的。

应用程序应该使用一个单一的缓冲池来为不同的连接分配缓冲区。

接口签名：

// BufferPool represents a pool of buffers. The *sync.Pool type satisfies this
// interface.  The type of the value stored in a pool is not specified.
type BufferPool interface {
    // Get gets a value from the pool or returns nil if the pool is empty.
    Get() interface{}
    // Put adds a value to the pool.
    Put(interface{})
}

示例代码：

type simpleBufferPool struct {
    v interface{}
}
​
func (p *simpleBufferPool) Get() interface{} {
    v := p.v
    p.v = nil
    return v
}
​
func (p *simpleBufferPool) Put(v interface{}) {
    p.v = v
}
​
var upgrader = websocket.HertzUpgrader{
    WriteBufferPool: &simpleBufferPool{},
}

Subprotocols

WebSocket 只是定义了一种交换任意消息的机制。这些消息是什么意思，客户端在任何特定的时间点可以期待什么样的消息，或者他们被允许发送什么样的消息，完全取决于实现应用程序。

所以你需要在服务器和客户端之间就这些事情达成协议。子协议参数只是让客户端和服务端正式地交换这些信息。你可以为你想要的任何协议编造任何名字。服务器可以简单地检查客户在握手过程中是否遵守了该协议。

Error

如果 Error 为 nil，则使用 Hertz 提供的 API 来生成 HTTP 错误响应。

函数签名：

func(ctx *app.RequestContext, status int, reason error)

示例代码：

var upgrader = websocket.HertzUpgrader{
    Error: func(ctx *app.RequestContext, status int, reason error) {
        ctx.Response.Header.Set("Sec-Websocket-Version", "13")
        ctx.AbortWithMsg(reason.Error(), status)
    },
}

CheckOrigin

如果 CheckOrigin 为nil，则使用一个安全的默认值：如果Origin请求头存在，并且源主机不等于请求主机头，则返回false。CheckOrigin 函数应该仔细验证请求的来源，以防止跨站请求伪造。

函数签名：

func(ctx *app.RequestContext) bool

默认实现：

func fastHTTPCheckSameOrigin(ctx *app.RequestContext) bool {
    origin := ctx.Request.Header.Peek("Origin")
    if len(origin) == 0 {
        return true
    }
    u, err := url.Parse(b2s(origin))
    if err != nil {
        return false
    }
    return equalASCIIFold(u.Host, b2s(ctx.Host()))
}

EnableCompression

服务端接受一个或者多个扩展字段，这些扩展字段是包含客户端请求的 Sec-WebSocket-Extensions 头字段扩展中的。当 EnableCompression 为 true 时，服务端根据当前自身支持的扩展与其进行匹配，如果匹配成功则支持压缩。

校验逻辑：

var strPermessageDeflate = []byte("permessage-deflate")
​
func (u *HertzUpgrader) isCompressionEnable(ctx *app.RequestContext) bool {
    extensions := parseDataHeader(ctx.Request.Header.Peek("Sec-WebSocket-Extensions"))
​
    // Negotiate PMCE
    if u.EnableCompression {
        for _, ext := range extensions {
            if bytes.HasPrefix(ext, strPermessageDeflate) {
                return true
            }
        }
    }
​
    return false
}

目前仅支持“无上下文接管”模式，详见上述 HertzUpgrader.Upgrade 代码部分。

Set Deadline

当使用 websocket 进行读写的时候，可以通过类似如下方式设置超时时间（在每次读写过程中都会生效）。

示例代码：

func echo(_ context.Context, c *app.RequestContext) {
    err := upgrader.Upgrade(c, func(conn *websocket.Conn) {
        defer conn.Close()
        // "github.com/cloudwego/hertz/pkg/network"
        conn.NetConn().(network.Conn).SetReadTimeout(1 * time.Second)
        ...
    })
    if err != nil {
        log.Print("upgrade:", err)
        return
    }
}

更多用法示例详见 examples 。

感兴趣可以关注公众号 「白泽talk」，白泽目前也打算打造一个氛围良好的行业交流群，文章的更新也会提前预告，欢迎加入：622383022。