fix(endpoint): use custom resolver if specified for [IP] tests

Author: ishanjain28

client/config.go

@@ -60,6 +60,8 @@ type Config struct {
 	// Expected format is {protocol}://{host}:{port}, e.g. tcp://8.8.8.8:53
 	DNSResolver string `yaml:"dns-resolver,omitempty"`
 
+	DNSResolverConfig *DNSResolverConfig `yaml:"-"`
+
 	// OAuth2Config is the OAuth2 configuration used for the client.
 	//
 	// If non-nil, the http.Client returned by getHTTPClient will automatically retrieve a token if necessary.
@@ -116,8 +118,10 @@ func (c *Config) ValidateAndSetDefaults() error {
 	}
 	if c.HasCustomDNSResolver() {
 		// Validate the DNS resolver now to make sure it will not return an error later.
-		if _, err := c.parseDNSResolver(); err != nil {
+		if resolver, err := c.parseDNSResolver(); err != nil {
 			return err
+		} else {
+			c.DNSResolverConfig = resolver
 		}
 	}
 	if c.HasOAuth2Config() && !c.OAuth2Config.isValid() {

config/endpoint/endpoint.go

@@ -2,6 +2,7 @@ package endpoint
 
 import (
 	"bytes"
+	"context"
 	"crypto/x509"
 	"encoding/json"
 	"errors"
@@ -432,11 +433,28 @@ func (e *Endpoint) getParsedBody() string {
 }
 
 func (e *Endpoint) getIP(result *Result) {
-	if ips, err := net.LookupIP(result.Hostname); err != nil {
+
+	resolver := net.DefaultResolver
+
+	if e.ClientConfig.HasCustomDNSResolver() {
+		dnsResolver := e.ClientConfig.DNSResolverConfig
+		resolver = &net.Resolver{
+			PreferGo: true,
+			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+				d := net.Dialer{}
+				return d.DialContext(ctx, dnsResolver.Protocol, dnsResolver.Host+":"+dnsResolver.Port)
+			},
+		}
+	}
+
+	addrs, err := resolver.LookupIP(context.Background(), e.ClientConfig.Network, result.Hostname)
+	if err != nil {
 		result.AddError(err.Error())
 		return
-	} else {
-		result.IP = ips[0].String()
+	}
+	for _, ia := range addrs {
+		result.IP = ia.String()
+		return
 	}
 }