From ba3b1e8ef53f3878005b6df2c820dddfe105010e Mon Sep 17 00:00:00 2001
From: "rui.zheng" <rui.zheng@tcl.com>
Date: Sat, 9 Jan 2016 12:52:59 +0800
Subject: [PATCH] support wss

---
 conn.go | 14 ++++++++++++--
 util.go |  2 +-
 ws.go   | 28 +++++++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/conn.go b/conn.go
index 2298c90..d20131b 100644
--- a/conn.go
+++ b/conn.go
@@ -35,6 +35,12 @@ func listenAndServe(arg Args) error {
 			glog.Infoln(err)
 		}
 		return err
+	case "wss": // websocket security connection
+		err = NewWs(arg).listenAndServeTLS()
+		if err != nil {
+			glog.Infoln(err)
+		}
+		return err
 	case "tls": // tls connection
 		ln, err = tls.Listen("tcp", arg.Addr,
 			&tls.Config{Certificates: []tls.Certificate{arg.Cert}})
@@ -253,8 +259,12 @@ func forward(conn net.Conn, arg Args) (net.Conn, error) {
 		if err != nil {
 			return nil, err
 		}
-	//case "wss": // websocket security
-	//	tlsUsed = true
+	case "wss": // websocket security
+		tlsUsed = true
+		conn, err = wssClient(conn, arg.Addr)
+		if err != nil {
+			return nil, err
+		}
 	case "tls": // tls connection
 		tlsUsed = true
 		conn = tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
diff --git a/util.go b/util.go
index bbf2e0d..63dc78e 100644
--- a/util.go
+++ b/util.go
@@ -72,7 +72,7 @@ func parseArgs(ss []string) (args []Args) {
 			arg.Protocol = "default"
 		}
 		switch arg.Transport {
-		case "ws", "tls", "tcp":
+		case "ws", "wss", "tls", "tcp":
 		default:
 			arg.Transport = "tcp"
 		}
diff --git a/ws.go b/ws.go
index ffcc924..dafb42c 100644
--- a/ws.go
+++ b/ws.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	//"github.com/ginuerzh/gosocks5"
+	"crypto/tls"
 	"github.com/golang/glog"
 	"github.com/gorilla/websocket"
 	"net"
@@ -17,7 +18,23 @@ type wsConn struct {
 }
 
 func wsClient(conn net.Conn, host string) (*wsConn, error) {
-	c, resp, err := websocket.NewClient(conn, &url.URL{Scheme: "ws", Host: host, Path: "/ws"}, nil, 1024, 1024)
+	c, resp, err := websocket.NewClient(conn, &url.URL{Scheme: "ws", Host: host, Path: "/ws"}, nil, 4096, 4096)
+	if err != nil {
+		return nil, err
+	}
+	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)
 	if err != nil {
 		return nil, err
 	}
@@ -115,3 +132,12 @@ func (s *ws) ListenAndServe() error {
 	http.HandleFunc("/ws", s.handle)
 	return http.ListenAndServe(s.arg.Addr, nil)
 }
+
+func (s *ws) listenAndServeTLS() error {
+	http.HandleFunc("/ws", s.handle)
+	server := &http.Server{
+		Addr:      s.arg.Addr,
+		TLSConfig: &tls.Config{Certificates: []tls.Certificate{s.arg.Cert}},
+	}
+	return server.ListenAndServeTLS("", "")
+}