fix wss bug

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
 		}

ws.go

@@ -17,24 +17,18 @@ type wsConn struct {
 	rb   []byte
 }
 
-func wsClient(conn net.Conn, host string) (*wsConn, error) {
+func wsClient(scheme string, conn net.Conn, host string) (*wsConn, error) {
-	c, resp, err := websocket.NewClient(conn, &url.URL{Scheme: "ws", Host: host, Path: "/ws"}, nil, 4096, 4096)
+	dialer := websocket.Dialer{
-	if err != nil {
+		ReadBufferSize:   4096,
-		return nil, err
+		WriteBufferSize:  4096,
-	}
+		TLSClientConfig:  &tls.Config{InsecureSkipVerify: true},
-	resp.Body.Close()
+		HandshakeTimeout: time.Second * 90,
-
+		NetDial: func(net, addr string) (net.Conn, error) {
-	return &wsConn{conn: c}, nil
+			return conn, 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
+	u := url.URL{Scheme: scheme, Host: host, Path: "/ws"}
-
+	c, resp, err := dialer.Dial(u.String(), nil)
-	c, resp, err := websocket.NewClient(conn, &url.URL{Scheme: "wss", Host: host, Path: "/ws"}, nil, 4096, 4096)
 	if err != nil {
 		return nil, err
 	}