dnswatcher/internal/resolver/iterative.go

package resolver

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

	"github.com/miekg/dns"
)

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

	// maxLookupDepth is how many lookups of nameserver addresses may be
	// under way one inside another. Looking up a nameserver's address
	// can meet a referral that names nameservers without their
	// addresses, which are then looked up in turn; without a limit,
	// delegations that point at each other would never end. Each level
	// multiplies the queries sent. pool.ntp.org needs three: the
	// address of its nameserver g.ntpns.org can need a.ntpns.org's,
	// which needs a bitnames.com nameserver's.
	maxLookupDepth = 3
)

// 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
}

// retryTCP returns the reply to msg over TCP when resp, its reply over
// UDP, is truncated. When that fails it returns resp, still truncated,
// which holds only the records that fit.
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, never asking it
// for recursion. A reply of REFUSED is returned as ErrRefused.
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 {
		return nil, fmt.Errorf(
			"query %s @%s: %w", name, serverIP, ErrRefused,
		)
	}

	return r.retryTCP(ctx, msg, addr, 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) {
	// servers are the root servers, the servers of zone ".".
	zone := "."

	var withoutAddresses []string

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

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

		nsSet := nsSetFrom(resp, domain)
		if len(nsSet) > 0 {
			return nsSet, 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)
		}

		servers, withoutAddresses = referralNameservers(resp)
		zone = referralZone(resp)
	}

	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 servers, the servers of zone, about name in a random
// order until one gives a usable reply. A server that times out, refuses
// or gives a reply that is not usable is passed over for the next. When
// every server refused, the error says so, and when they are the root
// servers it is ErrIntercepted.
func (r *Resolver) queryServers(
	ctx context.Context,
	servers []string,
	zone string,
	name string,
	qtype uint16,
) (*dns.Msg, error) {
	var lastErr error

	refused := 0

	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 && !usableReply(resp, zone, name) {
			err = fmt.Errorf(
				"query %s @%s: %w", name, ip, ErrUnusableReply,
			)
		}

		if err == nil {
			return resp, nil
		}

		if errors.Is(err, ErrRefused) {
			refused++
		}

		lastErr = err
	}

	if refused == len(servers) && zone == "." {
		return nil, fmt.Errorf(
			"every root server refused a query for %s: %w",
			name, ErrIntercepted,
		)
	}

	if refused == len(servers) {
		return nil, fmt.Errorf(
			"every server of %s refused a query for %s: %w",
			zone, name, ErrRefused,
		)
	}

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

// isErrorReply reports whether msg is an error reply: one with any code
// but NOERROR and NXDOMAIN, such as SERVFAIL, NOTIMP or FORMERR. An error
// reply says nothing about the name's records.
func isErrorReply(msg *dns.Msg) bool {
	return msg.Rcode != dns.RcodeSuccess && msg.Rcode != dns.RcodeNameError
}

// usableReply reports whether resp, a reply from one of the servers of
// zone to a query about name, is usable. An error reply such as SERVFAIL
// is not. Nor is a referral, unless it refers the query to a zone below
// zone that name is in: a server that refers it back to zone, up or
// sideways does not serve zone as it should.
func usableReply(resp *dns.Msg, zone string, name string) bool {
	if isErrorReply(resp) {
		return false
	}

	child := referralZone(resp)
	if resp.Authoritative || len(resp.Answer) > 0 || child == "" {
		return true
	}

	return child != zone && dns.IsSubDomain(zone, child) &&
		dns.IsSubDomain(child, name)
}

// referralZone returns the zone a referral refers the query to: the
// owner name of the NS records in resp's authority section, or "" when
// there are none.
func referralZone(resp *dns.Msg) string {
	for _, rr := range resp.Ns {
		if ns, ok := rr.(*dns.NS); ok {
			return strings.ToLower(ns.Hdr.Name)
		}
	}

	return ""
}

// nsSetFrom returns the NS set of domain that resp, a reply to a query
// for domain's NS records, gives: the delegation in a referral to domain
// itself, or else the NS records in the answer; empty when it gives
// neither. A referral to domain comes from its parent zone's servers,
// which all hold the same delegation, so the set does not depend on
// which of them answered. domain's own servers, which can disagree about
// their NS records, are then not asked.
func nsSetFrom(resp *dns.Msg, domain string) []string {
	if referralZone(resp) == domain {
		return extractNSSet(resp.Ns)
	}

	return extractNSSet(resp.Answer)
}

// referralNameservers returns the IPv4 addresses that resp, a referral,
// gives for the nameservers it names, and the names of the nameservers
// it gives no address for.
func referralNameservers(resp *dns.Msg) ([]string, []string) {
	glue := extractGlue(resp.Extra)

	var given, withoutAddresses []string

	for _, ns := range extractNSSet(resp.Ns) {
		ips := glueIPs([]string{ns}, glue)
		if len(ips) == 0 {
			withoutAddresses = append(withoutAddresses, ns)
		}

		given = append(given, ips...)
	}

	return given, withoutAddresses
}

// queryZone asks the servers of zone about name as queryServers does:
// first those at given, the addresses a referral gave, and only when
// none of them gives a usable reply, the nameservers named
// withoutAddresses, once their addresses are looked up. depth is how
// many lookups of a nameserver's address are under way, 0 in the walk
// to a domain's nameservers; at maxLookupDepth, no address is looked
// up. When the limit is why none was found, here or in a lookup this
// one started, the error is ErrLookupDepthExceeded.
func (r *Resolver) queryZone(
	ctx context.Context,
	given []string,
	withoutAddresses []string,
	zone string,
	name string,
	qtype uint16,
	depth int,
) (*dns.Msg, error) {
	err := fmt.Errorf(
		"no address for any nameserver of %s: %w", zone, ErrNoNameservers,
	)

	if len(given) > 0 {
		var resp *dns.Msg

		resp, err = r.queryServers(ctx, given, zone, name, qtype)
		if err == nil {
			return resp, nil
		}
	}

	if len(withoutAddresses) == 0 {
		return nil, err
	}

	if depth >= maxLookupDepth {
		return nil, fmt.Errorf(
			"addresses of the nameservers of %s not looked up: %w",
			zone, ErrLookupDepthExceeded,
		)
	}

	lookedUp, limitErr := r.resolveNSIPs(ctx, withoutAddresses, depth+1)
	if limitErr != nil {
		return nil, limitErr
	}

	if len(lookedUp) == 0 {
		return nil, err
	}

	return r.queryServers(ctx, lookedUp, zone, name, qtype)
}

// resolveNSIPs returns the addresses of every nameserver in nsNames
// whose name resolves, each looked up at depth (see resolveARecord).
// The walk can then go on to the zone's other nameservers when one
// gives no usable reply. When none resolves and the depth limit
// stopped one of the lookups, it returns that lookup's error.
func (r *Resolver) resolveNSIPs(
	ctx context.Context,
	nsNames []string,
	depth int,
) ([]string, error) {
	var (
		ips      []string
		limitErr error
	)

	for _, ns := range nsNames {
		resolved, err := r.resolveARecord(ctx, ns, depth)

		switch {
		case err == nil:
			ips = append(ips, resolved...)
		case errors.Is(err, ErrLookupDepthExceeded):
			limitErr = err
		}
	}

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

	return nil, limitErr
}

// 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()
	zone := "."

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

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

		nsNames := nsSetFrom(resp, domain)
		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
		zone = referralZone(resp)
	}

	return nil, ErrNoNameservers
}

// resolveARecord resolves a hostname, a nameserver's name, to IPv4
// addresses using iterative resolution through the delegation chain.
// depth is how many lookups of a nameserver's address are under way,
// this one included: 1 for a lookup that no other lookup started.
func (r *Resolver) resolveARecord(
	ctx context.Context,
	hostname string,
	depth int,
) ([]string, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

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

	var withoutAddresses []string

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

		resp, err := r.queryZone(
			ctx, servers, withoutAddresses, zone, hostname, dns.TypeA,
			depth,
		)
		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
		}

		servers, withoutAddresses = referralNameservers(resp)
		zone = referralZone(resp)
	}

	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, as the delegation from its parent zone's servers lists
// them. 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
		}

		// The root servers would refuse every parent name too.
		if errors.Is(err, ErrIntercepted) {
			return nil, err
		}
	}

	return nil, ErrNoNameservers
}

// recordTypes returns the record types a nameserver is asked for when a
// name is checked.
func recordTypes() []uint16 {
	return []uint16{
		dns.TypeA, dns.TypeAAAA, dns.TypeCNAME,
		dns.TypeMX, dns.TypeTXT, dns.TypeSRV,
		dns.TypeCAA, dns.TypeNS,
	}
}

// addressTypes returns the record types ResolveIPAddresses asks for,
// the only ones it reads.
func addressTypes() []uint16 {
	return []uint16{dns.TypeA, dns.TypeAAAA, dns.TypeCNAME}
}

// 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) {
	return r.queryNameserver(ctx, nsHostname, hostname, recordTypes())
}

// queryNameserver queries a specific nameserver for the record types
// in qtypes and builds a NameserverResponse.
func (r *Resolver) queryNameserver(
	ctx context.Context,
	nsHostname string,
	hostname string,
	qtypes []uint16,
) (*NameserverResponse, error) {
	if checkCtx(ctx) != nil {
		return nil, ErrContextCanceled
	}

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

	hostname = dns.Fqdn(hostname)

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

// 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.queryTypes(ctx, nsHostname, nsIP, hostname, recordTypes())
}

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

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

	return resp, nil
}

type queryState struct {
	gotNXDomain   bool
	gotErrorReply bool
	errorReply    string // its code, such as SERVFAIL, or number if unnamed
	gotRefused    bool
	gotTimeout    bool
	gotReferral   bool
	netErr        error
	hasRecords    bool
	answered      bool
}

// queryEachType asks the nameserver at nsIP about hostname once for each
// record type in qtypes, and lists in resp.FailedTypes the types whose
// query got no usable reply, logging each with the reason unless ctx was
// cancelled: shutdown cancels it, and a query it cut short did not fail.
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
		}

		err := r.querySingleType(ctx, nsIP, hostname, qtype, resp, &state)
		if err == nil {
			state.answered = true

			continue
		}

		rtype := dns.TypeToString[qtype]
		resp.FailedTypes = append(resp.FailedTypes, rtype)

		if errors.Is(ctx.Err(), context.Canceled) {
			continue
		}

		r.log.Warn(
			"record type query failed",
			"hostname", hostname,
			"nameserver", resp.Nameserver,
			"type", rtype,
			"error", err,
		)
	}

	// The reply about another type can carry the name's CNAME. When the
	// query for CNAME itself failed, that is left out too, so Records
	// holds nothing for a failed type.
	for _, rtype := range resp.FailedTypes {
		delete(resp.Records, rtype)
	}

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

	return state
}

// querySingleType asks the nameserver at nsIP about hostname's records
// of type qtype. It returns nil when the nameserver answered: with
// records, with none, or with NXDOMAIN; otherwise it returns why not.
func (r *Resolver) querySingleType(
	ctx context.Context,
	nsIP string,
	hostname string,
	qtype uint16,
	resp *NameserverResponse,
	state *queryState,
) error {
	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 err
	}

	return readReply(msg, resp, state)
}

// readReply adds to resp the records in msg, a nameserver's reply to a
// query about one record type. It returns nil when the nameserver
// answered: with records, with none, or with NXDOMAIN; otherwise it
// returns why not.
func readReply(
	msg *dns.Msg,
	resp *NameserverResponse,
	state *queryState,
) error {
	if msg.Rcode == dns.RcodeNameError {
		state.gotNXDomain = true

		return nil
	}

	if isErrorReply(msg) {
		state.gotErrorReply = true

		code, named := dns.RcodeToString[msg.Rcode]
		if !named {
			code = strconv.Itoa(msg.Rcode)
		}

		state.errorReply = code

		return fmt.Errorf(
			"server returned %s: %w", state.errorReply, ErrUnusableReply,
		)
	}

	// 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 server named in the delegation that
	// does not hold the zone may send one, as do a parent zone's servers
	// when FindAuthoritativeNameservers found no delegation for the
	// name's zone and moved on to a parent name.
	if !msg.Authoritative && len(msg.Answer) == 0 &&
		len(extractNSSet(msg.Ns)) > 0 {
		state.gotReferral = true

		return fmt.Errorf("server returned a referral: %w", ErrUnusableReply)
	}

	// A reply still truncated is one whose TCP retry failed, and holds
	// only the records that fit.
	if msg.Truncated {
		state.netErr = ErrTruncated

		return ErrTruncated
	}

	collectAnswerRecords(msg, resp, state)

	return nil
}

// collectAnswerRecords adds the records in msg's answer to resp, each
// value once per record type. For a name with a CNAME, a nameserver
// answers a query of any type with that CNAME, so the same value comes
// in the answer to every type asked for.
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]
		if !slices.Contains(resp.Records[typeName], val) {
			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
}

// classifyResponse sets the nameserver's status. One that answered no
// record type has failed, and Error says why; one that answered some has
// the status of those answers.
func classifyResponse(resp *NameserverResponse, state queryState) {
	switch {
	case state.gotNXDomain && !state.hasRecords:
		resp.Status = StatusNXDomain
	case state.gotTimeout && !state.answered:
		resp.Status = StatusTimeout
		resp.Error = "all queries timed out"
	case state.gotErrorReply && !state.answered:
		resp.Status = StatusError
		resp.Error = "server returned " + state.errorReply
	case state.gotRefused && !state.answered:
		resp.Status = StatusError
		resp.Error = "server returned REFUSED"
	case state.netErr != nil && !state.answered:
		resp.Status = StatusError
		resp.Error = "network error: " + state.netErr.Error()
	case state.gotReferral && !state.answered:
		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, recordTypes())
}

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

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

		// 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. It asks each
// nameserver of the name's zone for its A, AAAA and CNAME records only.
// 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
	}

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

	results, err := r.queryEachNS(ctx, nameservers, hostname, addressTypes())
	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. A nameserver whose query for one of the
// types failed gave only part of the addresses, and is left out. It
// returns ErrNoNameserverAnswered when every nameserver timed out,
// failed, returned a referral or was left out: 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 ||
			len(resp.FailedTypes) > 0 {
			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
}