Skip to content

Commit 505479b

Browse files
S-Saranya1S-Saranya1
committed
Use Ec2MetadataClient in defaults mode
1 parent cd49f1f commit 505479b

File tree

6 files changed

+759
-3
lines changed

6 files changed

+759
-3
lines changed

‎core/aws-core/pom.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,12 @@
206206
<artifactId>rxjava</artifactId>
207207
<scope>test</scope>
208208
</dependency>
209+
<dependency>
210+
<groupId>software.amazon.awssdk</groupId>
211+
<artifactId>imds</artifactId>
212+
<version>2.31.73-SNAPSHOT</version>
213+
<scope>compile</scope>
214+
</dependency>
209215
</dependencies>
210216
<build>
211217
<plugins>

‎core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/defaultsmode/AutoDefaultsModeDiscovery.java

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@
1919
import software.amazon.awssdk.annotations.SdkProtectedApi;
2020
import software.amazon.awssdk.awscore.defaultsmode.DefaultsMode;
2121
import software.amazon.awssdk.core.SdkSystemSetting;
22+
import software.amazon.awssdk.imds.Ec2MetadataClient;
23+
import software.amazon.awssdk.imds.Ec2MetadataRetryPolicy;
24+
import software.amazon.awssdk.imds.internal.Ec2MetadataSharedClient;
2225
import software.amazon.awssdk.regions.Region;
23-
import software.amazon.awssdk.regions.internal.util.EC2MetadataUtils;
2426
import software.amazon.awssdk.utils.JavaSystemSetting;
2527
import software.amazon.awssdk.utils.OptionalUtils;
2628
import software.amazon.awssdk.utils.SystemSetting;
@@ -81,9 +83,17 @@ private static DefaultsMode compareRegion(String region, Region clientRegion) {
8183
}
8284

8385
private static Optional<String> queryImdsV2() {
86+
87+
if (SdkSystemSetting.AWS_EC2_METADATA_DISABLED.getBooleanValueOrThrow()) {
88+
return Optional.empty();
89+
}
90+
8491
try {
85-
String ec2InstanceRegion = EC2MetadataUtils.fetchData(EC2_METADATA_REGION_PATH, false, 1);
86-
// ec2InstanceRegion could be null
92+
Ec2MetadataClient client = Ec2MetadataSharedClient.builder()
93+
.retryPolicy(Ec2MetadataRetryPolicy.none())
94+
.build();
95+
96+
String ec2InstanceRegion = client.get(EC2_METADATA_REGION_PATH).asString();
8797
return Optional.ofNullable(ec2InstanceRegion);
8898
} catch (Exception exception) {
8999
return Optional.empty();

‎core/aws-core/src/test/java/software/amazon/awssdk/awscore/internal/defaultsmode/AutoDefaultsModeDiscoveryEc2MetadataClientTest.java

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.awscore.internal.defaultsmode;
17+
18+
import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
19+
import static com.github.tomakehurst.wiremock.client.WireMock.get;
20+
import static com.github.tomakehurst.wiremock.client.WireMock.put;
21+
import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
22+
import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
23+
import static com.github.tomakehurst.wiremock.client.WireMock.putRequestedFor;
24+
import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
25+
import static com.github.tomakehurst.wiremock.client.WireMock.verify;
26+
import static org.assertj.core.api.Assertions.assertThat;
27+
28+
import com.github.tomakehurst.wiremock.junit.WireMockRule;
29+
import java.lang.reflect.Field;
30+
import org.junit.After;
31+
import org.junit.Before;
32+
import org.junit.Rule;
33+
import org.junit.Test;
34+
import software.amazon.awssdk.awscore.defaultsmode.DefaultsMode;
35+
import software.amazon.awssdk.core.SdkSystemSetting;
36+
import software.amazon.awssdk.http.SdkHttpClient;
37+
import software.amazon.awssdk.imds.internal.Ec2MetadataSharedClient;
38+
import software.amazon.awssdk.regions.Region;
39+
import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
40+
import software.amazon.awssdk.utils.Lazy;
41+
42+
/**
43+
* Tests specifically for AutoDefaultsModeDiscovery's migration to use Ec2MetadataClient.
44+
* These tests verify that the migration from EC2MetadataUtils to Ec2MetadataClient works correctly.
45+
*/
46+
public class AutoDefaultsModeDiscoveryEc2MetadataClientTest {
47+
private static final EnvironmentVariableHelper ENVIRONMENT_VARIABLE_HELPER = new EnvironmentVariableHelper();
48+
49+
@Rule
50+
public WireMockRule wireMock = new WireMockRule(0);
51+
52+
@Before
53+
public void setup() {
54+
System.setProperty(SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT.property(),
55+
"http://localhost:" + wireMock.port());
56+
ENVIRONMENT_VARIABLE_HELPER.remove(SdkSystemSetting.AWS_EC2_METADATA_DISABLED.environmentVariable());
57+
}
58+
59+
@After
60+
public void cleanup() {
61+
wireMock.resetAll();
62+
ENVIRONMENT_VARIABLE_HELPER.reset();
63+
System.clearProperty(SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT.property());
64+
}
65+
66+
@Test
67+
public void autoDefaultsModeDiscovery_shouldUseSharedHttpClient() throws Exception {
68+
// Stub successful IMDS responses
69+
stubFor(put("/latest/api/token")
70+
.willReturn(aResponse().withStatus(200).withBody("test-token")));
71+
stubFor(get("/latest/meta-data/placement/region")
72+
.willReturn(aResponse().withStatus(200).withBody("us-east-1")));
73+
74+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
75+
76+
AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
77+
DefaultsMode result = discovery.discover(Region.US_EAST_1);
78+
79+
// Should return IN_REGION since client region (us-east-1) matches IMDS region (us-east-1)
80+
assertThat(result).isEqualTo(DefaultsMode.IN_REGION);
81+
82+
// Verify that the shared HTTP client was used
83+
Field sharedClientField = Ec2MetadataSharedClient.class.getDeclaredField("SHARED_HTTP_CLIENT");
84+
sharedClientField.setAccessible(true);
85+
Lazy<SdkHttpClient> sharedHttpClient = (Lazy<SdkHttpClient>) sharedClientField.get(null);
86+
87+
// Verify the shared HTTP client was initialized
88+
assertThat(sharedHttpClient.hasValue()).isTrue();
89+
90+
// Verify IMDS requests were made
91+
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
92+
verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
93+
}
94+
95+
@Test
96+
public void multipleDiscoveryInstances_shouldShareSameHttpClient() throws Exception {
97+
stubFor(put("/latest/api/token")
98+
.willReturn(aResponse().withStatus(200).withBody("test-token")));
99+
stubFor(get("/latest/meta-data/placement/region")
100+
.willReturn(aResponse().withStatus(200).withBody("us-west-2")));
101+
102+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
103+
104+
// Create multiple discovery instances
105+
AutoDefaultsModeDiscovery discovery1 = new AutoDefaultsModeDiscovery();
106+
AutoDefaultsModeDiscovery discovery2 = new AutoDefaultsModeDiscovery();
107+
108+
// Both should use the same shared HTTP client
109+
DefaultsMode result1 = discovery1.discover(Region.US_EAST_1);
110+
DefaultsMode result2 = discovery2.discover(Region.US_EAST_1);
111+
112+
// Both should return CROSS_REGION (client: us-east-1, IMDS: us-west-2)
113+
assertThat(result1).isEqualTo(DefaultsMode.CROSS_REGION);
114+
assertThat(result2).isEqualTo(DefaultsMode.CROSS_REGION);
115+
116+
// Verify shared HTTP client was used
117+
Field sharedClientField = Ec2MetadataSharedClient.class.getDeclaredField("SHARED_HTTP_CLIENT");
118+
sharedClientField.setAccessible(true);
119+
Lazy<SdkHttpClient> sharedHttpClient = (Lazy<SdkHttpClient>) sharedClientField.get(null);
120+
121+
assertThat(sharedHttpClient.hasValue()).isTrue();
122+
123+
// Verify IMDS requests were made
124+
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
125+
verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
126+
}
127+
128+
@Test
129+
public void awsEc2MetadataDisabled_shouldSkipImdsAndUseStandardMode() {
130+
// Disable IMDS
131+
ENVIRONMENT_VARIABLE_HELPER.set(SdkSystemSetting.AWS_EC2_METADATA_DISABLED.environmentVariable(), "true");
132+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
133+
134+
135+
AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
136+
DefaultsMode result = discovery.discover(Region.US_EAST_1);
137+
138+
// Should return STANDARD mode without making IMDS calls
139+
assertThat(result).isEqualTo(DefaultsMode.STANDARD);
140+
141+
// Verify no IMDS requests were made
142+
wireMock.verify(0, putRequestedFor(urlEqualTo("/latest/api/token")));
143+
wireMock.verify(0, getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
144+
}
145+
146+
@Test
147+
public void imdsFailure_shouldFallbackToStandardMode() {
148+
// Stub IMDS to fail
149+
stubFor(put("/latest/api/token")
150+
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));
151+
stubFor(get("/latest/meta-data/placement/region")
152+
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));
153+
154+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
155+
156+
AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
157+
DefaultsMode result = discovery.discover(Region.US_EAST_1);
158+
159+
// Should fall back to STANDARD mode when IMDS fails
160+
assertThat(result).isEqualTo(DefaultsMode.STANDARD);
161+
162+
// Verify IMDS requests were attempted
163+
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
164+
}
165+
166+
@Test
167+
public void noRetryPolicy_shouldBeUsedByDefault() {
168+
// Stub token to succeed but region to fail with retryable error
169+
stubFor(put("/latest/api/token")
170+
.willReturn(aResponse().withStatus(200).withBody("test-token")));
171+
stubFor(get("/latest/meta-data/placement/region")
172+
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));
173+
174+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
175+
176+
AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
177+
DefaultsMode result = discovery.discover(Region.US_EAST_1);
178+
179+
// Should fail immediately without retries and fallback to STANDARD
180+
assertThat(result).isEqualTo(DefaultsMode.STANDARD);
181+
182+
// Verify requests were made but no retries
183+
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
184+
verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
185+
}
186+
187+
188+
@Test
189+
public void imdsV1Fallback_shouldWorkWhenTokenFails() {
190+
// Stub token request to fail
191+
stubFor(put("/latest/api/token")
192+
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));
193+
194+
// Stub successful IMDSv1 request
195+
stubFor(get("/latest/meta-data/placement/region")
196+
.willReturn(aResponse().withStatus(200).withBody("us-east-1")));
197+
198+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
199+
200+
AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
201+
DefaultsMode result = discovery.discover(Region.US_EAST_1);
202+
203+
// Should fall back to IMDSv1 and succeed, returning IN_REGION
204+
assertThat(result).isEqualTo(DefaultsMode.IN_REGION);
205+
206+
// Verify both token request (failed) and region request (succeeded) were made
207+
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
208+
verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/region")));
209+
}
210+
211+
@Test
212+
public void imdsV1Fallback_shouldNotWorkWhenV1Disabled() {
213+
// Disable IMDSv1 fallback
214+
ENVIRONMENT_VARIABLE_HELPER.set(SdkSystemSetting.AWS_EC2_METADATA_V1_DISABLED.environmentVariable(), "true");
215+
216+
// Stub token request to fail
217+
stubFor(put("/latest/api/token")
218+
.willReturn(aResponse().withStatus(500).withBody("Internal Server Error")));
219+
220+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
221+
222+
AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
223+
DefaultsMode result = discovery.discover(Region.US_EAST_1);
224+
225+
// Should fail without fallback to IMDSv1 and return STANDARD
226+
assertThat(result).isEqualTo(DefaultsMode.STANDARD);
227+
228+
// Verify only token request was made
229+
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
230+
}
231+
232+
@Test
233+
public void tokenRequest400Error_shouldNotFallbackToV1() {
234+
// Stub token request to fail with 400
235+
stubFor(put("/latest/api/token")
236+
.willReturn(aResponse().withStatus(400).withBody("Bad Request")));
237+
238+
ENVIRONMENT_VARIABLE_HELPER.set("AWS_DEFAULT_REGION", "us-west-2");
239+
240+
AutoDefaultsModeDiscovery discovery = new AutoDefaultsModeDiscovery();
241+
DefaultsMode result = discovery.discover(Region.US_EAST_1);
242+
243+
// Should fail without attempting IMDSv1 fallback and return STANDARD
244+
assertThat(result).isEqualTo(DefaultsMode.STANDARD);
245+
246+
// Verify only token request was made (no region request)
247+
verify(putRequestedFor(urlEqualTo("/latest/api/token")));
248+
}
249+
}

‎core/imds/src/main/java/software/amazon/awssdk/imds/internal/BaseEc2MetadataClient.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,10 @@ protected BaseEc2MetadataClient(DefaultEc2MetadataAsyncClient.Ec2MetadataAsyncBu
6060
this(builder.getRetryPolicy(), builder.getTokenTtl(), builder.getEndpoint(), builder.getEndpointMode());
6161
}
6262

63+
protected BaseEc2MetadataClient(DefaultEc2MetadataClientWithFallback.Ec2MetadataBuilder builder) {
64+
this(builder.getRetryPolicy(), builder.getTokenTtl(), builder.getEndpoint(), builder.getEndpointMode());
65+
}
66+
6367
private URI getEndpoint(URI builderEndpoint, EndpointMode builderEndpointMode) {
6468
Validate.mutuallyExclusive("Only one of 'endpoint' or 'endpointMode' must be specified, but not both",
6569
builderEndpoint, builderEndpointMode);

0 commit comments

Comments
 (0)