refactor: allow automatically wrap net/quic conn (#867)
See https://github.com/ooni/probe/issues/2219
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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{}
|
||||
|
||||
Reference in New Issue
Block a user