From a39bb9ecf84ede586bebe4706d327e1732e06654 Mon Sep 17 00:00:00 2001
From: "rui.zheng" <rui.zheng@tcl.com>
Date: Tue, 12 May 2015 13:33:45 +0800
Subject: [PATCH] fix bug

---
 client.go | 57 +++++++++++++--------------------------
 util.go   |  4 +++
 ws.go     | 80 ++++++++++++++++++++++++++++---------------------------
 3 files changed, 63 insertions(+), 78 deletions(-)

diff --git a/client.go b/client.go
index 8b3d1b7..aa72e7c 100644
--- a/client.go
+++ b/client.go
@@ -11,10 +11,10 @@ import (
 	"github.com/shadowsocks/shadowsocks-go/shadowsocks"
 	"io"
 	"io/ioutil"
-	"net/url"
 	"log"
 	"net"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
 	//"sync/atomic"
@@ -65,11 +65,9 @@ func clientMethodSelected(method uint8, conn net.Conn) (net.Conn, error) {
 		cipher, _ := shadowsocks.NewCipher(Methods[method], Password)
 		conn = shadowsocks.NewConn(conn, cipher)
 	case gosocks5.MethodNoAcceptable:
-		fallthrough
-	default:
 		return nil, gosocks5.ErrBadMethod
 	}
-	
+
 	return conn, nil
 }
 
@@ -81,54 +79,35 @@ func cliHandle(conn net.Conn) {
 			fmt.Println("session end", atomic.AddInt64(&sessionCount, -1))
 		}()
 	*/
-	addr := Saddr
-	if strings.HasPrefix(addr, "http://") {
-		u, err := url.Parse(addr)
-		if err != nil {
-			log.Println(err)
-			return
-		}
-		addr = u.Host
-		if !strings.Contains(addr, ":") {
-			addr += ":80"
-		}
-	}
-	if strings.HasPrefix(addr, "https://") {
-		u, err := url.Parse(addr)
-		if err != nil {
-			log.Println(err)
-			return
-		}
-		addr = u.Host
-		if !strings.Contains(addr, ":") {
-			addr += ":443"
-		}
-	}
-	log.Println("connect:", addr)
-	c, err := Connect(addr, Proxy)
+
+	//log.Println("connect:", Saddr)
+	c, err := Connect(Saddr, Proxy)
 	if err != nil {
 		log.Println(err)
 		return
 	}
 	defer c.Close()
-	
+
 	if Websocket {
-		u, err := url.Parse(addr)
-		if err != nil {
-			log.Println(err)
-			return
+		url := &url.URL{
+			Host: Saddr,
 		}
-		ws, resp, err := websocket.NewClient(c, u, nil, 8192, 8192)
+		ws, resp, err := websocket.NewClient(c, url, nil, 8192, 8192)
 		if err != nil {
 			log.Println(err)
 			return
 		}
 		resp.Body.Close()
-		
+
 		c = NewWSConn(ws)
 	}
-	
-	c = gosocks5.ClientConn(c, clientConfig)
+
+	sc := gosocks5.ClientConn(c, clientConfig)
+	if err := sc.Handleshake(); err != nil {
+		log.Println(err)
+		return
+	}
+	c = sc
 
 	if Shadows {
 		cipher, _ := shadowsocks.NewCipher(SMethod, SPassword)
@@ -186,7 +165,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:
diff --git a/util.go b/util.go
index 43a8953..478e393 100644
--- a/util.go
+++ b/util.go
@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"strings"
 )
 
 const (
@@ -51,6 +52,9 @@ func ToSocksAddr(addr net.Addr) *gosocks5.Addr {
 }
 
 func Connect(addr, proxy string) (net.Conn, error) {
+	if !strings.Contains(addr, ":") {
+		addr += ":80"
+	}
 	if len(proxy) == 0 {
 		taddr, err := net.ResolveTCPAddr("tcp", addr)
 		if err != nil {
diff --git a/ws.go b/ws.go
index 017fd4b..5ab6b71 100644
--- a/ws.go
+++ b/ws.go
@@ -1,69 +1,71 @@
 package main
 
 import (
-    "github.com/gorilla/websocket"
-    "net/http"
-    "log"
-    "time"
-    "github.com/ginuerzh/gosocks5"
+	"github.com/ginuerzh/gosocks5"
+	"github.com/gorilla/websocket"
+	"log"
+	"net/http"
+	"time"
 )
 
 type WSConn struct {
-    *websocket.Conn
+	*websocket.Conn
 }
 
 func NewWSConn(conn *websocket.Conn) *WSConn {
-    c := &WSConn{}
-    c.Conn = conn
-    
-    return c
+	c := &WSConn{}
+	c.Conn = conn
+
+	return c
 }
 
 func (conn *WSConn) Read(b []byte) (n int, err error) {
-    _, b, err = conn.ReadMessage()
-    n = len(b)
-    
-    return
+	_, r, err := conn.NextReader()
+	if err != nil {
+		return
+	}
+	return r.Read(b)
 }
 
 func (conn *WSConn) Write(b []byte) (n int, err error) {
-    n = len(b)
-    err = conn.WriteMessage(websocket.BinaryMessage, b)
-    return
+	w, err := conn.NextWriter(websocket.BinaryMessage)
+	if err != nil {
+		return
+	}
+	defer w.Close()
+	return w.Write(b)
 }
 
 func (conn *WSConn) SetDeadline(t time.Time) error {
-    if err := conn.SetReadDeadline(t); err != nil {
-        return err
-    }
-    return conn.SetWriteDeadline(t)
+	if err := conn.SetReadDeadline(t); err != nil {
+		return err
+	}
+	return conn.SetWriteDeadline(t)
 }
 
 type WSServer struct {
-    Addr string
+	Addr string
 }
 
 var upgrader = websocket.Upgrader{
-    ReadBufferSize:  8192,
-    WriteBufferSize: 8192,
-    CheckOrigin: func(r *http.Request) bool{ return true;},
+	ReadBufferSize:  8192,
+	WriteBufferSize: 8192,
+	CheckOrigin:     func(r *http.Request) bool { return true },
 }
 
-
 func (s *WSServer) handle(w http.ResponseWriter, r *http.Request) {
-    conn, err := upgrader.Upgrade(w, r, nil)
-    if err != nil {
-        log.Println(err)
-        return
-    }
-    defer conn.Close()
-    
-    c := NewWSConn(conn)
-    
-    socks5Handle(gosocks5.ServerConn(c, serverConfig))
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer conn.Close()
+
+	c := gosocks5.ServerConn(NewWSConn(conn), serverConfig)
+	socks5Handle(c)
 }
 
 func (s *WSServer) ListenAndServe() error {
-    http.HandleFunc("/", s.handle)
-    return http.ListenAndServe(s.Addr, nil)
-}
\ No newline at end of file
+	http.HandleFunc("/", s.handle)
+	return http.ListenAndServe(s.Addr, nil)
+}