From 80faa4aaf8a2fd747e17d7599647744abd59b57f Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Wed, 23 Nov 2022 21:19:52 +0000
Subject: [PATCH] feat: fast open

---
 .gitignore         |  4 +++-
 cmd/client.go      |  2 +-
 cmd/config.go      |  1 +
 pkg/core/client.go | 44 +++++++++++++++++++++++++++++++++-----------
 4 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/.gitignore b/.gitignore
index d1441e5..04b3409 100644
--- a/.gitignore
+++ b/.gitignore
@@ -183,4 +183,6 @@ cmd/relay/*.json
 hy_linux
 .vscode
 
-/build/
\ No newline at end of file
+/build/
+
+config*.json
\ No newline at end of file
diff --git a/cmd/client.go b/cmd/client.go
index 9bfbac9..cf3c4fa 100644
--- a/cmd/client.go
+++ b/cmd/client.go
@@ -134,7 +134,7 @@ func client(config *clientConfig) {
 	up, down, _ := config.Speed()
 	for {
 		try += 1
-		c, err := core.NewClient(config.Server, auth, tlsConfig, quicConfig, pktConnFunc, up, down,
+		c, err := core.NewClient(config.Server, auth, tlsConfig, quicConfig, pktConnFunc, up, down, config.FastOpen,
 			func(err error) {
 				if config.QuitOnDisconnect {
 					logrus.WithFields(logrus.Fields{
diff --git a/cmd/config.go b/cmd/config.go
index c4206d4..0356cc1 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -223,6 +223,7 @@ type clientConfig struct {
 	ReceiveWindowConn   uint64 `json:"recv_window_conn"`
 	ReceiveWindow       uint64 `json:"recv_window"`
 	DisableMTUDiscovery bool   `json:"disable_mtu_discovery"`
+	FastOpen            bool   `json:"fast_open"`
 	Resolver            string `json:"resolver"`
 	ResolvePreference   string `json:"resolve_preference"`
 }
diff --git a/pkg/core/client.go b/pkg/core/client.go
index e48260d..f6dc0ca 100644
--- a/pkg/core/client.go
+++ b/pkg/core/client.go
@@ -29,6 +29,7 @@ type Client struct {
 
 	sendBPS, recvBPS uint64
 	auth             []byte
+	fastOpen         bool
 
 	tlsConfig  *tls.Config
 	quicConfig *quic.Config
@@ -48,7 +49,8 @@ type Client struct {
 }
 
 func NewClient(serverAddr string, auth []byte, tlsConfig *tls.Config, quicConfig *quic.Config,
-	pktConnFunc pktconns.ClientPacketConnFunc, sendBPS uint64, recvBPS uint64, quicReconnectFunc func(err error),
+	pktConnFunc pktconns.ClientPacketConnFunc, sendBPS uint64, recvBPS uint64, fastOpen bool,
+	quicReconnectFunc func(err error),
 ) (*Client, error) {
 	quicConfig.DisablePathMTUDiscovery = quicConfig.DisablePathMTUDiscovery || pmtud.DisablePathMTUDiscovery
 	c := &Client{
@@ -56,6 +58,7 @@ func NewClient(serverAddr string, auth []byte, tlsConfig *tls.Config, quicConfig
 		sendBPS:           sendBPS,
 		recvBPS:           recvBPS,
 		auth:              auth,
+		fastOpen:          fastOpen,
 		tlsConfig:         tlsConfig,
 		quicConfig:        quicConfig,
 		pktConnFunc:       pktConnFunc,
@@ -221,21 +224,26 @@ func (c *Client) DialTCP(addr string) (net.Conn, error) {
 		_ = stream.Close()
 		return nil, err
 	}
-	// Read response
-	var sr serverResponse
-	err = struc.Unpack(stream, &sr)
-	if err != nil {
-		_ = stream.Close()
-		return nil, err
-	}
-	if !sr.OK {
-		_ = stream.Close()
-		return nil, fmt.Errorf("connection rejected: %s", sr.Message)
+	// If fast open is enabled, we return the stream immediately
+	// and defer the response handling to the first Read() call
+	if !c.fastOpen {
+		// Read response
+		var sr serverResponse
+		err = struc.Unpack(stream, &sr)
+		if err != nil {
+			_ = stream.Close()
+			return nil, err
+		}
+		if !sr.OK {
+			_ = stream.Close()
+			return nil, fmt.Errorf("connection rejected: %s", sr.Message)
+		}
 	}
 	return &hyTCPConn{
 		Orig:             stream,
 		PseudoLocalAddr:  session.LocalAddr(),
 		PseudoRemoteAddr: session.RemoteAddr(),
+		Established:      !c.fastOpen,
 	}, nil
 }
 
@@ -306,9 +314,23 @@ type hyTCPConn struct {
 	Orig             quic.Stream
 	PseudoLocalAddr  net.Addr
 	PseudoRemoteAddr net.Addr
+	Established      bool
 }
 
 func (w *hyTCPConn) Read(b []byte) (n int, err error) {
+	if !w.Established {
+		var sr serverResponse
+		err := struc.Unpack(w.Orig, &sr)
+		if err != nil {
+			_ = w.Close()
+			return 0, err
+		}
+		if !sr.OK {
+			_ = w.Close()
+			return 0, fmt.Errorf("connection rejected: %s", sr.Message)
+		}
+		w.Established = true
+	}
 	return w.Orig.Read(b)
 }