diff --git a/proxy/socks/tcp.go b/proxy/socks/tcp.go index afd0bc7..7d712fe 100644 --- a/proxy/socks/tcp.go +++ b/proxy/socks/tcp.go @@ -45,29 +45,30 @@ func (h *tcpHandler) relay(localConn, remoteConn net.Conn) { // Close defer closeOnce() - upCh := make(chan struct{}) + up := make(chan struct{}) + down := make(chan struct{}) // UpLink go func() { if _, err := io.Copy(remoteConn, localConn); err != nil { closeOnce() } - upCh <- struct{}{} + up <- struct{}{} }() // DownLink - if _, err := io.Copy(localConn, remoteConn); err != nil { - closeOnce() - } + go func() { + if _, err := io.Copy(localConn, remoteConn); err != nil { + closeOnce() + } + down <- struct{}{} + }() select { - case <-upCh: - case <-time.After(10 * time.Second): - closeOnce() + case <-up: // Wait for UpLink done. + case <-down: } - //<-upCh // Wait for UpLink done. - if h.sessionStater != nil { h.sessionStater.RemoveSession(localConn) }