ctrld/resolver.go

package ctrld

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/netip"
	"runtime"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"
	"github.com/rs/zerolog"
	"golang.org/x/sync/singleflight"
	"tailscale.com/net/netmon"
	"tailscale.com/net/tsaddr"

	"github.com/Control-D-Inc/ctrld/internal/dnscache"
	ctrldnet "github.com/Control-D-Inc/ctrld/internal/net"
)

const (
	// ResolverTypeDOH specifies DoH resolver.
	ResolverTypeDOH = "doh"
	// ResolverTypeDOH3 specifies DoH3 resolver.
	ResolverTypeDOH3 = "doh3"
	// ResolverTypeDOT specifies DoT resolver.
	ResolverTypeDOT = "dot"
	// ResolverTypeDOQ specifies DoQ resolver.
	ResolverTypeDOQ = "doq"
	// ResolverTypeOS specifies OS resolver.
	ResolverTypeOS = "os"
	// ResolverTypeLegacy specifies legacy resolver.
	ResolverTypeLegacy = "legacy"
	// ResolverTypePrivate is like ResolverTypeOS, but use for private resolver only.
	ResolverTypePrivate = "private"
	// ResolverTypeLocal is like ResolverTypeOS, but use for local resolver only.
	ResolverTypeLocal = "local"
	// ResolverTypeSDNS specifies resolver with information encoded using DNS Stamps.
	// See: https://dnscrypt.info/stamps-specifications/
	ResolverTypeSDNS = "sdns"
)

const controldPublicDns = "76.76.2.0"

const maxConcurrentOSResolverExchanges = 128

var controldPublicDnsWithPort = net.JoinHostPort(controldPublicDns, "53")

var osResolverExchangeSem = make(chan struct{}, maxConcurrentOSResolverExchanges)

var localResolver Resolver

func init() {
	// Initializing ProxyLogger here, so other places don't have to do nil check.
	l := zerolog.New(io.Discard)
	ProxyLogger.Store(&l)

	localResolver = newLocalResolver()
}

var (
	resolverMutex sync.Mutex
	or            *osResolver
	// osResolverPtr holds the same resolver as or, for readers alone. An
	// initialization holds resolverMutex across a scutil read of several
	// seconds, and a log header render must not wait for it.
	osResolverPtr    atomic.Pointer[osResolver]
	defaultLocalIPv4 atomic.Value // holds net.IP (IPv4)
	defaultLocalIPv6 atomic.Value // holds net.IP (IPv6)
)

// storeOsResolver keeps the resolver of the process and the copy that readers
// take. The caller holds resolverMutex.
func storeOsResolver(resolver *osResolver) {
	or = resolver
	osResolverPtr.Store(resolver)
}

// NameserversFn reads the nameservers of the system. A test of this package
// or of a package that drives the OS resolver replaces it to stay off the host.
var NameserversFn = nameservers

// osResolverReasonUnspecified marks a resolver read that no caller explained.
const osResolverReasonUnspecified = "unspecified"

// nameserverReads remembers the last nameserver list of one read. The system
// reports the same list many times, and only a change is news.
type nameserverReads struct {
	seen    bool
	last    []string
	repeats int
}

// record stores a new read. It reports whether the list changed, how many
// reads repeated the list before it, and the list that this read replaces.
func (n *nameserverReads) record(list []string) (changed bool, repeats int, before []string) {
	if n.seen && slices.Equal(n.last, list) {
		n.repeats++
		return false, n.repeats, n.last
	}
	before, repeats = n.last, n.repeats
	n.seen = true
	n.last = slices.Clone(list)
	n.repeats = 0
	return true, repeats, before
}

// osResolverLogState keeps the last nameserver reads and the reason of the
// initialization that runs. The read sits deep in the call chain, so the
// reason cannot travel as a parameter.
type osResolverLogState struct {
	mu     sync.Mutex
	system nameserverReads
	final  nameserverReads
	reason string
}

var osResolverLog = osResolverLogState{reason: osResolverReasonUnspecified}

func newLocalResolver() Resolver {
	var nss []string
	for _, addr := range Rfc1918Addresses() {
		nss = append(nss, net.JoinHostPort(addr, "53"))
	}
	return NewResolverWithNameserver(nss)
}

// LanQueryCtxKey is the context.Context key to indicate that the request is for LAN network.
type LanQueryCtxKey struct{}

// LanQueryCtx returns a context.Context with LanQueryCtxKey set.
func LanQueryCtx(ctx context.Context) context.Context {
	return context.WithValue(ctx, LanQueryCtxKey{}, true)
}

// LanOnlyQueryCtxKey is the context.Context key that limits an OS resolver
// query to its LAN nameservers.
type LanOnlyQueryCtxKey struct{}

// LanOnlyQueryCtx returns a context.Context with LanOnlyQueryCtxKey set.
//
// An OS resolver query made with it goes only to LAN-classified nameservers:
// private, loopback, link-local and CGNAT addresses. Every public nameserver
// is left out, the ones DHCP or a VPN supplied as well as ctrld's own public
// fallback, and with no LAN nameserver the query fails without being sent. It
// is for names that must not reach public DNS, such as an organization's
// Internal Domains. LanQueryCtx, by contrast, only drops ctrld's own public
// fallback.
func LanOnlyQueryCtx(ctx context.Context) context.Context {
	return context.WithValue(ctx, LanOnlyQueryCtxKey{}, true)
}

// isLanOnlyQuery reports whether ctx was made by LanOnlyQueryCtx.
func isLanOnlyQuery(ctx context.Context) bool {
	lanOnly, _ := ctx.Value(LanOnlyQueryCtxKey{}).(bool)
	return lanOnly
}

// defaultNameservers is like nameservers with each element formed "ip:53".
func defaultNameservers() []string {
	ns := nameservers()
	nss := make([]string, len(ns))
	for i := range ns {
		nss[i] = net.JoinHostPort(ns[i], "53")
	}
	return nss
}

// availableNameservers returns list of current available DNS servers of the system.
func availableNameservers() []string {
	var nss []string
	// Ignore local addresses to prevent loop.
	regularIPs, loopbackIPs, _ := netmon.LocalAddresses()
	machineIPsMap := make(map[string]struct{}, len(regularIPs))

	//load the logger
	logger := *ProxyLogger.Load()

	Log(context.Background(), logger.Debug(),
		"Got local addresses - regular IPs: %v, loopback IPs: %v", regularIPs, loopbackIPs)

	for _, v := range slices.Concat(regularIPs, loopbackIPs) {
		ipStr := v.String()
		machineIPsMap[ipStr] = struct{}{}
		Log(context.Background(), logger.Debug(),
			"Added local IP to OS resolverexclusion map: %s", ipStr)
	}

	systemNameservers := NameserversFn()
	logSystemNameservers(systemNameservers)

	for _, ns := range systemNameservers {
		if _, ok := machineIPsMap[ns]; ok {
			Log(context.Background(), logger.Debug(),
				"Skipping local nameserver: %s", ns)
			continue
		}
		nss = append(nss, ns)
		Log(context.Background(), logger.Debug(),
			"Added non-local nameserver: %s", ns)
	}

	logFinalNameservers(nss)
	return nss
}

// logSystemNameservers logs the system read at debug level on change only.
func logSystemNameservers(list []string) {
	osResolverLog.mu.Lock()
	changed, repeats, _ := osResolverLog.system.record(list)
	osResolverLog.mu.Unlock()
	if !changed {
		return
	}
	logger := *ProxyLogger.Load()
	Log(context.Background(), logger.Debug().Int("repeats", repeats),
		"Got system nameservers: %v", list)
}

// logFinalNameservers logs the effective read at debug level on change only,
// and sends each change to the journal, where an outage report needs it.
func logFinalNameservers(list []string) {
	osResolverLog.mu.Lock()
	changed, repeats, before := osResolverLog.final.record(list)
	reason := osResolverLog.reason
	osResolverLog.mu.Unlock()
	if !changed {
		return
	}
	logger := *ProxyLogger.Load()
	Log(context.Background(), logger.Debug().Int("repeats", repeats),
		"Final available nameservers: %v", list)
	Journal(logger.Info()).
		Str("source", osResolverSource()).
		Strs("before", before).
		Strs("after", list).
		Str("default_route", defaultRouteInterfaceName()).
		Str("reason", reason).
		Msg("OS resolver set changed")
}

// osResolverSource names the place where the platform holds the nameservers.
func osResolverSource() string {
	switch runtime.GOOS {
	case "darwin":
		return "scutil"
	case "windows":
		return "dhcp"
	case "linux", "dragonfly", "freebsd", "netbsd", "openbsd":
		return "resolv.conf"
	}
	return "unknown"
}

// defaultRouteInterfaceName returns the interface of the default route, and an
// empty string when the route table does not answer.
func defaultRouteInterfaceName() string {
	name, err := netmon.DefaultRouteInterface()
	if err != nil {
		return ""
	}
	return name
}

// setOsResolverReason keeps the reason of the initialization that runs.
func setOsResolverReason(reason string) {
	osResolverLog.mu.Lock()
	defer osResolverLog.mu.Unlock()
	osResolverLog.reason = reason
}

// InitializeOsResolver initializes OS resolver using the current system DNS settings.
// It returns the nameservers that is going to be used by the OS resolver.
//
// It's the caller's responsibility to ensure the system DNS is in a clean state before
// calling this function.
func InitializeOsResolver(guardAgainstNoNameservers bool) []string {
	return InitializeOsResolverWithReason(guardAgainstNoNameservers, osResolverReasonUnspecified)
}

// InitializeOsResolverWithReason is InitializeOsResolver with the reason that
// the journal reports when the resolver set changes.
func InitializeOsResolverWithReason(guardAgainstNoNameservers bool, reason string) []string {
	effective, _ := InitializeOsResolverWithSystemNameserversReason(guardAgainstNoNameservers, reason)
	return effective
}

// InitializeOsResolverWithSystemNameservers initializes the OS resolver and
// returns both the effective resolver list and the unmodified nameservers
// discovered from the system. The latter deliberately excludes synthetic
// fallbacks added by initializeOsResolver.
func InitializeOsResolverWithSystemNameservers(guardAgainstNoNameservers bool) (effective, system []string) {
	return InitializeOsResolverWithSystemNameserversReason(guardAgainstNoNameservers, osResolverReasonUnspecified)
}

// InitializeOsResolverWithSystemNameserversReason is
// InitializeOsResolverWithSystemNameservers with the reason that the journal
// reports when the resolver set changes.
func InitializeOsResolverWithSystemNameserversReason(guardAgainstNoNameservers bool, reason string) (effective, system []string) {
	resolverMutex.Lock()
	defer resolverMutex.Unlock()

	setOsResolverReason(reason)
	defer setOsResolverReason(osResolverReasonUnspecified)

	system = availableNameservers()
	if system == nil {
		// A non-nil empty slice means discovery completed and found no DNS.
		// Callers use nil to mean that discovery was not attempted.
		system = []string{}
	}
	effective, system, skip := osResolverNameserverSets(system, guardAgainstNoNameservers)
	if skip {
		return effective, system
	}
	storeOsResolver(newResolverWithNameserver(effective))
	return effective, system
}

func osResolverNameserverSets(system []string, guardAgainstNoNameservers bool) (effective, discovered []string, skip bool) {
	// if no nameservers, return empty slice so we dont remove all nameservers
	if len(system) == 0 && guardAgainstNoNameservers {
		return []string{}, system, true
	}
	return initializeOsResolver(system), system, false
}

// initializeOsResolver performs logic for choosing OS resolver nameserver.
// The logic:
//
// - First available LAN servers are saved and store.
// - Later calls, if no LAN servers available, the saved servers above will be used.
func initializeOsResolver(servers []string) []string {

	var lanNss, publicNss []string

	// First categorize servers
	for _, ns := range servers {
		addr, err := netip.ParseAddr(ns)
		if err != nil {
			continue
		}
		server := net.JoinHostPort(ns, "53")
		if isLanAddr(addr) {
			lanNss = append(lanNss, server)
		} else {
			publicNss = append(publicNss, server)
		}
	}

	if len(publicNss) == 0 {
		publicNss = []string{controldPublicDnsWithPort}
	}

	return slices.Concat(lanNss, publicNss)
}

// Resolver is the interface that wraps the basic DNS operations.
//
// Resolve resolves the DNS query, return the result and the corresponding error.
type Resolver interface {
	Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, error)
}

var errUnknownResolver = errors.New("unknown resolver")

// NewResolver creates a Resolver based on the given upstream config.
func NewResolver(uc *UpstreamConfig) (Resolver, error) {
	typ := uc.Type
	switch typ {
	case ResolverTypeDOH, ResolverTypeDOH3:
		return newDohResolver(uc), nil
	case ResolverTypeDOT:
		return &dotResolver{uc: uc}, nil
	case ResolverTypeDOQ:
		return &doqResolver{uc: uc}, nil
	case ResolverTypeOS:
		resolverMutex.Lock()
		if or == nil {
			ProxyLogger.Load().Debug().Msgf("Initialize new OS resolver")
			storeOsResolver(newResolverWithNameserver(defaultNameservers()))
		}
		resolverMutex.Unlock()
		return or, nil
	case ResolverTypeLegacy:
		return &legacyResolver{uc: uc}, nil
	case ResolverTypePrivate:
		return NewPrivateResolver(), nil
	case ResolverTypeLocal:
		return localResolver, nil
	}
	return nil, fmt.Errorf("%w: %s", errUnknownResolver, typ)
}

type osResolver struct {
	lanServers    atomic.Pointer[[]string]
	publicServers atomic.Pointer[[]string]
	group         *singleflight.Group
	cache         *sync.Map
	// Per-resolver seams let tests exercise the production Resolve path without
	// mutating process-wide resolver state.
	exchangeDNS dnsExchangeFunc
	localIP     func(string) net.IP
}

type osResolverResult struct {
	answer *dns.Msg
	err    error
	server string
	lan    bool
}

type publicResponse struct {
	answer *dns.Msg
	server string
}

// OsResolverNameservers returns the current OS resolver nameservers (host:port format).
// Returns nil if the OS resolver has not been initialized.
// The read takes no lock, so a log header render never waits for a resolver
// initialization.
func OsResolverNameservers() []string {
	r := osResolverPtr.Load()
	if r == nil {
		return nil
	}
	var nss []string
	if lan := r.lanServers.Load(); lan != nil {
		nss = append(nss, *lan...)
	}
	if pub := r.publicServers.Load(); pub != nil {
		nss = append(nss, *pub...)
	}
	return nss
}

// AppendOsResolverNameservers adds additional nameservers to the existing OS resolver
// without reinitializing it. This is used for late-arriving nameservers such as AD
// domain controller IPs discovered via background retry.
// Returns true if nameservers were actually added.
func AppendOsResolverNameservers(servers []string) bool {
	if len(servers) == 0 {
		return false
	}
	resolverMutex.Lock()
	defer resolverMutex.Unlock()
	if or == nil {
		return false
	}

	// Collect existing nameservers to avoid duplicates.
	existing := make(map[string]bool)
	if lan := or.lanServers.Load(); lan != nil {
		for _, s := range *lan {
			existing[s] = true
		}
	}
	if pub := or.publicServers.Load(); pub != nil {
		for _, s := range *pub {
			existing[s] = true
		}
	}

	var added bool
	for _, s := range servers {
		// Normalize to host:port format.
		if _, _, err := net.SplitHostPort(s); err != nil {
			s = net.JoinHostPort(s, "53")
		}
		if existing[s] {
			continue
		}
		existing[s] = true
		added = true

		ip, _, _ := net.SplitHostPort(s)
		addr, _ := netip.ParseAddr(ip)
		if isLanAddr(addr) {
			lan := or.lanServers.Load()
			newLan := append(append([]string{}, (*lan)...), s)
			or.lanServers.Store(&newLan)
		} else {
			pub := or.publicServers.Load()
			newPub := append(append([]string{}, (*pub)...), s)
			or.publicServers.Store(&newPub)
		}
	}
	return added
}

// SetDefaultLocalIPv4 updates the stored local IPv4.
func SetDefaultLocalIPv4(ip net.IP) {
	Log(context.Background(), ProxyLogger.Load().Debug(), "SetDefaultLocalIPv4: %s", ip)
	defaultLocalIPv4.Store(ip)
}

// SetDefaultLocalIPv6 updates the stored local IPv6.
func SetDefaultLocalIPv6(ip net.IP) {
	Log(context.Background(), ProxyLogger.Load().Debug(), "SetDefaultLocalIPv6: %s", ip)
	defaultLocalIPv6.Store(ip)
}

// GetDefaultLocalIPv4 returns the stored local IPv4 or nil if none.
func GetDefaultLocalIPv4() net.IP {
	if v := defaultLocalIPv4.Load(); v != nil {
		return v.(net.IP)
	}
	return nil
}

// GetDefaultLocalIPv6 returns the stored local IPv6 or nil if none.
func GetDefaultLocalIPv6() net.IP {
	if v := defaultLocalIPv6.Load(); v != nil {
		return v.(net.IP)
	}
	return nil
}

type dnsExchangeFunc func(context.Context, *dns.Msg, string, net.IP) (*dns.Msg, time.Duration, error)

func exchangeDNS(ctx context.Context, msg *dns.Msg, server string, localIP net.IP) (*dns.Msg, time.Duration, error) {
	baseDialer := &net.Dialer{
		Timeout:  3 * time.Second,
		Resolver: &net.Resolver{PreferGo: true},
	}
	if localIP != nil {
		baseDialer.LocalAddr = &net.UDPAddr{IP: localIP, Port: 0}
	}
	dnsClient := &dns.Client{Net: "udp"}
	dnsClient.Dialer = baseDialer
	return dnsClient.ExchangeContext(ctx, msg, server)
}

func defaultLocalIPForServer(server string) net.IP {
	if runtime.GOOS != "darwin" {
		return nil
	}
	host, _, err := net.SplitHostPort(server)
	if err != nil {
		return nil
	}
	addr, err := netip.ParseAddr(host)
	if err == nil && addr.Zone() != "" {
		// Scoped resolvers select their interface via the destination zone.
		// A cached default-interface source can have the wrong family or scope.
		return nil
	}
	if err == nil && addr.Is6() && !addr.Is4In6() {
		return GetDefaultLocalIPv6()
	}
	return GetDefaultLocalIPv4()
}

func preSendUnreachable(err error) bool {
	var opErr *net.OpError
	if !errors.As(err, &opErr) || (opErr.Op != "dial" && opErr.Op != "write") {
		return false
	}
	return ctrldnet.IsUnreachable(err)
}

// customDNSExchangeWith preserves the preferred source first. A route-selected
// retry is allowed only when the caller knows the server is an OS-selected resolver,
// not ctrld's synthetic public fallback. This includes public DNS pushed by a VPN:
// unbinding changes the source route, not the recipient.
func customDNSExchangeWith(ctx context.Context, msg *dns.Msg, server string, desiredLocalIP net.IP, allowRouteSelectedRetry bool, exchange dnsExchangeFunc) (*dns.Msg, time.Duration, error) {
	answer, rtt, err := exchange(ctx, msg, server, desiredLocalIP)
	if answer != nil || err == nil || ctx.Err() != nil || desiredLocalIP == nil || !allowRouteSelectedRetry || !preSendUnreachable(err) {
		return answer, rtt, err
	}

	Log(ctx, ProxyLogger.Load().Debug(), "OS resolver source binding is unreachable; retrying with route-selected source")
	return exchange(ctx, msg.Copy(), server, nil)
}

// allowRouteSelectedRetryForOSServer excludes only ctrld's synthetic public
// fallback. System-provided resolvers remain eligible even when their addresses
// are public, as with VPNs that push public DNS servers.
func allowRouteSelectedRetryForOSServer(server string) bool {
	return server != controldPublicDnsWithPort
}

const hotCacheTTL = time.Second

// Resolve resolves DNS queries using pre-configured nameservers.
// The Query is sent to all nameservers concurrently, and the first
// success response will be returned.
//
// To guard against unexpected DoS to upstreams, multiple queries of
// the same Qtype to a domain will be shared, so there's only 1 qps
// for each upstream at any time.
//
// Further, a hot cache will be used, so repeated queries will be cached
// for a short period (currently 1 second), reducing unnecessary traffics
// sent to upstreams.
func (o *osResolver) Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, error) {
	if err := validateMsg(msg); err != nil {
		return nil, err
	}
	if len(msg.Question) == 0 {
		return nil, errors.New("no question found")
	}
	domain := strings.TrimSuffix(msg.Question[0].Name, ".")
	qtype := msg.Question[0].Qtype

	// Unique key for the singleflight group. The EDNS Client Subnet is part of
	// the key so subnet-specific answers are neither coalesced nor hot-cached
	// across different subnets (RFC 7871 §7.3).
	key := fmt.Sprintf("%s:%d:%s", domain, qtype, dnscache.CanonicalECS(msg))
	// A LAN-only query must not join, or be answered from the hot cache of, an
	// ordinary query for the same name: that answer may have come from a public
	// nameserver.
	if isLanOnlyQuery(ctx) {
		key += ":lan-only"
	}

	// Checking the cache first.
	if val, ok := o.cache.Load(key); ok {
		if val, ok := val.(*dns.Msg); ok {
			Log(ctx, ProxyLogger.Load().Debug(), "hit hot cached result: %s - %s", domain, dns.TypeToString[qtype])
			res := val.Copy()
			SetCacheReply(res, msg, val.Rcode)
			return res, nil
		}
	}

	// Ensure only one DNS query is in flight for the key.
	v, err, shared := o.group.Do(key, func() (interface{}, error) {
		msg, err := o.resolve(ctx, msg)
		if err != nil {
			return nil, err
		}
		// If we got an answer, storing it to the hot cache for hotCacheTTL
		// This prevents possible DoS to upstream, ensuring there's only 1 QPS.
		o.cache.Store(key, msg)
		// Depends on go runtime scheduling, the result may end up in hot cache longer
		// than hotCacheTTL duration. However, this is fine since we only want to guard
		// against DoS attack. The result will be cleaned from the cache eventually.
		time.AfterFunc(hotCacheTTL, func() {
			o.removeCache(key)
		})
		return msg, nil
	})
	if err != nil {
		return nil, err
	}

	sharedMsg, ok := v.(*dns.Msg)
	if !ok {
		return nil, fmt.Errorf("invalid answer for key: %s", key)
	}
	res := sharedMsg.Copy()
	SetCacheReply(res, msg, sharedMsg.Rcode)
	if shared {
		Log(ctx, ProxyLogger.Load().Debug(), "shared result: %s - %s", domain, dns.TypeToString[qtype])
	}

	return res, nil
}

// resolve sends the query to current nameservers.
func (o *osResolver) resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, error) {
	publicServers := *o.publicServers.Load()
	var nss []string
	if p := o.lanServers.Load(); p != nil {
		nss = append(nss, (*p)...)
	}
	// A LAN-only query is never sent to a public nameserver.
	if isLanOnlyQuery(ctx) {
		publicServers = nil
	}
	numServers := len(nss) + len(publicServers)

	// If this is a LAN query, skip public DNS.
	lan, ok := ctx.Value(LanQueryCtxKey{}).(bool)

	// remove controldPublicDnsWithPort from publicServers for LAN queries
	// this is to prevent DoS for high frequency local requests
	if ok && lan {
		if index := slices.Index(publicServers, controldPublicDnsWithPort); index != -1 {
			publicServers = slices.Delete(publicServers, index, index+1)
			numServers--
		}
	}
	question := ""
	if msg != nil && len(msg.Question) > 0 {
		question = msg.Question[0].Name
	}
	Log(ctx, ProxyLogger.Load().Debug(), "os resolver query for %s with nameservers: %v public: %v", question, nss, publicServers)

	// New check: If no resolvers are available, return an error.
	if numServers == 0 {
		return nil, errors.New("no nameservers available for query")
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	ch := make(chan *osResolverResult, numServers)
	wg := &sync.WaitGroup{}
	exchange := o.exchangeDNS
	if exchange == nil {
		exchange = exchangeDNS
	}
	localIPForServer := o.localIP
	if localIPForServer == nil {
		localIPForServer = defaultLocalIPForServer
	}
	wg.Add(numServers)
	go func() {
		wg.Wait()
		close(ch)
	}()

	do := func(servers []string, isLan bool) {
		for _, server := range servers {
			go func(server string) {
				defer wg.Done()
				release, ok := acquireOSResolverExchangeSlot(ctx)
				if !ok {
					ch <- &osResolverResult{err: ctx.Err(), server: server, lan: isLan}
					return
				}
				defer release()

				var answer *dns.Msg
				var err error
				localOSResolverIP := localIPForServer(server)
				answer, _, err = customDNSExchangeWith(ctx, msg.Copy(), server, localOSResolverIP, allowRouteSelectedRetryForOSServer(server), exchange)
				ch <- &osResolverResult{answer: answer, err: err, server: server, lan: isLan}
			}(server)
		}
	}

	logAnswer := func(server string) {
		host, _, err := net.SplitHostPort(server)
		if err != nil {
			// If splitting fails, fallback to the original server string
			host = server
		}
		Log(ctx, ProxyLogger.Load().Debug(), "got answer from nameserver: %s", host)
	}

	// try local nameservers
	if len(nss) > 0 {
		do(nss, true)
	}

	// we must always try the public servers too, since DCHP may have only public servers
	// this is okay to do since we always prefer LAN nameserver responses
	if len(publicServers) > 0 {
		do(publicServers, false)
	}

	var (
		nonSuccessAnswer      *dns.Msg
		nonSuccessServer      string
		controldSuccessAnswer *dns.Msg
		publicResponses       []publicResponse
	)
	errs := make([]error, 0, numServers)
	for res := range ch {
		switch {
		case res.answer != nil && res.answer.Rcode == dns.RcodeSuccess:
			switch {
			case res.lan:
				// Always prefer LAN responses immediately
				Log(ctx, ProxyLogger.Load().Debug(), "using LAN answer from: %s", res.server)
				cancel()
				logAnswer(res.server)
				return res.answer, nil
			case res.server == controldPublicDnsWithPort:
				controldSuccessAnswer = res.answer
			case !res.lan:
				// if there are no LAN nameservers, we should not wait
				// just use the first response
				if len(nss) == 0 {
					Log(ctx, ProxyLogger.Load().Debug(), "using public answer from: %s", res.server)
					cancel()
					logAnswer(res.server)
					return res.answer, nil
				}
				publicResponses = append(publicResponses, publicResponse{
					answer: res.answer,
					server: res.server,
				})
			}
		case res.answer != nil:
			Log(ctx, ProxyLogger.Load().Debug(), "got non-success answer from: %s with code: %d",
				res.server, res.answer.Rcode)
			// When there are no LAN nameservers, we should not wait
			// for other nameservers to respond.
			if len(nss) == 0 {
				Log(ctx, ProxyLogger.Load().Debug(), "no lan nameservers using public non success answer")
				cancel()
				logAnswer(res.server)
				return res.answer, nil
			}
			nonSuccessAnswer = res.answer
			nonSuccessServer = res.server
		}
		errs = append(errs, res.err)
	}

	if len(publicResponses) > 0 {
		resp := publicResponses[0]
		Log(ctx, ProxyLogger.Load().Debug(), "using public answer from: %s", resp.server)
		logAnswer(resp.server)
		return resp.answer, nil
	}
	if controldSuccessAnswer != nil {
		Log(ctx, ProxyLogger.Load().Debug(), "using ControlD answer from: %s", controldPublicDnsWithPort)
		logAnswer(controldPublicDnsWithPort)
		return controldSuccessAnswer, nil
	}
	if nonSuccessAnswer != nil {
		Log(ctx, ProxyLogger.Load().Debug(), "using non-success answer from: %s", nonSuccessServer)
		logAnswer(nonSuccessServer)
		return nonSuccessAnswer, nil
	}
	return nil, errors.Join(errs...)
}

func acquireOSResolverExchangeSlot(ctx context.Context) (func(), bool) {
	select {
	case osResolverExchangeSem <- struct{}{}:
		return func() { <-osResolverExchangeSem }, true
	case <-ctx.Done():
		return nil, false
	}
}

func (o *osResolver) removeCache(key string) {
	o.cache.Delete(key)
}

type legacyResolver struct {
	uc *UpstreamConfig
}

func (r *legacyResolver) Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, error) {
	if err := validateMsg(msg); err != nil {
		return nil, err
	}

	// See comment in (*dotResolver).resolve method.
	dialer := newDialer(net.JoinHostPort(controldPublicDns, "53"))
	dnsTyp := uint16(0)
	if msg != nil && len(msg.Question) > 0 {
		dnsTyp = msg.Question[0].Qtype
	}
	_, udpNet := r.uc.netForDNSType(dnsTyp)
	dnsClient := &dns.Client{
		Net:    udpNet,
		Dialer: dialer,
	}
	endpoint := r.uc.Endpoint
	if r.uc.BootstrapIP != "" {
		dnsClient.Net = "udp"
		_, port, _ := net.SplitHostPort(endpoint)
		endpoint = net.JoinHostPort(r.uc.BootstrapIP, port)
	}

	answer, _, err := dnsClient.ExchangeContext(ctx, msg, endpoint)
	return answer, err
}

type dummyResolver struct{}

func (d dummyResolver) Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, error) {
	if err := validateMsg(msg); err != nil {
		return nil, err
	}
	ans := new(dns.Msg)
	ans.SetReply(msg)
	return ans, nil
}

// LookupIP looks up domain using current system nameservers settings.
// It returns a slice of that host's IPv4 and IPv6 addresses.
func LookupIP(domain string) []string {
	nss := initDefaultOsResolver()
	return lookupIP(domain, -1, nss)
}

// initDefaultOsResolver initializes the default OS resolver with system's default nameservers if it hasn't been initialized yet.
// It returns the combined list of LAN and public nameservers currently held by the resolver.
func initDefaultOsResolver() []string {
	resolverMutex.Lock()
	defer resolverMutex.Unlock()
	if or == nil {
		ProxyLogger.Load().Debug().Msgf("Initialize new OS resolver with default nameservers")
		storeOsResolver(newResolverWithNameserver(defaultNameservers()))
	}
	nss := *or.lanServers.Load()
	nss = append(nss, *or.publicServers.Load()...)
	return nss
}

// lookupIP looks up domain with given timeout and bootstrapDNS.
// If the timeout is negative, default timeout 2000 ms will be used.
// It returns nil if bootstrapDNS is nil or empty.
func lookupIP(domain string, timeout int, bootstrapDNS []string) (ips []string) {
	if net.ParseIP(domain) != nil {
		return []string{domain}
	}
	if bootstrapDNS == nil {
		ProxyLogger.Load().Debug().Msgf("empty bootstrap DNS")
		return nil
	}

	resolver := newResolverWithNameserver(bootstrapDNS)
	ProxyLogger.Load().Debug().Msgf("resolving %q using bootstrap DNS %q", domain, bootstrapDNS)
	timeoutMs := 2000
	if timeout > 0 && timeout < timeoutMs {
		timeoutMs = timeout
	}
	questionDomain := dns.Fqdn(domain)

	// Getting the real target domain name from CNAME if presents.
	targetDomain := func(answers []dns.RR) string {
		for _, a := range answers {
			switch ar := a.(type) {
			case *dns.CNAME:
				return ar.Target
			}
		}
		return questionDomain
	}
	// Getting ip address from A or AAAA record.
	ipFromRecord := func(record dns.RR, target string) string {
		switch ar := record.(type) {
		case *dns.A:
			if ar.Hdr.Name != target || len(ar.A) == 0 {
				return ""
			}
			return ar.A.String()
		case *dns.AAAA:
			if ar.Hdr.Name != target || len(ar.AAAA) == 0 {
				return ""
			}
			return ar.AAAA.String()
		}
		return ""
	}

	lookup := func(dnsType uint16) {
		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutMs)*time.Millisecond)
		defer cancel()
		m := new(dns.Msg)
		m.SetQuestion(questionDomain, dnsType)
		m.RecursionDesired = true

		r, err := resolver.Resolve(ctx, m)
		if err != nil {
			ProxyLogger.Load().Error().Err(err).Msgf("could not lookup %q record for domain %q", dns.TypeToString[dnsType], domain)
			return
		}
		if r.Rcode != dns.RcodeSuccess {
			ProxyLogger.Load().Error().Msgf("could not resolve domain %q, return code: %s", domain, dns.RcodeToString[r.Rcode])
			return
		}
		if len(r.Answer) == 0 {
			ProxyLogger.Load().Error().Msg("no answer from OS resolver")
			return
		}
		target := targetDomain(r.Answer)
		for _, a := range r.Answer {
			if ip := ipFromRecord(a, target); ip != "" {
				ips = append(ips, ip)
			}
		}
	}
	// Find all A, AAAA records of the domain.
	for _, dnsType := range []uint16{dns.TypeAAAA, dns.TypeA} {
		lookup(dnsType)
	}
	return ips
}

// NewBootstrapResolver returns an OS resolver, which use following nameservers:
//
//   - Gateway IP address (depends on OS).
//   - Input servers.
func NewBootstrapResolver(servers ...string) Resolver {
	logger := *ProxyLogger.Load()

	Log(context.Background(), logger.Debug(), "NewBootstrapResolver called with servers: %v", servers)
	nss := defaultNameservers()
	nss = append([]string{controldPublicDnsWithPort}, nss...)
	for _, ns := range servers {
		nss = append([]string{net.JoinHostPort(ns, "53")}, nss...)
	}
	return NewResolverWithNameserver(nss)
}

// NewPrivateResolver returns an OS resolver, which includes only private DNS servers,
// excluding:
//
// - Nameservers from /etc/resolv.conf file.
// - Nameservers which is local RFC1918 addresses.
//
// This is useful for doing PTR lookup in LAN network.
func NewPrivateResolver() Resolver {
	nss := initDefaultOsResolver()
	resolveConfNss := currentNameserversFromResolvconf()
	localRfc1918Addrs := Rfc1918Addresses()
	n := 0
	for _, ns := range nss {
		host, _, _ := net.SplitHostPort(ns)
		// Ignore nameserver from resolve.conf file, because the nameserver can be either:
		//
		//  - ctrld itself.
		//  - Direct listener that has ctrld as an upstream (e.g: dnsmasq).
		//
		// causing the query always succeed.
		if slices.Contains(resolveConfNss, host) {
			continue
		}
		// Ignoring local RFC 1918 addresses.
		if slices.Contains(localRfc1918Addrs, host) {
			continue
		}
		ip := net.ParseIP(host)
		if ip != nil && ip.IsPrivate() && !ip.IsLoopback() {
			nss[n] = ns
			n++
		}
	}
	nss = nss[:n]
	return newResolverWithNameserver(nss)
}

// NewResolverWithNameserver returns a Resolver which uses the given nameservers
// for resolving DNS queries. If nameservers is empty, a dummy resolver will be returned.
//
// Each nameserver must be form "host:port". It's the caller responsibility to ensure all
// nameservers are well formatted by using net.JoinHostPort function.
func NewResolverWithNameserver(nameservers []string) Resolver {
	if len(nameservers) == 0 {
		return &dummyResolver{}
	}
	return newResolverWithNameserver(nameservers)
}

// newResolverWithNameserver returns an OS resolver from given nameservers list.
// The caller must ensure each server in list is formed "ip:53".
func newResolverWithNameserver(nameservers []string) *osResolver {
	r := &osResolver{
		group: &singleflight.Group{},
		cache: &sync.Map{},
	}
	var publicNss []string
	var lanNss []string
	for _, ns := range slices.Sorted(slices.Values(nameservers)) {
		ip, _, _ := net.SplitHostPort(ns)
		addr, _ := netip.ParseAddr(ip)
		if isLanAddr(addr) {
			lanNss = append(lanNss, ns)
		} else {
			publicNss = append(publicNss, ns)
		}
	}
	r.lanServers.Store(&lanNss)
	r.publicServers.Store(&publicNss)
	return r
}

// Rfc1918Addresses returns the list of local physical interfaces private IP addresses
func Rfc1918Addresses() []string {
	vis := validInterfaces()
	var res []string
	netmon.ForeachInterface(func(i netmon.Interface, prefixes []netip.Prefix) {
		// Skip virtual interfaces.
		if _, existed := vis[i.Name]; !existed {
			return
		}
		addrs, _ := i.Addrs()
		for _, addr := range addrs {
			ipNet, ok := addr.(*net.IPNet)
			if !ok || !ipNet.IP.IsPrivate() {
				continue
			}
			res = append(res, ipNet.IP.String())
		}
	})
	return res
}

func newDialer(dnsAddress string) *net.Dialer {
	return &net.Dialer{
		Resolver: &net.Resolver{
			PreferGo: true,
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				d := net.Dialer{}
				return d.DialContext(ctx, network, dnsAddress)
			},
		},
	}
}

// isLanAddr reports whether addr is considered a LAN ip address.
func isLanAddr(addr netip.Addr) bool {
	return addr.IsPrivate() ||
		addr.IsLoopback() ||
		addr.IsLinkLocalUnicast() ||
		tsaddr.CGNATRange().Contains(addr)
}

func validateMsg(msg *dns.Msg) error {
	if msg == nil {
		return errors.New("nil DNS message")
	}
	if len(msg.Question) == 0 {
		return errors.New("no question found")
	}
	return nil
}