Skip to content

Commit 8249823

Browse files
committed
RAB staleness due to Oauth caching and flaky tests fixed.
1 parent 97996af commit 8249823

11 files changed

Lines changed: 135 additions & 13 deletions

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

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,11 @@ String getFileType() {
111111
private final String universeDomain;
112112
private final boolean isExplicitUniverseDomain;
113113

114+
// Note: this is for internal testing use use only.
115+
// TODO: Fix unit test mocks so this can be removed
116+
// Refer -> https://github.com/googleapis/google-auth-library-java/issues/1898
117+
@VisibleForTesting static boolean disableRabRefreshForTest = false;
118+
114119
transient RegionalAccessBoundaryManager regionalAccessBoundaryManager =
115120
new RegionalAccessBoundaryManager();
116121

@@ -361,6 +366,9 @@ void refreshRegionalAccessBoundaryIfExpired(
361366
@Nullable AccessToken token,
362367
@Nullable java.util.concurrent.Executor executor)
363368
throws IOException {
369+
if (disableRabRefreshForTest) {
370+
return;
371+
}
364372
if (!(this instanceof RegionalAccessBoundaryProvider) || !isDefaultUniverseDomain()) {
365373
return;
366374
}
@@ -526,22 +534,22 @@ static Map<String, List<String>> addQuotaProjectIdToRequestMetadata(
526534
}
527535

528536
/**
529-
* Adds Regional Access Boundary header to requestMetadata if available.
537+
* Adds Regional Access Boundary header to requestMetadata if available. Overwrites if present.
530538
*
531-
* @return a new map with Regional Access Boundary header added if needed
539+
* @return a new map with Regional Access Boundary header added or updated
532540
*/
533541
Map<String, List<String>> addRegionalAccessBoundaryToRequestMetadata(
534542
Map<String, List<String>> requestMetadata) {
535543
Preconditions.checkNotNull(requestMetadata);
536544
RegionalAccessBoundary rab = getRegionalAccessBoundary();
537-
if (rab != null
538-
&& !requestMetadata.containsKey(RegionalAccessBoundary.X_ALLOWED_LOCATIONS_HEADER_KEY)) {
539-
return ImmutableMap.<String, List<String>>builder()
540-
.putAll(requestMetadata)
541-
.put(
542-
RegionalAccessBoundary.X_ALLOWED_LOCATIONS_HEADER_KEY,
543-
Collections.singletonList(rab.getEncodedLocations()))
544-
.build();
545+
if (rab != null) {
546+
// Overwrite the header to ensure the most recent async update is used,
547+
// preventing staleness if the token itself hasn't expired yet.
548+
Map<String, List<String>> newMetadata = new HashMap<>(requestMetadata);
549+
newMetadata.put(
550+
RegionalAccessBoundary.X_ALLOWED_LOCATIONS_HEADER_KEY,
551+
Collections.singletonList(rab.getEncodedLocations()));
552+
return ImmutableMap.copyOf(newMetadata);
545553
}
546554
return requestMetadata;
547555
}

oauth2_http/javatests/com/google/auth/oauth2/AwsCredentialsTest.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,16 @@
6464
@RunWith(JUnit4.class)
6565
public class AwsCredentialsTest extends BaseSerializationTest {
6666

67+
@org.junit.Before
68+
public void setUp() {
69+
GoogleCredentials.disableRabRefreshForTest = true;
70+
}
71+
72+
@org.junit.After
73+
public void tearDown() {
74+
GoogleCredentials.disableRabRefreshForTest = false;
75+
}
76+
6777
private static final String STS_URL = "https://sts.googleapis.com/v1/token";
6878
private static final String AWS_CREDENTIALS_URL = "https://169.254.169.254";
6979
private static final String AWS_CREDENTIALS_URL_WITH_ROLE = "https://169.254.169.254/roleName";
@@ -1402,6 +1412,8 @@ public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext cont
14021412

14031413
@Test
14041414
public void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1415+
GoogleCredentials.disableRabRefreshForTest = false;
1416+
14051417
MockExternalAccountCredentialsTransportFactory transportFactory =
14061418
new MockExternalAccountCredentialsTransportFactory();
14071419

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,16 @@
7979
@RunWith(JUnit4.class)
8080
public class ComputeEngineCredentialsTest extends BaseSerializationTest {
8181

82+
@org.junit.Before
83+
public void setUp() {
84+
GoogleCredentials.disableRabRefreshForTest = true;
85+
}
86+
87+
@org.junit.After
88+
public void tearDown() {
89+
GoogleCredentials.disableRabRefreshForTest = false;
90+
}
91+
8292
private static final URI CALL_URI = URI.create("http://googleapis.com/testapi/v1/foo");
8393

8494
private static final String TOKEN_URL =
@@ -1148,6 +1158,8 @@ public void idTokenWithAudience_503StatusCode() {
11481158

11491159
@Test
11501160
public void refresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1161+
GoogleCredentials.disableRabRefreshForTest = false;
1162+
11511163
String defaultAccountEmail = "default@email.com";
11521164
MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
11531165
RegionalAccessBoundary regionalAccessBoundary =

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,9 +128,15 @@ public HttpTransport create() {
128128

129129
@Before
130130
public void setup() {
131+
GoogleCredentials.disableRabRefreshForTest = true;
131132
transportFactory = new MockExternalAccountAuthorizedUserCredentialsTransportFactory();
132133
}
133134

135+
@org.junit.After
136+
public void tearDown() {
137+
GoogleCredentials.disableRabRefreshForTest = false;
138+
}
139+
134140
@Test
135141
public void builder_allFields() throws IOException {
136142
ExternalAccountAuthorizedUserCredentials credentials =
@@ -1217,6 +1223,8 @@ public void toString_expectedFormat() {
12171223

12181224
@Test
12191225
public void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1226+
GoogleCredentials.disableRabRefreshForTest = false;
1227+
12201228
ExternalAccountAuthorizedUserCredentials credentials =
12211229
ExternalAccountAuthorizedUserCredentials.newBuilder()
12221230
.setClientId(CLIENT_ID)

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,9 +89,15 @@ public HttpTransport create() {
8989

9090
@Before
9191
public void setup() {
92+
GoogleCredentials.disableRabRefreshForTest = true;
9293
transportFactory = new MockExternalAccountCredentialsTransportFactory();
9394
}
9495

96+
@org.junit.After
97+
public void tearDown() {
98+
GoogleCredentials.disableRabRefreshForTest = false;
99+
}
100+
95101
@Test
96102
public void fromStream_identityPoolCredentials() throws IOException {
97103
GenericJson json = buildJsonIdentityPoolCredential();
@@ -1305,6 +1311,7 @@ public void getRegionalAccessBoundaryUrl_invalidAudience_throws() {
13051311
@Test
13061312
public void refresh_workload_regionalAccessBoundarySuccess()
13071313
throws IOException, InterruptedException {
1314+
GoogleCredentials.disableRabRefreshForTest = false;
13081315
String audience =
13091316
"//iam.googleapis.com/projects/12345/locations/global/workloadIdentityPools/my-pool/providers/my-provider";
13101317

@@ -1339,6 +1346,7 @@ public String retrieveSubjectToken() throws IOException {
13391346
@Test
13401347
public void refresh_workforce_regionalAccessBoundarySuccess()
13411348
throws IOException, InterruptedException {
1349+
GoogleCredentials.disableRabRefreshForTest = false;
13421350
String audience =
13431351
"//iam.googleapis.com/locations/global/workforcePools/my-pool/providers/my-provider";
13441352

@@ -1373,6 +1381,7 @@ public String retrieveSubjectToken() throws IOException {
13731381
@Test
13741382
public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
13751383
throws IOException, InterruptedException {
1384+
GoogleCredentials.disableRabRefreshForTest = false;
13761385
String projectNumber = "12345";
13771386
String poolId = "my-pool";
13781387
String providerId = "my-provider";
@@ -1433,6 +1442,7 @@ public void refresh_impersonated_workload_regionalAccessBoundarySuccess()
14331442
@Test
14341443
public void refresh_impersonated_workforce_regionalAccessBoundarySuccess()
14351444
throws IOException, InterruptedException {
1445+
GoogleCredentials.disableRabRefreshForTest = false;
14361446
String poolId = "my-pool";
14371447
String providerId = "my-provider";
14381448
String audience =

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,6 @@
5656
import java.util.*;
5757
import java.util.concurrent.atomic.AtomicLong;
5858
import java.util.concurrent.atomic.AtomicReference;
59-
import org.junit.After;
6059
import org.junit.Test;
6160
import org.junit.runner.RunWith;
6261
import org.junit.runners.JUnit4;
@@ -100,8 +99,14 @@ public class GoogleCredentialsTest extends BaseSerializationTest {
10099
private static final String GOOGLE_DEFAULT_UNIVERSE = "googleapis.com";
101100
private static final String TPC_UNIVERSE = "foo.bar";
102101

103-
@After
102+
@org.junit.Before
103+
public void setUp() {
104+
GoogleCredentials.disableRabRefreshForTest = true;
105+
}
106+
107+
@org.junit.After
104108
public void tearDown() {
109+
GoogleCredentials.disableRabRefreshForTest = false;
105110
RegionalAccessBoundary.setClockForTest(Clock.SYSTEM);
106111
RegionalAccessBoundaryManager.setClockForTest(Clock.SYSTEM);
107112
}
@@ -948,6 +953,7 @@ public void getCredentialInfo_impersonatedServiceAccount() throws IOException {
948953
@Test
949954
public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDataSuccessfully()
950955
throws IOException, InterruptedException {
956+
GoogleCredentials.disableRabRefreshForTest = false;
951957
MockTokenServerTransport transport = new MockTokenServerTransport();
952958
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
953959
RegionalAccessBoundary regionalAccessBoundary =
@@ -981,6 +987,8 @@ public void regionalAccessBoundary_shouldFetchAndReturnRegionalAccessBoundaryDat
981987
@Test
982988
public void regionalAccessBoundary_shouldRetryRegionalAccessBoundaryLookupOnFailure()
983989
throws IOException, InterruptedException {
990+
GoogleCredentials.disableRabRefreshForTest = false;
991+
984992
// This transport will be used for the regional access boundary lookup.
985993
// We will configure it to fail on the first attempt.
986994
MockTokenServerTransport regionalAccessBoundaryTransport = new MockTokenServerTransport();
@@ -1030,6 +1038,7 @@ public com.google.api.client.http.LowLevelHttpRequest buildRequest(
10301038
@Test
10311039
public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIsPassed()
10321040
throws IOException {
1041+
GoogleCredentials.disableRabRefreshForTest = false;
10331042
MockTokenServerTransport transport = new MockTokenServerTransport();
10341043
// Return an expired access token.
10351044
transport.addServiceAccount(SA_CLIENT_EMAIL, "expired-token");
@@ -1052,6 +1061,7 @@ public void regionalAccessBoundary_refreshShouldNotThrowWhenNoValidAccessTokenIs
10521061
@Test
10531062
public void regionalAccessBoundary_cooldownDoublingAndRefresh()
10541063
throws IOException, InterruptedException {
1064+
GoogleCredentials.disableRabRefreshForTest = false;
10551065
MockTokenServerTransport transport = new MockTokenServerTransport();
10561066
transport.addServiceAccount(SA_CLIENT_EMAIL, ACCESS_TOKEN);
10571067
// Always fail lookup for now.
@@ -1111,6 +1121,7 @@ public void regionalAccessBoundary_cooldownDoublingAndRefresh()
11111121

11121122
@Test
11131123
public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() throws IOException {
1124+
GoogleCredentials.disableRabRefreshForTest = false;
11141125
// Use a simple AccessToken-based credential that won't try to refresh.
11151126
GoogleCredentials credentials = GoogleCredentials.create(new AccessToken("some-token", null));
11161127

@@ -1122,6 +1133,7 @@ public void regionalAccessBoundary_shouldFailOpenWhenRefreshCannotBeStarted() th
11221133
@Test
11231134
public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
11241135
throws IOException, InterruptedException {
1136+
GoogleCredentials.disableRabRefreshForTest = false;
11251137
MockTokenServerTransport transport = new MockTokenServerTransport();
11261138
transport.setRegionalAccessBoundary(
11271139
new RegionalAccessBoundary("valid", Collections.singletonList("us-central1")));
@@ -1150,6 +1162,7 @@ public void regionalAccessBoundary_deduplicationOfConcurrentRefreshes()
11501162

11511163
@Test
11521164
public void regionalAccessBoundary_shouldSkipRefreshForRegionalEndpoints() throws IOException {
1165+
GoogleCredentials.disableRabRefreshForTest = false;
11531166
MockTokenServerTransport transport = new MockTokenServerTransport();
11541167
GoogleCredentials credentials = createTestCredentials(transport);
11551168

oauth2_http/javatests/com/google/auth/oauth2/IdentityPoolCredentialsTest.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,16 @@ public class IdentityPoolCredentialsTest extends BaseSerializationTest {
7272
private static final IdentityPoolSubjectTokenSupplier testProvider =
7373
(ExternalAccountSupplierContext context) -> "testSubjectToken";
7474

75+
@org.junit.Before
76+
public void setUp() {
77+
GoogleCredentials.disableRabRefreshForTest = true;
78+
}
79+
80+
@org.junit.After
81+
public void tearDown() {
82+
GoogleCredentials.disableRabRefreshForTest = false;
83+
}
84+
7585
@Test
7686
public void createdScoped_clonedCredentialWithAddedScopes() throws IOException {
7787
IdentityPoolCredentials credentials =
@@ -1307,6 +1317,8 @@ void setShouldThrowOnGetCertificatePath(boolean shouldThrow) {
13071317

13081318
@Test
13091319
public void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1320+
GoogleCredentials.disableRabRefreshForTest = false;
1321+
13101322
MockExternalAccountCredentialsTransportFactory transportFactory =
13111323
new MockExternalAccountCredentialsTransportFactory();
13121324
HttpTransportFactory testingHttpTransportFactory = transportFactory;

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,10 +165,16 @@ public class ImpersonatedCredentialsTest extends BaseSerializationTest {
165165

166166
@Before
167167
public void setup() throws IOException {
168+
GoogleCredentials.disableRabRefreshForTest = true;
168169
sourceCredentials = getSourceCredentials();
169170
mockTransportFactory = new MockIAMCredentialsServiceTransportFactory();
170171
}
171172

173+
@org.junit.After
174+
public void tearDown() {
175+
GoogleCredentials.disableRabRefreshForTest = false;
176+
}
177+
172178
static GoogleCredentials getSourceCredentials() throws IOException {
173179
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
174180
PrivateKey privateKey = OAuth2Utils.privateKeyFromPkcs8(SA_PRIVATE_KEY_PKCS8);
@@ -1311,6 +1317,7 @@ public void serialize() throws IOException, ClassNotFoundException {
13111317

13121318
@Test
13131319
public void refresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
1320+
GoogleCredentials.disableRabRefreshForTest = false;
13141321
// Mock regional access boundary response
13151322
RegionalAccessBoundary regionalAccessBoundary = REGIONAL_ACCESS_BOUNDARY;
13161323

oauth2_http/javatests/com/google/auth/oauth2/LoggingTest.java

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@
6464
import java.util.Map;
6565
import org.junit.BeforeClass;
6666
import org.junit.Test;
67+
import org.junit.runner.RunWith;
68+
import org.junit.runners.JUnit4;
6769
import org.slf4j.Logger;
6870
import org.slf4j.LoggerFactory;
6971
import org.slf4j.event.KeyValuePair;
@@ -73,6 +75,7 @@
7375
* credentials test classes with addition of test logging appender setup and test logic for logging.
7476
* This duplicates tests setups, but centralizes logging test setup in this class.
7577
*/
78+
@RunWith(JUnit4.class)
7679
public class LoggingTest {
7780

7881
private TestAppender setupTestLogger(Class<?> clazz) {
@@ -91,13 +94,24 @@ public static void setup() {
9194
LoggingUtils.setEnvironmentProvider(testEnvironmentProvider);
9295
}
9396

97+
@org.junit.Before
98+
public void setUp() {
99+
GoogleCredentials.disableRabRefreshForTest = true;
100+
}
101+
102+
@org.junit.After
103+
public void tearDown() {
104+
GoogleCredentials.disableRabRefreshForTest = false;
105+
}
106+
94107
@Test
95108
public void userCredentials_getRequestMetadata_fromRefreshToken_hasAccessToken()
96109
throws IOException {
97110
TestAppender testAppender = setupTestLogger(UserCredentials.class);
98111
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
99112
transportFactory.transport.addClient(CLIENT_ID, CLIENT_SECRET);
100113
transportFactory.transport.addRefreshToken(REFRESH_TOKEN, ACCESS_TOKEN);
114+
101115
UserCredentials userCredentials =
102116
UserCredentials.newBuilder()
103117
.setClientId(CLIENT_ID)
@@ -210,6 +224,7 @@ public void serviceAccountCredentials_idTokenWithAudience_iamFlow_targetAudience
210224
transportFactory.getTransport().setTargetPrincipal(CLIENT_EMAIL);
211225
transportFactory.getTransport().setIdToken(DEFAULT_ID_TOKEN);
212226
transportFactory.getTransport().addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
227+
213228
ServiceAccountCredentials credentials =
214229
createDefaultBuilder()
215230
.setScopes(SCOPES)
@@ -441,7 +456,7 @@ public void getRequestMetadata_hasAccessToken() throws IOException {
441456

442457
TestUtils.assertContainsBearerToken(metadata, ACCESS_TOKEN);
443458

444-
assertEquals(8, testAppender.events.size());
459+
assertEquals(6, testAppender.events.size());
445460

446461
ILoggingEvent defaultServiceAccountRequest = testAppender.events.get(0);
447462
assertEquals(

oauth2_http/javatests/com/google/auth/oauth2/PluggableAuthCredentialsTest.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,17 @@
5757
/** Tests for {@link PluggableAuthCredentials}. */
5858
@RunWith(JUnit4.class)
5959
public class PluggableAuthCredentialsTest extends BaseSerializationTest {
60+
61+
@org.junit.Before
62+
public void setUp() {
63+
GoogleCredentials.disableRabRefreshForTest = true;
64+
}
65+
66+
@org.junit.After
67+
public void tearDown() {
68+
GoogleCredentials.disableRabRefreshForTest = false;
69+
}
70+
6071
// The default timeout for waiting for the executable to finish (30 seconds).
6172
private static final int DEFAULT_EXECUTABLE_TIMEOUT_MS = 30 * 1000;
6273
// The minimum timeout for waiting for the executable to finish (5 seconds).
@@ -608,6 +619,8 @@ public void serialize() throws IOException, ClassNotFoundException {
608619

609620
@Test
610621
public void testRefresh_regionalAccessBoundarySuccess() throws IOException, InterruptedException {
622+
GoogleCredentials.disableRabRefreshForTest = false;
623+
611624
MockExternalAccountCredentialsTransportFactory transportFactory =
612625
new MockExternalAccountCredentialsTransportFactory();
613626
transportFactory.transport.setExpireTime(TestUtils.getDefaultExpireTime());

0 commit comments

Comments
 (0)