[[418263]] 

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

有一個(gè)服務(wù)器原來(lái)是 TCP 的私有協(xié)議，突然需求要支持 WebSocket，趕鴨子想在原來(lái)的端口上硬上 WebSocket。最后居然還比較簡(jiǎn)單地成功了，必須說(shuō) golang 很舒服。

1.  WebSocket 庫(kù)
2.  從 TCP 升級(jí)成 WebSocket
3.  使用方法

WebSocket 庫(kù)

websocket 庫(kù)選了官方的 http://golang.org/x/net/websocket，可以從 https://github.com/golang/net.git 克?。ǖ?go/src/golang.org/x/net）。

官方的 websocket 庫(kù)用起來(lái)還是挺簡(jiǎn)單的，接口文檔可以參考：websocket · pkg.go.dev。

客戶端 

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

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

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

雖然 *websocket.Conn 有 Read/Write 等方法，但使用 websocket.Message 更方便，因?yàn)榭梢员ＷC一個(gè)封包的完整性。

服務(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 兩個(gè)類(lèi)來(lái)升級(jí)（Upgrade） HTTP 請(qǐng)求，在回調(diào)中會(huì)收到一個(gè) *websocket.Conn 以供業(yè)務(wù)方使用。

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

連接收發(fā)

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

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

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

從 TCP 升級(jí)成 WebSocket

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

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

http.ResponseWriter

http.ResponseWriter 是一個(gè)接口，可以簡(jiǎn)單模擬，而且 WebSocket 會(huì)通過(guò) Hijack 轉(zhuǎn)走，所以可以暴力實(shí)現(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) {  
 // 處理升級(jí)失敗情況？？  
}  
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

對(duì)于 *http.Request，很幸運(yùn)地有 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 強(qiáng)大的接口自動(dòng)認(rèn)證的福，這個(gè)打包過(guò)程甚至不需要做太多，調(diào)用標(biāo)準(zhǔn)庫(kù)就滿足了。[]byte 可以變成 io.Reader，*net.TCPConn 自然就是 io.Reader，標(biāo)準(zhǔn)庫(kù)還能串聯(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)  
 }  
} 

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