refactor socks5 client, server and shadowsocks client, and add socks5…

… client unit test.

cmd/socksd/main.go

@@ -44,7 +44,7 @@ func BuildUpstreamRouter(conf Config) socks.Dialer {
 		switch strings.ToLower(upstreamConf.ServerType) {
 		case "socks5":
 			{
-				forward, err = socks.NewSocks5Client("tcp", upstreamConf.Addr, forward)
+				forward, err = socks.NewSocks5Client("tcp", upstreamConf.Addr, "", "", forward)
 			}
 		case "shadowsocks":
 			{

shadowsocks_client.go

@@ -1,8 +1,7 @@
 package socks
 
 import (
-	"bytes"
-	"encoding/binary"
+	"errors"
 	"net"
 	"strconv"
 )
@@ -26,50 +25,59 @@ func NewShadowSocksClient(network, address string, forward Dialer) (*ShadowSocks
 }
 
 func (s *ShadowSocksClient) Dial(network, address string) (net.Conn, error) {
-	conn, err := s.forward.Dial(s.network, s.address)
-	if err != nil {
-		return nil, err
+	switch network {
+	case "tcp", "tcp4", "tcp6":
+	default:
+		return nil, errors.New("socks: no support ShadowSocks proxy connections of type: " + network)
 	}
-	connClose := &conn
-	defer func() {
-		if connClose != nil {
-			(*connClose).Close()
-		}
-	}()
 
-	host, p, err := parseAddress(address)
+	host, portStr, err := net.SplitHostPort(address)
 	if err != nil {
 		return nil, err
 	}
-	port, err := strconv.Atoi(p)
+	port, err := strconv.Atoi(portStr)
 	if err != nil {
-		return nil, err
+		return nil, errors.New("socks: failed to parse port number:" + portStr)
+	}
+	if port < 1 || port > 0xffff {
+		return nil, errors.New("socks5: port number out of range:" + portStr)
 	}
 
-	req := bytes.NewBuffer(nil)
-	switch host.(type) {
-	case string:
-		domain := host.(string)
-		req.WriteByte(3)
-		req.WriteByte(byte(len(domain)))
-		req.WriteString(domain)
-		binary.Write(req, binary.BigEndian, uint16(port))
+	conn, err := s.forward.Dial(s.network, s.address)
+	if err != nil {
+		return nil, err
+	}
+	closeConn := &conn
+	defer func() {
+		if closeConn != nil {
+			(*closeConn).Close()
+		}
+	}()
 
-	case net.IP:
-		ip := host.(net.IP)
-		if len(ip) == net.IPv4len {
-			req.WriteByte(1)
+	buff := make([]byte, 0, 266)
+	if ip := net.ParseIP(host); ip != nil {
+		if ip4 := ip.To4(); ip4 != nil {
+			buff = append(buff, 1)
+			ip = ip4
 		} else {
-			req.WriteByte(4)
+			buff = append(buff, 4)
+		}
+		buff = append(buff, ip...)
+	} else {
+		if len(host) > 255 {
+			return nil, errors.New("socks: destination hostname too long: " + host)
 		}
-		req.WriteByte(byte(len(ip)))
-		req.Write(ip)
-		binary.Write(req, binary.BigEndian, uint16(port))
+		buff = append(buff, 3)
+		buff = append(buff, uint8(len(host)))
+		buff = append(buff, host...)
 	}
-	_, err = conn.Write(req.Bytes())
+	buff = append(buff, uint8(port>>8), uint8(port))
+
+	_, err = conn.Write(buff)
 	if err != nil {
 		return nil, err
 	}
-	connClose = nil
+
+	closeConn = nil
 	return conn, nil
 }

socks4.go

@@ -75,7 +75,7 @@ func (s *Socks4Client) Dial(network, address string) (net.Conn, error) {
 	switch network {
 	case "tcp", "tcp4", "tcp6":
 	default:
-		return nil, errors.New("socks: Socks4Client.Dial invalid network:" + network)
+		return nil, errors.New("socks: no support for SOCKS4 proxy connections of type:" + network)
 	}
 
 	host, portStr, err := net.SplitHostPort(address)
@@ -111,7 +111,7 @@ func (s *Socks4Client) Dial(network, address string) (net.Conn, error) {
 
 	buff := make([]byte, 0, 8+len(s.userID)+1)
 	buff = append(buff, socks4Version, socks4Connect)
-	buff = append(buff, byte(port)>>8, byte(port))
+	buff = append(buff, byte(port>>8), byte(port))
 	buff = append(buff, ip4...)
 	if len(s.userID) != 0 {
 		buff = append(buff, []byte(s.userID)...)