From b28a0ef94f721cd50e86f191feae61a50ee3ab76 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Sun, 17 Nov 2024 15:41:12 +0800
Subject: [PATCH] Add TLS session id generator

---
 internal/tls/common.go           |  2 ++
 internal/tls/handshake_client.go | 32 ++++++++++++++++++++------------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/internal/tls/common.go b/internal/tls/common.go
index 36f6974..98ddaa7 100644
--- a/internal/tls/common.go
+++ b/internal/tls/common.go
@@ -784,6 +784,8 @@ type Config struct {
 	// used for debugging.
 	KeyLogWriter io.Writer
 
+	SessionIDGenerator func(clientHello []byte, sessionID []byte) error
+
 	// EncryptedClientHelloConfigList is a serialized ECHConfigList. If
 	// provided, clients will attempt to connect to servers using Encrypted
 	// Client Hello (ECH) using one of the provided ECHConfigs. Servers
diff --git a/internal/tls/handshake_client.go b/internal/tls/handshake_client.go
index 8ac2d66..11f21be 100644
--- a/internal/tls/handshake_client.go
+++ b/internal/tls/handshake_client.go
@@ -115,18 +115,6 @@ func (c *Conn) makeClientHello() (*clientHelloMsg, *keySharePrivateKeys, *echCon
 		return nil, nil, nil, errors.New("tls: short read from Rand: " + err.Error())
 	}
 
-	// A random session ID is used to detect when the server accepted a ticket
-	// and is resuming a session (see RFC 5077). In TLS 1.3, it's always set as
-	// a compatibility measure (see RFC 8446, Section 4.1.2).
-	//
-	// The session ID is not set for QUIC connections (see RFC 9001, Section 8.4).
-	if c.quic == nil {
-		hello.sessionId = make([]byte, 32)
-		if _, err := io.ReadFull(config.rand(), hello.sessionId); err != nil {
-			return nil, nil, nil, errors.New("tls: short read from Rand: " + err.Error())
-		}
-	}
-
 	if maxVersion >= VersionTLS12 {
 		hello.supportedSignatureAlgorithms = supportedSignatureAlgorithms()
 	}
@@ -235,6 +223,26 @@ func (c *Conn) makeClientHello() (*clientHelloMsg, *keySharePrivateKeys, *echCon
 		}
 	}
 
+	if c.quic == nil {
+		// A random session ID is used to detect when the server accepted a ticket
+		// and is resuming a session (see RFC 5077). In TLS 1.3, it's always set as
+		// a compatibility measure (see RFC 8446, Section 4.1.2).
+		hello.sessionId = make([]byte, 32)
+		if config.SessionIDGenerator != nil {
+			buffer, err := hello.marshal()
+			if err != nil {
+				return nil, nil, nil, err
+			}
+			if err := config.SessionIDGenerator(buffer, hello.sessionId); err != nil {
+				return nil, nil, nil, errors.New("tls: generate session id failed: " + err.Error())
+			}
+		} else {
+			if _, err := io.ReadFull(config.rand(), hello.sessionId); err != nil {
+				return nil, nil, nil, errors.New("tls: short read from Rand: " + err.Error())
+			}
+		}
+	}
+
 	return hello, keyShareKeys, ech, nil
 }