From 9793537034bc6df8f1055b1e7807dfb243719927 Mon Sep 17 00:00:00 2001
From: Qian Wang <cbeuw.andy@gmail.com>
Date: Sun, 20 Jan 2019 12:13:29 +0000
Subject: [PATCH] Faster reconnection

---
 cmd/ck-client/ck-client.go    | 14 +++-----------
 internal/multiplex/session.go |  6 ++++++
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/cmd/ck-client/ck-client.go b/cmd/ck-client/ck-client.go
index 05a369e..016acfc 100644
--- a/cmd/ck-client/ck-client.go
+++ b/cmd/ck-client/ck-client.go
@@ -11,7 +11,6 @@ import (
 	"net"
 	"os"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/cbeuw/Cloak/internal/client"
@@ -167,6 +166,7 @@ func main() {
 	}
 
 start:
+	log.Println("Attemtping to start a new session")
 	var UNLIMITED int64 = 1e12
 	valve := mux.MakeValve(1e12, 1e12, &UNLIMITED, &UNLIMITED)
 	obfs := mux.MakeObfs(sta.UID)
@@ -190,10 +190,9 @@ start:
 	}
 	wg.Wait()
 
-	var broken uint32
 	for {
-		if atomic.LoadUint32(&broken) == 1 {
-			goto retry
+		if sesh.IsBroken() {
+			goto start
 		}
 		ssConn, err := listener.Accept()
 		if err != nil {
@@ -210,9 +209,6 @@ start:
 			}
 			stream, err := sesh.OpenStream()
 			if err != nil {
-				if err == mux.ErrBrokenSession {
-					atomic.StoreUint32(&broken, 1)
-				}
 				log.Println(err)
 				ssConn.Close()
 				return
@@ -228,9 +224,5 @@ start:
 			pipe(stream, ssConn)
 		}()
 	}
-retry:
-	time.Sleep(time.Second * 3)
-	log.Println("Reconnecting")
-	goto start
 
 }
diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index 8b3052e..0a37283 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -37,6 +37,7 @@ type Session struct {
 	// For accepting new streams
 	acceptCh chan *Stream
 
+	broken  uint32
 	die     chan struct{}
 	suicide sync.Once
 }
@@ -121,6 +122,7 @@ func (sesh *Session) getStream(id uint32, closingFrame bool) *Stream {
 func (sesh *Session) Close() error {
 	// Because closing a closed channel causes panic
 	sesh.suicide.Do(func() { close(sesh.die) })
+	atomic.StoreUint32(&sesh.broken, 1)
 	sesh.streamsM.Lock()
 	for id, stream := range sesh.streams {
 		// If we call stream.Close() here, streamsM will result in a deadlock
@@ -136,3 +138,7 @@ func (sesh *Session) Close() error {
 	return nil
 
 }
+
+func (sesh *Session) IsBroken() bool {
+	return atomic.LoadUint32(&sesh.broken) == 1
+}