Go WebSocket

使用 Go 语言创建 WebSocket 服务 | Go 技术论坛 (learnku.com)

go语言webSocket框架——gorilla_go websocket_尚墨1111的博客-CSDN博客

1.使用gorilla/websocket包构建websocket服务小例子

通过上面的章节可以看到用 Go 自带的 net/http 库实现 WebSocket 服务还是太复杂了。好在有很多对 WebSocket 支持良好的第三方库，能减少我们很多底层的编码工作。这里我们使用 gorilla web toolkit 家族的另外一个库 gorilla/websocket 来实现我们的 WebSocket 服务，构建一个简单的 Echo 服务（echo 意思是回音，就是客户端发什么，服务端再把消息发回给客户端）

我们在 http_demo 项目的 handler 目录下新建一个 ws 子目录用来存放 WebSocket 服务相关的路由对应的请求处理程序。

增加两个路由：

/ws/echo echo 应用的 WebSocket 服务的路由。
/ws/echo_display echo 应用的客户端页面的路由。

创建 WebSocket 服务端

这段程序的意思就是升级协议，然后把ReadMessage收到的内容全都WriteMessage 写回客户端

// handler/ws/echo.go
package ws

import (
    "fmt"
    "github.com/gorilla/websocket"
    "net/http"
)
//获取upgrader，设置读写缓存大小
var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
}

func EchoMessage(w http.ResponseWriter, r *http.Request) {
    //先用响应和请求升级协议，获取连接conn
    conn, _ := upgrader.Upgrade(w, r, nil) // 实际应用时记得做错误处理
    //一直收取对面发送消息
    for {
        // 读取客户端的消息
        msgType, msg, err := conn.ReadMessage()
        if err != nil {
            return
        }

        // 把消息打印到标准输出
        fmt.Printf("%s sent: %s\n", conn.RemoteAddr(), string(msg))

        // 把消息写回客户端，完成回音
        if err = conn.WriteMessage(msgType, msg); err != nil {
            return
        }
    }
}

conn 变量的类型是 *websocket.Conn, websocket.Conn 类型用来表示 WebSocket 连接。服务器应用程序从 HTTP 请求处理程序调用 Upgrader.Upgrade 方法以获取 *websocket.Conn

调用连接的 WriteMessage 和 ReadMessage 方法发送和接收消息。上面的 msg 接收到后在下面又回传给了客户端。msg 的类型是 []byte。

创建 WebSocket 客户端

前端页面路由对应的请求处理程序如下，直接返回 views/websockets.html 给到浏览器渲染页面即可。

// handler/ws/echo_display.go
package ws

import "net/http"

func DisplayEcho(w http.ResponseWriter, r *http.Request) {
    http.ServeFile(w, r, "views/websockets.html")
}

2.简易介绍

gorilla websocket简易介绍_gorilla/websocket 群发_shankusu2017的博客-CSDN博客

我们知道websocket由http升级而来，首先会发送附带Upgrade请求头的Http请求，所以我们需要在处理Http请求时拦截请求并判断其是否为websocket升级请求，如果是则调用gorilla/websocket库相应函数处理升级请求。

1）首先要创建Upgrader实例，该实例用于升级请求

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin:     checkOrigin,
}
func checkOrigin(r *http.Request) bool {
    return true
}

其中CheckOringin是一个函数，该函数用于拦截或放行跨域请求。函数返回值为bool类型，即true放行，false拦截。如果请求不是跨域请求可以不赋值，我这里是跨域请求并且为了方便直接返回true

 2）处理Http请求时拦截请求并判断其是否为websocket升级请求

服务器：是就升级协议，否则正常处理

//Http入口
func (e *Engine) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    //判断请求是否为websocket升级请求。
    if websocket.IsWebSocketUpgrade(r) {
        conn, err := upgrader.Upgrade(w, r, w.Header())
    } else {
        //处理普通请求
        c := newContext(w, r)
        e.router.handle(c)
    }
}

服务端：此时已经成功升级为websocket连接并获得一个conn实例，之后的发送接收操作皆由conn完成。其类型为websocket.Conn。

3)服务器升级后加上写消息

首先向客户端发送消息使用WriteMessage(messageType int, data []byte),参数1为消息类型，参数2消息内容

func (e *Engine) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    //判断请求是否为websocket升级请求。
    if websocket.IsWebSocketUpgrade(r) {
        conn, err := upgrader.Upgrade(w, r, w.Header())
        conn.WriteMessage(websocket.TextMessage, []byte("写入的消息"))
    } else {
        //处理普通请求
        c := newContext(w, r)
        e.router.handle(c)
    }
}

这个相比于上一块，就是多了一个写消息操作

4）服务器读客户端消息

接受客户端消息使用ReadMessage()该操作会阻塞线程，所以建议运行在其他协程上。该函数有三个返回值分别是，接收消息类型、接收消息内容、发生的错误，当然正常执行时错误为 nil。

一旦连接关闭消息返回值类型为-1，此时可以终止读操作。

func (e *Engine) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    //判断请求是否为websocket升级请求。
    if websocket.IsWebSocketUpgrade(r) {
        conn, err := upgrader.Upgrade(w, r, w.Header())
        conn.WriteMessage(websocket.TextMessage, []byte("xxxx"))
        go func() {
            for {
                //一直读取消息，如果关闭连接了就会读到-1，就退出
                t, c, _ := conn.ReadMessage()
                fmt.Println(t, string(c))
                if t == -1 {
                    return
                }
            }
        }()
    } else {
        //处理普通请求
        c := newContext(w, r)
        e.router.handle(c)
    }
}

这个在上一块基础上，再加了一个读取消息的协程

5）设置关闭连接监听

同时可以为连接设置关闭连接监听，函数为SetCloseHandler(h func(code int, text string) error)函数接收一个函数为参数，参数为nil时有一个默认实现，其源码为：

func (c *Conn) SetCloseHandler(h func(code int, text string) error) {
    if h == nil {
        h = func(code int, text string) error {
            message := FormatCloseMessage(code, "")
            c.WriteControl(CloseMessage, message, time.Now().Add(writeWait))
            return nil
        }
    }
    c.handleClose = h
}

可以看到函数的参数为int和string类型，前端关闭连接后js会把两个参数发送给后端，并传递给上述函数func(code int, text string) error使用。

6）整合

func (e *Engine) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    //判断请求是否为websocket升级请求。
    if websocket.IsWebSocketUpgrade(r) {
        //升级
        conn, err := upgrader.Upgrade(w, r, w.Header())
        //写消息
        conn.WriteMessage(websocket.TextMessage, []byte("xxxx"))
        //连接关闭 监听
        conn.SetCloseHandler(func(code int, text string) error {
            fmt.Println(code, text)
            return nil
        })
        //持续读消息，关闭后退出
        go func() {
            for {
                t, c, _ := conn.ReadMessage()
                fmt.Println(t, string(c))
                if t == -1 {
                    return
                }
            }
        }()
    } else {
        //处理普通请求
        c := newContext(w, r)
        e.router.handle(c)
    }
}