Golang 從 TCP 升級爲 WebSocket
【導讀】今天推送的文章中作者記錄了時間緊任務重且成功上線了 TCP 轉 WebSocket 升級的操作。
有一個服務器原來是 TCP 的私有協議,突然需求要支持 WebSocket,趕鴨子想在原來的端口上硬上 WebSocket。最後居然還比較簡單地成功了,必須說 golang 很舒服。
-
WebSocket 庫
-
從 TCP 升級成 WebSocket
-
使用方法
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。其中:
-
subprotocol 表示細分的協議格式(如多種不同的序列化方法),默認可爲空
-
origin 表示發起請求的網站(只需要 http:// 這樣)
雖然 *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 以供業務方使用。
-
Handler 或 Server 均可註冊到 net.http 中使用,但也可以自行調用 ServeHTTP 方法
-
Handler 只有一個簡單回調的函數接口,使用閉包可以使用更多上下文
連接收發
如前文所示,可以用 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(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 強大的接口自動認證的福,這個打包過程甚至不需要做太多,調用標準庫就滿足了。[]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(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)
}
}
以上只是粗略的使用方法,簡單的收發可以成功。只是未驗證過線程安全,Upgrade 失敗等情況。
本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源:https://mp.weixin.qq.com/s/t3PjaxJQaNcPr2kjMnco4A