feat: add support for public IP connections

cmd/root.go

@@ -571,7 +571,8 @@ status code.`)
 		`(*) Enables Unix sockets for all listeners using the provided directory.`)
 	localFlags.BoolVarP(&c.conf.AutoIAMAuthN, "auto-iam-authn", "i", false,
 		"(*) Enables Automatic IAM Authentication for all instances")
-
+	localFlags.BoolVar(&c.conf.PublicIP, "public-ip", false,
+		"(*) Connect to the public ip address for all instances")
 	v := viper.NewWithOptions(viper.EnvKeyReplacer(strings.NewReplacer("-", "_")))
 	v.SetEnvPrefix(envPrefix)
 	v.AutomaticEnv()
@@ -776,6 +777,10 @@ func parseConfig(cmd *Command, conf *proxy.Config, args []string) error {
 			if err != nil {
 				return err
 			}
+			ic.PublicIP, err = parseBoolOpt(q, "public-ip")
+			if err != nil {
+				return err
+			}
 		}
 		ics = append(ics, ic)
 	}

cmd/root_test.go

@@ -210,6 +210,61 @@ func TestNewCommandArguments(t *testing.T) {
 				}},
 			}),
 		},
+		{
+			desc: "Public IP",
+			args: []string{
+				"--public-ip",
+				"projects/proj/locations/region/clusters/clust/instances/inst",
+			},
+			want: withDefaults(&proxy.Config{
+				PublicIP:  true,
+				Instances: []proxy.InstanceConnConfig{{Name: "projects/proj/locations/region/clusters/clust/instances/inst"}},
+			}),
+		},
+		{
+			desc: "Public IP query param (key only)",
+			args: []string{
+				"projects/proj/locations/region/clusters/clust/instances/inst?public-ip",
+			},
+			want: withDefaults(&proxy.Config{
+				Instances: []proxy.InstanceConnConfig{{
+					PublicIP: pointer(true),
+					Name:     "projects/proj/locations/region/clusters/clust/instances/inst",
+				}},
+			}),
+		},
+		{
+			desc: "Public IP query param (t & f)",
+			args: []string{
+				"projects/proj/locations/region/clusters/clust/instances/inst1?public-ip=t",
+				"projects/proj/locations/region/clusters/clust/instances/inst2?public-ip=f",
+			},
+			want: withDefaults(&proxy.Config{
+				Instances: []proxy.InstanceConnConfig{{
+					PublicIP: pointer(true),
+					Name:     "projects/proj/locations/region/clusters/clust/instances/inst1",
+				}, {
+					PublicIP: pointer(false),
+					Name:     "projects/proj/locations/region/clusters/clust/instances/inst2",
+				}},
+			}),
+		},
+		{
+			desc: "Public IP query param (true & false)",
+			args: []string{
+				"projects/proj/locations/region/clusters/clust/instances/inst1?public-ip=true",
+				"projects/proj/locations/region/clusters/clust/instances/inst2?public-ip=false",
+			},
+			want: withDefaults(&proxy.Config{
+				Instances: []proxy.InstanceConnConfig{{
+					PublicIP: pointer(true),
+					Name:     "projects/proj/locations/region/clusters/clust/instances/inst1",
+				}, {
+					PublicIP: pointer(false),
+					Name:     "projects/proj/locations/region/clusters/clust/instances/inst2",
+				}},
+			}),
+		},
 		{
 			desc: "using the address flag",
 			args: []string{"--address", "0.0.0.0", "projects/proj/locations/region/clusters/clust/instances/inst"},
@@ -480,6 +535,14 @@ func TestNewCommandWithEnvironmentConfig(t *testing.T) {
 				GcloudAuth: true,
 			}),
 		},
+		{
+			desc:     "using the public-ip envvar",
+			envName:  "ALLOYDB_PROXY_PUBLIC_IP",
+			envValue: "true",
+			want: withDefaults(&proxy.Config{
+				PublicIP: true,
+			}),
+		},
 		{
 			desc:     "using the api-endpoint envvar",
 			envName:  "ALLOYDB_PROXY_ALLOYDBADMIN_API_ENDPOINT",

go.mod

@@ -3,7 +3,7 @@ module github.com/GoogleCloudPlatform/alloydb-auth-proxy
 go 1.20
 
 require (
-	cloud.google.com/go/alloydbconn v1.5.2
+	cloud.google.com/go/alloydbconn v1.6.0
 	contrib.go.opencensus.io/exporter/prometheus v0.4.2
 	contrib.go.opencensus.io/exporter/stackdriver v0.13.14
 	github.com/coreos/go-systemd/v22 v22.5.0

go.sum

@@ -33,8 +33,8 @@ cloud.google.com/go v0.111.0 h1:YHLKNupSD1KqjDbQ3+LVdQ81h/UJbJyZG203cEfnQgM=
 cloud.google.com/go v0.111.0/go.mod h1:0mibmpKP1TyOOFYQY5izo0LnT+ecvOQ0Sg3OdmMiNRU=
 cloud.google.com/go/alloydb v1.8.0 h1:jaUQ/0e/ikQ63YOu7dtogPY0l8NcHIENMqQmlvXMZpo=
 cloud.google.com/go/alloydb v1.8.0/go.mod h1:3cVvH8uiM4VrVTKMq+hsJ8YY5RiQfXxj6gEgc8bFIgg=
-cloud.google.com/go/alloydbconn v1.5.2 h1:nphqGdcIKzvtfc5hduq9DzdORTo14+0uJC9giOoSIcY=
-cloud.google.com/go/alloydbconn v1.5.2/go.mod h1:+vDE/+2UGr4xjya2h5j2K/ZFpRLEyooOFwdNpxnmRPs=
+cloud.google.com/go/alloydbconn v1.6.0 h1:hOh4zBgQ5A4XakZn5uKADGzXnW8ULfxgg4c8oQ5lzyc=
+cloud.google.com/go/alloydbconn v1.6.0/go.mod h1:+vDE/+2UGr4xjya2h5j2K/ZFpRLEyooOFwdNpxnmRPs=
 cloud.google.com/go/bigquery v1.0.1/go.mod h1:i/xbL2UlR5RvWAURpBYZTtm/cXjCha9lbfbpx4poX+o=
 cloud.google.com/go/bigquery v1.3.0/go.mod h1:PjpwJnslEMmckchkHFfq+HTD2DmtT67aNFKH1/VBDHE=
 cloud.google.com/go/bigquery v1.4.0/go.mod h1:S8dzgnTigyfTmLBfrtrhyYhwRxG72rYxvftPBK2Dvzc=

internal/proxy/proxy.go

@@ -60,6 +60,10 @@ type InstanceConnConfig struct {
 	// AutoIAMAuthN enables automatic IAM authentication on the instance only.
 	// See Config.AutoIAMAuthN for more details.
 	AutoIAMAuthN *bool
+
+	// PublicIP tells the proxy to attempt to connect to the db instance's
+	// public IP address instead of the private IP address
+	PublicIP *bool
 }
 
 // Config contains all the configuration provided by the caller.
@@ -68,11 +72,15 @@ type Config struct {
 	// API.
 	UserAgent string
 
-	// AutoIAMAuthN enabled automatic IAM authentication which results in the
+	// AutoIAMAuthN enables automatic IAM authentication which results in the
 	// Proxy sending the IAM principal's OAuth2 token to the backend to enable
 	// a passwordless login for callers.
 	AutoIAMAuthN bool
 
+	// PublicIP enables connections via the database server's public IP address
+	// for all instances.
+	PublicIP bool
+
 	// Token is the Bearer token used for authorization.
 	Token string
 
@@ -181,6 +189,19 @@ type Config struct {
 	RunConnectionTest bool
 }
 
+// dialOptions interprets appropriate dial options for a particular instance
+// configuration
+func dialOptions(c Config, i InstanceConnConfig) []alloydbconn.DialOption {
+	var opts []alloydbconn.DialOption
+
+	// If public IP is enabled at the instance level, or public IP is enabled
+	// globally, add the option.
+	if i.PublicIP != nil && *i.PublicIP || i.PublicIP == nil && c.PublicIP {
+		opts = append(opts, alloydbconn.WithPublicIP())
+	}
+	return opts
+}
+
 func parseImpersonationChain(chain string) (string, []string) {
 	accts := strings.Split(chain, ",")
 	target := accts[0]
@@ -456,11 +477,11 @@ func (c *Client) CheckConnections(ctx context.Context) (int, error) {
 	if c.fuseDir != "" {
 		mnts = c.fuseMounts()
 	}
-	for _, m := range mnts {
+	for _, mnt := range mnts {
 		wg.Add(1)
-		go func(inst string) {
+		go func(m *socketMount) {
 			defer wg.Done()
-			conn, err := c.dialer.Dial(ctx, inst)
+			conn, err := c.dialer.Dial(ctx, m.inst, m.dialOpts...)
 			if err != nil {
 				errCh <- err
 				return
@@ -469,10 +490,10 @@ func (c *Client) CheckConnections(ctx context.Context) (int, error) {
 			if cErr != nil {
 				c.logger.Errorf(
 					"connection check failed to close connection for %v: %v",
-					inst, cErr,
+					m.inst, cErr,
 				)
 			}
-		}(m.inst)
+		}(mnt)
 	}
 	wg.Wait()
 
@@ -648,7 +669,7 @@ func (c *Client) serveSocketMount(_ context.Context, s *socketMount) error {
 			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 			defer cancel()
 
-			sConn, err := c.dialer.Dial(ctx, s.inst)
+			sConn, err := c.dialer.Dial(ctx, s.inst, s.dialOpts...)
 			if err != nil {
 				c.logger.Errorf("[%s] failed to connect to instance: %v\n", s.instShort, err)
 				cConn.Close()
@@ -664,6 +685,7 @@ type socketMount struct {
 	inst      string
 	instShort string
 	listener  net.Listener
+	dialOpts  []alloydbconn.DialOption
 }
 
 func newSocketMount(ctx context.Context, conf *Config, pc *portConfig, inst InstanceConnConfig) (*socketMount, error) {
@@ -726,11 +748,12 @@ func newSocketMount(ctx context.Context, conf *Config, pc *portConfig, inst Inst
 		// access.
 		_ = os.Chmod(address, 0777)
 	}
-
+	opts := dialOptions(*conf, inst)
 	m := &socketMount{
 		inst:      inst.Name,
 		instShort: shortInst,
 		listener:  ln,
+		dialOpts:  opts,
 	}
 	return m, nil
 }

tests/alloydb_test.go

@@ -196,3 +196,16 @@ func TestAuthWithGcloudAuth(t *testing.T) {
 		[]string{"--gcloud-auth", *alloydbInstanceName, "--port=10005"},
 		"pgx", dsn)
 }
+
+func TestPostgresPublicIP(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping Postgres integration tests")
+	}
+	requirePostgresVars(t)
+
+	dsn := fmt.Sprintf(
+		"host=127.0.0.1 port=10006 user=%v password=%v database=%v sslmode=disable",
+		*alloydbUser, *alloydbPass, *alloydbDB,
+	)
+	proxyConnTest(t, []string{*alloydbInstanceName, "--public-ip", "--port=10006"}, "pgx", dsn)
+}