net: add (*Resolver).LookupIP

Previously, looking up only IPv4 or IPv6 addresses was only possible
with DefaultResolver via ResolveIPAddr. Add this functionality to the
Resolver type with a new method, LookupIP. This largely brings Resolver
functionally to parity with the global functions. The name LookupIP is
used over ResolveIPAddr to be consistent with the other Resolver
methods.

There are two main benefits to (*Resolver).LookupIP over
(*Resolver).LookupHost. First is an ergonomic benefit. Wanting a
specific family of address is common enough to justify a method, evident
by the existence of ResolveIPAddr. Second, this opens the possibility of
not performing unnecessary DNS requests when only a specific family of
addresses are needed. This optimization is left to follow up work.

Updates #30452

Change-Id: I241f61019588022a39738f8920b0ddba900cecdd
Reviewed-on: https://go-review.googlesource.com/c/go/+/228641
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Reviewed-by: Emmanuel Odeke <emm.odeke@gmail.com>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

2 files changed, 106 insertions, 0 deletions

diff --git a/src/net/lookup.go b/src/net/lookup.go
index 9cebd10c87..5f7119872a 100644
--- a/src/net/lookup.go
+++ b/src/net/lookup.go
@@ -204,6 +204,31 @@ func (r *Resolver) LookupIPAddr(ctx context.Context, host string) ([]IPAddr, err
 	return r.lookupIPAddr(ctx, "ip", host)
 }
 
+// LookupIP looks up host for the given network using the local resolver.
+// It returns a slice of that host's IP addresses of the type specified by
+// network.
+// network must be one of "ip", "ip4" or "ip6".
+func (r *Resolver) LookupIP(ctx context.Context, network, host string) ([]IP, error) {
+	afnet, _, err := parseNetwork(ctx, network, false)
+	if err != nil {
+		return nil, err
+	}
+	switch afnet {
+	case "ip", "ip4", "ip6":
+	default:
+		return nil, UnknownNetworkError(network)
+	}
+	addrs, err := r.internetAddrList(ctx, afnet, host)
+	if err != nil {
+		return nil, err
+	}
+	ips := make([]IP, 0, len(addrs))
+	for _, addr := range addrs {
+		ips = append(ips, addr.(*IPAddr).IP)
+	}
+	return ips, nil
+}
+
 // onlyValuesCtx is a context that uses an underlying context
 // for value lookup if the underlying context hasn't yet expired.
 type onlyValuesCtx struct {
diff --git a/src/net/lookup_test.go b/src/net/lookup_test.go
index 4956037f4b..68bffcab8f 100644
--- a/src/net/lookup_test.go
+++ b/src/net/lookup_test.go
@@ -913,6 +913,7 @@ func TestNilResolverLookup(t *testing.T) {
 	r.LookupCNAME(ctx, "google.com")
 	r.LookupHost(ctx, "google.com")
 	r.LookupIPAddr(ctx, "google.com")
+	r.LookupIP(ctx, "ip", "google.com")
 	r.LookupMX(ctx, "gmail.com")
 	r.LookupNS(ctx, "google.com")
 	r.LookupPort(ctx, "tcp", "smtp")
@@ -1185,3 +1186,83 @@ func TestLookupNullByte(t *testing.T) {
 	testenv.SkipFlakyNet(t)
 	LookupHost("foo\x00bar") // check that it doesn't panic; it used to on Windows
 }
+
+func TestResolverLookupIP(t *testing.T) {
+	testenv.MustHaveExternalNetwork(t)
+
+	v4Ok := supportsIPv4() && *testIPv4
+	v6Ok := supportsIPv6() && *testIPv6
+
+	defer dnsWaitGroup.Wait()
+
+	for _, impl := range []struct {
+		name string
+		fn   func() func()
+	}{
+		{"go", forceGoDNS},
+		{"cgo", forceCgoDNS},
+	} {
+		t.Run("implementation: "+impl.name, func(t *testing.T) {
+			fixup := impl.fn()
+			if fixup == nil {
+				t.Skip("not supported")
+			}
+			defer fixup()
+
+			for _, network := range []string{"ip", "ip4", "ip6"} {
+				t.Run("network: "+network, func(t *testing.T) {
+					switch {
+					case network == "ip4" && !v4Ok:
+						t.Skip("IPv4 is not supported")
+					case network == "ip6" && !v6Ok:
+						t.Skip("IPv6 is not supported")
+					}
+
+					// google.com has both A and AAAA records.
+					const host = "google.com"
+					ips, err := DefaultResolver.LookupIP(context.Background(), network, host)
+					if err != nil {
+						testenv.SkipFlakyNet(t)
+						t.Fatalf("DefaultResolver.LookupIP(%q, %q): failed with unexpected error: %v", network, host, err)
+					}
+
+					var v4Addrs []IP
+					var v6Addrs []IP
+					for _, ip := range ips {
+						switch {
+						case ip.To4() != nil:
+							// We need to skip the test below because To16 will
+							// convent an IPv4 address to an IPv4-mapped IPv6
+							// address.
+							v4Addrs = append(v4Addrs, ip)
+						case ip.To16() != nil:
+							v6Addrs = append(v6Addrs, ip)
+						default:
+							t.Fatalf("IP=%q is neither IPv4 nor IPv6", ip)
+						}
+					}
+
+					// Check that we got the expected addresses.
+					if network == "ip4" || network == "ip" && v4Ok {
+						if len(v4Addrs) == 0 {
+							t.Errorf("DefaultResolver.LookupIP(%q, %q): no IPv4 addresses", network, host)
+						}
+					}
+					if network == "ip6" || network == "ip" && v6Ok {
+						if len(v6Addrs) == 0 {
+							t.Errorf("DefaultResolver.LookupIP(%q, %q): no IPv6 addresses", network, host)
+						}
+					}
+
+					// Check that we didn't get any unexpected addresses.
+					if network == "ip6" && len(v4Addrs) > 0 {
+						t.Errorf("DefaultResolver.LookupIP(%q, %q): unexpected IPv4 addresses: %v", network, host, v4Addrs)
+					}
+					if network == "ip4" && len(v6Addrs) > 0 {
+						t.Errorf("DefaultResolver.LookupIP(%q, %q): unexpected IPv6 addresses: %v", network, host, v6Addrs)
+					}
+				})
+			}
+		})
+	}
+}