diff --git a/proxy/tcp.go b/proxy/tcp.go index 4e0bce1..9e47005 100644 --- a/proxy/tcp.go +++ b/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 } diff --git a/proxy/udp.go b/proxy/udp.go index 66d622f..6dac6f0 100644 --- a/proxy/udp.go +++ b/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) } diff --git a/proxy/utils.go b/proxy/utils.go index 50428e8..fcada58 100644 --- a/proxy/utils.go +++ b/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 -}