Golang 從 TCP 升級為 WebSocket

今天推送的文章中作者記錄了時間緊任務(wù)重且成功上線了 TCP 轉(zhuǎn) WebSocket升級的操作。

有一個服務(wù)器原來是 TCP 的私有協(xié)議，突然需求要支持 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, "") // 發(fā)送一個 string
var b []byte
websocket.Message.Receive(conn, &b) // 接收一個 []byte

客戶端用 websocket.Dial() 來創(chuàng)建連接 *websocket.Conn。其中：

• subprotocol 表示細分的協(xié)議格式（如多種不同的序列化方法），默認(rèn)可為空
• origin 表示發(fā)起請求的網(wǎng)站（只需要 http://<host> 這樣）

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

服務(wù)端

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 請求，在回調(diào)中會收到一個 *websocket.Conn 以供業(yè)務(wù)方使用。

• Handler 或 Server 均可注冊到 net.http 中使用，但也可以自行調(diào)用 ServeHTTP 方法
• Handler 只有一個簡單回調(diào)的函數(shù)接口，使用閉包可以使用更多上下文

連接收發(fā)

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

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

當(dāng)然 *websocket.Conn 本身也實現(xiàn)了 net.Conn，擁有 RemoteAddr、Read、Write 等方法。只是使用 Read/Write 會模糊 WebSocket 協(xié)議的封裝，沒有必要。

從 TCP 升級成 WebSocket

幸運的是，TCP 私有協(xié)議與 WebSocket 握手協(xié)議有完全不同的協(xié)議頭。所以判斷頭三個字節(jié)是不是 GET，就可以區(qū)分要不要轉(zhuǎn) WebSocket。

服務(wù)端創(chuàng)建 *websocket.Conn 可以通過 Handler.ServeHTTP()，但 TCP 協(xié)議嘛，只有一個 *net.TCPConn，而且已經(jīng)讀取了一些內(nèi)容了?，F(xiàn)在需要把一個 []byte + *net.TCPConn 變成 http.ResponseWriter + *http.Request。

http.ResponseWriter

http.ResponseWriter 是一個接口，可以簡單模擬，而且 WebSocket 會通過 Hijack 轉(zhuǎn)走，所以可以暴力實現(xiàn)之：

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(b []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 強大的接口自動認(rèn)證的福，這個打包過程甚至不需要做太多，調(diào)用標(biāo)準(zhǔn)庫就滿足了。[]byte 可以變成 io.Reader，*net.TCPConn 自然就是 io.Reader，標(biāo)準(zhǔn)庫還能串聯(lián) 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(b []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)
 }
}

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