Only start a session after a local connection has been made

cmd/ck-client/ck-client.go

@@ -88,6 +88,40 @@ func makeRemoteConn(sta *client.State) (net.Conn, error) {
 
 }
 
+func makeSession(sta *client.State) *mux.Session {
+	log.Println("Attemtping to start a new session")
+	// sessionID is usergenerated. There shouldn't be a security concern because the scope of
+	// sessionID is limited to its UID.
+	quad := make([]byte, 4)
+	rand.Read(quad)
+	sta.SessionID = binary.BigEndian.Uint32(quad)
+
+	sta.UpdateIntervalKeys()
+
+	_, tthKey, _ := sta.GetIntervalKeys()
+	sesh := mux.MakeSession(sta.SessionID, mux.UNLIMITED_VALVE, mux.MakeObfs(tthKey, sta.Cipher), mux.MakeDeobfs(tthKey, sta.Cipher), util.ReadTLS)
+
+	var wg sync.WaitGroup
+	for i := 0; i < sta.NumConn; i++ {
+		wg.Add(1)
+		go func() {
+		makeconn:
+			conn, err := makeRemoteConn(sta)
+			if err != nil {
+				log.Printf("Failed to establish new connections to remote: %v\n", err)
+				time.Sleep(time.Second * 3)
+				goto makeconn
+			}
+			sesh.AddConnection(conn)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+
+	log.Printf("Session %v established", sta.SessionID)
+	return sesh
+}
+
 func main() {
 	// Should be 127.0.0.1 to listen to a proxy client on this machine
 	var localHost string
@@ -173,71 +207,23 @@ func main() {
 		}
 	}
 
-start:
-	log.Println("Attemtping to start a new session")
-	// sessionID is usergenerated. There shouldn't be a security concern because the scope of
-	// sessionID is limited to its UID.
-	quad := make([]byte, 4)
-	rand.Read(quad)
-	sta.SessionID = binary.BigEndian.Uint32(quad)
-
-	sta.UpdateIntervalKeys()
-
 	if adminUID != nil {
 		sta.SessionID = 0
 		sta.UID = adminUID
 		sta.NumConn = 1
 	}
 
-	var crypto mux.Crypto
+	var sesh *mux.Session
-	switch sta.EncryptionMethod {
-	case 0x00:
-		crypto = &mux.Plain{}
-	case 0x01:
-		crypto, err = mux.MakeAESGCMCipher(sta.UID)
-		if err != nil {
-			log.Println(err)
-			return
-		}
-	case 0x02:
-		crypto, err = mux.MakeCPCipher(sta.UID)
-		if err != nil {
-			log.Println(err)
-			return
-		}
-	}
-
-	_, tthKey, _ := sta.GetIntervalKeys()
-	sesh := mux.MakeSession(sta.SessionID, mux.UNLIMITED_VALVE, mux.MakeObfs(tthKey, crypto), mux.MakeDeobfs(tthKey, crypto), util.ReadTLS)
-
-	var wg sync.WaitGroup
-	for i := 0; i < sta.NumConn; i++ {
-		wg.Add(1)
-		go func() {
-		makeconn:
-			conn, err := makeRemoteConn(sta)
-			if err != nil {
-				log.Printf("Failed to establish new connections to remote: %v\n", err)
-				time.Sleep(time.Second * 3)
-				goto makeconn
-			}
-			sesh.AddConnection(conn)
-			wg.Done()
-		}()
-	}
-	wg.Wait()
-
-	log.Printf("Session %v established", sta.SessionID)
 
 	for {
-		if sesh.IsBroken() {
-			goto start
-		}
 		localConn, err := listener.Accept()
 		if err != nil {
 			log.Println(err)
 			continue
 		}
+		if sesh == nil || sesh.IsBroken() {
+			sesh = makeSession(sta)
+		}
 		go func() {
 			data := make([]byte, 10240)
 			i, err := io.ReadAtLeast(localConn, data, 1)

internal/client/state.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/cbeuw/Cloak/internal/ecdh"
+	mux "github.com/cbeuw/Cloak/internal/multiplex"
 )
 
 type rawConfig struct {
@@ -51,6 +52,7 @@ type State struct {
 
 	ProxyMethod      string
 	EncryptionMethod byte
+	Cipher           mux.Crypto
 	TicketTimeHint   int
 	ServerName       string
 	BrowserSig       string
@@ -140,10 +142,19 @@ func (sta *State) ParseConfig(conf string) (err error) {
 	switch preParse.EncryptionMethod {
 	case "plain":
 		sta.EncryptionMethod = 0x00
+		sta.Cipher = &mux.Plain{}
 	case "aes-gcm":
 		sta.EncryptionMethod = 0x01
+		sta.Cipher, err = mux.MakeAESGCMCipher(sta.UID)
+		if err != nil {
+			return err
+		}
 	case "chacha20-poly1305":
 		sta.EncryptionMethod = 0x02
+		sta.Cipher, err = mux.MakeCPCipher(sta.UID)
+		if err != nil {
+			return err
+		}
 	default:
 		return errors.New("Unknown encryption method")
 	}