Golang 從 TCP 升級爲 WebSocket

【導讀】今天推送的文章中作者記錄了時間緊任務重且成功上線了 TCP 轉 WebSocket 升級的操作。

有一個服務器原來是 TCP 的私有協議,突然需求要支持 WebSocket,趕鴨子想在原來的端口上硬上 WebSocket。最後居然還比較簡單地成功了,必須說 golang 很舒服。

  1. WebSocket 庫

  2. 從 TCP 升級成 WebSocket

  3. 使用方法

WebSocket 庫

websocket 庫選了官方的 http://golang.org/x/net/websocket,可以從 https://github.com/golang/net.git 克隆(到 go/src/golang.org/x/net)。

官方的 websocket 庫用起來還是挺簡單的,接口文檔可以參考:websocket · pkg.go.dev。

客戶端

conn, err := websocket.Dial(url, subprotocol, origin)
websocket.Message.Send(conn, "") // 發送一個 string
var b []byte
websocket.Message.Receive(conn, &b) // 接收一個 []byte

客戶端用 websocket.Dial() 來創建連接 *websocket.Conn。其中:

雖然 *websocket.Conn 有 Read/Write 等方法,但使用 websocket.Message 更方便,因爲可以保證一個封包的完整性。

服務端

var recv func([]byte)
var err error
f := func(conn *websocket.Conn) {
    for {
        var b []byte
        err = websocket.Message.Receive(conn, &b)
        if err != nil {
            return
        } else {
            recv(b)
        }
    }
}

websocket.Handler(f).ServeHTTP(w, r)

用 websocket.Handler 或者 websocket.Server 兩個類來升級(Upgrade) HTTP 請求,在回調中會收到一個 *websocket.Conn 以供業務方使用。

連接收發

如前文所示,可以用 websocket.Message 來進行簡單的二進制或者字符串的收發,並且一次是一個完整的封包。

websocket.Codec 還可以支持序列化與反序列化,直接收發 golang 對象,只需要定義兩個函數就好了,一個序列化,一個反序列化。websocket.JSON 是預置的解碼器。另外 websocket.Message 也是一個解碼器。

當然 *websocket.Conn 本身也實現了 net.Conn,擁有 RemoteAddr、Read、Write 等方法。只是使用 Read/Write 會模糊 WebSocket 協議的封裝,沒有必要。

從 TCP 升級成 WebSocket

幸運的是,TCP 私有協議與 WebSocket 握手協議有完全不同的協議頭。所以判斷頭三個字節是不是 GET,就可以區分要不要轉 WebSocket。

服務端創建 *websocket.Conn 可以通過 Handler.ServeHTTP(),但 TCP 協議嘛,只有一個 *net.TCPConn,而且已經讀取了一些內容了。現在需要把一個 []byte + *net.TCPConn 變成 http.ResponseWriter + *http.Request。

http.ResponseWriter

http.ResponseWriter 是一個接口,可以簡單模擬,而且 WebSocket 會通過 Hijack 轉走,所以可以暴力實現之:

type wsFakeWriter struct {
 conn *net.TCPConn
 rw   *bufio.ReadWriter
}

func makeHttpResponseWriter(conn *net.TCPConn) *wsFakeWriter {
 w := new(wsFakeWriter)
 w.conn = conn
 w.rw = bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))

 return w
}

func (w *wsFakeWriter) Header() http.Header {
 return nil
}

func (w *wsFakeWriter) WriteHeader(int) {
 // 處理升級失敗情況??
}

func (w *wsFakeWriter) Write([]byte) (int, error) {
 return 0, nil
}

func (w *wsFakeWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 return w.conn, w.rw, nil
}

*http.Request

對於 *http.Request,很幸運地有 http.ReadRequest() 可用:

            func ReadRequest(b *bufio.Reader) (*Request, error)

只是需要把 []byte 和 *net.TCPConn 打包成 bufio.Reader:

func joinBufferAndReader(buffer []byte, reader io.Reader) io.Reader {
 return io.MultiReader(bytes.NewReader(buffer), reader)
}

func takeHttpRequest(buffer []byte, conn *net.TCPConn) (*http.Request, error) {
 r := joinBufferAndReader(buffer, conn)
 return http.ReadRequest(bufio.NewReader(r))
}

託 golang 強大的接口自動認證的福,這個打包過程甚至不需要做太多,調用標準庫就滿足了。[]byte 可以變成 io.Reader,*net.TCPConn 自然就是 io.Reader,標準庫還能串聯 io.Reader,一切都完美。

使用方法

func WebsocketOnTCP(buffer []byte, conn *net.TCPConn, recv func(*Package)) error {
 req, err := takeHttpRequest(buffer, conn)
 if err != nil {
  return err
 }

 f := func(ws *websocket.Conn) {
  err = doWebSocket(ws, recv)
 }

 w := makeHttpResponseWriter(conn)
 websocket.Handler(f).ServeHTTP(w, req)

 return err
}

func doWebSocket(conn *websocket.Conn, recv func(*Package)) error {
 remoteAddr := conn.RemoteAddr()
 reply := func([]byte) error {
  websocket.Message.Send(conn, b) // 此處線程不安全
 }

 for {
  var b []byte
  err := websocket.Message.Receive(conn, &b)
  if err != nil {
   return err
  }

  pack := new(Package)
  pack.Addr = remoteAddr
  pack.Content = b
  pack.Reply = reply

  recv(pack)
 }
}

以上只是粗略的使用方法,簡單的收發可以成功。只是未驗證過線程安全,Upgrade 失敗等情況。

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/t3PjaxJQaNcPr2kjMnco4A