diff --git a/proxy/echo/tcp.go b/proxy/echo/tcp.go index 1751936..675d8c9 100644 --- a/proxy/echo/tcp.go +++ b/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 } diff --git a/proxy/echo/udp.go b/proxy/echo/udp.go index a7238ff..cdeebce 100644 --- a/proxy/echo/udp.go +++ b/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) } diff --git a/proxy/redirect/tcp.go b/proxy/redirect/tcp.go index 832b5de..6399613 100644 --- a/proxy/redirect/tcp.go +++ b/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 diff --git a/proxy/redirect/udp.go b/proxy/redirect/udp.go index 0c2a57d..0aefcf9 100644 --- a/proxy/redirect/udp.go +++ b/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)) } }