From a7ae8b57f1a468c02a7c290f80699e145d5d443c Mon Sep 17 00:00:00 2001
From: RPRX <63339210+RPRX@users.noreply.github.com>
Date: Thu, 9 Mar 2023 20:51:17 +0800
Subject: [PATCH] Apply target's time-out to `hs.readClientFinished()`

---
 handshake_server_tls13.go |  2 ++
 tls.go                    | 28 ++++++++++++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/handshake_server_tls13.go b/handshake_server_tls13.go
index a0faede..a726e48 100644
--- a/handshake_server_tls13.go
+++ b/handshake_server_tls13.go
@@ -131,6 +131,8 @@ func (hs *serverHandshakeStateTLS13) handshake() error {
 	if _, err := c.flush(); err != nil {
 		return err
 	}
+	return nil
+
 	if err := hs.readClientCertificate(); err != nil {
 		return err
 	}
diff --git a/tls.go b/tls.go
index a45267e..f1561a4 100644
--- a/tls.go
+++ b/tls.go
@@ -32,6 +32,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pires/go-proxyproto"
@@ -354,13 +355,36 @@ func Server(ctx context.Context, conn net.Conn, config *Config) (*Conn, error) {
 				s2cSaved = s2cSaved[handshakeLen:]
 				handshakeLen = 0
 			}
+			start := time.Now()
 			err = hs.handshake()
 			if config.Show {
 				fmt.Printf("REALITY remoteAddr: %v\ths.handshake() err: %v\n", remoteAddr, err)
 			}
-			if err == nil {
-				handled = true
+			if err != nil {
+				break
 			}
+			go func() { // TODO: Probe target's maxUselessRecords and some time-outs in advance.
+				if handshakeLen-len(s2cSaved) > 0 {
+					io.ReadFull(target, buf[:handshakeLen-len(s2cSaved)])
+				}
+				if n, err := target.Read(buf); !hs.c.handshakeComplete() {
+					if err != nil {
+						conn.Close()
+					}
+					if config.Show {
+						fmt.Printf("REALITY remoteAddr: %v\ttime.Since(start): %v\tn: %v\terr: %v\n", remoteAddr, time.Since(start), n, err)
+					}
+				}
+			}()
+			err = hs.readClientFinished()
+			if config.Show {
+				fmt.Printf("REALITY remoteAddr: %v\ths.readClientFinished() err: %v\n", remoteAddr, err)
+			}
+			if err != nil {
+				break
+			}
+			atomic.StoreUint32(&hs.c.handshakeStatus, 1)
+			handled = true
 			break
 		}
 		if done {