[[418263]] 

今天推送的文章中作者記錄了時間緊任務(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é)議格式（如多種不同的序列化方法），默認可為空
•  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)容了。現(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 強大的接口自動認證的福，這個打包過程甚至不需要做太多，調(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 失敗等情況。