dnswatcher/internal/resolver/iterative.go

package resolver

import (
	"context"
	"errors"
	"fmt"
	"math/rand/v2"
	"net"
	"slices"
	"sort"
	"strings"
	"time"

	"github.com/miekg/dns"
)

const (
	queryTimeoutDuration = 2 * time.Second
	maxRetries           = 2
	maxDelegation        = 20
	timeoutMultiplier    = 2
)

// ErrRefused is returned when a DNS server refuses a query.
var ErrRefused = errors.New("dns query refused")

func rootServerList() []string {
	return []string{
		"198.41.0.4",     // a.root-servers.net
		"170.247.170.2",  // b
		"192.33.4.12",    // c
		"199.7.91.13",    // d
		"192.203.230.10", // e
		"192.5.5.241",    // f
		"192.112.36.4",   // g
		"198.97.190.53",  // h
		"192.36.148.17",  // i
		"192.58.128.30",  // j
		"193.0.14.129",   // k
		"199.7.83.42",    // l
		"202.12.27.33",   // m
	}
}

func checkCtx(ctx context.Context) error {
	err := ctx.Err()
	if err != nil {
		return ErrContextCanceled
	}

	return nil
}

func (r *Resolver) exchangeWithTimeout(
	ctx context.Context,
	msg *dns.Msg,
	addr string,
	attempt int,
) (*dns.Msg, error) {
	_ = attempt // timeout escalation handled by client config

	resp, _, err := r.client.ExchangeContext(ctx, msg, addr)

	return resp, err
}

func (r *Resolver) tryExchange(
	ctx context.Context,
	msg *dns.Msg,
	addr string,
) (*dns.Msg, error) {
	var resp *dns.Msg

	var err error

	for attempt := range maxRetries {
		if checkCtx(ctx) != nil {
			return nil, ErrContextCanceled
		}

		resp, err = r.exchangeWithTimeout(
			ctx, msg, addr, attempt,
		)
		if err == nil {
			break
		}
	}

	return resp, err
}

func (r *Resolver) retryTCP(
	ctx context.Context,
	msg *dns.Msg,
	addr string,
	resp *dns.Msg,
) *dns.Msg {
	if !resp.Truncated {
		return resp
	}

	tcpResp, _, tcpErr := r.tcp.ExchangeContext(ctx, msg, addr)
	if tcpErr == nil {
		return tcpResp
	}

	return resp
}

// queryDNS sends a DNS query to a specific server IP.
// Tries non-recursive first, falls back to recursive on
// REFUSED (handles DNS interception environments).
func (r *Resolver) queryDNS(
	ctx context.Context,
	serverIP string,
	name string,
	qtype uint16,
) (*dns.Msg, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	name = dns.Fqdn(name)
	addr := net.JoinHostPort(serverIP, "53")

	msg := new(dns.Msg)
	msg.SetQuestion(name, qtype)
	msg.RecursionDesired = false

	resp, err := r.tryExchange(ctx, msg, addr)
	if err != nil {
		return nil, fmt.Errorf("query %s @%s: %w", name, serverIP, err)
	}

	if resp.Rcode == dns.RcodeRefused {
		msg.RecursionDesired = true

		resp, err = r.tryExchange(ctx, msg, addr)
		if err != nil {
			return nil, fmt.Errorf(
				"query %s @%s: %w", name, serverIP, err,
			)
		}

		if resp.Rcode == dns.RcodeRefused {
			return nil, fmt.Errorf(
				"query %s @%s: %w", name, serverIP, ErrRefused,
			)
		}
	}

	resp = r.retryTCP(ctx, msg, addr, resp)

	return resp, nil
}

func extractNSSet(rrs []dns.RR) []string {
	nsSet := make(map[string]bool)

	for _, rr := range rrs {
		if ns, ok := rr.(*dns.NS); ok {
			nsSet[strings.ToLower(ns.Ns)] = true
		}
	}

	names := make([]string, 0, len(nsSet))
	for n := range nsSet {
		names = append(names, n)
	}

	sort.Strings(names)

	return names
}

func extractGlue(rrs []dns.RR) map[string][]net.IP {
	glue := make(map[string][]net.IP)

	for _, rr := range rrs {
		switch r := rr.(type) {
		case *dns.A:
			name := strings.ToLower(r.Hdr.Name)
			glue[name] = append(glue[name], r.A)
		case *dns.AAAA:
			name := strings.ToLower(r.Hdr.Name)
			glue[name] = append(glue[name], r.AAAA)
		}
	}

	return glue
}

func glueIPs(nsNames []string, glue map[string][]net.IP) []string {
	var ips []string

	for _, ns := range nsNames {
		for _, addr := range glue[ns] {
			if v4 := addr.To4(); v4 != nil {
				ips = append(ips, v4.String())
			}
		}
	}

	return ips
}

func (r *Resolver) followDelegation(
	ctx context.Context,
	domain string,
	servers []string,
) ([]string, error) {
	for range maxDelegation {
		if checkCtx(ctx) != nil {
			return nil, ErrContextCanceled
		}

		resp, err := r.queryServers(
			ctx, servers, domain, dns.TypeNS,
		)
		if err != nil {
			return nil, err
		}

		ansNS := extractNSSet(resp.Answer)
		if len(ansNS) > 0 {
			return ansNS, nil
		}

		// An authoritative reply comes from the servers of the zone
		// domain is in; it is not a referral, even when its authority
		// section lists that zone's NS records. Without NS records in
		// the answer, domain is not the zone's apex and has no
		// nameservers of its own.
		if resp.Authoritative {
			return nil, ErrNoNameservers
		}

		authNS := extractNSSet(resp.Ns)
		if len(authNS) == 0 {
			return r.resolveNSIterative(ctx, domain)
		}

		glue := extractGlue(resp.Extra)
		nextServers := glueIPs(authNS, glue)

		if len(nextServers) == 0 {
			nextServers = r.resolveNSIPs(ctx, authNS)
		}

		if len(nextServers) == 0 {
			return nil, ErrNoNameservers
		}

		servers = nextServers
	}

	return nil, ErrNoNameservers
}

// shuffled returns a copy of servers in the order shuffle puts them
// in. The resolver passes rand.Shuffle, so each time it walks a list of
// servers it starts at a random one, and no one server gets every
// first query.
func shuffled(
	servers []string,
	shuffle func(n int, swap func(i, j int)),
) []string {
	order := slices.Clone(servers)
	shuffle(len(order), func(i, j int) {
		order[i], order[j] = order[j], order[i]
	})

	return order
}

// queryServers asks the servers in a random order and returns the
// first reply; it fails only when every server has failed.
func (r *Resolver) queryServers(
	ctx context.Context,
	servers []string,
	name string,
	qtype uint16,
) (*dns.Msg, error) {
	var lastErr error

	for _, ip := range shuffled(servers, rand.Shuffle) {
		if checkCtx(ctx) != nil {
			return nil, ErrContextCanceled
		}

		resp, err := r.queryDNS(ctx, ip, name, qtype)
		if err == nil {
			return resp, nil
		}

		lastErr = err
	}

	return nil, fmt.Errorf("all servers failed: %w", lastErr)
}

// resolveNSIPs returns the addresses of one of the nameservers, trying
// their names in a random order until one resolves.
func (r *Resolver) resolveNSIPs(
	ctx context.Context,
	nsNames []string,
) []string {
	var ips []string

	for _, ns := range shuffled(nsNames, rand.Shuffle) {
		resolved, err := r.resolveARecord(ctx, ns)
		if err == nil {
			ips = append(ips, resolved...)
		}

		if len(ips) > 0 {
			break
		}
	}

	return ips
}

// resolveNSIterative queries for NS records using iterative
// resolution as a fallback when followDelegation finds no
// authoritative answer in the delegation chain.
func (r *Resolver) resolveNSIterative(
	ctx context.Context,
	domain string,
) ([]string, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	domain = dns.Fqdn(domain)
	servers := rootServerList()

	for range maxDelegation {
		if checkCtx(ctx) != nil {
			return nil, ErrContextCanceled
		}

		resp, err := r.queryServers(
			ctx, servers, domain, dns.TypeNS,
		)
		if err != nil {
			return nil, err
		}

		nsNames := extractNSSet(resp.Answer)
		if len(nsNames) > 0 {
			return nsNames, nil
		}

		// Follow delegation.
		authNS := extractNSSet(resp.Ns)
		if len(authNS) == 0 {
			break
		}

		glue := extractGlue(resp.Extra)
		nextServers := glueIPs(authNS, glue)

		if len(nextServers) == 0 {
			break
		}

		servers = nextServers
	}

	return nil, ErrNoNameservers
}

// resolveARecord resolves a hostname to IPv4 addresses using
// iterative resolution through the delegation chain.
func (r *Resolver) resolveARecord(
	ctx context.Context,
	hostname string,
) ([]string, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	hostname = dns.Fqdn(hostname)
	servers := rootServerList()

	for range maxDelegation {
		if checkCtx(ctx) != nil {
			return nil, ErrContextCanceled
		}

		resp, err := r.queryServers(
			ctx, servers, hostname, dns.TypeA,
		)
		if err != nil {
			return nil, fmt.Errorf(
				"resolving %s: %w", hostname, err,
			)
		}

		// Check for A records in the answer section.
		var ips []string

		for _, rr := range resp.Answer {
			if a, ok := rr.(*dns.A); ok {
				ips = append(ips, a.A.String())
			}
		}

		if len(ips) > 0 {
			return ips, nil
		}

		// Follow delegation if present.
		authNS := extractNSSet(resp.Ns)
		if len(authNS) == 0 {
			break
		}

		glue := extractGlue(resp.Extra)
		nextServers := glueIPs(authNS, glue)

		if len(nextServers) == 0 {
			// Resolve NS IPs iteratively — but guard
			// against infinite recursion by using only
			// already-resolved servers.
			break
		}

		servers = nextServers
	}

	return nil, fmt.Errorf(
		"cannot resolve %s: %w", hostname, ErrNoNameservers,
	)
}

// FindAuthoritativeNameservers traces the delegation chain from
// root servers to discover all authoritative nameservers for the
// given domain. For a name that is not a zone apex it tries each
// parent name in turn, so it returns the nameservers of the zone the
// name is in.
func (r *Resolver) FindAuthoritativeNameservers(
	ctx context.Context,
	domain string,
) ([]string, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	domain = dns.Fqdn(strings.ToLower(domain))
	labels := dns.SplitDomainName(domain)

	for i := range labels {
		if checkCtx(ctx) != nil {
			return nil, ErrContextCanceled
		}

		candidate := strings.Join(labels[i:], ".") + "."

		nsNames, err := r.followDelegation(
			ctx, candidate, rootServerList(),
		)
		if err == nil && len(nsNames) > 0 {
			sort.Strings(nsNames)

			return nsNames, nil
		}
	}

	return nil, ErrNoNameservers
}

// QueryNameserver queries a specific nameserver for all record
// types and builds a NameserverResponse.
func (r *Resolver) QueryNameserver(
	ctx context.Context,
	nsHostname string,
	hostname string,
) (*NameserverResponse, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	nsIPs, err := r.resolveARecord(ctx, nsHostname)
	if err != nil {
		return nil, fmt.Errorf("resolving NS %s: %w", nsHostname, err)
	}

	hostname = dns.Fqdn(hostname)

	return r.queryAllTypes(ctx, nsHostname, nsIPs[0], hostname)
}

// QueryNameserverIP queries a nameserver by its IP address directly,
// bypassing NS hostname resolution.
func (r *Resolver) QueryNameserverIP(
	ctx context.Context,
	nsHostname string,
	nsIP string,
	hostname string,
) (*NameserverResponse, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	hostname = dns.Fqdn(hostname)

	return r.queryAllTypes(ctx, nsHostname, nsIP, hostname)
}

func (r *Resolver) queryAllTypes(
	ctx context.Context,
	nsHostname string,
	nsIP string,
	hostname string,
) (*NameserverResponse, error) {
	resp := &NameserverResponse{
		Nameserver: nsHostname,
		Records:    make(map[string][]string),
		Status:     StatusOK,
	}

	qtypes := []uint16{
		dns.TypeA, dns.TypeAAAA, dns.TypeCNAME,
		dns.TypeMX, dns.TypeTXT, dns.TypeSRV,
		dns.TypeCAA, dns.TypeNS,
	}

	state := r.queryEachType(ctx, nsIP, hostname, qtypes, resp)
	classifyResponse(resp, state)

	return resp, nil
}

type queryState struct {
	gotNXDomain bool
	gotSERVFAIL bool
	gotRefused  bool
	gotTimeout  bool
	gotReferral bool
	netErr      error
	hasRecords  bool
}

func (r *Resolver) queryEachType(
	ctx context.Context,
	nsIP string,
	hostname string,
	qtypes []uint16,
	resp *NameserverResponse,
) queryState {
	var state queryState

	for _, qtype := range qtypes {
		if checkCtx(ctx) != nil {
			break
		}

		r.querySingleType(ctx, nsIP, hostname, qtype, resp, &state)
	}

	for k := range resp.Records {
		sort.Strings(resp.Records[k])
	}

	return state
}

func (r *Resolver) querySingleType(
	ctx context.Context,
	nsIP string,
	hostname string,
	qtype uint16,
	resp *NameserverResponse,
	state *queryState,
) {
	msg, err := r.queryDNS(ctx, nsIP, hostname, qtype)
	if err != nil {
		switch {
		case isTimeout(err):
			state.gotTimeout = true
		case errors.Is(err, ErrRefused):
			state.gotRefused = true
		default:
			state.netErr = err
		}

		return
	}

	if msg.Rcode == dns.RcodeNameError {
		state.gotNXDomain = true

		return
	}

	if msg.Rcode == dns.RcodeServerFailure {
		state.gotSERVFAIL = true

		return
	}

	// A reply with no answer that lists other nameservers, from a server
	// that does not hold the name's zone, is a referral and says nothing
	// about the name's records. A parent zone's servers send one when
	// every server of the name's own zone failed and
	// FindAuthoritativeNameservers moved on to the parent name.
	if !msg.Authoritative && len(msg.Answer) == 0 &&
		len(extractNSSet(msg.Ns)) > 0 {
		state.gotReferral = true

		return
	}

	collectAnswerRecords(msg, resp, state)
}

func collectAnswerRecords(
	msg *dns.Msg,
	resp *NameserverResponse,
	state *queryState,
) {
	for _, rr := range msg.Answer {
		val := extractRecordValue(rr)
		if val == "" {
			continue
		}

		typeName := dns.TypeToString[rr.Header().Rrtype]
		resp.Records[typeName] = append(
			resp.Records[typeName], val,
		)
		state.hasRecords = true
	}
}

// isTimeout checks whether an error is a network timeout.
func isTimeout(err error) bool {
	var netErr net.Error
	if errors.As(err, &netErr) {
		return netErr.Timeout()
	}

	return false
}

func classifyResponse(resp *NameserverResponse, state queryState) {
	switch {
	case state.gotNXDomain && !state.hasRecords:
		resp.Status = StatusNXDomain
	case state.gotTimeout && !state.hasRecords:
		resp.Status = StatusTimeout
		resp.Error = "all queries timed out"
	case state.gotSERVFAIL && !state.hasRecords:
		resp.Status = StatusError
		resp.Error = "server returned SERVFAIL"
	case state.gotRefused && !state.hasRecords:
		resp.Status = StatusError
		resp.Error = "server returned REFUSED"
	case state.netErr != nil && !state.hasRecords:
		resp.Status = StatusError
		resp.Error = "network error: " + state.netErr.Error()
	case state.gotReferral && !state.hasRecords:
		resp.Status = StatusError
		resp.Error = "server returned a referral"
	case !state.hasRecords && !state.gotNXDomain:
		resp.Status = StatusNoData
	}
}

// extractRecordValue formats a DNS RR value as a string. DNS names
// are case-insensitive and nameservers may answer in any letter case,
// so names are lower-cased to compare equal. TXT and CAA values keep
// their letter case.
func extractRecordValue(rr dns.RR) string {
	switch r := rr.(type) {
	case *dns.A:
		return r.A.String()
	case *dns.AAAA:
		return r.AAAA.String()
	case *dns.CNAME:
		return strings.ToLower(r.Target)
	case *dns.MX:
		return fmt.Sprintf("%d %s", r.Preference, strings.ToLower(r.Mx))
	case *dns.TXT:
		return strings.Join(r.Txt, "")
	case *dns.SRV:
		return fmt.Sprintf(
			"%d %d %d %s",
			r.Priority, r.Weight, r.Port, strings.ToLower(r.Target),
		)
	case *dns.CAA:
		return fmt.Sprintf(
			"%d %s \"%s\"", r.Flag, r.Tag, r.Value,
		)
	case *dns.NS:
		return strings.ToLower(r.Ns)
	default:
		return ""
	}
}

// QueryAllNameservers discovers the auth NSes of the zone the
// hostname is in, then queries each one independently.
func (r *Resolver) QueryAllNameservers(
	ctx context.Context,
	hostname string,
) (map[string]*NameserverResponse, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	nameservers, err := r.FindAuthoritativeNameservers(ctx, hostname)
	if err != nil {
		return nil, err
	}

	return r.queryEachNS(ctx, nameservers, hostname)
}

func (r *Resolver) queryEachNS(
	ctx context.Context,
	nameservers []string,
	hostname string,
) (map[string]*NameserverResponse, error) {
	results := make(map[string]*NameserverResponse)

	for _, ns := range nameservers {
		resp, err := r.QueryNameserver(ctx, ns, hostname)

		// A query the context cut short says nothing about the
		// nameserver, so it must not be returned as its failure.
		if checkCtx(ctx) != nil {
			return nil, ErrContextCanceled
		}

		if err != nil {
			results[ns] = &NameserverResponse{
				Nameserver: ns,
				Records:    make(map[string][]string),
				Status:     StatusError,
				Error:      err.Error(),
			}

			continue
		}

		results[ns] = resp
	}

	return results, nil
}

// LookupNS returns the NS record set for a domain.
func (r *Resolver) LookupNS(
	ctx context.Context,
	domain string,
) ([]string, error) {
	return r.FindAuthoritativeNameservers(ctx, domain)
}

// LookupAllRecords performs iterative resolution to find all DNS
// records for the given hostname, keyed by authoritative nameserver.
// Each nameserver's response carries its status and error with its
// records.
func (r *Resolver) LookupAllRecords(
	ctx context.Context,
	hostname string,
) (map[string]*NameserverResponse, error) {
	return r.QueryAllNameservers(ctx, hostname)
}

// ResolveIPAddresses resolves a hostname to all IPv4 and IPv6
// addresses, following CNAME chains up to MaxCNAMEDepth. When no
// nameserver of the name's zone answered, it returns an error rather
// than no addresses.
func (r *Resolver) ResolveIPAddresses(
	ctx context.Context,
	hostname string,
) ([]string, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

	return r.resolveIPWithCNAME(ctx, hostname, 0)
}

func (r *Resolver) resolveIPWithCNAME(
	ctx context.Context,
	hostname string,
	depth int,
) ([]string, error) {
	if depth > MaxCNAMEDepth {
		return nil, ErrCNAMEDepthExceeded
	}

	results, err := r.QueryAllNameservers(ctx, hostname)
	if err != nil {
		return nil, err
	}

	ips, cnameTarget, err := collectIPs(results)
	if err != nil {
		return nil, fmt.Errorf("resolving %s: %w", hostname, err)
	}

	if len(ips) == 0 && cnameTarget != "" {
		return r.resolveIPWithCNAME(ctx, cnameTarget, depth+1)
	}

	sort.Strings(ips)

	return ips, nil
}

// collectIPs returns the addresses in the nameservers' answers and the
// first CNAME target among them. It returns ErrNoNameserverAnswered when
// every nameserver timed out, failed or returned a referral: that is not
// a name with no addresses.
func collectIPs(
	results map[string]*NameserverResponse,
) ([]string, string, error) {
	seen := make(map[string]bool)

	var ips []string

	var cnameTarget string

	answered := false

	for _, resp := range results {
		if resp.Status == StatusTimeout || resp.Status == StatusError {
			continue
		}

		answered = true

		if resp.Status == StatusNXDomain {
			continue
		}

		for _, ip := range resp.Records["A"] {
			if !seen[ip] {
				seen[ip] = true
				ips = append(ips, ip)
			}
		}

		for _, ip := range resp.Records["AAAA"] {
			if !seen[ip] {
				seen[ip] = true
				ips = append(ips, ip)
			}
		}

		if len(resp.Records["CNAME"]) > 0 && cnameTarget == "" {
			cnameTarget = resp.Records["CNAME"][0]
		}
	}

	if !answered {
		return nil, "", ErrNoNameserverAnswered
	}

	return ips, cnameTarget, nil
}