Skip to content

Commit fa9583e

Browse files
vvermanvverman
committed
Using executor; RAB refresh not gated by env var; other fixes.
1 parent 30f2c14 commit fa9583e

12 files changed

Lines changed: 116 additions & 145 deletions

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

Lines changed: 34 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -40,10 +40,11 @@
4040
import com.google.api.client.http.HttpResponse;
4141
import com.google.api.client.http.HttpUnsuccessfulResponseHandler;
4242
import com.google.api.client.json.GenericJson;
43-
import com.google.api.client.json.JsonParser;
43+
import com.google.api.client.json.JsonObjectParser;
4444
import com.google.api.client.util.Clock;
4545
import com.google.api.client.util.ExponentialBackOff;
4646
import com.google.api.client.util.Key;
47+
import com.google.api.core.InternalApi;
4748
import com.google.auth.http.HttpTransportFactory;
4849
import com.google.common.annotations.VisibleForTesting;
4950
import com.google.common.base.MoreObjects;
@@ -53,7 +54,6 @@
5354
import java.io.Serializable;
5455
import java.util.Collections;
5556
import java.util.List;
56-
import javax.annotation.Nullable;
5757

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

70-
// Note: this is for internal testing use use only.
71-
// TODO: Fix unit test mocks so this can be removed
72-
// Refer -> https://github.com/googleapis/google-auth-library-java/issues/1898
73-
static final String ENABLE_EXPERIMENT_ENV_VAR = "GOOGLE_AUTH_TRUST_BOUNDARY_ENABLE_EXPERIMENT";
70+
private static final ThreadLocal<Boolean> DISABLE_RAB_FOR_TESTS =
71+
ThreadLocal.withInitial(() -> false);
72+
73+
@VisibleForTesting
74+
static void disableForTests() {
75+
DISABLE_RAB_FOR_TESTS.set(true);
76+
}
77+
78+
@VisibleForTesting
79+
static void enableForTests() {
80+
DISABLE_RAB_FOR_TESTS.set(false);
81+
}
82+
83+
@VisibleForTesting
84+
static void resetForTests() {
85+
DISABLE_RAB_FOR_TESTS.remove();
86+
}
87+
7488
static final long TTL_MILLIS = 6 * 60 * 60 * 1000L; // 6 hours
7589
static final long REFRESH_THRESHOLD_MILLIS = 1 * 60 * 60 * 1000L; // 1 hour
7690

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

82-
private static EnvironmentProvider environmentProvider = SystemEnvironmentProvider.getInstance();
83-
8496
/**
8597
* Creates a new RegionalAccessBoundary instance.
8698
*
@@ -172,28 +184,16 @@ public String toString() {
172184
}
173185
}
174186

175-
@VisibleForTesting
176-
static void setEnvironmentProviderForTest(@Nullable EnvironmentProvider provider) {
177-
environmentProvider = provider == null ? SystemEnvironmentProvider.getInstance() : provider;
178-
}
179-
180187
/**
181-
* Checks if the regional access boundary feature is enabled. The feature is enabled if the
182-
* environment variable or system property "GOOGLE_AUTH_TRUST_BOUNDARY_ENABLE_EXPERIMENT" is set
183-
* to "true" or "1" (case-insensitive).
188+
* Checks if the regional access boundary feature is enabled.
189+
*
190+
* <p>This method is for internal use only and may be changed or removed in future releases.
184191
*
185192
* @return True if the regional access boundary feature is enabled, false otherwise.
186193
*/
194+
@InternalApi
187195
static boolean isEnabled() {
188-
String enabled = environmentProvider.getEnv(ENABLE_EXPERIMENT_ENV_VAR);
189-
if (enabled == null) {
190-
enabled = System.getProperty(ENABLE_EXPERIMENT_ENV_VAR);
191-
}
192-
if (enabled == null) {
193-
return false;
194-
}
195-
String lowercased = enabled.toLowerCase();
196-
return "true".equals(lowercased) || "1".equals(enabled);
196+
return !DISABLE_RAB_FOR_TESTS.get();
197197
}
198198

199199
/**
@@ -249,15 +249,20 @@ static RegionalAccessBoundary refresh(
249249
HttpIOExceptionHandler ioExceptionHandler = new HttpBackOffIOExceptionHandler(backoff);
250250
request.setIOExceptionHandler(ioExceptionHandler);
251251

252+
request.setParser(new JsonObjectParser(OAuth2Utils.JSON_FACTORY));
253+
252254
RegionalAccessBoundaryResponse json;
255+
HttpResponse response = null;
253256
try {
254-
HttpResponse response = request.execute();
255-
String responseString = response.parseAsString();
256-
JsonParser parser = OAuth2Utils.JSON_FACTORY.createJsonParser(responseString);
257-
json = parser.parseAndClose(RegionalAccessBoundaryResponse.class);
257+
response = request.execute();
258+
json = response.parseAs(RegionalAccessBoundaryResponse.class);
258259
} catch (IOException e) {
259260
throw new IOException(
260261
"RegionalAccessBoundary: Failure while getting regional access boundaries:", e);
262+
} finally {
263+
if (response != null) {
264+
response.disconnect();
265+
}
261266
}
262267
String encodedLocations = json.getEncodedLocations();
263268
// The encodedLocations is the value attached to the x-allowed-locations header, and

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

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,9 @@
3636
import com.google.auth.http.HttpTransportFactory;
3737
import com.google.common.annotations.VisibleForTesting;
3838
import com.google.common.util.concurrent.SettableFuture;
39+
import java.util.concurrent.ExecutorService;
40+
import java.util.concurrent.Executors;
41+
import java.util.concurrent.atomic.AtomicInteger;
3942
import java.util.concurrent.atomic.AtomicReference;
4043
import java.util.logging.Level;
4144
import javax.annotation.Nullable;
@@ -78,6 +81,20 @@ final class RegionalAccessBoundaryManager {
7881
private final AtomicReference<CooldownState> cooldownState =
7982
new AtomicReference<>(new CooldownState(0, INITIAL_COOLDOWN_MILLIS));
8083

84+
// Unbounded thread creation is discouraged in library code to avoid resource
85+
// exhaustion. A shared, bounded executor service ensures a hard limit (5)
86+
// on concurrent refresh tasks, while threadCount provides unique names
87+
// for easier debugging.
88+
private static final AtomicInteger threadCount = new AtomicInteger(0);
89+
private static final ExecutorService EXECUTOR =
90+
Executors.newFixedThreadPool(
91+
5,
92+
r -> {
93+
Thread t = new Thread(r, "RAB-refresh-" + threadCount.getAndIncrement());
94+
t.setDaemon(true);
95+
return t;
96+
});
97+
8198
private final transient Clock clock;
8299
private final int maxRetryElapsedTimeMillis;
83100

@@ -161,14 +178,7 @@ void triggerAsyncRefresh(
161178
};
162179

163180
try {
164-
// We use new Thread() here instead of
165-
// CompletableFuture.runAsync() (which uses ForkJoinPool.commonPool()).
166-
// This avoids consuming CPU resources since
167-
// The common pool has a small, fixed number of threads designed for
168-
// CPU-bound tasks.
169-
Thread refreshThread = new Thread(refreshTask, "RAB-refresh-thread");
170-
refreshThread.setDaemon(true);
171-
refreshThread.start();
181+
EXECUTOR.submit(refreshTask);
172182
} catch (Exception | Error e) {
173183
// If scheduling fails (e.g., RejectedExecutionException, OutOfMemoryError for threads),
174184
// the task's finally block will never execute. We must release the lock here.

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -65,11 +65,6 @@ class AwsCredentialsTest extends BaseSerializationTest {
6565
@org.junit.jupiter.api.BeforeEach
6666
void setUp() {}
6767

68-
@org.junit.jupiter.api.AfterEach
69-
void tearDown() {
70-
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
71-
}
72-
7368
private static final String STS_URL = "https://sts.googleapis.com/v1/token";
7469
private static final String AWS_CREDENTIALS_URL = "https://169.254.169.254";
7570
private static final String AWS_CREDENTIALS_URL_WITH_ROLE = "https://169.254.169.254/roleName";
@@ -1369,9 +1364,6 @@ public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext cont
13691364

13701365
@Test
13711366
public void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1372-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1373-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1374-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
13751367

13761368
MockExternalAccountCredentialsTransportFactory transportFactory =
13771369
new MockExternalAccountCredentialsTransportFactory();

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -80,11 +80,6 @@ class ComputeEngineCredentialsTest extends BaseSerializationTest {
8080
@org.junit.jupiter.api.BeforeEach
8181
void setUp() {}
8282

83-
@org.junit.jupiter.api.AfterEach
84-
void tearDown() {
85-
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
86-
}
87-
8883
private static final URI CALL_URI = URI.create("http://googleapis.com/testapi/v1/foo");
8984

9085
private static final String TOKEN_URL =
@@ -1188,9 +1183,6 @@ void getProjectId_explicitSet_noMDsCall() {
11881183

11891184
@org.junit.jupiter.api.Test
11901185
void refresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1191-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1192-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1193-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
11941186

11951187
String defaultAccountEmail = "default@email.com";
11961188
MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();

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

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -130,9 +130,7 @@ void setup() {
130130
}
131131

132132
@org.junit.jupiter.api.AfterEach
133-
void tearDown() {
134-
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
135-
}
133+
void tearDown() {}
136134

137135
@Test
138136
void builder_allFields() throws IOException {
@@ -1243,9 +1241,6 @@ void serialize() throws IOException, ClassNotFoundException {
12431241

12441242
@org.junit.jupiter.api.Test
12451243
void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1246-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1247-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1248-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
12491244

12501245
ExternalAccountAuthorizedUserCredentials credentials =
12511246
ExternalAccountAuthorizedUserCredentials.newBuilder()

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

Lines changed: 4 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -92,11 +92,6 @@ void setup() {
9292
transportFactory = new MockExternalAccountCredentialsTransportFactory();
9393
}
9494

95-
@org.junit.jupiter.api.AfterEach
96-
void tearDown() {
97-
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
98-
}
99-
10095
@Test
10196
void fromStream_identityPoolCredentials() throws IOException {
10297
GenericJson json = buildJsonIdentityPoolCredential();
@@ -1302,9 +1297,7 @@ public void getRegionalAccessBoundaryUrl_invalidAudience_throws() {
13021297
@Test
13031298
public void refresh_workload_regionalAccessBoundarySuccess()
13041299
throws IOException, InterruptedException {
1305-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1306-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1307-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1300+
13081301
String audience =
13091302
"//iam.googleapis.com/projects/12345/locations/global/workloadIdentityPools/my-pool/providers/my-provider";
13101303

@@ -1339,9 +1332,7 @@ public String retrieveSubjectToken() throws IOException {
13391332
@Test
13401333
public void refresh_workforce_regionalAccessBoundarySuccess()
13411334
throws IOException, InterruptedException {
1342-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1343-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1344-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1335+
13451336
String audience =
13461337
"//iam.googleapis.com/locations/global/workforcePools/my-pool/providers/my-provider";
13471338

@@ -1376,9 +1367,7 @@ public String retrieveSubjectToken() throws IOException {
13761367
@Test
13771368
public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
13781369
throws IOException, InterruptedException {
1379-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1380-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1381-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1370+
13821371
String projectNumber = "12345";
13831372
String poolId = "my-pool";
13841373
String providerId = "my-provider";
@@ -1440,9 +1429,7 @@ public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
14401429
@Test
14411430
public void refresh_impersonated_workforce_regionalAccessBoundarySuccess()
14421431
throws IOException, InterruptedException {
1443-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1444-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1445-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1432+
14461433
String poolId = "my-pool";
14471434
String providerId = "my-provider";
14481435
String audience =

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

Lines changed: 7 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,7 @@ class GoogleCredentialsTest extends BaseSerializationTest {
109109
void setUp() {}
110110

111111
@org.junit.jupiter.api.AfterEach
112-
void tearDown() {
113-
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
114-
}
112+
void tearDown() {}
115113

116114
@Test
117115
void getApplicationDefault_nullTransport_throws() {
@@ -858,9 +856,6 @@ void serialize() throws IOException, ClassNotFoundException {
858856

859857
@Test
860858
public void serialize_removesStaleRabHeaders() throws Exception {
861-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
862-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
863-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
864859

865860
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
866861
RegionalAccessBoundary rab =
@@ -1046,9 +1041,7 @@ void getCredentialInfo_impersonatedServiceAccount() throws IOException {
10461041
@Test
10471042
public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDataSuccessfully()
10481043
throws IOException, InterruptedException {
1049-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1050-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1051-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1044+
10521045
MockTokenServerTransport transport = new MockTokenServerTransport();
10531046
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
10541047
RegionalAccessBoundary regionalAccessBoundary =
@@ -1083,9 +1076,6 @@ public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDat
10831076
@Test
10841077
public void regionalAccessBoundary_shouldRetryRegionalAccessBoundaryLookupOnFailure()
10851078
throws IOException, InterruptedException {
1086-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1087-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1088-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
10891079

10901080
// This transport will be used for the regional access boundary lookup.
10911081
// We will configure it to fail on the first attempt.
@@ -1137,9 +1127,7 @@ public com.google.api.client.http.LowLevelHttpRequest buildRequest(
11371127
@Test
11381128
public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIsPassed()
11391129
throws IOException {
1140-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1141-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1142-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1130+
11431131
MockTokenServerTransport transport = new MockTokenServerTransport();
11441132
// Return an expired access token.
11451133
transport.addServiceAccount(SA_CLIENT_EMAIL, "expired-token");
@@ -1162,9 +1150,7 @@ public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIs
11621150
@Test
11631151
public void regionalAccessBoundary_cooldownDoublingAndRefresh()
11641152
throws IOException, InterruptedException {
1165-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1166-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1167-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1153+
11681154
MockTokenServerTransport transport = new MockTokenServerTransport();
11691155
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
11701156
// Always fail lookup for now.
@@ -1224,9 +1210,7 @@ public void regionalAccessBoundary_cooldownDoublingAndRefresh()
12241210

12251211
@Test
12261212
public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() throws IOException {
1227-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1228-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1229-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1213+
12301214
// Use a simple AccessToken-based credential that won't try to refresh.
12311215
GoogleCredentials credentials = GoogleCredentials.create(new AccessToken("some-token", null));
12321216

@@ -1238,9 +1222,7 @@ public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() th
12381222
@Test
12391223
public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
12401224
throws IOException, InterruptedException {
1241-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1242-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1243-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1225+
12441226
MockTokenServerTransport transport = new MockTokenServerTransport();
12451227
transport.setRegionalAccessBoundary(
12461228
new RegionalAccessBoundary("valid", Collections.singletonList("us-central1"), null));
@@ -1269,9 +1251,7 @@ public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
12691251

12701252
@Test
12711253
public void regionalAccessBoundary_shouldSkipRefreshForRegionalEndpoints() throws IOException {
1272-
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
1273-
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1274-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
1254+
12751255
MockTokenServerTransport transport = new MockTokenServerTransport();
12761256
GoogleCredentials credentials = createTestCredentials(transport);
12771257

0 commit comments

Comments
 (0)