plugin/loadbalance: support prefer option (#7433)

Signed-off-by: Olli Janatuinen <olli.janatuinen@gmail.com>

plugin/loadbalance/prefer_test.go

@@ -0,0 +1,96 @@
+package loadbalance
+
+import (
+	"net"
+	"testing"
+
+	"github.com/coredns/coredns/plugin/test"
+
+	"github.com/miekg/dns"
+)
+
+func TestSortPreferred(t *testing.T) {
+	records := []dns.RR{
+		test.A("example.org. 300 IN A 10.9.30.1"),
+		test.A("example.org. 300 IN A 10.9.20.5"),
+		test.A("example.org. 300 IN A 192.168.1.2"),
+		test.A("example.org. 300 IN A 10.10.0.1"),
+		test.A("example.org. 300 IN A 10.9.20.3"),
+		test.A("example.org. 300 IN A 172.16.0.1"),
+		test.AAAA("example.org. 300 IN AAAA 2001:db8::1"),
+		test.AAAA("example.org. 300 IN AAAA 2001:db8:abcd::1"),
+		test.AAAA("example.org. 300 IN AAAA fd00::1"),
+		test.CNAME("example.org. 300 IN CNAME alias.example.org."),
+	}
+
+	subnets := []*net.IPNet{}
+	cidrs := []string{"2001:db8::/32", "10.9.20.0/24", "10.9.30.0/24"}
+	for _, cidr := range cidrs {
+		_, subnet, err := net.ParseCIDR(cidr)
+		if err != nil {
+			t.Fatalf("Failed to parse CIDR: %v", err)
+		}
+		subnets = append(subnets, subnet)
+	}
+
+	msg := &dns.Msg{Answer: records}
+	reorderPreferredSubnets(msg, subnets)
+	sorted := msg.Answer
+
+	expectedOrder := []string{
+		"alias.example.org.",
+		"2001:db8::1",
+		"2001:db8:abcd::1",
+		"10.9.20.5",
+		"10.9.20.3",
+		"10.9.30.1",
+		"192.168.1.2",
+		"10.10.0.1",
+		"172.16.0.1",
+		"fd00::1",
+	}
+
+	if len(sorted) != len(expectedOrder) {
+		t.Fatalf("Expected %d records, got %d", len(expectedOrder), len(sorted))
+	}
+
+	for i, rr := range sorted {
+		expected := expectedOrder[i]
+		switch r := rr.(type) {
+		case *dns.CNAME:
+			if r.Target != expected {
+				t.Errorf("Record %d: expected CNAME %s, got %s", i, expected, r.Target)
+			}
+		case *dns.A:
+			if r.A.String() != expected {
+				t.Errorf("Record %d: expected A IP %s, got %s", i, expected, r.A.String())
+			}
+		case *dns.AAAA:
+			if r.AAAA.String() != expected {
+				t.Errorf("Record %d: expected AAAA IP %s, got %s", i, expected, r.AAAA.String())
+			}
+		default:
+			t.Errorf("Record %d: unexpected RR type %T", i, r)
+		}
+	}
+}
+
+func TestExtractIP(t *testing.T) {
+	a := test.A("example.org. 300 IN A 10.0.0.1")
+	ip := extractIP(a)
+	if ip.String() != "10.0.0.1" {
+		t.Errorf("Expected 10.0.0.1, got %s", ip.String())
+	}
+
+	aaaa := test.AAAA("example.org. 300 IN AAAA ::1")
+	ip = extractIP(aaaa)
+	if ip.String() != "::1" {
+		t.Errorf("Expected ::1, got %s", ip.String())
+	}
+
+	cname := test.CNAME("example.org. 300 IN CNAME other.org.")
+	ip = extractIP(cname)
+	if ip != nil {
+		t.Errorf("Expected nil for CNAME, got %v", ip)
+	}
+}