From 925bfb39ab3c20aa1336ee718a20481399f9e4d1 Mon Sep 17 00:00:00 2001
From: adotkhan <61702862+adotkhan@users.noreply.github.com>
Date: Wed, 26 Jun 2024 23:12:28 -0400
Subject: [PATCH] Allow BuildHandshakeState to inspect ClientHello before
 setting SessionTicket/PSK (#301)

* Lock sessionController only on last call to BuildHandshakeState

* Add public getter/setter for SessionState.extMasterSecret

* Fix breaking exported MakeClientSessionState

* Revert `(*UConn).BuildHandshakeState` to lock session controller

This partially reverts ebe5d664d2fd5952a05e737447ceeefe62fc6c58
and introduces BuildHandshakeStateWithoutSession.

* fix: Marshal the Client Hello after loading session


---------

Signed-off-by: Gaukas Wang <i@gaukas.wang>
---
 u_conn.go   | 30 ++++++++++++++++++++++++------
 u_public.go | 14 ++++++++++++++
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/u_conn.go b/u_conn.go
index 6f531c5c..f683c05c 100644
--- a/u_conn.go
+++ b/u_conn.go
@@ -83,9 +83,23 @@ 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 {
+	return uconn.buildHandshakeState(true)
+}
+
+// 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 {
+	return uconn.buildHandshakeState(false)
+}
+
+func (uconn *UConn) buildHandshakeState(loadSession bool) error {
 	if uconn.ClientHelloID == HelloGolang {
 		if uconn.clientHelloBuildStatus == BuildByGoTLS {
 			return nil
@@ -125,9 +139,11 @@ func (uconn *UConn) BuildHandshakeState() error {
 			return err
 		}
 
-		err = uconn.uLoadSession()
-		if err != nil {
-			return err
+		if loadSession {
+			err = uconn.uLoadSession()
+			if err != nil {
+				return err
+			}
 		}
 
 		err = uconn.MarshalClientHello()
@@ -135,9 +151,11 @@ func (uconn *UConn) BuildHandshakeState() error {
 			return err
 		}
 
-		uconn.uApplyPatch()
+		if loadSession {
+			uconn.uApplyPatch()
+			uconn.sessionController.finalCheck()
+		}
 
-		uconn.sessionController.finalCheck()
 		uconn.clientHelloBuildStatus = BuildByUtls
 	}
 	return nil
diff --git a/u_public.go b/u_public.go
index 419b679f..cdcd6eeb 100644
--- a/u_public.go
+++ b/u_public.go
@@ -617,6 +617,9 @@ func (PSS PskIdentities) ToPrivate() []pskIdentity {
 
 // ClientSessionState is public, but all its fields are private. Let's add setters, getters and constructor
 
+// TODO! can we change this enought (or export SessionState),
+// such that we wouldn't need to fork crypto/tls?
+
 // ClientSessionState contains the state needed by clients to resume TLS sessions.
 func MakeClientSessionState(
 	SessionTicket []uint8,
@@ -625,6 +628,7 @@ func MakeClientSessionState(
 	MasterSecret []byte,
 	ServerCertificates []*x509.Certificate,
 	VerifiedChains [][]*x509.Certificate) *ClientSessionState {
+	// TODO: Add EMS to this constructor in uTLS v2
 	css := &ClientSessionState{
 		ticket: SessionTicket,
 		session: &SessionState{
@@ -658,6 +662,10 @@ func (css *ClientSessionState) MasterSecret() []byte {
 	return css.session.secret
 }
 
+func (css *ClientSessionState) EMS() bool {
+	return css.session.extMasterSecret
+}
+
 // Certificate chain presented by the server
 func (css *ClientSessionState) ServerCertificates() []*x509.Certificate {
 	return css.session.peerCertificates
@@ -689,6 +697,12 @@ func (css *ClientSessionState) SetMasterSecret(MasterSecret []byte) {
 	}
 	css.session.secret = MasterSecret
 }
+func (css *ClientSessionState) SetEMS(ems bool) {
+	if css.session == nil {
+		css.session = &SessionState{}
+	}
+	css.session.extMasterSecret = ems
+}
 func (css *ClientSessionState) SetServerCertificates(ServerCertificates []*x509.Certificate) {
 	if css.session == nil {
 		css.session = &SessionState{}