Skip to content

Commit a6f6b90

Browse files
vvermanvverman
committed
Using per-instance clocks.
1 parent b4dd7f6 commit a6f6b90

11 files changed

Lines changed: 85 additions & 57 deletions

‎oauth2_http/java/com/google/auth/oauth2/GoogleCredentials.java‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ String getFileType() {
117117
@VisibleForTesting static boolean disableRabRefreshForTest = false;
118118

119119
transient RegionalAccessBoundaryManager regionalAccessBoundaryManager =
120-
new RegionalAccessBoundaryManager();
120+
new RegionalAccessBoundaryManager(clock);
121121

122122
protected final String quotaProjectId;
123123

@@ -439,7 +439,7 @@ void refreshRegionalAccessBoundaryWithSelfSignedJwtIfExpired(
439439
@Override
440440
public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
441441
Map<String, List<String>> metadata = super.getRequestMetadata(uri);
442-
metadata = addRegionalAccessBoundaryToRequestMetadata(metadata);
442+
metadata = addRegionalAccessBoundaryToRequestMetadata(uri, metadata);
443443
try {
444444
// Sets off an async refresh for request-metadata.
445445
refreshRegionalAccessBoundaryIfExpired(uri, getAccessToken(), null);
@@ -470,7 +470,7 @@ public void getRequestMetadata(
470470
new RequestMetadataCallback() {
471471
@Override
472472
public void onSuccess(Map<String, List<String>> metadata) {
473-
metadata = addRegionalAccessBoundaryToRequestMetadata(metadata);
473+
metadata = addRegionalAccessBoundaryToRequestMetadata(uri, metadata);
474474
try {
475475
refreshRegionalAccessBoundaryIfExpired(uri, getAccessToken(), executor);
476476
} catch (IOException e) {
@@ -542,11 +542,21 @@ static Map<String, List<String>> addQuotaProjectIdToRequestMetadata(
542542
* Adds Regional Access Boundary header to requestMetadata if available. Overwrites if present.
543543
* If the current RAB is null, it removes any stale header that might have survived serialization.
544544
*
545+
* @param uri The URI of the request.
546+
* @param requestMetadata The request metadata.
545547
* @return a new map with Regional Access Boundary header added, updated, or removed
546548
*/
547549
Map<String, List<String>> addRegionalAccessBoundaryToRequestMetadata(
548-
Map<String, List<String>> requestMetadata) {
550+
URI uri, Map<String, List<String>> requestMetadata) {
549551
Preconditions.checkNotNull(requestMetadata);
552+
553+
if (uri != null && uri.getHost() != null) {
554+
String host = uri.getHost();
555+
if (host.endsWith(".rep.googleapis.com") || host.endsWith(".rep.sandbox.googleapis.com")) {
556+
return requestMetadata;
557+
}
558+
}
559+
550560
RegionalAccessBoundary rab = getRegionalAccessBoundary();
551561
if (rab != null) {
552562
// Overwrite the header to ensure the most recent async update is used,
@@ -684,7 +694,7 @@ public int hashCode() {
684694

685695
private void readObject(ObjectInputStream input) throws IOException, ClassNotFoundException {
686696
input.defaultReadObject();
687-
regionalAccessBoundaryManager = new RegionalAccessBoundaryManager();
697+
regionalAccessBoundaryManager = new RegionalAccessBoundaryManager(clock);
688698
}
689699

690700
public static Builder newBuilder() {

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

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -67,21 +67,25 @@ public final class RegionalAccessBoundary implements Serializable {
6767

6868
static final long TTL_MILLIS = 6 * 60 * 60 * 1000L; // 6 hours
6969
static final long REFRESH_THRESHOLD_MILLIS = 1 * 60 * 60 * 1000L; // 1 hour
70-
private static int maxRetryElapsedTimeMillis = 60000; // 1 minute
7170

7271
private final String encodedLocations;
7372
private final List<String> locations;
7473
private final long refreshTime;
75-
private static Clock clock = Clock.SYSTEM;
74+
private final transient Clock clock;
7675

7776
/**
7877
* Creates a new RegionalAccessBoundary instance.
7978
*
8079
* @param encodedLocations The encoded string representation of the allowed locations.
8180
* @param locations A list of human-readable location strings.
81+
* @param clock The clock used to set the creation time.
8282
*/
83-
RegionalAccessBoundary(String encodedLocations, List<String> locations) {
84-
this(encodedLocations, locations, clock.currentTimeMillis());
83+
RegionalAccessBoundary(String encodedLocations, List<String> locations, Clock clock) {
84+
this(
85+
encodedLocations,
86+
locations,
87+
clock != null ? clock.currentTimeMillis() : Clock.SYSTEM.currentTimeMillis(),
88+
clock);
8589
}
8690

8791
/**
@@ -90,14 +94,17 @@ public final class RegionalAccessBoundary implements Serializable {
9094
* @param encodedLocations The encoded string representation of the allowed locations.
9195
* @param locations A list of human-readable location strings.
9296
* @param refreshTime The time at which the information was last refreshed.
97+
* @param clock The clock to use for expiration checks.
9398
*/
94-
RegionalAccessBoundary(String encodedLocations, List<String> locations, long refreshTime) {
99+
RegionalAccessBoundary(
100+
String encodedLocations, List<String> locations, long refreshTime, Clock clock) {
95101
this.encodedLocations = encodedLocations;
96102
this.locations =
97103
locations == null
98104
? Collections.<String>emptyList()
99105
: Collections.unmodifiableList(locations);
100106
this.refreshTime = refreshTime;
107+
this.clock = clock != null ? clock : Clock.SYSTEM;
101108
}
102109

103110
/** Returns the encoded string representation of the allowed locations. */
@@ -157,28 +164,20 @@ public String toString() {
157164
}
158165
}
159166

160-
@VisibleForTesting
161-
static void setClockForTest(Clock testClock) {
162-
clock = testClock;
163-
}
164-
165-
@VisibleForTesting
166-
static void setMaxRetryElapsedTimeMillisForTest(int millis) {
167-
maxRetryElapsedTimeMillis = millis;
168-
}
169-
170167
/**
171168
* Refreshes the regional access boundary by making a network call to the lookup endpoint.
172169
*
173170
* @param transportFactory The HTTP transport factory to use for the network request.
174171
* @param url The URL of the regional access boundary endpoint.
175172
* @param accessToken The access token to authenticate the request.
173+
* @param clock The clock to use for expiration checks.
174+
* @param maxRetryElapsedTimeMillis The max duration to wait for retries.
176175
* @return A new RegionalAccessBoundary object containing the refreshed information.
177176
* @throws IllegalArgumentException If the provided access token is null or expired.
178177
* @throws IOException If a network error occurs or the response is malformed.
179178
*/
180179
static RegionalAccessBoundary refresh(
181-
HttpTransportFactory transportFactory, String url, AccessToken accessToken)
180+
HttpTransportFactory transportFactory, String url, AccessToken accessToken, Clock clock, int maxRetryElapsedTimeMillis)
182181
throws IOException {
183182
Preconditions.checkNotNull(accessToken, "The provided access token is null.");
184183
if (accessToken.getExpirationTimeMillis() != null
@@ -231,6 +230,6 @@ static RegionalAccessBoundary refresh(
231230
throw new IOException(
232231
"RegionalAccessBoundary: Malformed response from lookup endpoint - `encodedLocations` was null.");
233232
}
234-
return new RegionalAccessBoundary(encodedLocations, json.getLocations());
233+
return new RegionalAccessBoundary(encodedLocations, json.getLocations(), clock);
235234
}
236235
}

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

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,12 @@ final class RegionalAccessBoundaryManager {
5555
static final long INITIAL_COOLDOWN_MILLIS = 15 * 60 * 1000L; // 15 minutes
5656
static final long MAX_COOLDOWN_MILLIS = 6 * 60 * 60 * 1000L; // 6 hours
5757

58+
/**
59+
* The default maximum elapsed time in milliseconds for retrying Regional Access Boundary lookup
60+
* requests.
61+
*/
62+
private static final int DEFAULT_MAX_RETRY_ELAPSED_TIME_MILLIS = 60000;
63+
5864
/**
5965
* cachedRAB uses AtomicReference to provide thread-safe, lock-free access to the cached data for
6066
* high-concurrency request threads.
@@ -72,7 +78,23 @@ final class RegionalAccessBoundaryManager {
7278
private final AtomicReference<CooldownState> cooldownState =
7379
new AtomicReference<>(new CooldownState(0, INITIAL_COOLDOWN_MILLIS));
7480

75-
private static Clock clock = Clock.SYSTEM;
81+
private final transient Clock clock;
82+
private final int maxRetryElapsedTimeMillis;
83+
84+
/**
85+
* Creates a new RegionalAccessBoundaryManager with the default retry timeout of 60 seconds.
86+
*
87+
* @param clock The clock to use for cooldown and expiration checks.
88+
*/
89+
RegionalAccessBoundaryManager(Clock clock) {
90+
this(clock, DEFAULT_MAX_RETRY_ELAPSED_TIME_MILLIS);
91+
}
92+
93+
@VisibleForTesting
94+
RegionalAccessBoundaryManager(Clock clock, int maxRetryElapsedTimeMillis) {
95+
this.clock = clock != null ? clock : Clock.SYSTEM;
96+
this.maxRetryElapsedTimeMillis = maxRetryElapsedTimeMillis;
97+
}
7698

7799
/**
78100
* Returns the currently cached RegionalAccessBoundary, or null if none is available or if it has
@@ -125,7 +147,8 @@ void triggerAsyncRefresh(
125147
try {
126148
String url = provider.getRegionalAccessBoundaryUrl();
127149
RegionalAccessBoundary newRAB =
128-
RegionalAccessBoundary.refresh(transportFactory, url, accessToken);
150+
RegionalAccessBoundary.refresh(
151+
transportFactory, url, accessToken, clock, maxRetryElapsedTimeMillis);
129152
cachedRAB.set(newRAB);
130153
resetCooldown();
131154
// Complete the future so monitors (like unit tests) know we are done.
@@ -212,11 +235,6 @@ long getCurrentCooldownMillis() {
212235
return cooldownState.get().durationMillis;
213236
}
214237

215-
@VisibleForTesting
216-
static void setClockForTest(Clock testClock) {
217-
clock = testClock;
218-
}
219-
220238
private static class CooldownState {
221239
/** The time (in milliseconds from epoch) when the current cooldown period expires. */
222240
final long expiryTime;

‎oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1147,7 +1147,7 @@ private Map<String, List<String>> getRequestMetadataWithSelfSignedJwt(URI uri)
11471147
}
11481148

11491149
Map<String, List<String>> requestMetadata = jwtCredentials.getRequestMetadata(null);
1150-
requestMetadata = addRegionalAccessBoundaryToRequestMetadata(requestMetadata);
1150+
requestMetadata = addRegionalAccessBoundaryToRequestMetadata(uri, requestMetadata);
11511151
refreshRegionalAccessBoundaryWithSelfSignedJwtIfExpired(uri, requestMetadata, null);
11521152
return addQuotaProjectIdToRequestMetadata(quotaProjectId, requestMetadata);
11531153
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1165,7 +1165,8 @@ public void refresh_regionalAccessBoundarySuccess() throws IOException, Interrup
11651165
RegionalAccessBoundary regionalAccessBoundary =
11661166
new RegionalAccessBoundary(
11671167
TestUtils.REGIONAL_ACCESS_BOUNDARY_ENCODED_LOCATION,
1168-
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS);
1168+
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS,
1169+
null);
11691170
transportFactory.transport.setRegionalAccessBoundary(regionalAccessBoundary);
11701171
transportFactory.transport.setServiceAccountEmail(defaultAccountEmail);
11711172

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1397,7 +1397,7 @@ public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
13971397
String.format(
13981398
IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKLOAD_POOL, projectNumber, poolId);
13991399
RegionalAccessBoundary workloadRab =
1400-
new RegionalAccessBoundary("workload-encoded", Collections.singletonList("workload-loc"));
1400+
new RegionalAccessBoundary("workload-encoded", Collections.singletonList("workload-loc"), null);
14011401
transportFactory.transport.addRegionalAccessBoundary(workloadRabUrl, workloadRab);
14021402

14031403
String saEmail =
@@ -1406,7 +1406,7 @@ public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
14061406
String.format(IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT, saEmail);
14071407
RegionalAccessBoundary impersonatedRab =
14081408
new RegionalAccessBoundary(
1409-
"impersonated-encoded", Collections.singletonList("impersonated-loc"));
1409+
"impersonated-encoded", Collections.singletonList("impersonated-loc"), null);
14101410
transportFactory.transport.addRegionalAccessBoundary(impersonatedRabUrl, impersonatedRab);
14111411

14121412
// Use a URL-based source that the mock transport can handle, to avoid file IO.
@@ -1456,7 +1456,7 @@ public void refresh_impersonated_workforce_regionalAccessBoundarySuccess()
14561456
String workforceRabUrl =
14571457
String.format(IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL, poolId);
14581458
RegionalAccessBoundary workforceRab =
1459-
new RegionalAccessBoundary("workforce-encoded", Collections.singletonList("workforce-loc"));
1459+
new RegionalAccessBoundary("workforce-encoded", Collections.singletonList("workforce-loc"), null);
14601460
transportFactory.transport.addRegionalAccessBoundary(workforceRabUrl, workforceRab);
14611461

14621462
String saEmail =
@@ -1465,7 +1465,7 @@ public void refresh_impersonated_workforce_regionalAccessBoundarySuccess()
14651465
String.format(IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT, saEmail);
14661466
RegionalAccessBoundary impersonatedRab =
14671467
new RegionalAccessBoundary(
1468-
"impersonated-encoded", Collections.singletonList("impersonated-loc"));
1468+
"impersonated-encoded", Collections.singletonList("impersonated-loc"), null);
14691469
transportFactory.transport.addRegionalAccessBoundary(impersonatedRabUrl, impersonatedRab);
14701470

14711471
// Use a URL-based source that the mock transport can handle, to avoid file IO.

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,6 @@ public void setUp() {
109109
@org.junit.After
110110
public void tearDown() {
111111
GoogleCredentials.disableRabRefreshForTest = false;
112-
RegionalAccessBoundary.setClockForTest(Clock.SYSTEM);
113-
RegionalAccessBoundaryManager.setClockForTest(Clock.SYSTEM);
114112
}
115113

116114
@Test
@@ -810,7 +808,7 @@ public void serialize_removesStaleRabHeaders() throws Exception {
810808
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
811809
RegionalAccessBoundary rab =
812810
new RegionalAccessBoundary(
813-
"test-encoded", Collections.singletonList("test-loc"), System.currentTimeMillis());
811+
"test-encoded", Collections.singletonList("test-loc"), System.currentTimeMillis(), null);
814812
transportFactory.transport.setRegionalAccessBoundary(rab);
815813
transportFactory.transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
816814

@@ -1006,7 +1004,8 @@ public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDat
10061004
RegionalAccessBoundary regionalAccessBoundary =
10071005
new RegionalAccessBoundary(
10081006
TestUtils.REGIONAL_ACCESS_BOUNDARY_ENCODED_LOCATION,
1009-
Collections.singletonList("us-central1"));
1007+
Collections.singletonList("us-central1"),
1008+
null);
10101009
transport.setRegionalAccessBoundary(regionalAccessBoundary);
10111010

10121011
ServiceAccountCredentials credentials =
@@ -1044,7 +1043,8 @@ public void regionalAccessBoundary_shouldRetryRegionalAccessBoundaryLookupOnFail
10441043
RegionalAccessBoundary regionalAccessBoundary =
10451044
new RegionalAccessBoundary(
10461045
TestUtils.REGIONAL_ACCESS_BOUNDARY_ENCODED_LOCATION,
1047-
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS);
1046+
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS,
1047+
null);
10481048
regionalAccessBoundaryTransport.setRegionalAccessBoundary(regionalAccessBoundary);
10491049

10501050
// This transport will be used for the access token refresh.
@@ -1124,8 +1124,8 @@ public void regionalAccessBoundary_cooldownDoublingAndRefresh()
11241124
.build();
11251125

11261126
TestClock testClock = new TestClock();
1127-
RegionalAccessBoundaryManager.setClockForTest(testClock);
1128-
RegionalAccessBoundary.setMaxRetryElapsedTimeMillisForTest(100);
1127+
credentials.clock = testClock;
1128+
credentials.regionalAccessBoundaryManager = new RegionalAccessBoundaryManager(testClock, 100);
11291129

11301130
// First attempt: triggers lookup, fails, enters 15m cooldown.
11311131
credentials.getRequestMetadata();
@@ -1155,7 +1155,7 @@ public void regionalAccessBoundary_cooldownDoublingAndRefresh()
11551155

11561156
// Set successful response.
11571157
transport.setRegionalAccessBoundary(
1158-
new RegionalAccessBoundary("0x123", Collections.emptyList()));
1158+
new RegionalAccessBoundary("0x123", Collections.emptyList(), null));
11591159

11601160
// Fourth attempt: triggers lookup, succeeds, resets cooldown.
11611161
credentials.getRequestMetadata();
@@ -1183,7 +1183,7 @@ public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
11831183
GoogleCredentials.disableRabRefreshForTest = false;
11841184
MockTokenServerTransport transport = new MockTokenServerTransport();
11851185
transport.setRegionalAccessBoundary(
1186-
new RegionalAccessBoundary("valid", Collections.singletonList("us-central1")));
1186+
new RegionalAccessBoundary("valid", Collections.singletonList("us-central1"), null));
11871187
// Add delay to lookup to ensure threads overlap.
11881188
transport.setResponseDelayMillis(500);
11891189

‎oauth2_http/javatests/com/google/auth/oauth2/ImpersonatedCredentialsTest.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,8 @@ public class ImpersonatedCredentialsTest extends BaseSerializationTest {
158158
public static final RegionalAccessBoundary REGIONAL_ACCESS_BOUNDARY =
159159
new RegionalAccessBoundary(
160160
TestUtils.REGIONAL_ACCESS_BOUNDARY_ENCODED_LOCATION,
161-
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS);
161+
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS,
162+
null);
162163

163164
private GoogleCredentials sourceCredentials;
164165
private MockIAMCredentialsServiceTransportFactory mockTransportFactory;

‎oauth2_http/javatests/com/google/auth/oauth2/MockExternalAccountCredentialsTransport.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,8 @@ public LowLevelHttpResponse execute() throws IOException {
210210
rab =
211211
new RegionalAccessBoundary(
212212
TestUtils.REGIONAL_ACCESS_BOUNDARY_ENCODED_LOCATION,
213-
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS);
213+
TestUtils.REGIONAL_ACCESS_BOUNDARY_LOCATIONS,
214+
null);
214215
}
215216
GenericJson responseJson = new GenericJson();
216217
responseJson.setFactory(OAuth2Utils.JSON_FACTORY);

‎oauth2_http/javatests/com/google/auth/oauth2/RegionalAccessBoundaryTest.java‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -58,21 +58,17 @@ public class RegionalAccessBoundaryTest {
5858
@Before
5959
public void setUp() {
6060
testClock = new TestClock();
61-
RegionalAccessBoundary.setClockForTest(testClock);
62-
RegionalAccessBoundaryManager.setClockForTest(testClock);
6361
}
6462

6563
@After
6664
public void tearDown() {
67-
RegionalAccessBoundary.setClockForTest(Clock.SYSTEM);
68-
RegionalAccessBoundaryManager.setClockForTest(Clock.SYSTEM);
6965
}
7066

7167
@Test
7268
public void testIsExpired() {
7369
long now = testClock.currentTimeMillis();
7470
RegionalAccessBoundary rab =
75-
new RegionalAccessBoundary("encoded", Collections.singletonList("loc"), now);
71+
new RegionalAccessBoundary("encoded", Collections.singletonList("loc"), now, testClock);
7672

7773
assertFalse(rab.isExpired());
7874

@@ -87,7 +83,7 @@ public void testIsExpired() {
8783
public void testShouldRefresh() {
8884
long now = testClock.currentTimeMillis();
8985
RegionalAccessBoundary rab =
90-
new RegionalAccessBoundary("encoded", Collections.singletonList("loc"), now);
86+
new RegionalAccessBoundary("encoded", Collections.singletonList("loc"), now, testClock);
9187

9288
// Initial state: fresh
9389
assertFalse(rab.shouldRefresh());
@@ -127,7 +123,7 @@ public void testManagerTriggersRefreshInGracePeriod() throws InterruptedExceptio
127123
HttpTransportFactory transportFactory = () -> transport;
128124
RegionalAccessBoundaryProvider provider = () -> url;
129125

130-
RegionalAccessBoundaryManager manager = new RegionalAccessBoundaryManager();
126+
RegionalAccessBoundaryManager manager = new RegionalAccessBoundaryManager(testClock);
131127

132128
// 1. Let's first get a RAB into the cache
133129
manager.triggerAsyncRefresh(transportFactory, provider, token, null);
@@ -184,7 +180,7 @@ public void testManagerTriggersRefreshInGracePeriod() throws InterruptedExceptio
184180

185181
@Test
186182
public void testManagerReleasesLockOnSchedulingFailure() {
187-
RegionalAccessBoundaryManager manager = new RegionalAccessBoundaryManager();
183+
RegionalAccessBoundaryManager manager = new RegionalAccessBoundaryManager(testClock);
188184
HttpTransportFactory transportFactory = () -> new MockHttpTransport();
189185
RegionalAccessBoundaryProvider provider = () -> "https://dummy";
190186
AccessToken token =

0 commit comments

Comments
 (0)