fix(dns): validate address record lengths

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]))
 	}
 

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)
+			}
+		})
+	}
+}