From 739a68f2922f5ca3172972b922ca4d50330c6770 Mon Sep 17 00:00:00 2001
From: Seraphim Strub <sst@rievo.net>
Date: Wed, 17 Jul 2024 20:30:46 +0000
Subject: [PATCH] retry over all nameservers

---
 cmd/domaincheckbot/dns/domain.go | 109 ++++++++++++++++++-------------
 1 file changed, 62 insertions(+), 47 deletions(-)

diff --git a/cmd/domaincheckbot/dns/domain.go b/cmd/domaincheckbot/dns/domain.go
index 3126793..9f18dc5 100644
--- a/cmd/domaincheckbot/dns/domain.go
+++ b/cmd/domaincheckbot/dns/domain.go
@@ -13,63 +13,78 @@ type Domain struct {
 	NS   []string
 }
 
-func CheckDomain(domain, tld string) Domain {
+func CheckDomain(d, tld string) Domain {
 	nameservers := []string{}
-	// Create a new DNS message
-	m := new(dns.Msg)
-	// Set the question section of the message with the domain and query type
-	m.SetQuestion(dns.Fqdn(fmt.Sprintf("%v.%v", domain, tld)), dns.TypeNS)
+	fqdn := fmt.Sprintf("%v.%v", d, tld)
+
+	tldServers, err := getNameserverForTLD(tld)
+	if err != nil {
+		return Domain{fqdn, nameservers}
+	}
+
+	for _, tldServer := range tldServers {
+		nameservers, err = queryTLDNamserver(fqdn, tldServer)
+		if err == nil {
+			break
+		}
+	}
 
-	// Create a DNS client
-	c := new(dns.Client)
-	// Define the .com TLD server
-	ns, err := net.LookupNS(tld)
 	if err != nil {
 		return Domain{
-			Name: domain + "." + tld,
+			Name: fqdn,
 			NS:   nameservers,
 		}
 	}
-	var tldServer string
-	if len(ns) > 0 {
-		tldServer = strings.Trim(ns[0].Host, ".") + ":53"
-	} else {
-		//log.Fatalf("Could not resolve %v", tld)
-		return Domain{
-			Name: domain + "." + tld,
-			NS:   nameservers,
-		}
-	}
-
-	// Send the DNS query to the TLD server
-	r, _, err := c.Exchange(m, tldServer)
-	if err != nil {
-		//log.Fatalf("Failed to query the TLD server: %v", err)
-		return Domain{
-			Name: domain + "." + tld,
-			NS:   nameservers,
-		}
-	}
-
-	// Check for response errors
-	if r.Rcode != dns.RcodeSuccess {
-		//log.Fatalf("Failed to get a valid response: %v", r.Rcode)
-		return Domain{
-			Name: domain + "." + tld,
-			NS:   nameservers,
-		}
-	}
-
-	// Print the nameservers from the response
-	for _, ans := range r.Ns {
-		if ns, ok := ans.(*dns.NS); ok {
-			nameservers = append(nameservers, strings.ToLower(ns.Ns))
-		}
-	}
 
 	slices.Sort(nameservers)
 	return Domain{
-		Name: domain + "." + tld,
+		Name: fqdn,
 		NS:   nameservers,
 	}
 }
+
+func getNameserverForTLD(tld string) ([]string, error) {
+	ns, err := net.LookupNS(tld)
+	if err != nil {
+		return nil, err
+	}
+	var tldServer []string
+	if len(ns) > 0 {
+		for _, n := range ns {
+			tldServer = append(tldServer, strings.Trim(n.Host, ".")+":53")
+		}
+	} else {
+		return nil, fmt.Errorf("no nameserver found for TLD %v", tld)
+	}
+	return tldServer, nil
+}
+
+func queryTLDNamserver(fqdn, tldServer string) ([]string, error) {
+	// Create a new DNS message
+	m := new(dns.Msg)
+	// Set the question section of the message with the fqdn and query type
+	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeNS)
+
+	// Create a DNS client
+	c := new(dns.Client)
+
+	// Send the DNS query to the TLD server
+	r, _, err := c.Exchange(m, tldServer)
+	if err != nil {
+		return nil, fmt.Errorf("dns query error: %v", err)
+	}
+
+	// Check for response errors
+	if r.Rcode != dns.RcodeSuccess {
+		return nil, fmt.Errorf("failed to get a valid response for %v: %v", fqdn, r.Rcode)
+	}
+
+	// Print the nameservers from the response
+	nameservers := []string{}
+	for _, ans := range r.Ns {
+		if ns, ok := ans.(*dns.NS); ok {
+			nameservers = append(nameservers, strings.ToLower(ns.Ns))
+		}
+	}
+	return nameservers, nil
+}