From eeacf84daebdb15a558b65e67b0678a95324e840 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 29 Mar 2021 21:29:27 -0700
Subject: [PATCH] cmd/tailscale/cli: factor out tailscaleIPFromArg from ping
 command

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/tailscale/cli/ping.go | 60 ++++++++++++++++++++-------------------
 1 file changed, 31 insertions(+), 29 deletions(-)

diff --git a/cmd/tailscale/cli/ping.go b/cmd/tailscale/cli/ping.go
index 4d2a6f502..ff9af7e78 100644
--- a/cmd/tailscale/cli/ping.go
+++ b/cmd/tailscale/cli/ping.go
@@ -83,37 +83,11 @@ func runPing(ctx context.Context, args []string) error {
 	go pump(ctx, bc, c)
 
 	hostOrIP := args[0]
-
-	// If the argument is an IP address, use it directly without any resolution.
-	if net.ParseIP(hostOrIP) != nil {
-		ip = hostOrIP
+	ip, err := tailscaleIPFromArg(ctx, hostOrIP)
+	if err != nil {
+		return err
 	}
 
-	// Otherwise, try to resolve it first from the network peer list.
-	if ip == "" {
-		st, err := tailscale.Status(ctx)
-		if err != nil {
-			return err
-		}
-		for _, ps := range st.Peer {
-			if hostOrIP == dnsOrQuoteHostname(st, ps) || hostOrIP == ps.DNSName {
-				ip = ps.TailAddr
-				break
-			}
-		}
-	}
-
-	// Finally, use DNS.
-	if ip == "" {
-		var res net.Resolver
-		if addrs, err := res.LookupHost(ctx, hostOrIP); err != nil {
-			return fmt.Errorf("error looking up IP of %q: %v", hostOrIP, err)
-		} else if len(addrs) == 0 {
-			return fmt.Errorf("no IPs found for %q", hostOrIP)
-		} else {
-			ip = addrs[0]
-		}
-	}
 	if pingArgs.verbose && ip != hostOrIP {
 		log.Printf("lookup %q => %q", hostOrIP, ip)
 	}
@@ -166,3 +140,31 @@ func runPing(ctx context.Context, args []string) error {
 		}
 	}
 }
+
+func tailscaleIPFromArg(ctx context.Context, hostOrIP string) (ip string, err error) {
+	// If the argument is an IP address, use it directly without any resolution.
+	if net.ParseIP(hostOrIP) != nil {
+		return hostOrIP, nil
+	}
+
+	// Otherwise, try to resolve it first from the network peer list.
+	st, err := tailscale.Status(ctx)
+	if err != nil {
+		return "", err
+	}
+	for _, ps := range st.Peer {
+		if hostOrIP == dnsOrQuoteHostname(st, ps) || hostOrIP == ps.DNSName {
+			return ps.TailAddr, nil
+		}
+	}
+
+	// Finally, use DNS.
+	var res net.Resolver
+	if addrs, err := res.LookupHost(ctx, hostOrIP); err != nil {
+		return "", fmt.Errorf("error looking up IP of %q: %v", hostOrIP, err)
+	} else if len(addrs) == 0 {
+		return "", fmt.Errorf("no IPs found for %q", hostOrIP)
+	} else {
+		return addrs[0], nil
+	}
+}