Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ public class ImpersonatedCredentials extends GoogleCredentials
private static final long serialVersionUID = -2133257318957488431L;
private static final int TWELVE_HOURS_IN_SECONDS = 43200;
private static final int DEFAULT_LIFETIME_IN_SECONDS = 3600;
private GoogleCredentials sourceCredentials;
private final GoogleCredentials sourceCredentials;
private final String targetPrincipal;
private List<String> delegates;
private final List<String> scopes;
Expand Down Expand Up @@ -533,7 +533,18 @@ public ImpersonatedCredentials createWithCustomCalendar(Calendar calendar) {

private ImpersonatedCredentials(Builder builder) throws IOException {
super(builder);
this.sourceCredentials = builder.getSourceCredentials();
GoogleCredentials sourceCredentials = builder.getSourceCredentials();
if (sourceCredentials != null
&& !builder.sourceCredentialsScoped
&& sourceCredentials.getAccessToken() == null) {
// Apply the `CLOUD_PLATFORM_SCOPE` to access the iamcredentials endpoint
sourceCredentials =
firstNonNull(
sourceCredentials.createScoped(
Collections.singletonList(OAuth2Utils.CLOUD_PLATFORM_SCOPE)),
sourceCredentials);
Comment on lines +541 to +545

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.

medium

The use of firstNonNull here is redundant because GoogleCredentials.createScoped is guaranteed to return a non-null GoogleCredentials instance (either this or a newly created scoped instance). Additionally, since no new import for firstNonNull (e.g., from Guava's MoreObjects) was added in this file, this could lead to a compilation error if it is not already statically imported. We can safely simplify this by directly assigning the result of createScoped.

      sourceCredentials =
          sourceCredentials.createScoped(
              Collections.singletonList(OAuth2Utils.CLOUD_PLATFORM_SCOPE));

}
this.sourceCredentials = sourceCredentials;
this.targetPrincipal = builder.getTargetPrincipal();
this.delegates = builder.getDelegates();
this.scopes = ImmutableList.copyOf(builder.getScopes());
Expand Down Expand Up @@ -580,12 +591,6 @@ public String getUniverseDomain() throws IOException {

@Override
public AccessToken refreshAccessToken() throws IOException {
if (this.sourceCredentials.getAccessToken() == null) {
// Apply the `CLOUD_PLATFORM_SCOPE` to access the iamcredentials endpoint
this.sourceCredentials =
this.sourceCredentials.createScoped(
Collections.singletonList(OAuth2Utils.CLOUD_PLATFORM_SCOPE));
}

// skip for SA with SSJ flow because it uses self-signed JWT
// and will get refreshed at initialize request step
Expand Down Expand Up @@ -764,6 +769,7 @@ public static Builder newBuilder() {
public static class Builder extends GoogleCredentials.Builder {

private @Nullable GoogleCredentials sourceCredentials;
private boolean sourceCredentialsScoped;
private @Nullable String targetPrincipal;
private @Nullable List<String> delegates;
private @Nullable List<String> scopes;
Expand All @@ -789,6 +795,7 @@ protected Builder(GoogleCredentials sourceCredentials, String targetPrincipal) {
protected Builder(ImpersonatedCredentials credentials) {
super(credentials);
this.sourceCredentials = credentials.sourceCredentials;
this.sourceCredentialsScoped = true;
this.targetPrincipal = credentials.targetPrincipal;
this.delegates = credentials.delegates;
this.scopes = credentials.scopes;
Expand All @@ -800,6 +807,7 @@ protected Builder(ImpersonatedCredentials credentials) {
@CanIgnoreReturnValue
public Builder setSourceCredentials(GoogleCredentials sourceCredentials) {
this.sourceCredentials = sourceCredentials;
this.sourceCredentialsScoped = false;
return this;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,19 +47,23 @@
import com.google.api.client.http.HttpStatusCodes;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.json.GenericJson;
import com.google.api.client.json.Json;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.JsonGenerator;
import com.google.api.client.json.JsonParser;
import com.google.api.client.json.gson.GsonFactory;
import com.google.api.client.json.webtoken.JsonWebToken.Payload;
import com.google.api.client.testing.http.MockHttpTransport;
import com.google.api.client.testing.http.MockLowLevelHttpRequest;
import com.google.api.client.testing.http.MockLowLevelHttpResponse;
import com.google.api.client.util.Clock;
import com.google.auth.Credentials;
import com.google.auth.ServiceAccountSigner.SigningException;
import com.google.auth.TestUtils;
import com.google.auth.http.HttpTransportFactory;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.Uninterruptibles;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
Expand All @@ -70,9 +74,16 @@
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Calendar;
import java.util.Collection;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

Expand Down Expand Up @@ -1262,6 +1273,75 @@ void refreshAccessToken_afterSerialization_success() throws IOException, ClassNo
assertEquals(ACCESS_TOKEN, token.getTokenValue());
}

@Test
void refreshAccessToken_concurrentColdStart_scopesAndRefreshesSourceCredentialsOnce()
throws Exception {
int numThreads = 16;
AtomicInteger createScopedCount = new AtomicInteger(0);
AtomicInteger sourceRefreshCount = new AtomicInteger(0);

GoogleCredentials coldSourceCredentials =
new GoogleCredentials() {
@Override
public GoogleCredentials createScoped(Collection<String> scopes) {
createScopedCount.incrementAndGet();
return new GoogleCredentials() {
@Override
public AccessToken refreshAccessToken() {
sourceRefreshCount.incrementAndGet();
Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
return new AccessToken(
"source-token", new Date(System.currentTimeMillis() + 3600_000L));
}
};
}
};

ImpersonatedCredentials impersonatedCredentials =
(ImpersonatedCredentials)
ImpersonatedCredentials.create(
coldSourceCredentials,
IMPERSONATED_CLIENT_EMAIL,
null,
ImmutableList.of(),
VALID_LIFETIME,
() ->
new MockHttpTransport.Builder()
.setLowLevelHttpResponse(
new MockLowLevelHttpResponse()
.setContentType(Json.MEDIA_TYPE)
.setContent(
String.format(
"{\"accessToken\":\"%s\",\"expireTime\":\"%s\"}",
ACCESS_TOKEN, getDefaultExpireTime())))
.build())
.createScoped(IMMUTABLE_SCOPES_LIST);

assertEquals(1, createScopedCount.get());

CyclicBarrier barrier = new CyclicBarrier(numThreads);
ExecutorService executor = Executors.newFixedThreadPool(numThreads);
try {
List<Future<AccessToken>> futures = new ArrayList<>(numThreads);
for (int i = 0; i < numThreads; i++) {
futures.add(
executor.submit(
() -> {
barrier.await(5, TimeUnit.SECONDS);
return impersonatedCredentials.refreshAccessToken();
}));
}
for (Future<AccessToken> future : futures) {
assertEquals(ACCESS_TOKEN, future.get(10, TimeUnit.SECONDS).getTokenValue());
}
} finally {
executor.shutdownNow();
}

assertEquals(1, createScopedCount.get());
assertEquals(1, sourceRefreshCount.get());
}

public static String getDefaultExpireTime() {
return Instant.now().plusSeconds(VALID_LIFETIME).truncatedTo(ChronoUnit.SECONDS).toString();
}
Expand Down
Loading