shater/transport/v2rayquic/client_leak_lx_test.go

//go:build with_quic

package v2rayquic

import (
	"context"
	"net"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/sagernet/quic-go"
	"github.com/sagernet/sing-box/common/tls"
	"github.com/sagernet/sing-box/log"
	"github.com/sagernet/sing-box/option"
	qtls "github.com/sagernet/sing-quic"
	M "github.com/sagernet/sing/common/metadata"

	"github.com/stretchr/testify/require"
)

// quic-go does not take ownership of the packet conn handed to Dial: when the
// QUIC connection ends it only stops reading from it. offerNew() then dials a
// fresh one and overwrites c.rawConn, so the previous UDP socket was leaked for
// the lifetime of the process — one per reconnect, on a box with 512 MB and a
// health checker that reconnects constantly. Upstream 7067276170.

const testALPN = "shater-test"

type trackedConn struct {
	net.Conn
	closed atomic.Bool
}

func (c *trackedConn) Close() error {
	c.closed.Store(true)
	return c.Conn.Close()
}

type dialRecorder struct {
	access sync.Mutex
	conns  []*trackedConn
}

func (d *dialRecorder) DialContext(ctx context.Context, network string, destination M.Socksaddr) (net.Conn, error) {
	conn, err := new(net.Dialer).DialContext(ctx, network, destination.String())
	if err != nil {
		return nil, err
	}
	tracked := &trackedConn{Conn: conn}
	d.access.Lock()
	d.conns = append(d.conns, tracked)
	d.access.Unlock()
	return tracked, nil
}

func (d *dialRecorder) ListenPacket(ctx context.Context, destination M.Socksaddr) (net.PacketConn, error) {
	return nil, net.ErrClosed
}

func (d *dialRecorder) only(t *testing.T) *trackedConn {
	t.Helper()
	d.access.Lock()
	defer d.access.Unlock()
	require.Len(t, d.conns, 1, "client must have dialed exactly once")
	return d.conns[0]
}

// serveQUIC brings up a real QUIC listener on localhost with a self-signed
// certificate, runs handler for every accepted connection, and returns its
// address.
func serveQUIC(t *testing.T, handler func(conn *quic.Conn)) M.Socksaddr {
	t.Helper()
	ctx := context.Background()
	logger := log.NewNOPFactory().NewLogger("test")

	keyPem, certificatePem, err := tls.GenerateCertificate(nil, nil, time.Now, "localhost", time.Now().Add(time.Hour))
	require.NoError(t, err)
	serverTLSConfig, err := tls.NewSTDServer(ctx, logger, option.InboundTLSOptions{
		Enabled:     true,
		ServerName:  "localhost",
		ALPN:        []string{testALPN},
		Certificate: []string{string(certificatePem)},
		Key:         []string{string(keyPem)},
	})
	require.NoError(t, err)
	require.NoError(t, serverTLSConfig.Start())
	t.Cleanup(func() { serverTLSConfig.Close() })

	packetConn, err := net.ListenPacket("udp", "127.0.0.1:0")
	require.NoError(t, err)
	t.Cleanup(func() { packetConn.Close() })

	listener, err := qtls.Listen(packetConn, serverTLSConfig, &quic.Config{})
	require.NoError(t, err)
	t.Cleanup(func() { listener.Close() })

	go func() {
		for {
			conn, acceptErr := listener.Accept(ctx)
			if acceptErr != nil {
				return
			}
			go handler(conn)
		}
	}()
	return M.ParseSocksaddr(packetConn.LocalAddr().String())
}

func newTestClient(t *testing.T, dialer *dialRecorder, serverAddr M.Socksaddr) *Client {
	t.Helper()
	clientTLSConfig, err := tls.NewSTDClient(context.Background(), log.NewNOPFactory().NewLogger("test"), "localhost", option.OutboundTLSOptions{
		Enabled:    true,
		Insecure:   true,
		ServerName: "localhost",
		ALPN:       []string{testALPN},
	})
	require.NoError(t, err)
	transport, err := NewClient(context.Background(), dialer, serverAddr, option.V2RayQUICOptions{}, clientTLSConfig)
	require.NoError(t, err)
	client, isClient := transport.(*Client)
	require.True(t, isClient)
	return client
}

func TestClientClosesPacketConnWhenConnectionEnds(t *testing.T) {
	// Drop the connection right after the handshake: this is the server-side
	// reset / idle timeout the client must survive without leaking its socket.
	serverAddr := serveQUIC(t, func(conn *quic.Conn) {
		conn.CloseWithError(0, "bye")
	})
	dialer := &dialRecorder{}
	client := newTestClient(t, dialer, serverAddr)
	t.Cleanup(func() { client.Close() })

	quicConn, err := client.offer()
	require.NoError(t, err)
	require.NotNil(t, quicConn)

	// The server hangs up; the client keeps its Client alive (a health checker
	// would simply dial again later).
	select {
	case <-quicConn.Context().Done():
	case <-time.After(5 * time.Second):
		t.Fatal("server never closed the QUIC connection")
	}

	tracked := dialer.only(t)
	require.Eventually(t, tracked.closed.Load, 5*time.Second, 10*time.Millisecond,
		"the UDP socket behind a dead QUIC connection must be closed, not leaked until Client.Close()")
}

// quic-go's Stream.Close() does not unblock a Write parked on flow control. The
// writer goroutine (for us: the copy loop of a proxied connection) then survives
// its own connection forever. Closing has to push the write deadline into the
// past as well.
func TestStreamCloseUnblocksBlockedWrite(t *testing.T) {
	serverIsDone := make(chan struct{})
	t.Cleanup(func() { close(serverIsDone) })
	// Accept the stream but never read from it, so the client's writes fill the
	// receive window and block.
	serverAddr := serveQUIC(t, func(conn *quic.Conn) {
		_, err := conn.AcceptStream(context.Background())
		if err != nil {
			return
		}
		<-serverIsDone
	})
	dialer := &dialRecorder{}
	client := newTestClient(t, dialer, serverAddr)
	t.Cleanup(func() { client.Close() })

	stream, err := client.DialContext(context.Background())
	require.NoError(t, err)

	writeDone := make(chan error, 1)
	go func() {
		payload := make([]byte, 64*1024)
		// 32 MiB is far past any quic-go receive window, so this must park.
		for range 512 {
			_, writeErr := stream.Write(payload)
			if writeErr != nil {
				writeDone <- writeErr
				return
			}
		}
		writeDone <- nil
	}()

	select {
	case err = <-writeDone:
		t.Fatal("the write never blocked, the test proves nothing: ", err)
	case <-time.After(time.Second):
	}

	require.NoError(t, stream.Close())
	select {
	case <-writeDone:
	case <-time.After(5 * time.Second):
		t.Fatal("Write stayed blocked after Close: the writer goroutine is leaked")
	}
}