From 407e5cbfeb075549566960a5858fcfb0a1e711d3 Mon Sep 17 00:00:00 2001
From: mzack <marco.rivoli.nvh@gmail.com>
Date: Mon, 4 Oct 2021 15:31:14 +0200
Subject: [PATCH] Adding support for custom resolvers in DNS templates

---
 v2/pkg/protocols/dns/dns.go                      | 12 +++++++++---
 v2/pkg/protocols/dns/dnsclientpool/clientpool.go |  9 +++++++--
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/v2/pkg/protocols/dns/dns.go b/v2/pkg/protocols/dns/dns.go
index 4c2b8904f9..e4e7833e93 100644
--- a/v2/pkg/protocols/dns/dns.go
+++ b/v2/pkg/protocols/dns/dns.go
@@ -72,6 +72,8 @@ type Request struct {
 	// description: |
 	//   Recursion determines if resolver should recurse all records to get fresh results.
 	Recursion bool `yaml:"recursion,omitempty" jsonschema:"title=recurse all servers,description=Recursion determines if resolver should recurse all records to get fresh results"`
+	// Resolvers to use for the dns requests
+	Resolvers []string `yaml:"resolvers,omitempty" jsonschema:"title=Resolvers,description=Define resolvers to use within the template"`
 }
 
 // GetID returns the unique ID of the request if any.
@@ -81,10 +83,14 @@ func (r *Request) GetID() string {
 
 // Compile compiles the protocol request for further execution.
 func (r *Request) Compile(options *protocols.ExecuterOptions) error {
-	// Create a dns client for the class
-	client, err := dnsclientpool.Get(options.Options, &dnsclientpool.Configuration{
+	dnsClientOptions := &dnsclientpool.Configuration{
 		Retries: r.Retries,
-	})
+	}
+	if len(r.Resolvers) > 0 {
+		dnsClientOptions.Resolvers = r.Resolvers
+	}
+	// Create a dns client for the class
+	client, err := dnsclientpool.Get(options.Options, dnsClientOptions)
 	if err != nil {
 		return errors.Wrap(err, "could not get dns client")
 	}
diff --git a/v2/pkg/protocols/dns/dnsclientpool/clientpool.go b/v2/pkg/protocols/dns/dnsclientpool/clientpool.go
index 999f2a2c9b..46308c9fe6 100644
--- a/v2/pkg/protocols/dns/dnsclientpool/clientpool.go
+++ b/v2/pkg/protocols/dns/dnsclientpool/clientpool.go
@@ -44,21 +44,24 @@ func Init(options *types.Options) error {
 type Configuration struct {
 	// Retries contains the retries for the dns client
 	Retries int
+	// Resolvers contains the specific per request resolvers
+	Resolvers []string
 }
 
 // Hash returns the hash of the configuration to allow client pooling
 func (c *Configuration) Hash() string {
 	builder := &strings.Builder{}
-	builder.Grow(8)
 	builder.WriteString("r")
 	builder.WriteString(strconv.Itoa(c.Retries))
+	builder.WriteString("l")
+	builder.WriteString(strings.Join(c.Resolvers, ""))
 	hash := builder.String()
 	return hash
 }
 
 // Get creates or gets a client for the protocol based on custom configuration
 func Get(options *types.Options, configuration *Configuration) (*retryabledns.Client, error) {
-	if !(configuration.Retries > 1) {
+	if !(configuration.Retries > 1) && len(configuration.Resolvers) == 0 {
 		return normalClient, nil
 	}
 	hash := configuration.Hash()
@@ -72,6 +75,8 @@ func Get(options *types.Options, configuration *Configuration) (*retryabledns.Cl
 	resolvers := defaultResolvers
 	if options.ResolversFile != "" {
 		resolvers = options.InternalResolversList
+	} else if len(configuration.Resolvers) > 0 {
+		resolvers = configuration.Resolvers
 	}
 	client := retryabledns.New(resolvers, configuration.Retries)