From fdabebcabe727f8f6fd25dad42fdb04b1addbf42 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Sat, 16 May 2015 17:41:26 +0800
Subject: [PATCH] websocket tunnel support encryption

---
 client.go | 13 +++++++++----
 socks5.go | 14 ++++++++++----
 util.go   | 21 +++++++++++----------
 ws.go     | 17 ++++++++++++-----
 4 files changed, 42 insertions(+), 23 deletions(-)

diff --git a/client.go b/client.go
index 2213433..2ee82fc 100644
--- a/client.go
+++ b/client.go
@@ -62,7 +62,11 @@ func clientMethodSelected(method uint8, conn net.Conn) (net.Conn, error) {
 		conn = tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
 	case MethodAES128, MethodAES192, MethodAES256,
 		MethodDES, MethodBF, MethodCAST5, MethodRC4MD5, MethodRC4, MethodTable:
-		cipher, _ := shadowsocks.NewCipher(Methods[method], Password)
+		cipher, err := shadowsocks.NewCipher(Methods[method], Password)
+		if err != nil {
+			log.Println(err)
+			return nil, err
+		}
 		conn = shadowsocks.NewConn(conn, cipher)
 	case gosocks5.MethodNoAcceptable:
 		return nil, gosocks5.ErrBadMethod
@@ -104,7 +108,6 @@ func cliHandle(conn net.Conn) {
 
 	sc := gosocks5.ClientConn(c, clientConfig)
 	if err := sc.Handleshake(); err != nil {
-		log.Println(err)
 		return
 	}
 	c = sc
@@ -121,7 +124,6 @@ func cliHandle(conn net.Conn) {
 
 	n, err := io.ReadAtLeast(conn, b, 2)
 	if err != nil {
-		//log.Println(err)
 		return
 	}
 
@@ -148,6 +150,7 @@ func cliHandle(conn net.Conn) {
 
 		nn, err := conn.Read(b[n:])
 		if err != nil {
+			log.Println(err)
 			return
 		}
 		n += nn
@@ -155,6 +158,7 @@ func cliHandle(conn net.Conn) {
 
 	req, err := http.ReadRequest(bufio.NewReader(bytes.NewReader(b[:n])))
 	if err != nil {
+		log.Println(err)
 		return
 	}
 	handleHttp(req, conn, c)
@@ -165,7 +169,7 @@ func handleSocks5(conn net.Conn, sconn net.Conn) {
 	if err != nil {
 		return
 	}
-	log.Println(req)
+	//log.Println(req)
 
 	switch req.Cmd {
 	case gosocks5.CmdConnect, gosocks5.CmdBind:
@@ -202,6 +206,7 @@ func handleSocks5(conn net.Conn, sconn net.Conn) {
 
 		go cliTunnelUDP(uconn, sconn)
 
+		// block, waiting for client exit
 		ioutil.ReadAll(conn)
 	}
 }
diff --git a/socks5.go b/socks5.go
index b4e5812..aae3721 100644
--- a/socks5.go
+++ b/socks5.go
@@ -94,12 +94,19 @@ func (s *Socks5Server) ListenAndServe() error {
 }
 
 func serverSelectMethod(methods ...uint8) uint8 {
+	m := gosocks5.MethodNoAuth
+
 	for _, method := range methods {
 		if _, ok := Methods[method]; ok {
-			return method
+			m = method
 		}
 	}
-	return gosocks5.MethodNoAuth
+
+	if len(Method) == 0 || Methods[m] == Method {
+		return m
+	}
+
+	return gosocks5.MethodNoAcceptable
 }
 
 func serverMethodSelected(method uint8, conn net.Conn) (net.Conn, error) {
@@ -115,7 +122,6 @@ func serverMethodSelected(method uint8, conn net.Conn) (net.Conn, error) {
 		}
 
 		if err != nil {
-			log.Println(err)
 			return nil, err
 		}
 		conn = tls.Server(conn, &tls.Config{Certificates: []tls.Certificate{cert}})
@@ -144,7 +150,7 @@ func socks5Handle(conn net.Conn) {
 
 	switch req.Cmd {
 	case gosocks5.CmdConnect:
-		log.Println("connect", req.Addr.String())
+		//log.Println("connect", req.Addr.String())
 		tconn, err := Connect(req.Addr.String(), Proxy)
 		if err != nil {
 			gosocks5.NewReply(gosocks5.HostUnreachable, nil).Write(conn)
diff --git a/util.go b/util.go
index 478e393..075d129 100644
--- a/util.go
+++ b/util.go
@@ -28,16 +28,17 @@ const (
 )
 
 var Methods = map[uint8]string{
-	MethodTLS:    "tls",         // 0x80
-	MethodAES128: "aes-128-cfb", // 0x81
-	MethodAES192: "aes-192-cfb", // 0x82
-	MethodAES256: "aes-256-cfb", // 0x83
-	MethodDES:    "des-cfb",     // 0x84
-	MethodBF:     "bf-cfb",      // 0x85
-	MethodCAST5:  "cast5-cfb",   // 0x86
-	MethodRC4MD5: "rc4-md5",     // 8x87
-	MethodRC4:    "rc4",         // 0x88
-	MethodTable:  "table",       // 0x89
+	gosocks5.MethodNoAuth: "",            // 0x00
+	MethodTLS:             "tls",         // 0x80
+	MethodAES128:          "aes-128-cfb", // 0x81
+	MethodAES192:          "aes-192-cfb", // 0x82
+	MethodAES256:          "aes-256-cfb", // 0x83
+	MethodDES:             "des-cfb",     // 0x84
+	MethodBF:              "bf-cfb",      // 0x85
+	MethodCAST5:           "cast5-cfb",   // 0x86
+	MethodRC4MD5:          "rc4-md5",     // 8x87
+	MethodRC4:             "rc4",         // 0x88
+	MethodTable:           "table",       // 0x89
 }
 
 func ToSocksAddr(addr net.Addr) *gosocks5.Addr {
diff --git a/ws.go b/ws.go
index 2ade195..14d698a 100644
--- a/ws.go
+++ b/ws.go
@@ -10,20 +10,26 @@ import (
 
 type WSConn struct {
 	*websocket.Conn
+	rb []byte
 }
 
 func NewWSConn(conn *websocket.Conn) *WSConn {
-	c := &WSConn{}
-	c.Conn = conn
+	c := &WSConn{
+		Conn: conn,
+	}
 
 	return c
 }
 
 func (conn *WSConn) Read(b []byte) (n int, err error) {
-	_, p, err := conn.ReadMessage()
-	copy(b, p)
-	n = len(p)
+	if len(conn.rb) == 0 {
+		_, conn.rb, err = conn.ReadMessage()
+	}
+	n = copy(b, conn.rb)
+	conn.rb = conn.rb[n:]
+
 	//log.Println("ws r:", n)
+
 	return
 }
 
@@ -31,6 +37,7 @@ func (conn *WSConn) Write(b []byte) (n int, err error) {
 	err = conn.WriteMessage(websocket.BinaryMessage, b)
 	n = len(b)
 	//log.Println("ws w:", n)
+
 	return
 }