optimize proxy

proxy/tcp.go

@@ -1,13 +1,16 @@
 package proxy
 
 import (
+	"io"
 	"net"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/xjasonlyu/tun2socks/common/dns"
 	"github.com/xjasonlyu/tun2socks/common/log"
 	"github.com/xjasonlyu/tun2socks/common/lsof"
+	"github.com/xjasonlyu/tun2socks/common/pool"
 	"github.com/xjasonlyu/tun2socks/common/stats"
 	"github.com/xjasonlyu/tun2socks/core"
 	"github.com/xjasonlyu/tun2socks/proxy/socks"
@@ -30,7 +33,59 @@ func NewTCPHandler(proxyHost string, proxyPort int, fakeDns dns.FakeDns, session
 	}
 }
 
-func (h *tcpHandler) Handle(localConn net.Conn, target *net.TCPAddr) error {
+func (h *tcpHandler) relay(localConn, remoteConn net.Conn) {
+	var once sync.Once
+	closeOnce := func() {
+		once.Do(func() {
+			localConn.Close()
+			remoteConn.Close()
+		})
+	}
+
+	// Close
+	defer closeOnce()
+
+	// WaitGroup
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	// Up Link
+	go func() {
+		buf := pool.BufPool.Get().([]byte)
+		defer pool.BufPool.Put(buf[:cap(buf)])
+		if _, err := io.CopyBuffer(remoteConn, localConn, buf); err != nil {
+			closeOnce()
+		} else {
+			localConn.SetDeadline(time.Now())
+			remoteConn.SetDeadline(time.Now())
+			tcpCloseRead(remoteConn)
+		}
+		wg.Done()
+	}()
+
+	// Down Link
+	buf := pool.BufPool.Get().([]byte)
+	if _, err := io.CopyBuffer(localConn, remoteConn, buf); err != nil {
+		closeOnce()
+	} else {
+		localConn.SetDeadline(time.Now())
+		remoteConn.SetDeadline(time.Now())
+		tcpCloseRead(localConn)
+	}
+	pool.BufPool.Put(buf[:cap(buf)])
+
+	wg.Wait() // Wait for Up Link done
+
+	// Remove session
+	if h.sessionStater != nil {
+		h.sessionStater.RemoveSession(localConn)
+	}
+}
+
+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 {
@@ -66,20 +121,13 @@ func (h *tcpHandler) Handle(localConn net.Conn, target *net.TCPAddr) error {
 		remoteConn = stats.NewSessionConn(remoteConn, sess)
 	}
 
-	// set keepalive
+	// Set keepalive
 	tcpKeepAlive(localConn)
 	tcpKeepAlive(remoteConn)
 
-	go func() {
-		// relay connections
-		tcpRelay(localConn, remoteConn)
+	// Relay connections
+	go h.relay(localConn, remoteConn)
 
-		// remove session
-		if h.sessionStater != nil {
-			h.sessionStater.RemoveSession(localConn)
-		}
-	}()
-
-	log.Access(process, "proxy", target.Network(), localConn.LocalAddr().String(), targetAddr)
+	log.Access(process, "proxy", "tcp", localConn.LocalAddr().String(), targetAddr)
 	return nil
 }

proxy/udp.go

@@ -117,7 +117,14 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
 	return nil
 }
 
-func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr) error {
+func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr) (err error) {
+	// Close if return error
+	defer func() {
+		if err != nil {
+			h.Close(conn)
+		}
+	}()
+
 	var remoteAddr net.Addr
 	var remoteConn net.PacketConn
 
@@ -130,12 +137,10 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
 	}
 
 	if remoteAddr == nil || remoteConn == nil {
-		h.Close(conn)
 		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), addr))
 	}
 
-	if _, err := remoteConn.WriteTo(data, remoteAddr); err != nil {
-		h.Close(conn)
+	if _, err = remoteConn.WriteTo(data, remoteAddr); err != nil {
 		return errors.New(fmt.Sprintf("write remote failed: %v", err))
 	}
 
@@ -143,6 +148,7 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
 }
 
 func (h *udpHandler) Close(conn core.UDPConn) {
+	// Close
 	conn.Close()
 
 	if remoteConn, ok := h.remoteConnMap.Load(conn); ok {
@@ -152,6 +158,7 @@ func (h *udpHandler) Close(conn core.UDPConn) {
 
 	h.remoteAddrMap.Delete(conn)
 
+	// Remove session
 	if h.sessionStater != nil {
 		h.sessionStater.RemoveSession(conn)
 	}

proxy/utils.go

@@ -1,12 +1,8 @@
 package proxy
 
 import (
-	"io"
 	"net"
-	"sync"
 	"time"
-
-	"github.com/xjasonlyu/tun2socks/common/pool"
 )
 
 type duplexConn interface {
@@ -33,47 +29,3 @@ func tcpKeepAlive(conn net.Conn) {
 		tcp.SetKeepAlivePeriod(30 * time.Second)
 	}
 }
-
-func tcpRelay(localConn, remoteConn net.Conn) {
-	var once sync.Once
-	closeOnce := func() {
-		once.Do(func() {
-			localConn.Close()
-			remoteConn.Close()
-		})
-	}
-
-	// Close
-	defer closeOnce()
-
-	// WaitGroup
-	var wg sync.WaitGroup
-	wg.Add(1)
-
-	// Up Link
-	go func() {
-		buf := pool.BufPool.Get().([]byte)
-		defer pool.BufPool.Put(buf[:cap(buf)])
-		if _, err := io.CopyBuffer(remoteConn, localConn, buf); err != nil {
-			closeOnce()
-		} else {
-			localConn.SetDeadline(time.Now())
-			remoteConn.SetDeadline(time.Now())
-			tcpCloseRead(remoteConn)
-		}
-		wg.Done()
-	}()
-
-	// Down Link
-	buf := pool.BufPool.Get().([]byte)
-	if _, err := io.CopyBuffer(localConn, remoteConn, buf); err != nil {
-		closeOnce()
-	} else {
-		localConn.SetDeadline(time.Now())
-		remoteConn.SetDeadline(time.Now())
-		tcpCloseRead(localConn)
-	}
-	pool.BufPool.Put(buf[:cap(buf)])
-
-	wg.Wait() // Wait for Up Link done
-}