dcerpc: transport: add custom dns resolver (#127)

dcerpc/conn.go

@@ -122,7 +122,7 @@ func Dial(ctx context.Context, addr string, opts ...Option) (Conn, error) {
 		opts:       append(opts, WithGroup(group)),
 	}
 
-	ip, hostName, binding, err := ParseServerAddrWithDNSLookup(addr, true)
+	ip, hostName, binding, err := parseServerAddrWithDNSLookup(ctx, addr, settings.DNSResolver)
 	if err != nil {
 		return nil, fmt.Errorf("dial: %w", err)
 	}
@@ -156,6 +156,14 @@ func ParseServerAddr(addr string) (net.IP, string, *StringBinding, error) {
 // ParseServerAddrWithDNSLookup function parses the server address
 // and performs the DNS lookup if required.
 func ParseServerAddrWithDNSLookup(addr string, doLookup bool) (net.IP, string, *StringBinding, error) {
+	var resolver DNSResolver
+	if doLookup {
+		resolver = net.DefaultResolver
+	}
+	return parseServerAddrWithDNSLookup(context.Background(), addr, resolver)
+}
+
+func parseServerAddrWithDNSLookup(ctx context.Context, addr string, resolver DNSResolver) (net.IP, string, *StringBinding, error) {
 
 	if ip := net.ParseIP(addr); ip != nil {
 		// this is an ip address.
@@ -166,18 +174,18 @@ func ParseServerAddrWithDNSLookup(addr string, doLookup bool) (net.IP, string, *
 
 		// this is fqdn.
 
-		if !doLookup {
+		if resolver == nil {
 			return nil, addr, nil, nil
 		}
 
-		ips, err := net.LookupIP(addr)
+		ias, err := resolver.LookupIPAddr(ctx, addr)
 		if err != nil {
 			return nil, "", nil, fmt.Errorf("lookup server address: %w", err)
 		}
 
-		for _, ip := range ips {
-			if ip.To4() != nil {
-				return ip, addr, nil, nil
+		for _, ia := range ias {
+			if ia.IP.To4() != nil {
+				return ia.IP, addr, nil, nil
 			}
 		}
 

dcerpc/transport_settings.go

@@ -65,6 +65,8 @@ type Transport struct {
 	// If set to `true`, new connection will be established
 	// for every new client with matching binding.
 	NoReuseTransport bool
+	// DNS resolver.
+	DNSResolver DNSResolver
 }
 
 // The transport connection option.
@@ -116,6 +118,17 @@ func WithSMBDialer(dialer *smb2.Dialer) ConnectOption {
 	return func(o *Transport) { o.SMBDialer = dialer }
 }
 
+// DNSResolver interface is used to resolve the hostname to IP addresses.
+type DNSResolver interface {
+	// LookupIPAddr looks up the given host and returns a list of IP addresses.
+	LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error)
+}
+
+// WithDNSResolver function sets the DNS resolver for hostname resolution.
+func WithDNSResolver(resolver DNSResolver) ConnectOption {
+	return func(o *Transport) { o.DNSResolver = resolver }
+}
+
 // WithEndpointMapper option sets the endpoint mapper to find the endpoint
 // (port or named pipe) for the selected abstract syntax.
 //
@@ -144,6 +157,7 @@ func NewTransport() Transport {
 		Timeout:                      10 * time.Second,
 		Deadline:                     3 * time.Second,
 		SMBPort:                      445,
+		DNSResolver:                  net.DefaultResolver,
 	}
 }