update echo & redirect

proxy/echo/tcp.go

@@ -16,11 +16,11 @@ func NewTCPHandler() core.TCPConnHandler {
 	return &tcpHandler{}
 }
 
-func (h *tcpHandler) echoBack(conn net.Conn) {
-	io.Copy(conn, conn)
+func (h *tcpHandler) echoBack(localConn net.Conn) {
+	io.Copy(localConn, localConn)
 }
 
-func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
-	go h.echoBack(conn)
+func (h *tcpHandler) Handle(localConn net.Conn, target *net.TCPAddr) error {
+	go h.echoBack(localConn)
 	return nil
 }

proxy/echo/udp.go

@@ -14,15 +14,15 @@ func NewUDPHandler() core.UDPConnHandler {
 	return &udpHandler{}
 }
 
-func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
+func (h *udpHandler) Connect(localConn 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(localConn core.UDPConn, data []byte, addr *net.UDPAddr) error {
 	// Dispatch to another goroutine, otherwise will result in deadlock.
 	payload := append([]byte(nil), data...)
 	go func(b []byte) {
-		_, err := conn.WriteFrom(b, addr)
+		_, err := localConn.WriteFrom(b, addr)
 		if err != nil {
 			log.Warnf("failed to echo back data: %v", err)
 		}

proxy/redirect/tcp.go

@@ -1,14 +1,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"
+	. "github.com/xjasonlyu/tun2socks/proxy"
 )
 
 // To do a benchmark using iperf3 locally, you may follow these steps:
@@ -38,43 +35,17 @@ func NewTCPHandler(target string) core.TCPConnHandler {
 	return &tcpHandler{target: target}
 }
 
-func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
-	c, err := net.Dial("tcp", h.target)
+func (h *tcpHandler) Handle(localConn net.Conn, target *net.TCPAddr) error {
+	remoteConn, err := net.Dial("tcp", h.target)
 	if err != nil {
 		return err
 	}
 
-	// WaitGroup
-	var wg sync.WaitGroup
-	wg.Add(2)
+	// set keepalive
+	TCPKeepAlive(localConn)
+	TCPKeepAlive(remoteConn)
 
-	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)
+	go TCPRelay(localConn, remoteConn)
 
 	log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String())
 	return nil

proxy/redirect/udp.go

@@ -52,7 +52,7 @@ func (h *udpHandler) fetchUDPInput(conn core.UDPConn, pc *net.UDPConn) {
 	}
 }
 
-func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
+func (h *udpHandler) Connect(localConn core.UDPConn, target *net.UDPAddr) error {
 	bindAddr := &net.UDPAddr{IP: nil, Port: 0}
 	pc, err := net.ListenUDP("udp", bindAddr)
 	if err != nil {
@@ -61,24 +61,24 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
 	}
 	targetAddr, _ := net.ResolveUDPAddr("udp", h.target)
 
-	h.remoteAddrMap.Store(conn, targetAddr)
-	h.remoteUDPConnMap.Store(conn, pc)
+	h.remoteAddrMap.Store(localConn, targetAddr)
+	h.remoteUDPConnMap.Store(localConn, pc)
 
-	go h.fetchUDPInput(conn, pc)
+	go h.fetchUDPInput(localConn, pc)
 
 	log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String())
 	return nil
 }
 
-func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr) error {
+func (h *udpHandler) ReceiveTo(localConn core.UDPConn, data []byte, addr *net.UDPAddr) error {
 	var pc *net.UDPConn
 	var targetAddr *net.UDPAddr
 
-	if value, ok := h.remoteAddrMap.Load(conn); ok {
+	if value, ok := h.remoteAddrMap.Load(localConn); ok {
 		targetAddr = value.(*net.UDPAddr)
 	}
 
-	if value, ok := h.remoteUDPConnMap.Load(conn); ok {
+	if value, ok := h.remoteUDPConnMap.Load(localConn); ok {
 		pc = value.(*net.UDPConn)
 	}
 
@@ -90,7 +90,7 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
 		}
 		return nil
 	} else {
-		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), addr))
+		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", localConn.LocalAddr(), addr))
 	}
 }