为开源项目 go-gin-api 增加 WebSocket 模块

WebSocket 示例界面

第三方包

gorilla/websocket^[1]

websocket 链接时支持配置项：

type Upgrader struct {

    // 指定升级 websocket 握手完成的超时时间

    HandshakeTimeout time.Duration



    // 指定 io 操作的缓存大小，如果不指定就会自动分配。

    ReadBufferSize, WriteBufferSize int



    // 写数据操作的缓存池，如果没有设置值，write buffers 将会分配到链接生命周期里。

    WriteBufferPool BufferPool



    //按顺序指定服务支持的协议，如值存在，则服务会从第一个开始匹配客户端的协议。

    Subprotocols []string



    // 指定 http 的错误响应函数，如果没有设置 Error 则，会生成 http.Error 的错误响应。

    Error func(w http.ResponseWriter, r *http.Request, status int, reason error)



    // 请求检查函数，用于统一的链接检查，以防止跨站点请求伪造。如果不检查，就设置一个返回值为 true 的函数。

    // 如果请求 Origin 标头可以接受，CheckOrigin 将返回 true。 如果 CheckOrigin 为nil，则使用安全默认值：如果 Origin 请求头存在且原始主机不等于请求主机头，则返回 false

    CheckOrigin func(r *http.Request) bool



    // EnableCompression 指定服务器是否应尝试协商每个邮件压缩（RFC 7692）。 

    // 将此值设置为true并不能保证将支持压缩。 

    // 目前仅支持“无上下文接管”模式

    EnableCompression bool

}

Upgrade 函数可将 http 升级到 WebSocket 协议：

// responseHeader 包含在对客户端升级请求的响应中。 

// 使用 responseHeader 指定 cookie(Set-Cookie)和应用程序协商的子协议(Sec-WebSocket-Protocol)。

// 如果升级失败，则升级将使用 HTTP 错误响应回复客户端

// 返回一个 Conn 指针，拿到他后，可使用 Conn 读写数据与客户端通信。

func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*Conn, error)

代码封装

封装一个 socket_server，同时支持使用日志、数据库、缓存等操作：

type Server interface {



 // OnMessage 接收消息

 OnMessage()



 // OnSend 发送消息

 OnSend(message []byte) error



 // OnClose 关闭

 OnClose()

}



func New(logger *zap.Logger, db db.Repo, cache cache.Repo, conn *websocket.Conn) (Server, error) {

    

    ...

    

}

封装一个 socket_conn：

func (h *handler) Connect() core.HandlerFunc {

 var upGrader = websocket.Upgrader{

  HandshakeTimeout: 5 * time.Second,

  CheckOrigin: func(r *http.Request) bool {

   return true

  },

 }



 return func(ctx core.Context) {

  ws, err := upGrader.Upgrade(ctx.ResponseWriter(), ctx.Request(), nil)

  if err != nil {

   return

  }



  server, err = socket_server.New(h.logger, h.db, h.cache, ws)

  if err != nil {

   return

  }



  go server.OnMessage()

 }

}

socket_conn 是一个 HandlerFunc，可直接在路由中使用。

项目中 websocket 链接地址为：/socket/system/message，发送消息的接口为：/api/tool/send_message。

JavaScript 示例代码：

const ws = new WebSocket("ws://127.0.0.1:9999/socket/system/message");



//连接打开时触发

ws.onopen = function (evt) {

    ...

};



//接收到消息时触发

ws.onmessage = function (evt) {

    ...

};



//连接关闭时触发

ws.onclose = function (evt) {

    ...

};

在项目中实用工具箱 -> WebSocket 界面，右侧请求接口发送消息，在左侧可以实时收到消息。