google/internal/externalaccount: allow impersonation lifetime changes

Right now, impersonation tokens used for external accounts have a hardcoded lifetime of 1 hour (3600 seconds), but some of our customers want to be able to adjust this lifetime.  These changes (along with others in the gcloud cli) should allow this

Change-Id: I705f83dc2a092d8cdd0fcbfff83b014c220e28bb
GitHub-Last-Rev: 7e0ea92c8ef5f12b4a86ec5b389ff7a2055ad2ab
GitHub-Pull-Request: golang/oauth2#571
Reviewed-on: https://go-review.googlesource.com/c/oauth2/+/416797
Reviewed-by: Cody Oss <codyoss@google.com>
Reviewed-by: Shin Fan <shinfan@google.com>
Run-TryBot: Cody Oss <codyoss@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

4 files changed, 104 insertions, 48 deletions

diff --git a/google/google.go b/google/google.go
index ceddd5d..8df0c49 100644
--- a/google/google.go
+++ b/google/google.go
@@ -122,6 +122,7 @@ type credentialsFile struct {
 	TokenURLExternal               string                           `json:"token_url"`
 	TokenInfoURL                   string                           `json:"token_info_url"`
 	ServiceAccountImpersonationURL string                           `json:"service_account_impersonation_url"`
+	ServiceAccountImpersonation    serviceAccountImpersonationInfo  `json:"service_account_impersonation"`
 	Delegates                      []string                         `json:"delegates"`
 	CredentialSource               externalaccount.CredentialSource `json:"credential_source"`
 	QuotaProjectID                 string                           `json:"quota_project_id"`
@@ -131,6 +132,10 @@ type credentialsFile struct {
 	SourceCredentials *credentialsFile `json:"source_credentials"`
 }
 
+type serviceAccountImpersonationInfo struct {
+	TokenLifetimeSeconds int `json:"token_lifetime_seconds"`
+}
+
 func (f *credentialsFile) jwtConfig(scopes []string, subject string) *jwt.Config {
 	cfg := &jwt.Config{
 		Email:        f.ClientEmail,
@@ -178,12 +183,13 @@ func (f *credentialsFile) tokenSource(ctx context.Context, params CredentialsPar
 			TokenURL:                       f.TokenURLExternal,
 			TokenInfoURL:                   f.TokenInfoURL,
 			ServiceAccountImpersonationURL: f.ServiceAccountImpersonationURL,
-			ClientSecret:                   f.ClientSecret,
-			ClientID:                       f.ClientID,
-			CredentialSource:               f.CredentialSource,
-			QuotaProjectID:                 f.QuotaProjectID,
-			Scopes:                         params.Scopes,
-			WorkforcePoolUserProject:       f.WorkforcePoolUserProject,
+			ServiceAccountImpersonationLifetimeSeconds: f.ServiceAccountImpersonation.TokenLifetimeSeconds,
+			ClientSecret:             f.ClientSecret,
+			ClientID:                 f.ClientID,
+			CredentialSource:         f.CredentialSource,
+			QuotaProjectID:           f.QuotaProjectID,
+			Scopes:                   params.Scopes,
+			WorkforcePoolUserProject: f.WorkforcePoolUserProject,
 		}
 		return cfg.TokenSource(ctx)
 	case impersonatedServiceAccount:
diff --git a/google/internal/externalaccount/basecredentials.go b/google/internal/externalaccount/basecredentials.go
index b3d5fe2..2bf5391 100644
--- a/google/internal/externalaccount/basecredentials.go
+++ b/google/internal/externalaccount/basecredentials.go
@@ -39,6 +39,9 @@ type Config struct {
 	// ServiceAccountImpersonationURL is the URL for the service account impersonation request. This is only
 	// required for workload identity pools when APIs to be accessed have not integrated with UberMint.
 	ServiceAccountImpersonationURL string
+	// ServiceAccountImpersonationLifetimeSeconds is the number of seconds the service account impersonation
+	// token will be valid for.
+	ServiceAccountImpersonationLifetimeSeconds int
 	// ClientSecret is currently only required if token_info endpoint also
 	// needs to be called with the generated GCP access token. When provided, STS will be
 	// called with additional basic authentication using client_id as username and client_secret as password.
@@ -141,10 +144,11 @@ func (c *Config) tokenSource(ctx context.Context, tokenURLValidPats []*regexp.Re
 	scopes := c.Scopes
 	ts.conf.Scopes = []string{"https://www.googleapis.com/auth/cloud-platform"}
 	imp := ImpersonateTokenSource{
-		Ctx:    ctx,
-		URL:    c.ServiceAccountImpersonationURL,
-		Scopes: scopes,
-		Ts:     oauth2.ReuseTokenSource(nil, ts),
+		Ctx:                  ctx,
+		URL:                  c.ServiceAccountImpersonationURL,
+		Scopes:               scopes,
+		Ts:                   oauth2.ReuseTokenSource(nil, ts),
+		TokenLifetimeSeconds: c.ServiceAccountImpersonationLifetimeSeconds,
 	}
 	return oauth2.ReuseTokenSource(nil, imp), nil
 }
diff --git a/google/internal/externalaccount/impersonate.go b/google/internal/externalaccount/impersonate.go
index 8251fc8..54c8f20 100644
--- a/google/internal/externalaccount/impersonate.go
+++ b/google/internal/externalaccount/impersonate.go
@@ -48,12 +48,19 @@ type ImpersonateTokenSource struct {
 	// Each service account must be granted roles/iam.serviceAccountTokenCreator
 	// on the next service account in the chain. Optional.
 	Delegates []string
+	// TokenLifetimeSeconds is the number of seconds the impersonation token will
+	// be valid for.
+	TokenLifetimeSeconds int
 }
 
 // Token performs the exchange to get a temporary service account token to allow access to GCP.
 func (its ImpersonateTokenSource) Token() (*oauth2.Token, error) {
+	lifetimeString := "3600s"
+	if its.TokenLifetimeSeconds != 0 {
+		lifetimeString = fmt.Sprintf("%ds", its.TokenLifetimeSeconds)
+	}
 	reqBody := generateAccessTokenReq{
-		Lifetime:  "3600s",
+		Lifetime:  lifetimeString,
 		Scope:     its.Scopes,
 		Delegates: its.Delegates,
 	}
diff --git a/google/internal/externalaccount/impersonate_test.go b/google/internal/externalaccount/impersonate_test.go
index 6fed7b9..17e2f6d 100644
--- a/google/internal/externalaccount/impersonate_test.go
+++ b/google/internal/externalaccount/impersonate_test.go
@@ -13,28 +13,18 @@ import (
 	"testing"
 )
 
-var testImpersonateConfig = Config{
-	Audience:         "32555940559.apps.googleusercontent.com",
-	SubjectTokenType: "urn:ietf:params:oauth:token-type:jwt",
-	TokenInfoURL:     "http://localhost:8080/v1/tokeninfo",
-	ClientSecret:     "notsosecret",
-	ClientID:         "rbrgnognrhongo3bi4gb9ghg9g",
-	CredentialSource: testBaseCredSource,
-	Scopes:           []string{"https://www.googleapis.com/auth/devstorage.full_control"},
-}
-
 var (
 	baseImpersonateCredsReqBody  = "audience=32555940559.apps.googleusercontent.com&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange&requested_token_type=urn%3Aietf%3Aparams%3Aoauth%3Atoken-type%3Aaccess_token&scope=https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fcloud-platform&subject_token=street123&subject_token_type=urn%3Aietf%3Aparams%3Aoauth%3Atoken-type%3Ajwt"
 	baseImpersonateCredsRespBody = `{"accessToken":"Second.Access.Token","expireTime":"2020-12-28T15:01:23Z"}`
 )
 
-func TestImpersonation(t *testing.T) {
-	impersonateServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if got, want := r.URL.String(), "/"; got != want {
+func createImpersonationServer(urlWanted, authWanted, bodyWanted, response string, t *testing.T) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if got, want := r.URL.String(), urlWanted; got != want {
 			t.Errorf("URL.String(): got %v but want %v", got, want)
 		}
 		headerAuth := r.Header.Get("Authorization")
-		if got, want := headerAuth, "Bearer Sample.Access.Token"; got != want {
+		if got, want := headerAuth, authWanted; got != want {
 			t.Errorf("got %v but want %v", got, want)
 		}
 		headerContentType := r.Header.Get("Content-Type")
@@ -45,14 +35,16 @@ func TestImpersonation(t *testing.T) {
 		if err != nil {
 			t.Fatalf("Failed reading request body: %v.", err)
 		}
-		if got, want := string(body), "{\"lifetime\":\"3600s\",\"scope\":[\"https://www.googleapis.com/auth/devstorage.full_control\"]}"; got != want {
+		if got, want := string(body), bodyWanted; got != want {
 			t.Errorf("Unexpected impersonation payload: got %v but want %v", got, want)
 		}
 		w.Header().Set("Content-Type", "application/json")
-		w.Write([]byte(baseImpersonateCredsRespBody))
+		w.Write([]byte(response))
 	}))
-	testImpersonateConfig.ServiceAccountImpersonationURL = impersonateServer.URL
-	targetServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+}
+
+func createTargetServer(t *testing.T) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if got, want := r.URL.String(), "/"; got != want {
 			t.Errorf("URL.String(): got %v but want %v", got, want)
 		}
@@ -74,27 +66,74 @@ func TestImpersonation(t *testing.T) {
 		w.Header().Set("Content-Type", "application/json")
 		w.Write([]byte(baseCredsResponseBody))
 	}))
-	defer targetServer.Close()
+}
 
-	testImpersonateConfig.TokenURL = targetServer.URL
-	allURLs := regexp.MustCompile(".+")
-	ourTS, err := testImpersonateConfig.tokenSource(context.Background(), []*regexp.Regexp{allURLs}, []*regexp.Regexp{allURLs}, "http")
-	if err != nil {
-		t.Fatalf("Failed to create TokenSource: %v", err)
-	}
+var impersonationTests = []struct {
+	name                      string
+	config                    Config
+	expectedImpersonationBody string
+}{
+	{
+		name: "Base Impersonation",
+		config: Config{
+			Audience:         "32555940559.apps.googleusercontent.com",
+			SubjectTokenType: "urn:ietf:params:oauth:token-type:jwt",
+			TokenInfoURL:     "http://localhost:8080/v1/tokeninfo",
+			ClientSecret:     "notsosecret",
+			ClientID:         "rbrgnognrhongo3bi4gb9ghg9g",
+			CredentialSource: testBaseCredSource,
+			Scopes:           []string{"https://www.googleapis.com/auth/devstorage.full_control"},
+		},
+		expectedImpersonationBody: "{\"lifetime\":\"3600s\",\"scope\":[\"https://www.googleapis.com/auth/devstorage.full_control\"]}",
+	},
+	{
+		name: "With TokenLifetime Set",
+		config: Config{
+			Audience:         "32555940559.apps.googleusercontent.com",
+			SubjectTokenType: "urn:ietf:params:oauth:token-type:jwt",
+			TokenInfoURL:     "http://localhost:8080/v1/tokeninfo",
+			ClientSecret:     "notsosecret",
+			ClientID:         "rbrgnognrhongo3bi4gb9ghg9g",
+			CredentialSource: testBaseCredSource,
+			Scopes:           []string{"https://www.googleapis.com/auth/devstorage.full_control"},
+			ServiceAccountImpersonationLifetimeSeconds: 10000,
+		},
+		expectedImpersonationBody: "{\"lifetime\":\"10000s\",\"scope\":[\"https://www.googleapis.com/auth/devstorage.full_control\"]}",
+	},
+}
 
-	oldNow := now
-	defer func() { now = oldNow }()
-	now = testNow
+func TestImpersonation(t *testing.T) {
+	for _, tt := range impersonationTests {
+		t.Run(tt.name, func(t *testing.T) {
+			testImpersonateConfig := tt.config
+			impersonateServer := createImpersonationServer("/", "Bearer Sample.Access.Token", tt.expectedImpersonationBody, baseImpersonateCredsRespBody, t)
+			defer impersonateServer.Close()
+			testImpersonateConfig.ServiceAccountImpersonationURL = impersonateServer.URL
 
-	tok, err := ourTS.Token()
-	if err != nil {
-		t.Fatalf("Unexpected error: %e", err)
-	}
-	if got, want := tok.AccessToken, "Second.Access.Token"; got != want {
-		t.Errorf("Unexpected access token: got %v, but wanted %v", got, want)
-	}
-	if got, want := tok.TokenType, "Bearer"; got != want {
-		t.Errorf("Unexpected TokenType: got %v, but wanted %v", got, want)
+			targetServer := createTargetServer(t)
+			defer targetServer.Close()
+			testImpersonateConfig.TokenURL = targetServer.URL
+
+			allURLs := regexp.MustCompile(".+")
+			ourTS, err := testImpersonateConfig.tokenSource(context.Background(), []*regexp.Regexp{allURLs}, []*regexp.Regexp{allURLs}, "http")
+			if err != nil {
+				t.Fatalf("Failed to create TokenSource: %v", err)
+			}
+
+			oldNow := now
+			defer func() { now = oldNow }()
+			now = testNow
+
+			tok, err := ourTS.Token()
+			if err != nil {
+				t.Fatalf("Unexpected error: %e", err)
+			}
+			if got, want := tok.AccessToken, "Second.Access.Token"; got != want {
+				t.Errorf("Unexpected access token: got %v, but wanted %v", got, want)
+			}
+			if got, want := tok.TokenType, "Bearer"; got != want {
+				t.Errorf("Unexpected TokenType: got %v, but wanted %v", got, want)
+			}
+		})
 	}
 }