feat: implement DNS server fingerprinting (HAN-384) (#10)

internal/fingerprint/fingerprint.go

@@ -1 +1,138 @@
 package fingerprint
+
+import (
+	"context"
+	"net"
+	"sync"
+	"time"
+
+	miekgdns "github.com/miekg/dns"
+)
+
+const defaultTimeout = 2 * time.Second
+
+// Fingerprinter queries DNS servers for their software version via the
+// version.bind CHAOS TXT query. Results are cached per server IP.
+type Fingerprinter struct {
+	mu       sync.Mutex
+	cache    map[string]string
+	timeout  time.Duration
+	exchange func(ctx context.Context, addr string, m *miekgdns.Msg) (*miekgdns.Msg, error)
+}
+
+// New returns a Fingerprinter with a 2-second per-query timeout.
+func New() *Fingerprinter {
+	return NewWithTimeout(defaultTimeout)
+}
+
+// NewWithTimeout returns a Fingerprinter using the given per-query timeout.
+func NewWithTimeout(timeout time.Duration) *Fingerprinter {
+	return &Fingerprinter{
+		cache:   make(map[string]string),
+		timeout: timeout,
+	}
+}
+
+// Query returns the version string for ip, or "" if the server doesn't
+// respond or doesn't support the version.bind CHAOS query.
+// Results are cached: subsequent calls for the same IP return immediately.
+func (f *Fingerprinter) Query(ctx context.Context, ip net.IP) string {
+	key := ip.String()
+
+	f.mu.Lock()
+	if v, ok := f.cache[key]; ok {
+		f.mu.Unlock()
+		return v
+	}
+	f.mu.Unlock()
+
+	version := f.probe(ctx, ip)
+
+	f.mu.Lock()
+	f.cache[key] = version
+	f.mu.Unlock()
+
+	return version
+}
+
+// FingerprintAll queries all ips concurrently and returns a map of
+// IP string → version string. IPs that don't respond map to "".
+// Already-cached IPs are returned from cache without a network round-trip.
+func (f *Fingerprinter) FingerprintAll(ctx context.Context, ips []net.IP) map[string]string {
+	results := make(map[string]string, len(ips))
+
+	var (
+		wg      sync.WaitGroup
+		mu      sync.Mutex
+		toQuery []net.IP
+	)
+
+	f.mu.Lock()
+	for _, ip := range ips {
+		key := ip.String()
+		if v, ok := f.cache[key]; ok {
+			results[key] = v
+		} else {
+			toQuery = append(toQuery, ip)
+		}
+	}
+	f.mu.Unlock()
+
+	for _, ip := range toQuery {
+		wg.Add(1)
+		go func(ip net.IP) {
+			defer wg.Done()
+			version := f.probe(ctx, ip)
+			key := ip.String()
+
+			f.mu.Lock()
+			f.cache[key] = version
+			f.mu.Unlock()
+
+			mu.Lock()
+			results[key] = version
+			mu.Unlock()
+		}(ip)
+	}
+
+	wg.Wait()
+	return results
+}
+
+// probe sends a version.bind CHAOS TXT query and returns the version string,
+// or "" on any error or non-success response.
+func (f *Fingerprinter) probe(ctx context.Context, ip net.IP) string {
+	m := new(miekgdns.Msg)
+	m.SetQuestion("version.bind.", miekgdns.TypeTXT)
+	m.Question[0].Qclass = miekgdns.ClassCHAOS
+	m.RecursionDesired = false
+
+	target := net.JoinHostPort(ip.String(), "53")
+
+	queryCtx, cancel := context.WithTimeout(ctx, f.timeout)
+	defer cancel()
+
+	var resp *miekgdns.Msg
+	var err error
+
+	if f.exchange != nil {
+		resp, err = f.exchange(queryCtx, target, m)
+	} else {
+		client := &miekgdns.Client{
+			Net:     "udp",
+			Timeout: f.timeout,
+		}
+		resp, _, err = client.ExchangeContext(queryCtx, m, target)
+	}
+
+	if err != nil || resp == nil || resp.Rcode != miekgdns.RcodeSuccess {
+		return ""
+	}
+
+	for _, rr := range resp.Answer {
+		if txt, ok := rr.(*miekgdns.TXT); ok && len(txt.Txt) > 0 {
+			return txt.Txt[0]
+		}
+	}
+	return ""
+}