shater/dns/transport/quic/packetconn_leak_test.go

package quic

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"net/url"
	"sync"
	"testing"
	"time"

	"github.com/sagernet/quic-go"
	"github.com/sagernet/quic-go/http3"
	sbTLS "github.com/sagernet/sing-box/common/tls"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/dns"
	"github.com/sagernet/sing-box/dns/transport"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing/common"
	"github.com/sagernet/sing/common/logger"
	M "github.com/sagernet/sing/common/metadata"
	N "github.com/sagernet/sing/common/network"

	mDNS "github.com/miekg/dns"
)

var _ N.Dialer = (*trackingDialer)(nil)

// These tests pin down who owns the UDP socket handed to quic-go.
//
// quic-go's Dial/DialEarly take a net.PacketConn but do NOT take ownership of
// it: quic.setupTransport() builds a Transport with createdConn=false, and
// Transport.Close() then only calls conn.SetReadDeadline(time.Now()) instead of
// conn.Close(). So every QUIC connection torn down here — idle timeout, a
// retryable error, an engine reload calling Reset() — used to strand the UDP
// socket that carried it for the rest of the process's life. On a router that
// resolves through DoQ/DoH3 for months that is an unbounded fd leak.
//
// Both tests reconnect once and assert the socket from the FIRST connection is
// actually closed. Without the `<-conn.Context().Done() -> rawConn.Close()`
// watchdogs in quic.go / http3.go they fail on that assertion.

type trackedConn struct {
	net.Conn
	closeOnce sync.Once
	closed    chan struct{}
}

func (c *trackedConn) Close() error {
	c.closeOnce.Do(func() { close(c.closed) })
	return c.Conn.Close()
}

// trackingDialer hands out real UDP sockets and remembers every one of them.
type trackingDialer struct {
	access sync.Mutex
	conns  []*trackedConn
}

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

func (d *trackingDialer) ListenPacket(ctx context.Context, destination M.Socksaddr) (net.PacketConn, error) {
	return net.ListenUDP("udp", nil)
}

func (d *trackingDialer) count() int {
	d.access.Lock()
	defer d.access.Unlock()
	return len(d.conns)
}

func (d *trackingDialer) at(index int) *trackedConn {
	d.access.Lock()
	defer d.access.Unlock()
	return d.conns[index]
}

func (d *trackingDialer) closeAll() {
	d.access.Lock()
	defer d.access.Unlock()
	for _, conn := range d.conns {
		conn.Close()
	}
}

func requireClosed(t *testing.T, conn *trackedConn, what string) {
	t.Helper()
	select {
	case <-conn.closed:
	case <-time.After(5 * time.Second):
		t.Fatalf("%s: the UDP socket of the retired QUIC connection was never closed — quic-go does not own it, we must", what)
	}
}

func requireDialed(t *testing.T, dialer *trackingDialer, want int) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if dialer.count() >= want {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("expected at least %d dial(s), got %d", want, dialer.count())
}

func testServerTLSConfig(t *testing.T, nextProtos []string) *tls.Config {
	t.Helper()
	certificate, err := sbTLS.GenerateKeyPair(nil, nil, nil, "localhost")
	if err != nil {
		t.Fatal(err)
	}
	return &tls.Config{
		Certificates: []tls.Certificate{*certificate},
		NextProtos:   nextProtos,
		MinVersion:   tls.VersionTLS13,
	}
}

func testClientTLSConfig(t *testing.T, nextProtos []string) sbTLS.Config {
	t.Helper()
	config, err := sbTLS.NewClient(context.Background(), logger.NOP(), "localhost", option.OutboundTLSOptions{
		Enabled:    true,
		Insecure:   true,
		ServerName: "localhost",
	})
	if err != nil {
		t.Fatal(err)
	}
	config.SetNextProtos(nextProtos)
	return config
}

// startDoQServer serves a minimal DoQ responder and returns its address.
func startDoQServer(t *testing.T) M.Socksaddr {
	t.Helper()
	listener, err := quic.ListenAddr("127.0.0.1:0", testServerTLSConfig(t, []string{"doq"}), nil)
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(func() {
		cancel()
		listener.Close()
	})
	go func() {
		for {
			conn, acceptErr := listener.Accept(ctx)
			if acceptErr != nil {
				return
			}
			go func(conn *quic.Conn) {
				for {
					stream, streamErr := conn.AcceptStream(ctx)
					if streamErr != nil {
						return
					}
					go func(stream *quic.Stream) {
						defer stream.Close()
						request, readErr := transport.ReadMessage(stream)
						if readErr != nil {
							return
						}
						response := new(mDNS.Msg)
						response.SetReply(request)
						transport.WriteMessage(stream, 0, response)
					}(stream)
				}
			}(conn)
		}
	}()
	return M.ParseSocksaddr(listener.Addr().String())
}

func testQuery() *mDNS.Msg {
	message := new(mDNS.Msg)
	message.SetQuestion("example.com.", mDNS.TypeA)
	return message
}

func TestQUICTransportClosesPacketConnOnReconnect(t *testing.T) {
	t.Parallel()
	serverAddr := startDoQServer(t)
	dialer := &trackingDialer{}
	t.Cleanup(dialer.closeAll)

	dnsTransport := &Transport{
		TransportAdapter: dns.NewTransportAdapter(C.DNSTypeQUIC, "test-doq", nil),
		dialer:           dialer,
		serverAddr:       serverAddr,
		tlsConfig:        testClientTLSConfig(t, []string{"doq"}),
		connection: transport.NewConnPool(transport.ConnPoolOptions[*quic.Conn]{
			Mode: transport.ConnPoolSingle,
			IsAlive: func(conn *quic.Conn) bool {
				return conn != nil && !common.Done(conn.Context())
			},
			Close: func(conn *quic.Conn, _ error) {
				conn.CloseWithError(0, "")
			},
		}),
	}
	t.Cleanup(func() { dnsTransport.Close() })

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	if _, err := dnsTransport.Exchange(ctx, testQuery()); err != nil {
		t.Fatal("first exchange: ", err)
	}
	requireDialed(t, dialer, 1)
	first := dialer.at(0)

	// Retire the connection the way a retryable error or an engine reload does.
	dnsTransport.Reset()
	requireClosed(t, first, "Reset()")

	// The reconnect must still work, on a fresh socket.
	if _, err := dnsTransport.Exchange(ctx, testQuery()); err != nil {
		t.Fatal("second exchange: ", err)
	}
	requireDialed(t, dialer, 2)
	second := dialer.at(1)
	if second == first {
		t.Fatal("expected a new UDP socket for the reconnect")
	}

	if err := dnsTransport.Close(); err != nil {
		t.Fatal(err)
	}
	requireClosed(t, second, "Close()")
}

func TestHTTP3TransportClosesPacketConnOnReconnect(t *testing.T) {
	t.Parallel()
	mux := http.NewServeMux()
	mux.HandleFunc("/dns-query", func(writer http.ResponseWriter, request *http.Request) {
		message, err := readRequestMessage(request)
		if err != nil {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		response := new(mDNS.Msg)
		response.SetReply(message)
		rawResponse, err := response.Pack()
		if err != nil {
			writer.WriteHeader(http.StatusInternalServerError)
			return
		}
		writer.Header().Set("Content-Type", transport.MimeType)
		writer.Write(rawResponse)
	})
	listener, err := quic.ListenAddrEarly("127.0.0.1:0", testServerTLSConfig(t, []string{http3.NextProtoH3}), nil)
	if err != nil {
		t.Fatal(err)
	}
	server := &http3.Server{Handler: mux}
	go server.ServeListener(listener)
	t.Cleanup(func() {
		server.Close()
		listener.Close()
	})
	serverAddr := M.ParseSocksaddr(listener.Addr().String())

	dialer := &trackingDialer{}
	t.Cleanup(dialer.closeAll)

	stdConfig := &tls.Config{
		InsecureSkipVerify: true,
		ServerName:         "localhost",
		NextProtos:         []string{http3.NextProtoH3},
		MinVersion:         tls.VersionTLS13,
	}
	dnsTransport := &HTTP3Transport{
		TransportAdapter: dns.NewTransportAdapter(C.DNSTypeHTTP3, "test-doh3", nil),
		logger:           logger.NOP(),
		dialer:           dialer,
		destination:      &url.URL{Scheme: "https", Host: "localhost", Path: "/dns-query"},
		headers:          http.Header{},
		serverAddr:       serverAddr,
		tlsConfig:        stdConfig,
	}
	dnsTransport.transport = dnsTransport.newTransport()
	t.Cleanup(func() { dnsTransport.Close() })

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	if _, err = dnsTransport.Exchange(ctx, testQuery()); err != nil {
		t.Fatal("first exchange: ", err)
	}
	requireDialed(t, dialer, 1)
	first := dialer.at(0)

	dnsTransport.Reset()
	requireClosed(t, first, "Reset()")

	if _, err = dnsTransport.Exchange(ctx, testQuery()); err != nil {
		t.Fatal("second exchange: ", err)
	}
	requireDialed(t, dialer, 2)
	second := dialer.at(1)
	if second == first {
		t.Fatal("expected a new UDP socket for the reconnect")
	}

	if err = dnsTransport.Close(); err != nil {
		t.Fatal(err)
	}
	requireClosed(t, second, "Close()")
}

func readRequestMessage(request *http.Request) (*mDNS.Msg, error) {
	defer request.Body.Close()
	rawMessage := make([]byte, 4096)
	n, err := readFull(request.Body, rawMessage)
	if err != nil {
		return nil, err
	}
	var message mDNS.Msg
	err = message.Unpack(rawMessage[:n])
	if err != nil {
		return nil, err
	}
	return &message, nil
}

func readFull(reader interface{ Read([]byte) (int, error) }, buffer []byte) (int, error) {
	var total int
	for total < len(buffer) {
		n, err := reader.Read(buffer[total:])
		total += n
		if err != nil {
			if total > 0 {
				return total, nil
			}
			return total, err
		}
	}
	return total, nil
}