From 097926c51fa634676308a0b8efdd1ee85143bfb8 Mon Sep 17 00:00:00 2001
From: DecFox <33030671+DecFox@users.noreply.github.com>
Date: Thu, 18 Aug 2022 00:28:06 +0530
Subject: [PATCH] refactor: allow automatically wrap net/quic conn (#867)

See https://github.com/ooni/probe/issues/2219
---
 .../simplequicping/simplequicping.go          |  3 +-
 internal/engine/experiment/tlsping/tlsping.go |  1 -
 internal/measurexlite/conn.go                 |  8 +--
 internal/measurexlite/conn_test.go            | 20 +++----
 internal/measurexlite/quic.go                 | 24 --------
 internal/measurexlite/quic_test.go            | 59 -------------------
 internal/model/mocks/trace.go                 | 13 ++++
 internal/model/mocks/trace_test.go            | 27 +++++++++
 internal/model/netx.go                        | 15 +++++
 internal/netxlite/dialer.go                   |  2 +-
 internal/netxlite/quic.go                     |  1 +
 internal/netxlite/trace.go                    | 11 ++++
 12 files changed, 83 insertions(+), 101 deletions(-)

diff --git a/internal/engine/experiment/simplequicping/simplequicping.go b/internal/engine/experiment/simplequicping/simplequicping.go
index 89cbed7..d8a1ec4 100644
--- a/internal/engine/experiment/simplequicping/simplequicping.go
+++ b/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 := trace.WrapQUICListener(quicListener)
+	listener := netxlite.NewQUICListener()
 	dialer := trace.NewQUICDialerWithoutResolver(listener, logger)
 	tlsConfig := &tls.Config{
 		NextProtos: alpn,
diff --git a/internal/engine/experiment/tlsping/tlsping.go b/internal/engine/experiment/tlsping/tlsping.go
index a89bdc7..eda8023 100644
--- a/internal/engine/experiment/tlsping/tlsping.go
+++ b/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,
diff --git a/internal/measurexlite/conn.go b/internal/measurexlite/conn.go
index c6e4714..0236ee3 100644
--- a/internal/measurexlite/conn.go
+++ b/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.
-func (tx *Trace) WrapNetConn(conn net.Conn) net.Conn {
+// MaybeWrapNetConn implements model.Trace.MaybeWrapNetConn.
+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.
-func (tx *Trace) WrapUDPLikeConn(conn model.UDPLikeConn) model.UDPLikeConn {
+// MaybeWrapUDPLikeConn implements model.Trace.MaybeWrapUDPLikeConn.
+func (tx *Trace) MaybeWrapUDPLikeConn(conn model.UDPLikeConn) model.UDPLikeConn {
 	return &udpLikeConnTrace{
 		UDPLikeConn: conn,
 		tx:          tx,
diff --git a/internal/measurexlite/conn_test.go b/internal/measurexlite/conn_test.go
index bf2dbea..ff59e71 100644
--- a/internal/measurexlite/conn_test.go
+++ b/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{
diff --git a/internal/measurexlite/quic.go b/internal/measurexlite/quic.go
index b3c23f4..8cb8a91 100644
--- a/internal/measurexlite/quic.go
+++ b/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 {
diff --git a/internal/measurexlite/quic_test.go b/internal/measurexlite/quic_test.go
index 29e148e..9aa3deb 100644
--- a/internal/measurexlite/quic_test.go
+++ b/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{}
diff --git a/internal/model/mocks/trace.go b/internal/model/mocks/trace.go
index 11136e9..aca9962 100644
--- a/internal/model/mocks/trace.go
+++ b/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)
diff --git a/internal/model/mocks/trace_test.go b/internal/model/mocks/trace_test.go
index b56a461..2972bf7 100644
--- a/internal/model/mocks/trace_test.go
+++ b/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{
diff --git a/internal/model/netx.go b/internal/model/netx.go
index 8dbd659..bfb40ed 100644
--- a/internal/model/netx.go
+++ b/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.
 	//
diff --git a/internal/netxlite/dialer.go b/internal/netxlite/dialer.go
index c766fe3..b9d0dcd 100644
--- a/internal/netxlite/dialer.go
+++ b/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)
 	}
diff --git a/internal/netxlite/quic.go b/internal/netxlite/quic.go
index 8810cbd..9e7f028 100644
--- a/internal/netxlite/quic.go
+++ b/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(
diff --git a/internal/netxlite/trace.go b/internal/netxlite/trace.go
index f364ae6..8409c01 100644
--- a/internal/netxlite/trace.go
+++ b/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) {