From cbf02ab408c2a1e622facce48124d1f7a95c95e3 Mon Sep 17 00:00:00 2001
From: Amir Khan <a.khan@psiphon.ca>
Date: Wed, 26 Jun 2024 16:46:26 -0400
Subject: [PATCH] Revert `(*UConn).BuildHandshakeState` to lock session
 controller

This partially reverts ebe5d664d2fd5952a05e737447ceeefe62fc6c58
and introduces BuildHandshakeStateWithoutSession.
---
 u_conn.go      | 50 ++++++++++++++++++++++++++------------------------
 u_conn_test.go |  5 -----
 2 files changed, 26 insertions(+), 29 deletions(-)

diff --git a/u_conn.go b/u_conn.go
index babb0835..3b519de7 100644
--- a/u_conn.go
+++ b/u_conn.go
@@ -83,9 +83,33 @@ func UClient(conn net.Conn, config *Config, clientHelloID ClientHelloID) *UConn
 //	[each call] marshal ClientHello.
 //
 // BuildHandshakeState is automatically called before uTLS performs handshake,
-// amd should only be called explicitly to inspect/change fields of
+// and should only be called explicitly to inspect/change fields of
 // default/mimicked ClientHello.
+// With the excpetion of session ticket and psk extensions, which cannot be changed
+// after calling BuildHandshakeState, all other fields can be modified.
 func (uconn *UConn) BuildHandshakeState() error {
+	err := uconn.BuildHandshakeStateWithoutSession()
+	if err != nil {
+		return err
+	}
+	if uconn.ClientHelloID != HelloGolang {
+		err := uconn.uLoadSession()
+		if err != nil {
+			return err
+		}
+
+		uconn.uApplyPatch()
+
+		uconn.sessionController.finalCheck()
+	}
+	return nil
+}
+
+// BuildHandshakeStateWithoutSession is the same as BuildHandshakeState, but does not
+// set the session. This is only useful when you want to inspect the ClientHello before
+// setting the session manually through SetSessionTicketExtension or SetPSKExtension.
+// BuildHandshakeState is automatically called before uTLS performs handshake.
+func (uconn *UConn) BuildHandshakeStateWithoutSession() error {
 	if uconn.ClientHelloID == HelloGolang {
 		if uconn.clientHelloBuildStatus == BuildByGoTLS {
 			return nil
@@ -129,23 +153,8 @@ func (uconn *UConn) BuildHandshakeState() error {
 		if err != nil {
 			return err
 		}
-
-	}
-	return nil
-}
-
-func (uconn *UConn) lockSessionState() error {
-
-	err := uconn.uLoadSession()
-	if err != nil {
-		return err
+		uconn.clientHelloBuildStatus = BuildByUtls
 	}
-
-	uconn.uApplyPatch()
-
-	uconn.sessionController.finalCheck()
-	uconn.clientHelloBuildStatus = BuildByUtls
-
 	return nil
 }
 
@@ -364,10 +373,6 @@ 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)
@@ -993,9 +998,6 @@ 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 a8166cab..5709841b 100644
--- a/u_conn_test.go
+++ b/u_conn_test.go
@@ -512,11 +512,6 @@ 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 {