refactor: allow automatically wrap net/quic conn (#867)

See https://github.com/ooni/probe/issues/2219

internal/engine/experiment/simplequicping/simplequicping.go

@@ -172,8 +172,7 @@ func (m *Measurer) quicHandshake(ctx context.Context, index int64,
 	alpn := strings.Split(m.config.alpn(), " ")
 	trace := measurexlite.NewTrace(index, zeroTime)
 	ol := measurexlite.NewOperationLogger(logger, "SimpleQUICPing #%d %s %s %v", index, address, sni, alpn)
-	quicListener := netxlite.NewQUICListener()
+	listener := netxlite.NewQUICListener()
-	listener := trace.WrapQUICListener(quicListener)
 	dialer := trace.NewQUICDialerWithoutResolver(listener, logger)
 	tlsConfig := &tls.Config{
 		NextProtos: alpn,

internal/engine/experiment/tlsping/tlsping.go

@@ -182,7 +182,6 @@ func (m *Measurer) tlsConnectAndHandshake(ctx context.Context, index int64,
 		return sp
 	}
 	defer conn.Close()
-	conn = trace.WrapNetConn(conn)
 	thx := trace.NewTLSHandshakerStdlib(logger)
 	config := &tls.Config{
 		NextProtos: alpn,

internal/measurexlite/conn.go

@@ -21,8 +21,8 @@ func MaybeClose(conn net.Conn) (err error) {
 	return
 }
 
-// WrapNetConn returns a wrapped conn that saves network events into this trace.
+// MaybeWrapNetConn implements model.Trace.MaybeWrapNetConn.
-func (tx *Trace) WrapNetConn(conn net.Conn) net.Conn {
+func (tx *Trace) MaybeWrapNetConn(conn net.Conn) net.Conn {
 	return &connTrace{
 		Conn: conn,
 		tx:   tx,
@@ -77,8 +77,8 @@ func MaybeCloseUDPLikeConn(conn model.UDPLikeConn) (err error) {
 	return
 }
 
-// WrapUDPLikeConn returns a wrapped conn that saves network events into this trace.
+// MaybeWrapUDPLikeConn implements model.Trace.MaybeWrapUDPLikeConn.
-func (tx *Trace) WrapUDPLikeConn(conn model.UDPLikeConn) model.UDPLikeConn {
+func (tx *Trace) MaybeWrapUDPLikeConn(conn model.UDPLikeConn) model.UDPLikeConn {
 	return &udpLikeConnTrace{
 		UDPLikeConn: conn,
 		tx:          tx,

internal/measurexlite/conn_test.go

@@ -40,7 +40,7 @@ func TestWrapNetConn(t *testing.T) {
 		underlying := &mocks.Conn{}
 		zeroTime := time.Now()
 		trace := NewTrace(0, zeroTime)
-		conn := trace.WrapNetConn(underlying)
+		conn := trace.MaybeWrapNetConn(underlying)
 		ct := conn.(*connTrace)
 		if ct.Conn != underlying {
 			t.Fatal("invalid underlying")
@@ -70,7 +70,7 @@ func TestWrapNetConn(t *testing.T) {
 		td := testingx.NewTimeDeterministic(zeroTime)
 		trace := NewTrace(0, zeroTime)
 		trace.TimeNowFn = td.Now // deterministic time counting
-		conn := trace.WrapNetConn(underlying)
+		conn := trace.MaybeWrapNetConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		count, err := conn.Read(buffer)
@@ -118,7 +118,7 @@ func TestWrapNetConn(t *testing.T) {
 		zeroTime := time.Now()
 		trace := NewTrace(0, zeroTime)
 		trace.networkEvent = make(chan *model.ArchivalNetworkEvent) // no buffer
-		conn := trace.WrapNetConn(underlying)
+		conn := trace.MaybeWrapNetConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		count, err := conn.Read(buffer)
@@ -154,7 +154,7 @@ func TestWrapNetConn(t *testing.T) {
 		td := testingx.NewTimeDeterministic(zeroTime)
 		trace := NewTrace(0, zeroTime)
 		trace.TimeNowFn = td.Now // deterministic time tracking
-		conn := trace.WrapNetConn(underlying)
+		conn := trace.MaybeWrapNetConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		count, err := conn.Write(buffer)
@@ -202,7 +202,7 @@ func TestWrapNetConn(t *testing.T) {
 		zeroTime := time.Now()
 		trace := NewTrace(0, zeroTime)
 		trace.networkEvent = make(chan *model.ArchivalNetworkEvent) // no buffer
-		conn := trace.WrapNetConn(underlying)
+		conn := trace.MaybeWrapNetConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		count, err := conn.Write(buffer)
@@ -224,7 +224,7 @@ func TestWrapUDPLikeConn(t *testing.T) {
 		underlying := &mocks.UDPLikeConn{}
 		zeroTime := time.Now()
 		trace := NewTrace(0, zeroTime)
-		conn := trace.WrapUDPLikeConn(underlying)
+		conn := trace.MaybeWrapUDPLikeConn(underlying)
 		ct := conn.(*udpLikeConnTrace)
 		if ct.UDPLikeConn != underlying {
 			t.Fatal("invalid underlying")
@@ -248,7 +248,7 @@ func TestWrapUDPLikeConn(t *testing.T) {
 		td := testingx.NewTimeDeterministic(zeroTime)
 		trace := NewTrace(0, zeroTime)
 		trace.TimeNowFn = td.Now // deterministic time counting
-		conn := trace.WrapUDPLikeConn(underlying)
+		conn := trace.MaybeWrapUDPLikeConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		count, addr, err := conn.ReadFrom(buffer)
@@ -293,7 +293,7 @@ func TestWrapUDPLikeConn(t *testing.T) {
 		zeroTime := time.Now()
 		trace := NewTrace(0, zeroTime)
 		trace.networkEvent = make(chan *model.ArchivalNetworkEvent) // no buffer
-		conn := trace.WrapUDPLikeConn(underlying)
+		conn := trace.MaybeWrapUDPLikeConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		count, addr, err := conn.ReadFrom(buffer)
@@ -322,7 +322,7 @@ func TestWrapUDPLikeConn(t *testing.T) {
 		td := testingx.NewTimeDeterministic(zeroTime)
 		trace := NewTrace(0, zeroTime)
 		trace.TimeNowFn = td.Now // deterministic time tracking
-		conn := trace.WrapUDPLikeConn(underlying)
+		conn := trace.MaybeWrapUDPLikeConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		addr := &mocks.Addr{
@@ -365,7 +365,7 @@ func TestWrapUDPLikeConn(t *testing.T) {
 		zeroTime := time.Now()
 		trace := NewTrace(0, zeroTime)
 		trace.networkEvent = make(chan *model.ArchivalNetworkEvent) // no buffer
-		conn := trace.WrapUDPLikeConn(underlying)
+		conn := trace.MaybeWrapUDPLikeConn(underlying)
 		const bufsiz = 128
 		buffer := make([]byte, bufsiz)
 		addr := &mocks.Addr{

internal/measurexlite/quic.go

@@ -7,7 +7,6 @@ package measurexlite
 import (
 	"context"
 	"crypto/tls"
-	"net"
 	"time"
 
 	"github.com/lucas-clemente/quic-go"
@@ -15,29 +14,6 @@ import (
 	"github.com/ooni/probe-cli/v3/internal/netxlite"
 )
 
-// WrapQUICListener returns a wrapped model.QUICListener that uses this trace.
-func (tx *Trace) WrapQUICListener(listener model.QUICListener) model.QUICListener {
-	return &quicListenerTrace{
-		QUICListener: listener,
-		tx:           tx,
-	}
-}
-
-// quicListenerTrace is a trace-aware QUIC listener.
-type quicListenerTrace struct {
-	model.QUICListener
-	tx *Trace
-}
-
-// Listen implements model.QUICListener.Listen
-func (ql *quicListenerTrace) Listen(addr *net.UDPAddr) (model.UDPLikeConn, error) {
-	pconn, err := ql.QUICListener.Listen(addr)
-	if err != nil {
-		return nil, err
-	}
-	return ql.tx.WrapUDPLikeConn(pconn), nil
-}
-
 // NewQUICDialerWithoutResolver is equivalent to netxlite.NewQUICDialerWithoutResolver
 // except that it returns a model.QUICDialer that uses this trace.
 func (tx *Trace) NewQUICDialerWithoutResolver(listener model.QUICListener, dl model.DebugLogger) model.QUICDialer {

internal/measurexlite/quic_test.go

@@ -17,65 +17,6 @@ import (
 	"github.com/ooni/probe-cli/v3/internal/testingx"
 )
 
-func TestNewQUICListener(t *testing.T) {
-	t.Run("NewQUICListenerTrace creates a wrapped listener", func(t *testing.T) {
-		underlying := &mocks.QUICListener{}
-		zeroTime := time.Now()
-		trace := NewTrace(0, zeroTime)
-		listenert := trace.WrapQUICListener(underlying).(*quicListenerTrace)
-		if listenert.QUICListener != underlying {
-			t.Fatal("invalid quic dialer")
-		}
-		if listenert.tx != trace {
-			t.Fatal("invalid trace")
-		}
-	})
-
-	t.Run("Listen works as intended", func(t *testing.T) {
-		t.Run("with error", func(t *testing.T) {
-			zeroTime := time.Now()
-			trace := NewTrace(0, zeroTime)
-			mockedErr := errors.New("mocked")
-			mockListener := &mocks.QUICListener{
-				MockListen: func(addr *net.UDPAddr) (model.UDPLikeConn, error) {
-					return nil, mockedErr
-				},
-			}
-			listener := trace.WrapQUICListener(mockListener)
-			pconn, err := listener.Listen(&net.UDPAddr{})
-			if !errors.Is(err, mockedErr) {
-				t.Fatal("unexpected err", err)
-			}
-			if pconn != nil {
-				t.Fatal("expected nil conn")
-			}
-		})
-
-		t.Run("without error", func(t *testing.T) {
-			zeroTime := time.Now()
-			trace := NewTrace(0, zeroTime)
-			mockConn := &mocks.UDPLikeConn{}
-			mockListener := &mocks.QUICListener{
-				MockListen: func(addr *net.UDPAddr) (model.UDPLikeConn, error) {
-					return mockConn, nil
-				},
-			}
-			listener := trace.WrapQUICListener(mockListener)
-			pconn, err := listener.Listen(&net.UDPAddr{})
-			if err != nil {
-				t.Fatal("unexpected err", err)
-			}
-			conn := pconn.(*udpLikeConnTrace)
-			if conn.UDPLikeConn != mockConn {
-				t.Fatal("invalid conn")
-			}
-			if conn.tx != trace {
-				t.Fatal("invalid trace")
-			}
-		})
-	})
-}
-
 func TestNewQUICDialerWithoutResolver(t *testing.T) {
 	t.Run("NewQUICDialerWithoutResolver creates a wrapped dialer", func(t *testing.T) {
 		underlying := &mocks.QUICDialer{}

internal/model/mocks/trace.go

@@ -6,6 +6,7 @@ package mocks
 
 import (
 	"crypto/tls"
+	"net"
 	"time"
 
 	"github.com/lucas-clemente/quic-go"
@@ -16,6 +17,10 @@ import (
 type Trace struct {
 	MockTimeNow func() time.Time
 
+	MockMaybeWrapNetConn func(conn net.Conn) net.Conn
+
+	MockMaybeWrapUDPLikeConn func(conn model.UDPLikeConn) model.UDPLikeConn
+
 	MockOnDNSRoundTripForLookupHost func(started time.Time, reso model.Resolver, query model.DNSQuery,
 		response model.DNSResponse, addrs []string, err error, finished time.Time)
 
@@ -39,6 +44,14 @@ func (t *Trace) TimeNow() time.Time {
 	return t.MockTimeNow()
 }
 
+func (t *Trace) MaybeWrapNetConn(conn net.Conn) net.Conn {
+	return t.MockMaybeWrapNetConn(conn)
+}
+
+func (t *Trace) MaybeWrapUDPLikeConn(conn model.UDPLikeConn) model.UDPLikeConn {
+	return t.MockMaybeWrapUDPLikeConn(conn)
+}
+
 func (t *Trace) OnDNSRoundTripForLookupHost(started time.Time, reso model.Resolver, query model.DNSQuery,
 	response model.DNSResponse, addrs []string, err error, finished time.Time) {
 	t.MockOnDNSRoundTripForLookupHost(started, reso, query, response, addrs, err, finished)

internal/model/mocks/trace_test.go

@@ -2,6 +2,7 @@ package mocks
 
 import (
 	"crypto/tls"
+	"net"
 	"testing"
 	"time"
 
@@ -22,6 +23,32 @@ func TestTrace(t *testing.T) {
 		}
 	})
 
+	t.Run("MaybeWrapNetConn", func(t *testing.T) {
+		expect := &Conn{}
+		tx := &Trace{
+			MockMaybeWrapNetConn: func(conn net.Conn) net.Conn {
+				return expect
+			},
+		}
+		got := tx.MaybeWrapNetConn(&Conn{})
+		if got != expect {
+			t.Fatal("not working as intended")
+		}
+	})
+
+	t.Run("MaybeWrapUDPLikeConn", func(t *testing.T) {
+		expect := &UDPLikeConn{}
+		tx := &Trace{
+			MockMaybeWrapUDPLikeConn: func(conn model.UDPLikeConn) model.UDPLikeConn {
+				return expect
+			},
+		}
+		got := tx.MaybeWrapUDPLikeConn(&UDPLikeConn{})
+		if got != expect {
+			t.Fatal("not working as intended")
+		}
+	})
+
 	t.Run("OnDNSRoundTripForLookupHost", func(t *testing.T) {
 		var called bool
 		tx := &Trace{

internal/model/netx.go

@@ -303,6 +303,21 @@ type Trace interface {
 	// can use functionality exported by the ./internal/testingx pkg.
 	TimeNow() time.Time
 
+	// MaybeWrapNetConn possibly wraps a net.Conn with the caller trace. If there's no
+	// desire to wrap the net.Conn, this function just returns the original net.Conn.
+	//
+	// Arguments:
+	//
+	// - conn is the non-nil underlying net.Conn to be wrapped
+	MaybeWrapNetConn(conn net.Conn) net.Conn
+
+	// MaybeWrapUDPLikeConn is like MaybeWrapNetConn but for UDPLikeConn.
+	//
+	// Arguments:
+	//
+	// - conn is the non-nil underlying UDPLikeConn to be wrapped
+	MaybeWrapUDPLikeConn(conn UDPLikeConn) UDPLikeConn
+
 	// OnDNSRoundTripForLookupHost is used with a DNSTransport and called
 	// when the RoundTrip terminates.
 	//

internal/netxlite/dialer.go

@@ -224,7 +224,7 @@ func (d *dialerResolverWithTracing) DialContext(ctx context.Context, network, ad
 		trace.OnConnectDone(started, network, onlyhost, target, err, finished)
 		if err == nil {
 			conn = &dialerErrWrapperConn{conn}
-			return conn, nil
+			return trace.MaybeWrapNetConn(conn), nil
 		}
 		errorslist = append(errorslist, err)
 	}

internal/netxlite/quic.go

@@ -135,6 +135,7 @@ func (d *quicDialerQUICGo) DialContext(ctx context.Context, network string,
 	}
 	tlsConfig = d.maybeApplyTLSDefaults(tlsConfig, udpAddr.Port)
 	trace := ContextTraceOrDefault(ctx)
+	pconn = trace.MaybeWrapUDPLikeConn(pconn)
 	started := trace.TimeNow()
 	trace.OnQUICHandshakeStart(started, address, quicConfig)
 	qconn, err := d.dialEarlyContext(

internal/netxlite/trace.go

@@ -7,6 +7,7 @@ package netxlite
 import (
 	"context"
 	"crypto/tls"
+	"net"
 	"time"
 
 	"github.com/lucas-clemente/quic-go"
@@ -50,6 +51,16 @@ func (*traceDefault) TimeNow() time.Time {
 	return time.Now()
 }
 
+// MaybeWrapNetConn implements model.Trace.MaybeWrapNetConn
+func (*traceDefault) MaybeWrapNetConn(conn net.Conn) net.Conn {
+	return conn
+}
+
+// MaybeWrapUDPLikeConn implements model.Trace.MaybeWrapUDPLikeConn
+func (*traceDefault) MaybeWrapUDPLikeConn(conn model.UDPLikeConn) model.UDPLikeConn {
+	return conn
+}
+
 // OnDNSRoundTripForLookupHost implements model.Trace.OnDNSRoundTripForLookupHost.
 func (*traceDefault) OnDNSRoundTripForLookupHost(started time.Time, reso model.Resolver, query model.DNSQuery,
 	response model.DNSResponse, addrs []string, err error, finished time.Time) {