ExploreDNS/internal/traverse/cache.go

package traverse

import (
	"net"
	"strings"
	"sync"

	miekgdns "github.com/miekg/dns"
)

type InfoCache struct {
	parent   *InfoCache
	mu       sync.RWMutex
	ns       map[string][]string
	glue     map[string][]net.IP
}

func NewInfoCache(parent *InfoCache) *InfoCache {
	return &InfoCache{
		parent: parent,
		ns:     make(map[string][]string),
		glue:   make(map[string][]net.IP),
	}
}

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()
}

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
}

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()
}

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
}

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

func (c *InfoCache) NSCount() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.ns)
}

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