resolver/dns: support custom dns authority

wip add custom resolver test
grpc · Sep 25, 2018 · 5094edd · 5094edd
1 parent 30212c8
commit 5094edd
Show file tree

Hide file tree

Showing 9 changed files with 301 additions and 132 deletions.
diff --git a/resolver/dns/dns_resolver.go b/resolver/dns/dns_resolver.go
@@ -1,6 +1,6 @@
 /*
  *
- * Copyright 2017 gRPC authors.
+ * Copyright 2018 gRPC authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -73,10 +73,7 @@ type dnsBuilder struct {
 
 // Build creates and starts a DNS resolver that watches the name resolution of the target.
 func (b *dnsBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOption) (resolver.Resolver, error) {
-	if target.Authority != "" {
-		return nil, fmt.Errorf("default DNS resolver does not support custom DNS server")
-	}
-	host, port, err := parseTarget(target.Endpoint)
+	host, port, err := parseTarget(target.Endpoint, defaultPort)
 	if err != nil {
 		return nil, err
 	}
@@ -111,6 +108,15 @@ func (b *dnsBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts
 		disableServiceConfig: opts.DisableServiceConfig,
 	}
 
+	if target.Authority == "" {
+		d.resolver = defaultResolver
+	} else {
+		d.resolver, err = customAuthorityResolver(target.Authority)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	d.wg.Add(1)
 	go d.watcher()
 	return d, nil
@@ -121,6 +127,12 @@ func (b *dnsBuilder) Scheme() string {
 	return "dns"
 }
 
+type netResolver interface {
+	LookupHost(ctx context.Context, host string) (addrs []string, err error)
+	LookupSRV(ctx context.Context, service, proto, name string) (cname string, addrs []*net.SRV, err error)
+	LookupTXT(ctx context.Context, name string) (txts []string, err error)
+}
+
 // ipResolver watches for the name resolution update for an IP address.
 type ipResolver struct {
 	cc resolver.ClientConn
@@ -161,6 +173,7 @@ type dnsResolver struct {
 	retryCount int
 	host       string
 	port       string
+	resolver   netResolver
 	ctx        context.Context
 	cancel     context.CancelFunc
 	cc         resolver.ClientConn
@@ -218,13 +231,13 @@ func (d *dnsResolver) watcher() {
 
 func (d *dnsResolver) lookupSRV() []resolver.Address {
 	var newAddrs []resolver.Address
-	_, srvs, err := lookupSRV(d.ctx, "grpclb", "tcp", d.host)
+	_, srvs, err := d.resolver.LookupSRV(d.ctx, "grpclb", "tcp", d.host)
 	if err != nil {
 		grpclog.Infof("grpc: failed dns SRV record lookup due to %v.\n", err)
 		return nil
 	}
 	for _, s := range srvs {
-		lbAddrs, err := lookupHost(d.ctx, s.Target)
+		lbAddrs, err := d.resolver.LookupHost(d.ctx, s.Target)
 		if err != nil {
 			grpclog.Infof("grpc: failed load balancer address dns lookup due to %v.\n", err)
 			continue
@@ -243,7 +256,7 @@ func (d *dnsResolver) lookupSRV() []resolver.Address {
 }
 
 func (d *dnsResolver) lookupTXT() string {
-	ss, err := lookupTXT(d.ctx, d.host)
+	ss, err := d.resolver.LookupTXT(d.ctx, d.host)
 	if err != nil {
 		grpclog.Infof("grpc: failed dns TXT record lookup due to %v.\n", err)
 		return ""
@@ -263,7 +276,7 @@ func (d *dnsResolver) lookupTXT() string {
 
 func (d *dnsResolver) lookupHost() []resolver.Address {
 	var newAddrs []resolver.Address
-	addrs, err := lookupHost(d.ctx, d.host)
+	addrs, err := d.resolver.LookupHost(d.ctx, d.host)
 	if err != nil {
 		grpclog.Warningf("grpc: failed dns A record lookup due to %v.\n", err)
 		return nil
@@ -305,16 +318,16 @@ func formatIP(addr string) (addrIP string, ok bool) {
 	return "[" + addr + "]", true
 }
 
-// parseTarget takes the user input target string, returns formatted host and port info.
+// parseTarget takes the user input target string and default port, returns formatted host and port info.
 // If target doesn't specify a port, set the port to be the defaultPort.
 // If target is in IPv6 format and host-name is enclosed in sqarue brackets, brackets
 // are strippd when setting the host.
 // examples:
-// target: "www.google.com" returns host: "www.google.com", port: "443"
-// target: "ipv4-host:80" returns host: "ipv4-host", port: "80"
-// target: "[ipv6-host]" returns host: "ipv6-host", port: "443"
-// target: ":80" returns host: "localhost", port: "80"
-func parseTarget(target string) (host, port string, err error) {
+// target: "www.google.com" defaultPort: "443" returns host: "www.google.com", port: "443"
+// target: "ipv4-host:80" defaultPort: "443" returns host: "ipv4-host", port: "80"
+// target: "[ipv6-host]" defaultPort: "443" returns host: "ipv6-host", port: "443"
+// target: ":80" defaultPort: "443" returns host: "localhost", port: "80"
+func parseTarget(target, defaultPort string) (host, port string, err error) {
 	if target == "" {
 		return "", "", errMissingAddr
 	}

diff --git a/resolver/dns/dns_resolver_go19_test.go b/resolver/dns/dns_resolver_go19_test.go
@@ -0,0 +1,135 @@
+// +build go1.9
+
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package dns
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"testing"
+
+	"golang.org/x/net/context"
+	"google.golang.org/grpc/internal/leakcheck"
+	"google.golang.org/grpc/resolver"
+)
+
+func TestCustomAuthority(t *testing.T) {
+	defer leakcheck.Check(t)
+
+	tests := []struct {
+		authority     string
+		authorityWant string
+		expectError   bool
+	}{
+		{
+			"4.3.2.1:" + defaultDNSSvrPort,
+			"4.3.2.1:" + defaultDNSSvrPort,
+			false,
+		},
+		{
+			"4.3.2.1:123",
+			"4.3.2.1:123",
+			false,
+		},
+		{
+			"4.3.2.1",
+			"4.3.2.1:" + defaultDNSSvrPort,
+			false,
+		},
+		{
+			"::1",
+			"[::1]:" + defaultDNSSvrPort,
+			false,
+		},
+		{
+			"[::1]",
+			"[::1]:" + defaultDNSSvrPort,
+			false,
+		},
+		{
+			"[::1]:123",
+			"[::1]:123",
+			false,
+		},
+		{
+			"dnsserver.com",
+			"dnsserver.com:" + defaultDNSSvrPort,
+			false,
+		},
+		{
+			":123",
+			"localhost:123",
+			false,
+		},
+		{
+			":",
+			"",
+			true,
+		},
+		{
+			"[::1]:",
+			"",
+			true,
+		},
+		{
+			"dnsserver.com:",
+			"",
+			true,
+		},
+	}
+	oldCustomAuthorityDialler := customAuthorityDialler
+	defer func() {
+		customAuthorityDialler = oldCustomAuthorityDialler
+	}()
+
+	for _, a := range tests {
+		errChan := make(chan error, 1)
+		customAuthorityDialler = func(authority string) func(ctx context.Context, network, address string) (net.Conn, error) {
+			if authority != a.authorityWant {
+				errChan <- fmt.Errorf("wrong custom authority passed to resolver. input: %s expected: %s actual: %s", a.authority, a.authorityWant, authority)
+			} else {
+				errChan <- nil
+			}
+			return func(ctx context.Context, network, address string) (net.Conn, error) {
+				return nil, errors.New("no need to dial")
+			}
+		}
+
+		b := NewBuilder()
+		cc := &testClientConn{target: "foo.bar.com"}
+		r, err := b.Build(resolver.Target{Endpoint: "foo.bar.com", Authority: a.authority}, cc, resolver.BuildOption{})
+
+		if err == nil {
+			r.Close()
+
+			err = <-errChan
+			if err != nil {
+				t.Errorf(err.Error())
+			}
+
+			if a.expectError {
+				t.Errorf("custom authority should have caused an error: %s", a.authority)
+			}
+		} else if !a.expectError {
+			t.Errorf("unexpected error using custom authority %s: %s", a.authority, err)
+		}
+	}
+}
diff --git a/resolver/dns/dns_resolver_test.go b/resolver/dns/dns_resolver_test.go
@@ -1,6 +1,6 @@
 /*
  *
- * Copyright 2017 gRPC authors.
+ * Copyright 2018 gRPC authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,6 +27,7 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/net/context"
 	"google.golang.org/grpc/internal/leakcheck"
 	"google.golang.org/grpc/resolver"
 )
@@ -78,6 +79,30 @@ func (t *testClientConn) getSc() (string, int) {
 	return t.sc, t.s
 }
 
+type testResolver struct {
+}
+
+func (*testResolver) LookupHost(ctx context.Context, host string) ([]string, error) {
+	return hostLookup(host)
+}
+
+func (*testResolver) LookupSRV(ctx context.Context, service, proto, name string) (string, []*net.SRV, error) {
+	return srvLookup(service, proto, name)
+}
+
+func (*testResolver) LookupTXT(ctx context.Context, host string) ([]string, error) {
+	return txtLookup(host)
+}
+
+func replaceNetFunc() func() {
+	oldResolver := defaultResolver
+	defaultResolver = &testResolver{}
+
+	return func() {
+		defaultResolver = oldResolver
+	}
+}
+
 var hostLookupTbl = struct {
 	sync.Mutex
 	tbl map[string][]string

diff --git a/resolver/dns/go17.go b/resolver/dns/go17.go
diff --git a/resolver/dns/go17_test.go b/resolver/dns/go17_test.go