Update tcp.go

This commit is contained in:
Jason
2019-08-11 19:49:53 +08:00
parent 575eb72e95
commit 513f6f2966

View File

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