Refactor post-handshake records detection & imitation again

https://github.com/XTLS/Xray-core/issues/4843#issuecomment-3013196642

record_detect.go

@@ -12,49 +12,47 @@ import (
 	utls "github.com/refraction-networking/utls"
 )
 
-var GlobalPostHandshakeRecordsLock sync.Mutex
+var GlobalPostHandshakeRecordsLens sync.Map
 
-var GlobalPostHandshakeRecordsLens map[*Config]map[string][]int
+func DetectPostHandshakeRecordsLens(config *Config) {
-
-func DetectPostHandshakeRecordsLens(config *Config) map[string][]int {
-	GlobalPostHandshakeRecordsLock.Lock()
-	defer GlobalPostHandshakeRecordsLock.Unlock()
-	if GlobalPostHandshakeRecordsLens == nil {
-		GlobalPostHandshakeRecordsLens = make(map[*Config]map[string][]int)
-	}
-	if GlobalPostHandshakeRecordsLens[config] == nil {
-		GlobalPostHandshakeRecordsLens[config] = make(map[string][]int)
 	for sni := range config.ServerNames {
+		key := config.Dest + " " + sni
+		if _, loaded := GlobalPostHandshakeRecordsLens.LoadOrStore(key, false); !loaded {
+			go func() {
+				defer func() {
+					val, _ := GlobalPostHandshakeRecordsLens.Load(key)
+					if _, ok := val.(bool); ok {
+						GlobalPostHandshakeRecordsLens.Store(key, []int{})
+					}
+				}()
 				target, err := net.Dial("tcp", config.Dest)
 				if err != nil {
-				continue
+					return
 				}
 				if config.Xver == 1 || config.Xver == 2 {
 					if _, err = proxyproto.HeaderProxyFromAddrs(config.Xver, target.LocalAddr(), target.RemoteAddr()).WriteTo(target); err != nil {
-					continue
+						return
 					}
 				}
 				detectConn := &DetectConn{
 					Conn: target,
-				PostHandshakeRecordsLens: GlobalPostHandshakeRecordsLens[config],
+					Key:  key,
-				Sni:                      sni,
 				}
 				uConn := utls.UClient(detectConn, &utls.Config{
-				ServerName: sni,
+					ServerName: sni, // needs new loopvar behaviour
 				}, utls.HelloChrome_Auto)
 				if err = uConn.Handshake(); err != nil {
-				continue
+					return
 				}
 				io.Copy(io.Discard, uConn)
+			}()
 		}
 	}
-	return GlobalPostHandshakeRecordsLens[config]
 }
 
 type DetectConn struct {
 	net.Conn
-	PostHandshakeRecordsLens map[string][]int
+	Key     string
-	Sni                      string
 	CcsSent bool
 }
 
@@ -71,14 +69,16 @@ func (c *DetectConn) Read(b []byte) (n int, err error) {
 	}
 	c.Conn.SetReadDeadline(time.Now().Add(5 * time.Second))
 	data, _ := io.ReadAll(c.Conn)
+	var postHandshakeRecordsLens []int
 	for {
 		if len(data) >= 5 && bytes.Equal(data[:3], []byte{23, 3, 3}) {
 			length := int(binary.BigEndian.Uint16(data[3:5])) + 5
-			c.PostHandshakeRecordsLens[c.Sni] = append(c.PostHandshakeRecordsLens[c.Sni], length)
+			postHandshakeRecordsLens = append(postHandshakeRecordsLens, length)
 			data = data[length:]
 		} else {
 			break
 		}
 	}
+	GlobalPostHandshakeRecordsLens.Store(c.Key, postHandshakeRecordsLens)
 	return 0, io.EOF
 }

tls.go

@@ -157,13 +157,9 @@ func Value(vals ...byte) (value int) {
 	return
 }
 
-// Server returns a new TLS server side connection
+// You MUST call `DetectPostHandshakeRecordsLens(config)` in advance manually
-// using conn as the underlying transport.
+// if you don't use REALITY's listener, e.g., Xray-core's RAW transport.
-// The configuration config must be non-nil and must include
-// at least one certificate or else set GetCertificate.
 func Server(ctx context.Context, conn net.Conn, config *Config) (*Conn, error) {
-	postHandshakeRecordsLens := DetectPostHandshakeRecordsLens(config)
-
 	remoteAddr := conn.RemoteAddr().String()
 	if config.Show {
 		fmt.Printf("REALITY remoteAddr: %v\n", remoteAddr)
@@ -374,7 +370,10 @@ func Server(ctx context.Context, conn net.Conn, config *Config) (*Conn, error) {
 			if err != nil {
 				break
 			}
-			for _, length := range postHandshakeRecordsLens[hs.clientHello.serverName] {
+			for {
+				if val, ok := GlobalPostHandshakeRecordsLens.Load(config.Dest + " " + hs.clientHello.serverName); ok {
+					if postHandshakeRecordsLens, ok := val.([]int); ok {
+						for _, length := range postHandshakeRecordsLens {
 							plainText := make([]byte, length-16)
 							plainText[0] = 23
 							plainText[1] = 3
@@ -389,6 +388,11 @@ func Server(ctx context.Context, conn net.Conn, config *Config) (*Conn, error) {
 								fmt.Printf("REALITY remoteAddr: %v\tlen(postHandshakeRecord): %v\n", remoteAddr, len(postHandshakeRecord))
 							}
 						}
+						break
+					}
+				}
+				time.Sleep(5 * time.Second)
+			}
 			hs.c.isHandshakeComplete.Store(true)
 			break
 		}