From e32dc646b17c22e2811649b4eae154d1ced3819e Mon Sep 17 00:00:00 2001
From: perp <perp@tcp.direct>
Date: Tue, 9 Jul 2024 19:17:34 +0100
Subject: [PATCH] Moved A & AAAA record lookups into one function

---
 v1/cmd/blink/main.go          |  4 ++
 v1/pkg/dns/aaaa.go            | 70 -----------------------------------
 v1/pkg/dns/dns.go             |  1 +
 v1/pkg/dns/{a.go => query.go} | 43 +++++++++++++++------
 v1/pkg/runner/runner.go       |  8 +---
 5 files changed, 39 insertions(+), 87 deletions(-)
 delete mode 100644 v1/pkg/dns/aaaa.go
 rename v1/pkg/dns/{a.go => query.go} (55%)

diff --git a/v1/cmd/blink/main.go b/v1/cmd/blink/main.go
index d2385b8..e1c0122 100644
--- a/v1/cmd/blink/main.go
+++ b/v1/cmd/blink/main.go
@@ -162,10 +162,14 @@ var rootCmd = &cobra.Command{
 		// Print warnings
 		if opts.UDP {
 			log.Warn().Msg("UDP is enabled")
+		} else {
+			log.Warn().Msg("TCP is enabled")
 		}
 
 		if opts.IPv6 {
 			log.Warn().Msg("IPv6 is enabled")
+		} else {
+			log.Warn().Msg("IPv4 is enabled")
 		}
 
 		// Wildcard enabled
diff --git a/v1/pkg/dns/aaaa.go b/v1/pkg/dns/aaaa.go
deleted file mode 100644
index 04e987f..0000000
--- a/v1/pkg/dns/aaaa.go
+++ /dev/null
@@ -1,70 +0,0 @@
-package dns
-
-import (
-	"errors"
-	"fmt"
-	"math/rand"
-	"strings"
-
-	"github.com/miekg/dns"
-)
-
-// Query AAAA record (ipv6)
-func (q *Query) AAAA() {
-	for {
-		// Create result
-		result := &Result{
-			Domain:    q.Domain,
-			Subdomain: q.Subdomain,
-		}
-
-		// Create message
-		message := new(dns.Msg)
-
-		// Set question
-		message.SetQuestion(dns.Fqdn(q.Subdomain), dns.TypeAAAA)
-
-		// Choose random resolver
-		resolver := q.Resolvers[rand.Intn(len(q.Resolvers))]
-
-		// Query response
-		resp, _, err := q.Client.Exchange(message, resolver)
-		if err != nil {
-			// Ratelimited
-			if strings.ContainsAny(err.Error(), "i/o timeout") {
-				continue
-			}
-
-			result.Error = err
-			q.Results <- result
-			return
-		}
-
-		// Store IPs
-		var ips []string
-
-		// Go through answers
-		for _, answer := range resp.Answer {
-			// Map record
-			record, ok := answer.(*dns.AAAA)
-			if !ok {
-				continue
-			}
-
-			// Append IPv6
-			ips = append(ips, record.AAAA.String())
-		}
-
-		// No IPs found
-		if len(ips) < 1 {
-			result.Error = errors.New(fmt.Sprintf("No IPv6s found for %s", q.Domain))
-			q.Results <- result
-			return
-		}
-
-		// Send result
-		result.IPv6 = ips
-		q.Results <- result
-		return
-	}
-}
diff --git a/v1/pkg/dns/dns.go b/v1/pkg/dns/dns.go
index 0e98a46..316af38 100644
--- a/v1/pkg/dns/dns.go
+++ b/v1/pkg/dns/dns.go
@@ -10,6 +10,7 @@ import (
 // DNS query
 type Query struct {
 	Client    *dns.Client    // DNS client
+	IPv6      bool           // IPv6 mode
 	Resolvers []string       // Target resolvers
 	Domain    string         // Target domain
 	Subdomain string         // Target subdomain
diff --git a/v1/pkg/dns/a.go b/v1/pkg/dns/query.go
similarity index 55%
rename from v1/pkg/dns/a.go
rename to v1/pkg/dns/query.go
index a0502d7..f0b9f0f 100644
--- a/v1/pkg/dns/a.go
+++ b/v1/pkg/dns/query.go
@@ -9,8 +9,8 @@ import (
 	"github.com/miekg/dns"
 )
 
-// Query A record (ipv4)
-func (q *Query) A() {
+// Query records
+func (q *Query) Lookup() {
 	for {
 		// Create result
 		result := &Result{
@@ -22,7 +22,12 @@ func (q *Query) A() {
 		message := new(dns.Msg)
 
 		// Set question
-		message.SetQuestion(dns.Fqdn(q.Subdomain), dns.TypeA)
+		switch q.IPv6 {
+		case false:
+			message.SetQuestion(dns.Fqdn(q.Subdomain), dns.TypeA)
+		case true:
+			message.SetQuestion(dns.Fqdn(q.Subdomain), dns.TypeAAAA)
+		}
 
 		// Choose random resolver
 		resolver := q.Resolvers[rand.Intn(len(q.Resolvers))]
@@ -46,24 +51,40 @@ func (q *Query) A() {
 		// Go through answers
 		for _, answer := range resp.Answer {
 			// Map record
-			record, ok := answer.(*dns.A)
-			if !ok {
-				continue
-			}
+			switch q.IPv6 {
+			case false:
+				record, ok := answer.(*dns.A)
+				if !ok {
+					continue
+				}
 
-			// Append IPv4
-			ips = append(ips, record.A.String())
+				// Append IPv4
+				ips = append(ips, record.A.String())
+			case true:
+				record, ok := answer.(*dns.AAAA)
+				if !ok {
+					continue
+				}
+
+				// Append IPv4
+				ips = append(ips, record.AAAA.String())
+			}
 		}
 
 		// No IPs found
 		if len(ips) < 1 {
-			result.Error = errors.New(fmt.Sprintf("No IPv4s found for %s", q.Domain))
+			result.Error = errors.New(fmt.Sprintf("No IPs found for %s", q.Domain))
 			q.Results <- result
 			return
 		}
 
 		// Send result
-		result.IPv4 = ips
+		switch q.IPv6 {
+		case false:
+			result.IPv4 = ips
+		case true:
+			result.IPv6 = ips
+		}
 		q.Results <- result
 		return
 	}
diff --git a/v1/pkg/runner/runner.go b/v1/pkg/runner/runner.go
index 9874513..a77c26d 100644
--- a/v1/pkg/runner/runner.go
+++ b/v1/pkg/runner/runner.go
@@ -63,6 +63,7 @@ func (r *Runner) Submit() {
 			// Create query
 			query := &dns.Query{
 				Client:    r.client,
+				IPv6:      r.options.IPv6,
 				Resolvers: r.options.Resolvers,
 				Domain:    domain,
 				Subdomain: word + "." + domain,
@@ -71,12 +72,7 @@ func (r *Runner) Submit() {
 
 			// Submit query
 			r.pool.Submit(func() {
-				// Check IPv6
-				if !r.options.IPv6 {
-					query.A()
-				} else {
-					query.AAAA()
-				}
+				query.Lookup()
 			})
 		}
 	}