feat(dns): add resource size helpers

Co-Authored-By: Virgil <virgil@lethean.io>

pkg/dns/resource.go

@@ -122,6 +122,20 @@ type TXTRecordJSON struct {
 	Entries []string `json:"txt"`
 }
 
+// GetSize returns the encoded size of the resource payload.
+func (r *Resource) GetSize() int {
+	if r == nil {
+		return 0
+	}
+
+	size := 1
+	for _, record := range r.Records {
+		size += 1 + resourceRecordSize(record)
+	}
+
+	return size
+}
+
 // NewResource constructs a resource with the reference default TTL.
 func NewResource() *Resource {
 	return &Resource{TTL: DEFAULT_TTL}
@@ -358,6 +372,45 @@ func (r *Resource) FromJSON(jsonView ResourceJSON) error {
 	return nil
 }
 
+// GetSize returns the encoded size of the DS record.
+func (r DSRecord) GetSize() int {
+	return 5 + len(r.Digest)
+}
+
+// GetSize returns the encoded size of the NS record.
+func (r NSRecord) GetSize() int {
+	return resourceNameSize(r.NS)
+}
+
+// GetSize returns the encoded size of the GLUE4 record.
+func (r GLUE4Record) GetSize() int {
+	return resourceNameSize(r.NS) + 4
+}
+
+// GetSize returns the encoded size of the GLUE6 record.
+func (r GLUE6Record) GetSize() int {
+	return resourceNameSize(r.NS) + 16
+}
+
+// GetSize returns the encoded size of the SYNTH4 record.
+func (r SYNTH4Record) GetSize() int {
+	return 4
+}
+
+// GetSize returns the encoded size of the SYNTH6 record.
+func (r SYNTH6Record) GetSize() int {
+	return 16
+}
+
+// GetSize returns the encoded size of the TXT record.
+func (r TXTRecord) GetSize() int {
+	size := 1
+	for _, entry := range r.Entries {
+		size += 1 + len(entry)
+	}
+	return size
+}
+
 // DecodeResource decodes a raw DNS resource payload.
 func DecodeResource(raw []byte) (*Resource, error) {
 	resource := NewResource()
@@ -408,6 +461,80 @@ func resourceRecordJSON(record ResourceRecord) any {
 	}
 }
 
+func resourceRecordSize(record ResourceRecord) int {
+	switch rr := record.(type) {
+	case DSRecord:
+		return rr.GetSize()
+	case *DSRecord:
+		if rr == nil {
+			return 0
+		}
+		return rr.GetSize()
+	case NSRecord:
+		return rr.GetSize()
+	case *NSRecord:
+		if rr == nil {
+			return 0
+		}
+		return rr.GetSize()
+	case GLUE4Record:
+		return rr.GetSize()
+	case *GLUE4Record:
+		if rr == nil {
+			return 0
+		}
+		return rr.GetSize()
+	case GLUE6Record:
+		return rr.GetSize()
+	case *GLUE6Record:
+		if rr == nil {
+			return 0
+		}
+		return rr.GetSize()
+	case SYNTH4Record:
+		return rr.GetSize()
+	case *SYNTH4Record:
+		if rr == nil {
+			return 0
+		}
+		return rr.GetSize()
+	case SYNTH6Record:
+		return rr.GetSize()
+	case *SYNTH6Record:
+		if rr == nil {
+			return 0
+		}
+		return rr.GetSize()
+	case TXTRecord:
+		return rr.GetSize()
+	case *TXTRecord:
+		if rr == nil {
+			return 0
+		}
+		return rr.GetSize()
+	default:
+		return 0
+	}
+}
+
+func resourceNameSize(name string) int {
+	name = fqdn(name)
+	if name == "." {
+		return 1
+	}
+
+	trimmed := trimFQDN(name)
+	if trimmed == "" {
+		return 1
+	}
+
+	size := 1
+	for _, label := range strings.Split(trimmed, ".") {
+		size += 1 + len(label)
+	}
+	return size
+}
+
 type resourceRecordProbe struct {
 	Type string `json:"type"`
 }

pkg/dns/resource_test.go

@@ -162,6 +162,61 @@ func TestResourceJSONRoundTrip(t *testing.T) {
 	}
 }
 
+func TestResourceSizeHelpers(t *testing.T) {
+	resource := NewResource()
+	resource.Records = []ResourceRecord{
+		DSRecord{
+			KeyTag:     1,
+			Algorithm:  2,
+			DigestType: 3,
+			Digest:     []byte{0xaa, 0xbb, 0xcc},
+		},
+		NSRecord{NS: "ns1.example."},
+		GLUE4Record{NS: "ns1.example.", Address: netip.MustParseAddr("192.0.2.1")},
+		GLUE6Record{NS: "ns2.example.", Address: netip.MustParseAddr("2001:db8::1")},
+		SYNTH4Record{Address: netip.MustParseAddr("198.51.100.9")},
+		SYNTH6Record{Address: netip.MustParseAddr("2001:db8::9")},
+		TXTRecord{Entries: []string{"hello", "world"}},
+	}
+
+	raw, err := resource.Encode()
+	if err != nil {
+		t.Fatalf("Encode returned error: %v", err)
+	}
+
+	if got := resource.GetSize(); got != len(raw) {
+		t.Fatalf("Resource.GetSize() = %d, want %d", got, len(raw))
+	}
+
+	if got := (DSRecord{Digest: []byte{1, 2, 3}}).GetSize(); got != 8 {
+		t.Fatalf("DSRecord.GetSize() = %d, want 8", got)
+	}
+
+	if got := (NSRecord{NS: "ns1.example."}).GetSize(); got != resourceNameSize("ns1.example.") {
+		t.Fatalf("NSRecord.GetSize() = %d, want %d", got, resourceNameSize("ns1.example."))
+	}
+
+	if got := (GLUE4Record{NS: "ns1.example."}).GetSize(); got != resourceNameSize("ns1.example.")+4 {
+		t.Fatalf("GLUE4Record.GetSize() = %d, want %d", got, resourceNameSize("ns1.example.")+4)
+	}
+
+	if got := (GLUE6Record{NS: "ns1.example."}).GetSize(); got != resourceNameSize("ns1.example.")+16 {
+		t.Fatalf("GLUE6Record.GetSize() = %d, want %d", got, resourceNameSize("ns1.example.")+16)
+	}
+
+	if got := (SYNTH4Record{}).GetSize(); got != 4 {
+		t.Fatalf("SYNTH4Record.GetSize() = %d, want 4", got)
+	}
+
+	if got := (SYNTH6Record{}).GetSize(); got != 16 {
+		t.Fatalf("SYNTH6Record.GetSize() = %d, want 16", got)
+	}
+
+	if got := (TXTRecord{Entries: []string{"hello", "world"}}).GetSize(); got != 13 {
+		t.Fatalf("TXTRecord.GetSize() = %d, want 13", got)
+	}
+}
+
 func TestResourceTypeHelpers(t *testing.T) {
 	resource := NewResource()
 	resource.Records = []ResourceRecord{