feat: add support for public IP (#474)

Author: jackwotherspoon

README.md

@@ -29,6 +29,7 @@ language. Using an AlloyDB connector provides the following benefits:
 ## Installation
 
 You can install this repo with `go get`:
+
 ```sh
 go get cloud.google.com/go/alloydbconn
 ```
@@ -39,9 +40,11 @@ This package provides several functions for authorizing and encrypting
 connections. These functions can be used with your database driver to connect to
 your AlloyDB instance.
 
-AlloyDB supports network connectivity through private, internal IP addresses only. 
-This package must be run in an environment that is connected to the
-[VPC Network][vpc] that hosts your AlloyDB private IP address.
+AlloyDB supports network connectivity through public IP addresses and private,
+internal IP addresses. By default this package will attempt to connect over a
+private IP connection. When doing so, this package must be run in an
+environment that is connected to the [VPC Network][vpc] that hosts your
+AlloyDB private IP address.
 
 Please see [Configuring AlloyDB Connectivity][alloydb-connectivity] for more details.
 
@@ -52,12 +55,12 @@ Please see [Configuring AlloyDB Connectivity][alloydb-connectivity] for more det
 
 This package requires the following to connect successfully:
 
-- IAM principal (user, service account, etc.) with the [AlloyDB
+* IAM principal (user, service account, etc.) with the [AlloyDB
   Client and Service Usage Consumer][client-role] roles or equivalent
   permissions. [Credentials](#credentials) for the IAM principal are
   used to authorize connections to an AlloyDB instance.
 
-- The [AlloyDB Admin API][admin-api] to be enabled within your Google Cloud
+* The [AlloyDB Admin API][admin-api] to be enabled within your Google Cloud
   Project. By default, the API will be called in the project associated with the
   IAM principal.
 
@@ -136,14 +139,14 @@ For a full list of customizable behavior, see alloydbconn.Option.
 
 ### Using DialOptions
 
-If you want to customize things about how the connection is created, use
-`DialOption`:
+If you want to customize things about how the connection is created, such as
+connecting to AlloyDB over a public IP, use a `DialOption`:
 
 ```go
 conn, err := d.Dial(
     ctx,
     "projects/<PROJECT>/locations/<REGION>/clusters/<CLUSTER>/instances/<INSTANCE>",
-    alloydbconn.WithTCPKeepAlive(30*time.Second),
+    alloydbconn.WithPublicIP(),
 )
 ```
 
@@ -154,7 +157,7 @@ be used by default:
 d, err := alloydbconn.NewDialer(
     ctx,
     alloydbconn.WithDefaultDialOptions(
-        alloydbconn.WithTCPKeepAlive(30*time.Second),
+        alloydbconn.WithPublicIP(),
     ),
 )
 ```

dialer.go

@@ -30,8 +30,8 @@ import (
 	"sync/atomic"
 	"time"
 
-	alloydbadmin "cloud.google.com/go/alloydb/apiv1beta"
-	"cloud.google.com/go/alloydb/connectors/apiv1beta/connectorspb"
+	alloydbadmin "cloud.google.com/go/alloydb/apiv1alpha"
+	"cloud.google.com/go/alloydb/connectors/apiv1alpha/connectorspb"
 	"cloud.google.com/go/alloydbconn/errtype"
 	"cloud.google.com/go/alloydbconn/internal/alloydb"
 	"cloud.google.com/go/alloydbconn/internal/trace"
@@ -75,7 +75,7 @@ func getDefaultKeys() (*rsa.PrivateKey, error) {
 
 type connectionInfoCache interface {
 	OpenConns() *uint64
-	ConnectInfo(context.Context) (string, *tls.Config, error)
+	ConnectInfo(context.Context, string) (string, *tls.Config, error)
 	ForceRefresh()
 	io.Closer
 }
@@ -156,6 +156,7 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
 	}
 
 	dialCfg := dialCfg{
+		ipType:       alloydb.PrivateIP,
 		tcpKeepAlive: defaultTCPKeepAlive,
 	}
 	for _, opt := range cfg.dialOpts {
@@ -211,7 +212,7 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 		endInfo(err)
 		return nil, err
 	}
-	addr, tlsCfg, err := i.ConnectInfo(ctx)
+	addr, tlsCfg, err := i.ConnectInfo(ctx, cfg.ipType)
 	if err != nil {
 		d.lock.Lock()
 		defer d.lock.Unlock()
@@ -231,7 +232,7 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 	if invalidClientCert(tlsCfg) {
 		i.ForceRefresh()
 		// Block on refreshed connection info
-		addr, tlsCfg, err = i.ConnectInfo(ctx)
+		addr, tlsCfg, err = i.ConnectInfo(ctx, cfg.ipType)
 		if err != nil {
 			d.lock.Lock()
 			defer d.lock.Unlock()

dialer_test.go

@@ -28,7 +28,7 @@ import (
 	"testing"
 	"time"
 
-	alloydbadmin "cloud.google.com/go/alloydb/apiv1beta"
+	alloydbadmin "cloud.google.com/go/alloydb/apiv1alpha"
 	"cloud.google.com/go/alloydbconn/errtype"
 	"cloud.google.com/go/alloydbconn/internal/alloydb"
 	"cloud.google.com/go/alloydbconn/internal/mock"
@@ -341,7 +341,7 @@ type spyConnectionInfoCache struct {
 	connectionInfoCache
 }
 
-func (s *spyConnectionInfoCache) ConnectInfo(_ context.Context) (string, *tls.Config, error) {
+func (s *spyConnectionInfoCache) ConnectInfo(_ context.Context, _ string) (string, *tls.Config, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	res := s.connectInfoCalls[s.connectInfoIndex]

e2e_test.go

@@ -227,3 +227,40 @@ func TestAutoIAMAuthN(t *testing.T) {
 	}
 	t.Log(tt)
 }
+
+func TestPublicIP(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+	ctx := context.Background()
+
+	d, err := alloydbconn.NewDialer(ctx)
+	if err != nil {
+		t.Fatalf("failed to init Dialer: %v", err)
+	}
+
+	dsn := fmt.Sprintf(
+		"user=%s password=%s dbname=%s sslmode=disable",
+		alloydbUser, alloydbPass, alloydbDB,
+	)
+	config, err := pgx.ParseConfig(dsn)
+	if err != nil {
+		t.Fatalf("failed to parse pgx config: %v", err)
+	}
+
+	config.DialFunc = func(ctx context.Context, network string, instance string) (net.Conn, error) {
+		return d.Dial(ctx, alloydbInstanceName, alloydbconn.WithPublicIP())
+	}
+
+	conn, connErr := pgx.ConnectConfig(ctx, config)
+	if connErr != nil {
+		t.Fatalf("failed to connect: %s", connErr)
+	}
+	defer conn.Close(ctx)
+
+	var tt time.Time
+	if err := conn.QueryRow(context.Background(), "SELECT NOW()").Scan(&tt); err != nil {
+		t.Fatal(err)
+	}
+	t.Log(tt)
+}

internal/alloydb/instance.go

@@ -23,7 +23,7 @@ import (
 	"sync"
 	"time"
 
-	alloydbadmin "cloud.google.com/go/alloydb/apiv1beta"
+	alloydbadmin "cloud.google.com/go/alloydb/apiv1alpha"
 	"cloud.google.com/go/alloydbconn/errtype"
 	"golang.org/x/time/rate"
 )
@@ -197,13 +197,26 @@ func (i *Instance) Close() error {
 	return nil
 }
 
-// ConnectInfo returns an IP address of the AlloyDB instance.
-func (i *Instance) ConnectInfo(ctx context.Context) (string, *tls.Config, error) {
+// ConnectInfo returns an IP address specified by ipType (i.e., public or
+// private) of the AlloyDB instance.
+func (i *Instance) ConnectInfo(ctx context.Context, ipType string) (string, *tls.Config, error) {
 	res, err := i.result(ctx)
 	if err != nil {
 		return "", nil, err
 	}
-	return res.result.instanceIPAddr, res.result.conf, nil
+	var (
+		addr string
+		ok   bool
+	)
+	addr, ok = res.result.ipAddrs[ipType]
+	if !ok {
+		err := errtype.NewConfigError(
+			fmt.Sprintf("instance does not have IP of type %q", ipType),
+			i.instanceURI.String(),
+		)
+		return "", nil, err
+	}
+	return addr, res.result.conf, nil
 }
 
 // ForceRefresh triggers an immediate refresh operation to be scheduled and

internal/alloydb/instance_test.go

@@ -22,7 +22,7 @@ import (
 	"testing"
 	"time"
 
-	alloydbadmin "cloud.google.com/go/alloydb/apiv1beta"
+	alloydbadmin "cloud.google.com/go/alloydb/apiv1alpha"
 	"cloud.google.com/go/alloydbconn/errtype"
 	"cloud.google.com/go/alloydbconn/internal/mock"
 	"golang.org/x/oauth2"
@@ -127,7 +127,7 @@ func TestConnectInfo(t *testing.T) {
 	wantAddr := "0.0.0.0"
 	inst := mock.NewFakeInstance(
 		"my-project", "my-region", "my-cluster", "my-instance",
-		mock.WithIPAddr(wantAddr),
+		mock.WithPrivateIP(wantAddr),
 	)
 	mc, url, cleanup := mock.HTTPClient(
 		mock.InstanceGetSuccess(inst, 1),
@@ -156,7 +156,7 @@ func TestConnectInfo(t *testing.T) {
 		t.Fatalf("failed to create mock instance: %v", err)
 	}
 
-	gotAddr, _, err := i.ConnectInfo(ctx)
+	gotAddr, _, err := i.ConnectInfo(ctx, PrivateIP)
 	if err != nil {
 		t.Fatalf("failed to retrieve connect info: %v", err)
 	}
@@ -190,11 +190,17 @@ func TestConnectInfoErrors(t *testing.T) {
 		t.Fatalf("failed to initialize Instance: %v", err)
 	}
 
-	_, _, err = i.ConnectInfo(ctx)
+	_, _, err = i.ConnectInfo(ctx, PrivateIP)
 	var wantErr *errtype.DialError
 	if !errors.As(err, &wantErr) {
 		t.Fatalf("when connect info fails, want = %T, got = %v", wantErr, err)
 	}
+
+	// when client asks for wrong IP address type
+	gotAddr, _, err := i.ConnectInfo(ctx, PublicIP)
+	if err == nil {
+		t.Fatalf("expected ConnectInfo to fail but returned IP address = %v", gotAddr)
+	}
 }
 
 func TestClose(t *testing.T) {
@@ -214,7 +220,7 @@ func TestClose(t *testing.T) {
 	}
 	i.Close()
 
-	_, _, err = i.ConnectInfo(ctx)
+	_, _, err = i.ConnectInfo(ctx, PrivateIP)
 	if !errors.Is(err, context.Canceled) {
 		t.Fatalf("failed to retrieve connect info: %v", err)
 	}

internal/alloydb/refresh.go

@@ -26,16 +26,23 @@ import (
 	"strings"
 	"time"
 
-	alloydbadmin "cloud.google.com/go/alloydb/apiv1beta"
-	"cloud.google.com/go/alloydb/apiv1beta/alloydbpb"
+	alloydbadmin "cloud.google.com/go/alloydb/apiv1alpha"
+	"cloud.google.com/go/alloydb/apiv1alpha/alloydbpb"
 	"cloud.google.com/go/alloydbconn/errtype"
 	"cloud.google.com/go/alloydbconn/internal/trace"
 	"google.golang.org/protobuf/types/known/durationpb"
 )
 
+const (
+	// PublicIP is the value for public IP connections.
+	PublicIP = "PUBLIC"
+	// PrivateIP is the value for private IP connections.
+	PrivateIP = "PRIVATE"
+)
+
 type connectInfo struct {
-	// ipAddr is the instance's IP addresses
-	ipAddr string
+	// ipAddrs is the instance's IP addresses
+	ipAddrs map[string]string
 	// uid is the instance UID
 	uid string
 }
@@ -56,7 +63,23 @@ func fetchMetadata(ctx context.Context, cl *alloydbadmin.AlloyDBAdminClient, ins
 	if err != nil {
 		return connectInfo{}, errtype.NewRefreshError("failed to get instance metadata", inst.String(), err)
 	}
-	return connectInfo{ipAddr: resp.IpAddress, uid: resp.InstanceUid}, nil
+
+	// parse any ip addresses that might be used to connect
+	ipAddrs := make(map[string]string)
+	if resp.GetIpAddress() != "" {
+		ipAddrs[PrivateIP] = resp.GetIpAddress()
+	}
+	if resp.GetPublicIpAddress() != "" {
+		ipAddrs[PublicIP] = resp.GetPublicIpAddress()
+	}
+
+	if len(ipAddrs) == 0 {
+		return connectInfo{}, errtype.NewConfigError(
+			"cannot connect to instance - it has no supported IP addresses",
+			inst.String(),
+		)
+	}
+	return connectInfo{ipAddrs: ipAddrs, uid: resp.InstanceUid}, nil
 }
 
 var errInvalidPEM = errors.New("certificate is not a valid PEM")
@@ -184,9 +207,9 @@ type refresher struct {
 }
 
 type refreshResult struct {
-	instanceIPAddr string
-	conf           *tls.Config
-	expiry         time.Time
+	ipAddrs map[string]string
+	conf    *tls.Config
+	expiry  time.Time
 }
 
 type certs struct {
@@ -254,9 +277,9 @@ func (r refresher) performRefresh(ctx context.Context, cn InstanceURI, k *rsa.Pr
 	c := &tls.Config{
 		Certificates: []tls.Certificate{cc.certChain},
 		RootCAs:      caCerts,
-		ServerName:   info.ipAddr,
+		ServerName:   info.ipAddrs[PrivateIP],
 		MinVersion:   tls.VersionTLS13,
 	}
 
-	return refreshResult{instanceIPAddr: info.ipAddr, conf: c, expiry: cc.expiry}, nil
+	return refreshResult{ipAddrs: info.ipAddrs, conf: c, expiry: cc.expiry}, nil
 }

internal/alloydb/refresh_test.go

@@ -20,15 +20,16 @@ import (
 	"testing"
 	"time"
 
-	alloydbadmin "cloud.google.com/go/alloydb/apiv1beta"
+	alloydbadmin "cloud.google.com/go/alloydb/apiv1alpha"
 	"cloud.google.com/go/alloydbconn/internal/mock"
 	"google.golang.org/api/option"
 )
 
 const testDialerID = "some-dialer-id"
 
 func TestRefresh(t *testing.T) {
-	wantIP := "10.0.0.1"
+	wantPrivateIP := "10.0.0.1"
+	wantPublicIP := "127.0.0.1"
 	wantExpiry := time.Now().Add(time.Hour).UTC().Round(time.Second)
 	wantInstURI := "/projects/my-project/locations/my-region/clusters/my-cluster/instances/my-instance"
 	cn, err := ParseInstURI(wantInstURI)
@@ -37,7 +38,8 @@ func TestRefresh(t *testing.T) {
 	}
 	inst := mock.NewFakeInstance(
 		"my-project", "my-region", "my-cluster", "my-instance",
-		mock.WithIPAddr(wantIP),
+		mock.WithPrivateIP(wantPrivateIP),
+		mock.WithPublicIP(wantPublicIP),
 		mock.WithCertExpiry(wantExpiry),
 	)
 	mc, url, cleanup := mock.HTTPClient(
@@ -64,8 +66,19 @@ func TestRefresh(t *testing.T) {
 		t.Fatalf("performRefresh unexpectedly failed with error: %v", err)
 	}
 
-	if got := res.instanceIPAddr; wantIP != got {
-		t.Fatalf("metadata IP mismatch, want = %v, got = %v", wantIP, got)
+	gotIP, ok := res.ipAddrs[PrivateIP]
+	if !ok {
+		t.Fatal("metadata IP addresses did not include private address")
+	}
+	if wantPrivateIP != gotIP {
+		t.Fatalf("metadata IP mismatch, want = %v, got = %v", wantPrivateIP, gotIP)
+	}
+	gotIP, ok = res.ipAddrs[PublicIP]
+	if !ok {
+		t.Fatal("metadata IP addresses did not include public address")
+	}
+	if wantPublicIP != gotIP {
+		t.Fatalf("metadata IP mismatch, want = %v, got = %v", wantPublicIP, gotIP)
 	}
 	if got := res.expiry; wantExpiry != got {
 		t.Fatalf("expiry mismatch, want = %v, got = %v", wantExpiry, got)