ExploreDNS/internal/traverse/cache.go

package traverse

import (
	"net"
	"strings"
	"sync"

	miekgdns "github.com/miekg/dns"
)

// InfoCache is a two-level (parent/child) concurrent cache for NS records and
// glue addresses discovered during traversal.
//
// Lookups walk the parent chain: a child cache falls back to its parent when
// no local entry is found.  Writes always go to the local cache, never to the
// parent.  This makes it safe to give sibling branches separate child caches
// that share the root cache read-only in fast mode.
type InfoCache struct {
	parent *InfoCache
	mu     sync.RWMutex
	ns     map[string][]string
	glue   map[string][]net.IP
}

// NewInfoCache creates an InfoCache with an optional parent.
// Pass nil for a root-level cache with no parent.
func NewInfoCache(parent *InfoCache) *InfoCache {
	return &InfoCache{
		parent: parent,
		ns:     make(map[string][]string),
		glue:   make(map[string][]net.IP),
	}
}

// StoreNS records the nameserver names for zone in the local cache.
// Duplicate names within a zone are deduplicated.
func (c *InfoCache) StoreNS(zone string, nameservers []string) {
	if len(nameservers) == 0 {
		return
	}
	zone = normalize(zone)
	c.mu.Lock()
	seen := make(map[string]bool)
	for _, ns := range nameservers {
		ns = normalize(ns)
		if !seen[ns] {
			seen[ns] = true
			c.ns[zone] = append(c.ns[zone], ns)
		}
	}
	c.mu.Unlock()
}

// LookupNS returns the cached nameserver names for zone,
// walking the parent chain when no local entry is found.
func (c *InfoCache) LookupNS(zone string) []string {
	zone = normalize(zone)
	if names := c.localNS(zone); len(names) > 0 {
		return names
	}
	if c.parent != nil {
		return c.parent.LookupNS(zone)
	}
	return nil
}

func (c *InfoCache) localNS(zone string) []string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	names, ok := c.ns[zone]
	if !ok {
		return nil
	}
	result := make([]string, len(names))
	copy(result, names)
	return result
}

// StoreGlue records the IP addresses for a nameserver hostname in the local cache.
// Duplicate addresses are deduplicated.
func (c *InfoCache) StoreGlue(name string, addrs []net.IP) {
	if len(addrs) == 0 {
		return
	}
	name = normalize(name)
	c.mu.Lock()
	seen := make(map[string]bool)
	for _, addr := range addrs {
		key := addr.String()
		if !seen[key] {
			seen[key] = true
			c.glue[name] = append(c.glue[name], addr)
		}
	}
	c.mu.Unlock()
}

// LookupGlue returns the cached IP addresses for a nameserver hostname,
// walking the parent chain when no local entry is found.
func (c *InfoCache) LookupGlue(name string) []net.IP {
	name = normalize(name)
	if addrs := c.localGlue(name); len(addrs) > 0 {
		return addrs
	}
	if c.parent != nil {
		return c.parent.LookupGlue(name)
	}
	return nil
}

func (c *InfoCache) localGlue(name string) []net.IP {
	c.mu.RLock()
	defer c.mu.RUnlock()
	addrs, ok := c.glue[name]
	if !ok {
		return nil
	}
	result := make([]net.IP, len(addrs))
	copy(result, addrs)
	return result
}

// Child creates a new InfoCache that inherits from c.
// The child reads from c when a local lookup misses, but never writes to c.
func (c *InfoCache) Child() *InfoCache {
	return NewInfoCache(c)
}

// NSCount returns the number of zone→nameservers entries in the local cache.
func (c *InfoCache) NSCount() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.ns)
}

// GlueCount returns the number of nameserver→addresses entries in the local cache.
func (c *InfoCache) GlueCount() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.glue)
}

func normalize(name string) string {
	return strings.ToLower(miekgdns.Fqdn(name))
}