feat: Automatially configure connections using DNS

Author: hessjcg

dialer.go

@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"sort"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -153,6 +154,10 @@ type Dialer struct {
 
 	// iamTokenSource supplies the OAuth2 token used for IAM DB Authn.
 	iamTokenSource oauth2.TokenSource
+
+	// resolver does SRV record DNS lookups when resolving DNS name dialer
+	// configuration.
+	resolver NetResolver
 }
 
 var (
@@ -253,6 +258,13 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
 	if err != nil {
 		return nil, err
 	}
+	var r NetResolver
+	if cfg.resolver != nil {
+		r = cfg.resolver
+	} else {
+		r = net.DefaultResolver
+	}
+
 	d := &Dialer{
 		closed:            make(chan struct{}),
 		cache:             make(map[instance.ConnName]monitoredCache),
@@ -265,10 +277,25 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
 		dialerID:          uuid.New().String(),
 		iamTokenSource:    cfg.iamLoginTokenSource,
 		dialFunc:          cfg.dialFunc,
+		resolver:          r,
 	}
 	return d, nil
 }
 
+func (d *Dialer) resolveInstanceName(ctx context.Context, icn string) (instance.ConnName, error) {
+	cn, err := instance.ParseConnName(icn)
+	if err != nil {
+		// The connection name was not project:region:instance
+		// Attempt to query a SRV record and see if it works instead.
+		cn, err = d.queryDNS(ctx, icn)
+		if err != nil {
+			return instance.ConnName{}, err
+		}
+	}
+
+	return cn, nil
+}
+
 // Dial returns a net.Conn connected to the specified Cloud SQL instance. The
 // icn argument must be the instance's connection name, which is in the format
 // "project-name:region:instance-name".
@@ -288,9 +315,14 @@ func (d *Dialer) Dial(ctx context.Context, icn string, opts ...DialOption) (conn
 		go trace.RecordDialError(context.Background(), icn, d.dialerID, err)
 		endDial(err)
 	}()
-	cn, err := instance.ParseConnName(icn)
+	cn, err := d.resolveInstanceName(ctx, icn)
 	if err != nil {
-		return nil, err
+		// The connection name was not project:region:instance
+		// Attempt to query a SRV record and see if it works instead.
+		cn, err = d.queryDNS(ctx, icn)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	cfg := d.defaultDialConfig
@@ -429,7 +461,7 @@ func validClientCert(
 // the instance:
 // https://cloud.google.com/sql/docs/mysql/admin-api/rest/v1beta4/SqlDatabaseVersion
 func (d *Dialer) EngineVersion(ctx context.Context, icn string) (string, error) {
-	cn, err := instance.ParseConnName(icn)
+	cn, err := d.resolveInstanceName(ctx, icn)
 	if err != nil {
 		return "", err
 	}
@@ -449,7 +481,7 @@ func (d *Dialer) EngineVersion(ctx context.Context, icn string) (string, error)
 // Use Warmup to start the refresh process early if you don't know when you'll
 // need to call "Dial".
 func (d *Dialer) Warmup(ctx context.Context, icn string, opts ...DialOption) error {
-	cn, err := instance.ParseConnName(icn)
+	cn, err := d.resolveInstanceName(ctx, icn)
 	if err != nil {
 		return err
 	}
@@ -565,3 +597,43 @@ func (d *Dialer) connectionInfoCache(
 
 	return c, nil
 }
+
+func (d *Dialer) queryDNS(ctx context.Context, domainName string) (instance.ConnName, error) {
+	// Attempt to query the SRV records.
+	// This could return a partial error where both err != nil && len(records) > 0.
+	_, records, err := d.resolver.LookupSRV(ctx, "", "", domainName)
+
+	// Process the records returning the first valid SRV record.
+
+	// Sort the record slice so that lowest priority comes first.
+	sort.Slice(records, func(i, j int) bool {
+		return records[i].Priority < records[j].Priority
+	})
+	var perr error
+	// Attempt to parse records, returning the first valid record.
+	for _, record := range records {
+		// Remove trailing '.' from target value.
+		target := strings.TrimRight(record.Target, ".")
+
+		// Parse the target as a CN
+		cn, parseErr := instance.ParseConnName(target)
+		if parseErr != nil {
+			perr = fmt.Errorf("unable to parse SRV for %s: %v", domainName, parseErr)
+			continue
+		}
+		return cn, nil
+	}
+
+	// If resolve failed and no records were found, return the error.
+	if err != nil {
+		return instance.ConnName{}, fmt.Errorf("unable to resolve SRV record for %s: %v", domainName, err)
+	}
+
+	// If all the records failed to parse, return one of the parse errors
+	if perr != nil {
+		return instance.ConnName{}, perr
+	}
+
+	// No records were found, return an error.
+	return instance.ConnName{}, fmt.Errorf("no valid SRV records found for %s", domainName)
+}

dialer_test.go

@@ -1016,3 +1016,92 @@ func TestDialerInitializesLazyCache(t *testing.T) {
 		t.Fatalf("dialer was initialized with non-lazy type: %T", tt)
 	}
 }
+
+type fakeResolver struct{}
+
+func (r *fakeResolver) LookupSRV(_ context.Context, _, _, name string) (cname string, addrs []*net.SRV, err error) {
+	// For TestDialerSuccessfullyDialsDnsSrvRecord
+	if name == "db.example.com" {
+		return "", []*net.SRV{
+			&net.SRV{Target: "my-project:my-region:my-instance."},
+		}, nil
+	}
+	if name == "db2.example.com" {
+		return "", []*net.SRV{
+			&net.SRV{Target: "my-project:my-region:my-instance"},
+		}, nil
+	}
+	// For TestDialerFailsDnsSrvRecordMalformed
+	if name == "malformed.example.com" {
+		return "", []*net.SRV{
+			&net.SRV{Target: "an-invalid-instance-name"},
+		}, nil
+	}
+	return "", nil, fmt.Errorf("no resolution for %v", name)
+}
+
+func TestDialerSuccessfullyDialsDnsSrvRecord(t *testing.T) {
+	inst := mock.NewFakeCSQLInstance(
+		"my-project", "my-region", "my-instance",
+	)
+	d := setupDialer(t, setupConfig{
+		testInstance: inst,
+		reqs: []*mock.Request{
+			mock.InstanceGetSuccess(inst, 1),
+			mock.CreateEphemeralSuccess(inst, 1),
+		},
+		dialerOptions: []Option{
+			WithTokenSource(mock.EmptyTokenSource{}),
+			WithResolver(&fakeResolver{}),
+		},
+	})
+
+	// Target has a trailing '.'
+	testSuccessfulDial(
+		context.Background(), t, d,
+		"db.example.com",
+	)
+	// Target does not have a trailing '.'
+	testSuccessfulDial(
+		context.Background(), t, d,
+		"db2.example.com",
+	)
+}
+
+func TestDialerFailsDnsSrvRecordMissing(t *testing.T) {
+	inst := mock.NewFakeCSQLInstance(
+		"my-project", "my-region", "my-instance",
+	)
+	d := setupDialer(t, setupConfig{
+		testInstance: inst,
+		reqs:         []*mock.Request{},
+		dialerOptions: []Option{
+			WithTokenSource(mock.EmptyTokenSource{}),
+			WithResolver(&fakeResolver{}),
+		},
+	})
+	_, err := d.Dial(context.Background(), "doesnt-exist.example.com")
+	wantMsg := "unable to resolve SRV record for doesnt-exist.example.com"
+	if !strings.Contains(err.Error(), wantMsg) {
+		t.Fatalf("want = %v, got = %v", wantMsg, err)
+	}
+}
+
+func TestDialerFailsDnsSrvRecordMalformed(t *testing.T) {
+	inst := mock.NewFakeCSQLInstance(
+		"my-project", "my-region", "my-instance",
+	)
+	d := setupDialer(t, setupConfig{
+		testInstance: inst,
+		reqs:         []*mock.Request{},
+		dialerOptions: []Option{
+			WithTokenSource(mock.EmptyTokenSource{}),
+			WithResolver(&fakeResolver{}),
+		},
+	})
+	_, err := d.Dial(context.Background(), "malformed.example.com")
+	wantMsg := "unable to parse SRV for malformed.example.com"
+	if !strings.Contains(err.Error(), wantMsg) {
+		t.Fatalf("want = %v, got = %v", wantMsg, err)
+	}
+}

options.go

@@ -34,6 +34,13 @@ import (
 // An Option is an option for configuring a Dialer.
 type Option func(d *dialerConfig)
 
+// NetResolver groups the methods on net.Resolver that are used by the DNS
+// resolver implementation. This allows the default net.Resolver instance to be
+// overridden from tests.
+type NetResolver interface {
+	LookupSRV(ctx context.Context, service, proto, name string) (cname string, addrs []*net.SRV, err error)
+}
+
 type dialerConfig struct {
 	rsaKey                 *rsa.PrivateKey
 	sqladminOpts           []apiopt.ClientOption
@@ -52,6 +59,7 @@ type dialerConfig struct {
 	setCredentials         bool
 	setTokenSource         bool
 	setIAMAuthNTokenSource bool
+	resolver               NetResolver
 	// err tracks any dialer options that may have failed.
 	err error
 }
@@ -234,6 +242,15 @@ func WithIAMAuthN() Option {
 	}
 }
 
+// WithResolver replaces the default DNS resolver with an alternate
+// implementation to use when resolving SRV records containing the
+// instance name. By default, the dialer will use net.DefaultResolver.
+func WithResolver(r NetResolver) Option {
+	return func(d *dialerConfig) {
+		d.resolver = r
+	}
+}
+
 type debugLoggerWithoutContext struct {
 	logger debug.Logger
 }