housecarl/internal/names/names.go

// Package names finds the names of addresses by reverse lookups, for
// showing them beside the addresses (connections, the log, quotas,
// traffic per device). Names are cached and never stored: statistics stay
// keyed by address.
package names

import (
	"context"
	"errors"
	"net"
	"net/netip"
	"strings"
	"sync"
	"time"
)

// Lookups running at once, and the most names kept.
const (
	parallel = 16
	maxCache = 20000
)

// Resolver looks up and caches addresses' names.
type Resolver struct {
	lookup func(ctx context.Context, addr string) ([]string, error)
	// ttl is how long a name (or its absence) is kept; a lookup that
	// failed (a timeout, a server error) is tried again after retry.
	ttl, retry time.Duration
	timeout    time.Duration
	now        func() time.Time

	mu      sync.Mutex
	cache   map[netip.Addr]entry
	pending map[netip.Addr]chan struct{}
	sem     chan struct{}
}

type entry struct {
	name    string
	expires time.Time
}

// New makes a resolver asking server ("10.10.0.1:53"; "" for the
// system's, /etc/hosts then /etc/resolv.conf), keeping names for ttl.
func New(server string, ttl time.Duration) *Resolver {
	res := &net.Resolver{PreferGo: true}
	if server != "" {
		res.Dial = func(ctx context.Context, network, _ string) (net.Conn, error) {
			var d net.Dialer
			return d.DialContext(ctx, network, server)
		}
	}
	return NewFunc(res.LookupAddr, ttl)
}

// NewFunc makes a resolver with its own lookup (tests).
func NewFunc(lookup func(ctx context.Context, addr string) ([]string, error), ttl time.Duration) *Resolver {
	return &Resolver{lookup: lookup, ttl: ttl, retry: time.Minute, timeout: 2 * time.Second, now: time.Now,
		cache: map[netip.Addr]entry{}, pending: map[netip.Addr]chan struct{}{}, sem: make(chan struct{}, parallel)}
}

// Names returns the names of addrs: those cached, and those looked up
// within wait. Lookups still running then carry on and are cached for
// the next call. An address without a name is left out.
func (r *Resolver) Names(ctx context.Context, addrs []netip.Addr, wait time.Duration) map[netip.Addr]string {
	out := map[netip.Addr]string{}
	var waiting []netip.Addr
	var done []chan struct{}
	now := r.now()
	r.mu.Lock()
	for _, a := range addrs {
		if e, ok := r.cache[a]; ok && now.Before(e.expires) {
			if e.name != "" {
				out[a] = e.name
			}
			continue
		}
		ch, ok := r.pending[a]
		if !ok {
			ch = make(chan struct{})
			r.pending[a] = ch
			go r.resolve(a, ch)
		}
		waiting = append(waiting, a)
		done = append(done, ch)
	}
	r.mu.Unlock()
	if len(waiting) == 0 {
		return out
	}
	t := time.NewTimer(wait)
	defer t.Stop()
	for _, ch := range done {
		select {
		case <-ch:
		case <-t.C:
			return r.collect(out, waiting)
		case <-ctx.Done():
			return r.collect(out, waiting)
		}
	}
	return r.collect(out, waiting)
}

// collect adds what the lookups found for addrs.
func (r *Resolver) collect(out map[netip.Addr]string, addrs []netip.Addr) map[netip.Addr]string {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, a := range addrs {
		if e, ok := r.cache[a]; ok && e.name != "" {
			out[a] = e.name
		}
	}
	return out
}

func (r *Resolver) resolve(a netip.Addr, done chan struct{}) {
	r.sem <- struct{}{}
	ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
	ns, err := r.lookup(ctx, a.String())
	cancel()
	<-r.sem
	e := entry{expires: r.now().Add(r.ttl)}
	var dnsErr *net.DNSError
	switch {
	case err == nil && len(ns) > 0:
		e.name = strings.TrimSuffix(ns[0], ".")
	case err != nil && !(errors.As(err, &dnsErr) && dnsErr.IsNotFound):
		e.expires = r.now().Add(r.retry)
	}
	r.mu.Lock()
	if len(r.cache) >= maxCache {
		r.prune()
	}
	r.cache[a] = e
	delete(r.pending, a)
	r.mu.Unlock()
	close(done)
}

// prune drops expired names, or all of them when none has expired.
func (r *Resolver) prune() {
	now := r.now()
	for a, e := range r.cache {
		if !now.Before(e.expires) {
			delete(r.cache, a)
		}
	}
	if len(r.cache) >= maxCache {
		clear(r.cache)
	}
}