Fix sts web identity provider not respecting prefetch time and stale time

Author: S-Saranya1

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

@@ -106,11 +106,16 @@ private StsWebIdentityTokenFileCredentialsProvider(Builder builder) {
                                                         .assumeRoleWithWebIdentityRequest(assumeRoleWithWebIdentityRequest.get())
                                                         .webIdentityTokenFile(credentialProperties.webIdentityTokenFile())
                                                         .build();
-            credentialsProviderLocal =
+
+            StsAssumeRoleWithWebIdentityCredentialsProvider.Builder internalBuilder =
                 StsAssumeRoleWithWebIdentityCredentialsProvider.builder()
                                                                .stsClient(builder.stsClient)
-                                                               .refreshRequest(supplier)
-                                                               .build();
+                                                               .refreshRequest(supplier);
+
+            internalBuilder.staleTime(this.staleTime())
+                           .prefetchTime(this.prefetchTime());
+
+            credentialsProviderLocal = internalBuilder.build();
         } catch (RuntimeException e) {
             // If we couldn't load the credentials provider for some reason, save an exception describing why. This exception
             // will only be raised on calls to getCredentials. We don't want to raise an exception here because it may be
@@ -181,7 +186,7 @@ private Builder() {
         }
 
         private Builder(StsWebIdentityTokenFileCredentialsProvider provider) {
-            super(StsWebIdentityTokenFileCredentialsProvider::new);
+            super(StsWebIdentityTokenFileCredentialsProvider::new, provider);
             this.roleArn = provider.roleArn;
             this.roleSessionName = provider.roleSessionName;
             this.webIdentityTokenFile = provider.webIdentityTokenFile;

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

@@ -30,10 +30,12 @@
 import software.amazon.awssdk.services.sts.model.Credentials;
 
 import java.nio.file.Paths;
+import java.time.Duration;
 import java.time.Instant;
 import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
 
 import static org.mockito.Mockito.when;
+import static org.assertj.core.api.Assertions.assertThat;
 
 @ExtendWith(MockitoExtension.class)
 class StsWebIdentityTokenCredentialProviderTest {
@@ -57,6 +59,18 @@ public void cleanUp(){
         ENVIRONMENT_VARIABLE_HELPER.reset();
     }
 
+    private void mockStsClientResponse(Instant expiration) {
+        when(stsClient.assumeRoleWithWebIdentity(Mockito.any(AssumeRoleWithWebIdentityRequest.class)))
+            .thenReturn(AssumeRoleWithWebIdentityResponse.builder()
+                                                         .credentials(Credentials.builder()
+                                                                                 .accessKeyId("key")
+                                                                                 .expiration(expiration)
+                                                                                 .sessionToken("session")
+                                                                                 .secretAccessKey("secret")
+                                                                                 .build())
+                                                         .build());
+    }
+
     @Test
     void createAssumeRoleWithWebIdentityTokenCredentialsProviderWithoutStsClient_throws_Exception() {
 
@@ -111,4 +125,57 @@ void createAssumeRoleWithWebIdentityTokenCredentialsProvider_raisesInResolveCred
         // exception should be raised lazily when resolving credentials, not at creation time.
         Assert.assertThrows(IllegalStateException.class, provider::resolveCredentials);
     }
+
+    @Test
+    void prefetchTimeAndStaleTime_withCustomConfiguration_shouldReturnConfiguredValues() {
+        mockStsClientResponse(Instant.now().plusSeconds(3600));
+
+        StsWebIdentityTokenFileCredentialsProvider provider =
+            StsWebIdentityTokenFileCredentialsProvider.builder()
+                                                      .stsClient(stsClient)
+                                                      .prefetchTime(Duration.ofMinutes(10))
+                                                      .staleTime(Duration.ofMinutes(2))
+                                                      .build();
+
+        provider.resolveCredentials();
+        
+        assertThat(provider.prefetchTime()).isEqualTo(Duration.ofMinutes(10));
+        assertThat(provider.staleTime()).isEqualTo(Duration.ofMinutes(2));
+
+    }
+
+    @Test
+    void prefetchTimeAndStaleTime_withoutConfiguration_shouldReturnDefaultValues() {
+        mockStsClientResponse(Instant.now().plusSeconds(3600));
+
+        StsWebIdentityTokenFileCredentialsProvider provider =
+            StsWebIdentityTokenFileCredentialsProvider.builder()
+                                                      .stsClient(stsClient)
+                                                      .build();
+
+
+        provider.resolveCredentials();
+        
+        assertThat(provider.prefetchTime()).isEqualTo(Duration.ofMinutes(5));
+        assertThat(provider.staleTime()).isEqualTo(Duration.ofMinutes(1));
+    }
+
+    @Test
+    void toBuilder_withTimingConfiguration_shouldPreserveConfiguration() {
+        mockStsClientResponse(Instant.now().plusSeconds(3600));
+        StsWebIdentityTokenFileCredentialsProvider originalProvider =
+            StsWebIdentityTokenFileCredentialsProvider.builder()
+                                                      .stsClient(stsClient)
+                                                      .prefetchTime(Duration.ofMinutes(8))
+                                                      .staleTime(Duration.ofMinutes(3))
+                                                      .build();
+
+
+        StsWebIdentityTokenFileCredentialsProvider copiedProvider = originalProvider.toBuilder().build();
+
+        copiedProvider.resolveCredentials();
+
+        assertThat(copiedProvider.prefetchTime()).isEqualTo(Duration.ofMinutes(8));
+        assertThat(copiedProvider.staleTime()).isEqualTo(Duration.ofMinutes(3));
+    }
 }