Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,11 @@
import com.google.api.client.http.HttpResponse;
import com.google.api.client.http.HttpUnsuccessfulResponseHandler;
import com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonParser;
import com.google.api.client.json.JsonObjectParser;
import com.google.api.client.util.Clock;
import com.google.api.client.util.ExponentialBackOff;
import com.google.api.client.util.Key;
import com.google.api.core.InternalApi;
import com.google.auth.http.HttpTransportFactory;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
Expand All @@ -53,7 +54,6 @@
import java.io.Serializable;
import java.util.Collections;
import java.util.List;
import javax.annotation.Nullable;

/**
* Represents the regional access boundary configuration for a credential. This class holds the
Expand All @@ -67,10 +67,24 @@ final class RegionalAccessBoundary implements Serializable {
static final String X_ALLOWED_LOCATIONS_HEADER_KEY = "x-allowed-locations";
private static final long serialVersionUID = -2428522338274020302L;

// Note: this is for internal testing use use only.
// TODO: Fix unit test mocks so this can be removed
// Refer -> https://github.com/googleapis/google-auth-library-java/issues/1898
Comment thread
vverman marked this conversation as resolved.
static final String ENABLE_EXPERIMENT_ENV_VAR = "GOOGLE_AUTH_TRUST_BOUNDARY_ENABLE_EXPERIMENT";
private static final ThreadLocal<Boolean> DISABLE_RAB_FOR_TESTS =
ThreadLocal.withInitial(() -> false);

@VisibleForTesting
static void disableForTests() {
DISABLE_RAB_FOR_TESTS.set(true);
}

@VisibleForTesting
static void enableForTests() {
DISABLE_RAB_FOR_TESTS.set(false);
}

@VisibleForTesting
static void resetForTests() {
DISABLE_RAB_FOR_TESTS.remove();
}
Comment thread
vverman marked this conversation as resolved.
Outdated

static final long TTL_MILLIS = 6 * 60 * 60 * 1000L; // 6 hours
static final long REFRESH_THRESHOLD_MILLIS = 1 * 60 * 60 * 1000L; // 1 hour

Expand All @@ -79,8 +93,6 @@ final class RegionalAccessBoundary implements Serializable {
private final long refreshTime;
private transient Clock clock;

private static EnvironmentProvider environmentProvider = SystemEnvironmentProvider.getInstance();

/**
* Creates a new RegionalAccessBoundary instance.
*
Expand Down Expand Up @@ -172,28 +184,16 @@ public String toString() {
}
}

@VisibleForTesting
static void setEnvironmentProviderForTest(@Nullable EnvironmentProvider provider) {
Comment thread
vverman marked this conversation as resolved.
environmentProvider = provider == null ? SystemEnvironmentProvider.getInstance() : provider;
}

/**
* Checks if the regional access boundary feature is enabled. The feature is enabled if the
* environment variable or system property "GOOGLE_AUTH_TRUST_BOUNDARY_ENABLE_EXPERIMENT" is set
* to "true" or "1" (case-insensitive).
* Checks if the regional access boundary feature is enabled.
*
* <p>This method is for internal use only and may be changed or removed in future releases.
*
* @return True if the regional access boundary feature is enabled, false otherwise.
*/
Comment thread
vverman marked this conversation as resolved.
Outdated
@InternalApi
static boolean isEnabled() {
String enabled = environmentProvider.getEnv(ENABLE_EXPERIMENT_ENV_VAR);
if (enabled == null) {
enabled = System.getProperty(ENABLE_EXPERIMENT_ENV_VAR);
}
if (enabled == null) {
return false;
}
String lowercased = enabled.toLowerCase();
return "true".equals(lowercased) || "1".equals(enabled);
return !DISABLE_RAB_FOR_TESTS.get();
Comment thread
lsirac marked this conversation as resolved.
Outdated
}

/**
Expand Down Expand Up @@ -249,15 +249,20 @@ static RegionalAccessBoundary refresh(
HttpIOExceptionHandler ioExceptionHandler = new HttpBackOffIOExceptionHandler(backoff);
request.setIOExceptionHandler(ioExceptionHandler);

request.setParser(new JsonObjectParser(OAuth2Utils.JSON_FACTORY));

RegionalAccessBoundaryResponse json;
HttpResponse response = null;
try {
HttpResponse response = request.execute();
String responseString = response.parseAsString();
JsonParser parser = OAuth2Utils.JSON_FACTORY.createJsonParser(responseString);
json = parser.parseAndClose(RegionalAccessBoundaryResponse.class);
response = request.execute();
json = response.parseAs(RegionalAccessBoundaryResponse.class);
} catch (IOException e) {
throw new IOException(
"RegionalAccessBoundary: Failure while getting regional access boundaries:", e);
} finally {
if (response != null) {
response.disconnect();
}
}
String encodedLocations = json.getEncodedLocations();
// The encodedLocations is the value attached to the x-allowed-locations header, and
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,9 @@
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.Executors;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Level;
import javax.annotation.Nullable;
Expand Down Expand Up @@ -78,6 +81,20 @@ final class RegionalAccessBoundaryManager {
private final AtomicReference<CooldownState> cooldownState =
new AtomicReference<>(new CooldownState(0, INITIAL_COOLDOWN_MILLIS));

// 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
// for easier debugging.
private static final AtomicInteger threadCount = new AtomicInteger(0);
private static final ExecutorService EXECUTOR =
Executors.newFixedThreadPool(
5,
r -> {
Thread t = new Thread(r, "RAB-refresh-" + threadCount.getAndIncrement());
t.setDaemon(true);
return t;
});
Comment thread
vverman marked this conversation as resolved.
Outdated

private final transient Clock clock;
private final int maxRetryElapsedTimeMillis;

Expand Down Expand Up @@ -161,14 +178,7 @@ void triggerAsyncRefresh(
};

try {
// We use new Thread() here instead of
// CompletableFuture.runAsync() (which uses ForkJoinPool.commonPool()).
// This avoids consuming CPU resources since
// The common pool has a small, fixed number of threads designed for
// CPU-bound tasks.
Thread refreshThread = new Thread(refreshTask, "RAB-refresh-thread");
refreshThread.setDaemon(true);
refreshThread.start();
EXECUTOR.submit(refreshTask);
Comment thread
vverman marked this conversation as resolved.
Outdated
} catch (Exception | Error e) {
// If scheduling fails (e.g., RejectedExecutionException, OutOfMemoryError for threads),
// the task's finally block will never execute. We must release the lock here.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -65,11 +65,6 @@ class AwsCredentialsTest extends BaseSerializationTest {
@org.junit.jupiter.api.BeforeEach
void setUp() {}

@org.junit.jupiter.api.AfterEach
void tearDown() {
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
}

private static final String STS_URL = "https://sts.googleapis.com/v1/token";
private static final String AWS_CREDENTIALS_URL = "https://169.254.169.254";
private static final String AWS_CREDENTIALS_URL_WITH_ROLE = "https://169.254.169.254/roleName";
Expand Down Expand Up @@ -1369,9 +1364,6 @@ public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext cont

@Test
public void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

MockExternalAccountCredentialsTransportFactory transportFactory =
new MockExternalAccountCredentialsTransportFactory();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -80,11 +80,6 @@ class ComputeEngineCredentialsTest extends BaseSerializationTest {
@org.junit.jupiter.api.BeforeEach
void setUp() {}

@org.junit.jupiter.api.AfterEach
void tearDown() {
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
}

private static final URI CALL_URI = URI.create("http://googleapis.com/testapi/v1/foo");

private static final String TOKEN_URL =
Expand Down Expand Up @@ -1188,9 +1183,6 @@ void getProjectId_explicitSet_noMDsCall() {

@org.junit.jupiter.api.Test
void refresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

String defaultAccountEmail = "default@email.com";
MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -130,9 +130,7 @@ void setup() {
}

@org.junit.jupiter.api.AfterEach
void tearDown() {
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
}
void tearDown() {}

@Test
void builder_allFields() throws IOException {
Expand Down Expand Up @@ -1243,9 +1241,6 @@ void serialize() throws IOException, ClassNotFoundException {

@org.junit.jupiter.api.Test
void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

ExternalAccountAuthorizedUserCredentials credentials =
ExternalAccountAuthorizedUserCredentials.newBuilder()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -92,11 +92,6 @@ void setup() {
transportFactory = new MockExternalAccountCredentialsTransportFactory();
}

@org.junit.jupiter.api.AfterEach
void tearDown() {
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
}

@Test
void fromStream_identityPoolCredentials() throws IOException {
GenericJson json = buildJsonIdentityPoolCredential();
Expand Down Expand Up @@ -1302,9 +1297,7 @@ public void getRegionalAccessBoundaryUrl_invalidAudience_throws() {
@Test
public void refresh_workload_regionalAccessBoundarySuccess()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

String audience =
"//iam.googleapis.com/projects/12345/locations/global/workloadIdentityPools/my-pool/providers/my-provider";

Expand Down Expand Up @@ -1339,9 +1332,7 @@ public String retrieveSubjectToken() throws IOException {
@Test
public void refresh_workforce_regionalAccessBoundarySuccess()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

String audience =
"//iam.googleapis.com/locations/global/workforcePools/my-pool/providers/my-provider";

Expand Down Expand Up @@ -1376,9 +1367,7 @@ public String retrieveSubjectToken() throws IOException {
@Test
public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

String projectNumber = "12345";
String poolId = "my-pool";
String providerId = "my-provider";
Expand Down Expand Up @@ -1440,9 +1429,7 @@ public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
@Test
public void refresh_impersonated_workforce_regionalAccessBoundarySuccess()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

String poolId = "my-pool";
String providerId = "my-provider";
String audience =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -109,9 +109,7 @@ class GoogleCredentialsTest extends BaseSerializationTest {
void setUp() {}

@org.junit.jupiter.api.AfterEach
void tearDown() {
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
}
void tearDown() {}

@Test
void getApplicationDefault_nullTransport_throws() {
Expand Down Expand Up @@ -858,9 +856,6 @@ void serialize() throws IOException, ClassNotFoundException {

@Test
public void serialize_removesStaleRabHeaders() throws Exception {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
RegionalAccessBoundary rab =
Expand Down Expand Up @@ -1046,9 +1041,7 @@ void getCredentialInfo_impersonatedServiceAccount() throws IOException {
@Test
public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDataSuccessfully()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

MockTokenServerTransport transport = new MockTokenServerTransport();
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
RegionalAccessBoundary regionalAccessBoundary =
Expand Down Expand Up @@ -1083,9 +1076,6 @@ public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDat
@Test
public void regionalAccessBoundary_shouldRetryRegionalAccessBoundaryLookupOnFailure()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

// This transport will be used for the regional access boundary lookup.
// We will configure it to fail on the first attempt.
Expand Down Expand Up @@ -1137,9 +1127,7 @@ public com.google.api.client.http.LowLevelHttpRequest buildRequest(
@Test
public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIsPassed()
throws IOException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

MockTokenServerTransport transport = new MockTokenServerTransport();
// Return an expired access token.
transport.addServiceAccount(SA_CLIENT_EMAIL, "expired-token");
Expand All @@ -1162,9 +1150,7 @@ public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIs
@Test
public void regionalAccessBoundary_cooldownDoublingAndRefresh()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

MockTokenServerTransport transport = new MockTokenServerTransport();
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
// Always fail lookup for now.
Expand Down Expand Up @@ -1224,9 +1210,7 @@ public void regionalAccessBoundary_cooldownDoublingAndRefresh()

@Test
public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() throws IOException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

// Use a simple AccessToken-based credential that won't try to refresh.
GoogleCredentials credentials = GoogleCredentials.create(new AccessToken("some-token", null));

Expand All @@ -1238,9 +1222,7 @@ public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() th
@Test
public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
throws IOException, InterruptedException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

MockTokenServerTransport transport = new MockTokenServerTransport();
transport.setRegionalAccessBoundary(
new RegionalAccessBoundary("valid", Collections.singletonList("us-central1"), null));
Expand Down Expand Up @@ -1269,9 +1251,7 @@ public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()

@Test
public void regionalAccessBoundary_shouldSkipRefreshForRegionalEndpoints() throws IOException {
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");

MockTokenServerTransport transport = new MockTokenServerTransport();
GoogleCredentials credentials = createTestCredentials(transport);

Expand Down
Loading
Loading