Skip to content

Commit 93bb045

Browse files
committed
Switched to executor and removed env variable
1 parent 0044eca commit 93bb045

13 files changed

+117
-50
lines changed

google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/JwtClaims.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
import com.google.common.collect.ImmutableMap;
3636
import java.io.Serializable;
3737
import java.util.Map;
38-
import javax.annotation.Nullable;
38+
import org.jspecify.annotations.Nullable;
3939

4040
/**
4141
* Value class representing the set of fields used as the payload of a JWT token.

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

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,6 @@ 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";
7470
static final long TTL_MILLIS = 6 * 60 * 60 * 1000L; // 6 hours
7571
static final long REFRESH_THRESHOLD_MILLIS = 1 * 60 * 60 * 1000L; // 1 hour
7672

@@ -81,6 +77,25 @@ final class RegionalAccessBoundary implements Serializable {
8177

8278
private static EnvironmentProvider environmentProvider = SystemEnvironmentProvider.getInstance();
8379

80+
// Static thread-isolated flag for granular testing setups
81+
private static final ThreadLocal<Boolean> DISABLE_RAB_FOR_TESTS =
82+
ThreadLocal.withInitial(() -> false);
83+
84+
@VisibleForTesting
85+
static void disableForTests() {
86+
DISABLE_RAB_FOR_TESTS.set(true);
87+
}
88+
89+
@VisibleForTesting
90+
static void enableForTests() {
91+
DISABLE_RAB_FOR_TESTS.set(false);
92+
}
93+
94+
@VisibleForTesting
95+
static void resetForTests() {
96+
DISABLE_RAB_FOR_TESTS.remove();
97+
}
98+
8499
/**
85100
* Creates a new RegionalAccessBoundary instance.
86101
*
@@ -178,22 +193,17 @@ static void setEnvironmentProviderForTest(@Nullable EnvironmentProvider provider
178193
}
179194

180195
/**
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).
196+
* Checks if the regional access boundary feature is enabled.
184197
*
185198
* @return True if the regional access boundary feature is enabled, false otherwise.
186199
*/
187200
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) {
201+
// 1. Check if granular opt-out flag is active for THIS thread
202+
if (DISABLE_RAB_FOR_TESTS.get()) {
193203
return false;
194204
}
195-
String lowercased = enabled.toLowerCase();
196-
return "true".equals(lowercased) || "1".equals(enabled);
205+
// 2. Fallback to standard GA behavior (enabled by default)
206+
return true;
197207
}
198208

199209
/**

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

Lines changed: 15 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.ThreadFactory;
3942
import java.util.concurrent.atomic.AtomicReference;
4043
import java.util.logging.Level;
4144
import javax.annotation.Nullable;
@@ -52,6 +55,17 @@ final class RegionalAccessBoundaryManager {
5255
private static final LoggerProvider LOGGER_PROVIDER =
5356
LoggerProvider.forClazz(RegionalAccessBoundaryManager.class);
5457

58+
private static final ExecutorService REFRESH_EXECUTOR =
59+
Executors.newCachedThreadPool(
60+
new ThreadFactory() {
61+
@Override
62+
public Thread newThread(Runnable r) {
63+
Thread t = new Thread(r, "RAB-refresh-thread");
64+
t.setDaemon(true);
65+
return t;
66+
}
67+
});
68+
5569
static final long INITIAL_COOLDOWN_MILLIS = 15 * 60 * 1000L; // 15 minutes
5670
static final long MAX_COOLDOWN_MILLIS = 6 * 60 * 60 * 1000L; // 6 hours
5771

@@ -161,14 +175,7 @@ void triggerAsyncRefresh(
161175
};
162176

163177
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();
178+
REFRESH_EXECUTOR.execute(refreshTask);
172179
} catch (Exception | Error e) {
173180
// If scheduling fails (e.g., RejectedExecutionException, OutOfMemoryError for threads),
174181
// 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: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,10 +62,13 @@
6262
class AwsCredentialsTest extends BaseSerializationTest {
6363

6464
@org.junit.jupiter.api.BeforeEach
65-
void setUp() {}
65+
void setUp() {
66+
RegionalAccessBoundary.disableForTests();
67+
}
6668

6769
@org.junit.jupiter.api.AfterEach
6870
void tearDown() {
71+
RegionalAccessBoundary.resetForTests();
6972
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
7073
}
7174

@@ -1404,9 +1407,9 @@ public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext cont
14041407

14051408
@Test
14061409
public void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1410+
RegionalAccessBoundary.enableForTests();
14071411
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
14081412
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1409-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
14101413

14111414
MockExternalAccountCredentialsTransportFactory transportFactory =
14121415
new MockExternalAccountCredentialsTransportFactory();

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,10 +76,13 @@
7676
class ComputeEngineCredentialsTest extends BaseSerializationTest {
7777

7878
@org.junit.jupiter.api.BeforeEach
79-
void setUp() {}
79+
void setUp() {
80+
RegionalAccessBoundary.disableForTests();
81+
}
8082

8183
@org.junit.jupiter.api.AfterEach
8284
void tearDown() {
85+
RegionalAccessBoundary.resetForTests();
8386
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
8487
}
8588

@@ -1205,9 +1208,9 @@ void getProjectId_explicitSet_noMDsCall() {
12051208

12061209
@org.junit.jupiter.api.Test
12071210
void refresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1211+
RegionalAccessBoundary.enableForTests();
12081212
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
12091213
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1210-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
12111214

12121215
String defaultAccountEmail = "default@email.com";
12131216
MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import com.google.api.client.http.HttpTransport;
4545
import com.google.api.client.json.GenericJson;
4646
import com.google.api.client.testing.http.MockLowLevelHttpRequest;
47+
import com.google.api.client.util.Clock;
4748
import com.google.auth.TestUtils;
4849
import com.google.auth.http.AuthHttpConstants;
4950
import com.google.auth.http.HttpTransportFactory;
@@ -126,11 +127,13 @@ public HttpTransport create() {
126127

127128
@BeforeEach
128129
void setup() {
130+
RegionalAccessBoundary.disableForTests();
129131
transportFactory = new MockExternalAccountAuthorizedUserCredentialsTransportFactory();
130132
}
131133

132134
@org.junit.jupiter.api.AfterEach
133135
void tearDown() {
136+
RegionalAccessBoundary.resetForTests();
134137
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
135138
}
136139

@@ -1247,9 +1250,9 @@ void serialize() throws IOException, ClassNotFoundException {
12471250

12481251
@org.junit.jupiter.api.Test
12491252
void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1253+
RegionalAccessBoundary.enableForTests();
12501254
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
12511255
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1252-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
12531256

12541257
ExternalAccountAuthorizedUserCredentials credentials =
12551258
ExternalAccountAuthorizedUserCredentials.newBuilder()

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

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
import com.google.api.client.json.GenericJson;
4848
import com.google.api.client.util.Clock;
4949
import com.google.auth.TestUtils;
50+
import java.util.Collections;
5051
import com.google.auth.http.HttpTransportFactory;
5152
import com.google.auth.oauth2.ExternalAccountCredentials.SubjectTokenTypes;
5253
import com.google.auth.oauth2.ExternalAccountCredentialsTest.TestExternalAccountCredentials.TestCredentialSource;
@@ -91,11 +92,13 @@ public HttpTransport create() {
9192

9293
@BeforeEach
9394
void setup() {
95+
RegionalAccessBoundary.disableForTests();
9496
transportFactory = new MockExternalAccountCredentialsTransportFactory();
9597
}
9698

9799
@org.junit.jupiter.api.AfterEach
98100
void tearDown() {
101+
RegionalAccessBoundary.resetForTests();
99102
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
100103
}
101104

@@ -1311,9 +1314,9 @@ public void getRegionalAccessBoundaryUrl_invalidAudience_throws() {
13111314
@Test
13121315
public void refresh_workload_regionalAccessBoundarySuccess()
13131316
throws IOException, InterruptedException {
1317+
RegionalAccessBoundary.enableForTests();
13141318
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
13151319
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1316-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
13171320
String audience =
13181321
"//iam.googleapis.com/projects/12345/locations/global/workloadIdentityPools/my-pool/providers/my-provider";
13191322

@@ -1348,9 +1351,9 @@ public String retrieveSubjectToken() throws IOException {
13481351
@Test
13491352
public void refresh_workforce_regionalAccessBoundarySuccess()
13501353
throws IOException, InterruptedException {
1354+
RegionalAccessBoundary.enableForTests();
13511355
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
13521356
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1353-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
13541357
String audience =
13551358
"//iam.googleapis.com/locations/global/workforcePools/my-pool/providers/my-provider";
13561359

@@ -1385,9 +1388,9 @@ public String retrieveSubjectToken() throws IOException {
13851388
@Test
13861389
public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
13871390
throws IOException, InterruptedException {
1391+
RegionalAccessBoundary.enableForTests();
13881392
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
13891393
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1390-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
13911394
String projectNumber = "12345";
13921395
String poolId = "my-pool";
13931396
String providerId = "my-provider";
@@ -1449,9 +1452,9 @@ public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
14491452
@Test
14501453
public void refresh_impersonated_workforce_regionalAccessBoundarySuccess()
14511454
throws IOException, InterruptedException {
1455+
RegionalAccessBoundary.enableForTests();
14521456
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
14531457
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1454-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
14551458
String poolId = "my-pool";
14561459
String providerId = "my-provider";
14571460
String audience =

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

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@
4545
import com.google.api.client.json.GenericJson;
4646
import com.google.api.client.util.Clock;
4747
import com.google.auth.Credentials;
48+
import com.google.auth.RequestMetadataCallback;
49+
import javax.annotation.Nullable;
4850
import com.google.auth.TestUtils;
4951
import com.google.auth.http.HttpTransportFactory;
5052
import com.google.auth.oauth2.ExternalAccountAuthorizedUserCredentialsTest.MockExternalAccountAuthorizedUserCredentialsTransportFactory;
@@ -59,6 +61,7 @@
5961
import java.util.Collections;
6062
import java.util.List;
6163
import java.util.Map;
64+
import java.util.concurrent.atomic.AtomicLong;
6265
import java.util.concurrent.atomic.AtomicReference;
6366
import org.junit.jupiter.api.Test;
6467

@@ -101,10 +104,13 @@ class GoogleCredentialsTest extends BaseSerializationTest {
101104
private static final String TPC_UNIVERSE = "foo.bar";
102105

103106
@org.junit.jupiter.api.BeforeEach
104-
void setUp() {}
107+
void setUp() {
108+
RegionalAccessBoundary.disableForTests();
109+
}
105110

106111
@org.junit.jupiter.api.AfterEach
107112
void tearDown() {
113+
RegionalAccessBoundary.resetForTests();
108114
RegionalAccessBoundary.setEnvironmentProviderForTest(null);
109115
}
110116

@@ -793,9 +799,9 @@ void serialize() throws IOException, ClassNotFoundException {
793799

794800
@Test
795801
public void serialize_removesStaleRabHeaders() throws Exception {
802+
RegionalAccessBoundary.enableForTests();
796803
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
797804
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
798-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
799805

800806
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
801807
RegionalAccessBoundary rab =
@@ -992,9 +998,9 @@ void getCredentialInfo_impersonatedServiceAccount() throws IOException {
992998
@Test
993999
public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDataSuccessfully()
9941000
throws IOException, InterruptedException {
1001+
RegionalAccessBoundary.enableForTests();
9951002
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
9961003
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
997-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
9981004
MockTokenServerTransport transport = new MockTokenServerTransport();
9991005
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
10001006
RegionalAccessBoundary regionalAccessBoundary =
@@ -1029,9 +1035,9 @@ public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDat
10291035
@Test
10301036
public void regionalAccessBoundary_shouldRetryRegionalAccessBoundaryLookupOnFailure()
10311037
throws IOException, InterruptedException {
1038+
RegionalAccessBoundary.enableForTests();
10321039
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
10331040
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1034-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
10351041

10361042
// This transport will be used for the regional access boundary lookup.
10371043
// We will configure it to fail on the first attempt.
@@ -1083,9 +1089,9 @@ public com.google.api.client.http.LowLevelHttpRequest buildRequest(
10831089
@Test
10841090
public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIsPassed()
10851091
throws IOException {
1092+
RegionalAccessBoundary.enableForTests();
10861093
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
10871094
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1088-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
10891095
MockTokenServerTransport transport = new MockTokenServerTransport();
10901096
// Return an expired access token.
10911097
transport.addServiceAccount(SA_CLIENT_EMAIL, "expired-token");
@@ -1108,9 +1114,9 @@ public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIs
11081114
@Test
11091115
public void regionalAccessBoundary_cooldownDoublingAndRefresh()
11101116
throws IOException, InterruptedException {
1117+
RegionalAccessBoundary.enableForTests();
11111118
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
11121119
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1113-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
11141120
MockTokenServerTransport transport = new MockTokenServerTransport();
11151121
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
11161122
// Always fail lookup for now.
@@ -1170,9 +1176,9 @@ public void regionalAccessBoundary_cooldownDoublingAndRefresh()
11701176

11711177
@Test
11721178
public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() throws IOException {
1179+
RegionalAccessBoundary.enableForTests();
11731180
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
11741181
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1175-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
11761182
// Use a simple AccessToken-based credential that won't try to refresh.
11771183
GoogleCredentials credentials = GoogleCredentials.create(new AccessToken("some-token", null));
11781184

@@ -1184,9 +1190,9 @@ public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() th
11841190
@Test
11851191
public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
11861192
throws IOException, InterruptedException {
1193+
RegionalAccessBoundary.enableForTests();
11871194
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
11881195
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1189-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
11901196
MockTokenServerTransport transport = new MockTokenServerTransport();
11911197
transport.setRegionalAccessBoundary(
11921198
new RegionalAccessBoundary("valid", Collections.singletonList("us-central1"), null));
@@ -1215,9 +1221,9 @@ public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
12151221

12161222
@Test
12171223
public void regionalAccessBoundary_shouldSkipRefreshForRegionalEndpoints() throws IOException {
1224+
RegionalAccessBoundary.enableForTests();
12181225
TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
12191226
RegionalAccessBoundary.setEnvironmentProviderForTest(environmentProvider);
1220-
environmentProvider.setEnv(RegionalAccessBoundary.ENABLE_EXPERIMENT_ENV_VAR, "1");
12211227
MockTokenServerTransport transport = new MockTokenServerTransport();
12221228
GoogleCredentials credentials = createTestCredentials(transport);
12231229

0 commit comments

Comments
 (0)