Golang 从 TCP 升级为 WebSocket-golang 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://<host> 这样）

虽然 *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 失败等情况。