resolver/dns: support custom dns authority

wip add custom resolver test
grpc · Sep 17, 2018 · 91db3ef · 91db3ef
1 parent f2aaa9b
commit 91db3ef
Show file tree

Hide file tree

Showing 10 changed files with 359 additions and 98 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,9 +73,6 @@ 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)
 	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

diff --git a/resolver/dns/dns_resolver_go19_test.go b/resolver/dns/dns_resolver_go19_test.go
@@ -0,0 +1,129 @@
+// +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 (
+	"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
+		parseError    bool
+	}{
+		{
+			"4.3.2.1:53",
+			"4.3.2.1:53",
+			false,
+		},
+		{
+			"4.3.2.1:123",
+			"4.3.2.1:123",
+			false,
+		},
+		{
+			"4.3.2.1",
+			"4.3.2.1:53",
+			false,
+		},
+		{
+			"::1",
+			"[::1]:53",
+			false,
+		},
+		{
+			"[::1]",
+			"[::1]:53",
+			false,
+		},
+		{
+			"[::1]:123",
+			"[::1]:123",
+			false,
+		},
+		{
+			"dnsserver.com",
+			"dnsserver.com:53",
+			false,
+		},
+		{
+			":123",
+			"localhost:123",
+			false,
+		},
+		{
+			":",
+			"",
+			true,
+		},
+		{
+			"[::1]:",
+			"",
+			true,
+		},
+		{
+			"dnsserver.com:",
+			"",
+			true,
+		},
+	}
+	oldCustomAuthorityDialler := customAuthorityDialler
+	defer func() {
+		customAuthorityDialler = oldCustomAuthorityDialler
+	}()
+	for _, a := range tests {
+
+		parsed, err := ensureAuthorityPort(a.authority)
+		if err != nil {
+			if !a.parseError {
+				t.Errorf("unexpected authority parse error. input: %s error: %s", a.authority, err)
+			}
+			continue
+		}
+		if parsed != a.authorityWant {
+			t.Errorf("authority did not parse correctly. input: %s expected: %s actual: %s", a.authority, a.authorityWant, parsed)
+		}
+
+		customAuthorityDialler = func(authority string) func(ctx context.Context, network, address string) (net.Conn, error) {
+			if authority != a.authorityWant {
+				t.Errorf("wrong custom authority passed to resolver. input: %s expected: %s actual: %s", a.authority, a.authorityWant, authority)
+			}
+			return func(ctx context.Context, network, address string) (net.Conn, error) {
+				var dialer net.Dialer
+				return dialer.DialContext(ctx, network, authority)
+			}
+		}
+
+		b := NewBuilder()
+		cc := &testClientConn{target: "foo.bar.com"}
+		r, _ := b.Build(resolver.Target{Endpoint: "foo.bar.com", Authority: a.authority}, cc, resolver.BuildOption{})
+		r.Close()
+	}
+
+}
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,28 @@ 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_test.go b/resolver/dns/go17_test.go
diff --git a/resolver/dns/go18.go b/resolver/dns/go18.go
@@ -2,7 +2,7 @@
 
 /*
  *
- * 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.
@@ -20,10 +20,26 @@
 
 package dns
 
-import "net"
+import (
+	"net"
+
+	"golang.org/x/net/context"
+)
 
 var (
-	lookupHost = net.DefaultResolver.LookupHost
-	lookupSRV  = net.DefaultResolver.LookupSRV
-	lookupTXT  = net.DefaultResolver.LookupTXT
+	defaultResolver netResolver = &go18upResolver{resolver: net.DefaultResolver}
 )
+
+type go18upResolver struct {
+	resolver *net.Resolver
+}
+
+func (r *go18upResolver) LookupHost(ctx context.Context, host string) ([]string, error) {
+	return r.resolver.LookupHost(ctx, host)
+}
+func (r *go18upResolver) LookupSRV(ctx context.Context, service, proto, name string) (string, []*net.SRV, error) {
+	return r.resolver.LookupSRV(ctx, service, proto, name)
+}
+func (r *go18upResolver) LookupTXT(ctx context.Context, name string) ([]string, error) {
+	return r.resolver.LookupTXT(ctx, name)
+}
diff --git a/resolver/dns/go18_test.go b/resolver/dns/go18_test.go
@@ -2,7 +2,7 @@
 
 /*
  *
- * 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.
@@ -21,29 +21,7 @@
 package dns
 
 import (
-	"context"
 	"fmt"
-	"net"
 )
 
 var errForInvalidTarget = fmt.Errorf("invalid target address [2001:db8:a0b:12f0::1, error info: address [2001:db8:a0b:12f0::1:443: missing ']' in address")
-
-func replaceNetFunc() func() {
-	oldLookupHost := lookupHost
-	oldLookupSRV := lookupSRV
-	oldLookupTXT := lookupTXT
-	lookupHost = func(ctx context.Context, host string) ([]string, error) {
-		return hostLookup(host)
-	}
-	lookupSRV = func(ctx context.Context, service, proto, name string) (string, []*net.SRV, error) {
-		return srvLookup(service, proto, name)
-	}
-	lookupTXT = func(ctx context.Context, host string) ([]string, error) {
-		return txtLookup(host)
-	}
-	return func() {
-		lookupHost = oldLookupHost
-		lookupSRV = oldLookupSRV
-		lookupTXT = oldLookupTXT
-	}
-}