Skip to content

Commit a80dc94

Browse files
S-Saranya1S-Saranya1
authored
Fix StsWebIdentityTokenFileCredentialsProvider not respecting prefetchTime and staleTime configuration (#6650)
* Fix sts web identity provider not respecting prefetch time and stale time * Adding additional tests * Fix failed test * Add changelog * Address PR feedback * Change wait buffer to to reduce test flakiness * Add asyncCredentialUpdateEnabled getter in builder
1 parent b18ba2a commit a80dc94

File tree

4 files changed

+132
-3
lines changed

4 files changed

+132
-3
lines changed

‎.changes/next-release/bugfix-AWSSTS-bc7dfee.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "bugfix",
3+
"category": "AWS STS",
4+
"contributor": "",
5+
"description": "Fix `StsWebIdentityTokenFileCredentialsProvider` not respecting custom `prefetchTime` and `staleTime` configurations."
6+
}

‎services/sts/src/main/java/software/amazon/awssdk/services/sts/auth/StsCredentialsProvider.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,5 +230,13 @@ public B prefetchTime(Duration prefetchTime) {
230230
public T build() {
231231
return providerConstructor.apply((B) this);
232232
}
233+
234+
/**
235+
* Whether the provider should fetch credentials asynchronously in the background.
236+
* <p>By default, this is false.</p>
237+
*/
238+
Boolean asyncCredentialUpdateEnabled() {
239+
return asyncCredentialUpdateEnabled;
240+
}
233241
}
234242
}

‎services/sts/src/main/java/software/amazon/awssdk/services/sts/auth/StsWebIdentityTokenFileCredentialsProvider.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -106,11 +106,19 @@ private StsWebIdentityTokenFileCredentialsProvider(Builder builder) {
106106
.assumeRoleWithWebIdentityRequest(assumeRoleWithWebIdentityRequest.get())
107107
.webIdentityTokenFile(credentialProperties.webIdentityTokenFile())
108108
.build();
109-
credentialsProviderLocal =
109+
110+
StsAssumeRoleWithWebIdentityCredentialsProvider.Builder providerBuilder =
110111
StsAssumeRoleWithWebIdentityCredentialsProvider.builder()
111112
.stsClient(builder.stsClient)
112113
.refreshRequest(supplier)
113-
.build();
114+
.staleTime(this.staleTime())
115+
.prefetchTime(this.prefetchTime());
116+
117+
if (builder.asyncCredentialUpdateEnabled() != null) {
118+
providerBuilder.asyncCredentialUpdateEnabled(builder.asyncCredentialUpdateEnabled());
119+
}
120+
121+
credentialsProviderLocal = providerBuilder.build();
114122
} catch (RuntimeException e) {
115123
// If we couldn't load the credentials provider for some reason, save an exception describing why. This exception
116124
// will only be raised on calls to getCredentials. We don't want to raise an exception here because it may be
@@ -181,7 +189,7 @@ private Builder() {
181189
}
182190

183191
private Builder(StsWebIdentityTokenFileCredentialsProvider provider) {
184-
super(StsWebIdentityTokenFileCredentialsProvider::new);
192+
super(StsWebIdentityTokenFileCredentialsProvider::new, provider);
185193
this.roleArn = provider.roleArn;
186194
this.roleSessionName = provider.roleSessionName;
187195
this.webIdentityTokenFile = provider.webIdentityTokenFile;

‎services/sts/src/test/java/software/amazon/awssdk/services/sts/auth/StsWebIdentityTokenCredentialProviderTest.java‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,16 @@
2727
import software.amazon.awssdk.services.sts.StsClient;
2828
import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest;
2929
import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityResponse;
30+
import software.amazon.awssdk.services.sts.model.AssumedRoleUser;
3031
import software.amazon.awssdk.services.sts.model.Credentials;
3132

3233
import java.nio.file.Paths;
34+
import java.time.Duration;
3335
import java.time.Instant;
3436
import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
3537

3638
import static org.mockito.Mockito.when;
39+
import static org.assertj.core.api.Assertions.assertThat;
3740

3841
@ExtendWith(MockitoExtension.class)
3942
class StsWebIdentityTokenCredentialProviderTest {
@@ -111,4 +114,108 @@ void createAssumeRoleWithWebIdentityTokenCredentialsProvider_raisesInResolveCred
111114
// exception should be raised lazily when resolving credentials, not at creation time.
112115
Assert.assertThrows(IllegalStateException.class, provider::resolveCredentials);
113116
}
117+
118+
@Test
119+
void customPrefetchTime_actuallyTriggersRefreshEarly() throws InterruptedException {
120+
Mockito.reset(stsClient);
121+
122+
Instant tokenExpiration = Instant.now().plusSeconds(5);
123+
Duration customPrefetchTime = Duration.ofSeconds(2);
124+
Duration customStaleTime = Duration.ofSeconds(1);
125+
126+
when(stsClient.assumeRoleWithWebIdentity(Mockito.any(AssumeRoleWithWebIdentityRequest.class)))
127+
.thenReturn(AssumeRoleWithWebIdentityResponse.builder()
128+
.credentials(Credentials.builder()
129+
.accessKeyId("key1")
130+
.secretAccessKey("secret1")
131+
.sessionToken("session1")
132+
.expiration(tokenExpiration)
133+
.build())
134+
.assumedRoleUser(AssumedRoleUser.builder()
135+
.arn("arn:aws:iam::123456789012:role/test-role")
136+
.assumedRoleId("role:session")
137+
.build())
138+
.build())
139+
140+
.thenReturn(AssumeRoleWithWebIdentityResponse.builder()
141+
.credentials(Credentials.builder()
142+
.accessKeyId("key2")
143+
.secretAccessKey("secret2")
144+
.sessionToken("session2")
145+
.expiration(Instant.now().plusSeconds(8))
146+
.build())
147+
.assumedRoleUser(AssumedRoleUser.builder()
148+
.arn("arn:aws:iam::123456789012:role/test-role")
149+
.assumedRoleId("role:session")
150+
.build())
151+
.build());
152+
153+
154+
StsWebIdentityTokenFileCredentialsProvider provider =
155+
StsWebIdentityTokenFileCredentialsProvider.builder()
156+
.stsClient(stsClient)
157+
.asyncCredentialUpdateEnabled(true)
158+
.prefetchTime(customPrefetchTime)
159+
.staleTime(customStaleTime)
160+
.build();
161+
162+
try {
163+
assertThat(provider.prefetchTime()).isEqualTo(customPrefetchTime);
164+
assertThat(provider.staleTime()).isEqualTo(customStaleTime);
165+
166+
provider.resolveCredentials();
167+
Mockito.verify(stsClient, Mockito.times(1)).assumeRoleWithWebIdentity(Mockito.any(AssumeRoleWithWebIdentityRequest.class));
168+
169+
// Wait 5 seconds to ensure prefetch completes
170+
Thread.sleep(5_000);
171+
Mockito.verify(stsClient, Mockito.times(2)).assumeRoleWithWebIdentity(Mockito.any(AssumeRoleWithWebIdentityRequest.class));
172+
173+
} finally {
174+
provider.close();
175+
}
176+
}
177+
178+
@Test
179+
void defaultTiming_usesStandardValues() {
180+
StsWebIdentityTokenFileCredentialsProvider provider =
181+
StsWebIdentityTokenFileCredentialsProvider.builder()
182+
.stsClient(stsClient)
183+
.build();
184+
185+
try {
186+
assertThat(provider.prefetchTime()).isEqualTo(Duration.ofMinutes(5));
187+
assertThat(provider.staleTime()).isEqualTo(Duration.ofMinutes(1));
188+
} finally {
189+
provider.close();
190+
}
191+
}
192+
193+
@Test
194+
void toBuilder_preservesCustomTimingConfiguration() {
195+
Duration customPrefetch = Duration.ofMinutes(10);
196+
Duration customStale = Duration.ofMinutes(3);
197+
198+
StsWebIdentityTokenFileCredentialsProvider originalProvider =
199+
StsWebIdentityTokenFileCredentialsProvider.builder()
200+
.stsClient(stsClient)
201+
.prefetchTime(customPrefetch)
202+
.staleTime(customStale)
203+
.build();
204+
205+
try {
206+
assertThat(originalProvider.prefetchTime()).isEqualTo(customPrefetch);
207+
assertThat(originalProvider.staleTime()).isEqualTo(customStale);
208+
209+
StsWebIdentityTokenFileCredentialsProvider copiedProvider = originalProvider.toBuilder().build();
210+
211+
try {
212+
assertThat(copiedProvider.prefetchTime()).isEqualTo(customPrefetch);
213+
assertThat(copiedProvider.staleTime()).isEqualTo(customStale);
214+
} finally {
215+
copiedProvider.close();
216+
}
217+
} finally {
218+
originalProvider.close();
219+
}
220+
}
114221
}

0 commit comments

Comments
 (0)