From 9bdfaeabd9f0c261eada3346bb54ad5db500f00a Mon Sep 17 00:00:00 2001
From: fhgffy <102001626+fhgffy@users.noreply.github.com>
Date: Wed, 14 Oct 2026 00:30:17 +0800
Subject: [PATCH] fix(dns): validate address record lengths

---
 dns/message.go      |  6 +++++
 dns/message_test.go | 59 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 dns/message_test.go

diff --git a/dns/message.go b/dns/message.go
index 863f987..4560ec4 100644
--- a/dns/message.go
+++ b/dns/message.go
@@ -430,8 +430,14 @@ func (m *Message) UnmarshalRR(start int, rr *RR) (n int, err error) {
 	rr.RDATA = p[n+10 : n+10+int(rr.RDLENGTH)]
 
 	if rr.TYPE == QTypeA {
+		if rr.RDLENGTH != 4 {
+			return 0, errors.New("UnmarshalRR: invalid A record length")
+		}
 		rr.IP = netip.AddrFrom4(*(*[4]byte)(rr.RDATA[:4]))
 	} else if rr.TYPE == QTypeAAAA {
+		if rr.RDLENGTH != 16 {
+			return 0, errors.New("UnmarshalRR: invalid AAAA record length")
+		}
 		rr.IP = netip.AddrFrom16(*(*[16]byte)(rr.RDATA[:16]))
 	}
 
diff --git a/dns/message_test.go b/dns/message_test.go
new file mode 100644
index 0000000..db93f2b
--- /dev/null
+++ b/dns/message_test.go
@@ -0,0 +1,59 @@
+package dns
+
+import (
+	"fmt"
+	"net/netip"
+	"testing"
+)
+
+func TestUnmarshalAddressRecordLength(t *testing.T) {
+	for _, tt := range []struct {
+		qtype uint16
+		size  int
+		valid bool
+	}{
+		{QTypeA, 0, false},
+		{QTypeA, 1, false},
+		{QTypeA, 3, false},
+		{QTypeA, 4, true},
+		{QTypeA, 5, false},
+		{QTypeAAAA, 0, false},
+		{QTypeAAAA, 1, false},
+		{QTypeAAAA, 15, false},
+		{QTypeAAAA, 16, true},
+		{QTypeAAAA, 17, false},
+	} {
+		t.Run(fmt.Sprintf("type_%d_length_%d", tt.qtype, tt.size), func(t *testing.T) {
+			defer func() {
+				if value := recover(); value != nil {
+					t.Errorf("UnmarshalMessage panicked: %v", value)
+				}
+			}()
+			data := make([]byte, tt.size)
+			for i := range data {
+				data[i] = byte(i + 1)
+			}
+			message := NewMessage(1, ResponseMsg)
+			message.SetQuestion(NewQuestion(tt.qtype, "example.com"))
+			message.AddAnswer(&RR{NAME: "example.com", TYPE: tt.qtype, CLASS: ClassINET, TTL: 60, RDLENGTH: uint16(len(data)), RDATA: data})
+			wire, err := message.Marshal()
+			if err != nil {
+				t.Fatal(err)
+			}
+			parsed, err := UnmarshalMessage(wire[:len(wire):len(wire)])
+			if !tt.valid {
+				if err == nil {
+					t.Fatal("accepted an invalid address record length")
+				}
+				return
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			want, _ := netip.AddrFromSlice(data)
+			if len(parsed.Answers) != 1 || parsed.Answers[0].IP != want {
+				t.Fatalf("decoded address = %v, want %s", parsed.Answers, want)
+			}
+		})
+	}
+}