From f4d5f30b918e8cbaceb49f939f07edaac28b663f Mon Sep 17 00:00:00 2001
From: Karan Sharma <hello@mrkaran.dev>
Date: Sat, 2 Jul 2022 11:19:47 +0530
Subject: [PATCH] feat: add retry flag

---
 cmd/doggo/cli.go    |  8 +++-----
 internal/app/app.go | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/cmd/doggo/cli.go b/cmd/doggo/cli.go
index 010e363..97548fc 100644
--- a/cmd/doggo/cli.go
+++ b/cmd/doggo/cli.go
@@ -39,7 +39,8 @@ func main() {
 	f.BoolP("reverse", "x", false, "Performs a DNS Lookup for an IPv4 or IPv6 address. Sets the query type and class to PTR and IN respectively.")
 
 	// Resolver Options
-	f.Int("timeout", 5, "Sets the timeout for a query to T seconds. The default timeout is 5 seconds.")
+	f.Int("timeout", 2, "Sets the timeout for a query to T seconds. The default timeout is 2 seconds.")
+	f.Int("retry", 3, "Number of times to retry DNS lookup")
 	f.Bool("search", true, "Use the search list provided in resolv.conf. It sets the `ndots` parameter as well unless overridden by `ndots` flag.")
 	f.Int("ndots", -1, "Specify the ndots parameter. Default value is taken from resolv.conf and fallbacks to 1 if ndots statement is missing in resolv.conf")
 	f.BoolP("ipv4", "4", false, "Use IPv4 only")
@@ -75,7 +76,6 @@ func main() {
 
 	// Set log level.
 	if k.Bool("debug") {
-		// Set logger level
 		app.Logger.SetLevel(logf.DebugLevel)
 	}
 	// Unmarshall flags to the app.
@@ -96,7 +96,6 @@ func main() {
 	// query type as PTR and query class as IN.
 	// Modify query name like 94.2.0.192.in-addr.arpa if it's an IPv4 address.
 	// Use IP6.ARPA nibble format otherwise.
-
 	if app.QueryFlags.ReverseLookup {
 		app.ReverseLookup()
 	}
@@ -131,7 +130,6 @@ func main() {
 	}
 	app.Resolvers = rslvrs
 
-	// Run the app.
 	app.Logger.Debug("Starting doggo 🐶")
 	if len(app.QueryFlags.QNames) == 0 {
 		f.Usage()
@@ -142,7 +140,7 @@ func main() {
 	var responses []resolvers.Response
 	for _, q := range app.Questions {
 		for _, rslv := range app.Resolvers {
-			resp, err := rslv.Lookup(q)
+			resp, err := app.LookupWithRetry(app.QueryFlags.RetryCount, rslv, q)
 			if err != nil {
 				app.Logger.WithError(err).Error("error looking up DNS records")
 			}
diff --git a/internal/app/app.go b/internal/app/app.go
index f3e9c17..a0cac0b 100644
--- a/internal/app/app.go
+++ b/internal/app/app.go
@@ -1,6 +1,9 @@
 package app
 
 import (
+	"math/rand"
+	"time"
+
 	"github.com/miekg/dns"
 	"github.com/mr-karan/doggo/pkg/models"
 	"github.com/mr-karan/doggo/pkg/resolvers"
@@ -33,3 +36,20 @@ func New(logger *logf.Logger, buildVersion string) App {
 	}
 	return app
 }
+
+// Attempts a DNS Lookup with retries for a given Question.
+func (app *App) LookupWithRetry(attempts int, resolver resolvers.Resolver, ques dns.Question) (resolvers.Response, error) {
+	resp, err := resolver.Lookup(ques)
+	if err != nil {
+		// Retry lookup.
+		attempts--
+		if attempts > 0 {
+			// Add some random delay.
+			time.Sleep(time.Millisecond*300 + (time.Duration(rand.Int63n(int64(time.Millisecond*100))))/2)
+			app.Logger.Debug("retrying lookup")
+			return app.LookupWithRetry(attempts, resolver, ques)
+		}
+		return resolvers.Response{}, err
+	}
+	return resp, nil
+}