Skip to content

Use Ec2MetadataClient in defaults mode #6301

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 11 commits into
base: feature/master/use-imds-client
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .changes/next-release/feature-AWSSDKforJavav2-8c69562.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
{
"type": "feature",
"category": "AWS SDK for Java v2",
"contributor": "",
"description": "Updated AutoDefaultsModeDiscovery from using EC2MetadataUtils to Ec2MetadataClient"
}
6 changes: 6 additions & 0 deletions core/aws-core/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,12 @@
<artifactId>rxjava</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>imds</artifactId>
<version>2.31.73-SNAPSHOT</version>
<scope>compile</scope>
</dependency>
</dependencies>
<build>
<plugins>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,10 @@
import software.amazon.awssdk.annotations.SdkProtectedApi;
import software.amazon.awssdk.awscore.defaultsmode.DefaultsMode;
import software.amazon.awssdk.core.SdkSystemSetting;
import software.amazon.awssdk.imds.Ec2MetadataClient;
import software.amazon.awssdk.imds.Ec2MetadataRetryPolicy;
import software.amazon.awssdk.imds.internal.Ec2MetadataSharedClient;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.regions.internal.util.EC2MetadataUtils;
import software.amazon.awssdk.utils.JavaSystemSetting;
import software.amazon.awssdk.utils.OptionalUtils;
import software.amazon.awssdk.utils.SystemSetting;
Expand Down Expand Up @@ -81,9 +83,17 @@ private static DefaultsMode compareRegion(String region, Region clientRegion) {
}

private static Optional<String> queryImdsV2() {

if (SdkSystemSetting.AWS_EC2_METADATA_DISABLED.getBooleanValueOrThrow()) {
return Optional.empty();
}

try {
String ec2InstanceRegion = EC2MetadataUtils.fetchData(EC2_METADATA_REGION_PATH, false, 1);
// ec2InstanceRegion could be null
Ec2MetadataClient client = Ec2MetadataSharedClient.builder()
.retryPolicy(Ec2MetadataRetryPolicy.none())
.build();

String ec2InstanceRegion = client.get(EC2_METADATA_REGION_PATH).asString();
return Optional.ofNullable(ec2InstanceRegion);
} catch (Exception exception) {
return Optional.empty();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,247 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/apache2.0
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/

package software.amazon.awssdk.awscore.internal.defaultsmode;

import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
import static com.github.tomakehurst.wiremock.client.WireMock.get;
import static com.github.tomakehurst.wiremock.client.WireMock.put;
import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
import static com.github.tomakehurst.wiremock.client.WireMock.putRequestedFor;
import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
import static com.github.tomakehurst.wiremock.client.WireMock.verify;
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
import static org.assertj.core.api.Assertions.assertThat;

import com.github.tomakehurst.wiremock.junit.WireMockRule;
import java.lang.reflect.Field;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import software.amazon.awssdk.awscore.defaultsmode.DefaultsMode;
import software.amazon.awssdk.core.SdkSystemSetting;
import software.amazon.awssdk.http.SdkHttpClient;
import software.amazon.awssdk.imds.internal.Ec2MetadataSharedClient;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
import software.amazon.awssdk.utils.Lazy;

/**
* Tests specifically for AutoDefaultsModeDiscovery's migration to use Ec2MetadataClient.
* These tests verify that the migration from EC2MetadataUtils to Ec2MetadataClient works correctly.
*/
public class AutoDefaultsModeDiscoveryEc2MetadataClientTest {
private static final EnvironmentVariableHelper ENVIRONMENT_VARIABLE_HELPER = new EnvironmentVariableHelper();

@Rule
public WireMockRule wireMock = new WireMockRule(wireMockConfig()
.port(0)
.httpsPort(-1));

@Before
public void setup() {
System.setProperty(SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT.property(),
"http://localhost:" + wireMock.port());

clearEnvironmentVariable("AWS_EXECUTION_ENV");
clearEnvironmentVariable("AWS_REGION");
clearEnvironmentVariable("AWS_DEFAULT_REGION");
}

@After
public void cleanup() {
wireMock.resetAll();
ENVIRONMENT_VARIABLE_HELPER.reset();
System.clearProperty(SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT.property());
}

// Clear an environment variable by setting it to null.
private void clearEnvironmentVariable(String name) {
try {
ENVIRONMENT_VARIABLE_HELPER.set(name, null);
} catch (Exception e) {
// Ignore
}
}

@Test
public void autoDefaultsModeDiscovery_shouldUseSharedHttpClient() throws Exception {
// Stub successful IMDS responses
stubFor(put("/latest/api/token")
.willReturn(aResponse().withStatus(200).withBody("test-token")));
stubFor(get("/latest/meta-data/placement/region")
.willReturn(aResponse().withStatus(200).withBody("us-east-1")));

AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
DefaultsMode result = discovery.discover(Region.US_EAST_1);

// Should return IN_REGION since client region matches IMDS region
assertThat(result).isEqualTo(DefaultsMode.IN_REGION);

// Verify that the shared HTTP client was used
Field sharedClientField = Ec2MetadataSharedClient.class.getDeclaredField("SHARED_HTTP_CLIENT");
sharedClientField.setAccessible(true);
Lazy<SdkHttpClient> sharedHttpClient = (Lazy<SdkHttpClient>) sharedClientField.get(null);

// Verify the shared HTTP client was initialized
assertThat(sharedHttpClient.hasValue()).isTrue();

// Verify IMDS requests were made
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
}

@Test
public void multipleDiscoveryInstances_shouldShareSameHttpClient() throws Exception {
stubFor(put("/latest/api/token")
.willReturn(aResponse().withStatus(200).withBody("test-token")));
stubFor(get("/latest/meta-data/placement/region")
.willReturn(aResponse().withStatus(200).withBody("us-west-2")));

// Create multiple discovery instances
AutoDefaultsModeDiscovery discovery1 = new AutoDefaultsModeDiscovery();
AutoDefaultsModeDiscovery discovery2 = new AutoDefaultsModeDiscovery();

// Both should use the same shared HTTP client
DefaultsMode result1 = discovery1.discover(Region.US_EAST_1);
DefaultsMode result2 = discovery2.discover(Region.US_EAST_1);

// Both should return CROSS_REGION
assertThat(result1).isEqualTo(DefaultsMode.CROSS_REGION);
assertThat(result2).isEqualTo(DefaultsMode.CROSS_REGION);

// Verify shared HTTP client was used
Field sharedClientField = Ec2MetadataSharedClient.class.getDeclaredField("SHARED_HTTP_CLIENT");
sharedClientField.setAccessible(true);
Lazy<SdkHttpClient> sharedHttpClient = (Lazy<SdkHttpClient>) sharedClientField.get(null);

assertThat(sharedHttpClient.hasValue()).isTrue();

// Verify IMDS requests were made
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
}

@Test
public void awsEc2MetadataDisabled_shouldSkipImdsAndUseStandardMode() {
// Disable IMDS
ENVIRONMENT_VARIABLE_HELPER.set(SdkSystemSetting.AWS_EC2_METADATA_DISABLED.environmentVariable(), "true");

AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
DefaultsMode result = discovery.discover(Region.US_EAST_1);

// Should return STANDARD mode without making IMDS calls
assertThat(result).isEqualTo(DefaultsMode.STANDARD);

// Verify no IMDS requests were made
verify(0, putRequestedFor(urlEqualTo("/latest/api/token")));
verify(0, getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
}

@Test
public void imdsFailure_shouldFallbackToStandardMode() {
// Stub IMDS to fail
stubFor(put("/latest/api/token")
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));
stubFor(get("/latest/meta-data/placement/region")
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));

AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
DefaultsMode result = discovery.discover(Region.US_EAST_1);

// Should fall back to STANDARD mode when IMDS fails
assertThat(result).isEqualTo(DefaultsMode.STANDARD);

// Verify IMDS requests were attempted
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
}

@Test
public void noRetryPolicy_shouldBeUsedByDefault() {
// Stub token to succeed but region to fail with retryable error
stubFor(put("/latest/api/token")
.willReturn(aResponse().withStatus(200).withBody("test-token")));
stubFor(get("/latest/meta-data/placement/region")
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));

AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
DefaultsMode result = discovery.discover(Region.US_EAST_1);

// Should fail immediately without retries and fallback to STANDARD
assertThat(result).isEqualTo(DefaultsMode.STANDARD);

// Verify requests were made once (no retries)
verify(1, putRequestedFor(urlEqualTo("/latest/api/token")));
verify(1, getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
}

@Test
public void imdsV1Fallback_shouldWorkWhenTokenFails() {
// Stub token request to fail
stubFor(put("/latest/api/token")
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));

// Stub successful IMDSv1 request
stubFor(get("/latest/meta-data/placement/region")
.willReturn(aResponse().withStatus(200).withBody("us-east-1")));

AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
DefaultsMode result = discovery.discover(Region.US_EAST_1);

// Should fall back to IMDSv1 and return IN_REGION
assertThat(result).isEqualTo(DefaultsMode.IN_REGION);

// Verify both token request and region request were made
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
}

@Test
public void imdsV1Fallback_shouldNotWorkWhenV1Disabled() {
// Disable IMDSv1 fallback
ENVIRONMENT_VARIABLE_HELPER.set(SdkSystemSetting.AWS_EC2_METADATA_V1_DISABLED.environmentVariable(), "true");

// Stub token request to fail
stubFor(put("/latest/api/token")
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));

AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
DefaultsMode result = discovery.discover(Region.US_EAST_1);

// Should fail without fallback to IMDSv1 and return STANDARD
assertThat(result).isEqualTo(DefaultsMode.STANDARD);

// Verify only token request was made
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
}

@Test
public void tokenRequest400Error_shouldNotFallbackToV1() {
// Stub token request to fail with 400
stubFor(put("/latest/api/token")
.willReturn(aResponse().withStatus(400).withBody("Bad Request")));

AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
DefaultsMode result = discovery.discover(Region.US_EAST_1);

// Should fail without attempting IMDSv1 fallback and return STANDARD
assertThat(result).isEqualTo(DefaultsMode.STANDARD);

// Verify only token request was made
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,10 @@ protected BaseEc2MetadataClient(DefaultEc2MetadataAsyncClient.Ec2MetadataAsyncBu
this(builder.getRetryPolicy(), builder.getTokenTtl(), builder.getEndpoint(), builder.getEndpointMode());
}

protected BaseEc2MetadataClient(DefaultEc2MetadataClientWithFallback.Ec2MetadataBuilder builder) {
this(builder.getRetryPolicy(), builder.getTokenTtl(), builder.getEndpoint(), builder.getEndpointMode());
}

private URI getEndpoint(URI builderEndpoint, EndpointMode builderEndpointMode) {
Validate.mutuallyExclusive("Only one of 'endpoint' or 'endpointMode' must be specified, but not both",
builderEndpoint, builderEndpointMode);
Expand Down
Loading
Loading