fix(xdns): address code review findings

- Change resolverIdx from plain uint32 to atomic.Uint32 for safe
  round-robin counter access, matching the original design intent
- Track direct-mode recvLoop goroutine with recvWg so Close() waits
  for it to finish before returning
- Convert TestParseResolverAddr to use t.Run subtests and add error
  test cases for empty string and empty host inputs
- Add TestMultiResolverDistribution to verify queries distribute
  across multiple resolvers via round-robin

Author: nnemirovsky

transport/internet/finalmask/xdns/client.go

@@ -58,7 +58,7 @@ type xdnsConnClient struct {
 	domain   Name
 
 	resolverConns []*resolverConn
-	resolverIdx   uint32
+	resolverIdx   atomic.Uint32
 	recvWg        sync.WaitGroup
 	sendWg        sync.WaitGroup
 
@@ -114,7 +114,11 @@ func NewConnClient(c *Config, raw net.PacketConn) (net.PacketConn, error) {
 			}(rc.conn)
 		}
 	} else {
-		go conn.recvLoop()
+		conn.recvWg.Add(1)
+		go func() {
+			defer conn.recvWg.Done()
+			conn.recvLoop()
+		}()
 	}
 	conn.sendWg.Add(1)
 	go func() {
@@ -251,8 +255,7 @@ func (c *xdnsConnClient) sendLoop() {
 		if p != nil {
 			var err error
 			if len(c.resolverConns) > 0 {
-				idx := c.resolverIdx
-				c.resolverIdx++
+				idx := c.resolverIdx.Add(1)
 				rc := c.resolverConns[idx%uint32(len(c.resolverConns))]
 				_, err = rc.conn.WriteTo(p.p, rc.addr)
 			} else {
@@ -312,8 +315,8 @@ func (c *xdnsConnClient) Close() error {
 			rc.conn.Close()
 		}
 		c.closeErr = c.PacketConn.Close()
+		c.recvWg.Wait()
 		if len(c.resolverConns) > 0 {
-			c.recvWg.Wait()
 			close(c.pollChan)
 			close(c.readQueue)
 			c.mutex.Lock()

transport/internet/finalmask/xdns/dns_test.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -600,26 +601,38 @@ func TestRDataTXTRoundTrip(t *testing.T) {
 
 func TestParseResolverAddr(t *testing.T) {
 	tests := []struct {
+		name     string
 		input    string
 		wantIP   string
 		wantPort int
+		wantErr  bool
 	}{
-		{"1.1.1.1", "1.1.1.1", 53},
-		{"8.8.8.8:53", "8.8.8.8", 53},
-		{"8.8.8.8:5353", "8.8.8.8", 5353},
-		{"[2606:4700:4700::1111]:53", "2606:4700:4700::1111", 53},
+		{"bare_ipv4", "1.1.1.1", "1.1.1.1", 53, false},
+		{"ipv4_with_port", "8.8.8.8:53", "8.8.8.8", 53, false},
+		{"ipv4_custom_port", "8.8.8.8:5353", "8.8.8.8", 5353, false},
+		{"ipv6_with_port", "[2606:4700:4700::1111]:53", "2606:4700:4700::1111", 53, false},
+		{"empty_string", "", "", 0, true},
+		{"empty_host_with_port", ":53", "", 0, true},
 	}
 	for _, tt := range tests {
-		addr, err := parseResolverAddr(tt.input)
-		if err != nil {
-			t.Fatalf("parseResolverAddr(%q): %v", tt.input, err)
-		}
-		if addr.IP.String() != tt.wantIP {
-			t.Errorf("parseResolverAddr(%q).IP = %v, want %v", tt.input, addr.IP, tt.wantIP)
-		}
-		if addr.Port != tt.wantPort {
-			t.Errorf("parseResolverAddr(%q).Port = %v, want %v", tt.input, addr.Port, tt.wantPort)
-		}
+		t.Run(tt.name, func(t *testing.T) {
+			addr, err := parseResolverAddr(tt.input)
+			if tt.wantErr {
+				if err == nil {
+					t.Fatalf("parseResolverAddr(%q): want error, got nil", tt.input)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("parseResolverAddr(%q): %v", tt.input, err)
+			}
+			if addr.IP.String() != tt.wantIP {
+				t.Errorf("parseResolverAddr(%q).IP = %v, want %v", tt.input, addr.IP, tt.wantIP)
+			}
+			if addr.Port != tt.wantPort {
+				t.Errorf("parseResolverAddr(%q).Port = %v, want %v", tt.input, addr.Port, tt.wantPort)
+			}
+		})
 	}
 }
 
@@ -716,3 +729,70 @@ func TestResolverModeRoundTrip(t *testing.T) {
 		t.Errorf("server received %q, want %q", serverBuf[:readN], testPayload)
 	}
 }
+
+func TestMultiResolverDistribution(t *testing.T) {
+	const numResolvers = 3
+
+	// Create mock resolvers that count received packets
+	resolvers := make([]net.PacketConn, numResolvers)
+	var counts [numResolvers]atomic.Int32
+	for i := range resolvers {
+		r, err := net.ListenPacket("udp", "127.0.0.1:0")
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer r.Close()
+		resolvers[i] = r
+		idx := i
+		go func() {
+			buf := make([]byte, 4096)
+			for {
+				_, _, err := resolvers[idx].ReadFrom(buf)
+				if err != nil {
+					return
+				}
+				counts[idx].Add(1)
+			}
+		}()
+	}
+
+	resolverAddrs := make([]string, numResolvers)
+	for i, r := range resolvers {
+		resolverAddrs[i] = r.LocalAddr().String()
+	}
+
+	config := &Config{
+		Domain:    "t.example.com",
+		Resolvers: resolverAddrs,
+	}
+	rawConn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rawConn.Close()
+
+	client, err := NewConnClient(config, rawConn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	// Send enough messages to hit all resolvers via round-robin
+	for i := 0; i < numResolvers*3; i++ {
+		payload := []byte(fmt.Sprintf("msg-%d", i))
+		_, err = client.WriteTo(payload, rawConn.LocalAddr())
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// Allow sendLoop to process the queue
+	time.Sleep(500 * time.Millisecond)
+
+	for i := 0; i < numResolvers; i++ {
+		c := counts[i].Load()
+		if c == 0 {
+			t.Errorf("resolver %d received no queries, want at least 1", i)
+		}
+	}
+}