initial commit

internal/engine/reverse.go

@@ -0,0 +1,202 @@
+package engine
+
+import (
+	"context"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/miekg/dns"
+)
+
+// ReverseCache maps client IPs to reverse-DNS names so the ACL can match on
+// names without a synchronous PTR lookup on the query path.
+//
+// A miss schedules a background lookup and, for the duration of that lookup,
+// matches on nothing. That is the right trade: a blocking PTR on the hot path
+// would make a broken reverse zone a latency incident, whereas a miss only
+// means an ACL rule keyed on a name does not fire for the first few queries.
+type ReverseCache struct {
+	mu      sync.RWMutex
+	entries map[string]reverseEntry
+	// negative caches the "no PTR exists" answer so we do not re-query for
+	// every query from the same client.
+	max      int
+	ttl      time.Duration
+	negative time.Duration
+	resolver *dns.Client
+	names    string
+	timeout  time.Duration
+	inflight map[string]bool
+}
+
+type reverseEntry struct {
+	names    []string
+	fetched  time.Time
+	negative bool
+}
+
+// ReverseOptions configures a ReverseCache.
+type ReverseOptions struct {
+	// Max is the number of cached clients.
+	Max int
+	// TTL is how long a successful PTR is trusted.
+	TTL time.Duration
+	// NegativeTTL is how long a "no PTR" answer is trusted.
+	NegativeTTL time.Duration
+	// Nameserver is the server to ask, e.g. "9.9.9.9:53". Empty uses the
+	// system resolver via net.LookupAddr.
+	Nameserver string
+	// Timeout bounds one PTR lookup.
+	Timeout time.Duration
+}
+
+// NewReverseCache builds a cache.
+func NewReverseCache(opt ReverseOptions) *ReverseCache {
+	if opt.Max <= 0 {
+		opt.Max = 16384
+	}
+	if opt.TTL <= 0 {
+		opt.TTL = 10 * time.Minute
+	}
+	if opt.NegativeTTL <= 0 {
+		opt.NegativeTTL = 30 * time.Minute
+	}
+	if opt.Timeout <= 0 {
+		opt.Timeout = 2 * time.Second
+	}
+	rc := &ReverseCache{
+		entries:  make(map[string]reverseEntry, 1024),
+		max:      opt.Max,
+		ttl:      opt.TTL,
+		negative: opt.NegativeTTL,
+		names:    opt.Nameserver,
+		timeout:  opt.Timeout,
+		inflight: make(map[string]bool),
+	}
+	if opt.Nameserver != "" {
+		rc.resolver = &dns.Client{Net: "udp", Timeout: opt.Timeout}
+	}
+	return rc
+}
+
+// Names returns the cached reverse names for an IP, and whether the entry was
+// present. The second return is false for a miss (including a cached negative).
+func (r *ReverseCache) Names(ip string) ([]string, bool) {
+	if r == nil {
+		return nil, false
+	}
+	r.mu.RLock()
+	e, ok := r.entries[ip]
+	r.mu.RUnlock()
+	if !ok {
+		r.Prime(ip)
+		return nil, false
+	}
+	ttl := r.ttl
+	if e.negative {
+		ttl = r.negative
+	}
+	if time.Since(e.fetched) > ttl {
+		r.Prime(ip)
+		return e.names, true // serve the stale value while refreshing
+	}
+	return e.names, !e.negative
+}
+
+// Prime schedules a background PTR lookup for ip, at most one in flight.
+func (r *ReverseCache) Prime(ip string) {
+	if r == nil || ip == "" {
+		return
+	}
+	r.mu.Lock()
+	if r.inflight[ip] {
+		r.mu.Unlock()
+		return
+	}
+	r.inflight[ip] = true
+	r.mu.Unlock()
+	go r.lookup(ip)
+}
+
+func (r *ReverseCache) lookup(ip string) {
+	defer func() {
+		r.mu.Lock()
+		delete(r.inflight, ip)
+		r.mu.Unlock()
+	}()
+	names, err := r.query(ip)
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if len(r.entries) >= r.max {
+		// Simple bound: drop the oldest quarter. A full LRU here would cost
+		// more than it saves, since the working set is usually small.
+		cutoff := time.Now().Add(-r.ttl)
+		removed := 0
+		for k, e := range r.entries {
+			if e.fetched.Before(cutoff) || removed >= r.max/4 {
+				delete(r.entries, k)
+				removed++
+			}
+		}
+		if removed == 0 {
+			for k := range r.entries {
+				delete(r.entries, k)
+				if removed >= r.max/4 {
+					break
+				}
+			}
+		}
+	}
+	r.entries[ip] = reverseEntry{
+		names:    names,
+		fetched:  time.Now(),
+		negative: err != nil || len(names) == 0,
+	}
+}
+
+func (r *ReverseCache) query(ip string) ([]string, error) {
+	if r.resolver == nil {
+		ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
+		defer cancel()
+		return net.DefaultResolver.LookupAddr(ctx, ip)
+	}
+	name, err := dns.ReverseAddr(ip)
+	if err != nil {
+		return nil, err
+	}
+	m := new(dns.Msg)
+	m.SetQuestion(name, dns.TypePTR)
+	m.RecursionDesired = true
+	resp, _, err := r.resolver.Exchange(m, r.names)
+	if err != nil {
+		return nil, err
+	}
+	var out []string
+	for _, rr := range resp.Answer {
+		if ptr, ok := rr.(*dns.PTR); ok {
+			out = append(out, ptr.Ptr)
+		}
+	}
+	return out, nil
+}
+
+// Len is the number of cached clients.
+func (r *ReverseCache) Len() int {
+	if r == nil {
+		return 0
+	}
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return len(r.entries)
+}
+
+// Flush empties the cache.
+func (r *ReverseCache) Flush() {
+	if r == nil {
+		return
+	}
+	r.mu.Lock()
+	r.entries = make(map[string]reverseEntry, 1024)
+	r.mu.Unlock()
+}