feat: add support for per-dial IAM AuthN config (#713)

This commit adds a new option WithDialIAMAuthN that allows callers to
enable or disable IAM AuthN on a per-dial basis, regardless of how the
underlying Dialer is configured.

Most clients will configure IAM AuthN on a dialer basis and not need
this option. However, in some advanced uses, it is necessary to enable
or disable IAM AuthN without recreating a dialer. This new option
supports that use case.

Author: enocom

dialer.go

@@ -170,7 +170,6 @@ type Dialer struct {
 	// network. By default it is golang.org/x/net/proxy#Dial.
 	dialFunc func(cxt context.Context, network, addr string) (net.Conn, error)
 
-	useIAMAuthN    bool
 	iamTokenSource oauth2.TokenSource
 	userAgent      string
 
@@ -201,6 +200,7 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
 	dialCfg := dialCfg{
 		ipType:       alloydb.PrivateIP,
 		tcpKeepAlive: defaultTCPKeepAlive,
+		iamAuthN:     cfg.useIAMAuthN, // Use dialer configuration as default
 	}
 	for _, opt := range cfg.dialOpts {
 		opt(&dialCfg)
@@ -233,7 +233,6 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
 		dialerID:                dialerID,
 		metricRecorders:         map[alloydb.InstanceURI]telv2.MetricRecorder{},
 		dialFunc:                cfg.dialFunc,
-		useIAMAuthN:             cfg.useIAMAuthN,
 		iamTokenSource:          cfg.tokenSource,
 		userAgent:               userAgent,
 		buffer:                  newBuffer(),
@@ -271,6 +270,10 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 		return nil, ErrDialerClosed
 	default:
 	}
+	cfg := d.defaultDialCfg
+	for _, opt := range opts {
+		opt(&cfg)
+	}
 
 	inst, err := alloydb.ParseInstURI(instance)
 	if err != nil {
@@ -282,7 +285,7 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 		startTime = time.Now()
 		endDial   tel.EndSpanFunc
 		attrs     = telv2.Attributes{
-			IAMAuthN:    d.useIAMAuthN,
+			IAMAuthN:    cfg.iamAuthN,
 			UserAgent:   d.userAgent,
 			RefreshType: telv2.RefreshAheadType,
 		}
@@ -300,11 +303,6 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 		endDial(err)
 	}()
 
-	cfg := d.defaultDialCfg
-	for _, opt := range opts {
-		opt(&cfg)
-	}
-
 	var endInfo tel.EndSpanFunc
 	ctx, endInfo = tel.StartSpan(ctx, "cloud.google.com/go/alloydbconn/internal.InstanceInfo")
 	cache, cacheHit, err := d.connectionInfoCache(ctx, inst, mr)
@@ -400,7 +398,7 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 	if !d.disableMetadataExchange {
 		// The metadata exchange must occur after the TLS connection is established
 		// to avoid leaking sensitive information.
-		err = d.metadataExchange(tlsConn)
+		err = d.metadataExchange(tlsConn, cfg.iamAuthN)
 		if err != nil {
 			_ = tlsConn.Close() // best effort close attempt
 			attrs.DialStatus = telv2.DialMDXError
@@ -480,13 +478,13 @@ func invalidClientCert(
 //     metadata exchange has succeeded and the connection is complete.
 //
 // Subsequent interactions with the server use the database protocol.
-func (d *Dialer) metadataExchange(conn net.Conn) error {
+func (d *Dialer) metadataExchange(conn net.Conn, useIAMAuthN bool) error {
 	tok, err := d.iamTokenSource.Token()
 	if err != nil {
 		return err
 	}
 	authType := connectorspb.MetadataExchangeRequest_DB_NATIVE
-	if d.useIAMAuthN {
+	if useIAMAuthN {
 		authType = connectorspb.MetadataExchangeRequest_AUTO_IAM
 	}
 	req := &connectorspb.MetadataExchangeRequest{

e2e_test.go

@@ -312,35 +312,83 @@ func TestAutoIAMAuthN(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping integration test")
 	}
-	ctx := context.Background()
-
-	d, err := alloydbconn.NewDialer(ctx, alloydbconn.WithIAMAuthN(), alloydbconn.WithOptOutOfBuiltInTelemetry())
-	if err != nil {
-		t.Fatalf("failed to init Dialer: %v", err)
+	tcs := []struct {
+		desc         string
+		opts         []alloydbconn.Option
+		dialOpts     []alloydbconn.DialOption
+		wantIAMAuthN bool
+	}{
+		{
+			desc: "default behavior",
+			opts: []alloydbconn.Option{
+				alloydbconn.WithIAMAuthN(),
+				alloydbconn.WithOptOutOfBuiltInTelemetry(),
+			},
+			wantIAMAuthN: true,
+		},
+		{
+			desc: "dialer with IAM authn off, dial on",
+			opts: []alloydbconn.Option{
+				alloydbconn.WithOptOutOfBuiltInTelemetry(),
+			},
+			dialOpts: []alloydbconn.DialOption{
+				alloydbconn.WithDialIAMAuthN(true),
+			},
+			wantIAMAuthN: true,
+		},
+		{
+			desc: "dialer with IAM authn on, dial off",
+			opts: []alloydbconn.Option{
+				alloydbconn.WithIAMAuthN(),
+				alloydbconn.WithOptOutOfBuiltInTelemetry(),
+			},
+			dialOpts: []alloydbconn.DialOption{
+				alloydbconn.WithDialIAMAuthN(false),
+			},
+			wantIAMAuthN: false,
+		},
 	}
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			ctx := context.Background()
 
-	dsn := fmt.Sprintf(
-		"user=%s dbname=%s sslmode=disable",
-		alloydbIAMUser, alloydbDB,
-	)
-	config, err := pgx.ParseConfig(dsn)
-	if err != nil {
-		t.Fatalf("failed to parse pgx config: %v", err)
-	}
+			d, err := alloydbconn.NewDialer(ctx, tc.opts...)
+			if err != nil {
+				t.Fatalf("failed to init Dialer: %v", err)
+			}
 
-	config.DialFunc = func(ctx context.Context, _, _ string) (net.Conn, error) {
-		return d.Dial(ctx, alloydbInstanceName)
-	}
+			var (
+				password = "ignored"
+				user     = alloydbIAMUser
+			)
+			if !tc.wantIAMAuthN {
+				user = alloydbUser
+				password = alloydbPass
+			}
+			dsn := fmt.Sprintf(
+				"user=%s dbname=%s password=%s sslmode=disable",
+				user, alloydbDB, password,
+			)
+			config, err := pgx.ParseConfig(dsn)
+			if err != nil {
+				t.Fatalf("failed to parse pgx config: %v", err)
+			}
 
-	conn, connErr := pgx.ConnectConfig(ctx, config)
-	if connErr != nil {
-		t.Fatalf("failed to connect: %s", connErr)
-	}
-	defer conn.Close(ctx)
+			config.DialFunc = func(ctx context.Context, _, _ string) (net.Conn, error) {
+				return d.Dial(ctx, alloydbInstanceName, tc.dialOpts...)
+			}
 
-	var tt time.Time
-	if err := conn.QueryRow(context.Background(), "SELECT NOW()").Scan(&tt); err != nil {
-		t.Fatal(err)
+			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)
+		})
 	}
-	t.Log(tt)
 }

options.go

@@ -364,6 +364,7 @@ type dialCfg struct {
 	dialFunc     func(ctx context.Context, network, addr string) (net.Conn, error)
 	ipType       string
 	tcpKeepAlive time.Duration
+	iamAuthN     bool
 }
 
 // DialOptions turns a list of DialOption instances into an DialOption.
@@ -415,3 +416,12 @@ func WithPSC() DialOption {
 		cfg.ipType = alloydb.PSC
 	}
 }
+
+// WithDialIAMAuthN allows calls to Dial to enable or disable IAM AuthN on a
+// one-off basis, regardless whether the dialer itself is configured with IAM
+// AuthN. There is no performance penalty to using this option.
+func WithDialIAMAuthN(enabled bool) DialOption {
+	return func(cfg *dialCfg) {
+		cfg.iamAuthN = enabled
+	}
+}