From bd20b6bfe746cfea778b9e1a9702de28047e5950 Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Mon, 17 Oct 2022 14:53:52 -0600
Subject: [PATCH] feat: add WithAutoIP option (#346)

This commit adds an option to support a legacy behavior where a public
IP is selected if available, and otherwise the private IP is used.

Fixes 244.
---
 internal/cloudsql/instance.go      | 16 +++++++-
 internal/cloudsql/instance_test.go | 61 ++++++++++++++++++++++++++++++
 internal/cloudsql/refresh.go       |  2 +
 internal/cloudsql/refresh_test.go  |  2 +-
 internal/mock/cloudsql.go          |  4 +-
 options.go                         |  9 +++++
 6 files changed, 90 insertions(+), 4 deletions(-)

diff --git a/internal/cloudsql/instance.go b/internal/cloudsql/instance.go
index fa48f567..230cb5b6 100644
--- a/internal/cloudsql/instance.go
+++ b/internal/cloudsql/instance.go
@@ -199,7 +199,21 @@ func (i *Instance) ConnectInfo(ctx context.Context, ipType string) (string, *tls
 	if err != nil {
 		return "", nil, err
 	}
-	addr, ok := res.md.ipAddrs[ipType]
+	var (
+		addr string
+		ok   bool
+	)
+	switch ipType {
+	case AutoIP:
+		// Try Public first
+		addr, ok = res.md.ipAddrs[PublicIP]
+		if !ok {
+			// Try Private second
+			addr, ok = res.md.ipAddrs[PrivateIP]
+		}
+	default:
+		addr, ok = res.md.ipAddrs[ipType]
+	}
 	if !ok {
 		err := errtype.NewConfigError(
 			fmt.Sprintf("instance does not have IP of type %q", ipType),
diff --git a/internal/cloudsql/instance_test.go b/internal/cloudsql/instance_test.go
index a1d10515..59e92d41 100644
--- a/internal/cloudsql/instance_test.go
+++ b/internal/cloudsql/instance_test.go
@@ -149,6 +149,67 @@ func TestConnectInfo(t *testing.T) {
 	}
 }
 
+func TestConnectInfoAutoIP(t *testing.T) {
+	tcs := []struct {
+		desc   string
+		ips    []mock.FakeCSQLInstanceOption
+		wantIP string
+	}{
+		{
+			desc: "when public IP is enabled",
+			ips: []mock.FakeCSQLInstanceOption{
+				mock.WithPublicIP("8.8.8.8"),
+				mock.WithPrivateIP("10.0.0.1"),
+			},
+			wantIP: "8.8.8.8",
+		},
+		{
+			desc: "when only private IP is enabled",
+			ips: []mock.FakeCSQLInstanceOption{
+				mock.WithPrivateIP("10.0.0.1"),
+			},
+			wantIP: "10.0.0.1",
+		},
+	}
+
+	for _, tc := range tcs {
+		var opts []mock.FakeCSQLInstanceOption
+		opts = append(opts, mock.WithNoIPAddrs())
+		opts = append(opts, tc.ips...)
+		inst := mock.NewFakeCSQLInstance("p", "r", "i", opts...)
+		client, cleanup, err := mock.NewSQLAdminService(
+			context.Background(),
+			mock.InstanceGetSuccess(inst, 1),
+			mock.CreateEphemeralSuccess(inst, 1),
+		)
+		if err != nil {
+			t.Fatalf("%s", err)
+		}
+		defer func() {
+			if cErr := cleanup(); cErr != nil {
+				t.Fatalf("%v", cErr)
+			}
+		}()
+
+		i, err := NewInstance("p:r:i", client, RSAKey, 30*time.Second, nil, "", RefreshCfg{})
+		if err != nil {
+			t.Fatalf("failed to create mock instance: %v", err)
+		}
+
+		got, _, err := i.ConnectInfo(context.Background(), AutoIP)
+		if err != nil {
+			t.Fatalf("failed to retrieve connect info: %v", err)
+		}
+
+		if got != tc.wantIP {
+			t.Fatalf(
+				"ConnectInfo returned unexpected IP address, want = %v, got = %v",
+				tc.wantIP, got,
+			)
+		}
+	}
+}
+
 func TestConnectInfoErrors(t *testing.T) {
 	ctx := context.Background()
 
diff --git a/internal/cloudsql/refresh.go b/internal/cloudsql/refresh.go
index c5414747..b35a0d2c 100644
--- a/internal/cloudsql/refresh.go
+++ b/internal/cloudsql/refresh.go
@@ -36,6 +36,8 @@ const (
 	PublicIP = "PUBLIC"
 	// PrivateIP is the value for private IP Cloud SQL instances.
 	PrivateIP = "PRIVATE"
+	// AutoIP selects public IP if available and otherwise selects private IP.
+	AutoIP = "AutoIP"
 )
 
 // metadata contains information about a Cloud SQL instance needed to create connections.
diff --git a/internal/cloudsql/refresh_test.go b/internal/cloudsql/refresh_test.go
index abb65f4f..6d263398 100644
--- a/internal/cloudsql/refresh_test.go
+++ b/internal/cloudsql/refresh_test.go
@@ -272,7 +272,7 @@ func TestRefreshMetadataConfigError(t *testing.T) {
 				mock.NewFakeCSQLInstance(
 					cn.project, cn.region, cn.name,
 					mock.WithRegion("my-region"),
-					mock.WithMissingIPAddrs(),
+					mock.WithNoIPAddrs(),
 				), 1),
 			wantErr: &errtype.ConfigError{},
 			desc:    "When the instance has no supported IP addresses",
diff --git a/internal/mock/cloudsql.go b/internal/mock/cloudsql.go
index d4e5f021..a68b0610 100644
--- a/internal/mock/cloudsql.go
+++ b/internal/mock/cloudsql.go
@@ -134,9 +134,9 @@ func WithClientCertSigner(s ClientSignFunc) FakeCSQLInstanceOption {
 	}
 }
 
-// WithMissingIPAddrs configures a Fake Cloud SQL instance to have no IP
+// WithNoIPAddrs configures a Fake Cloud SQL instance to have no IP
 // addresses.
-func WithMissingIPAddrs() FakeCSQLInstanceOption {
+func WithNoIPAddrs() FakeCSQLInstanceOption {
 	return func(f *FakeCSQLInstance) {
 		f.ipAddrs = map[string]string{}
 	}
diff --git a/options.go b/options.go
index 0d13d359..f5caf661 100644
--- a/options.go
+++ b/options.go
@@ -240,6 +240,15 @@ func WithPrivateIP() DialOption {
 	}
 }
 
+// WithAutoIP returns a DialOption that selects the public IP if available and
+// otherwise falls back to private IP. This option is present for backwards
+// compatibility only and is not recommended for use in production.
+func WithAutoIP() DialOption {
+	return func(cfg *dialCfg) {
+		cfg.ipType = cloudsql.AutoIP
+	}
+}
+
 // WithDialIAMAuthN allows you to enable or disable IAM Authentication for this
 // instance as descibed in the documentation for WithIAMAuthN. This value will
 // overide the Dialer-level configuration set with WithIAMAuthN.