diff --git a/proxy/util.go b/proxy/util.go index c29afcc..0a0e5f0 100644 --- a/proxy/util.go +++ b/proxy/util.go @@ -4,6 +4,7 @@ import ( "io" "net" "sync" + "sync/atomic" "time" "github.com/golang/glog" @@ -20,19 +21,26 @@ const ( func relay(conn1, conn2 net.Conn) { wg := &sync.WaitGroup{} + exitFlag := new(int32) wg.Add(2) - go redirect(conn2, conn1, wg) - redirect(conn1, conn2, wg) + go redirect(conn2, conn1, wg, exitFlag) + redirect(conn1, conn2, wg, exitFlag) wg.Wait() } -func redirect(dst, src net.Conn, wg *sync.WaitGroup) { +func redirect(dst, src net.Conn, wg *sync.WaitGroup, exitFlag *int32) { if _, err := io.Copy(dst, src); err != nil { glog.V(1).Infof("%s<>%s -> %s<>%s: %s", src.RemoteAddr(), src.LocalAddr(), dst.LocalAddr(), dst.RemoteAddr(), err) } - now := time.Now() - src.SetReadDeadline(now) - dst.SetWriteDeadline(now) + if atomic.CompareAndSwapInt32(exitFlag, 0, 1) { + now := time.Now() + src.SetReadDeadline(now) + dst.SetWriteDeadline(now) + } else { + src.Close() + dst.Close() + } + wg.Done() }