migration to new dns library

3 files changed, 179 insertions, 66 deletions

diff --git a/internal/resolver/resolver.go b/internal/resolver/resolver.go
index 9ee9cd6..5aa7bc1 100644
--- a/internal/resolver/resolver.go
+++ b/internal/resolver/resolver.go
@@ -7,7 +7,7 @@ import (
 	"net"
 	"time"
 
-	"github.com/miekg/dns"
+	"codeberg.org/miekg/dns"
 )
 
 var (
@@ -26,16 +26,15 @@ type Resolver struct {
 type Option func(*Resolver)
 
 func New(opts ...Option) *Resolver {
+	transport := dns.NewTransport()
+	transport.ReadTimeout = 2 * time.Second
+
 	r := &Resolver{
 		roots:          loadRootServers(),
 		maxDelegations: 30,
 		timeout:        2 * time.Second,
 		retries:        2,
-		client: &dns.Client{
-			Net: "udp",
-			UDPSize: 4096,
-			Timeout: 2 * time.Second,
-		},
+		client:         &dns.Client{Transport: transport},
 	}
 	for _, opt := range opts {
 		opt(r)
@@ -52,6 +51,7 @@ func WithRootAddresses(addrs []string) Option {
 func WithTimeout(d time.Duration) Option {
 	return func(r *Resolver) {
 		r.timeout = d
+		r.client.Transport.ReadTimeout = d
 	}
 }
 
@@ -81,6 +81,9 @@ func (r *Resolver) resolve(ctx context.Context, qname string, qtype uint16) (*dn
 		}
 		switch {
 		case reply.Rcode == dns.RcodeSuccess && len(reply.Answer) > 0:
+			if needsCNAMEResolution(reply, qtype) {
+				return r.resolveCNAME(ctx, reply, qtype)
+			}
 			return reply, nil
 		case reply.Rcode == dns.RcodeNameError:
 			return reply, nil
@@ -104,6 +107,62 @@ func (r *Resolver) resolve(ctx context.Context, qname string, qtype uint16) (*dn
 	return nil, ErrMaxDelegations
 }
 
+func needsCNAMEResolution(reply *dns.Msg, qtype uint16) bool {
+	hasCNAME := false
+	hasTarget := false
+	for _, rr := range reply.Answer {
+		if _, ok := rr.(*dns.CNAME); ok {
+			hasCNAME = true
+		}
+		if dns.RRToType(rr) == qtype && dns.RRToType(rr) != dns.TypeCNAME {
+			hasTarget = true
+		}
+	}
+	return hasCNAME && !hasTarget
+}
+func (r *Resolver) resolveCNAME(ctx context.Context, reply *dns.Msg, qtype uint16) (*dns.Msg, error) {
+	var target string
+	for _, rr := range reply.Answer {
+		if cn, ok := rr.(*dns.CNAME); ok {
+			target = cn.Target
+		}
+	}
+	if target == "" {
+        return reply, nil
+    }
+
+    const maxChain = 10
+    seen := map[string]bool{reply.Question[0].Header().Name: true}
+    for i := 0; i < maxChain; i++ {
+        select {
+        case <-ctx.Done():
+            return nil, ctx.Err()
+        default:
+        }
+        if seen[target] {
+            break
+        }
+        seen[target] = true
+
+        chainReply, err := r.resolve(ctx, target, qtype)
+        if err != nil {
+            return reply, nil
+        }
+        reply.Answer = append(reply.Answer, chainReply.Answer...)
+
+        nextTarget := ""
+        for _, rr := range chainReply.Answer {
+            if cn, ok := rr.(*dns.CNAME); ok && cn.Header().Name == target {
+                nextTarget = cn.Target
+            }
+        }
+        if nextTarget == "" {
+            break
+        }
+        target = nextTarget
+    }
+    return reply, nil
+}
 func isReferral(msg *dns.Msg) bool {
 	return !msg.Authoritative && len(msg.Ns) > 0
 }
@@ -122,12 +181,12 @@ func (r *Resolver) nextServers(ctx context.Context, msg *dns.Msg) ([]string, err
 	for _, rr := range msg.Extra {
 		switch v := rr.(type) {
 		case *dns.A:
-			if _, exists := glue[v.Hdr.Name]; !exists {
-				glue[v.Hdr.Name] = v.A.String()
+			if _, exists := glue[v.Header().Name]; !exists {
+				glue[v.Header().Name] = v.A.Addr.String()
 			}
 		case *dns.AAAA:
-			if _, exists := glue[v.Hdr.Name]; !exists {
-				glue[v.Hdr.Name] = v.AAAA.String()
+			if _, exists := glue[v.Header().Name]; !exists {
+				glue[v.Header().Name] = v.AAAA.Addr.String()
 			}
 		}
 	}
@@ -156,7 +215,7 @@ func (r *Resolver) nextServers(ctx context.Context, msg *dns.Msg) ([]string, err
 		}
 		for _, rr := range reply.Answer {
 			if a, ok := rr.(*dns.A); ok {
-				addrs = append(addrs, a.A.String())
+				addrs = append(addrs, a.A.Addr.String())
 				break
 			}
 		}
@@ -171,9 +230,8 @@ func (r *Resolver) nextServers(ctx context.Context, msg *dns.Msg) ([]string, err
 func (r *Resolver) exchangeWithRetries(ctx context.Context, servers []string,
 	qname string, qtype uint16) (*dns.Msg, error) {
 
-	msg := new(dns.Msg)
-	msg.SetQuestion(qname, qtype)
-	msg.SetEdns0(4096, false)
+	msg := dns.NewMsg(qname, qtype)
+	msg.UDPSize = 4096
 	msg.RecursionDesired = false
 
 	type result struct {
@@ -189,11 +247,10 @@ func (r *Resolver) exchangeWithRetries(ctx context.Context, servers []string,
 		}
 		go func(addr string) {
 			for attempt := 0; attempt < r.retries; attempt++ {
-				reply, _, err := r.client.ExchangeContext(ctx, msg, addr)
+				reply, _, err := r.client.Exchange(ctx, msg, "udp", addr)
 				if err == nil {
 					if reply.Truncated {
-						tcpClient := &dns.Client{Net: "tcp", Timeout: r.timeout}
-						reply, _, err = tcpClient.ExchangeContext(ctx, msg, addr)
+						reply, _, err = r.client.Exchange(ctx, msg, "tcp", addr)
 						if err == nil {
 							ch <- result{reply: reply, err: nil}
 							return
diff --git a/internal/resolver/resolver_test.go b/internal/resolver/resolver_test.go
index 0bd0402..54f727e 100644
--- a/internal/resolver/resolver_test.go
+++ b/internal/resolver/resolver_test.go
@@ -2,54 +2,53 @@ package resolver
 
 import (
 	"context"
-	"net"
+	"io"
+	"net/netip"
 	"testing"
 	"time"
 
-	"github.com/miekg/dns"
+	"codeberg.org/miekg/dns"
+	"codeberg.org/miekg/dns/dnsutil"
+	"codeberg.org/miekg/dns/rdata"
 )
 
 func startTestServer(t *testing.T, addr string, handler dns.Handler) *dns.Server {
 	t.Helper()
+	ready := make(chan struct{})
 	srv := &dns.Server{
 		Addr:    addr,
 		Net:     "udp",
 		Handler: handler,
 		UDPSize: 4096,
+		NotifyStartedFunc: func(ctx context.Context) {
+			close(ready)
+		},
 	}
 	go func() {
-		if err := srv.ListenAndServe(); err != nil {
-		}
+		_ = srv.ListenAndServe()
 	}()
+	<-ready
 	return srv
 }
 
 func TestLookupDirectAnswer(t *testing.T) {
-	mux := dns.NewServeMux()
-	mux.HandleFunc(".", func(w dns.ResponseWriter, req *dns.Msg) {
-		resp := new(dns.Msg)
-		resp.SetReply(req)
-		resp.Authoritative = true
-		if req.Question[0].Name == "example.com." &&
-			req.Question[0].Qtype == dns.TypeA {
+	srv := startTestServer(t, "127.0.0.1:15353",
+		dns.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
+			resp := new(dns.Msg)
+			dnsutil.SetReply(resp, req)
+			resp.Authoritative = true
 			resp.Answer = append(resp.Answer, &dns.A{
-				Hdr: dns.RR_Header{
-					Name:   "example.com.",
-					Rrtype: dns.TypeA,
-					Class:  dns.ClassINET,
-					Ttl:    60,
+				Hdr: dns.Header{
+					Name:  "example.com.",
+					Class: dns.ClassINET,
+					TTL:   60,
 				},
-				A: net.IPv4(127, 0, 0, 1),
+				A: rdata.A{Addr: netip.AddrFrom4([4]byte{127, 0, 0, 1})},
 			})
-		} else {
-			resp.Rcode = dns.RcodeNameError
-		}
-		w.WriteMsg(resp)
-	})
-
-	srv := startTestServer(t, "127.0.0.1:15353", mux)
-	defer srv.Shutdown()
-	time.Sleep(50 * time.Millisecond)
+			io.Copy(w, resp)
+		}))
+	defer srv.Shutdown(context.Background())
+	time.Sleep(100 * time.Millisecond)
 
 	r := New(
 		WithRootAddresses([]string{"127.0.0.1:15353"}),
@@ -73,23 +72,21 @@ func TestLookupDirectAnswer(t *testing.T) {
 	if !ok {
 		t.Fatal("expected A record")
 	}
-	if !a.A.Equal(net.IPv4(127, 0, 0, 1)) {
-		t.Fatalf("expected 127.0.0.1, got %s", a.A)
+	if a.A.Addr != netip.AddrFrom4([4]byte{127, 0, 0, 1}) {
+		t.Fatalf("expected 127.0.0.1, got %s", a.A.Addr)
 	}
 }
 
 func TestLookupNXDOMAIN(t *testing.T) {
-	mux := dns.NewServeMux()
-	mux.HandleFunc(".", func(w dns.ResponseWriter, req *dns.Msg) {
-		resp := new(dns.Msg)
-		resp.SetReply(req)
-		resp.Rcode = dns.RcodeNameError
-		w.WriteMsg(resp)
-	})
-
-	srv := startTestServer(t, "127.0.0.1:15354", mux)
-	defer srv.Shutdown()
-	time.Sleep(50 * time.Millisecond)
+	srv := startTestServer(t, "127.0.0.1:15354",
+		dns.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
+			resp := new(dns.Msg)
+			dnsutil.SetReply(resp, req)
+			resp.Rcode = dns.RcodeNameError
+			io.Copy(w, resp)
+		}))
+	defer srv.Shutdown(context.Background())
+	time.Sleep(100 * time.Millisecond)
 
 	r := New(
 		WithRootAddresses([]string{"127.0.0.1:15354"}),
@@ -111,16 +108,16 @@ func TestNextServersWithGlue(t *testing.T) {
 	msg := new(dns.Msg)
 	msg.Authoritative = false
 	msg.Ns = append(msg.Ns, &dns.NS{
-		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeNS, Ttl: 300},
-		Ns:  "ns1.example.com.",
+		Hdr: dns.Header{Name: "example.com.", Class: dns.ClassINET, TTL: 300},
+		NS:  rdata.NS{Ns: "ns1.example.com."},
 	})
 	msg.Extra = append(msg.Extra, &dns.A{
-		Hdr: dns.RR_Header{Name: "ns1.example.com.", Rrtype: dns.TypeA, Ttl: 300},
-		A:   net.ParseIP("192.0.2.1").To4(),
+		Hdr: dns.Header{Name: "ns1.example.com.", Class: dns.ClassINET, TTL: 300},
+		A:   rdata.A{Addr: netip.MustParseAddr("192.0.2.1")},
 	})
 	msg.Extra = append(msg.Extra, &dns.AAAA{
-		Hdr:  dns.RR_Header{Name: "ns1.example.com.", Rrtype: dns.TypeAAAA, Ttl: 300},
-		AAAA: net.ParseIP("2001:db8::1"),
+		Hdr:   dns.Header{Name: "ns1.example.com.", Class: dns.ClassINET, TTL: 300},
+		AAAA: rdata.AAAA{Addr: netip.MustParseAddr("2001:db8::1")},
 	})
 
 	r := &Resolver{}
@@ -137,8 +134,8 @@ func TestNextServersNoGlue(t *testing.T) {
 	msg := new(dns.Msg)
 	msg.Authoritative = false
 	msg.Ns = append(msg.Ns, &dns.NS{
-		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeNS, Ttl: 300},
-		Ns:  "ns1.example.com.",
+		Hdr: dns.Header{Name: "example.com.", Class: dns.ClassINET, TTL: 300},
+		NS:  rdata.NS{Ns: "ns1.example.com."},
 	})
 
 	r := &Resolver{maxDelegations: 30, timeout: time.Second, retries: 1}
@@ -147,3 +144,62 @@ func TestNextServersNoGlue(t *testing.T) {
 		t.Fatal("expected error when no glue and no roots")
 	}
 }
+
+func TestLookupCNAME(t *testing.T) {
+	callCount := 0
+	srv := startTestServer(t, "127.0.0.1:15355",
+		dns.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
+			callCount++
+			resp := new(dns.Msg)
+			dnsutil.SetReply(resp, req)
+			resp.Authoritative = true
+			resp.Answer = append(resp.Answer,
+				&dns.CNAME{
+					Hdr:   dns.Header{Name: "alias.example.com.", Class: dns.ClassINET, TTL: 60},
+					CNAME: rdata.CNAME{Target: "real.example.com."},
+				},
+			)
+			if callCount > 1 {
+				resp.Answer = append(resp.Answer,
+					&dns.A{
+						Hdr: dns.Header{Name: "real.example.com.", Class: dns.ClassINET, TTL: 60},
+						A:   rdata.A{Addr: netip.AddrFrom4([4]byte{127, 0, 0, 1})},
+					},
+				)
+			}
+			io.Copy(w, resp)
+		}))
+	defer srv.Shutdown(context.Background())
+	time.Sleep(100 * time.Millisecond)
+
+	r := New(
+		WithRootAddresses([]string{"127.0.0.1:15355"}),
+		WithTimeout(time.Second),
+	)
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+
+	resp, err := r.Lookup(ctx, "alias.example.com.", dns.TypeA)
+	if err != nil {
+		t.Fatalf("CNAME lookup failed: %v", err)
+	}
+	if resp.Rcode != dns.RcodeSuccess {
+		t.Fatalf("expected NOERROR, got %d", resp.Rcode)
+	}
+	hasCNAME := false
+	hasA := false
+	for _, rr := range resp.Answer {
+		if _, ok := rr.(*dns.CNAME); ok {
+			hasCNAME = true
+		}
+		if _, ok := rr.(*dns.A); ok {
+			hasA = true
+		}
+	}
+	if !hasCNAME {
+		t.Fatal("expected CNAME record in answer")
+	}
+	if !hasA {
+		t.Fatal("expected A record (CNAME target) in answer")
+	}
+}
diff --git a/internal/resolver/root.go b/internal/resolver/root.go
index 9dac31c..0557dac 100644
--- a/internal/resolver/root.go
+++ b/internal/resolver/root.go
@@ -2,7 +2,7 @@ package resolver
 
 import (
 	_ "embed"
-	"github.com/miekg/dns"
+	"codeberg.org/miekg/dns"
 	"strings"
 )
 
@@ -23,7 +23,7 @@ func loadRootServers() []string {
 		if !ok {
 			continue
 		}
-		ip := a.A.String()
+		ip := a.A.Addr.String()
 		if !seen[ip] {
 			seen[ip] = true
 			addrs = append(addrs, ip)