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 or container credentials, if the account ID is returned in the credentials response. Endpoint host names can change for services that support it, such as DynamoDB. To opt out, set `accountIdEndpointMode` to `disabled`. See https://docs.aws.amazon.com/sdkref/latest/guide/feature-account-endpoints.html"
+}

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/ContainerCredentialsProviderTest.java

@@ -50,6 +50,7 @@ public class ContainerCredentialsProviderTest {
     private static final String SECRET_ACCESS_KEY = "SECRET_ACCESS_KEY";
     private static final String TOKEN = "TOKEN_TOKEN_TOKEN";
     private static final String EXPIRATION_TIME = "3000-05-03T04:55:54Z";
+    private static final String ACCOUNT_ID = "123456789012";
     private ContainerCredentialsProvider credentialsProvider;
     private static EnvironmentVariableHelper helper = new EnvironmentVariableHelper();
 
@@ -100,6 +101,34 @@ public void testGetCredentialsReturnsValidResponseFromEcsEndpoint() {
         assertThat(credentials.providerName()).isPresent().contains(BusinessMetricFeatureId.CREDENTIALS_HTTP.value());
     }
 
+    /**
+     * Tests that the account ID is attached to the resolved credentials when the endpoint returns it.
+     */
+    @Test
+    public void resolveCredentials_withAccountIdInResponse_returnsCredentialsWithAccountId() {
+        stubFor200Response(getSuccessfulBodyWithAccountId());
+
+        AwsCredentials credentials = credentialsProvider.resolveCredentials();
+
+        assertThat(credentials.accessKeyId()).isEqualTo(ACCESS_KEY_ID);
+        assertThat(credentials.secretAccessKey()).isEqualTo(SECRET_ACCESS_KEY);
+        assertThat(credentials.accountId()).isPresent().hasValue(ACCOUNT_ID);
+    }
+
+    /**
+     * Tests that credentials still resolve without an account ID when the endpoint does not return it.
+     */
+    @Test
+    public void resolveCredentials_withoutAccountIdInResponse_returnsCredentialsWithoutAccountId() {
+        stubForSuccessResponse();
+
+        AwsCredentials credentials = credentialsProvider.resolveCredentials();
+
+        assertThat(credentials.accessKeyId()).isEqualTo(ACCESS_KEY_ID);
+        assertThat(credentials.secretAccessKey()).isEqualTo(SECRET_ACCESS_KEY);
+        assertThat(credentials.accountId()).isNotPresent();
+    }
+
     /**
      * Tests that the getCredentials won't leak credentials data if the response from ECS is corrupted.
      */
@@ -140,4 +169,12 @@ private String getSuccessfulBody() {
                "\"Token\":\"TOKEN_TOKEN_TOKEN\"," +
                "\"Expiration\":\"3000-05-03T04:55:54Z\"}";
     }
+
+    private String getSuccessfulBodyWithAccountId() {
+        return "{\"AccessKeyId\":\"ACCESS_KEY_ID\"," +
+               "\"SecretAccessKey\":\"SECRET_ACCESS_KEY\"," +
+               "\"Token\":\"TOKEN_TOKEN_TOKEN\"," +
+               "\"Expiration\":\"3000-05-03T04:55:54Z\"," +
+               "\"AccountId\":\"" + ACCOUNT_ID + "\"}";
+    }
 }

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,57 @@ 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.
+     */
+    @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"
+}