Code refactor and cleanups

cache/cache.go

@@ -26,74 +26,80 @@ import (
 	"strings"
 	"sync"
 
+	"filippo.io/mostly-harmless/cryptosource"
+
 	"github.com/miekg/dns"
-	"github.com/sean-/seed"
 )
 
 type Cache struct {
 	items map[string][]*dns.A
 	m     sync.RWMutex
+	r     *rand.Rand
 }
 
-func init() {
-	_, _ = seed.Init()
-}
+const (
+	defaultSize = 100
+)
 
-// New returns a new and initialized A dns cache
+// New returns a new and initialized A dns cache.
 func New() *Cache {
-	return &Cache{items: make(map[string][]*dns.A)}
+	return &Cache{
+		items: make(map[string][]*dns.A, defaultSize),
+		r:     rand.New(cryptosource.New()), //nolint:gosec
+	}
 }
 
-// Set adds a single item to the cache, replacing all existing items
+// Set adds a single item to the cache, replacing all existing items.
 func (c *Cache) Set(qname string, rr *dns.A) {
 	c.m.Lock()
+	defer c.m.Unlock()
+
 	c.items[strings.ToLower(qname)] = append([]*dns.A(nil), rr)
-	c.m.Unlock()
 }
 
-// Add an item to the cache only if item doesn't exist for a given key
+// Add an item to the cache only if item doesn't exist for a given key.
 func (c *Cache) Add(qname string, rr *dns.A) {
 	c.m.Lock()
+	defer c.m.Unlock()
 
 	// Check if rr is duplicate and skip adding if true
 	if v, ok := c.items[strings.ToLower(qname)]; ok {
 		for _, rr1 := range v {
 			if dns.IsDuplicate(rr, rr1) {
-				c.m.Unlock()
 				return
 			}
 		}
 	}
 
 	c.items[strings.ToLower(qname)] = append(c.items[strings.ToLower(qname)], rr)
-	c.m.Unlock()
 }
 
-// Get an item (list) from the cache, returning an item and a boolean indicating if the key has been found
+// Get an item (list) from the cache, returning an item and a boolean indicating if the key has been found.
 func (c *Cache) Get(qname string) ([]*dns.A, bool) {
 	c.m.RLock()
+	defer c.m.RUnlock()
+
 	v, ok := c.items[strings.ToLower(qname)]
 	if !ok {
-		c.m.RUnlock()
 		return nil, false
 	}
 
-	c.m.RUnlock()
 	return v, true
 }
 
-// Get a randomized item (single item) from the cache, returning also a boolean indicating if the key has been found
+// GetRand gets a randomized item (single item) from the cache, returning also a boolean indicating if the key has been
+// found.
 func (c *Cache) GetRand(qname string) (*dns.A, bool) {
 	c.m.RLock()
+	defer c.m.RUnlock()
+
 	v, ok := c.items[strings.ToLower(qname)]
 	if !ok {
-		c.m.RUnlock()
 		return nil, false
 	}
 
-	c.m.RUnlock()
-
 	// Randomized item for a given key
-	n := rand.Int() % len(v)
+	n := c.r.Int() % len(v)
+
 	return v[n], true
 }

go.mod

@@ -3,9 +3,9 @@ module github.com/dkorunic/dnstrace
 go 1.19
 
 require (
+	filippo.io/mostly-harmless/cryptosource v0.0.0-20221105134001-969191f8eda1
 	github.com/fatih/color v1.13.0
 	github.com/miekg/dns v1.1.50
-	github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529
 )
 
 require (

go.sum

@@ -1,3 +1,5 @@
+filippo.io/mostly-harmless/cryptosource v0.0.0-20221105134001-969191f8eda1 h1:AJbN8AhtHIwACnnRiY9nZf8JZg0b6IsvXWDIIPkrk8c=
+filippo.io/mostly-harmless/cryptosource v0.0.0-20221105134001-969191f8eda1/go.mod h1:hIhrS2FFqXOAlsXLtU3FZ3R5TMD47Mhh/8UQ+vBylQk=
 github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
@@ -9,8 +11,6 @@ github.com/mattn/go-isatty v0.0.16 h1:bq3VjFmv/sOjHtdEhmkEV4x1AJtvUvOJ2PFAZ5+peK
 github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
 github.com/miekg/dns v1.1.50 h1:DQUfb9uc6smULcREF09Uc+/Gd46YWqJd5DbpPE9xkcA=
 github.com/miekg/dns v1.1.50/go.mod h1:e3IlAVfNqAllflbibAZEWOXOQ+Ynzk/dDozDxY7XnME=
-github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529 h1:nn5Wsu0esKSJiIVhscUtVbo7ada43DJhG55ua/hjS5I=
-github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc=
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=

hints/hints.go

@@ -24,7 +24,7 @@ package hints
 import (
 	"math/rand"
 
-	"github.com/sean-/seed"
+	"filippo.io/mostly-harmless/cryptosource"
 )
 
 type Root struct {
@@ -35,29 +35,29 @@ type Root struct {
 
 type Hints struct {
 	hints []Root
-}
-
-func init() {
-	seed.Init()
+	r     *rand.Rand
 }
 
 // New returns a new and initialized root nameserver hints db
 func New() *Hints {
-	return &Hints{hints: []Root{
-		{"a.root-servers.net.", "198.41.0.4", "2001:503:ba3e::2:30"},
-		{"b.root-servers.net.", "192.228.79.201", "2001:478:65::53"},
-		{"c.root-servers.net.", "192.33.4.12", "2001:500:2::c"},
-		{"d.root-servers.net.", "199.7.91.13", "2001:500:2d::d"},
-		{"e.root-servers.net.", "192.203.230.10", "NASA"},
-		{"f.root-servers.net.", "192.5.5.241", "2001:500:2f::f"},
-		{"g.root-servers.net.", "192.112.36.4", "U.S."},
-		{"h.root-servers.net.", "128.63.2.53", "2001:500:1::803f:235"},
-		{"i.root-servers.net.", "192.36.148.17", "2001:7FE::53"},
-		{"j.root-servers.net.", "192.58.128.30", "2001:503:c27::2:30"},
-		{"k.root-servers.net.", "193.0.14.129", "2001:7fd::1"},
-		{"l.root-servers.net.", "199.7.83.42", "2001:500:3::42"},
-		{"m.root-servers.net.", "202.12.27.33", "2001:dc3::35"},
-	}}
+	return &Hints{
+		hints: []Root{
+			{"a.root-servers.net.", "198.41.0.4", "2001:503:ba3e::2:30"},
+			{"b.root-servers.net.", "192.228.79.201", "2001:478:65::53"},
+			{"c.root-servers.net.", "192.33.4.12", "2001:500:2::c"},
+			{"d.root-servers.net.", "199.7.91.13", "2001:500:2d::d"},
+			{"e.root-servers.net.", "192.203.230.10", "NASA"},
+			{"f.root-servers.net.", "192.5.5.241", "2001:500:2f::f"},
+			{"g.root-servers.net.", "192.112.36.4", "U.S."},
+			{"h.root-servers.net.", "128.63.2.53", "2001:500:1::803f:235"},
+			{"i.root-servers.net.", "192.36.148.17", "2001:7FE::53"},
+			{"j.root-servers.net.", "192.58.128.30", "2001:503:c27::2:30"},
+			{"k.root-servers.net.", "193.0.14.129", "2001:7fd::1"},
+			{"l.root-servers.net.", "199.7.83.42", "2001:500:3::42"},
+			{"m.root-servers.net.", "202.12.27.33", "2001:dc3::35"},
+		},
+		r: rand.New(cryptosource.New()), //nolint:gosec
+	}
 }
 
 // Get returns an array of root nameserver hints
@@ -67,6 +67,7 @@ func (h *Hints) Get() []Root {
 
 // GetRand returns a randomized item (root nameserver) from root hints
 func (h *Hints) GetRand() (string, string, string) {
-	n := rand.Int() % len(h.hints)
+	n := h.r.Int() % len(h.hints)
+
 	return h.hints[n].Name, h.hints[n].IPv4Address, h.hints[n].IPv6Address
 }

main.go

@@ -40,7 +40,7 @@ var (
 	ignoresub = flag.Bool("ignoresub", false, "Ignore tracing sub-requests when missing glue")
 	tcp       = flag.Bool("tcp", false, "Use TCP when querying DNS servers")
 	client    = flag.String("client", "", "Sends EDNS Client Subnet option with specified IP address")
-	port      = flag.Int("port", 53, "Use to send DNS queries to non-standard ports")
+	port      = flag.Uint("port", 53, "Use to send DNS queries to non-standard ports")
 )
 
 func main() {
@@ -69,12 +69,14 @@ doesn't necessarily reflect real life.
 		// Nameserver starts with '@'
 		if arg[0] == '@' {
 			nsIP = arg
+
 			continue
 		}
 
 		// Presume next argument is qtype and attempt to match
 		if v, ok := dns.StringToType[strings.ToUpper(arg)]; ok {
 			qtype = append(qtype, v)
+
 			continue
 		}
 

resolver.go

@@ -22,6 +22,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"net"
 	"strconv"
@@ -34,11 +35,16 @@ import (
 	"github.com/miekg/dns"
 )
 
-const defaultDNSTimeout = 2000 * time.Millisecond
+const (
+	defaultDNSTimeout = 2000 * time.Millisecond
+)
 
 var (
-	roots  *hints.Hints
-	aCache *cache.Cache
+	roots                *hints.Hints
+	aCache               *cache.Cache
+	ErrParseClientSubnet = errors.New("failure to parse client-subnet IP")
+	ErrIdMismatch        = errors.New("id mismatch")
+	ErrResolve           = errors.New("unable to resolve")
 )
 
 func doDNSQuery(qname string, qtype uint16, nsIP, nsLabel, zone string, rttIn time.Duration, sub bool) (time.Duration, error) {
@@ -57,6 +63,7 @@ func doDNSQuery(qname string, qtype uint16, nsIP, nsLabel, zone string, rttIn ti
 			for _, rr := range v {
 				color.Green("%v", rr)
 			}
+
 			return rttIn, nil
 		}
 	}
@@ -90,7 +97,7 @@ func doDNSQuery(qname string, qtype uint16, nsIP, nsLabel, zone string, rttIn ti
 			}
 
 			if e.Address == nil {
-				return rttIn, fmt.Errorf("failure to parse client-subnet IP: %s", *client)
+				return rttIn, fmt.Errorf("%w: %v", ErrParseClientSubnet, *client)
 			}
 
 			if e.Address.To4() == nil {
@@ -103,14 +110,13 @@ func doDNSQuery(qname string, qtype uint16, nsIP, nsLabel, zone string, rttIn ti
 		m.Extra = append(m.Extra, o)
 	}
 
-	r, rtt, err := c.Exchange(m, net.JoinHostPort(nsIP, strconv.Itoa(*port)))
+	r, rtt, err := c.Exchange(m, net.JoinHostPort(nsIP, strconv.Itoa(int(*port))))
 
 Retry:
 	// Truncated responses and UDP timeouts are candidates for retry
 	if m.Truncated ||
 		(err != nil && strings.HasPrefix(err.Error(), "read udp") &&
 			strings.HasSuffix(err.Error(), "i/o timeout")) {
-
 		if *fallback {
 			// Enable EDNS and 4096 buffer size if previously not enabled
 			if !*edns {
@@ -122,7 +128,7 @@ Retry:
 				o.SetUDPSize(dns.DefaultMsgSize)
 				m.Extra = append(m.Extra, o)
 
-				r, rtt, err = c.Exchange(m, net.JoinHostPort(nsIP, strconv.Itoa(*port)))
+				r, rtt, err = c.Exchange(m, net.JoinHostPort(nsIP, strconv.Itoa(int(*port))))
 				*edns = true
 
 				goto Retry
@@ -131,7 +137,7 @@ Retry:
 				color.Red("! Answer truncated, retrying with TCP")
 
 				c.Net = "tcp"
-				r, rtt, err = c.Exchange(m, net.JoinHostPort(nsIP, strconv.Itoa(*port)))
+				r, rtt, err = c.Exchange(m, net.JoinHostPort(nsIP, strconv.Itoa(int(*port))))
 				*fallback = false
 
 				goto Retry
@@ -143,16 +149,18 @@ Retry:
 
 	// Response ID mismatch
 	if r.Id != m.Id {
-		return rttIn, fmt.Errorf("id mismatch")
+		return rttIn, fmt.Errorf("%w", ErrIdMismatch)
 	}
 
 	fmt.Printf("Answer RTT: %v from %v\n", rtt, nsIP)
 
 	// Update A cache from A type RRs in ANSWER section
 	for _, rr := range r.Answer {
 		if rr.Header().Rrtype == dns.TypeA {
-			rrn := strings.ToLower(rr.Header().Name)
-			aCache.Add(rrn, rr.(*dns.A))
+			if a, ok := rr.(*dns.A); ok {
+				rrn := strings.ToLower(rr.Header().Name)
+				aCache.Add(rrn, a)
+			}
 		}
 	}
 
@@ -196,13 +204,15 @@ Retry:
 					// In-zone CNAME target so we can continue with current nameserver
 					if strings.HasSuffix(c.Target, "."+zone) {
 						fmt.Printf("\n")
+
 						return doDNSQuery(c.Target, qtype, nsIP, nsLabel, zone, rtt+rttIn, sub)
 					}
 
 					// Start sub-query from the root nameservers
 					nsLabel, nsIP, _ = roots.GetRand()
 					color.Yellow("~ Out of zone CNAME target, sub-query will restart from \".\"")
 					fmt.Printf("\n")
+
 					return doDNSQuery(c.Target, qtype, nsIP, nsLabel, ".", rtt+rttIn, sub)
 				}
 			}
@@ -244,7 +254,7 @@ Retry:
 					nextNs = v.A.String()
 					rtt += rtt2
 				} else {
-					return rttIn, fmt.Errorf("unable to resolve %q/A (following authority NS)", nextNs)
+					return rttIn, fmt.Errorf("%w: %q/A (following authority NS)", ErrResolve, nextNs)
 				}
 			}
 
@@ -255,7 +265,7 @@ Retry:
 		}
 	}
 
-	return rtt + rttIn, fmt.Errorf("unable to resolve %q/%v @ %v(%v)", qname, dns.TypeToString[qtype],
+	return rtt + rttIn, fmt.Errorf("%w: %q/%v @ %v(%v)", ErrResolve, qname, dns.TypeToString[qtype],
 		nsIP, nsLabel)
 }
 
@@ -267,5 +277,6 @@ func getRRset(rr []dns.RR, qname string, qtype uint16) []dns.RR {
 			rr1 = append(rr1, rr)
 		}
 	}
+
 	return rr1
 }