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

See https://github.com/ooni/probe/issues/2219
This commit is contained in:
DecFox
2022-08-18 00:28:06 +05:30
committed by GitHub
parent e1d014e826
commit 097926c51f
12 changed files with 83 additions and 101 deletions
+4 -4
View File
@@ -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,
+10 -10
View File
@@ -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{
-24
View File
@@ -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 {
-59
View File
@@ -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{}