fix bugs & cleanup

proxy/socks/client.go

@@ -1,110 +0,0 @@
-package socks
-
-import (
-	"fmt"
-	"io"
-	"io/ioutil"
-	"net"
-	"time"
-)
-
-var tcpTimeout = 30 * time.Second
-
-func Dial(proxy, target string) (net.Conn, error) {
-	c, err := net.DialTimeout("tcp", proxy, tcpTimeout)
-	if err != nil {
-		return nil, fmt.Errorf("%s connect error", proxy)
-	}
-
-	targetAddr := ParseAddr(target)
-	if targetAddr == nil {
-		return nil, fmt.Errorf("target address parse error")
-	}
-
-	if _, err := ClientHandshake(c, targetAddr, CmdConnect); err != nil {
-		return nil, err
-	}
-	return c, nil
-}
-
-func DialUDP(proxy, target string) (_ net.PacketConn, _ net.Addr, err error) {
-	c, err := net.DialTimeout("tcp", proxy, tcpTimeout)
-	if err != nil {
-		err = fmt.Errorf("%s connect error", proxy)
-		return
-	}
-
-	// tcp set keepalive
-	c.(*net.TCPConn).SetKeepAlive(true)
-	c.(*net.TCPConn).SetKeepAlivePeriod(tcpTimeout)
-
-	defer func() {
-		if err != nil {
-			c.Close()
-		}
-	}()
-
-	targetAddr := ParseAddr(target)
-	if targetAddr == nil {
-		err = fmt.Errorf("target address parse error")
-		return
-	}
-
-	bindAddr, err := ClientHandshake(c, targetAddr, CmdUDPAssociate)
-	if err != nil {
-		err = fmt.Errorf("%v client hanshake error", err)
-		return
-	}
-
-	addr, err := net.ResolveUDPAddr("udp", bindAddr.String())
-	if err != nil {
-		return
-	}
-
-	pc, err := net.ListenPacket("udp", "")
-	if err != nil {
-		return
-	}
-
-	go func() {
-		io.Copy(ioutil.Discard, c)
-		c.Close()
-		// A UDP association terminates when the TCP connection that the UDP
-		// ASSOCIATE request arrived on terminates. RFC1928
-		pc.Close()
-	}()
-
-	return &socksUDPConn{PacketConn: pc, tcpConn: c, targetAddr: targetAddr}, addr, nil
-}
-
-type socksUDPConn struct {
-	net.PacketConn
-	tcpConn    net.Conn
-	targetAddr Addr
-}
-
-func (c *socksUDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
-	packet, err := EncodeUDPPacket(c.targetAddr, b)
-	if err != nil {
-		return
-	}
-	return c.PacketConn.WriteTo(packet, addr)
-}
-
-func (c *socksUDPConn) ReadFrom(b []byte) (int, net.Addr, error) {
-	n, a, e := c.PacketConn.ReadFrom(b)
-	if e != nil {
-		return 0, nil, e
-	}
-	addr, payload, err := DecodeUDPPacket(b)
-	if err != nil {
-		return 0, nil, err
-	}
-	copy(b, payload)
-	return n - len(addr) - 3, a, nil
-}
-
-func (c *socksUDPConn) Close() error {
-	c.tcpConn.Close()
-	return c.PacketConn.Close()
-}

proxy/tcp.go

@@ -13,7 +13,6 @@ import (
 	"github.com/xjasonlyu/tun2socks/common/pool"
 	"github.com/xjasonlyu/tun2socks/common/stats"
 	"github.com/xjasonlyu/tun2socks/core"
-	"github.com/xjasonlyu/tun2socks/proxy/socks"
 )
 
 type tcpHandler struct {
@@ -86,18 +85,17 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 	// Alias
 	var localConn = conn
 
-	// Replace with a domain name if target address IP is a fake IP.
-	var targetHost = target.IP.String()
-	if h.fakeDns != nil {
-		if host, exist := h.fakeDns.IPToHost(target.IP); exist {
-			targetHost = host
-		}
+	// Lookup fakeDNS host record
+	targetHost, err := lookupHost(h.fakeDns, target)
+	if err != nil {
+		log.Warnf("lookup target host error: %v", err)
+		return err
 	}
 
 	proxyAddr := net.JoinHostPort(h.proxyHost, strconv.Itoa(h.proxyPort))
 	targetAddr := net.JoinHostPort(targetHost, strconv.Itoa(target.Port))
 	// Dial
-	remoteConn, err := socks.Dial(proxyAddr, targetAddr)
+	remoteConn, err := dial(proxyAddr, targetAddr)
 	if err != nil {
 		log.Warnf("Dial %v error: %v", proxyAddr, err)
 		return err

proxy/udp.go

@@ -1,7 +1,6 @@
 package proxy
 
 import (
-	"errors"
 	"fmt"
 	"net"
 	"strconv"
@@ -14,7 +13,6 @@ import (
 	"github.com/xjasonlyu/tun2socks/common/pool"
 	"github.com/xjasonlyu/tun2socks/common/stats"
 	"github.com/xjasonlyu/tun2socks/core"
-	"github.com/xjasonlyu/tun2socks/proxy/socks"
 )
 
 type udpHandler struct {
@@ -65,26 +63,17 @@ func (h *udpHandler) fetchUDPInput(conn core.UDPConn, input net.PacketConn, addr
 }
 
 func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
-	if target == nil {
-		log.Warnf("UDP target is invalid: %s", conn.LocalAddr().String())
-		return errors.New("UDP target is invalid")
-	}
-
-	// Replace with a domain name if target address IP is a fake IP.
-	var targetHost = target.IP.String()
-	if h.fakeDns != nil {
-		if host, exist := h.fakeDns.IPToHost(target.IP); exist {
-			targetHost = host
-		}
-	}
-	targetAddr := net.JoinHostPort(targetHost, strconv.Itoa(target.Port))
-	if len(targetAddr) == 0 {
-		return errors.New("target address is empty")
+	// Lookup fakeDNS host record
+	targetHost, err := lookupHost(h.fakeDns, target)
+	if err != nil {
+		log.Warnf("lookup target host error: %v", err)
+		return err
 	}
 
 	proxyAddr := net.JoinHostPort(h.proxyHost, strconv.Itoa(h.proxyPort))
+	targetAddr := net.JoinHostPort(targetHost, strconv.Itoa(target.Port))
 	// Dial
-	remoteConn, remoteAddr, err := socks.DialUDP(proxyAddr, targetAddr)
+	remoteConn, remoteAddr, err := dialUDP(proxyAddr, targetAddr)
 	if err != nil {
 		log.Warnf("DialUDP %v error: %v", proxyAddr, err)
 		return err
@@ -137,11 +126,11 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
 	}
 
 	if remoteAddr == nil || remoteConn == nil {
-		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), addr))
+		return fmt.Errorf("proxy connection %v->%v does not exists", conn.LocalAddr(), addr)
 	}
 
 	if _, err = remoteConn.WriteTo(data, remoteAddr); err != nil {
-		return errors.New(fmt.Sprintf("write remote failed: %v", err))
+		return fmt.Errorf("write remote failed: %v", err)
 	}
 
 	return nil

proxy/utils.go

@@ -1,10 +1,41 @@
 package proxy
 
 import (
+	"errors"
+	"fmt"
+	"io"
+	"io/ioutil"
 	"net"
 	"time"
+
+	"github.com/xjasonlyu/tun2socks/common/dns"
+	"github.com/xjasonlyu/tun2socks/proxy/socks"
 )
 
+// DNS lookup
+func lookupHost(fakeDns dns.FakeDns, target net.Addr) (targetHost string, err error) {
+	var targetIP net.IP
+	switch addr := target.(type) {
+	case *net.TCPAddr:
+		targetIP = addr.IP
+	case *net.UDPAddr:
+		targetIP = addr.IP
+	default:
+		err = errors.New("invalid target type")
+		return
+	}
+
+	targetHost = targetIP.String()
+	// Replace with a domain name if target address IP is a fake IP
+	if fakeDns != nil {
+		if host, exist := fakeDns.IPToHost(targetIP); exist {
+			targetHost = host
+		}
+	}
+	return
+}
+
+// TCP functions
 type duplexConn interface {
 	net.Conn
 	CloseRead() error
@@ -29,3 +60,103 @@ func tcpKeepAlive(conn net.Conn) {
 		tcp.SetKeepAlivePeriod(30 * time.Second)
 	}
 }
+
+// Socks dialer
+func dial(proxy, target string) (net.Conn, error) {
+	c, err := net.DialTimeout("tcp", proxy, 30*time.Second)
+	if err != nil {
+		return nil, fmt.Errorf("%s connect error", proxy)
+	}
+
+	targetAddr := socks.ParseAddr(target)
+	if targetAddr == nil {
+		return nil, fmt.Errorf("target address parse error")
+	}
+
+	if _, err := socks.ClientHandshake(c, targetAddr, socks.CmdConnect); err != nil {
+		return nil, err
+	}
+	return c, nil
+}
+
+func dialUDP(proxy, target string) (_ net.PacketConn, _ net.Addr, err error) {
+	c, err := net.DialTimeout("tcp", proxy, 30*time.Second)
+	if err != nil {
+		err = fmt.Errorf("%s connect error", proxy)
+		return
+	}
+
+	// tcp set keepalive
+	tcpKeepAlive(c)
+
+	defer func() {
+		if err != nil {
+			c.Close()
+		}
+	}()
+
+	targetAddr := socks.ParseAddr(target)
+	if targetAddr == nil {
+		err = fmt.Errorf("target address parse error")
+		return
+	}
+
+	bindAddr, err := socks.ClientHandshake(c, targetAddr, socks.CmdUDPAssociate)
+	if err != nil {
+		err = fmt.Errorf("%v client hanshake error", err)
+		return
+	}
+
+	addr, err := net.ResolveUDPAddr("udp", bindAddr.String())
+	if err != nil {
+		return
+	}
+
+	pc, err := net.ListenPacket("udp", "")
+	if err != nil {
+		return
+	}
+
+	go func() {
+		io.Copy(ioutil.Discard, c)
+		c.Close()
+		// A UDP association terminates when the TCP connection that the UDP
+		// ASSOCIATE request arrived on terminates. RFC1928
+		pc.Close()
+	}()
+
+	return &socksUDPConn{PacketConn: pc, tcpConn: c, targetAddr: targetAddr}, addr, nil
+}
+
+// Socks wrapped UDPConn
+type socksUDPConn struct {
+	net.PacketConn
+	tcpConn    net.Conn
+	targetAddr socks.Addr
+}
+
+func (c *socksUDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
+	packet, err := socks.EncodeUDPPacket(c.targetAddr, b)
+	if err != nil {
+		return
+	}
+	return c.PacketConn.WriteTo(packet, addr)
+}
+
+func (c *socksUDPConn) ReadFrom(b []byte) (int, net.Addr, error) {
+	n, a, e := c.PacketConn.ReadFrom(b)
+	if e != nil {
+		return 0, nil, e
+	}
+	addr, payload, err := socks.DecodeUDPPacket(b)
+	if err != nil {
+		return 0, nil, err
+	}
+	copy(b, payload)
+	return n - len(addr) - 3, a, nil
+}
+
+func (c *socksUDPConn) Close() error {
+	c.tcpConn.Close()
+	return c.PacketConn.Close()
+}