refactor(dns): centralize tree root selection

Co-Authored-By: Virgil <virgil@lethean.io>

service.go

@@ -1069,6 +1069,16 @@ func (service *Service) refreshDerivedStateLocked() {
 	service.zoneApex = computeZoneApex(service.records)
 }
 
+func (service *Service) currentTreeRootLocked() string {
+	if service == nil {
+		return ""
+	}
+	if service.chainTreeRoot != "" {
+		return service.chainTreeRoot
+	}
+	return service.treeRoot
+}
+
 // ResolveAddresses returns A and AAAA values merged into one address list.
 //
 //	addresses, ok := service.ResolveAddresses("gateway.charon.lthn")
@@ -1182,18 +1192,29 @@ func (service *Service) Health() HealthResult {
 	service.mutex.RLock()
 	defer service.mutex.RUnlock()
 
-	treeRoot := service.treeRoot
-	if service.chainTreeRoot != "" {
-		treeRoot = service.chainTreeRoot
-	}
-
 	return HealthResult{
 		Status:      "ready",
 		NamesCached: len(service.records),
-		TreeRoot:    treeRoot,
+		TreeRoot:    service.currentTreeRootLocked(),
 	}
 }
 
+// CurrentTreeRoot returns the active tree root used by health and snapshot data.
+//
+//	treeRoot := service.CurrentTreeRoot()
+//	fmt.Println(treeRoot)
+func (service *Service) CurrentTreeRoot() string {
+	if service == nil {
+		return ""
+	}
+	service.pruneExpiredRecords()
+
+	service.mutex.RLock()
+	defer service.mutex.RUnlock()
+
+	return service.currentTreeRootLocked()
+}
+
 // Snapshot returns a structured view of the service state.
 //
 //	snapshot := service.Snapshot()
@@ -1210,11 +1231,6 @@ func (service *Service) Snapshot() ServiceDescription {
 	service.mutex.RLock()
 	defer service.mutex.RUnlock()
 
-	treeRoot := service.treeRoot
-	if service.chainTreeRoot != "" {
-		treeRoot = service.chainTreeRoot
-	}
-
 	dnsPort := service.dnsPort
 	if dnsPort <= 0 {
 		dnsPort = DefaultDNSPort
@@ -1229,7 +1245,7 @@ func (service *Service) Snapshot() ServiceDescription {
 		Status:     "ready",
 		Records:    len(service.records),
 		ZoneApex:   service.zoneApex,
-		TreeRoot:   treeRoot,
+		TreeRoot:   service.currentTreeRootLocked(),
 		DNSPort:    dnsPort,
 		HealthPort: httpPort,
 		HTTPPort:   httpPort,

service_test.go

@@ -3857,6 +3857,31 @@ func TestServiceMethodsHandleNilReceiverWithoutPanicking(t *testing.T) {
 	if got := service.Health().Status; got != "not_ready" {
 		t.Fatalf("expected nil service health status \"not_ready\", got %q", got)
 	}
+	if got := service.CurrentTreeRoot(); got != "" {
+		t.Fatalf("expected nil service current tree root to be empty, got %q", got)
+	}
+}
+
+func TestServiceCurrentTreeRootPrefersChainRootWhenPresent(t *testing.T) {
+	service := NewService(ServiceOptions{
+		Records: map[string]NameRecords{
+			"gateway.charon.lthn": {
+				A: []string{"10.10.10.10"},
+			},
+		},
+	})
+
+	if got := service.CurrentTreeRoot(); got == "" {
+		t.Fatal("expected current tree root to be populated from cached records")
+	}
+
+	service.mutex.Lock()
+	service.chainTreeRoot = "chain-root-1"
+	service.mutex.Unlock()
+
+	if got := service.CurrentTreeRoot(); got != "chain-root-1" {
+		t.Fatalf("expected current tree root to prefer chain root, got %q", got)
+	}
 }
 
 func TestServiceServeReturnsErrorOnNilReceiver(t *testing.T) {