refactor: cleaner way of passing a UDPConn around (#421)

* refactor: cleaner way of passing a UDPConn around

Also part of https://github.com/ooni/probe/issues/1505

* Update internal/engine/netx/quicdialer/connectionstate.go

internal/engine/netx/quicdialer/connectionstate.go

@@ -6,7 +6,7 @@ import (
 	"github.com/lucas-clemente/quic-go"
 )
 
-// ConnectionState returns the ConnectionState of a QUIC Session.
+// connectionState returns the ConnectionState of a QUIC Session.
-func ConnectionState(sess quic.EarlySession) tls.ConnectionState {
+func connectionState(sess quic.EarlySession) tls.ConnectionState {
 	return sess.ConnectionState().TLS.ConnectionState
 }

internal/engine/netx/quicdialer/quicdialer.go

@@ -14,12 +14,6 @@ type ContextDialer interface {
 		tlsCfg *tls.Config, cfg *quic.Config) (quic.EarlySession, error)
 }
 
-// Dialer dials QUIC connections.
-type Dialer interface {
-	// Note: assumes that tlsCfg and cfg are not nil.
-	Dial(network, addr string, tlsCfg *tls.Config, cfg *quic.Config) (quic.EarlySession, error)
-}
-
 // Resolver is the interface we expect from a resolver.
 type Resolver interface {
 	LookupHost(ctx context.Context, hostname string) (addrs []string, err error)

internal/engine/netx/quicdialer/saver.go

@@ -45,7 +45,7 @@ func (h HandshakeSaver) DialContext(ctx context.Context, network string,
 		})
 		return nil, err
 	}
-	state := ConnectionState(sess)
+	state := connectionState(sess)
 	h.Saver.Write(trace.Event{
 		Duration:           stop.Sub(start),
 		Name:               "quic_handshake_done",

internal/engine/netx/quicdialer/system.go

@@ -1,18 +1,18 @@
 package quicdialer
 
 import (
-	"errors"
 	"net"
 	"time"
 
 	"github.com/ooni/probe-cli/v3/internal/engine/netx/trace"
 	"github.com/ooni/probe-cli/v3/internal/errorsx"
+	"github.com/ooni/probe-cli/v3/internal/quicx"
 )
 
 // QUICListener listens for QUIC connections.
 type QUICListener interface {
-	// Listen creates a new listening PacketConn.
+	// Listen creates a new listening UDPConn.
-	Listen(addr *net.UDPAddr) (net.PacketConn, error)
+	Listen(addr *net.UDPAddr) (quicx.UDPConn, error)
 }
 
 // QUICListenerSaver is a QUICListener that also implements saving events.
@@ -25,21 +25,16 @@ type QUICListenerSaver struct {
 }
 
 // Listen implements QUICListener.Listen.
-func (qls *QUICListenerSaver) Listen(addr *net.UDPAddr) (net.PacketConn, error) {
+func (qls *QUICListenerSaver) Listen(addr *net.UDPAddr) (quicx.UDPConn, error) {
 	pconn, err := qls.QUICListener.Listen(addr)
 	if err != nil {
 		return nil, err
 	}
-	// TODO(bassosimone): refactor to remove this restriction.
+	return saverUDPConn{UDPConn: pconn, saver: qls.Saver}, nil
-	udpConn, ok := pconn.(*net.UDPConn)
-	if !ok {
-		return nil, errors.New("quicdialer: cannot convert to udpConn")
-	}
-	return saverUDPConn{UDPConn: udpConn, saver: qls.Saver}, nil
 }
 
 type saverUDPConn struct {
-	*net.UDPConn
+	quicx.UDPConn
 	saver *trace.Saver
 }
 

internal/netxlite/quic.go

@@ -8,6 +8,7 @@ import (
 	"strconv"
 
 	"github.com/lucas-clemente/quic-go"
+	"github.com/ooni/probe-cli/v3/internal/quicx"
 )
 
 // QUICContextDialer is a dialer for QUIC using Context.
@@ -21,8 +22,8 @@ type QUICContextDialer interface {
 
 // QUICListener listens for QUIC connections.
 type QUICListener interface {
-	// Listen creates a new listening PacketConn.
+	// Listen creates a new listening UDPConn.
-	Listen(addr *net.UDPAddr) (net.PacketConn, error)
+	Listen(addr *net.UDPAddr) (quicx.UDPConn, error)
 }
 
 // QUICListenerStdlib is a QUICListener using the standard library.
@@ -31,7 +32,7 @@ type QUICListenerStdlib struct{}
 var _ QUICListener = &QUICListenerStdlib{}
 
 // Listen implements QUICListener.Listen.
-func (qls *QUICListenerStdlib) Listen(addr *net.UDPAddr) (net.PacketConn, error) {
+func (qls *QUICListenerStdlib) Listen(addr *net.UDPAddr) (quicx.UDPConn, error) {
 	return net.ListenUDP("udp", addr)
 }
 

internal/netxlite/quic_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/lucas-clemente/quic-go"
 	"github.com/ooni/probe-cli/v3/internal/netxmocks"
+	"github.com/ooni/probe-cli/v3/internal/quicx"
 )
 
 func TestQUICDialerQUICGoCannotSplitHostPort(t *testing.T) {
@@ -75,7 +76,7 @@ func TestQUICDialerQUICGoCannotListen(t *testing.T) {
 	}
 	systemdialer := QUICDialerQUICGo{
 		QUICListener: &netxmocks.QUICListener{
-			MockListen: func(addr *net.UDPAddr) (net.PacketConn, error) {
+			MockListen: func(addr *net.UDPAddr) (quicx.UDPConn, error) {
 				return nil, expected
 			},
 		},

internal/netxmocks/quic.go

@@ -6,15 +6,16 @@ import (
 	"net"
 
 	"github.com/lucas-clemente/quic-go"
+	"github.com/ooni/probe-cli/v3/internal/quicx"
 )
 
 // QUICListener is a mockable netxlite.QUICListener.
 type QUICListener struct {
-	MockListen func(addr *net.UDPAddr) (net.PacketConn, error)
+	MockListen func(addr *net.UDPAddr) (quicx.UDPConn, error)
 }
 
 // Listen calls MockListen.
-func (ql *QUICListener) Listen(addr *net.UDPAddr) (net.PacketConn, error) {
+func (ql *QUICListener) Listen(addr *net.UDPAddr) (quicx.UDPConn, error) {
 	return ql.MockListen(addr)
 }
 

internal/netxmocks/quic_test.go

@@ -9,12 +9,13 @@ import (
 	"testing"
 
 	"github.com/lucas-clemente/quic-go"
+	"github.com/ooni/probe-cli/v3/internal/quicx"
 )
 
 func TestQUICListenerListen(t *testing.T) {
 	expected := errors.New("mocked error")
 	ql := &QUICListener{
-		MockListen: func(addr *net.UDPAddr) (net.PacketConn, error) {
+		MockListen: func(addr *net.UDPAddr) (quicx.UDPConn, error) {
 			return nil, expected
 		},
 	}

internal/quicx/quicx.go

@@ -0,0 +1,13 @@
+// Package quicx contains definitions useful to implement QUIC.
+package quicx
+
+import "net"
+
+// UDPConn is an UDP connection used by quic.
+type UDPConn interface {
+	// PacketConn is the underlying base interface.
+	net.PacketConn
+
+	// ReadMsgUDP behaves like net.UDPConn.ReadMsgUDP.
+	ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
+}