gorilla websocket簡易介紹

shankusu2017發表於2020-12-11

以下內容轉載自 https://studygolang.com/articles/30074

前言

最近打算為我的網站新增一個伺服器資源監視功能,需要服務端主動向前端推動資源佔用資料。這時Http則不能達到要求。所以自然想到採用websocket。

不自量力

以前使用SpringBoot時使用websocket很簡單,只需要將ServerEndpointExporter注入到bean容器並配合相應註解即可建立一個websocket服務。這裡要感謝各位前輩的封裝讓我們能儘快實現相應的功能,但本次出於學習目並不是公司專案(效率&穩定性至上)同時使用的開發語言為Golang,其web開發生態也不會像Java那樣豐富,所以起初是想自己實現一個websocket服務的,但當我真正查閱websocket協議RFC文件時發現以我現在的水平可能無法順利完成該項任務。最後還是選擇了開源實現gorilla/websocket專案地址

正文

注意:本文所有程式碼示例均在本人個人網站副本上擷取,如果有幸觀看本文的不是我自己需要注意例項程式碼可能無法正常執行因其缺少其他元件特此說明以防耽誤大家時間。

新增依賴

本文使用go mod管理依賴

執行go get github.com/gorilla/websocket新增依賴

使用

我們知道websocket由http升級而來,首先會傳送附帶Upgrade請求頭的Http請求,所以我們需要在處理Http請求時攔截請求並判斷其是否為websocket升級請求,如果是則呼叫gorilla/websocket庫相應函式處理升級請求。

首相要建立Upgrader例項,該例項用於升級請求

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

其中CheckOringin是一個函式,該函式用於攔截或放行跨域請求。函式返回值為bool型別,即true放行,false攔截。如果請求不是跨域請求可以不賦值,我這裡是跨域請求並且為了方便直接返回true

//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。

首先向客戶端傳送訊息使用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("wxm.alming"))
    } else {
        //處理普通請求
        c := newContext(w, r)
        e.router.handle(c)
    }
}

接受客戶端訊息使用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("wxm.alming"))
        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)
    }
}

同時可以為連線設定關閉連線監聽,函式為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型別正好和前端的close(long string)對應即前端呼叫close(long string)關閉連線後兩個引數會被髮送給後端並最終被func(code int, text string) error所使用。
示例:

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("wxm.alming"))
        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)
    }
}

則斷開連線時將列印code和text

注意:要想使斷連處理生效必須要有ReadMessage()操作否則不會觸發斷連處理操作。
以上是常用基礎操作點選官方API手冊學習更多。

附錄:前端測試程式碼

<!DOCTYPE html>
<html>
    <head>
        <meta charset="utf-8">
        <title></title>
        <style type="text/css">
            .body{text-align: center;}
            #open{width: 120px;height: 35px;}
            #close{width: 120px;height: 35px;}
            #text{display: inline-block;margin: auto;margin-top: 10px;margin-bottom: 10px;width: 240px;}
        </style>
    </head>
    <body class="body">
        <button id="open">開啟連線</button>
        <button id="close">關閉連線</button>
        <br/>
        <input type="text" name="text" id="text" value="" />
        <br/>
        <button id="send">傳送</button>
        <div id="msg">
            
        </div>
    </body>
    <script>
        var openbtn = document.getElementById("open")
        var closebtn = document.getElementById("close")
        var text = document.getElementById("text")
        var send = document.getElementById("send")
        var msg = document.getElementById("msg")
        var websocket
        openbtn.onclick = function(){
            websocket = new WebSocket("ws://your ip and port/")
            websocket.onopen=function(){
                console.log("connected");
            }
            websocket.onmessage = function(e){
                console.log(e);
                msg.innerHTML = '接收:'+e.data
            }
            websocket.onclose=function(e){
                console.log("closed",e);
            }
        }
        closebtn.onclick=function(){
            websocket.close(1000,"close")
        }
        send.onclick=function(){
            var msg = text.value
            websocket.send(msg)
        }
    </script>
</html>