Add ConnRF to make HTTPUpgrade 0-RTT

#3128 (comment)
XTLS · Mar 13, 2024 · 5c41292 · 5c41292
1 parent d3a218f
commit 5c41292
Showing 1 changed file with 24 additions and 12 deletions.
diff --git a/transport/internet/httpupgrade/dialer.go b/transport/internet/httpupgrade/dialer.go
@@ -15,6 +15,29 @@ import (
 	"github.com/xtls/xray-core/transport/internet/tls"
 )
 
+type ConnRF struct {
+	net.Conn
+	Req   *http.Request
+	First bool
+}
+
+func (c *ConnRF) Read(b []byte) (int, error) {
+	if c.First {
+		c.First = false
+		// TODO The bufio usage here is unreliable
+		resp, err := http.ReadResponse(bufio.NewReader(c.Conn), c.Req) // nolint:bodyclose
+		if err != nil {
+			return 0, err
+		}
+		if resp.Status != "101 Switching Protocols" ||
+			strings.ToLower(resp.Header.Get("Upgrade")) != "websocket" ||
+			strings.ToLower(resp.Header.Get("Connection")) != "upgrade" {
+			return 0, newError("unrecognized reply")
+		}
+	}
+	return c.Conn.Read(b)
+}
+
 func dialhttpUpgrade(ctx context.Context, dest net.Destination, streamSettings *internet.MemoryStreamConfig) (net.Conn, error) {
 	transportConfiguration := streamSettings.ProtocolSettings.(*Config)
 
@@ -58,18 +81,7 @@ func dialhttpUpgrade(ctx context.Context, dest net.Destination, streamSettings *
 		return nil, err
 	}
 
-	// TODO The bufio usage here is unreliable
-	resp, err := http.ReadResponse(bufio.NewReader(conn), req) // nolint:bodyclose
-	if err != nil {
-		return nil, err
-	}
-
-	if resp.Status == "101 Switching Protocols" &&
-		strings.ToLower(resp.Header.Get("Upgrade")) == "websocket" &&
-		strings.ToLower(resp.Header.Get("Connection")) == "upgrade" {
-		return conn, nil
-	}
-	return nil, newError("unrecognized reply")
+	return &ConnRF{Conn: conn, Req: req, First: true}, nil
 }
 
 func dial(ctx context.Context, dest net.Destination, streamSettings *internet.MemoryStreamConfig) (stat.Connection, error) {