optimize tcp

proxy/d/tcp.go

@@ -4,9 +4,12 @@ import (
 	"io"
 	"net"
 	"strconv"
+	"sync"
+	"time"
 
 	"github.com/xjasonlyu/tun2socks/common/log"
 	"github.com/xjasonlyu/tun2socks/common/lsof"
+	"github.com/xjasonlyu/tun2socks/common/pool"
 	"github.com/xjasonlyu/tun2socks/core"
 )
 
@@ -28,16 +31,16 @@ import (
 // It's very important to have two default routes, and the default route to TUN should has the highest priority.
 
 type tcpHandler struct {
-	proxyHandler  core.TCPConnHandler
 	exceptionApps []string
 	sendThrough   net.Addr
+	proxyHandler  core.TCPConnHandler
 }
 
 func NewTCPHandler(proxyHandler core.TCPConnHandler, exceptionApps []string, sendThrough net.Addr) core.TCPConnHandler {
 	return &tcpHandler{
-		proxyHandler,
-		exceptionApps,
-		sendThrough,
+		exceptionApps: exceptionApps,
+		sendThrough:   sendThrough,
+		proxyHandler:  proxyHandler,
 	}
 }
 
@@ -50,19 +53,44 @@ func (h *tcpHandler) isExceptionApp(name string) bool {
 	return false
 }
 
-func (h *tcpHandler) relay(lhs, rhs net.Conn) {
-	cls := func() {
-		rhs.Close()
-		lhs.Close()
+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)
+
 	go func() {
-		io.Copy(rhs, lhs)
-		cls()
+		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())
+		}
+		wg.Done()
 	}()
 
-	io.Copy(lhs, rhs)
-	cls()
+	buf := pool.BufPool.Get().([]byte)
+	if _, err := io.CopyBuffer(localConn, remoteConn, buf); err != nil {
+		closeOnce()
+	} else {
+		localConn.SetDeadline(time.Now())
+		remoteConn.SetDeadline(time.Now())
+	}
+	pool.BufPool.Put(buf[:cap(buf)])
+
+	wg.Wait() // Wait for Up Link done
 }
 
 func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
@@ -70,7 +98,7 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 	localPortInt, _ := strconv.Atoi(localPortStr)
 	cmd, err := lsof.GetCommandNameBySocket("tcp", localHost, uint16(localPortInt))
 	if err != nil {
-		cmd = "unknown process"
+		cmd = "N/A"
 	}
 
 	if h.isExceptionApp(cmd) {

proxy/direct/tcp.go

@@ -3,8 +3,11 @@ package direct
 import (
 	"io"
 	"net"
+	"sync"
+	"time"
 
 	"github.com/xjasonlyu/tun2socks/common/log"
+	"github.com/xjasonlyu/tun2socks/common/pool"
 	"github.com/xjasonlyu/tun2socks/core"
 )
 
@@ -14,29 +17,44 @@ func NewTCPHandler() core.TCPConnHandler {
 	return &tcpHandler{}
 }
 
-func (h *tcpHandler) handleInput(conn net.Conn, input io.ReadCloser) {
-	defer func() {
-		conn.Close()
-		input.Close()
-	}()
-	io.Copy(conn, input)
-}
-
-func (h *tcpHandler) handleOutput(conn net.Conn, output io.WriteCloser) {
-	defer func() {
-		conn.Close()
-		output.Close()
-	}()
-	io.Copy(output, conn)
-}
-
 func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 	c, err := net.DialTCP("tcp", nil, target)
 	if err != nil {
 		return err
 	}
-	go h.handleInput(conn, c)
-	go h.handleOutput(conn, c)
+
+	// WaitGroup
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	var once sync.Once
+	relayCopy := func(dst, src net.Conn) {
+		closeOnce := func() {
+			once.Do(func() {
+				src.Close()
+				dst.Close()
+			})
+		}
+
+		// Close
+		defer closeOnce()
+
+		buf := pool.BufPool.Get().([]byte)
+		defer pool.BufPool.Put(buf[:cap(buf)])
+		if _, err := io.CopyBuffer(dst, src, buf); err != nil {
+			closeOnce()
+		} else {
+			src.SetDeadline(time.Now())
+			dst.SetDeadline(time.Now())
+		}
+		wg.Done()
+
+		wg.Wait() // Wait for another goroutine
+	}
+
+	go relayCopy(conn, c)
+	go relayCopy(c, conn)
+
 	log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String())
 	return nil
 }

proxy/redirect/tcp.go

@@ -3,8 +3,11 @@ package redirect
 import (
 	"io"
 	"net"
+	"sync"
+	"time"
 
 	"github.com/xjasonlyu/tun2socks/common/log"
+	"github.com/xjasonlyu/tun2socks/common/pool"
 	"github.com/xjasonlyu/tun2socks/core"
 )
 
@@ -41,23 +44,6 @@ func NewTCPHandler(target string) core.TCPConnHandler {
 	return &tcpHandler{target: target}
 }
 
-func (h *tcpHandler) handleInput(conn net.Conn, input io.ReadCloser) {
-	defer func() {
-		if tcpConn, ok := conn.(core.TCPConn); ok {
-			tcpConn.CloseWrite()
-		} else {
-			conn.Close()
-		}
-		if tcpInput, ok := input.(duplexConn); ok {
-			tcpInput.CloseRead()
-		} else {
-			input.Close()
-		}
-	}()
-
-	io.Copy(conn, input)
-}
-
 func (h *tcpHandler) handleOutput(conn net.Conn, output io.WriteCloser) {
 	defer func() {
 		if tcpConn, ok := conn.(core.TCPConn); ok {
@@ -72,7 +58,9 @@ func (h *tcpHandler) handleOutput(conn net.Conn, output io.WriteCloser) {
 		}
 	}()
 
-	io.Copy(output, conn)
+	buf := pool.BufPool.Get().([]byte)
+	io.CopyBuffer(output, conn, buf)
+	pool.BufPool.Put(buf[:cap(buf)])
 }
 
 func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
@@ -80,8 +68,39 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 	if err != nil {
 		return err
 	}
-	go h.handleInput(conn, c)
-	go h.handleOutput(conn, c)
+
+	// WaitGroup
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	var once sync.Once
+	relayCopy := func(dst, src net.Conn) {
+		closeOnce := func() {
+			once.Do(func() {
+				src.Close()
+				dst.Close()
+			})
+		}
+
+		// Close
+		defer closeOnce()
+
+		buf := pool.BufPool.Get().([]byte)
+		defer pool.BufPool.Put(buf[:cap(buf)])
+		if _, err := io.CopyBuffer(dst, src, buf); err != nil {
+			closeOnce()
+		} else {
+			src.SetDeadline(time.Now())
+			dst.SetDeadline(time.Now())
+		}
+		wg.Done()
+
+		wg.Wait() // Wait for another goroutine
+	}
+
+	go relayCopy(conn, c)
+	go relayCopy(c, conn)
+
 	log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String())
 	return nil
 }