给Go的Gin web框架增加 WebSocket 功能,让WebSocket 更好用

Gin 是一个 go 的 web 框架，它具有轻量级，高性能，运行速度快，分组的路由器，良好的崩溃捕获和错误处理，非常好的支持中间件，rest api和json。

总之在 Go语言开发领域是一款值得好好研究的 Web 框架。官方地址：https://github.com/gin-gonic/gin

但是最近想尝试下websocket功能，发现Gin框架里没有。

Go 官方没有提供对 WebSocket 的支持，必须选择第三方提供的包。

常用的有两种，golang.org/x/net/websocket和 https://github.com/gorilla/websocket

《Go Web 编程》一书中的例子使用的是golang.org/x/net/websocket

其中gorilla/websocket更常用些，Apache的Paho GO mqtt client库中和go的另外一个web框架iris中，就使用的是gorilla/websocket库。

是不是直接把gorilla/websocket引入进gin框架就可以了？

但是还得需要自己完成封装，若不花功夫封装好，不是很好用，稳定性更难保证。

网上找到gin中使用gorilla的websocket库的例子，都只是一些简单的测试，一点儿都不好用。

接下来，为了让websocket在gin中更好用些，移植iris框架中的websocket功能到gin框架中，使用起来就简单啦，使用如下：

github 地址: https://github.com/yangyongzhen/gin-websocket.git

package main

import (
	"fmt"
	//"github.com/gin-contrib/cors"
	"github.com/gin-gonic/gin"
	"net/http"
	"strings"
	"websockTest/websocket"
)

func main() {
	ws := websocket.New(websocket.Config{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	})

	ws.OnConnection(handleConnection)

	r := gin.Default()
	//允许跨域
	//config := cors.DefaultConfig()
	//config.AllowOrigins = []string{"http://127.0.0.1:9090"}
	//r.Use(Cors())
	//静态资源
	r.Static("/static", "./static")
	r.LoadHTMLGlob("views/*")
	r.GET("/api/v3/device", ws.Handler())
	r.GET("/test", func(c *gin.Context) {
		c.HTML(http.StatusOK, "test.html", gin.H{
			"title": "this is a test",
		})
	})
	r.Run(":9090")
}

func handleConnection(c websocket.Connection) {
	fmt.Println("client connected,id=", c.ID())
	c.Write(1, []byte("welcome client"))
	// 测试从浏览器中读取事件
	c.On("chat", func(msg string) {
		// 将消息打印到控制台
		fmt.Printf("%s sent: %s\n", c.Context().ClientIP(), msg)
		// 将消息写回客户端消息所有者：
		// c.Emit("chat", msg)
		c.To(websocket.All).Emit("chat", msg)
	})

	c.OnMessage(func(msg []byte) {
		fmt.Println("received msg:", string(msg))
		c.Write(1, []byte("hello aa"))
		c.To(websocket.All).Emit("chat", msg)
	})

	c.OnDisconnect(func() {
		fmt.Println("client Disconnect,id=", c.ID())
	})
}

附，网上找到的一些测试代码如下： 

package server

import (
	"fmt"
	"net"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second
	// Maximum message size allowed from peer.
	maxMessageSize = 8192
	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second
)

type WsServer struct {
	listener net.Listener
	addr     string
	upgrade  *websocket.Upgrader
}

func NewWsServer() *WsServer {
	ws := new(WsServer)
	ws.addr = "0.0.0.0:9090"
	ws.upgrade = &websocket.Upgrader{
		ReadBufferSize:  4096,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			if r.Method != "GET" {
				fmt.Println("method is not GET")
				return false
			}
			if r.URL.Path != "/ws" {
				fmt.Println("path error")
				return false
			}
			return true
		},
	}
	return ws
}

func (self *WsServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	if r.URL.Path != "/ws" {
		httpCode := http.StatusInternalServerError
		reasePhrase := http.StatusText(httpCode)
		fmt.Println("path error ", reasePhrase)
		fmt.Println("path error,url=", r.URL.Path)
		http.Error(w, reasePhrase, httpCode)
		return
	}

	conn, err := self.upgrade.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println("websocket error:", err)
		return
	}
	fmt.Println("client connect :", conn.RemoteAddr())
	go self.connHandle(conn)

}

func (self *WsServer) connHandle(conn *websocket.Conn) {
	defer func() {
		conn.Close()
	}()
	stopCh := make(chan int)
	go self.send(conn, stopCh)
	conn.SetReadLimit(maxMessageSize)
	conn.SetReadDeadline(time.Now().Add(pongWait))
	conn.SetPongHandler(func(string) error {
		fmt.Println("this is PongHandler")
		conn.SetReadDeadline(time.Now().Add(pongWait))
		if err := conn.WriteMessage(websocket.PongMessage, nil); err != nil {
			fmt.Println("resp pong error", err)
		} else {
			fmt.Println("resp pong ok")
		}
		return nil
	})
	for {
		conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(60000)))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			// 判断是不是超时
			if netErr, ok := err.(net.Error); ok {
				if netErr.Timeout() {
					fmt.Printf("ReadMessage timeout remote: %v\n", conn.RemoteAddr())
					close(stopCh)
					return
				}
			}
			// 其他错误，如果是 1001 和 1000 就不打印日志
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				fmt.Printf("ReadMessage other remote:%v error: %v \n", conn.RemoteAddr(), err)
			}
			//close(stopCh)
			return
		}
		fmt.Println("收到消息：", string(msg))
	}
}

//测试一次性发送 10万条数据给 client, 如果不使用 time.Sleep browser 过了超时时间会断开
func (self *WsServer) send10(conn *websocket.Conn) {
	for i := 0; i < 1; i++ {
		data := fmt.Sprintf("hello websocket test from server %v,count=%d", time.Now().UnixNano(), i+1)
		err := conn.WriteMessage(1, []byte(data))
		if err != nil {
			fmt.Println("send msg faild ", err)
			return
		}
		//time.Sleep(time.Millisecond * 1)
	}
}

func (self *WsServer) send(conn *websocket.Conn, stopCh chan int) {
	self.send10(conn)
	for {
		select {
		case <-stopCh:
			fmt.Println("connect closed")
			return
			// case <-time.After(time.Second * 1):
			// 	fmt.Println("time after....")
		}
	}
}

func (w *WsServer) Start() (err error) {
	w.listener, err = net.Listen("tcp", w.addr)
	if err != nil {
		fmt.Println("net listen error:", err)
		return
	}

	err = http.Serve(w.listener, w)
	if err != nil {
		fmt.Println("http serve error:", err)
		return
	}

	return nil
}