fix(dns): decode up to four answer records (#186)

* fix(dns): decode up to four answer records

A single DNS question can return multiple answer records. Use an answer
decode limit of four in Client.StartResolve and add a regression test
covering a single-question response with multiple A records.

* fix(dns): add MaxResponseAnswers to ResolveConfig

This allows callers to explicitly declare the maximum number of answer
records to retain, removing the hardcoded limit in Client.StartResolve.

xnet.StackAsync is updated to set this limit to match the length of its
lookup-result buffer. This maintains the zero-allocation design while
fixing the issue where multiple A records were ignored.

dns/client.go

@@ -24,6 +24,9 @@ type ResolveConfig struct {
 	Questions       []Question
 	Additional      []Resource
 	EnableRecursion bool
+	// MaxResponseAnswers limits how many answer records are decoded from the
+	// DNS response. If zero it defaults to the number of Questions.
+	MaxResponseAnswers uint16
 }
 
 func (sudp *Client) Protocol() uint64 { return uint64(lneto.IPProtoUDP) }
@@ -37,8 +40,12 @@ func (c *Client) StartResolve(localPort, txid uint16, cfg ResolveConfig) error {
 	if nd > math.MaxUint16 {
 		return lneto.ErrInvalidConfig
 	}
+	maxAns := cfg.MaxResponseAnswers
+	if maxAns == 0 {
+		maxAns = uint16(nd)
+	}
 	c.reset(localPort, txid, CQueryPending, cfg.EnableRecursion)
-	c.msg.LimitResourceDecoding(uint16(nd), uint16(nd), 0, 0)
+	c.msg.LimitResourceDecoding(uint16(nd), maxAns, 0, 0)
 	c.msg.AddQuestions(cfg.Questions)
 	c.msg.AddAdditionals(cfg.Additional)
 	return nil

dns/dns_test.go

@@ -243,76 +243,100 @@ func TestClient_ReceivesDNSResponse(t *testing.T) {
 	const hostname = "example.com"
 	const txid = uint16(12345)
 	const clientPort = uint16(54321)
-	wantIP := [4]byte{93, 184, 216, 34}
+	const maxAnswers = 4
-
+	allIPs := [5][4]byte{
-	// Build a DNS response message.
+		{192, 0, 2, 1},
-	name := MustNewName(hostname)
+		{192, 0, 2, 2},
-	responseMsg := Message{
+		{192, 0, 2, 3},
-		Questions: []Question{{
+		{192, 0, 2, 4},
-			Name:  name,
+		{192, 0, 2, 5},
-			Type:  TypeA,
-			Class: ClassINET,
-		}},
-		Answers: []Resource{
-			NewResource(name, TypeA, ClassINET, 300, wantIP[:]),
-		},
 	}
-
+	tests := []struct {
-	// Response flags: QR=1 (response), RD=1, RA=1.
+		name        string
-	responseFlags := HeaderFlags(1<<15 | 1<<8 | 1<<7)
+		responseIPs [][4]byte
-
+		wantAnswers int
-	var buf [512]byte
+	}{
-	dnsPayload, err := responseMsg.AppendTo(buf[:0], txid, responseFlags)
+		{name: "single_answer", responseIPs: allIPs[:1], wantAnswers: 1},
-	if err != nil {
+		{name: "multiple_answers", responseIPs: allIPs[:4], wantAnswers: 4},
-		t.Fatal("failed to build DNS response:", err)
+		{name: "answer_limit", responseIPs: allIPs[:5], wantAnswers: maxAnswers},
 	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			name := MustNewName(hostname)
+			responseMsg := Message{
+				Questions: []Question{{
+					Name:  name,
+					Type:  TypeA,
+					Class: ClassINET,
+				}},
+				Answers: make([]Resource, len(tt.responseIPs)),
+			}
+			for i := range tt.responseIPs {
+				responseMsg.Answers[i] = NewResource(name, TypeA, ClassINET, 300, tt.responseIPs[i][:])
+			}
 
-	// Set up the DNS client.
+			// Response flags: QR=1 (response), RD=1, RA=1.
-	var client Client
+			responseFlags := HeaderFlags(1<<15 | 1<<8 | 1<<7)
-	client.StartResolve(clientPort, txid, ResolveConfig{
+			var responseBuf [512]byte
-		Questions: []Question{{
+			dnsPayload, err := responseMsg.AppendTo(responseBuf[:0], txid, responseFlags)
-			Name:  MustNewName(hostname),
+			if err != nil {
-			Type:  TypeA,
+				t.Fatal("failed to build DNS response:", err)
-			Class: ClassINET,
+			}
-		}},
-		EnableRecursion: true,
-	})
 
-	// Simulate sending by calling Encapsulate (changes state to AwaitResponse).
+			var client Client
-	var dummy [512]byte
+			err = client.StartResolve(clientPort, txid, ResolveConfig{
-	client.Encapsulate(dummy[:], 0, 0)
+				Questions: []Question{{
+					Name:  name,
+					Type:  TypeA,
+					Class: ClassINET,
+				}},
+				EnableRecursion:    true,
+				MaxResponseAnswers: maxAnswers,
+			})
+			if err != nil {
+				t.Fatal("failed to start DNS resolve:", err)
+			}
 
-	// Call Demux with DNS payload.
+			// Encapsulate the query to move the client into the outstanding state.
-	err = client.Demux(dnsPayload, 0)
+			var queryBuf [512]byte
-	if err != nil {
+			_, err = client.Encapsulate(queryBuf[:], 0, 0)
-		t.Fatal("Client Demux error:", err)
+			if err != nil {
-	}
+				t.Fatal("failed to encapsulate DNS query:", err)
+			}
+			if err := client.Demux(dnsPayload, 0); err != nil {
+				t.Fatal("failed to demux DNS response:", err)
+			}
 
-	// Check the client received the answer.
+			var addrs [maxAnswers]netip.Addr
-	var addrs [4]netip.Addr
+			answers, err := client.ResponseAnswerLookup(addrs[:], hostname)
-	answers, err := client.ResponseAnswerLookup(addrs[:], hostname)
+			if err != nil {
-	if answers != 1 {
+				t.Fatal("failed to look up DNS response answers:", err)
-		t.Fatalf("expected 1 answer, got %d", answers)
+			}
-	}
+			if answers != uint16(tt.wantAnswers) {
-	addr := addrs[0]
+				t.Fatalf("expected %d answers, got %d", tt.wantAnswers, answers)
-	if !addr.Is4() {
+			}
-		t.Fatalf("expected 4 bytes in answer, got %d", addr.BitLen()/8)
+			for i := 0; i < tt.wantAnswers; i++ {
-	}
+				addr := addrs[i]
-	if addr.As4() != wantIP {
+				if !addr.Is4() {
-		t.Errorf("expected IP %v, got %v", wantIP, addr.String())
+					t.Errorf("answer %d: expected IPv4 address, got %v", i, addr)
-	}
+					continue
+				}
+				if addr.As4() != tt.responseIPs[i] {
+					t.Errorf("answer %d: expected IP %v, got %v", i, tt.responseIPs[i], addr)
+				}
+			}
 
-	// Test MessageCopyTo as well.
+			var lookup Message
-	var lookup Message
+			done, err := client.ResponseCopyTo(&lookup)
-	lookup.LimitResourceDecoding(1, 1, 0, 0)
+			if err != nil {
-	done, err := client.ResponseCopyTo(&lookup)
+				t.Fatal("failed to copy DNS response:", err)
-	if err != nil {
+			}
-		t.Fatal("MessageCopyTo error:", err)
+			if !done {
-	}
+				t.Fatal("expected done=true")
-	if !done {
+			}
-		t.Fatal("expected done=true")
+			if len(lookup.Answers) != tt.wantAnswers {
-	}
+				t.Fatalf("expected %d copied answers, got %d", tt.wantAnswers, len(lookup.Answers))
-	if len(lookup.Answers) != 1 {
+			}
-		t.Fatalf("MessageCopyTo: expected 1 answer, got %d", len(lookup.Answers))
+		})
 	}
 }

x/xnet/stack-async.go

@@ -629,7 +629,8 @@ func (s *StackAsync) StartLookupIPType(host string, qtype dns.Type) error {
 		Additional: []dns.Resource{
 			s.ednsopt,
 		},
-		EnableRecursion: true,
+		EnableRecursion:    true,
+		MaxResponseAnswers: uint16(len(s.addrbufnip)),
 	})
 	if err != nil {
 		return err