From 986b27475f358e75e51e5c846411a6a7e9cd661c Mon Sep 17 00:00:00 2001
From: "rui.zheng" <rui.zheng@tcl.com>
Date: Thu, 15 Sep 2016 12:28:07 +0800
Subject: [PATCH] fix wss bug

---
 conn.go |  4 ++--
 ws.go   | 28 +++++++++++-----------------
 2 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/conn.go b/conn.go
index cec2338..14e8bcf 100644
--- a/conn.go
+++ b/conn.go
@@ -440,13 +440,13 @@ func forward(conn net.Conn, arg Args) (net.Conn, error) {
 
 	switch arg.Transport {
 	case "ws": // websocket connection
-		conn, err = wsClient(conn, arg.Addr)
+		conn, err = wsClient("ws", conn, arg.Addr)
 		if err != nil {
 			return nil, err
 		}
 	case "wss": // websocket security
 		tlsUsed = true
-		conn, err = wssClient(conn, arg.Addr)
+		conn, err = wsClient("wss", conn, arg.Addr)
 		if err != nil {
 			return nil, err
 		}
diff --git a/ws.go b/ws.go
index 39c1490..e549723 100644
--- a/ws.go
+++ b/ws.go
@@ -17,24 +17,18 @@ type wsConn struct {
 	rb   []byte
 }
 
-func wsClient(conn net.Conn, host string) (*wsConn, error) {
-	c, resp, err := websocket.NewClient(conn, &url.URL{Scheme: "ws", Host: host, Path: "/ws"}, nil, 4096, 4096)
-	if err != nil {
-		return nil, err
+func wsClient(scheme string, conn net.Conn, host string) (*wsConn, error) {
+	dialer := websocket.Dialer{
+		ReadBufferSize:   4096,
+		WriteBufferSize:  4096,
+		TLSClientConfig:  &tls.Config{InsecureSkipVerify: true},
+		HandshakeTimeout: time.Second * 90,
+		NetDial: func(net, addr string) (net.Conn, error) {
+			return conn, nil
+		},
 	}
-	resp.Body.Close()
-
-	return &wsConn{conn: c}, nil
-}
-
-func wssClient(conn net.Conn, host string) (*wsConn, error) {
-	tlsConn := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
-	if err := tlsConn.Handshake(); err != nil {
-		return nil, err
-	}
-	conn = tlsConn
-
-	c, resp, err := websocket.NewClient(conn, &url.URL{Scheme: "wss", Host: host, Path: "/ws"}, nil, 4096, 4096)
+	u := url.URL{Scheme: scheme, Host: host, Path: "/ws"}
+	c, resp, err := dialer.Dial(u.String(), nil)
 	if err != nil {
 		return nil, err
 	}