bucketeer/internal/dns/resolver.go

// Package dns provides a bounded, cached DNS enricher that follows CNAME chains
// with loop detection.
package dns

import (
	"context"
	"errors"
	"net"
	"strings"
	"time"

	"tea.karlcom.de/karl/bucketeer/internal/metrics"
)

// ErrCNAMELoop is returned when a CNAME chain revisits a host.
var ErrCNAMELoop = errors.New("cname loop detected")

// Resolver is the subset of DNS operations the enricher needs.
type Resolver interface {
	LookupCNAME(ctx context.Context, host string) (string, error)
	LookupIP(ctx context.Context, host string) ([]net.IP, error)
}

// Result is a resolved hostname: its CNAME chain (chain[0] is the queried host)
// and any A/AAAA addresses.
type Result struct {
	Chain []string
	IPs   []net.IP
}

// GoResolver uses the standard library resolver.
//
// ponytail: net.Resolver.LookupCNAME returns the final canonical name, so the
// chain has at most two entries and intermediate CNAME hops are not observed.
// Swap in miekg/dns if the full chain is ever needed.
type GoResolver struct {
	R *net.Resolver
}

func NewGoResolver() *GoResolver {
	return &GoResolver{R: net.DefaultResolver}
}

func (g *GoResolver) LookupCNAME(ctx context.Context, host string) (string, error) {
	return g.R.LookupCNAME(ctx, host)
}

func (g *GoResolver) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
	return g.R.LookupIP(ctx, "ip", host)
}

// Follow walks a CNAME chain up to maxDepth hops, detecting loops.
func Follow(ctx context.Context, r Resolver, host string, maxDepth int) ([]string, error) {
	if maxDepth < 1 {
		maxDepth = 1
	}
	chain := []string{host}
	cur := host
	for i := 0; i < maxDepth; i++ {
		next, err := r.LookupCNAME(ctx, cur)
		if err != nil {
			return chain, err
		}
		next = strings.TrimSuffix(next, ".")
		if next == "" || next == cur {
			break
		}
		for _, seen := range chain {
			if seen == next {
				return chain, ErrCNAMELoop
			}
		}
		chain = append(chain, next)
		cur = next
	}
	return chain, nil
}

// Enricher resolves hostnames with caching and bounded CNAME depth. It is safe
// for concurrent use.
type Enricher struct {
	Resolver    Resolver
	Cache       *TTLCache[Result]
	PositiveTTL time.Duration
	NegativeTTL time.Duration
	Timeout     time.Duration
	MaxDepth    int
	Metrics     *metrics.Metrics
}

// ResolveCNAME follows just the CNAME chain (no A/AAAA), for the high-volume
// enrichment path where only the targets are classified. Cached separately from
// full results.
func (e *Enricher) ResolveCNAME(ctx context.Context, host string) ([]string, bool) {
	if e.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, e.Timeout)
		defer cancel()
	}
	key := "c:" + host
	if e.Cache != nil {
		if res, ok := e.Cache.Get(key); ok {
			if e.Metrics != nil {
				e.Metrics.DNSCacheHits.Add(1)
			}
			return res.Chain, true
		}
	}
	if e.Metrics != nil {
		e.Metrics.DNSQueries.Add(1)
	}
	chain, err := Follow(ctx, e.Resolver, host, e.MaxDepth)
	if err != nil && len(chain) == 1 {
		if e.Metrics != nil && !definitive(err) {
			e.Metrics.DNSErrors.Add(1)
		}
		if e.Cache != nil {
			e.Cache.Set(key, Result{}, e.NegativeTTL)
		}
		return nil, false
	}
	if e.Cache != nil {
		e.Cache.Set(key, Result{Chain: chain}, e.PositiveTTL)
	}
	return chain, true
}

// ResolveIPs resolves A/AAAA for a host, cached.
func (e *Enricher) ResolveIPs(ctx context.Context, host string) ([]net.IP, bool) {
	if e.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, e.Timeout)
		defer cancel()
	}
	key := "ip:" + host
	if e.Cache != nil {
		if res, ok := e.Cache.Get(key); ok {
			if e.Metrics != nil {
				e.Metrics.DNSCacheHits.Add(1)
			}
			return res.IPs, true
		}
	}
	ips, err := e.Resolver.LookupIP(ctx, host)
	if err != nil || len(ips) == 0 {
		if e.Metrics != nil {
			e.Metrics.DNSErrors.Add(1)
		}
		if e.Cache != nil {
			e.Cache.Set(key, Result{}, e.NegativeTTL)
		}
		return nil, false
	}
	if e.Cache != nil {
		e.Cache.Set(key, Result{IPs: ips}, e.PositiveTTL)
	}
	return ips, true
}

// Resolve returns the result for host. ok is false when the lookup failed;
// failures (including "no such host") are negatively cached for NegativeTTL.
func (e *Enricher) Resolve(ctx context.Context, host string) (Result, bool) {
	if e.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, e.Timeout)
		defer cancel()
	}
	if e.Cache != nil {
		if res, ok := e.Cache.Get(host); ok {
			if e.Metrics != nil {
				e.Metrics.DNSCacheHits.Add(1)
			}
			return res, len(res.Chain) > 0 || len(res.IPs) > 0
		}
	}
	if e.Metrics != nil {
		e.Metrics.DNSQueries.Add(1)
	}

	chain, cnameErr := Follow(ctx, e.Resolver, host, e.MaxDepth)
	if cnameErr != nil && len(chain) == 1 {
		// A definitive NXDOMAIN has no A/AAAA to fetch either, so skip the
		// second lookup. NXDOMAIN is a valid answer, not an infrastructure error.
		if e.Metrics != nil && !definitive(cnameErr) {
			e.Metrics.DNSErrors.Add(1)
		}
		if e.Cache != nil {
			e.Cache.Set(host, Result{}, e.NegativeTTL)
		}
		return Result{}, false
	}
	res := Result{Chain: chain}
	if ips, err := e.Resolver.LookupIP(ctx, host); err == nil {
		res.IPs = ips
	}
	if e.Cache != nil {
		e.Cache.Set(host, res, e.PositiveTTL)
	}
	return res, true
}