extract.go

package main

import (
	"bufio"
	"fmt"
	"io"
	"log/slog"
	"sort"
	"strings"
)

// hostrec holds the currently-best hostname chosen for a given ip, along with
// the sigil it was defined with (needed to apply the preference rules as more
// records for the same ip are seen).
type hostrec struct {
	hostname string
	sigil    byte // '=' or '+'
}

// supersededBy reports whether cand should replace the existing hostrec for an
// ip. The rules:
//   - a '=' hostname beats a '+' hostname; among '=' entries the first seen wins
//   - a '+' hostname beats another '+' only if it is strictly shorter
func (existing hostrec) supersededBy(cand hostrec) bool {
	if cand.sigil == '=' {
		// '=' beats '+', but an existing '=' keeps its place (first '=' wins).
		return existing.sigil != '='
	}
	// cand is '+': it can never displace an existing '='.
	if existing.sigil == '=' {
		return false
	}
	return len(cand.hostname) < len(existing.hostname)
}

// matchesDomain reports whether hostname belongs to domain: either an exact
// match (the apex) or a subdomain of it. An empty domain matches everything.
// So example.com matches "example.com" and "foo.example.com", but not
// "anotherexample.com".
func matchesDomain(hostname, domain string) bool {
	if domain == "" {
		return true
	}
	return hostname == domain || strings.HasSuffix(hostname, "."+domain)
}

// extract reads tinydns data line by line and builds a map of ip -> best
// hostname. Lines not beginning with '=' or '+' are skipped silently; lines
// with the right sigil but not exactly three colon-separated fields are warned
// about and discarded. Wildcard hostnames (beginning with "*.") are never
// treated as candidates. If domain is non-empty, only hostnames belonging to
// that domain (see matchesDomain) are considered.
func extract(r io.Reader, domain string) (map[string]hostrec, error) {
	m := make(map[string]hostrec)

	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		line := scanner.Text()
		if line == "" {
			continue
		}

		sigil := line[0]
		if sigil != '=' && sigil != '+' {
			continue
		}

		// Format: {sigil}{hostname}:{ip}:{ttl}
		fields := strings.Split(line[1:], ":")
		if len(fields) != 3 {
			slog.Warn("discarding line without exactly three fields",
				"line", line, "fields", len(fields))
			continue
		}

		// Wildcard hostnames (e.g. *.example.com) are never candidates.
		if strings.HasPrefix(fields[0], "*.") {
			continue
		}

		if !matchesDomain(fields[0], domain) {
			continue
		}

		cand := hostrec{hostname: fields[0], sigil: sigil}
		ip := fields[1]
		if existing, ok := m[ip]; !ok || existing.supersededBy(cand) {
			m[ip] = cand
		}
	}

	return m, scanner.Err()
}

// applyOverrides reads "hostname,ip" records from r (the same shape as the
// input) and relocates matching entries in m to their override ip.
//
// Motivation: the advertised/public ip for a host is sometimes not configured
// on the server itself but on a gateway or bastion that fronts it. In that case
// the public ip can't be validated by talking to the server directly. An
// override lets us swap the public ip for a server ip that we *can* reach and
// check, while keeping the hostname the record is emitted under.
//
// The map is keyed by ip, so each override is resolved by hostname: for every
// override record we look for the entry whose hostname matches. Exactly one
// match is required — zero matches (unknown hostname) or more than one
// (ambiguous hostname) are both fatal errors, and applyOverrides returns
// leaving m partially modified. When the single match is found, its record is
// re-keyed under the override ip (keeping the original hostname and sigil) and
// the original entry is removed.
//
// When domain is non-empty, override records whose hostname does not belong to
// domain (see matchesDomain) are skipped: those hostnames were filtered out of
// the input by extract, so treating them as "not found" would turn a shared,
// multi-domain overrides file into a guaranteed failure under -d.
//
// Blank lines and lines beginning with '#' are ignored. A line without exactly
// two comma-separated non-empty fields is also a fatal error.
func applyOverrides(r io.Reader, m map[string]hostrec, domain string) error {
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		fields := strings.Split(line, ",")
		if len(fields) != 2 {
			return fmt.Errorf("override %q: expected 'hostname,ip'", line)
		}
		hostname := strings.TrimSpace(fields[0])
		ip := strings.TrimSpace(fields[1])
		if hostname == "" || ip == "" {
			return fmt.Errorf("override %q: empty hostname or ip", line)
		}

		// Out-of-domain overrides were never in m (extract filtered them out),
		// so skip them rather than failing the run on a "not found".
		if !matchesDomain(hostname, domain) {
			continue
		}

		oldIP, rec, matches := "", hostrec{}, 0
		for eip, erec := range m {
			if erec.hostname == hostname {
				oldIP, rec = eip, erec
				matches++
			}
		}
		if matches == 0 {
			return fmt.Errorf("override %q: hostname %q not found in input", line, hostname)
		}
		if matches > 1 {
			return fmt.Errorf("override %q: hostname %q is ambiguous, found %d entries", line, hostname, matches)
		}

		// Re-key under the override ip. Delete first so that an override whose
		// ip is unchanged still leaves the entry in place.
		delete(m, oldIP)
		m[ip] = rec
	}

	return scanner.Err()
}

// formatResults renders the map as "{hostname},{ip}" lines, sorted for stable
// output (the spec permits any order).
func formatResults(m map[string]hostrec) []string {
	out := make([]string, 0, len(m))
	for ip, rec := range m {
		out = append(out, rec.hostname+","+ip)
	}
	sort.Strings(out)
	return out
}