From 2077cbb1d56e247274d2edb31d7262f7e094300c Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Fri, 3 Apr 2026 19:51:43 +0000
Subject: [PATCH] feat(dns): add discover callback-based cache refresh

Co-Authored-By: Virgil <virgil@lethean.io>
---
 service.go      | 33 ++++++++++++++++++++++++++-
 service_test.go | 59 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/service.go b/service.go
index ff65646..d47dbee 100644
--- a/service.go
+++ b/service.go
@@ -33,10 +33,12 @@ type Service struct {
 	records      map[string]NameRecords
 	reverseIndex map[string][]string
 	treeRoot     string
+	discoverer   func() (map[string]NameRecords, error)
 }
 
 type ServiceOptions struct {
-	Records map[string]NameRecords
+	Records    map[string]NameRecords
+	Discoverer func() (map[string]NameRecords, error)
 }
 
 func NewService(options ServiceOptions) *Service {
@@ -49,9 +51,38 @@ func NewService(options ServiceOptions) *Service {
 		records:      cached,
 		reverseIndex: buildReverseIndex(cached),
 		treeRoot:     treeRoot,
+		discoverer:   options.Discoverer,
 	}
 }
 
+func (service *Service) Discover() error {
+	discoverer := service.discoverer
+	if discoverer == nil {
+		return nil
+	}
+
+	discovered, err := discoverer()
+	if err != nil {
+		return err
+	}
+
+	cached := make(map[string]NameRecords, len(discovered))
+	for name, record := range discovered {
+		normalizedName := normalizeName(name)
+		if normalizedName == "" {
+			continue
+		}
+		cached[normalizedName] = record
+	}
+
+	service.mu.Lock()
+	defer service.mu.Unlock()
+	service.records = cached
+	service.reverseIndex = buildReverseIndex(service.records)
+	service.treeRoot = computeTreeRoot(service.records)
+	return nil
+}
+
 func (service *Service) SetRecord(name string, record NameRecords) {
 	service.mu.Lock()
 	defer service.mu.Unlock()
diff --git a/service_test.go b/service_test.go
index c3d06d8..b775660 100644
--- a/service_test.go
+++ b/service_test.go
@@ -184,3 +184,62 @@ func TestServiceHealthUsesDeterministicTreeRootAndUpdatesOnMutations(t *testing.
 		t.Fatalf("expected updated tree root after RemoveRecord, got %s", removedRoot)
 	}
 }
+
+func TestServiceDiscoverReplacesRecordsFromDiscoverer(t *testing.T) {
+	records := []map[string]NameRecords{
+		{
+			"gateway.charon.lthn": {
+				A: []string{"10.10.10.10"},
+			},
+			"*.lthn": {
+				A: []string{"10.0.0.1"},
+			},
+		},
+	}
+	index := 0
+	service := NewService(ServiceOptions{
+		Records: map[string]NameRecords{
+			"legacy.charon.lthn": {
+				A: []string{"10.11.11.11"},
+			},
+		},
+		Discoverer: func() (map[string]NameRecords, error) {
+			next := records[index%len(records)]
+			index++
+			return next, nil
+		},
+	})
+
+	if _, ok := service.Resolve("legacy.charon.lthn"); !ok {
+		t.Fatal("expected baseline record before discovery")
+	}
+
+	if err := service.Discover(); err != nil {
+		t.Fatalf("unexpected discover error: %v", err)
+	}
+
+	result, ok := service.Resolve("gateway.charon.lthn")
+	if !ok {
+		t.Fatal("expected discovered exact record")
+	}
+	if len(result.A) != 1 || result.A[0] != "10.10.10.10" {
+		t.Fatalf("unexpected discovered resolve result: %#v", result.A)
+	}
+	if _, ok := service.Resolve("legacy.unknown"); ok {
+		t.Fatal("expected replaced cache not to resolve via old record")
+	}
+}
+
+func TestServiceDiscoverReturnsNilWithoutDiscoverer(t *testing.T) {
+	service := NewService(ServiceOptions{
+		Records: map[string]NameRecords{
+			"gateway.charon.lthn": {
+				A: []string{"10.10.10.10"},
+			},
+		},
+	})
+
+	if err := service.Discover(); err != nil {
+		t.Fatalf("expected no error when discoverer is missing: %v", err)
+	}
+}