From c97bff01294928676587eee33f4e425fbcd526a9 Mon Sep 17 00:00:00 2001
From: wweir <wweir@foxmail.com>
Date: Sun, 25 Nov 2018 22:04:31 +0800
Subject: [PATCH] Fix client session rotate logic

---
 proxy/client.go | 73 +++++++++++++++++++++++++++++++++----------------
 1 file changed, 49 insertions(+), 24 deletions(-)

diff --git a/proxy/client.go b/proxy/client.go
index a0a6b57..4186d85 100644
--- a/proxy/client.go
+++ b/proxy/client.go
@@ -3,6 +3,8 @@ package proxy
 import (
 	"crypto/tls"
 	"net"
+	"sync/atomic"
+	"time"
 
 	"github.com/golang/glog"
 	"github.com/lucas-clemente/quic-go"
@@ -10,8 +12,10 @@ import (
 
 func StartClient(server string) {
 	connCh := listenLocal([]string{":80", ":443"})
+	reDialCh := make(chan net.Conn, 10)
+	var conn net.Conn
 
-	for conn := range connCh {
+	for {
 		sess, err := quic.DialAddr(server, &tls.Config{InsecureSkipVerify: true}, nil)
 		if err != nil {
 			if sess, err = quic.DialAddr(server, &tls.Config{InsecureSkipVerify: true}, nil); err != nil {
@@ -21,18 +25,55 @@ func StartClient(server string) {
 		}
 		glog.Infoln("new session to", sess.RemoteAddr())
 
-		reDial := false // no cas action, no need add lock
-		for !reDial {
-			go openStream(conn, sess, &reDial)
-			conn = <-connCh
-		}
+		for { // session rotate logic
+			select {
+			case conn = <-connCh:
+			case conn = <-reDialCh:
+			}
 
-		sess.Close()
+			if !openStream(conn, sess, reDialCh) {
+				sess.Close()
+				break
+			}
+		}
 	}
 }
 
+func openStream(conn net.Conn, sess quic.Session, reDialCh chan<- net.Conn) bool {
+	glog.V(1).Infoln("new request from", conn.RemoteAddr())
+
+	sessStat := new(int32) // 0:waiting 1:ok 2:timeout
+	go func() {
+		stream, err := sess.OpenStream()
+		if err != nil {
+			glog.Warningf("connect to remote(%s) fail:%s\n", sess.RemoteAddr(), err)
+			reDialCh <- conn
+			return
+		}
+		defer stream.Close()
+
+		if !atomic.CompareAndSwapInt32(sessStat, 0, 1) {
+			reDialCh <- conn // timeout
+		}
+
+		conn.(*net.TCPConn).SetKeepAlive(true)
+		relay(&streamConn{stream, sess}, conn)
+		conn.Close()
+	}()
+
+	// wait timeout: 10ms * 100 => 1s
+	for i := 0; i < 100; i++ {
+		if atomic.LoadInt32(sessStat) == 1 {
+			return true
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	return !atomic.CompareAndSwapInt32(sessStat, 0, 2)
+}
+
 func listenLocal(ports []string) <-chan net.Conn {
-	connCh := make(chan net.Conn)
+	connCh := make(chan net.Conn, 10)
 	for i := range ports {
 		go func(port string) {
 			ln, err := net.Listen("tcp", port)
@@ -54,19 +95,3 @@ func listenLocal(ports []string) <-chan net.Conn {
 	glog.Infoln("listening ports:", ports)
 	return connCh
 }
-
-func openStream(conn net.Conn, sess quic.Session, reDial *bool) {
-	defer conn.Close()
-	conn.(*net.TCPConn).SetKeepAlive(true)
-
-	glog.V(1).Infoln("new request from", conn.RemoteAddr())
-	stream, err := sess.OpenStream()
-	if err != nil {
-		glog.Warningf("connect to remote(%s) fail:%s\n", sess.RemoteAddr(), err)
-		*reDial = true
-		return
-	}
-	defer stream.Close()
-
-	relay(&streamConn{stream, sess}, conn)
-}