ExploreDNS/internal/traverse/cache.go

package traverse

import (
	"fmt"
	"net"
	"strings"
	"sync"

	miekgdns "github.com/miekg/dns"
)

// StartServer is one entry returned by GetStartServers: a nameserver hostname
// plus its cached IPv4 addresses. IPs == nil means no addresses are cached and
// the caller must resolve the name itself (glueless).
type StartServer struct {
	Name string
	IPs  []string
}

// InfoCache is the hierarchical per-branch record cache (info_cache.rb). Each
// response wraps its parent's cache in a child so sibling branches never see
// each other's records; lookups recurse towards the root cache.
type InfoCache struct {
	parent *InfoCache
	mu     sync.RWMutex
	data   map[string][]miekgdns.RR
}

func NewInfoCache(parent *InfoCache) *InfoCache {
	return &InfoCache{
		parent: parent,
		data:   make(map[string][]miekgdns.RR),
	}
}

func (c *InfoCache) Child() *InfoCache {
	return NewInfoCache(c)
}

// canonicalName lowercases a DNS name and strips the trailing dot; the root
// (and empty string) canonicalises to "", matching the Ruby engine's
// representation of "no bailiwick".
func canonicalName(name string) string {
	return strings.TrimSuffix(strings.ToLower(name), ".")
}

func cacheKey(name string, qclass, qtype uint16) string {
	return fmt.Sprintf("%s:%d:%d", canonicalName(name), qclass, qtype)
}

func rrCacheKey(rr miekgdns.RR) string {
	h := rr.Header()
	return cacheKey(h.Name, h.Class, h.Rrtype)
}

// Add stores resource records, REPLACING any existing entries that share a
// name:class:type key (info_cache.rb add: clear pass, then append pass, so
// several records under one key in a single call are all kept).
func (c *InfoCache) Add(rrs []miekgdns.RR) {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, rr := range rrs {
		c.data[rrCacheKey(rr)] = nil
	}
	for _, rr := range rrs {
		key := rrCacheKey(rr)
		c.data[key] = append(c.data[key], rr)
	}
}

// AddHints seeds NS records for domain ("" = root hints) plus A/AAAA records
// for each server that has known addresses (info_cache.rb add_hints).
func (c *InfoCache) AddHints(domain string, servers []StartServer) {
	var rrs []miekgdns.RR
	owner := miekgdns.Fqdn(canonicalName(domain))
	for _, srv := range servers {
		name := miekgdns.Fqdn(canonicalName(srv.Name))
		rrs = append(rrs, &miekgdns.NS{
			Hdr: miekgdns.RR_Header{Name: owner, Rrtype: miekgdns.TypeNS, Class: miekgdns.ClassINET},
			Ns:  name,
		})
		for _, ip := range srv.IPs {
			addr := net.ParseIP(ip)
			if addr == nil {
				continue
			}
			if v4 := addr.To4(); v4 != nil {
				rrs = append(rrs, &miekgdns.A{
					Hdr: miekgdns.RR_Header{Name: name, Rrtype: miekgdns.TypeA, Class: miekgdns.ClassINET},
					A:   v4,
				})
			} else {
				rrs = append(rrs, &miekgdns.AAAA{
					Hdr:  miekgdns.RR_Header{Name: name, Rrtype: miekgdns.TypeAAAA, Class: miekgdns.ClassINET},
					AAAA: addr,
				})
			}
		}
	}
	c.Add(rrs)
}

// Get returns the cached RRset for name/class/type, consulting parent caches
// on a local miss. Returns nil when nothing is cached anywhere in the chain.
func (c *InfoCache) Get(name string, qclass, qtype uint16) []miekgdns.RR {
	key := cacheKey(name, qclass, qtype)
	c.mu.RLock()
	rrs, ok := c.data[key]
	c.mu.RUnlock()
	if ok {
		out := make([]miekgdns.RR, len(rrs))
		copy(out, rrs)
		return out
	}
	if c.parent != nil {
		return c.parent.Get(name, qclass, qtype)
	}
	return nil
}

// getNS finds the nearest cached NS RRset at or above domain, walking labels
// upward to the root (info_cache.rb get_ns?).
func (c *InfoCache) getNS(domain string) ([]miekgdns.RR, error) {
	domain = canonicalName(domain)
	for {
		if rrs := c.Get(domain, miekgdns.ClassINET, miekgdns.TypeNS); len(rrs) > 0 {
			return rrs, nil
		}
		if domain == "" {
			return nil, fmt.Errorf("no nameservers available for %q -- no root hints set??", domain)
		}
		if i := strings.Index(domain, "."); i >= 0 {
			domain = domain[i+1:]
		} else {
			domain = ""
		}
	}
}

// GetStartServers returns the servers to start querying for domain: the
// nearest cached NS RRset walking labels upward, each nameserver paired with
// its cached A addresses (nil when unknown). newbailiwick is the owner name of
// that NS RRset ("" for root).
func (c *InfoCache) GetStartServers(domain string) (starters []StartServer, newbailiwick string, err error) {
	ns, err := c.getNS(domain)
	if err != nil {
		return nil, "", err
	}
	for _, rr := range ns {
		nsrr, ok := rr.(*miekgdns.NS)
		if !ok {
			continue
		}
		name := canonicalName(nsrr.Ns)
		var ips []string
		for _, iprr := range c.Get(name, miekgdns.ClassINET, miekgdns.TypeA) {
			if a, ok := iprr.(*miekgdns.A); ok {
				ips = append(ips, a.A.String())
			}
		}
		starters = append(starters, StartServer{Name: name, IPs: ips})
	}
	newbailiwick = canonicalName(ns[0].Header().Name)
	return starters, newbailiwick, nil
}