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

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

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.
 	//