fix: ensure separate token source with auto-iam-authn

cmd/root.go

@@ -175,9 +175,9 @@ Instance Level Configuration
 
     When necessary, you may specify the full path to a Unix socket. Set the
     unix-socket-path query parameter to the absolute path of the Unix socket for
-    the database instance. The parent directory of the unix-socket-path must 
+    the database instance. The parent directory of the unix-socket-path must
     exist when the proxy starts or else socket creation will fail. For Postgres
-    instances, the proxy will ensure that the last path element is 
+    instances, the proxy will ensure that the last path element is
     '.s.PGSQL.5432' appending it if necessary. For example,
 
         ./cloud-sql-proxy \
@@ -371,6 +371,8 @@ func NewCommand(opts ...Option) *Command {
 		"Space separated list of additional user agents, e.g. cloud-sql-proxy-operator/0.0.1")
 	pflags.StringVarP(&c.conf.Token, "token", "t", "",
 		"Use bearer token as a source of IAM credentials.")
+	pflags.StringVar(&c.conf.LoginToken, "login-token", "",
+		"Use bearer token as a database password (used with token and auto-iam-authn only)")
 	pflags.StringVarP(&c.conf.CredentialsFile, "credentials-file", "c", "",
 		"Use service account key file as a source of IAM credentials.")
 	pflags.StringVarP(&c.conf.CredentialsJSON, "json-credentials", "j", "",
@@ -509,6 +511,15 @@ func parseConfig(cmd *Command, conf *proxy.Config, args []string) error {
 		return newBadCommandError("cannot specify --json-credentials and --gcloud-auth flags at the same time")
 	}
 
+	// When using token with auto-iam-authn, login-token must also be set.
+	// All three are required together.
+	if conf.IAMAuthN && conf.Token != "" && conf.LoginToken == "" {
+		return newBadCommandError("cannot specify --auto-iam-authn and --token without --login-token")
+	}
+	if conf.LoginToken != "" && (conf.Token == "" || !conf.IAMAuthN) {
+		return newBadCommandError("cannot specify --login-token without --token and --auto-iam-authn")
+	}
+
 	if userHasSet("http-port") && !userHasSet("prometheus") && !userHasSet("health-check") {
 		cmd.logger.Infof("Ignoring --http-port because --prometheus or --health-check was not set")
 	}

cmd/root_test.go

@@ -393,6 +393,18 @@ func TestNewCommandArguments(t *testing.T) {
 				QuitQuitQuit: true,
 			}),
 		},
+		{
+			desc: "using the login-token flag",
+			args: []string{
+				"--auto-iam-authn",
+				"--token", "MYTOK",
+				"--login-token", "MYLOGINTOKEN", "proj:region:inst"},
+			want: withDefaults(&proxy.Config{
+				IAMAuthN:   true,
+				Token:      "MYTOK",
+				LoginToken: "MYLOGINTOKEN",
+			}),
+		},
 	}
 
 	for _, tc := range tcs {
@@ -1020,6 +1032,20 @@ func TestNewCommandWithErrors(t *testing.T) {
 			desc: "using fuse-tmp-dir without fuse",
 			args: []string{"--fuse-tmp-dir", "/mydir"},
 		},
+		{
+			desc: "using --auto-iam-authn with just token flag",
+			args: []string{"--auto-iam-authn", "--token", "MYTOKEN", "p:r:i"},
+		},
+		{
+			desc: "using the --login-token without --token and --auto-iam-authn",
+			args: []string{"--login-token", "MYTOKEN", "p:r:i"},
+		},
+		{
+			desc: "using --token and --login-token without --auto-iam-authn",
+			args: []string{
+				"--token", "MYTOKEN",
+				"--login-token", "MYLOGINTOKEN", "p:r:i"},
+		},
 	}
 
 	for _, tc := range tcs {

internal/proxy/proxy.go

@@ -112,6 +112,10 @@ type Config struct {
 	// Token is the Bearer token used for authorization.
 	Token string
 
+	// LoginToken is the Bearer token used for Auto IAM AuthN. Used only in
+	// conjunction with Token.
+	LoginToken string
+
 	// CredentialsFile is the path to a service account key.
 	CredentialsFile string
 
@@ -319,30 +323,36 @@ func credentialsOpt(c Config, l cloudsql.Logger) (cloudsqlconn.Option, error) {
 	}
 
 	// Otherwise, configure credentials as usual.
+	var opt cloudsqlconn.Option
 	switch {
 	case c.Token != "":
 		l.Infof("Authorizing with OAuth2 token")
-		return cloudsqlconn.WithTokenSource(
-			oauth2.StaticTokenSource(&oauth2.Token{AccessToken: c.Token}),
-		), nil
+		ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: c.Token})
+		if c.IAMAuthN {
+			lts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: c.LoginToken})
+			opt = cloudsqlconn.WithIAMAuthNTokenSources(ts, lts)
+		} else {
+			opt = cloudsqlconn.WithTokenSource(ts)
+		}
 	case c.CredentialsFile != "":
 		l.Infof("Authorizing with the credentials file at %q", c.CredentialsFile)
-		return cloudsqlconn.WithCredentialsFile(c.CredentialsFile), nil
+		opt = cloudsqlconn.WithCredentialsFile(c.CredentialsFile)
 	case c.CredentialsJSON != "":
 		l.Infof("Authorizing with JSON credentials environment variable")
-		return cloudsqlconn.WithCredentialsJSON([]byte(c.CredentialsJSON)), nil
+		opt = cloudsqlconn.WithCredentialsJSON([]byte(c.CredentialsJSON))
 	case c.GcloudAuth:
 		l.Infof("Authorizing with gcloud user credentials")
 		ts, err := gcloud.TokenSource()
 		if err != nil {
 			return nil, err
 		}
-		return cloudsqlconn.WithTokenSource(ts), nil
+		opt = cloudsqlconn.WithTokenSource(ts)
 	default:
 		l.Infof("Authorizing with Application Default Credentials")
 		// Return no-op options to avoid having to handle nil in caller code
-		return cloudsqlconn.WithOptions(), nil
+		opt = cloudsqlconn.WithOptions()
 	}
+	return opt, nil
 }
 
 // DialerOptions builds appropriate list of options from the Config