fix(client): follow referrals in Trace

client.go

@@ -11,6 +11,7 @@ import (
 	"math/rand"
 	"net"
 	"net/url"
+	"slices"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -633,7 +634,9 @@ func (c *Client) QueryParallel(host string, requestType uint16, resolvers []stri
 	return dnsdatas, nil
 }
 
-// Trace the requested domain with the provided query type
+// Trace the requested domain with the provided query type, following referrals
+// from the root servers. Besides its reply's answer, each step's NS lists the
+// name servers that the reply referred to.
 func (c *Client) Trace(host string, requestType uint16, maxrecursion int) (*TraceData, error) {
 	var tracedata TraceData
 	host = dns.CanonicalName(host)
@@ -663,18 +666,17 @@ func (c *Client) Trace(host string, requestType uint16, maxrecursion int) (*Trac
 			}
 		}
 
-		var newNSResolvers []string
+		var nsNames []string
+		glue := make(map[string][]string)
 		var nextCname string
 		for _, d := range dnsdatas {
-			// Add ns records as new resolvers
-			for _, ns := range d.NS {
-				nsData, err := c.A(ns)
-				if err != nil {
-					continue
-				}
-				for _, ip := range nsData.A {
-					newNSResolvers = append(newNSResolvers, net.JoinHostPort(ip, "53"))
-				}
+			// A referral lists the next name servers in the authority section,
+			// which ParseFromMsg leaves out of d.NS
+			names, addrs := referral(d.RawResp)
+			d.NS = sliceutil.Dedupe(append(d.NS, names...))
+			nsNames = append(nsNames, d.NS...)
+			for name, ips := range addrs {
+				glue[name] = append(glue[name], ips...)
 			}
 			// Follow CNAME - should happen at the final step of the trace
 			for _, cname := range d.CNAME {
@@ -684,6 +686,22 @@ func (c *Client) Trace(host string, requestType uint16, maxrecursion int) (*Trac
 				}
 			}
 		}
+		// Every root server returns the same referral, so resolve each name
+		// server once, from its glue when the referral has some
+		var newNSResolvers []string
+		for _, ns := range sliceutil.Dedupe(nsNames) {
+			ips := glue[ns]
+			if len(ips) == 0 {
+				nsData, err := c.A(ns)
+				if err != nil {
+					continue
+				}
+				ips = nsData.A
+			}
+			for _, ip := range ips {
+				newNSResolvers = append(newNSResolvers, net.JoinHostPort(ip, "53"))
+			}
+		}
 		newNSResolvers = sliceutil.Dedupe(newNSResolvers)
 
 		// if we have no new resolvers => return
@@ -713,6 +731,31 @@ func (c *Client) Trace(host string, requestType uint16, maxrecursion int) (*Trac
 	return &tracedata, nil
 }
 
+// referral returns the name servers in msg's authority section, where a
+// referral lists them, and the IPv4 glue that its additional section gives for
+// them, keyed by name. Glue for any other name is left out.
+func referral(msg *dns.Msg) (names []string, glue map[string][]string) {
+	if msg == nil {
+		return nil, nil
+	}
+	for _, rr := range msg.Ns {
+		if ns, ok := rr.(*dns.NS); ok {
+			names = append(names, trimChars(ns.Ns))
+		}
+	}
+	glue = make(map[string][]string)
+	for _, rr := range msg.Extra {
+		a, ok := rr.(*dns.A)
+		if !ok {
+			continue
+		}
+		if name := trimChars(a.Hdr.Name); slices.Contains(names, name) {
+			glue[name] = append(glue[name], a.A.String())
+		}
+	}
+	return names, glue
+}
+
 func (c *Client) axfr(host string) (*AXFRData, error) {
 	// obtain ns servers
 	dnsData, err := c.NS(host)

client_test.go

@@ -1,6 +1,7 @@
 package retryabledns
 
 import (
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -186,6 +187,90 @@ func TestTrace(t *testing.T) {
 	require.NoError(t, err)
 }
 
+// TestTraceFollowsReferrals ensures Trace follows a referral, which leaves the
+// answer section empty, lists the next name servers in the authority section
+// and may give their addresses (glue) in the additional section. Each step
+// must list the name servers its reply referred to, and Trace must look up a
+// name server only when the referral has no glue for it, once however many
+// root servers returned the referral.
+func TestTraceFollowsReferrals(t *testing.T) {
+	const nsName = "ns1.trace.test."
+	for name, glue := range map[string]bool{"without glue": false, "with glue": true} {
+		t.Run(name, func(t *testing.T) {
+			var nsLookups atomic.Int64
+			respond := func(req *dns.Msg) *dns.Msg {
+				m := &dns.Msg{}
+				m.SetReply(req)
+				if req.Question[0].Name == nsName {
+					nsLookups.Add(1)
+					m.Answer = []dns.RR{arecord(nsName, "127.0.0.1")}
+					return m
+				}
+				m.Ns = []dns.RR{&dns.NS{
+					Hdr: dns.RR_Header{Name: "test.", Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: 60},
+					Ns:  nsName,
+				}}
+				if glue {
+					m.Extra = []dns.RR{arecord(nsName, "127.0.0.1")}
+				}
+				return m
+			}
+			root1 := newFakeResolverFunc(t, respond)
+			root2 := newFakeResolverFunc(t, respond)
+			roots := RootDNSServersIPv4
+			RootDNSServersIPv4 = []string{root1.Addr(), root2.Addr()}
+			t.Cleanup(func() { RootDNSServersIPv4 = roots })
+
+			client, err := NewWithOptions(Options{
+				BaseResolvers: []string{root1.Addr()},
+				MaxRetries:    1,
+				Timeout:       time.Second,
+			})
+			require.NoError(t, err)
+
+			// A max recursion of 2 runs one step: query the roots, then find
+			// the address of the name server they referred to.
+			trace, err := client.Trace(fakeHost, dns.TypeA, 2)
+			require.NoError(t, err)
+			require.Len(t, trace.DNSData, 2, "one step per root server")
+			for _, d := range trace.DNSData {
+				require.Equal(t, []string{"ns1.trace.test"}, d.NS, "a step must list the name servers its reply referred to")
+			}
+			wantLookups := 1
+			if glue {
+				wantLookups = 0
+			}
+			require.EqualValues(t, wantLookups, nsLookups.Load(), "name server lookups")
+		})
+	}
+}
+
+// TestReferral ensures referral reads the name servers from the authority
+// section and keeps only the IPv4 glue for them.
+func TestReferral(t *testing.T) {
+	ns := func(target string) dns.RR {
+		return &dns.NS{
+			Hdr: dns.RR_Header{Name: "test.", Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: 60},
+			Ns:  target,
+		}
+	}
+	msg := &dns.Msg{
+		Ns: []dns.RR{ns("ns1.trace.test."), ns("ns2.trace.test.")},
+		Extra: []dns.RR{
+			arecord("ns1.trace.test", "192.0.2.1"),
+			aaaarecord("ns1.trace.test", "2001:db8::1"),
+			arecord("unrelated.test", "192.0.2.9"),
+		},
+	}
+	names, glue := referral(msg)
+	require.Equal(t, []string{"ns1.trace.test", "ns2.trace.test"}, names)
+	require.Equal(t, map[string][]string{"ns1.trace.test": {"192.0.2.1"}}, glue)
+
+	names, glue = referral(nil)
+	require.Empty(t, names)
+	require.Empty(t, glue)
+}
+
 func TestAXFRWithUDPResolvers(t *testing.T) {
 	// Resolver address configured without a tcp: prefix defaults to
 	// UDP. AXFR is TCP-only, so axfr() must force TCP on the fallback

proxy_test.go

@@ -400,8 +400,8 @@ func TestTraceUsesProxy(t *testing.T) {
 			m.Answer = []dns.RR{arecord(nsName, "127.0.0.1")}
 			return m
 		}
-		// The trace reads name servers from the answer section only.
-		m.Answer = []dns.RR{&dns.NS{
+		// A referral: no answer, and the name server in the authority section.
+		m.Ns = []dns.RR{&dns.NS{
 			Hdr: dns.RR_Header{Name: req.Question[0].Name, Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: 60},
 			Ns:  nsName,
 		}}