Skip to content
Open
Original file line number Diff line number Diff line change
Expand Up @@ -107,10 +107,16 @@ public class PluggableAuthCredentials extends ExternalAccountCredentials {

private final ExecutableHandler handler;

private final @Nullable String impersonatedServiceAccountEmail;

/** Internal constructor. See {@link Builder}. */
PluggableAuthCredentials(Builder builder) {
super(builder);
this.config = (PluggableAuthCredentialSource) builder.credentialSource;
this.impersonatedServiceAccountEmail =
getServiceAccountEmail() != null
? getServiceAccountEmail()
: builder.impersonatedServiceAccountEmail;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: I believe that falling back to builder.impersonatedServiceAccountEmail whenever getServiceAccountEmail() == null means that an external caller doing creds.toBuilder().setServiceAccountImpersonationUrl(null).build() retains
the stale impersonatedServiceAccountEmail and still sends GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL in retrieveSubjectToken().

Perhaps we want to address this at the same time as b/563846018 since it is very similar in nature? We'd just to make sure we also take care of resetting impersonatedServiceAccountEmail as well.


if (builder.handler != null) {
handler = builder.handler;
Expand All @@ -121,6 +127,12 @@ public class PluggableAuthCredentials extends ExternalAccountCredentials {

@Override
public AccessToken refreshAccessToken() throws IOException {
if (getServiceAccountImpersonationUrl() != null) {
if (this.impersonatedCredentials == null) {
this.impersonatedCredentials = this.buildImpersonatedCredentials();
}
return this.impersonatedCredentials.refreshAccessToken();
}
String credential = retrieveSubjectToken();
StsTokenExchangeRequest.Builder stsTokenExchangeRequest =
StsTokenExchangeRequest.newBuilder(credential, getSubjectTokenType())
Expand Down Expand Up @@ -150,8 +162,8 @@ public String retrieveSubjectToken() throws IOException {
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_TOKEN_TYPE", getSubjectTokenType());
// Always set to 0 for Workload Identity Federation.
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_INTERACTIVE", "0");
if (getServiceAccountEmail() != null) {
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL", getServiceAccountEmail());
if (impersonatedServiceAccountEmail != null) {
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL", impersonatedServiceAccountEmail);
}
if (outputFilePath != null && !outputFilePath.isEmpty()) {
envMap.put("GOOGLE_EXTERNAL_ACCOUNT_OUTPUT_FILE", outputFilePath);
Expand Down Expand Up @@ -217,12 +229,14 @@ ExecutableHandler getExecutableHandler() {
public static class Builder extends ExternalAccountCredentials.Builder {

private @Nullable ExecutableHandler handler;
private @Nullable String impersonatedServiceAccountEmail;

Builder() {}

Builder(PluggableAuthCredentials credentials) {
super(credentials);
this.handler = credentials.handler;
this.impersonatedServiceAccountEmail = credentials.impersonatedServiceAccountEmail;
}

@CanIgnoreReturnValue
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import static com.google.auth.oauth2.MockExternalAccountCredentialsTransport.SERVICE_ACCOUNT_IMPERSONATION_URL;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertThrows;

import com.google.api.client.http.HttpTransport;
Expand All @@ -48,6 +49,7 @@
import java.io.NotSerializableException;
import java.math.BigDecimal;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand All @@ -63,6 +65,7 @@ class PluggableAuthCredentialsTest extends BaseSerializationTest {
// The maximum timeout for waiting for the executable to finish (120 seconds).
private static final int MAXIMUM_EXECUTABLE_TIMEOUT_MS = 120 * 1000;
private static final String STS_URL = "https://sts.googleapis.com";
private static final String IMPERSONATED_EMAIL = "testn@test.iam.gserviceaccount.com";

private static final PluggableAuthCredentials CREDENTIAL =
PluggableAuthCredentials.newBuilder()
Expand Down Expand Up @@ -212,6 +215,15 @@ void refreshAccessToken_withServiceAccountImpersonation() throws IOException {

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

final int[] invocationCount = {0};
final ExecutableOptions[] providedOptions = {null};
ExecutableHandler executableHandler =
options -> {
invocationCount[0]++;
providedOptions[0] = options;
return "pluggableAuthToken";
};

PluggableAuthCredentials credential =
PluggableAuthCredentials.newBuilder()
.setAudience(
Expand All @@ -220,18 +232,19 @@ void refreshAccessToken_withServiceAccountImpersonation() throws IOException {
.setTokenInfoUrl("tokenInfoUrl")
.setTokenUrl(transportFactory.transport.getStsUrl())
.setCredentialSource(buildCredentialSource())
.setExecutableHandler(executableHandler)
.setServiceAccountImpersonationUrl(
transportFactory.transport.getServiceAccountImpersonationUrl())
.setHttpTransportFactory(transportFactory)
.build();

credential =
PluggableAuthCredentials.newBuilder(credential)
.setExecutableHandler(options -> "pluggableAuthToken")
.build();

AccessToken accessToken = credential.refreshAccessToken();

// Validate that the executable was invoked once with the impersonated email.
assertEquals(1, invocationCount[0]);
assertEquals(
IMPERSONATED_EMAIL,
providedOptions[0].getEnvironmentMap().get("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL"));
assertEquals(
transportFactory.transport.getServiceAccountAccessToken(), accessToken.getTokenValue());

Expand All @@ -244,6 +257,11 @@ void refreshAccessToken_withServiceAccountImpersonation() throws IOException {
Map<String, List<String>> headers =
transportFactory.transport.getRequests().get(0).getHeaders();
ExternalAccountCredentialsTest.validateMetricsHeader(headers, "executable", true, false);

// Validate that refreshing a second time reuses cached impersonatedCredentials and does not
// re-invoke the executable while the source STS token is still unexpired.
credential.refreshAccessToken();
assertEquals(1, invocationCount[0]);
}

@Test
Expand All @@ -253,6 +271,15 @@ void refreshAccessToken_withServiceAccountImpersonationOptions() throws IOExcept

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

final int[] invocationCount = {0};
final ExecutableOptions[] providedOptions = {null};
ExecutableHandler executableHandler =
options -> {
invocationCount[0]++;
providedOptions[0] = options;
return "pluggableAuthToken";
};

PluggableAuthCredentials credential =
PluggableAuthCredentials.newBuilder()
.setAudience(
Expand All @@ -261,20 +288,21 @@ void refreshAccessToken_withServiceAccountImpersonationOptions() throws IOExcept
.setTokenInfoUrl("tokenInfoUrl")
.setTokenUrl(transportFactory.transport.getStsUrl())
.setCredentialSource(buildCredentialSource())
.setExecutableHandler(executableHandler)
.setServiceAccountImpersonationUrl(
transportFactory.transport.getServiceAccountImpersonationUrl())
.setServiceAccountImpersonationOptions(
ExternalAccountCredentialsTest.buildServiceAccountImpersonationOptions())
.setHttpTransportFactory(transportFactory)
.build();

credential =
PluggableAuthCredentials.newBuilder(credential)
.setExecutableHandler(options -> "pluggableAuthToken")
.build();

AccessToken accessToken = credential.refreshAccessToken();

// Validate that the executable was invoked once with the impersonated email.
assertEquals(1, invocationCount[0]);
assertEquals(
IMPERSONATED_EMAIL,
providedOptions[0].getEnvironmentMap().get("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL"));
assertEquals(
transportFactory.transport.getServiceAccountAccessToken(), accessToken.getTokenValue());

Expand All @@ -292,6 +320,36 @@ void refreshAccessToken_withServiceAccountImpersonationOptions() throws IOExcept
ExternalAccountCredentialsTest.validateMetricsHeader(headers, "executable", true, true);
}

@Test
void refreshAccessToken_withServiceAccountImpersonation_executableFailure() {
MockExternalAccountCredentialsTransportFactory transportFactory =
new MockExternalAccountCredentialsTransportFactory();

PluggableAuthException expectedException =
new PluggableAuthException("INVALID_EXECUTABLE", "Executable failed.");
ExecutableHandler executableHandler =
options -> {
throw expectedException;
};

PluggableAuthCredentials credential =
PluggableAuthCredentials.newBuilder()
.setAudience(
"//iam.googleapis.com/projects/123/locations/global/workloadIdentityPools/pool/providers/provider")
.setSubjectTokenType("subjectTokenType")
.setTokenInfoUrl("tokenInfoUrl")
.setTokenUrl(transportFactory.transport.getStsUrl())
.setCredentialSource(buildCredentialSource())
.setExecutableHandler(executableHandler)
.setServiceAccountImpersonationUrl(
transportFactory.transport.getServiceAccountImpersonationUrl())
.setHttpTransportFactory(transportFactory)
.build();

IOException exception = assertThrows(IOException.class, credential::refreshAccessToken);
assertSame(expectedException, exception.getCause());
}

@Test
void pluggableAuthCredentialSource_allFields() {
Map<String, Object> source = new HashMap<>();
Expand Down Expand Up @@ -585,6 +643,35 @@ void createdScoped_clonedCredentialWithAddedScopes() {
assertEquals("universeDomain", newCredentials.getUniverseDomain());
}

@Test
void createScoped_preservesImpersonatedServiceAccountEmail() throws IOException {
final String[] recordedImpersonatedEmail = {null};
PluggableAuthCredentials outerCredentials =
PluggableAuthCredentials.newBuilder(CREDENTIAL)
.setExecutableHandler(
options -> {
recordedImpersonatedEmail[0] =
options.getEnvironmentMap().get("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL");
return "pluggableAuthToken";
})
.setServiceAccountImpersonationUrl(
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/"
+ IMPERSONATED_EMAIL
+ ":generateAccessToken")
.build();
PluggableAuthCredentials sourceCredentials =
(PluggableAuthCredentials)
outerCredentials.buildImpersonatedCredentials().getSourceCredentials();

PluggableAuthCredentials scopedCredentials =
sourceCredentials.createScoped(Collections.singletonList("scope1"));
scopedCredentials.retrieveSubjectToken();

assertNull(scopedCredentials.getServiceAccountImpersonationUrl());
assertNull(scopedCredentials.getServiceAccountEmail());
assertEquals(IMPERSONATED_EMAIL, recordedImpersonatedEmail[0]);
}

@Test
void serialize() {
PluggableAuthCredentials testCredentials =
Expand Down
Loading