Pull request #412: AGDNS-3434-impl-0RTT-tests

Merge in GO/dnsproxy from AGDNS-3434-impl-0RTT-tests to master

Squashed commit of the following:

commit 3413457
Author: f.setrakov <[email protected]>
Date:   Fri Nov 21 15:46:18 2025 +0300

    upstream: use defer for mutext unlocking

commit b148fd5
Author: f.setrakov <[email protected]>
Date:   Fri Nov 21 12:03:27 2025 +0300

    upstream: imp 0RTT-tests

commit c3c0025
Author: f.setrakov <[email protected]>
Date:   Thu Nov 20 12:34:02 2025 +0300

    upstream: impl 0RTT tests

Author: masterkusok

upstream/doh.go

@@ -98,13 +98,19 @@ func newDoH(addr *url.URL, opts *Options) (u Upstream, err error) {
 		httpVersions = DefaultHTTPVersions
 	}
 
+	quicConf := &quic.Config{
+		KeepAlivePeriod: QUICKeepAlivePeriod,
+		TokenStore:      newQUICTokenStore(),
+	}
+
+	if opts.QUICTracer != nil {
+		quicConf.Tracer = opts.QUICTracer.TraceForConnection
+	}
+
 	ups := &dnsOverHTTPS{
-		getDialer: newDialerInitializer(addr, opts),
-		addr:      addr,
-		quicConf: &quic.Config{
-			KeepAlivePeriod: QUICKeepAlivePeriod,
-			TokenStore:      newQUICTokenStore(),
-		},
+		getDialer:  newDialerInitializer(addr, opts),
+		addr:       addr,
+		quicConf:   quicConf,
 		quicConfMu: &sync.Mutex{},
 		tlsConf: &tls.Config{
 			ServerName:   addr.Hostname(),

upstream/doh_internal_test.go

@@ -22,7 +22,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// TODO(f.setrakov): Implement 0RTT tests.
 func TestUpstreamDoH(t *testing.T) {
 	t.Parallel()
 
@@ -269,6 +268,60 @@ func TestUpstreamDoH_serverRestart(t *testing.T) {
 	}
 }
 
+func TestUpstreamDoH_0RTT(t *testing.T) {
+	t.Parallel()
+
+	// Run the first server instance.
+	srv := startDoHServer(t, testDoHServerOptions{
+		http3Enabled: true,
+	})
+
+	// Create a DNS-over-HTTPS upstream.
+	tracer := &testTracer{}
+	address := fmt.Sprintf("h3://%s/dns-query", srv.addr)
+	u, err := AddressToUpstream(address, &Options{
+		Logger:             testLogger,
+		InsecureSkipVerify: true,
+		QUICTracer:         tracer,
+	})
+	require.NoError(t, err)
+	testutil.CleanupAndRequireSuccess(t, u.Close)
+
+	uh := testutil.RequireTypeAssert[*dnsOverHTTPS](t, u)
+	req := createTestMessage()
+
+	// Trigger connection to a DoH3 server.
+	resp, err := uh.Exchange(req)
+	require.NoError(t, err)
+	requireResponse(t, req, resp)
+
+	// Close the active connection to make sure we'll reconnect.
+	func() {
+		uh.clientMu.Lock()
+		defer uh.clientMu.Unlock()
+
+		err = uh.closeClient(uh.client)
+		require.NoError(t, err)
+
+		uh.client = nil
+	}()
+
+	// Trigger second connection.
+	resp, err = uh.Exchange(req)
+	require.NoError(t, err)
+	requireResponse(t, req, resp)
+
+	// Check traced connections info.
+	conns := tracer.connectionsInfo()
+	require.Len(t, conns, 2)
+
+	// Examine the first connection (no 0-RTT there).
+	require.False(t, conns[0].is0RTT())
+
+	// Examine the second connection (the one that used 0-RTT).
+	require.True(t, conns[1].is0RTT())
+}
+
 // testDoHServerOptions allows customizing testDoHServer behavior.
 type testDoHServerOptions struct {
 	// handler is an HTTP handler that should be used by the server.  The

upstream/doq.go

@@ -99,13 +99,19 @@ type dnsOverQUIC struct {
 func newDoQ(addr *url.URL, opts *Options) (u Upstream, err error) {
 	addPort(addr, defaultPortDoQ)
 
+	quicConf := &quic.Config{
+		KeepAlivePeriod: QUICKeepAlivePeriod,
+		TokenStore:      newQUICTokenStore(),
+	}
+
+	if opts.QUICTracer != nil {
+		quicConf.Tracer = opts.QUICTracer.TraceForConnection
+	}
+
 	u = &dnsOverQUIC{
-		getDialer: newDialerInitializer(addr, opts),
-		addr:      addr,
-		quicConfig: &quic.Config{
-			KeepAlivePeriod: QUICKeepAlivePeriod,
-			TokenStore:      newQUICTokenStore(),
-		},
+		getDialer:  newDialerInitializer(addr, opts),
+		addr:       addr,
+		quicConfig: quicConf,
 		tlsConf: &tls.Config{
 			ServerName:   addr.Hostname(),
 			RootCAs:      opts.RootCAs,

upstream/doq_internal_test.go

@@ -24,7 +24,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// TODO(f.setrakov): Implement 0RTT tests.
 func TestUpstreamDoQ(t *testing.T) {
 	tlsConf, rootCAs := createServerTLSConfig(t, "127.0.0.1")
 
@@ -183,6 +182,56 @@ func TestUpstreamDoQ_serverRestart(t *testing.T) {
 	})
 }
 
+func TestUpstreamDoQ_0RTT(t *testing.T) {
+	tlsConf, rootCAs := createServerTLSConfig(t, "127.0.0.1")
+
+	srv := startDoQServer(t, tlsConf, 0)
+
+	tracer := &testTracer{}
+	address := fmt.Sprintf("quic://%s", srv.addr)
+	u, err := AddressToUpstream(address, &Options{
+		Logger:     testLogger,
+		QUICTracer: tracer,
+		RootCAs:    rootCAs,
+	})
+	require.NoError(t, err)
+	testutil.CleanupAndRequireSuccess(t, u.Close)
+
+	uq := testutil.RequireTypeAssert[*dnsOverQUIC](t, u)
+	req := createTestMessage()
+
+	// Trigger connection to a QUIC server.
+	resp, err := uq.Exchange(req)
+	require.NoError(t, err)
+	requireResponse(t, req, resp)
+
+	// Close the active connection to make sure we'll reconnect.
+	func() {
+		uq.connMu.Lock()
+		defer uq.connMu.Unlock()
+
+		err = uq.conn.CloseWithError(QUICCodeNoError, "")
+		require.NoError(t, err)
+
+		uq.conn = nil
+	}()
+
+	// Trigger second connection.
+	resp, err = uq.Exchange(req)
+	require.NoError(t, err)
+	requireResponse(t, req, resp)
+
+	// Check traced connections info.
+	conns := tracer.connectionsInfo()
+	require.Len(t, conns, 2)
+
+	// Examine the first connection (no 0-RTT there).
+	require.False(t, conns[0].is0RTT())
+
+	// Examine the second connection (the one that used 0-RTT).
+	require.True(t, conns[1].is0RTT())
+}
+
 // testDoHServer is an instance of a test DNS-over-QUIC server.
 type testDoQServer struct {
 	// listener is the QUIC connections listener.

upstream/upstream.go

@@ -24,6 +24,8 @@ import (
 	"github.com/ameshkov/dnscrypt/v2"
 	"github.com/ameshkov/dnsstamps"
 	"github.com/miekg/dns"
+	"github.com/quic-go/quic-go"
+	"github.com/quic-go/quic-go/qlogwriter"
 )
 
 // Upstream is an interface for a DNS resolver.  All the methods must be safe
@@ -43,6 +45,17 @@ type Upstream interface {
 	io.Closer
 }
 
+// QUICTracer creates [qlogwriter.Trace] instances for QUIC connection tracing.
+type QUICTracer interface {
+	// TraceForConnection creates a [qlogwriter.Trace] specific for a given
+	// role and connection ID.
+	TraceForConnection(
+		ctx context.Context,
+		isClient bool,
+		connID quic.ConnectionID,
+	) (trace qlogwriter.Trace)
+}
+
 // Options for AddressToUpstream func.  With these options we can configure the
 // upstream properties.
 type Options struct {
@@ -63,6 +76,10 @@ type Options struct {
 	// Upstream.Exchange method returns any error caused by it.
 	VerifyDNSCryptCertificate func(cert *dnscrypt.Cert) error
 
+	// QUICTracer allows tracing every QUIC connection and logging every packet
+	// that goes through.
+	QUICTracer QUICTracer
+
 	// RootCAs is the CertPool that must be used by all upstreams.  Redefining
 	// RootCAs makes sense on iOS to overcome the 15MB memory limit of the
 	// NEPacketTunnelProvider.
@@ -102,6 +119,7 @@ func (o *Options) Clone() (clone *Options) {
 		VerifyDNSCryptCertificate: o.VerifyDNSCryptCertificate,
 		InsecureSkipVerify:        o.InsecureSkipVerify,
 		PreferIPv6:                o.PreferIPv6,
+		QUICTracer:                o.QUICTracer,
 		RootCAs:                   o.RootCAs,
 		CipherSuites:              o.CipherSuites,
 		Logger:                    o.Logger,

upstream/upstream_internal_test.go

@@ -1,6 +1,7 @@
 package upstream
 
 import (
+	"context"
 	"crypto/ecdsa"
 	"crypto/rand"
 	"crypto/rsa"
@@ -24,6 +25,9 @@ import (
 	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/ameshkov/dnsstamps"
 	"github.com/miekg/dns"
+	"github.com/quic-go/quic-go"
+	"github.com/quic-go/quic-go/qlog"
+	"github.com/quic-go/quic-go/qlogwriter"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -772,3 +776,110 @@ func publicKey(priv any) (pub any) {
 		return nil
 	}
 }
+
+// testTracer collects QUIC connection traces for testing.
+type testTracer struct {
+	tracers []*quicTracer
+}
+
+// TraceForConnection creates a tracer for a QUIC connection.
+func (t *testTracer) TraceForConnection(
+	_ context.Context,
+	_ bool,
+	_ quic.ConnectionID,
+) (tracer qlogwriter.Trace) {
+	newTracer := &quicTracer{recorder: &headerRecorder{}}
+	t.tracers = append(t.tracers, newTracer)
+
+	return newTracer
+}
+
+// connectionsInfo returns info for all traced connections.
+func (t *testTracer) connectionsInfo() (res []*connInfo) {
+	res = make([]*connInfo, 0, len(t.tracers))
+	for _, tracer := range t.tracers {
+		hdrs := tracer.recorder.headersWithLock()
+
+		res = append(res, &connInfo{
+			headers: hdrs,
+		})
+	}
+
+	return res
+}
+
+// connInfo contains all trace event headers recorded for single connection.
+type connInfo struct {
+	headers []qlog.PacketHeader
+}
+
+// is0RTT returns true if the connection used 0-RTT packets.
+func (c *connInfo) is0RTT() (ok bool) {
+	for _, hdr := range c.headers {
+		if hdr.PacketType == qlog.PacketType0RTT {
+			return true
+		}
+	}
+
+	return false
+}
+
+// quicTracer is an implementation of [qlogwriter.Trace] for testing.
+type quicTracer struct {
+	// recorder is used for recording trace events.  It must not be nil.
+	recorder *headerRecorder
+}
+
+// type check
+var _ qlogwriter.Trace = (*quicTracer)(nil)
+
+// AddProducer implements the [qlogwriter.Trace] interface for *quicTracer.
+func (q *quicTracer) AddProducer() (recorder qlogwriter.Recorder) {
+	return q.recorder
+}
+
+// SupportsSchemas implements the [qlogwriter.Trace] interface for *quicTracer.
+func (q *quicTracer) SupportsSchemas(string) (ok bool) {
+	return false
+}
+
+// Recorder is an implementation of [qlogwriter.Recorder] that records
+// [qlog.PacketSent] events headers.
+type headerRecorder struct {
+	headers []qlog.PacketHeader
+	mx      sync.Mutex
+}
+
+// type check
+var _ qlogwriter.Recorder = (*headerRecorder)(nil)
+
+// RecordEvent implements the [qlogwriter.Recorder] interface for
+// *headerRecorder.
+func (r *headerRecorder) RecordEvent(ev qlogwriter.Event) {
+	event, ok := ev.(qlog.PacketSent)
+	if !ok {
+		return
+	}
+
+	r.mx.Lock()
+	defer r.mx.Unlock()
+
+	r.headers = append(r.headers, event.Header)
+}
+
+// headersWithLock returns copy of recorded headers.  It is safe for concurrent
+// use.
+func (r *headerRecorder) headersWithLock() (res []qlog.PacketHeader) {
+	r.mx.Lock()
+	defer r.mx.Unlock()
+
+	res = r.headers
+
+	return res
+}
+
+// Close implements the [qlogwriter.Recorder] interface for
+// *headerRecorder.
+func (*headerRecorder) Close() (err error) {
+	return nil
+}