Skip to content

Commit e42ad38

Browse files
S-Saranya1S-Saranya1
committed
Additional Changes:
-Adding additional tests -Updating to use AtomicReference
1 parent c3d02a5 commit e42ad38

File tree

3 files changed

+68
-105
lines changed

3 files changed

+68
-105
lines changed

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

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import java.util.Collections;
2828
import java.util.Map;
2929
import java.util.Optional;
30+
import java.util.concurrent.atomic.AtomicReference;
3031
import java.util.function.Supplier;
3132
import software.amazon.awssdk.annotations.SdkPublicApi;
3233
import software.amazon.awssdk.annotations.SdkTestInternalApi;
@@ -82,8 +83,8 @@ private enum ApiVersion {
8283

8384
private static final String EC2_METADATA_TOKEN_TTL_HEADER = "x-aws-ec2-metadata-token-ttl-seconds";
8485
private static final String DEFAULT_TOKEN_TTL = "21600";
85-
private volatile ApiVersion apiVersion = ApiVersion.UNKNOWN;
86-
private String resolvedProfile = null;
86+
private final AtomicReference<ApiVersion> apiVersion = new AtomicReference<>(ApiVersion.UNKNOWN);
87+
private final AtomicReference<String> resolvedProfile = new AtomicReference<>();
8788

8889
private final Clock clock;
8990
private final String endpoint;
@@ -176,9 +177,9 @@ private RefreshResult<AwsCredentials> refreshCredentials() {
176177
.prefetchTime(prefetchTime(expiration))
177178
.build();
178179
} catch (Ec2MetadataClientException e) {
179-
if (e.statusCode() == 404 && apiVersion == ApiVersion.EXTENDED) {
180-
apiVersion = ApiVersion.LEGACY;
181-
resolvedProfile = null;
180+
if (e.statusCode() == 404 && apiVersion.compareAndSet(ApiVersion.EXTENDED, ApiVersion.LEGACY)) {
181+
log.debug(() -> "Unable to load credential path from extended API. Falling back to legacy API.");
182+
resolvedProfile.set(null);
182183
return refreshCredentials();
183184
}
184185
throw SdkClientException.create("Failed to load credentials from IMDS.", e);
@@ -226,7 +227,7 @@ public String toString() {
226227
}
227228

228229
private String getSecurityCredentialsResource() {
229-
return apiVersion == ApiVersion.LEGACY ?
230+
return apiVersion.get() == ApiVersion.LEGACY ?
230231
SECURITY_CREDENTIALS_RESOURCE :
231232
SECURITY_CREDENTIALS_EXTENDED_RESOURCE;
232233
}
@@ -309,8 +310,9 @@ private boolean isInsecureFallbackDisabled() {
309310
}
310311

311312
private String[] getSecurityCredentials(String imdsHostname, String metadataToken) {
312-
if (resolvedProfile != null) {
313-
return new String[]{resolvedProfile};
313+
String profile = resolvedProfile.get();
314+
if (profile != null) {
315+
return new String[]{profile};
314316
}
315317

316318
String urlBase = getSecurityCredentialsResource();
@@ -330,15 +332,13 @@ private String[] getSecurityCredentials(String imdsHostname, String metadataToke
330332
throw SdkClientException.builder().message("Unable to load credentials path").build();
331333
}
332334

333-
if (apiVersion == ApiVersion.UNKNOWN) {
334-
apiVersion = ApiVersion.EXTENDED;
335-
}
336-
resolvedProfile = securityCredentials[0];
335+
apiVersion.compareAndSet(ApiVersion.UNKNOWN, ApiVersion.EXTENDED);
336+
resolvedProfile.set(securityCredentials[0]);
337337
return securityCredentials;
338338

339339
} catch (Ec2MetadataClientException e) {
340-
if (apiVersion == ApiVersion.UNKNOWN) {
341-
apiVersion = ApiVersion.LEGACY;
340+
if (e.statusCode() == 404 && apiVersion.compareAndSet(ApiVersion.UNKNOWN, ApiVersion.LEGACY)) {
341+
log.debug(() -> "Unable to load credential path from extended API. Falling back to legacy API.");
342342
return getSecurityCredentials(imdsHostname, metadataToken);
343343
}
344344
throw SdkClientException.create("Failed to load credentials from IMDS.", e);

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -70,11 +70,11 @@ public LoadedCredentials loadCredentials(ResourcesEndpointProvider endpoint) {
7070
Validate.notNull(secretKey, "Failed to load secret key from metadata service.");
7171

7272
return new LoadedCredentials(accessKey.text(),
73-
secretKey.text(),
74-
token != null ? token.text() : null,
75-
expiration != null ? expiration.text() : null,
76-
accountId != null ? accountId.text() : null,
77-
providerName);
73+
secretKey.text(),
74+
token != null ? token.text() : null,
75+
expiration != null ? expiration.text() : null,
76+
accountId != null ? accountId.text() : null,
77+
providerName);
7878
} catch (SdkClientException e) {
7979
throw e;
8080
} catch (RuntimeException | IOException e) {

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

Lines changed: 49 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ void resolveCredentials_fallsBackToLegacy_noAccountId() {
105105
"\"SecretAccessKey\":\"SECRET_ACCESS_KEY\"," +
106106
"\"Token\":\"SESSION_TOKEN\"," +
107107
"\"Expiration\":\"%s\"," +
108-
"\"Code\":\"Success\"}", // No AccountId field at all
108+
"\"Code\":\"Success\"}",
109109
DateUtils.formatIso8601Date(Instant.now().plus(Duration.ofDays(1)))
110110
);
111111

@@ -159,119 +159,82 @@ void resolveCredentials_withInvalidProfile_throwsException() {
159159
}
160160

161161
@Test
162-
void resolveCredentials_withUnstableProfile_noAccountId_refreshesCredentials() {
163-
String firstCredentials = String.format(
164-
"{\"AccessKeyId\":\"ASIAIOSFODNN7EXAMPLE\"," +
165-
"\"SecretAccessKey\":\"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\"," +
166-
"\"Token\":\"AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKw\"," +
167-
"\"Expiration\":\"%s\"," +
168-
"\"Code\":\"Success\"," +
169-
"\"Type\":\"AWS-HMAC\"," +
170-
"\"LastUpdated\":\"2025-03-18T20:53:17.832308Z\"," +
171-
"\"UnexpectedElement7\":{\"Name\":\"ignore-me-7\"}}",
172-
DateUtils.formatIso8601Date(Instant.now().plus(Duration.ofDays(1)))
173-
);
174-
175-
String secondCredentials = String.format(
176-
"{\"AccessKeyId\":\"ASIAIOSFODNN7EXAMPLE\"," +
177-
"\"SecretAccessKey\":\"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\"," +
178-
"\"Token\":\"AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKw\"," +
162+
void resolveCredentials_cachesProfile_maintainsAccountId() {
163+
String credentialsWithAccountId = String.format(
164+
"{\"AccessKeyId\":\"ACCESS_KEY_ID\"," +
165+
"\"SecretAccessKey\":\"SECRET_ACCESS_KEY\"," +
166+
"\"Token\":\"SESSION_TOKEN\"," +
179167
"\"Expiration\":\"%s\"," +
180-
"\"Code\":\"Success\"," +
181-
"\"Type\":\"AWS-HMAC\"," +
182-
"\"LastUpdated\":\"2025-03-18T20:53:17.832308Z\"," +
183-
"\"UnexpectedElement7\":{\"Name\":\"ignore-me-7\"}}",
184-
DateUtils.formatIso8601Date(Instant.now().plus(Duration.ofDays(1)))
168+
"\"AccountId\":\"%s\"}",
169+
DateUtils.formatIso8601Date(Instant.now().plus(Duration.ofDays(1))),
170+
ACCOUNT_ID
185171
);
186172

187-
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH))
188-
.inScenario("Profile Change No AccountId")
189-
.whenScenarioStateIs("Started")
190-
.willReturn(aResponse().withBody("my-profile-0007"))
191-
.willSetStateTo("First Profile"));
192-
193-
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + "my-profile-0007"))
194-
.inScenario("Profile Change No AccountId")
195-
.whenScenarioStateIs("First Profile")
196-
.willReturn(aResponse().withBody(firstCredentials))
197-
.willSetStateTo("First Profile Done"));
198-
199-
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + "my-profile-0007"))
200-
.inScenario("Profile Change No AccountId")
201-
.whenScenarioStateIs("First Profile Done")
202-
.willReturn(aResponse().withStatus(404))
203-
.willSetStateTo("Profile Changed"));
204-
205-
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH))
206-
.inScenario("Profile Change No AccountId")
207-
.whenScenarioStateIs("Profile Changed")
208-
.willReturn(aResponse().withBody("my-profile-0007-b")));
209-
210-
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + "my-profile-0007-b"))
211-
.willReturn(aResponse().withBody(secondCredentials)));
212-
213-
wireMockServer.stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
214-
.willReturn(aResponse().withBody(TOKEN_STUB)));
215-
173+
stubSecureCredentialsResponse(aResponse().withBody(credentialsWithAccountId), true);
216174
InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
217175

176+
// First call
218177
AwsCredentials creds1 = provider.resolveCredentials();
219-
assertThat(creds1.accountId()).isEmpty();
220-
assertThat(creds1.accessKeyId()).isEqualTo("ASIAIOSFODNN7EXAMPLE");
178+
assertThat(creds1.accountId()).hasValue(ACCOUNT_ID);
221179

180+
// Second call - should use cached profile
222181
AwsCredentials creds2 = provider.resolveCredentials();
223-
assertThat(creds2.accountId()).isEmpty();
224-
assertThat(creds2.accessKeyId()).isEqualTo("ASIAIOSFODNN7EXAMPLE");
182+
assertThat(creds2.accountId()).hasValue(ACCOUNT_ID);
183+
184+
// Verify profile discovery only called once
185+
verify(1, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH)));
225186
}
226187

227188
@Test
228-
void resolveCredentials_withDiscoveredInvalidProfile_noAccountId_throwsException() {
229-
String invalidProfile = "my-profile-0008";
189+
void resolveCredentials_withNon404Error_doesNotFallbackToLegacy() {
190+
wireMockServer.stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
191+
.willReturn(aResponse().withBody(TOKEN_STUB)));
230192

231193
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH))
232-
.willReturn(aResponse().withBody(invalidProfile)));
194+
.willReturn(aResponse().withBody(PROFILE_NAME)));
233195

234-
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + invalidProfile))
235-
.willReturn(aResponse().withStatus(404)));
196+
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + PROFILE_NAME))
197+
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));
236198

237-
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH + invalidProfile))
238-
.willReturn(aResponse().withStatus(404)));
239-
240-
wireMockServer.stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
241-
.willReturn(aResponse().withBody(TOKEN_STUB)));
242199

243200
InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
244201

245202
assertThatThrownBy(() -> provider.resolveCredentials())
246203
.isInstanceOf(SdkClientException.class)
247204
.hasMessageContaining("Failed to load credentials from IMDS");
248-
}
249205

206+
// Verify extended endpoint was called
207+
verify(1, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH)));
208+
verify(1, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + PROFILE_NAME)));
209+
210+
// Verify legacy endpoint was NOT called
211+
verify(0, getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH)));
212+
verify(0, getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH + PROFILE_NAME)));
213+
}
214+
250215
@Test
251-
void resolveCredentials_cachesProfile_maintainsAccountId() {
252-
String credentialsWithAccountId = String.format(
253-
"{\"AccessKeyId\":\"ACCESS_KEY_ID\"," +
254-
"\"SecretAccessKey\":\"SECRET_ACCESS_KEY\"," +
255-
"\"Token\":\"SESSION_TOKEN\"," +
256-
"\"Expiration\":\"%s\"," +
257-
"\"AccountId\":\"%s\"}",
258-
DateUtils.formatIso8601Date(Instant.now().plus(Duration.ofDays(1))),
259-
ACCOUNT_ID
260-
);
216+
void resolveCredentials_withNon404ErrorOnProfileDiscovery_doesNotFallbackToLegacy() {
217+
wireMockServer.stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
218+
.willReturn(aResponse().withBody(TOKEN_STUB)));
261219

262-
stubSecureCredentialsResponse(aResponse().withBody(credentialsWithAccountId), true);
263-
InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
220+
wireMockServer.stubFor(get(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH))
221+
.willReturn(aResponse().withStatus(403).withBody("Forbidden")));
264222

265-
// First call
266-
AwsCredentials creds1 = provider.resolveCredentials();
267-
assertThat(creds1.accountId()).hasValue(ACCOUNT_ID);
223+
InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
268224

269-
// Second call - should use cached profile
270-
AwsCredentials creds2 = provider.resolveCredentials();
271-
assertThat(creds2.accountId()).hasValue(ACCOUNT_ID);
225+
assertThatThrownBy(() -> provider.resolveCredentials())
226+
.isInstanceOf(SdkClientException.class)
227+
.hasMessageContaining("Failed to load credentials from IMDS");
272228

273-
// Verify profile discovery only called once
229+
// Verify extended endpoint was called
274230
verify(1, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH)));
231+
232+
// Verify profile-specific endpoint was NOT called
233+
verify(0, getRequestedFor(urlPathEqualTo(CREDENTIALS_EXTENDED_RESOURCE_PATH + PROFILE_NAME)));
234+
235+
// Verify legacy endpoint was NOT called
236+
verify(0, getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH)));
237+
verify(0, getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH + PROFILE_NAME)));
275238
}
276239

277240
private void stubSecureCredentialsResponse(com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder responseDefinitionBuilder, boolean useExtended) {

0 commit comments

Comments
 (0)