URLTest delay includes the whole connection setup when WebSocket early data is enabled

package main

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/netip"
	"testing"
	"time"

	"github.com/sagernet/sing-box/adapter"
	"github.com/sagernet/sing-box/common/urltest"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing/common"
	"github.com/sagernet/sing/common/json/badoption"

	"github.com/gofrs/uuid/v5"
	"github.com/stretchr/testify/require"
)

func vlessWSInbound(port uint16, user string, certPem, keyPem string, early uint32) option.Inbound {
	transport := &option.V2RayTransportOptions{
		Type:             C.V2RayTransportTypeWebsocket,
		WebsocketOptions: option.V2RayWebsocketOptions{Path: "/"},
	}
	if early > 0 {
		transport.WebsocketOptions.MaxEarlyData = early
		transport.WebsocketOptions.EarlyDataHeaderName = "Sec-WebSocket-Protocol"
	}
	return option.Inbound{
		Type: C.TypeVLESS,
		Options: &option.VLESSInboundOptions{
			ListenOptions: option.ListenOptions{
				Listen:     common.Ptr(badoption.Addr(netip.IPv4Unspecified())),
				ListenPort: port,
			},
			Users: []option.VLESSUser{{UUID: user}},
			InboundTLSOptionsContainer: option.InboundTLSOptionsContainer{
				TLS: &option.InboundTLSOptions{
					Enabled: true, ServerName: "example.org",
					CertificatePath: certPem, KeyPath: keyPem,
				},
			},
			Transport: transport,
		},
	}
}

func vlessWSOutbound(tag string, port uint16, user string, certPem string, early uint32) option.Outbound {
	transport := &option.V2RayTransportOptions{
		Type:             C.V2RayTransportTypeWebsocket,
		WebsocketOptions: option.V2RayWebsocketOptions{Path: "/"},
	}
	if early > 0 {
		transport.WebsocketOptions.MaxEarlyData = early
		transport.WebsocketOptions.EarlyDataHeaderName = "Sec-WebSocket-Protocol"
	}
	return option.Outbound{
		Type: C.TypeVLESS,
		Tag:  tag,
		Options: &option.VLESSOutboundOptions{
			ServerOptions: option.ServerOptions{Server: "127.0.0.1", ServerPort: port},
			UUID:          user,
			OutboundTLSOptionsContainer: option.OutboundTLSOptionsContainer{
				TLS: &option.OutboundTLSOptions{
					Enabled: true, ServerName: "example.org", CertificatePath: certPem,
				},
			},
			Transport: transport,
		},
	}
}

func latencyPipe(dst io.Writer, src io.Reader, delay time.Duration) {
	buffer := make([]byte, 32*1024)
	for {
		n, err := src.Read(buffer)
		if n > 0 {
			time.Sleep(delay)
			if _, writeErr := dst.Write(buffer[:n]); writeErr != nil {
				return
			}
		}
		if err != nil {
			return
		}
	}
}

// latencyProxy forwards listenPort to targetPort, adding delay to the connect
// and to every chunk in each direction, so one round trip costs 2*delay.
func latencyProxy(t *testing.T, listenPort, targetPort uint16, delay time.Duration) {
	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", listenPort))
	require.NoError(t, err)
	t.Cleanup(func() { listener.Close() })
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			go func() {
				defer conn.Close()
				time.Sleep(delay)
				upstream, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", targetPort))
				if err != nil {
					return
				}
				defer upstream.Close()
				go latencyPipe(conn, upstream, delay)
				latencyPipe(upstream, conn, delay)
			}()
		}
	}()
}

func TestURLTestEarlyDataDelay(t *testing.T) {
	user, err := uuid.DefaultGenerator.NewV4()
	require.NoError(t, err)
	_, certPem, keyPem := createSelfSignedCertificate(t, "example.org")

	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", testPort))
	require.NoError(t, err)
	server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})}
	go server.Serve(listener)
	defer server.Close()

	const hop = 25 * time.Millisecond
	latencyProxy(t, otherClientPort, serverPort, hop)
	latencyProxy(t, otherClientPort+1, otherPort, hop)

	instance := startInstance(t, option.Options{
		Inbounds: []option.Inbound{
			vlessWSInbound(serverPort, user.String(), certPem, keyPem, 0),
			vlessWSInbound(otherPort, user.String(), certPem, keyPem, 2560),
		},
		Outbounds: []option.Outbound{
			{Type: C.TypeDirect},
			vlessWSOutbound("plain", otherClientPort, user.String(), certPem, 0),
			vlessWSOutbound("early", otherClientPort+1, user.String(), certPem, 2560),
		},
	})

	link := fmt.Sprintf("http://127.0.0.1:%d/", testPort)
	for _, tag := range []string{"plain", "early", "plain", "early", "plain", "early"} {
		out, ok := instance.Outbound().Outbound(tag)
		require.True(t, ok)
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		delay, err := urltest.URLTest(ctx, link, out.(adapter.Outbound))
		cancel()
		require.NoError(t, err)
		t.Logf("URLTest delay via %-6s = %d ms", tag, delay)
	}
}