Skip to content

Commit 4d4f171

Browse files
lsiraclsirac
committed
fix(auth): avoid redundant executable run and preserve impersonated email in PluggableAuthCredentials
1 parent 48e5ce5 commit 4d4f171

3 files changed

Lines changed: 89 additions & 11 deletions

File tree

‎google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,9 @@ protected ExternalAccountCredentials(ExternalAccountCredentials.Builder builder)
286286
if (serviceAccountImpersonationUrl == null) {
287287
return null;
288288
}
289+
String targetPrincipal =
290+
ImpersonatedCredentials.extractTargetPrincipal(serviceAccountImpersonationUrl);
291+
289292
// Create a copy of this instance without service account impersonation.
290293
ExternalAccountCredentials sourceCredentials;
291294
if (this instanceof AwsCredentials) {
@@ -297,6 +300,7 @@ protected ExternalAccountCredentials(ExternalAccountCredentials.Builder builder)
297300
sourceCredentials =
298301
PluggableAuthCredentials.newBuilder((PluggableAuthCredentials) this)
299302
.setServiceAccountImpersonationUrl(null)
303+
.setImpersonatedServiceAccountEmail(targetPrincipal)
300304
.build();
301305
} else {
302306
sourceCredentials =
@@ -305,8 +309,6 @@ protected ExternalAccountCredentials(ExternalAccountCredentials.Builder builder)
305309
.build();
306310
}
307311

308-
String targetPrincipal =
309-
ImpersonatedCredentials.extractTargetPrincipal(serviceAccountImpersonationUrl);
310312
return ImpersonatedCredentials.newBuilder()
311313
.setSourceCredentials(sourceCredentials)
312314
.setHttpTransportFactory(transportFactory)
@@ -524,6 +526,13 @@ private boolean shouldBuildImpersonatedCredential() {
524526
return this.serviceAccountImpersonationUrl != null && this.impersonatedCredentials == null;
525527
}
526528

529+
@Nullable ImpersonatedCredentials getImpersonatedCredentials() {
530+
if (this.shouldBuildImpersonatedCredential()) {
531+
this.impersonatedCredentials = this.buildImpersonatedCredentials();
532+
}
533+
return this.impersonatedCredentials;
534+
}
535+
527536
/**
528537
* Exchanges the external credential for a Google Cloud access token.
529538
*
@@ -534,11 +543,9 @@ private boolean shouldBuildImpersonatedCredential() {
534543
protected AccessToken exchangeExternalCredentialForAccessToken(
535544
StsTokenExchangeRequest stsTokenExchangeRequest) throws IOException {
536545
// Handle service account impersonation if necessary.
537-
if (this.shouldBuildImpersonatedCredential()) {
538-
this.impersonatedCredentials = this.buildImpersonatedCredentials();
539-
}
540-
if (this.impersonatedCredentials != null) {
541-
return this.impersonatedCredentials.refreshAccessToken();
546+
ImpersonatedCredentials impersonated = getImpersonatedCredentials();
547+
if (impersonated != null) {
548+
return impersonated.refreshAccessToken();
542549
}
543550

544551
StsRequestHandler.Builder requestHandler =

‎google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/PluggableAuthCredentials.java‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,13 @@ public class PluggableAuthCredentials extends ExternalAccountCredentials {
107107

108108
private final ExecutableHandler handler;
109109

110+
private final @Nullable String impersonatedServiceAccountEmail;
111+
110112
/** Internal constructor. See {@link Builder}. */
111113
PluggableAuthCredentials(Builder builder) {
112114
super(builder);
113115
this.config = (PluggableAuthCredentialSource) builder.credentialSource;
116+
this.impersonatedServiceAccountEmail = builder.impersonatedServiceAccountEmail;
114117

115118
if (builder.handler != null) {
116119
handler = builder.handler;
@@ -121,6 +124,10 @@ public class PluggableAuthCredentials extends ExternalAccountCredentials {
121124

122125
@Override
123126
public AccessToken refreshAccessToken() throws IOException {
127+
ImpersonatedCredentials impersonated = getImpersonatedCredentials();
128+
if (impersonated != null) {
129+
return impersonated.refreshAccessToken();
130+
}
124131
String credential = retrieveSubjectToken();
125132
StsTokenExchangeRequest.Builder stsTokenExchangeRequest =
126133
StsTokenExchangeRequest.newBuilder(credential, getSubjectTokenType())
@@ -150,8 +157,9 @@ public String retrieveSubjectToken() throws IOException {
150157
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_TOKEN_TYPE", getSubjectTokenType());
151158
// Always set to 0 for Workload Identity Federation.
152159
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_INTERACTIVE", "0");
153-
if (getServiceAccountEmail() != null) {
154-
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL", getServiceAccountEmail());
160+
String serviceAccountEmail = getServiceAccountEmail();
161+
if (serviceAccountEmail != null) {
162+
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL", serviceAccountEmail);
155163
}
156164
if (outputFilePath != null && !outputFilePath.isEmpty()) {
157165
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_OUTPUT_FILE", outputFilePath);
@@ -185,6 +193,15 @@ public String getOutputFilePath() {
185193
return this.handler.retrieveTokenFromExecutable(options);
186194
}
187195

196+
@Override
197+
public @Nullable String getServiceAccountEmail() {
198+
String email = super.getServiceAccountEmail();
199+
if (email != null) {
200+
return email;
201+
}
202+
return impersonatedServiceAccountEmail;
203+
}
204+
188205
/** Clones the PluggableAuthCredentials with the specified scopes. */
189206
@Override
190207
public PluggableAuthCredentials createScoped(Collection<String> newScopes) {
@@ -217,12 +234,14 @@ ExecutableHandler getExecutableHandler() {
217234
public static class Builder extends ExternalAccountCredentials.Builder {
218235

219236
private @Nullable ExecutableHandler handler;
237+
private @Nullable String impersonatedServiceAccountEmail;
220238

221239
Builder() {}
222240

223241
Builder(PluggableAuthCredentials credentials) {
224242
super(credentials);
225243
this.handler = credentials.handler;
244+
this.impersonatedServiceAccountEmail = credentials.impersonatedServiceAccountEmail;
226245
}
227246

228247
@CanIgnoreReturnValue
@@ -277,6 +296,13 @@ public Builder setCredentialSource(PluggableAuthCredentialSource credentialSourc
277296
public Builder setServiceAccountImpersonationUrl(
278297
@Nullable String serviceAccountImpersonationUrl) {
279298
super.setServiceAccountImpersonationUrl(serviceAccountImpersonationUrl);
299+
this.impersonatedServiceAccountEmail = null;
300+
return this;
301+
}
302+
303+
@CanIgnoreReturnValue
304+
Builder setImpersonatedServiceAccountEmail(@Nullable String impersonatedServiceAccountEmail) {
305+
this.impersonatedServiceAccountEmail = impersonatedServiceAccountEmail;
280306
return this;
281307
}
282308

‎google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/PluggableAuthCredentialsTest.java‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,15 @@ void refreshAccessToken_withServiceAccountImpersonation() throws IOException {
212212

213213
transportFactory.transport.setExpireTime(TestUtils.getDefaultExpireTime());
214214

215+
final int[] invocationCount = {0};
216+
final ExecutableOptions[] providedOptions = {null};
217+
ExecutableHandler executableHandler =
218+
options -> {
219+
invocationCount[0]++;
220+
providedOptions[0] = options;
221+
return "pluggableAuthToken";
222+
};
223+
215224
PluggableAuthCredentials credential =
216225
PluggableAuthCredentials.newBuilder()
217226
.setAudience(
@@ -227,11 +236,15 @@ void refreshAccessToken_withServiceAccountImpersonation() throws IOException {
227236

228237
credential =
229238
PluggableAuthCredentials.newBuilder(credential)
230-
.setExecutableHandler(options -> "pluggableAuthToken")
239+
.setExecutableHandler(executableHandler)
231240
.build();
232241

233242
AccessToken accessToken = credential.refreshAccessToken();
234243

244+
assertEquals(1, invocationCount[0]);
245+
assertEquals(
246+
credential.getServiceAccountEmail(),
247+
providedOptions[0].getEnvironmentMap().get("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL"));
235248
assertEquals(
236249
transportFactory.transport.getServiceAccountAccessToken(), accessToken.getTokenValue());
237250

@@ -253,6 +266,15 @@ void refreshAccessToken_withServiceAccountImpersonationOptions() throws IOExcept
253266

254267
transportFactory.transport.setExpireTime(TestUtils.getDefaultExpireTime());
255268

269+
final int[] invocationCount = {0};
270+
final ExecutableOptions[] providedOptions = {null};
271+
ExecutableHandler executableHandler =
272+
options -> {
273+
invocationCount[0]++;
274+
providedOptions[0] = options;
275+
return "pluggableAuthToken";
276+
};
277+
256278
PluggableAuthCredentials credential =
257279
PluggableAuthCredentials.newBuilder()
258280
.setAudience(
@@ -270,11 +292,15 @@ void refreshAccessToken_withServiceAccountImpersonationOptions() throws IOExcept
270292

271293
credential =
272294
PluggableAuthCredentials.newBuilder(credential)
273-
.setExecutableHandler(options -> "pluggableAuthToken")
295+
.setExecutableHandler(executableHandler)
274296
.build();
275297

276298
AccessToken accessToken = credential.refreshAccessToken();
277299

300+
assertEquals(1, invocationCount[0]);
301+
assertEquals(
302+
credential.getServiceAccountEmail(),
303+
providedOptions[0].getEnvironmentMap().get("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL"));
278304
assertEquals(
279305
transportFactory.transport.getServiceAccountAccessToken(), accessToken.getTokenValue());
280306

@@ -585,6 +611,25 @@ void createdScoped_clonedCredentialWithAddedScopes() {
585611
assertEquals("universeDomain", newCredentials.getUniverseDomain());
586612
}
587613

614+
@Test
615+
void createScoped_preservesImpersonatedServiceAccountEmail() {
616+
PluggableAuthCredentials sourceCredentials =
617+
PluggableAuthCredentials.newBuilder(CREDENTIAL)
618+
.setServiceAccountImpersonationUrl(null)
619+
.setImpersonatedServiceAccountEmail("testn@test.iam.gserviceaccount.com")
620+
.build();
621+
622+
PluggableAuthCredentials scopedCredentials =
623+
sourceCredentials.createScoped(Arrays.asList("scope1"));
624+
625+
assertNull(scopedCredentials.getServiceAccountImpersonationUrl());
626+
assertEquals("testn@test.iam.gserviceaccount.com", scopedCredentials.getServiceAccountEmail());
627+
628+
PluggableAuthCredentials clearedCredentials =
629+
scopedCredentials.toBuilder().setServiceAccountImpersonationUrl(null).build();
630+
assertNull(clearedCredentials.getServiceAccountEmail());
631+
}
632+
588633
@Test
589634
void serialize() {
590635
PluggableAuthCredentials testCredentials =

0 commit comments

Comments
 (0)