From ebe5d664d2fd5952a05e737447ceeefe62fc6c58 Mon Sep 17 00:00:00 2001
From: Amir Khan <a.khan@psiphon.ca>
Date: Wed, 19 Jun 2024 13:41:08 -0400
Subject: [PATCH] Lock sessionController only on last call to
 BuildHandshakeState

---
 u_conn.go      | 29 +++++++++++++++++++++--------
 u_conn_test.go |  5 +++++
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/u_conn.go b/u_conn.go
index 6f531c5c..babb0835 100644
--- a/u_conn.go
+++ b/u_conn.go
@@ -125,21 +125,27 @@ func (uconn *UConn) BuildHandshakeState() error {
 			return err
 		}
 
-		err = uconn.uLoadSession()
-		if err != nil {
-			return err
-		}
-
 		err = uconn.MarshalClientHello()
 		if err != nil {
 			return err
 		}
 
-		uconn.uApplyPatch()
+	}
+	return nil
+}
 
-		uconn.sessionController.finalCheck()
-		uconn.clientHelloBuildStatus = BuildByUtls
+func (uconn *UConn) lockSessionState() error {
+
+	err := uconn.uLoadSession()
+	if err != nil {
+		return err
 	}
+
+	uconn.uApplyPatch()
+
+	uconn.sessionController.finalCheck()
+	uconn.clientHelloBuildStatus = BuildByUtls
+
 	return nil
 }
 
@@ -358,6 +364,10 @@ func (c *UConn) handshakeContext(ctx context.Context) (ret error) {
 		if err != nil {
 			return err
 		}
+		err = c.lockSessionState()
+		if err != nil {
+			return err
+		}
 	}
 	// [uTLS section ends]
 	c.handshakeErr = c.handshakeFn(handshakeCtx)
@@ -983,6 +993,9 @@ func (c *UConn) handleRenegotiation() error {
 	if err = c.BuildHandshakeState(); err != nil {
 		return err
 	}
+	if err = c.lockSessionState(); err != nil {
+		return err
+	}
 	// [uTLS section ends]
 	if c.handshakeErr = c.clientHandshake(context.Background()); c.handshakeErr == nil {
 		c.handshakes++
diff --git a/u_conn_test.go b/u_conn_test.go
index 5709841b..a8166cab 100644
--- a/u_conn_test.go
+++ b/u_conn_test.go
@@ -512,6 +512,11 @@ func (test *clientTest) runUTLS(t *testing.T, write bool, hello helloStrategy, o
 			t.Errorf("Client.BuildHandshakeState() failed: %s", err)
 			return
 		}
+		err = client.lockSessionState()
+		if err != nil {
+			t.Errorf("Client.lockSessionState() failed: %s", err)
+			return
+		}
 		// TODO: fix this name hack if we ever decide to use non-standard testing object
 		err = client.SetClientRandom([]byte("Custom ClientRandom h^xbw8bf0sn3"))
 		if err != nil {