fix(webconnectivitylte): count bytes sent and received (#1488)

Closes ooni/probe#2655

internal/bytecounter/context.go

@@ -11,7 +11,7 @@ import (
 
 type byteCounterSessionKey struct{}
 
-// ContextSessionByteCounter retrieves the session byte counter from the context
+// ContextSessionByteCounter retrieves the possibly-nil session byte counter from the context.
 func ContextSessionByteCounter(ctx context.Context) *Counter {
 	counter, _ := ctx.Value(byteCounterSessionKey{}).(*Counter)
 	return counter
@@ -24,7 +24,7 @@ func WithSessionByteCounter(ctx context.Context, counter *Counter) context.Conte
 
 type byteCounterExperimentKey struct{}
 
-// ContextExperimentByteCounter retrieves the experiment byte counter from the context
+// ContextExperimentByteCounter retrieves the possibly-nil experiment byte counter from the context.
 func ContextExperimentByteCounter(ctx context.Context) *Counter {
 	counter, _ := ctx.Value(byteCounterExperimentKey{}).(*Counter)
 	return counter

internal/bytecounter/dialer.go

@@ -14,16 +14,15 @@ import (
 // MaybeWrapWithContextAwareDialer wraps the given dialer with a ContextAwareDialer
 // if the enabled argument is true and otherwise just returns the given dialer.
 //
-// # Bug
+// # Caveat
 //
 // This implementation cannot properly account for the bytes that are sent by
 // persistent connections, because they stick to the counters set when the
 // connection was established. This typically means we miss the bytes sent and
 // received when submitting a measurement. Such bytes are specifically not
 // seen by the experiment specific byte counter.
 //
-// For this reason, this implementation may be heavily changed/removed
-// in the future (<- this message is now ~two years old, though).
+// As such, this implementation should only be used when measuring.
 func MaybeWrapWithContextAwareDialer(enabled bool, dialer model.Dialer) model.Dialer {
 	if !enabled {
 		return dialer

internal/bytecounter/resolver.go

@@ -8,11 +8,59 @@ import (
 	"github.com/ooni/probe-cli/v3/internal/netxlite"
 )
 
+// WrapWithContextAwareSystemResolver wraps the given resolver with a resolver that
+// is aware of context-byte counting. See MaybeWrapSystemResolver for a list of caveats.
+func WrapWithContextAwareSystemResolver(reso model.Resolver) model.Resolver {
+	return &ContextAwareSystemResolver{reso}
+}
+
+// ContextAwareSystemResolver is a [model.Resolver] that knows how to count bytes
+// sent and received. We typically use this for the system resolver only because for
+// other resolvers we are better off just wrapping their connections.
+type ContextAwareSystemResolver struct {
+	R model.Resolver
+}
+
+// Address implements model.Resolver.
+func (r *ContextAwareSystemResolver) Address() string {
+	return r.R.Address()
+}
+
+// CloseIdleConnections implements model.Resolver.
+func (r *ContextAwareSystemResolver) CloseIdleConnections() {
+	r.R.CloseIdleConnections()
+}
+
+func (r *ContextAwareSystemResolver) wrap(ctx context.Context) model.Resolver {
+	return MaybeWrapSystemResolver(MaybeWrapSystemResolver(
+		r.R, ContextSessionByteCounter(ctx)), ContextExperimentByteCounter(ctx))
+}
+
+// LookupHTTPS implements model.Resolver.
+func (r *ContextAwareSystemResolver) LookupHTTPS(ctx context.Context, domain string) (*model.HTTPSSvc, error) {
+	return r.wrap(ctx).LookupHTTPS(ctx, domain)
+}
+
+// LookupHost implements model.Resolver.
+func (r *ContextAwareSystemResolver) LookupHost(ctx context.Context, hostname string) (addrs []string, err error) {
+	return r.wrap(ctx).LookupHost(ctx, hostname)
+}
+
+// LookupNS implements model.Resolver.
+func (r *ContextAwareSystemResolver) LookupNS(ctx context.Context, domain string) ([]*net.NS, error) {
+	return r.wrap(ctx).LookupNS(ctx, domain)
+}
+
+// Network implements model.Resolver.
+func (r *ContextAwareSystemResolver) Network() string {
+	return r.R.Network()
+}
+
 // MaybeWrapSystemResolver takes in input a Resolver and either wraps it
 // to perform byte counting, if this counter is not nil, or just returns to the
 // caller the original resolver, when the counter is nil.
 //
-// # Bug
+// # Caveat
 //
 // The returned resolver will only approximately estimate the bytes
 // sent and received by this resolver if this resolver is the system

internal/bytecounter/resolver_test.go

@@ -283,3 +283,274 @@ func TestMaybeWrapSystemResolver(t *testing.T) {
 		}
 	})
 }
+
+func TestWrapWithContextAwareSystemResolver(t *testing.T) {
+	t.Run("Address works as intended", func(t *testing.T) {
+		underlying := &mocks.Resolver{
+			MockAddress: func() string {
+				return "8.8.8.8:53"
+			},
+		}
+		reso := WrapWithContextAwareSystemResolver(underlying)
+		if reso.Address() != "8.8.8.8:53" {
+			t.Fatal("unexpected result")
+		}
+	})
+
+	t.Run("CloseIdleConnections works as intended", func(t *testing.T) {
+		var called bool
+		underlying := &mocks.Resolver{
+			MockCloseIdleConnections: func() {
+				called = true
+			},
+		}
+		reso := WrapWithContextAwareSystemResolver(underlying)
+		reso.CloseIdleConnections()
+		if !called {
+			t.Fatal("not called")
+		}
+	})
+
+	t.Run("LookupHTTPS works as intended", func(t *testing.T) {
+		t.Run("on success", func(t *testing.T) {
+			expected := &model.HTTPSSvc{}
+			underlying := &mocks.Resolver{
+				MockLookupHTTPS: func(ctx context.Context, domain string) (*model.HTTPSSvc, error) {
+					return expected, nil
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupHTTPS(ctx, "dns.google")
+			if err != nil {
+				t.Fatal("unexpected error", err)
+			}
+			if got != expected {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 10 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 256 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+
+		t.Run("on non-DNS failure", func(t *testing.T) {
+			expected := errors.New("mocked error")
+			underlying := &mocks.Resolver{
+				MockLookupHTTPS: func(ctx context.Context, domain string) (*model.HTTPSSvc, error) {
+					return nil, expected
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupHTTPS(ctx, "dns.google")
+			if !errors.Is(err, expected) {
+				t.Fatal("unexpected error", err)
+			}
+			if got != nil {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 10 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 0 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+
+		t.Run("on DNS failure", func(t *testing.T) {
+			expected := errors.New(netxlite.FailureDNSNXDOMAINError)
+			underlying := &mocks.Resolver{
+				MockLookupHTTPS: func(ctx context.Context, domain string) (*model.HTTPSSvc, error) {
+					return nil, expected
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupHTTPS(ctx, "dns.google")
+			if !errors.Is(err, expected) {
+				t.Fatal("unexpected error", err)
+			}
+			if got != nil {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 10 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 128 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+	})
+
+	t.Run("LookupNS works as intended", func(t *testing.T) {
+		t.Run("on success", func(t *testing.T) {
+			underlying := &mocks.Resolver{
+				MockLookupNS: func(ctx context.Context, domain string) ([]*net.NS, error) {
+					out := make([]*net.NS, 3)
+					return out, nil
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupNS(ctx, "dns.google")
+			if err != nil {
+				t.Fatal("unexpected error", err)
+			}
+			if len(got) != 3 {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 10 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 256 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+
+		t.Run("on non-DNS failure", func(t *testing.T) {
+			expected := errors.New("mocked error")
+			underlying := &mocks.Resolver{
+				MockLookupNS: func(ctx context.Context, domain string) ([]*net.NS, error) {
+					return nil, expected
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupNS(ctx, "dns.google")
+			if !errors.Is(err, expected) {
+				t.Fatal("unexpected error", err)
+			}
+			if len(got) != 0 {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 10 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 0 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+
+		t.Run("on DNS failure", func(t *testing.T) {
+			expected := errors.New(netxlite.FailureDNSNXDOMAINError)
+			underlying := &mocks.Resolver{
+				MockLookupNS: func(ctx context.Context, domain string) ([]*net.NS, error) {
+					return nil, expected
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupNS(ctx, "dns.google")
+			if !errors.Is(err, expected) {
+				t.Fatal("unexpected error", err)
+			}
+			if len(got) != 0 {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 10 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 128 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+	})
+
+	t.Run("LookupHost works as intended", func(t *testing.T) {
+		t.Run("on success", func(t *testing.T) {
+			underlying := &mocks.Resolver{
+				MockLookupHost: func(ctx context.Context, domain string) ([]string, error) {
+					out := make([]string, 3)
+					return out, nil
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupHost(ctx, "dns.google")
+			if err != nil {
+				t.Fatal("unexpected error", err)
+			}
+			if len(got) != 3 {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 20 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 256 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+
+		t.Run("on non-DNS failure", func(t *testing.T) {
+			expected := errors.New("mocked error")
+			underlying := &mocks.Resolver{
+				MockLookupHost: func(ctx context.Context, domain string) ([]string, error) {
+					return nil, expected
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupHost(ctx, "dns.google")
+			if !errors.Is(err, expected) {
+				t.Fatal("unexpected error", err)
+			}
+			if len(got) != 0 {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 20 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 0 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+
+		t.Run("on DNS failure", func(t *testing.T) {
+			expected := errors.New(netxlite.FailureDNSNXDOMAINError)
+			underlying := &mocks.Resolver{
+				MockLookupHost: func(ctx context.Context, domain string) ([]string, error) {
+					return nil, expected
+				},
+			}
+			counter := New()
+			reso := WrapWithContextAwareSystemResolver(underlying)
+			ctx := WithSessionByteCounter(context.Background(), counter)
+			got, err := reso.LookupHost(ctx, "dns.google")
+			if !errors.Is(err, expected) {
+				t.Fatal("unexpected error", err)
+			}
+			if len(got) != 0 {
+				t.Fatal("invalid result")
+			}
+			if nsent := counter.BytesSent(); nsent != 20 {
+				t.Fatal("unexpected nsent", nsent)
+			}
+			if nrecv := counter.BytesReceived(); nrecv != 128 {
+				t.Fatal("unexpected nrecv")
+			}
+		})
+	})
+
+	t.Run("Network works as intended", func(t *testing.T) {
+		underlying := &mocks.Resolver{
+			MockNetwork: func() string {
+				return "udp"
+			},
+		}
+		reso := WrapWithContextAwareSystemResolver(underlying)
+		if reso.Network() != "udp" {
+			t.Fatal("unexpected result")
+		}
+	})
+}