Skip to content
Open
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 @@ -72,6 +72,7 @@
import java.util.Objects;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.regex.Pattern;

/**
* OAuth2 credentials representing the built-in service account for a Google Compute Engine VM.
Expand Down Expand Up @@ -117,6 +118,7 @@ public class ComputeEngineCredentials extends GoogleCredentials

private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
private static final String PARSE_ERROR_ACCOUNT = "Error parsing service account response. ";
private static final Pattern EMAIL_PATTERN = Pattern.compile("^[^@]+@[^@]+\\.[^@]+$");
private static final long serialVersionUID = -4113476462526554235L;

private final String transportFactoryClassName;
Expand Down Expand Up @@ -800,8 +802,20 @@ public String getAccount() {
@InternalApi
@Override
public String getRegionalAccessBoundaryUrl() throws IOException {
String account = getAccount();
// In GKE environments, the default service account might return a non-email placeholder.
// Since RAB lookup requires a valid email-based service account, we skip RAB lookup
// in non-email scenarios by returning null.
Comment on lines +806 to +808
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

nit: Since this is a temp change, can you link to the internal ticket tracking this (b/XXXX)

if (account == null || !EMAIL_PATTERN.matcher(account).matches()) {
LoggingUtils.log(
LOGGER_PROVIDER,
Level.INFO,
Collections.emptyMap(),
"Regional Access Boundary lookup is skipped for this instance because it is a non-email instance.");
Comment on lines +810 to +814
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

two small nits:

  1. I know the spec says 1-time INFO log (I think this should be fine given the skipRAB flag). For Java SDK, I think that might be too noisy for an internal/ hidden impl. Is it possible to change the warning level?
  2. Can we try a bit softer wording without mentioning RAB? It may be me, but lookup skipped ... non-email instance sounds a bit like an error and not super actionable for users.

Maybe something like (ff to change this to something clearer): Unable to retrieve this instance's email and will skip the regional request routing. Proceeding with request

return null;
}
return String.format(
OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT, getAccount());
OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT, account);
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,11 @@
import com.google.api.core.InternalApi;
import com.google.auth.http.HttpTransportFactory;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.SettableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Level;
Expand Down Expand Up @@ -75,16 +75,16 @@ final class RegionalAccessBoundaryManager {
private final AtomicReference<RegionalAccessBoundary> cachedRAB = new AtomicReference<>();

/**
* refreshFuture acts as an atomic gate for request de-duplication. If a future is present, it
* indicates a background refresh is already in progress. It also provides a handle for
* observability and unit testing to track the background task's lifecycle.
* isRefreshing acts as an atomic gate for request de-duplication. If true, it indicates a
* background refresh is already in progress.
*/
private final AtomicReference<SettableFuture<RegionalAccessBoundary>> refreshFuture =
new AtomicReference<>();
private final AtomicBoolean isRefreshing = new AtomicBoolean(false);

private final AtomicReference<CooldownState> cooldownState =
new AtomicReference<>(new CooldownState(0, INITIAL_COOLDOWN_MILLIS));

private final AtomicBoolean skipRAB = new AtomicBoolean(false);

// Unbounded thread creation is discouraged in library code to avoid resource
// exhaustion. A shared, bounded executor service ensures a hard limit (5)
// on concurrent refresh tasks, while threadCount provides unique names
Expand Down Expand Up @@ -178,7 +178,7 @@ void triggerAsyncRefresh(
final HttpTransportFactory transportFactory,
final RegionalAccessBoundaryProvider provider,
final AccessToken accessToken) {
if (isCooldownActive()) {
if (skipRAB.get() || isCooldownActive()) {
return;
}

Expand All @@ -187,28 +187,28 @@ void triggerAsyncRefresh(
return;
}

SettableFuture<RegionalAccessBoundary> future = SettableFuture.create();
// Atomically check if a refresh is already running. If compareAndSet returns true,
// this thread "won the race" and is responsible for starting the background task.
// All other concurrent threads will return false and exit immediately.
if (refreshFuture.compareAndSet(null, future)) {
if (isRefreshing.compareAndSet(false, true)) {
Runnable refreshTask =
() -> {
try {
String url = provider.getRegionalAccessBoundaryUrl();
if (url == null) {
skipRAB.set(true);
return;
}
Comment on lines +198 to +201
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

qq, can you remind me again why we need to do future.set(null);?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I was using a SettableFuture as a gate to deduplicate RAB refreshes which let's us set the future to indicate to dependent threads as to the result of the refresh.

However, I realized the same can be done with AtomicBoolean which is more easier to maintain. Hence updated the logic.

RegionalAccessBoundary newRAB =
RegionalAccessBoundary.refresh(
transportFactory, url, accessToken, clock, maxRetryElapsedTimeMillis);
cachedRAB.set(newRAB);
resetCooldown();
// Complete the future so monitors (like unit tests) know we are done.
future.set(newRAB);
} catch (Exception e) {
handleRefreshFailure(e);
future.setException(e);
} finally {
// Open the gate again for future refresh requests.
refreshFuture.set(null);
isRefreshing.set(false);
}
};

Expand All @@ -224,8 +224,7 @@ void triggerAsyncRefresh(
"Could not submit background refresh task for Regional Access Boundary. "
+ "This is non-blocking and the library will attempt to refresh on the next access. Error: "
+ e.getMessage());
future.setException(e);
refreshFuture.set(null);
isRefreshing.set(false);
}
}
}
Expand Down Expand Up @@ -279,6 +278,11 @@ long getCurrentCooldownMillis() {
return cooldownState.get().durationMillis;
}

@VisibleForTesting
boolean isSkipRAB() {
return skipRAB.get();
}

private static class CooldownState {
/** The time (in milliseconds from epoch) when the current cooldown period expires. */
final long expiryTime;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1213,7 +1213,7 @@ void getProjectId_explicitSet_noMDsCall() {
assertEquals(0, transportFactory.transport.getRequestCount());
}

@org.junit.jupiter.api.Test
@Test
void refresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {

String defaultAccountEmail = "default@email.com";
Expand Down Expand Up @@ -1242,6 +1242,76 @@ void refresh_regionalAccessBoundarySuccess() throws IOException, InterruptedExce
Arrays.asList(TestUtils.REGIONAL_ACCESS_BOUNDARY_ENCODED_LOCATION));
}

@Test
void refresh_regionalAccessBoundaryNonEmail_skipsRABLookup()
throws IOException, InterruptedException {
String nonEmailAccount = "non-email-account-value";
MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
RegionalAccessBoundary regionalAccessBoundary =
new RegionalAccessBoundary(
TestUtils.REGIONAL_ACCESS_BOUNDARY_ENCODED_LOCATION,
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS,
null);
transportFactory.transport.setRegionalAccessBoundary(regionalAccessBoundary);
transportFactory.transport.setServiceAccountEmail(nonEmailAccount);

ComputeEngineCredentials credentials =
ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();

// Before any call, skipRAB flag should be false
assertFalse(credentials.regionalAccessBoundaryManager.isSkipRAB());

// First call: triggers lookup which determines non-email, returns null, and sets skipRAB to
// true
Map<String, List<String>> headers = credentials.getRequestMetadata();
assertNull(headers.get(X_ALLOWED_LOCATIONS_HEADER_KEY));

// Since the task is scheduled asynchronously on the shared executor, wait for it to complete
long deadline = System.currentTimeMillis() + 5000;
while (!credentials.regionalAccessBoundaryManager.isSkipRAB()
&& System.currentTimeMillis() < deadline) {
Thread.sleep(50);
}

// Verify skipRAB flag has been set to true
assertTrue(credentials.regionalAccessBoundaryManager.isSkipRAB());

// Verify RAB is still null
assertNull(credentials.getRegionalAccessBoundary());

// Second call: should bypass triggerAsyncRefresh completely and remain null
headers = credentials.getRequestMetadata();
assertNull(headers.get(X_ALLOWED_LOCATIONS_HEADER_KEY));
}

@Test
void getRegionalAccessBoundaryUrl_validEmail_returnsUrl() throws IOException {
MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
String defaultAccountEmail = "mail@mail.com";

transportFactory.transport.setServiceAccountEmail(defaultAccountEmail);
ComputeEngineCredentials credentials =
ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();

String expectedUrl =
String.format(
OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT,
defaultAccountEmail);
assertEquals(expectedUrl, credentials.getRegionalAccessBoundaryUrl());
}

@Test
void getRegionalAccessBoundaryUrl_invalidEmail_returnsNull() throws IOException {
MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
String defaultAccountEmail = "default"; // non-email account format

transportFactory.transport.setServiceAccountEmail(defaultAccountEmail);
ComputeEngineCredentials credentials =
ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();

assertNull(credentials.getRegionalAccessBoundaryUrl());
}

private void waitForRegionalAccessBoundary(GoogleCredentials credentials)
throws InterruptedException {
long deadline = System.currentTimeMillis() + 5000;
Expand Down
Loading