Add stale time config to InstanceProfileCredentialsProvider (#5758)

* Add stale time config to InstanceProfileCredentialsProvider to allow for refreshing credentials earlier due to stale value. This will help prevent returning invalid credentials when an error is encountered during asynchronous refresh.

* fix spotbug

* add configurable retry policy

* javadoc

* remove InstanceProfileCredentialsRetryPolicy

* default retryPolicy set in StaticResourcesEndpointProvider ctor

Author: L-Applin

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/InstanceProfileCredentialsProvider.java

@@ -88,6 +88,8 @@ public final class InstanceProfileCredentialsProvider
 
     private final String profileName;
 
+    private final Duration staleTime;
+
     /**
      * @see #builder()
      */
@@ -108,6 +110,8 @@ private InstanceProfileCredentialsProvider(BuilderImpl builder) {
                                      .profileName(profileName)
                                      .build();
 
+        this.staleTime = Validate.getOrDefault(builder.staleTime, () -> Duration.ofSeconds(1));
+
         if (Boolean.TRUE.equals(builder.asyncCredentialUpdateEnabled)) {
             Validate.paramNotBlank(builder.asyncThreadName, "asyncThreadName");
             this.credentialsCache = CachedSupplier.builder(this::refreshCredentials)
@@ -174,7 +178,7 @@ private Instant staleTime(Instant expiration) {
             return null;
         }
 
-        return expiration.minusSeconds(1);
+        return expiration.minus(staleTime);
     }
 
     private Instant prefetchTime(Instant expiration) {
@@ -340,6 +344,18 @@ public interface Builder extends HttpCredentialsProvider.Builder<InstanceProfile
          */
         Builder profileName(String profileName);
 
+        /**
+         * Configure the amount of time before the moment of expiration of credentials for which to consider the credentials to
+         * be stale. A higher value can lead to a higher rate of request being made to the Amazon EC2 Instance Metadata Service.
+         * The default is 1 sec.
+         * <p>Increasing this value to a higher value (10s or more) may help with situations where a higher load on the instance
+         * metadata service causes it to return 503s error, for which the SDK may not be able to recover fast enough and
+         * returns expired credentials.
+         *
+         * @param duration the amount of time before expiration for when to consider the credentials to be stale and need refresh
+         */
+        Builder staleTime(Duration duration);
+
         /**
          * Build a {@link InstanceProfileCredentialsProvider} from the provided configuration.
          */
@@ -355,6 +371,7 @@ static final class BuilderImpl implements Builder {
         private String asyncThreadName;
         private Supplier<ProfileFile> profileFile;
         private String profileName;
+        private Duration staleTime;
 
         private BuilderImpl() {
             asyncThreadName("instance-profile-credentials-provider");
@@ -367,6 +384,7 @@ private BuilderImpl(InstanceProfileCredentialsProvider provider) {
             this.asyncThreadName = provider.asyncThreadName;
             this.profileFile = provider.profileFile;
             this.profileName = provider.profileName;
+            this.staleTime = provider.staleTime;
         }
 
         Builder clock(Clock clock) {
@@ -435,6 +453,16 @@ public void setProfileName(String profileName) {
             profileName(profileName);
         }
 
+        @Override
+        public Builder staleTime(Duration duration) {
+            this.staleTime = duration;
+            return this;
+        }
+
+        public void setStaleTime(Duration duration) {
+            staleTime(duration);
+        }
+
         @Override
         public InstanceProfileCredentialsProvider build() {
             return new InstanceProfileCredentialsProvider(this);

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/internal/ContainerCredentialsRetryPolicy.java

@@ -17,10 +17,15 @@
 
 import java.io.IOException;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.auth.credentials.ContainerCredentialsProvider;
+import software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider;
 import software.amazon.awssdk.http.HttpStatusFamily;
 import software.amazon.awssdk.regions.util.ResourcesEndpointRetryParameters;
 import software.amazon.awssdk.regions.util.ResourcesEndpointRetryPolicy;
 
+/**
+ * Retry policy shared by {@link InstanceProfileCredentialsProvider} and {@link ContainerCredentialsProvider#}.
+ */
 @SdkInternalApi
 public final class ContainerCredentialsRetryPolicy implements ResourcesEndpointRetryPolicy {
 

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/internal/StaticResourcesEndpointProvider.java

@@ -24,23 +24,27 @@
 import java.util.Optional;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.regions.util.ResourcesEndpointProvider;
+import software.amazon.awssdk.regions.util.ResourcesEndpointRetryPolicy;
 import software.amazon.awssdk.utils.Validate;
 
 @SdkInternalApi
 public final class StaticResourcesEndpointProvider implements ResourcesEndpointProvider {
     private final URI endpoint;
     private final Map<String, String> headers;
     private final Duration connectionTimeout;
+    private final ResourcesEndpointRetryPolicy retryPolicy;
 
     private StaticResourcesEndpointProvider(URI endpoint,
-                                           Map<String, String> additionalHeaders,
-                                           Duration customTimeout) {
+                                            Map<String, String> additionalHeaders,
+                                            Duration customTimeout,
+                                            ResourcesEndpointRetryPolicy retryPolicy) {
         this.endpoint = Validate.paramNotNull(endpoint, "endpoint");
         this.headers = ResourcesEndpointProvider.super.headers();
         if (additionalHeaders != null) {
             this.headers.putAll(additionalHeaders);
         }
         this.connectionTimeout = customTimeout;
+        this.retryPolicy = Validate.getOrDefault(retryPolicy, () -> ResourcesEndpointRetryPolicy.NO_RETRY);
     }
 
     @Override
@@ -58,10 +62,16 @@ public Map<String, String> headers() {
         return Collections.unmodifiableMap(headers);
     }
 
+    @Override
+    public ResourcesEndpointRetryPolicy retryPolicy() {
+        return this.retryPolicy;
+    }
+
     public static class Builder {
         private URI endpoint;
         private Map<String, String> additionalHeaders = new HashMap<>();
         private Duration customTimeout;
+        private ResourcesEndpointRetryPolicy retryPolicy;
 
         public Builder endpoint(URI endpoint) {
             this.endpoint = Validate.paramNotNull(endpoint, "endpoint");
@@ -80,8 +90,13 @@ public Builder connectionTimeout(Duration timeout) {
             return this;
         }
 
+        public Builder retryPolicy(ResourcesEndpointRetryPolicy retryPolicy) {
+            this.retryPolicy = retryPolicy;
+            return this;
+        }
+
         public StaticResourcesEndpointProvider build() {
-            return new StaticResourcesEndpointProvider(endpoint, additionalHeaders, customTimeout);
+            return new StaticResourcesEndpointProvider(endpoint, additionalHeaders, customTimeout, retryPolicy);
         }
     }
 

core/auth/src/test/java/software/amazon/awssdk/auth/credentials/InstanceProfileCredentialsProviderTest.java

@@ -17,13 +17,15 @@
 
 import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
 import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
+import static com.github.tomakehurst.wiremock.client.WireMock.exactly;
 import static com.github.tomakehurst.wiremock.client.WireMock.get;
 import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.put;
 import static com.github.tomakehurst.wiremock.client.WireMock.putRequestedFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
+import static com.github.tomakehurst.wiremock.stubbing.Scenario.STARTED;
 import static java.time.temporal.ChronoUnit.HOURS;
 import static java.time.temporal.ChronoUnit.MINUTES;
 import static java.time.temporal.ChronoUnit.SECONDS;
@@ -55,6 +57,7 @@
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.Mockito;
 import software.amazon.awssdk.core.SdkSystemSetting;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.util.SdkUserAgent;
@@ -596,13 +599,74 @@ void imdsCallFrequencyIsLimited() {
         }
     }
 
+    @Test
+    void testErrorWhileCacheIsStale_shouldRecover() {
+        AdjustableClock clock = new AdjustableClock();
+
+        Instant now = Instant.now();
+        Instant expiration = now.plus(Duration.ofHours(6));
+
+        String successfulCredentialsResponse =
+            "{"
+            + "\"AccessKeyId\":\"ACCESS_KEY_ID\","
+            + "\"SecretAccessKey\":\"SECRET_ACCESS_KEY\","
+            + "\"Expiration\":\"" + DateUtils.formatIso8601Date(expiration) + '"'
+            + "}";
+
+        String staleResponse =
+            "{"
+            + "\"AccessKeyId\":\"ACCESS_KEY_ID_2\","
+            + "\"SecretAccessKey\":\"SECRET_ACCESS_KEY_2\","
+            + "\"Expiration\":\"" + DateUtils.formatIso8601Date(Instant.now()) + '"'
+            + "}";
+
+
+        Duration staleTime = Duration.ofMinutes(5);
+        AwsCredentialsProvider provider = credentialsProviderWithClock(clock, staleTime);
+
+        // cache expiration with expiration = 6 hours
+        clock.time = now;
+        stubSecureCredentialsResponse(aResponse().withBody(successfulCredentialsResponse));
+        AwsCredentials validCreds = provider.resolveCredentials();
+
+        // failure while cache is stale
+        clock.time = expiration.minus(staleTime.minus(Duration.ofMinutes(2)));
+        stubTokenFetchErrorResponse(aResponse().withFixedDelay(2000).withBody(STUB_CREDENTIALS), 500);
+        stubSecureCredentialsResponse(aResponse().withBody(staleResponse));
+        AwsCredentials refreshedWhileStale = provider.resolveCredentials();
+
+        assertThat(refreshedWhileStale).isNotEqualTo(validCreds);
+        assertThat(refreshedWhileStale.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY_2");
+    }
+
+    @Test
+    void shouldNotRetry_whenSucceeds() {
+        stubSecureCredentialsResponse(aResponse().withBody(STUB_CREDENTIALS));
+        InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
+        AwsCredentials credentials = provider.resolveCredentials();
+        assertThat(credentials.accessKeyId()).isEqualTo("ACCESS_KEY_ID");
+        assertThat(credentials.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY");
+        assertThat(credentials.providerName()).isPresent().contains("InstanceProfileCredentialsProvider");
+        verifyImdsCallWithToken();
+        WireMock.verify(exactly(1), getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH + "some-profile")));
+    }
+
     private AwsCredentialsProvider credentialsProviderWithClock(Clock clock) {
         InstanceProfileCredentialsProvider.BuilderImpl builder =
             (InstanceProfileCredentialsProvider.BuilderImpl) InstanceProfileCredentialsProvider.builder();
         builder.clock(clock);
         return builder.build();
     }
 
+    private AwsCredentialsProvider credentialsProviderWithClock(Clock clock, Duration staleTime) {
+        InstanceProfileCredentialsProvider.BuilderImpl builder =
+            (InstanceProfileCredentialsProvider.BuilderImpl) InstanceProfileCredentialsProvider.builder();
+        builder.clock(clock);
+        builder.staleTime(staleTime);
+        return builder.build();
+    }
+
+
     private static class AdjustableClock extends Clock {
         private Instant time;
 

core/auth/src/test/resources/log4j2.properties

@@ -36,3 +36,9 @@ rootLogger.appenderRef.stdout.ref = ConsoleAppender
 #
 #logger.netty.name = io.netty.handler.logging
 #logger.netty.level = debug
+
+#logger.cache.name = software.amazon.awssdk.utils.cache.CachedSupplier
+#logger.cache.level = DEBUG
+
+#logger.instance.name = software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider
+#logger.instance.level = DEBUG