routewatch/internal/ptr/cache.go

// Package ptr looks up the PTR names of addresses over DNS-over-HTTPS. It
// keeps each answer in memory for as long as DNS allows, and in cache.db so
// that the answers survive a restart.
package ptr

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"maps"
	"net/http"
	"net/netip"
	"slices"
	"strings"
	"sync"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/bogon"
	"github.com/miekg/dns"
)

const (
	// flushInterval is how often new answers are written to cache.db.
	flushInterval = 5 * time.Second

	// sweepInterval is how often expired answers are deleted.
	sweepInterval = 10 * time.Minute

	// nameSeparator separates the names in a row of the hostnames table.
	nameSeparator = " "
)

// Params is what New needs.
type Params struct {
	// DB is the open cache.db.
	DB *sql.DB

	// Resolvers are the URLs of the DNS-over-HTTPS resolvers to ask, in order.
	Resolvers []string

	Logger *slog.Logger
}

// entry is one address's cached answer.
type entry struct {
	// names are the address's PTR names, nil when it has none.
	names   []string
	expires time.Time
}

// Cache answers PTR lookups from memory, and asks the resolvers for an address
// it holds no current answer for. Start loads it from cache.db.
type Cache struct {
	db        *sql.DB
	resolvers []string
	timeLimit time.Duration
	logger    *slog.Logger
	client    *http.Client
	now       func() time.Time

	// entries holds every answer by address. pending holds those that
	// cache.db does not have yet.
	mu      sync.RWMutex
	entries map[string]entry
	pending map[string]entry

	// stop ends the goroutine Start runs, and done is closed when it returns.
	stop context.CancelFunc
	done chan struct{}
}

// New returns an empty Cache. Start loads it.
func New(params Params) *Cache {
	return &Cache{
		db:        params.DB,
		resolvers: params.Resolvers,
		timeLimit: LookupTimeout,
		logger:    params.Logger.With("component", "ptr"),
		client:    &http.Client{},
		now:       time.Now,
		entries:   make(map[string]entry),
		pending:   make(map[string]entry),
		done:      make(chan struct{}),
	}
}

// LookupAddr returns the PTR names of addr, fully qualified, or none when it
// has none. A reserved address has none, and is not sent to the resolvers. A
// lookup takes at most LookupTimeout. A cached answer makes no outside request.
// It returns an error when no resolver answered, and nothing is cached then.
func (c *Cache) LookupAddr(ctx context.Context, addr string) ([]string, error) {
	parsed, err := netip.ParseAddr(addr)
	if err != nil {
		return nil, fmt.Errorf("parsing the address: %w", err)
	}

	parsed = parsed.Unmap()

	_, reserved := bogon.Prefix(netip.PrefixFrom(parsed, parsed.BitLen()))
	if reserved {
		return []string{}, nil
	}

	address := parsed.String()

	cached, ok := c.current(address)
	if ok {
		return slices.Clone(cached.names), nil
	}

	reverseName, err := dns.ReverseAddr(address)
	if err != nil {
		return nil, fmt.Errorf("building the reverse name: %w", err)
	}

	result, err := c.lookup(ctx, reverseName)
	if err != nil {
		return nil, err
	}

	if result.ttl > 0 {
		c.store(address, entry{names: result.names, expires: c.now().Add(result.ttl)})
	}

	return slices.Clone(result.names), nil
}

// Start loads the answers that have not expired from cache.db. It then starts
// the goroutine that writes new answers there and deletes expired ones. The
// goroutine runs until Stop: the end of ctx does not stop it.
func (c *Cache) Start(ctx context.Context) error {
	err := c.load(ctx)
	if err != nil {
		return err
	}

	c.logger.Info("PTR cache loaded", "entries", len(c.entries))

	ctx, c.stop = context.WithCancel(context.WithoutCancel(ctx))

	go c.run(ctx)

	return nil
}

// Stop ends the goroutine, which first writes the answers cache.db does not
// have yet, and waits until it has returned or ctx ends.
func (c *Cache) Stop(ctx context.Context) {
	c.stop()

	select {
	case <-c.done:
		c.logger.Info("PTR cache stopped")
	case <-ctx.Done():
		c.logger.Warn("Stop deadline passed before the PTR cache was written")
	}
}

// current returns the answer held for address, unless it has expired.
func (c *Cache) current(address string) (entry, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	found, ok := c.entries[address]
	if !ok || !c.now().Before(found.expires) {
		return entry{}, false
	}

	return found, true
}

// store keeps an answer in memory and queues it for cache.db.
func (c *Cache) store(address string, answer entry) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.entries[address] = answer
	c.pending[address] = answer
}

// run writes new answers and deletes expired ones until ctx ends.
func (c *Cache) run(ctx context.Context) {
	defer close(c.done)

	flushes := time.NewTicker(flushInterval)
	defer flushes.Stop()

	sweeps := time.NewTicker(sweepInterval)
	defer sweeps.Stop()

	for {
		select {
		case <-ctx.Done():
			c.flush(context.WithoutCancel(ctx))

			return
		case <-flushes.C:
			c.flush(ctx)
		case <-sweeps.C:
			c.sweep(ctx)
		}
	}
}

// load reads the rows that have not expired into memory.
func (c *Cache) load(ctx context.Context) error {
	rows, err := c.db.QueryContext(ctx,
		"SELECT address, names, expires_at FROM hostnames WHERE expires_at > ?",
		c.now().Unix())
	if err != nil {
		return fmt.Errorf("reading the PTR cache: %w", err)
	}

	defer func() { _ = rows.Close() }()

	for rows.Next() {
		var (
			address, names string
			expires        int64
		)

		err = rows.Scan(&address, &names, &expires)
		if err != nil {
			return fmt.Errorf("reading a PTR cache row: %w", err)
		}

		c.entries[address] = entry{
			names:   strings.Fields(names),
			expires: time.Unix(expires, 0).UTC(),
		}
	}

	err = rows.Err()
	if err != nil {
		return fmt.Errorf("reading the PTR cache: %w", err)
	}

	return nil
}

// flush writes the pending answers to cache.db. Those it cannot write stay
// pending, unless a newer answer for the address has arrived meanwhile.
func (c *Cache) flush(ctx context.Context) {
	c.mu.Lock()
	batch := c.pending
	c.pending = make(map[string]entry)
	c.mu.Unlock()

	if len(batch) == 0 {
		return
	}

	err := c.writeRows(ctx, batch)
	if err == nil {
		return
	}

	c.logger.Error("Failed to write the PTR cache", "error", err)

	c.mu.Lock()
	defer c.mu.Unlock()

	for address, answer := range batch {
		_, newer := c.pending[address]
		if !newer {
			c.pending[address] = answer
		}
	}
}

// writeRows writes batch to cache.db in one transaction.
func (c *Cache) writeRows(ctx context.Context, batch map[string]entry) error {
	tx, err := c.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}

	defer func() { _ = tx.Rollback() }()

	for address, answer := range batch {
		_, err = tx.ExecContext(ctx,
			"INSERT OR REPLACE INTO hostnames (address, names, expires_at) "+
				"VALUES (?, ?, ?)",
			address,
			strings.Join(answer.names, nameSeparator),
			answer.expires.Unix())
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

// sweep deletes the expired answers from memory and cache.db.
func (c *Cache) sweep(ctx context.Context) {
	now := c.now()
	expired := func(_ string, answer entry) bool {
		return !now.Before(answer.expires)
	}

	c.mu.Lock()
	maps.DeleteFunc(c.entries, expired)
	maps.DeleteFunc(c.pending, expired)
	c.mu.Unlock()

	_, err := c.db.ExecContext(ctx,
		"DELETE FROM hostnames WHERE expires_at <= ?", now.Unix())
	if err != nil {
		c.logger.Error("Failed to delete expired PTR cache rows", "error", err)
	}
}