From fd7089db3af2b47df4275acca9ddae8916e85860 Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Mon, 31 Oct 2022 16:12:40 -0600
Subject: [PATCH] fix: impersonated user uses downscoped token

When an impersonated user logs in with Auto IAM AuthN, the login token
is downscoped to support login only.

Fixes #1519.
---
 internal/proxy/proxy.go | 17 ++++++++++++++++
 tests/postgres_test.go  | 43 ++++++++++++++++++++++++++++++++---------
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index 4546e7fef..406b32c07 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -199,6 +199,8 @@ func dialOptions(c Config, i InstanceConnConfig) []cloudsqlconn.DialOption {
 	return opts
 }
 
+const iamLoginScope = "https://www.googleapis.com/auth/sqlservice.login"
+
 func credentialsOpt(c Config, l cloudsql.Logger) (cloudsqlconn.Option, error) {
 	// If service account impersonation is configured, set up an impersonated
 	// credentials token source.
@@ -238,6 +240,21 @@ func credentialsOpt(c Config, l cloudsql.Logger) (cloudsqlconn.Option, error) {
 		if err != nil {
 			return nil, err
 		}
+		if c.IAMAuthN {
+			iamLoginTS, err := impersonate.CredentialsTokenSource(
+				context.Background(),
+				impersonate.CredentialsConfig{
+					TargetPrincipal: c.ImpersonateTarget,
+					Delegates:       c.ImpersonateDelegates,
+					Scopes:          []string{iamLoginScope},
+				},
+				iopts...,
+			)
+			if err != nil {
+				return nil, err
+			}
+			return cloudsqlconn.WithIAMAuthNTokenSources(ts, iamLoginTS), nil
+		}
 		return cloudsqlconn.WithTokenSource(ts), nil
 	}
 
diff --git a/tests/postgres_test.go b/tests/postgres_test.go
index 9e6999050..891e7ef6c 100644
--- a/tests/postgres_test.go
+++ b/tests/postgres_test.go
@@ -19,6 +19,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/GoogleCloudPlatform/cloud-sql-proxy/v2/internal/proxy"
@@ -196,16 +197,40 @@ func TestPostgresIAMDBAuthn(t *testing.T) {
 		t.Fatal("'postgres_user_iam' not set")
 	}
 
-	dsn := fmt.Sprintf("host=localhost user=%s database=%s sslmode=disable",
+	defaultDSN := fmt.Sprintf("host=localhost user=%s database=%s sslmode=disable",
 		*postgresIAMUser, *postgresDB)
-	// using the global flag
-	proxyConnTest(t,
-		[]string{"--auto-iam-authn", *postgresConnName},
-		"pgx", dsn)
-	// using the instance-level query param
-	proxyConnTest(t,
-		[]string{fmt.Sprintf("%s?auto-iam-authn=true", *postgresConnName)},
-		"pgx", dsn)
+	impersonatedIAMUser := strings.ReplaceAll(*impersonatedUser, ".gserviceaccount.com", "")
+
+	tcs := []struct {
+		desc string
+		dsn  string
+		args []string
+	}{
+		{
+			desc: "using default flag",
+			args: []string{"--auto-iam-authn", *postgresConnName},
+			dsn:  defaultDSN,
+		},
+		{
+			desc: "using query param",
+			args: []string{fmt.Sprintf("%s?auto-iam-authn=true", *postgresConnName)},
+			dsn:  defaultDSN,
+		},
+		{
+			desc: "using impersonation",
+			args: []string{
+				"--auto-iam-authn",
+				"--impersonate-service-account", *impersonatedUser,
+				*postgresConnName},
+			dsn: fmt.Sprintf("host=localhost user=%s database=%s sslmode=disable",
+				impersonatedIAMUser, *postgresDB),
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			proxyConnTest(t, tc.args, "pgx", tc.dsn)
+		})
+	}
 }
 
 func TestPostgresHealthCheck(t *testing.T) {