From 6c3b17ba0d0fed8b0e0746d79ea78842317ef722 Mon Sep 17 00:00:00 2001
From: Karan Sharma <hello@mrkaran.dev>
Date: Sun, 3 Jul 2022 11:31:33 +0530
Subject: [PATCH] feat: Add support for extra protocol tweaks

Adds support to set protocol tweaks like aa,ad,cd flags to be
set when preparing a dns.Message
---
 cmd/doggo/cli.go       | 18 +++++++++++++++---
 cmd/doggo/help.go      |  8 ++++----
 cmd/doggo/parse.go     | 10 ++++++++++
 pkg/models/models.go   |  1 +
 pkg/resolvers/utils.go |  2 +-
 5 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/cmd/doggo/cli.go b/cmd/doggo/cli.go
index 97548fc..eb5b661 100644
--- a/cmd/doggo/cli.go
+++ b/cmd/doggo/cli.go
@@ -48,6 +48,7 @@ func main() {
 	f.String("strategy", "all", "Strategy to query nameservers in resolv.conf file (`all`, `random`, `first`)")
 	f.String("tls-hostname", "", "Provide a hostname for doing verification of the certificate if the provided DoT nameserver is an IP")
 	f.Bool("skip-hostname-verification", false, "Skip TLS Hostname Verification")
+	f.StringSliceP("tweaks", "Z", []string{}, "Specify protocol tweaks. Set flags like aa,ad,cd")
 
 	// Output Options
 	f.BoolP("json", "J", false, "Set the output format as JSON")
@@ -112,8 +113,7 @@ func main() {
 		app.Logger.WithError(err).Fatal("error loading nameservers")
 	}
 
-	// Load Resolvers.
-	rslvrs, err := resolvers.LoadResolvers(resolvers.Options{
+	ropts := resolvers.Options{
 		Nameservers:        app.Nameservers,
 		UseIPv4:            app.QueryFlags.UseIPv4,
 		UseIPv6:            app.QueryFlags.UseIPv6,
@@ -124,7 +124,19 @@ func main() {
 		Strategy:           app.QueryFlags.Strategy,
 		InsecureSkipVerify: app.QueryFlags.InsecureSkipVerify,
 		TLSHostname:        app.QueryFlags.TLSHostname,
-	})
+	}
+
+	if contains(app.QueryFlags.Tweaks, "aa") {
+		ropts.Authoritative = true
+	}
+	if contains(app.QueryFlags.Tweaks, "ad") {
+		ropts.AuthenticatedData = true
+	}
+	if contains(app.QueryFlags.Tweaks, "cd") {
+		ropts.CheckingDisabled = true
+	}
+	// Load Resolvers.
+	rslvrs, err := resolvers.LoadResolvers(ropts)
 	if err != nil {
 		app.Logger.WithError(err).Fatal("error loading resolver")
 	}
diff --git a/cmd/doggo/help.go b/cmd/doggo/help.go
index a4e719c..511531a 100644
--- a/cmd/doggo/help.go
+++ b/cmd/doggo/help.go
@@ -19,10 +19,10 @@ var appHelpTextTemplate = `{{ "NAME" | color "" "heading" }}:
   {{.Version | color "red" "" }}
 
 {{ "EXAMPLES" | color "" "heading" }}:
-  {{ .Name | color "green" "bold" }} {{ "mrkaran.dev" | color "cyan" "" }}                            	Query a domain using defaults.
-  {{ .Name | color "green" "bold" }} {{ "mrkaran.dev CNAME" | color "cyan" "" }}                 	Looks up for a CNAME record.
-  {{ .Name | color "green" "bold" }} {{ "mrkaran.dev MX @9.9.9.9" | color "cyan" "" }}              	Uses a custom DNS resolver.
-  {{ .Name | color "green" "bold" }} {{"-q mrkaran.dev -t MX -n 1.1.1.1" | color "yellow" ""}}   	Using named arguments.
+  {{ .Name | color "green" "bold" }} {{ "mrkaran.dev" | color "cyan" "" }} {{"\t"}} Query a domain using defaults.
+  {{ .Name | color "green" "bold" }} {{ "mrkaran.dev CNAME" | color "cyan" "" }} {{"\t"}} Looks up for a CNAME record.
+  {{ .Name | color "green" "bold" }} {{ "mrkaran.dev MX @9.9.9.9" | color "cyan" "" }} {{"\t"}} Uses a custom DNS resolver.
+  {{ .Name | color "green" "bold" }} {{"-q mrkaran.dev -t MX -n 1.1.1.1" | color "yellow" ""}} {{"\t"}} Using named arguments.
 
 {{ "Free Form Arguments" | color "" "heading" }}:
   Supply hostnames, query types, classes without any flag. For eg:
diff --git a/cmd/doggo/parse.go b/cmd/doggo/parse.go
index 52ac09a..7b5a324 100644
--- a/cmd/doggo/parse.go
+++ b/cmd/doggo/parse.go
@@ -33,3 +33,13 @@ func loadUnparsedArgs(args []string) ([]string, []string, []string, []string) {
 	}
 	return ns, qt, qc, qn
 }
+
+// contains is a helper method to check if a paritcular element exists in the slice.
+func contains(s []string, e string) bool {
+	for _, a := range s {
+		if a == e {
+			return true
+		}
+	}
+	return false
+}
diff --git a/pkg/models/models.go b/pkg/models/models.go
index cdcf6c1..2df74e4 100644
--- a/pkg/models/models.go
+++ b/pkg/models/models.go
@@ -40,6 +40,7 @@ type QueryFlags struct {
 	InsecureSkipVerify bool          `koanf:"skip-hostname-verification" skip-hostname-verification:"-"`
 	TLSHostname        string        `koanf:"tls-hostname" tls-hostname:"-"`
 	RetryCount         int           `koanf:"retry" retry:"-"`
+	Tweaks             []string      `koanf:"tweaks" json:"-"`
 }
 
 // Nameserver represents the type of Nameserver
diff --git a/pkg/resolvers/utils.go b/pkg/resolvers/utils.go
index 8a20492..a53f80a 100644
--- a/pkg/resolvers/utils.go
+++ b/pkg/resolvers/utils.go
@@ -24,7 +24,7 @@ func prepareMessages(q dns.Question, opts Options) []dns.Msg {
 				Authoritative:     opts.Authoritative,
 				AuthenticatedData: opts.AuthenticatedData,
 				CheckingDisabled:  opts.CheckingDisabled,
-				RecursionDesired:  opts.RecursionDesired,
+				RecursionDesired:  true,
 			},
 		}
 		// It's recommended to only send 1 question for 1 DNS message.