Support account-ID based endpoints for IMDS and container credentials

.changes/next-release/feature-AWSSDKforJavav2-f54dd50.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Add support for account-ID based endpoints when using IMDS credential provider. The SDK now sends the `X-amzn-extended: 1` header on IMDS credential requests and parses the `AccountId` field from the response when available."
+}

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

@@ -25,6 +25,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -76,6 +77,8 @@ public final class InstanceProfileCredentialsProvider
     private static final String TOKEN_RESOURCE = "/latest/api/token";
     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 String EXTENDED_CREDENTIALS_HEADER = "X-amzn-extended";
+    private static final String EXTENDED_CREDENTIALS_VALUE = "1";
 
     private final Clock clock;
     private final String endpoint;
@@ -225,7 +228,7 @@ private ResourcesEndpointProvider createEndpointProvider() {
         return StaticResourcesEndpointProvider.builder()
                                               .endpoint(URI.create(imdsHostname + SECURITY_CREDENTIALS_RESOURCE
                                                                    + securityCredentials[0]))
-                                              .headers(getTokenHeaders(token))
+                                              .headers(getCredentialHeaders(token))
                                               .connectionTimeout(Duration.ofMillis(
                                                   this.configProvider.serviceTimeout()))
                                               .build();
@@ -320,6 +323,19 @@ private Map<String, String> getTokenHeaders(String metadataToken) {
         return Collections.singletonMap(EC2_METADATA_TOKEN_HEADER, metadataToken);
     }
 
+    /**
+     * Returns headers for the credential GET request. Includes the extended credentials header to signal IMDS to
+     * return the AccountId field in the response, and the IMDSv2 token if available.
+     */
+    private Map<String, String> getCredentialHeaders(String metadataToken) {
+        Map<String, String> headers = new HashMap<>();
+        if (metadataToken != null) {
+            headers.put(EC2_METADATA_TOKEN_HEADER, metadataToken);
+        }
+        headers.put(EXTENDED_CREDENTIALS_HEADER, EXTENDED_CREDENTIALS_VALUE);
+        return Collections.unmodifiableMap(headers);
+    }
+
     @Override
     public Builder toBuilder() {
         return new BuilderImpl(this);

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

@@ -64,6 +64,7 @@ public LoadedCredentials loadCredentials(ResourcesEndpointProvider endpoint) {
             JsonNode secretKey = node.get("SecretAccessKey");
             JsonNode token = node.get("Token");
             JsonNode expiration = node.get("Expiration");
+            JsonNode accountId = node.get("AccountId");
 
             Validate.notNull(accessKey, "Failed to load access key from metadata service.");
             Validate.notNull(secretKey, "Failed to load secret key from metadata service.");
@@ -72,6 +73,7 @@ public LoadedCredentials loadCredentials(ResourcesEndpointProvider endpoint) {
                                          secretKey.text(),
                                          token != null ? token.text() : null,
                                          expiration != null ? expiration.text() : null,
+                                         accountId != null ? accountId.text() : null,
                                          providerName);
         } catch (SdkClientException e) {
             throw e;
@@ -88,13 +90,16 @@ public static final class LoadedCredentials {
         private final String secretKey;
         private final String token;
         private final Instant expiration;
+        private final String accountId;
         private final String providerName;
 
-        private LoadedCredentials(String accessKeyId, String secretKey, String token, String expiration, String providerName) {
+        private LoadedCredentials(String accessKeyId, String secretKey, String token, String expiration,
+                                  String accountId, String providerName) {
             this.accessKeyId = Validate.paramNotBlank(accessKeyId, "accessKeyId");
             this.secretKey = Validate.paramNotBlank(secretKey, "secretKey");
             this.token = token;
             this.expiration = expiration == null ? null : parseExpiration(expiration);
+            this.accountId = accountId;
             this.providerName = providerName;
         }
 
@@ -105,11 +110,13 @@ public AwsCredentials getAwsCredentials() {
                                         .secretAccessKey(secretKey)
                                         .sessionToken(token)
                                         .expirationTime(expiration)
+                                        .accountId(accountId)
                                         .providerName(providerName)
                                         .build() :
                    AwsBasicCredentials.builder()
                                       .accessKeyId(accessKeyId)
                                       .secretAccessKey(secretKey)
+                                      .accountId(accountId)
                                       .providerName(providerName)
                                       .build();
         }

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

@@ -51,14 +51,24 @@ public class HttpCredentialsLoaderTest {
 
     private static String successResponseWithInvalidBody;
 
+    private static String successResponseWithAccountId;
+
+    private static String successResponseWithUnexpectedFields;
+
     @BeforeClass
     public static void setup() throws IOException {
         try (InputStream successInputStream = HttpCredentialsLoaderTest.class.getResourceAsStream
             ("/resources/wiremock/successResponse.json");
              InputStream responseWithInvalidBodyInputStream = HttpCredentialsLoaderTest.class.getResourceAsStream
-                 ("/resources/wiremock/successResponseWithInvalidBody.json")) {
+                 ("/resources/wiremock/successResponseWithInvalidBody.json");
+             InputStream responseWithAccountIdInputStream = HttpCredentialsLoaderTest.class.getResourceAsStream
+                 ("/resources/wiremock/successResponseWithAccountId.json");
+             InputStream responseWithUnexpectedFieldsInputStream = HttpCredentialsLoaderTest.class.getResourceAsStream
+                 ("/resources/wiremock/successResponseWithUnexpectedFields.json")) {
             successResponse = IoUtils.toUtf8String(successInputStream);
             successResponseWithInvalidBody = IoUtils.toUtf8String(responseWithInvalidBodyInputStream);
+            successResponseWithAccountId = IoUtils.toUtf8String(responseWithAccountIdInputStream);
+            successResponseWithUnexpectedFields = IoUtils.toUtf8String(responseWithUnexpectedFieldsInputStream);
         }
     }
 
@@ -76,6 +86,58 @@ public void testLoadCredentialsParsesJsonResponseProperly() {
         assertThat(credentials.accessKeyId()).isEqualTo("ACCESS_KEY_ID");
         assertThat(credentials.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY");
         assertThat(credentials.sessionToken()).isEqualTo("TOKEN_TOKEN_TOKEN");
+        assertThat(credentials.accountId()).isNotPresent();
+    }
+
+    /**
+     * Test that loadCredentials parses AccountId when present in the response.
+     */
+    @Test
+    public void testLoadCredentialsParsesAccountIdWhenPresent() {
+        stubForSuccessResponseWithCustomBody(successResponseWithAccountId);
+
+        HttpCredentialsLoader credentialsProvider = HttpCredentialsLoader.create(PROVIDER_NAME);
+        AwsSessionCredentials credentials = (AwsSessionCredentials) credentialsProvider.loadCredentials(testEndpointProvider())
+                                                                                       .getAwsCredentials();
+
+        assertThat(credentials.accessKeyId()).isEqualTo("ACCESS_KEY_ID");
+        assertThat(credentials.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY");
+        assertThat(credentials.sessionToken()).isEqualTo("TOKEN_TOKEN_TOKEN");
+        assertThat(credentials.accountId()).isPresent().hasValue("123456789012");
+    }
+
+    /**
+     * Test that loadCredentials succeeds without AccountId when it's absent from the response.
+     */
+    @Test
+    public void testLoadCredentialsSucceedsWithoutAccountId() {
+        stubForSuccessResponseWithCustomBody(successResponse);
+
+        HttpCredentialsLoader credentialsProvider = HttpCredentialsLoader.create(PROVIDER_NAME);
+        AwsSessionCredentials credentials = (AwsSessionCredentials) credentialsProvider.loadCredentials(testEndpointProvider())
+                                                                                       .getAwsCredentials();
+
+        assertThat(credentials.accessKeyId()).isEqualTo("ACCESS_KEY_ID");
+        assertThat(credentials.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY");
+        assertThat(credentials.accountId()).isNotPresent();
+    }
+
+    /**
+     * Test that loadCredentials does not fail when response contains unexpected/unknown fields.
+     * This verifies the SEP requirement: "Parsing MUST NOT fail in the presence of unexpected fields."
+     */
+    @Test
+    public void testLoadCredentialsIgnoresUnexpectedFields() {
+        stubForSuccessResponseWithCustomBody(successResponseWithUnexpectedFields);
+
+        HttpCredentialsLoader credentialsProvider = HttpCredentialsLoader.create(PROVIDER_NAME);
+        AwsSessionCredentials credentials = (AwsSessionCredentials) credentialsProvider.loadCredentials(testEndpointProvider())
+                                                                                       .getAwsCredentials();
+
+        assertThat(credentials.accessKeyId()).isEqualTo("ACCESS_KEY_ID");
+        assertThat(credentials.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY");
+        assertThat(credentials.sessionToken()).isEqualTo("TOKEN_TOKEN_TOKEN");
+        assertThat(credentials.accountId()).isPresent().hasValue("123456789012");
     }
 
     /**

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

@@ -81,7 +81,13 @@ public class InstanceProfileCredentialsProviderTest {
     private static final String STUB_CREDENTIALS = "{\"AccessKeyId\":\"ACCESS_KEY_ID\",\"SecretAccessKey\":\"SECRET_ACCESS_KEY\","
             + "\"Expiration\":\"" + DateUtils.formatIso8601Date(Instant.now().plus(Duration.ofDays(1)))
             + "\"}";
+    private static final String STUB_CREDENTIALS_WITH_ACCOUNT_ID =
+            "{\"AccessKeyId\":\"ACCESS_KEY_ID\",\"SecretAccessKey\":\"SECRET_ACCESS_KEY\","
+            + "\"Token\":\"TOKEN_TOKEN_TOKEN\","
+            + "\"Expiration\":\"" + DateUtils.formatIso8601Date(Instant.now().plus(Duration.ofDays(1)))
+            + "\",\"AccountId\":\"123456789012\"}";
     private static final String TOKEN_HEADER = "x-aws-ec2-metadata-token";
+    private static final String EXTENDED_CREDENTIALS_HEADER = "X-amzn-extended";
     private static final String USER_AGENT_HEADER = "User-Agent";
     private static final String TOKEN_STUB = "some-token";
     private static final String PROFILE_NAME = "some-profile";
@@ -161,6 +167,47 @@ void resolveCredentials_usesTokenByDefault() {
         verifyImdsCallWithToken();
     }
 
+    @Test
+    void resolveCredentials_sendsExtendedHeaderOnCredentialGet() {
+        stubSecureCredentialsResponse(aResponse().withBody(STUB_CREDENTIALS_WITH_ACCOUNT_ID));
+        InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
+        provider.resolveCredentials();
+
+        // Verify extended header is sent on credential GET
+        WireMock.verify(getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH + PROFILE_NAME))
+                            .withHeader(EXTENDED_CREDENTIALS_HEADER, equalTo("1")));
+
+        // Verify extended header is NOT sent on profile discovery GET
+        WireMock.verify(getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH))
+                            .withoutHeader(EXTENDED_CREDENTIALS_HEADER));
+
+        // Verify extended header is NOT sent on token PUT
+        WireMock.verify(putRequestedFor(urlPathEqualTo(TOKEN_RESOURCE_PATH))
+                            .withoutHeader(EXTENDED_CREDENTIALS_HEADER));
+    }
+
+    @Test
+    void resolveCredentials_parsesAccountIdWhenPresent() {
+        stubSecureCredentialsResponse(aResponse().withBody(STUB_CREDENTIALS_WITH_ACCOUNT_ID));
+        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.accountId()).isPresent().hasValue("123456789012");
+    }
+
+    @Test
+    void resolveCredentials_succeedsWithoutAccountId() {
+        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.accountId()).isNotPresent();
+    }
+
     @Test
     void resolveCredentials_WhenConnectionDelaySetToHighValue() {
         environmentVariableHelper.set(SdkSystemSetting.AWS_METADATA_SERVICE_TIMEOUT, "10");

core/auth/src/test/resources/resources/wiremock/successResponseWithAccountId.json

@@ -0,0 +1,7 @@
+{
+  "AccessKeyId": "ACCESS_KEY_ID",
+  "SecretAccessKey": "SECRET_ACCESS_KEY",
+  "Token": "TOKEN_TOKEN_TOKEN",
+  "Expiration": "3000-05-03T04:55:54Z",
+  "AccountId": "123456789012"
+}

core/auth/src/test/resources/resources/wiremock/successResponseWithUnexpectedFields.json

@@ -0,0 +1,11 @@
+{
+  "AccessKeyId": "ACCESS_KEY_ID",
+  "SecretAccessKey": "SECRET_ACCESS_KEY",
+  "Token": "TOKEN_TOKEN_TOKEN",
+  "Expiration": "3000-05-03T04:55:54Z",
+  "AccountId": "123456789012",
+  "UnexpectedElement1": {
+    "Name": "ignore-me-1"
+  },
+  "AnotherUnexpected": "should-be-ignored"
+}