Adding retry counter to the fallback logic

Author: S-Saranya1

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

@@ -82,10 +82,12 @@ private enum ApiVersion {
 
     private static final String EC2_METADATA_TOKEN_TTL_HEADER = "x-aws-ec2-metadata-token-ttl-seconds";
     private static final String DEFAULT_TOKEN_TTL = "21600";
+    private static final int MAX_PROFILE_RETRIES = 3;
 
     // These fields are accessed from methods called by CachedSupplier which provides thread safety through its ReentrantLock
     private ApiVersion apiVersion = ApiVersion.UNKNOWN;
     private String resolvedProfile = null;
+    private int profileRetryCount = 0;
 
     private final Clock clock;
     private final String endpoint;
@@ -173,6 +175,8 @@ private RefreshResult<AwsCredentials> refreshCredentials() {
             Instant expiration = credentials.getExpiration().orElse(null);
             log.debug(() -> "Loaded credentials from IMDS with expiration time of " + expiration);
 
+            profileRetryCount = 0;
+
             return RefreshResult.builder(credentials.getAwsCredentials())
                                 .staleTime(staleTime(expiration))
                                 .prefetchTime(prefetchTime(expiration))
@@ -181,6 +185,10 @@ private RefreshResult<AwsCredentials> refreshCredentials() {
             if (e.statusCode() == 404) {
                 log.debug(() -> "Resolved profile is no longer available. Resetting it and trying again.");
                 resolvedProfile = null;
+                profileRetryCount++;
+                if (profileRetryCount > MAX_PROFILE_RETRIES) {
+                    throw SdkClientException.create("Failed to load credentials from IMDS.", e);
+                }
                 return refreshCredentials();
             }
             throw SdkClientException.create("Failed to load credentials from IMDS.", e);

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

@@ -26,6 +26,7 @@
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 
 import com.github.tomakehurst.wiremock.junit5.WireMockExtension;
 import com.github.tomakehurst.wiremock.junit5.WireMockTest;
@@ -212,6 +213,82 @@ void resolveCredentials_withNon404ErrorOnProfileDiscovery_doesNotFallbackToLegac
         verify(0, getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH + PROFILE_NAME)));
     }
 
+    @Test
+    void resolveCredentials_withUnstableProfile_ReturnsCredentials() {
+        String initialProfile = "my-profile-0007";
+        String newProfile = "my-profile-0007-b";
+        String credentialsJson = String.format(
+            "{\"Code\":\"Success\"," +
+            "\"LastUpdated\":\"2025-03-18T20:53:17.832308Z\"," +
+            "\"Type\":\"AWS-HMAC\"," +
+            "\"AccessKeyId\":\"ASIAIOSFODNN7EXAMPLE\"," +
+            "\"SecretAccessKey\":\"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\"," +
+            "\"Token\":\"AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKw...\"," +
+            "\"Expiration\":\"2025-03-18T21:53:17.832308Z\"," +
+            "\"UnexpectedElement7\":{\"Name\":\"ignore-me-7\"}}");
+
+        wireMockServer.stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
+            .willReturn(aResponse().withBody(TOKEN_STUB)));
+
+        wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH))
+            .inScenario("unstable-profile")
+            .whenScenarioStateIs("Started")
+            .willReturn(aResponse().withBody(initialProfile))
+            .willSetStateTo("initial-profile-discovered"));
+
+
+        wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + initialProfile))
+            .inScenario("unstable-profile")
+            .whenScenarioStateIs("initial-profile-discovered")
+            .willReturn(aResponse().withStatus(404))
+            .willSetStateTo("profile-not-found"));
+
+        wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH))
+            .inScenario("unstable-profile")
+            .whenScenarioStateIs("profile-not-found")
+            .willReturn(aResponse().withBody(newProfile))
+            .willSetStateTo("new-profile-discovered"));
+
+        wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + newProfile))
+            .inScenario("unstable-profile")
+            .whenScenarioStateIs("new-profile-discovered")
+            .willReturn(aResponse().withBody(credentialsJson)));
+
+        InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
+
+        AwsCredentials creds1 = provider.resolveCredentials();
+        assertThat(creds1.accessKeyId()).isEqualTo("ASIAIOSFODNN7EXAMPLE");
+
+        AwsCredentials creds2 = assertDoesNotThrow(() -> provider.resolveCredentials());
+        assertThat(creds2.accessKeyId()).isEqualTo("ASIAIOSFODNN7EXAMPLE");
+
+        verify(1, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + initialProfile)));
+        verify(1, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + newProfile)));
+    }
+
+    @Test
+    void resolveCredentials_withTooManyProfileFailures_throwsException() {
+        String profile = "unstable-profile";
+
+        wireMockServer.stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
+            .willReturn(aResponse().withBody(TOKEN_STUB)));
+
+        wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH))
+            .willReturn(aResponse().withBody(profile)));
+
+        wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + profile))
+            .willReturn(aResponse().withStatus(404)));
+
+        InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
+
+        assertThatThrownBy(() -> provider.resolveCredentials())
+            .isInstanceOf(SdkClientException.class)
+            .hasMessageContaining("Failed to load credentials from IMDS.");
+
+        verify(4, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH)));
+        verify(4, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + profile)));
+    }
+
     private void stubSecureCredentialsResponse(com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder responseDefinitionBuilder, boolean useExtended) {
         wireMockServer.stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH)).willReturn(aResponse().withBody(TOKEN_STUB)));
         String path = useExtended ? CREDENTIALS_EXTENDED_RESOURCE_PATH : CREDENTIALS_RESOURCE_PATH;