From d60f9a6b769b5c0156c100c90d743eded0098614 Mon Sep 17 00:00:00 2001
From: Gerard Nguyen <gerard@replicated.com>
Date: Mon, 30 Sep 2024 11:22:10 +1000
Subject: [PATCH] change resolvedFromSearch content (#1629)

---
 pkg/collect/host_dns.go      | 12 +++++++++++-
 pkg/collect/host_dns_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)
 create mode 100644 pkg/collect/host_dns_test.go

diff --git a/pkg/collect/host_dns.go b/pkg/collect/host_dns.go
index b362a802c..b8aad47dd 100644
--- a/pkg/collect/host_dns.go
+++ b/pkg/collect/host_dns.go
@@ -177,7 +177,7 @@ func queryDNS(name, query, server string) DNSEntry {
 
 		// remember the search domain that resolved the query
 		// e.g. foo.test.com -> test.com
-		entry.Search = strings.Replace(query, name, "", 1)
+		entry.Search = extractSearchFromFQDN(query, name)
 
 		// populate record detail
 		switch rec {
@@ -207,3 +207,13 @@ func queryDNS(name, query, server string) DNSEntry {
 func (c *CollectHostDNS) RemoteCollect(progressChan chan<- interface{}) (map[string][]byte, error) {
 	return nil, ErrRemoteCollectorNotImplemented
 }
+
+func extractSearchFromFQDN(fqdn, name string) string {
+	// no search domain
+	if fqdn == name {
+		return ""
+	}
+	search := strings.TrimPrefix(fqdn, name+".") // remove name
+	search = strings.TrimSuffix(search, ".")     // remove root dot
+	return search
+}
diff --git a/pkg/collect/host_dns_test.go b/pkg/collect/host_dns_test.go
new file mode 100644
index 000000000..9cd14fd33
--- /dev/null
+++ b/pkg/collect/host_dns_test.go
@@ -0,0 +1,26 @@
+package collect
+
+import (
+	"testing"
+)
+
+func TestExtractSearchFromFQDN(t *testing.T) {
+	tests := []struct {
+		fqdn     string
+		name     string
+		expected string
+	}{
+		{"foo.com.", "foo.com", ""},
+		{"bar.com", "bar.com", ""},
+		{"*.foo.testcluster.net.", "*", "foo.testcluster.net"},
+	}
+
+	for _, test := range tests {
+		t.Run(test.fqdn, func(t *testing.T) {
+			result := extractSearchFromFQDN(test.fqdn, test.name)
+			if result != test.expected {
+				t.Errorf("extractSearchFromFQDN(%q, %q) = %q; want %q", test.fqdn, test.name, result, test.expected)
+			}
+		})
+	}
+}