fix(proxydial): preserve buffered CONNECT tunnel data

internal/proxydial/proxydial.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"net"
 	"net/http"
@@ -25,6 +26,23 @@ type connectDialer struct {
 	direct      *net.Dialer
 }
 
+// bufferedConn 保留 CONNECT 响应解析器预读的隧道数据，其余操作沿用底层连接。
+type bufferedConn struct {
+	net.Conn
+	reader *bufio.Reader
+}
+
+func (connection *bufferedConn) Read(data []byte) (int, error) {
+	return connection.reader.Read(data)
+}
+
+func (connection *bufferedConn) CloseWrite() error {
+	if writer, ok := connection.Conn.(interface{ CloseWrite() error }); ok {
+		return writer.CloseWrite()
+	}
+	return errors.ErrUnsupported
+}
+
 // Address 返回代理的拨号地址，HTTP 与 HTTPS 代理省略端口时补齐默认端口
 func Address(proxyURL *url.URL) string {
 	if proxyURL.Port() != "" {
@@ -125,7 +143,8 @@ func (dialer *connectDialer) DialContext(ctx context.Context, network, address s
 	if err := request.Write(connection); err != nil {
 		return nil, err
 	}
-	response, err := http.ReadResponse(bufio.NewReader(connection), request)
+	reader := bufio.NewReader(connection)
+	response, err := http.ReadResponse(reader, request)
 	if err != nil {
 		return nil, err
 	}
@@ -137,5 +156,8 @@ func (dialer *connectDialer) DialContext(ctx context.Context, network, address s
 		return nil, err
 	}
 	success = true
+	if reader.Buffered() > 0 {
+		return &bufferedConn{Conn: connection, reader: reader}, nil
+	}
 	return connection, nil
 }

internal/proxydial/relay_test.go

@@ -2,6 +2,8 @@ package proxydial
 
 import (
 	"bufio"
+	"bytes"
+	"context"
 	"encoding/base64"
 	"io"
 	"net"
@@ -13,6 +15,133 @@ import (
 	"golang.org/x/net/proxy"
 )
 
+func TestCONNECTPreservesDataAfterResponseHeaders(t *testing.T) {
+	for _, throughRelay := range []bool{false, true} {
+		name := "direct"
+		if throughRelay {
+			name = "relay"
+		}
+		t.Run(name, func(t *testing.T) {
+			listener, err := net.Listen("tcp", "127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer listener.Close()
+			payload := bytes.Repeat([]byte{0, 1, 2, 0xff}, 64)
+			serverDone := make(chan error, 1)
+			go func() {
+				conn, err := listener.Accept()
+				if err != nil {
+					serverDone <- err
+					return
+				}
+				defer conn.Close()
+				_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
+				if _, err = http.ReadRequest(bufio.NewReader(conn)); err == nil {
+					response := append([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"), payload...)
+					_, err = conn.Write(response)
+				}
+				serverDone <- err
+			}()
+			upstream := &url.URL{Scheme: "http", Host: listener.Addr().String()}
+			dialer, err := Dialer(upstream)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if throughRelay {
+				relay, err := StartRelay(upstream)
+				if err != nil {
+					t.Fatal(err)
+				}
+				defer relay.Close()
+				socks, err := proxy.SOCKS5("tcp", relay.Addr().String(), nil, &net.Dialer{})
+				if err != nil {
+					t.Fatal(err)
+				}
+				dialer = socks.(proxy.ContextDialer)
+			}
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+			conn, err := dialer.DialContext(ctx, "tcp", "example.com:443")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer conn.Close()
+			_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
+			got := make([]byte, len(payload))
+			n, readErr := io.ReadFull(conn, got)
+			if err := <-serverDone; err != nil {
+				t.Fatal(err)
+			}
+			if readErr != nil || !bytes.Equal(got, payload) {
+				t.Fatalf("tunnel read %d of %d bytes, error = %v", n, len(payload), readErr)
+			}
+		})
+	}
+}
+
+func TestCONNECTBufferedConnectionKeepsHalfClose(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listener.Close()
+	serverDone := make(chan error, 1)
+	go func() {
+		conn, err := listener.Accept()
+		if err != nil {
+			serverDone <- err
+			return
+		}
+		defer conn.Close()
+		_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
+		_, err = http.ReadRequest(bufio.NewReader(conn))
+		if err == nil {
+			_, err = io.WriteString(conn, "HTTP/1.1 200 Connection Established\r\n\r\nhello")
+		}
+		if err == nil {
+			var request []byte
+			request, err = io.ReadAll(conn)
+			if err == nil {
+				_, err = conn.Write(request)
+			}
+		}
+		serverDone <- err
+	}()
+	relay, err := StartRelay(&url.URL{Scheme: "http", Host: listener.Addr().String()})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer relay.Close()
+	dialer, err := proxy.SOCKS5("tcp", relay.Addr().String(), nil, &net.Dialer{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	conn, err := dialer.Dial("tcp", "example.com:443")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
+	greeting := make([]byte, 5)
+	if _, err := io.ReadFull(conn, greeting); err != nil || string(greeting) != "hello" {
+		t.Fatalf("greeting = %q, error = %v", greeting, err)
+	}
+	if _, err := io.WriteString(conn, "request"); err != nil {
+		t.Fatal(err)
+	}
+	if err := conn.(interface{ CloseWrite() error }).CloseWrite(); err != nil {
+		t.Fatal(err)
+	}
+	reply, err := io.ReadAll(conn)
+	if err != nil || string(reply) != "request" {
+		t.Fatalf("reply after half-close = %q, error = %v", reply, err)
+	}
+	if err := <-serverDone; err != nil {
+		t.Fatal(err)
+	}
+}
+
 // TestRelayAuthenticatedUpstream 验证中转经带账号密码的 CONNECT 上游转发域名目标，凭据错误时返回 SOCKS5 失败
 func TestRelayAuthenticatedUpstream(t *testing.T) {
 	echo, err := net.Listen("tcp", "127.0.0.1:0")