feat(netxlite): support extracting the CNAME (#875)

* feat(netxlite): support extracting the CNAME

Closes https://github.com/ooni/probe/issues/2225

* fix(netxlite): attempt to increase coverage and improve tests

1. dnsovergetaddrinfo: specify the behavior of a DNSResponse returned
by this file to make it line with normal responses and write unit tests
to make sure we adhere to expectations;

2. dnsoverudp: make sure we wait to deferred responses also w/o a
custom context and post on a private channel and test that;

3. utls: recognize that we can actually write a test for NetConn and
what needs to change when we'll use go1.19 by default will just be
a cast that at that point can be removed.

internal/model/mocks/dnsresponse.go

@@ -18,6 +18,7 @@ type DNSResponse struct {
 	MockDecodeHTTPS      func() (*model.HTTPSSvc, error)
 	MockDecodeLookupHost func() ([]string, error)
 	MockDecodeNS         func() ([]*net.NS, error)
+	MockDecodeCNAME      func() (string, error)
 }
 
 var _ model.DNSResponse = &DNSResponse{}
@@ -45,3 +46,7 @@ func (r *DNSResponse) DecodeLookupHost() ([]string, error) {
 func (r *DNSResponse) DecodeNS() ([]*net.NS, error) {
 	return r.MockDecodeNS()
 }
+
+func (r *DNSResponse) DecodeCNAME() (string, error) {
+	return r.MockDecodeCNAME()
+}

internal/model/mocks/dnsresponse_test.go

@@ -102,4 +102,20 @@ func TestDNSResponse(t *testing.T) {
 			t.Fatal("unexpected out")
 		}
 	})
+
+	t.Run("DecodeCNAME", func(t *testing.T) {
+		expected := errors.New("mocked error")
+		r := &DNSResponse{
+			MockDecodeCNAME: func() (string, error) {
+				return "", expected
+			},
+		}
+		out, err := r.DecodeCNAME()
+		if !errors.Is(err, expected) {
+			t.Fatal("unexpected err", err)
+		}
+		if out != "" {
+			t.Fatal("unexpected out")
+		}
+	})
 }

internal/model/netx.go

@@ -36,6 +36,9 @@ type DNSResponse interface {
 
 	// DecodeNS returns all the NS entries in this response.
 	DecodeNS() ([]*net.NS, error)
+
+	// DecodeCNAME returns the first CNAME entry in this response.
+	DecodeCNAME() (string, error)
 }
 
 // The DNSDecoder decodes DNS responses.

internal/netxlite/dnsdecoder.go

@@ -166,5 +166,19 @@ func (r *dnsResponse) DecodeNS() ([]*net.NS, error) {
 	return out, nil
 }
 
+// DecodeCNAME implements model.DNSResponse.DecodeCNAME.
+func (r *dnsResponse) DecodeCNAME() (string, error) {
+	if err := r.rcodeToError(); err != nil {
+		return "", err
+	}
+	for _, answer := range r.msg.Answer {
+		switch avalue := answer.(type) {
+		case *dns.CNAME:
+			return avalue.Target, nil
+		}
+	}
+	return "", ErrOODNSNoAnswer
+}
+
 var _ model.DNSDecoder = &DNSDecoderMiekg{}
 var _ model.DNSResponse = &dnsResponse{}

internal/netxlite/dnsdecoder_test.go

@@ -44,7 +44,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 				queryID     = 17
 				unrelatedID = 14
 			)
-			reply := dnsGenLookupHostReplySuccess(dnsGenQuery(dns.TypeA, queryID))
+			reply := dnsGenLookupHostReplySuccess(dnsGenQuery(dns.TypeA, queryID), nil)
 			resp, err := d.DecodeResponse(reply, &mocks.DNSQuery{
 				MockID: func() uint16 {
 					return unrelatedID
@@ -62,7 +62,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 			d := &DNSDecoderMiekg{}
 			queryID := dns.Id()
 			rawQuery := dnsGenQuery(dns.TypeA, queryID)
-			rawResponse := dnsGenLookupHostReplySuccess(rawQuery)
+			rawResponse := dnsGenLookupHostReplySuccess(rawQuery, nil)
 			query := &mocks.DNSQuery{
 				MockID: func() uint16 {
 					return queryID
@@ -81,7 +81,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 			d := &DNSDecoderMiekg{}
 			queryID := dns.Id()
 			rawQuery := dnsGenQuery(dns.TypeA, queryID)
-			rawResponse := dnsGenLookupHostReplySuccess(rawQuery)
+			rawResponse := dnsGenLookupHostReplySuccess(rawQuery, nil)
 			query := &mocks.DNSQuery{
 				MockID: func() uint16 {
 					return queryID
@@ -323,7 +323,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 			})
 		})
 
-		t.Run("dnsResponse.LookupHost", func(t *testing.T) {
+		t.Run("dnsResponse.DecodeLookupHost", func(t *testing.T) {
 			t.Run("with failure", func(t *testing.T) {
 				// Ensure that we're not trying to decode if rcode != 0
 				d := &DNSDecoderMiekg{}
@@ -352,7 +352,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 				d := &DNSDecoderMiekg{}
 				queryID := dns.Id()
 				rawQuery := dnsGenQuery(dns.TypeA, queryID)
-				rawResponse := dnsGenLookupHostReplySuccess(rawQuery)
+				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, nil)
 				query := &mocks.DNSQuery{
 					MockID: func() uint16 {
 						return queryID
@@ -375,7 +375,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 				d := &DNSDecoderMiekg{}
 				queryID := dns.Id()
 				rawQuery := dnsGenQuery(dns.TypeA, queryID)
-				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, "1.1.1.1", "8.8.8.8")
+				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, nil, "1.1.1.1", "8.8.8.8")
 				query := &mocks.DNSQuery{
 					MockID: func() uint16 {
 						return queryID
@@ -407,7 +407,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 				d := &DNSDecoderMiekg{}
 				queryID := dns.Id()
 				rawQuery := dnsGenQuery(dns.TypeAAAA, queryID)
-				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, "::1", "fe80::1")
+				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, nil, "::1", "fe80::1")
 				query := &mocks.DNSQuery{
 					MockID: func() uint16 {
 						return queryID
@@ -439,7 +439,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 				d := &DNSDecoderMiekg{}
 				queryID := dns.Id()
 				rawQuery := dnsGenQuery(dns.TypeAAAA, queryID)
-				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, "1.1.1.1", "8.8.8.8")
+				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, nil, "1.1.1.1", "8.8.8.8")
 				query := &mocks.DNSQuery{
 					MockID: func() uint16 {
 						return queryID
@@ -465,7 +465,7 @@ func TestDNSDecoderMiekg(t *testing.T) {
 				d := &DNSDecoderMiekg{}
 				queryID := dns.Id()
 				rawQuery := dnsGenQuery(dns.TypeA, queryID)
-				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, "::1", "fe80::1")
+				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, nil, "::1", "fe80::1")
 				query := &mocks.DNSQuery{
 					MockID: func() uint16 {
 						return queryID
@@ -487,6 +487,82 @@ func TestDNSDecoderMiekg(t *testing.T) {
 				}
 			})
 		})
+
+		t.Run("dnsResponse.DecodeCNAME", func(t *testing.T) {
+			t.Run("with failure", func(t *testing.T) {
+				// Ensure that we're not trying to decode if rcode != 0
+				d := &DNSDecoderMiekg{}
+				queryID := dns.Id()
+				rawQuery := dnsGenQuery(dns.TypeA, queryID)
+				rawResponse := dnsGenReplyWithError(rawQuery, dns.RcodeRefused)
+				query := &mocks.DNSQuery{
+					MockID: func() uint16 {
+						return queryID
+					},
+				}
+				resp, err := d.DecodeResponse(rawResponse, query)
+				if err != nil {
+					t.Fatal(err)
+				}
+				cname, err := resp.DecodeCNAME()
+				if !errors.Is(err, ErrOODNSRefused) {
+					t.Fatal("unexpected err", err)
+				}
+				if cname != "" {
+					t.Fatal("expected empty cname result")
+				}
+			})
+
+			t.Run("with empty answer", func(t *testing.T) {
+				d := &DNSDecoderMiekg{}
+				queryID := dns.Id()
+				rawQuery := dnsGenQuery(dns.TypeA, queryID)
+				var expectedCNAME *dnsCNAMEAnswer = nil // explicity not set
+				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, expectedCNAME, "8.8.8.8")
+				query := &mocks.DNSQuery{
+					MockID: func() uint16 {
+						return queryID
+					},
+				}
+				resp, err := d.DecodeResponse(rawResponse, query)
+				if err != nil {
+					t.Fatal(err)
+				}
+				cname, err := resp.DecodeCNAME()
+				if !errors.Is(err, ErrOODNSNoAnswer) {
+					t.Fatal("unexpected err", err)
+				}
+				if cname != "" {
+					t.Fatal("expected empty cname result")
+				}
+			})
+
+			t.Run("with full answer", func(t *testing.T) {
+				expectedCNAME := &dnsCNAMEAnswer{
+					CNAME: "dns.google.",
+				}
+				d := &DNSDecoderMiekg{}
+				queryID := dns.Id()
+				rawQuery := dnsGenQuery(dns.TypeA, queryID)
+				rawResponse := dnsGenLookupHostReplySuccess(rawQuery, expectedCNAME, "8.8.8.8")
+				query := &mocks.DNSQuery{
+					MockID: func() uint16 {
+						return queryID
+					},
+				}
+				resp, err := d.DecodeResponse(rawResponse, query)
+				if err != nil {
+					t.Fatal(err)
+				}
+				cname, err := resp.DecodeCNAME()
+				if err != nil {
+					t.Fatal(err)
+				}
+				if cname != expectedCNAME.CNAME {
+					t.Fatal("unexpected cname", cname)
+				}
+			})
+		})
 	})
 }
 
@@ -522,9 +598,18 @@ func dnsGenReplyWithError(rawQuery []byte, code int) []byte {
 	return data
 }
 
+// ImplementationNote: dnsCNAMEAnswer could have been a string but then
+// dnsGenLookupHostReplySuccess invocations would have been confusing to read,
+// because they would not have had a boundary between CNAME and addrs.
+
+// dnsCNAMEAnswer is the DNS cname answer to include into a response.
+type dnsCNAMEAnswer struct {
+	CNAME string
+}
+
 // dnsGenLookupHostReplySuccess generates a successful DNS reply containing the given ips...
 // in the answers where each answer's type depends on the IP's type (A/AAAA).
-func dnsGenLookupHostReplySuccess(rawQuery []byte, ips ...string) []byte {
+func dnsGenLookupHostReplySuccess(rawQuery []byte, cname *dnsCNAMEAnswer, ips ...string) []byte {
 	query := new(dns.Msg)
 	err := query.Unpack(rawQuery)
 	runtimex.PanicOnError(err, "query.Unpack failed")
@@ -562,6 +647,17 @@ func dnsGenLookupHostReplySuccess(rawQuery []byte, ips ...string) []byte {
 			})
 		}
 	}
+	if cname != nil {
+		reply.Answer = append(reply.Answer, &dns.CNAME{
+			Hdr: dns.RR_Header{
+				Name:   question.Name,
+				Rrtype: dns.TypeCNAME,
+				Class:  dns.ClassINET,
+				Ttl:    0,
+			},
+			Target: cname.CNAME,
+		})
+	}
 	data, err := reply.Pack()
 	runtimex.PanicOnError(err, "reply.Pack failed")
 	return data

internal/netxlite/dnsovergetaddrinfo.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/miekg/dns"
 	"github.com/ooni/probe-cli/v3/internal/model"
+	"github.com/ooni/probe-cli/v3/internal/runtimex"
 )
 
 // dnsOverGetaddrinfoTransport is a DNSTransport using getaddrinfo.
@@ -33,6 +34,7 @@ func (txp *dnsOverGetaddrinfoTransport) RoundTrip(
 	}
 	resp := &dnsOverGetaddrinfoResponse{
 		addrs: addrs,
+		cname: "", // TODO: implement this functionality
 		query: query,
 	}
 	return resp, nil
@@ -40,6 +42,7 @@ func (txp *dnsOverGetaddrinfoTransport) RoundTrip(
 
 type dnsOverGetaddrinfoResponse struct {
 	addrs []string
+	cname string
 	query model.DNSQuery
 }
 
@@ -101,6 +104,7 @@ func (txp *dnsOverGetaddrinfoTransport) CloseIdleConnections() {
 }
 
 func (r *dnsOverGetaddrinfoResponse) Query() model.DNSQuery {
+	runtimex.PanicIfNil(r.query, "dnsOverGetaddrinfoResponse with nil query")
 	return r.query
 }
 
@@ -117,9 +121,19 @@ func (r *dnsOverGetaddrinfoResponse) DecodeHTTPS() (*model.HTTPSSvc, error) {
 }
 
 func (r *dnsOverGetaddrinfoResponse) DecodeLookupHost() ([]string, error) {
+	if len(r.addrs) <= 0 {
+		return nil, ErrOODNSNoAnswer
+	}
 	return r.addrs, nil
 }
 
 func (r *dnsOverGetaddrinfoResponse) DecodeNS() ([]*net.NS, error) {
 	return nil, ErrNoDNSTransport
 }
+
+func (r *dnsOverGetaddrinfoResponse) DecodeCNAME() (string, error) {
+	if r.cname == "" {
+		return "", ErrOODNSNoAnswer
+	}
+	return r.cname, nil
+}

internal/netxlite/dnsovergetaddrinfo_test.go

@@ -8,7 +8,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	"github.com/miekg/dns"
+	"github.com/ooni/probe-cli/v3/internal/model/mocks"
 )
 
 func TestDNSOverGetaddrinfo(t *testing.T) {
@@ -187,3 +189,157 @@ func TestDNSOverGetaddrinfo(t *testing.T) {
 		})
 	})
 }
+
+func TestDNSOverGetaddrinfoResponse(t *testing.T) {
+	t.Run("Query works as intended", func(t *testing.T) {
+		t.Run("when query is not nil", func(t *testing.T) {
+			resp := &dnsOverGetaddrinfoResponse{
+				addrs: []string{},
+				cname: "",
+				query: &mocks.DNSQuery{},
+			}
+			out := resp.Query()
+			if out != resp.query {
+				t.Fatal("unexpected query")
+			}
+		})
+
+		t.Run("when query is nil", func(t *testing.T) {
+			resp := &dnsOverGetaddrinfoResponse{
+				addrs: []string{},
+				cname: "",
+				query: nil, // oops
+			}
+			panicked := false
+			func() {
+				defer func() {
+					if recover() != nil {
+						panicked = true
+					}
+				}()
+				_ = resp.Query()
+			}()
+			if !panicked {
+				t.Fatal("did not panic")
+			}
+		})
+	})
+
+	t.Run("Bytes works as intended", func(t *testing.T) {
+		resp := &dnsOverGetaddrinfoResponse{
+			addrs: []string{},
+			cname: "",
+			query: nil,
+		}
+		if len(resp.Bytes()) > 0 {
+			t.Fatal("unexpected bytes")
+		}
+	})
+
+	t.Run("Rcode works as intended", func(t *testing.T) {
+		resp := &dnsOverGetaddrinfoResponse{
+			addrs: []string{},
+			cname: "",
+			query: nil,
+		}
+		if resp.Rcode() != 0 {
+			t.Fatal("unexpected rcode")
+		}
+	})
+
+	t.Run("DecodeHTTPS works as intended", func(t *testing.T) {
+		resp := &dnsOverGetaddrinfoResponse{
+			addrs: []string{},
+			cname: "",
+			query: nil,
+		}
+		out, err := resp.DecodeHTTPS()
+		if !errors.Is(err, ErrNoDNSTransport) {
+			t.Fatal("unexpected err")
+		}
+		if out != nil {
+			t.Fatal("unexpected result")
+		}
+	})
+
+	t.Run("DecodeLookupHost works as intended", func(t *testing.T) {
+		t.Run("on success", func(t *testing.T) {
+			resp := &dnsOverGetaddrinfoResponse{
+				addrs: []string{
+					"1.1.1.1", "1.0.0.1",
+				},
+				cname: "",
+				query: nil,
+			}
+			out, err := resp.DecodeLookupHost()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if diff := cmp.Diff(resp.addrs, out); diff != "" {
+				t.Fatal(diff)
+			}
+		})
+
+		t.Run("on failure", func(t *testing.T) {
+			resp := &dnsOverGetaddrinfoResponse{
+				addrs: []string{},
+				cname: "",
+				query: nil,
+			}
+			out, err := resp.DecodeLookupHost()
+			if !errors.Is(err, ErrOODNSNoAnswer) {
+				t.Fatal("unexpected err")
+			}
+			if len(out) > 0 {
+				t.Fatal("unexpected addrs")
+			}
+		})
+	})
+
+	t.Run("DecodeNS works as intended", func(t *testing.T) {
+		resp := &dnsOverGetaddrinfoResponse{
+			addrs: []string{},
+			cname: "",
+			query: nil,
+		}
+		out, err := resp.DecodeNS()
+		if !errors.Is(err, ErrNoDNSTransport) {
+			t.Fatal("unexpected err")
+		}
+		if len(out) != 0 {
+			t.Fatal("unexpected result")
+		}
+	})
+
+	t.Run("DecodeCNAME works as intended", func(t *testing.T) {
+		t.Run("on success", func(t *testing.T) {
+			resp := &dnsOverGetaddrinfoResponse{
+				addrs: []string{},
+				cname: "antani",
+				query: nil,
+			}
+			out, err := resp.DecodeCNAME()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if out != resp.cname {
+				t.Fatal("unexpected cname")
+			}
+		})
+
+		t.Run("on failure", func(t *testing.T) {
+			resp := &dnsOverGetaddrinfoResponse{
+				addrs: []string{},
+				cname: "",
+				query: nil,
+			}
+			out, err := resp.DecodeCNAME()
+			if !errors.Is(err, ErrOODNSNoAnswer) {
+				t.Fatal("unexpected err")
+			}
+			if out != "" {
+				t.Fatal("unexpected cname")
+			}
+		})
+	})
+}

internal/netxlite/dnsoverudp.go

@@ -22,10 +22,10 @@ import (
 // endpoint for several seconds when you query for blocked domains. We could also
 // have used an unconnected UDP socket here, but:
 //
-// 1. connected sockets are great because they get some ICMP errors to be
+// 1. connected UDP sockets are great because they get some ICMP errors to be
 // translated into socket errors (among them, host_unreachable);
 //
-// 2. connected sockets ignore responses from illegitimate IP addresses but
+// 2. connected UDP sockets ignore responses from illegitimate IP addresses but
 // most if not all DNS resolvers also do that, therefore this does not seem to
 // be a realistic censorship vector. At the same time, connected sockets
 // provide us for free with the feature that we don't need to bother with checking
@@ -34,8 +34,8 @@ import (
 // Being able to observe some ICMP errors is good because it could possibly
 // make this code suitable to implement parasitic traceroute.
 //
-// This transport is capable of collecting additional responses after the first
+// This transport by default listens for additional responses after the first
-// response. To see these responses, use the AsyncRoundTrip method.
+// one and makes them available using the context-configured trace.
 type DNSOverUDPTransport struct {
 	// Decoder is the MANDATORY DNSDecoder to use.
 	Decoder model.DNSDecoder
@@ -45,6 +45,11 @@ type DNSOverUDPTransport struct {
 
 	// Endpoint is the MANDATORY server's endpoint (e.g., 1.1.1.1:53)
 	Endpoint string
+
+	// lateResponses is posted in nonblocking mode each time this
+	// transport detects there was a late response for a query that had
+	// already been answered. Use this channel for testing.
+	lateResponses chan any
 }
 
 // NewUnwrappedDNSOverUDPTransport creates a DNSOverUDPTransport instance
@@ -63,20 +68,23 @@ type DNSOverUDPTransport struct {
 // have less control over which IP address is being used.
 func NewUnwrappedDNSOverUDPTransport(dialer model.Dialer, address string) *DNSOverUDPTransport {
 	return &DNSOverUDPTransport{
-		Decoder:  &DNSDecoderMiekg{},
+		Decoder:       &DNSDecoderMiekg{},
-		Dialer:   dialer,
+		Dialer:        dialer,
-		Endpoint: address,
+		Endpoint:      address,
+		lateResponses: nil, // not interested by default
 	}
 }
 
 // RoundTrip sends a query and receives a response.
 func (t *DNSOverUDPTransport) RoundTrip(
 	ctx context.Context, query model.DNSQuery) (model.DNSResponse, error) {
-	// QUIRK: the original code had a five seconds timeout, which is
+	// QUIRK: the original DNS-over-UDP code had a five seconds timeout, which is
-	// consistent with the Bionic implementation.
+	// consistent with the Bionic implementation. Let's try to preserve such a
+	// behavior by combining dialing and I/O timeout together.
 	//
 	// See https://labs.ripe.net/Members/baptiste_jonglez_1/persistent-dns-connections-for-reliability-and-performance
 	const opTimeout = 5 * time.Second
+	deadline := time.Now().Add(opTimeout)
 	ctx, cancel := context.WithTimeout(ctx, opTimeout)
 	defer cancel()
 	rawQuery, err := query.Bytes()
@@ -87,7 +95,7 @@ func (t *DNSOverUDPTransport) RoundTrip(
 	if err != nil {
 		return nil, err
 	}
-	conn.SetDeadline(time.Now().Add(opTimeout))
+	conn.SetDeadline(deadline) // time to dial (usually ~zero) already factored in
 	joinedch := make(chan bool)
 	myaddr := conn.LocalAddr().String()
 	if _, err := conn.Write(rawQuery); err != nil {
@@ -153,6 +161,13 @@ func (t *DNSOverUDPTransport) ownConnAndSendRecvLoop(ctx context.Context, conn n
 			// this seems how censorship is implemented in, e.g., China.
 			return
 		}
+		// if there's testing code waiting to be unblocked because we
+		// received a delayed response, unblock it
+		select {
+		case t.lateResponses <- true:
+		default:
+			// there's no one waiting and it does not matter
+		}
 		addrs, err := resp.DecodeLookupHost()
 		if err := trace.OnDelayedDNSResponse(started, t, query, resp, addrs, err, finished); err != nil {
 			// This error typically indicates that the buffer on which we're

internal/netxlite/dnsoverudp_test.go

@@ -284,6 +284,38 @@ func TestDNSOverUDPTransport(t *testing.T) {
 	})
 
 	t.Run("recording delayed DNS responses", func(t *testing.T) {
+		t.Run("without any context-injected traces", func(t *testing.T) {
+			srvr := &filtering.DNSServer{
+				OnQuery: func(domain string) filtering.DNSAction {
+					return filtering.DNSActionLocalHostPlusCache
+				},
+				Cache: map[string][]string{
+					"dns.google.": {"8.8.8.8"},
+				},
+			}
+			listener, err := srvr.Start("127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer listener.Close()
+			dialer := NewDialerWithoutResolver(model.DiscardLogger)
+			expectedAddress := listener.LocalAddr().String()
+			txp := NewUnwrappedDNSOverUDPTransport(dialer, expectedAddress)
+			txp.lateResponses = make(chan any, 1) // with buffer to avoid deadlocks
+			encoder := &DNSEncoderMiekg{}
+			query := encoder.Encode("dns.google.", dns.TypeA, false)
+			rch, err := txp.RoundTrip(context.Background(), query)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if _, err := rch.DecodeLookupHost(); err != nil {
+				t.Fatal(err)
+			}
+			// Now wait for the delayed response to arrive. We don't care much
+			// about observing it here, rather we want to know it happened.
+			<-txp.lateResponses
+		})
+
 		t.Run("uses a context-injected custom trace (success case)", func(t *testing.T) {
 			var (
 				delayedDNSResponseCalled bool
@@ -354,6 +386,7 @@ func TestDNSOverUDPTransport(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
+			defer mu.Unlock()
 			mu.Lock()
 			if diff := cmp.Diff(addrs, []string{"127.0.0.1"}); diff != "" {
 				t.Fatal(diff)
@@ -376,7 +409,6 @@ func TestDNSOverUDPTransport(t *testing.T) {
 			if !goodError {
 				t.Fatal("unexpected error encountered")
 			}
-			mu.Unlock()
 		})
 
 		t.Run("uses a context-injected custom trace (failure case)", func(t *testing.T) {
@@ -445,6 +477,7 @@ func TestDNSOverUDPTransport(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
+			defer mu.Unlock()
 			mu.Lock()
 			if diff := cmp.Diff(addrs, []string{"127.0.0.1"}); diff != "" {
 				t.Fatal(diff)
@@ -467,7 +500,6 @@ func TestDNSOverUDPTransport(t *testing.T) {
 			if !goodError {
 				t.Fatal("unexpected error encountered")
 			}
-			mu.Unlock()
 		})
 	})
 

internal/netxlite/utls_test.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/apex/log"
+	"github.com/ooni/probe-cli/v3/internal/model/mocks"
 	utls "gitlab.com/yawning/utls.git"
 )
 
@@ -93,23 +94,22 @@ func TestUTLSConn(t *testing.T) {
 		})
 	})
 
-	// TODO(https://github.com/ooni/probe/issues/2222): we cannot enable
+	t.Run("NetConn", func(t *testing.T) {
-	// this test until we use oocrypto >= v0.2 which uses go1.19. In turn,
+		factory := newConnUTLS(&utls.HelloChrome_70)
-	// we cannot use go1.19 as our main version until we upgrade psiphon
+		conn := &mocks.Conn{}
-	// such that it builds using go1.19, which is the issue in #2222.
+		tconn, err := factory(conn, &tls.Config{})
-	/*
+		if err != nil {
-		t.Run("NetConn", func(t *testing.T) {
+			t.Fatal(err)
-			factory := newConnUTLS(&utls.HelloChrome_70)
+		}
-			conn := &mocks.Conn{}
+		// TODO(https://github.com/ooni/probe/issues/2222): we cannot avoid
-			tconn, err := factory(conn, &tls.Config{})
+		// castedTconn until we use oocrypto >= v0.2 which uses go1.19. In turn,
-			if err != nil {
+		// we cannot use go1.19 as our main version until we upgrade psiphon
-				t.Fatal(err)
+		// such that it builds using go1.19, which is the issue in #2222.
-			}
+		castedTconn := tconn.(*utlsConn)
-			if tconn.NetConn() != conn {
+		if castedTconn.NetConn() != conn {
-				t.Fatal("NetConn is not WAI")
+			t.Fatal("NetConn is not WAI")
-			}
+		}
-		})
+	})
-	*/
 }
 
 func Test_newConnUTLSWithHelloID(t *testing.T) {