Ability to pass custom DNS resolver to TCPDialer (#689)

* Ability to pass custom DNS resolver to TCPdialer

* Update tcpdialer.go

Co-Authored-By: Erik Dubbelboer <erik@dubbelboer.com>

1 files changed, 26 insertions, 7 deletions

diff --git a/tcpdialer.go b/tcpdialer.go
index 6a5cd3a..8d84983 100644
--- a/tcpdialer.go
+++ b/tcpdialer.go
@@ -1,6 +1,7 @@
 package fasthttp
 
 import (
+	"context"
 	"errors"
 	"net"
 	"strconv"
@@ -129,6 +130,19 @@ type TCPDialer struct {
 	// Changes made after the first Dial will not affect anything.
 	Concurrency int
 
+	// This may be used to override DNS resolving policy, like this:
+	// var dialer = &fasthttp.TCPDialer{
+	// 	Resolver: &net.Resolver{
+	// 		PreferGo:     true,
+	// 		StrictErrors: false,
+	// 		Dial: func (ctx context.Context, network, address string) (net.Conn, error) {
+	// 			d := net.Dialer{}
+	// 			return d.DialContext(ctx, "udp", "8.8.8.8:53")
+	// 		},
+	// 	},
+	// }
+	Resolver *net.Resolver
+
 	tcpAddrsLock sync.Mutex
 	tcpAddrsMap  map[string]*tcpAddrEntry
 
@@ -387,7 +401,7 @@ func (d *TCPDialer) getTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, uin
 	d.tcpAddrsLock.Unlock()
 
 	if e == nil {
-		addrs, err := resolveTCPAddrs(addr, dualStack)
+		addrs, err := resolveTCPAddrs(addr, dualStack, d.Resolver)
 		if err != nil {
 			d.tcpAddrsLock.Lock()
 			e = d.tcpAddrsMap[addr]
@@ -412,7 +426,7 @@ func (d *TCPDialer) getTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, uin
 	return e.addrs, idx, nil
 }
 
-func resolveTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, error) {
+func resolveTCPAddrs(addr string, dualStack bool, resolver *net.Resolver) ([]net.TCPAddr, error) {
 	host, portS, err := net.SplitHostPort(addr)
 	if err != nil {
 		return nil, err
@@ -422,20 +436,25 @@ func resolveTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, error) {
 		return nil, err
 	}
 
-	ips, err := net.LookupIP(host)
+	if resolver == nil {
+		resolver = net.DefaultResolver
+	}
+
+	ctx := context.Background()
+	ipaddrs, err := resolver.LookupIPAddr(ctx, host)
 	if err != nil {
 		return nil, err
 	}
 
-	n := len(ips)
+	n := len(ipaddrs)
 	addrs := make([]net.TCPAddr, 0, n)
 	for i := 0; i < n; i++ {
-		ip := ips[i]
-		if !dualStack && ip.To4() == nil {
+		ip := ipaddrs[i]
+		if !dualStack && ip.IP.To4() == nil {
 			continue
 		}
 		addrs = append(addrs, net.TCPAddr{
-			IP:   ip,
+			IP:   ip.IP,
 			Port: port,
 		})
 	}