From 0742784136fa0d6702491a6829babe347f3b7cae Mon Sep 17 00:00:00 2001
From: Thor <thansen@digitalocean.com>
Date: Thu, 11 Apr 2019 10:57:15 -0500
Subject: [PATCH] use LookupIPAddr to support context timeout

---
 prober/utils.go | 68 ++++++++++++++++++++++++-------------------------
 1 file changed, 34 insertions(+), 34 deletions(-)

diff --git a/prober/utils.go b/prober/utils.go
index 8c9a72b5..d530aac5 100644
--- a/prober/utils.go
+++ b/prober/utils.go
@@ -61,45 +61,45 @@ func chooseProtocol(ctx context.Context, IPProtocol string, fallbackIPProtocol b
 		probeDNSLookupTimeSeconds.Add(lookupTime)
 	}()
 
-	ipC := make(chan *net.IPAddr, 1)
-	errors := make(chan error, 1)
-	go func() {
-		defer close(ipC)
-		defer close(errors)
-		ip, err := net.ResolveIPAddr(IPProtocol, target)
-		if err != nil {
-			if !fallbackIPProtocol {
-				level.Error(logger).Log("msg", "Resolution with IP protocol failed (fallback_ip_protocol is false):", "err", err)
-			} else {
-				level.Warn(logger).Log("msg", "Resolution with IP protocol failed, attempting fallback protocol", "fallback_protocol", fallbackProtocol, "err", err)
-				ip, err = net.ResolveIPAddr(fallbackProtocol, target)
+	resolver := &net.Resolver{}
+	ips, err := resolver.LookupIPAddr(ctx, target)
+	if err != nil {
+		level.Error(logger).Log("msg", "Resolution with IP protocol failed", "err", err)
+		return nil, 0.0, err
+	}
+
+	// Return the IP in the requested protocol format
+	var fallback net.IPAddr
+	for _, ip := range ips {
+		switch IPProtocol {
+		case "ip4":
+			if ip.IP.To4() != nil {
+				level.Info(logger).Log("msg", "Resolved target address", "ip", ip)
+				probeIPProtocolGauge.Set(4)
+				return &ip, lookupTime, nil
 			}
 
-			if err != nil {
-				if IPProtocol == "ip6" {
-					probeIPProtocolGauge.Set(6)
-				} else {
-					probeIPProtocolGauge.Set(4)
-				}
-				errors <- err
-				return
+			// ip4 as fallback
+			fallback = ip
+
+		case "ip6":
+
+			if ip.IP.To4() == nil {
+				level.Info(logger).Log("msg", "Resolved target address", "ip", ip)
+				probeIPProtocolGauge.Set(6)
+				return &ip, lookupTime, nil
 			}
-		}
 
-		if ip.IP.To4() == nil {
-			probeIPProtocolGauge.Set(6)
-		} else {
-			probeIPProtocolGauge.Set(4)
+			// ip6 as fallback
+			fallback = ip
 		}
+	}
 
-		ipC <- ip
-	}()
-
-	select {
-	case <-ctx.Done():
-		return nil, lookupTime, ctx.Err()
-	case ip := <-ipC:
-		level.Info(logger).Log("msg", "Resolved target address", "ip", ip)
-		return ip, lookupTime, nil
+	// Use fallback ip protocol
+	if fallbackProtocol == "ip4" {
+		probeIPProtocolGauge.Set(4)
+	} else {
+		probeIPProtocolGauge.Set(6)
 	}
+	return &fallback, lookupTime, nil
 }