Fix relay logic

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()
 }