Skip to content

Commit 2b6e15d

Browse files
ChrisHegartyChrisHegarty
authored
Merge branch 'main' into lucene_snapshot_10_1
2 parents 1e77a8a + 34ec706 commit 2b6e15d

File tree

3 files changed

+99
-21
lines changed

3 files changed

+99
-21
lines changed

‎test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsHttpHandler.java‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,11 @@
1414
import org.elasticsearch.ExceptionsHelper;
1515
import org.elasticsearch.common.Strings;
1616
import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
17+
import org.elasticsearch.core.Nullable;
1718
import org.elasticsearch.core.SuppressForbidden;
19+
import org.elasticsearch.core.TimeValue;
1820
import org.elasticsearch.rest.RestStatus;
21+
import org.elasticsearch.xcontent.ToXContent;
1922

2023
import java.io.IOException;
2124
import java.nio.charset.StandardCharsets;
@@ -45,21 +48,38 @@ public class Ec2ImdsHttpHandler implements HttpHandler {
4548

4649
private final BiConsumer<String, String> newCredentialsConsumer;
4750
private final Map<String, String> instanceAddresses;
48-
private final Set<String> validCredentialsEndpoints = ConcurrentCollections.newConcurrentSet();
51+
private final Set<String> validCredentialsEndpoints;
52+
private final boolean dynamicProfileNames;
4953
private final Supplier<String> availabilityZoneSupplier;
54+
@Nullable // if instance identity document not available
55+
private final ToXContent instanceIdentityDocument;
5056

5157
public Ec2ImdsHttpHandler(
5258
Ec2ImdsVersion ec2ImdsVersion,
5359
BiConsumer<String, String> newCredentialsConsumer,
5460
Collection<String> alternativeCredentialsEndpoints,
5561
Supplier<String> availabilityZoneSupplier,
62+
@Nullable ToXContent instanceIdentityDocument,
5663
Map<String, String> instanceAddresses
5764
) {
5865
this.ec2ImdsVersion = Objects.requireNonNull(ec2ImdsVersion);
5966
this.newCredentialsConsumer = Objects.requireNonNull(newCredentialsConsumer);
6067
this.instanceAddresses = instanceAddresses;
61-
this.validCredentialsEndpoints.addAll(alternativeCredentialsEndpoints);
68+
69+
if (alternativeCredentialsEndpoints.isEmpty()) {
70+
dynamicProfileNames = true;
71+
validCredentialsEndpoints = ConcurrentCollections.newConcurrentSet();
72+
} else if (ec2ImdsVersion == Ec2ImdsVersion.V2) {
73+
throw new IllegalArgumentException(
74+
Strings.format("alternative credentials endpoints %s requires IMDSv1", alternativeCredentialsEndpoints)
75+
);
76+
} else {
77+
dynamicProfileNames = false;
78+
validCredentialsEndpoints = Set.copyOf(alternativeCredentialsEndpoints);
79+
}
80+
6281
this.availabilityZoneSupplier = availabilityZoneSupplier;
82+
this.instanceIdentityDocument = instanceIdentityDocument;
6383
}
6484

6585
@Override
@@ -78,6 +98,8 @@ public void handle(final HttpExchange exchange) throws IOException {
7898
validImdsTokens.add(token);
7999
final var responseBody = token.getBytes(StandardCharsets.UTF_8);
80100
exchange.getResponseHeaders().add("Content-Type", "text/plain");
101+
exchange.getResponseHeaders()
102+
.add("x-aws-ec2-metadata-token-ttl-seconds", Long.toString(TimeValue.timeValueDays(1).seconds()));
81103
exchange.sendResponseHeaders(RestStatus.OK.getStatus(), responseBody.length);
82104
exchange.getResponseBody().write(responseBody);
83105
}
@@ -98,7 +120,7 @@ public void handle(final HttpExchange exchange) throws IOException {
98120
}
99121

100122
if ("GET".equals(requestMethod)) {
101-
if (path.equals(IMDS_SECURITY_CREDENTIALS_PATH)) {
123+
if (path.equals(IMDS_SECURITY_CREDENTIALS_PATH) && dynamicProfileNames) {
102124
final var profileName = randomIdentifier();
103125
validCredentialsEndpoints.add(IMDS_SECURITY_CREDENTIALS_PATH + profileName);
104126
sendStringResponse(exchange, profileName);
@@ -107,6 +129,9 @@ public void handle(final HttpExchange exchange) throws IOException {
107129
final var availabilityZone = availabilityZoneSupplier.get();
108130
sendStringResponse(exchange, availabilityZone);
109131
return;
132+
} else if (instanceIdentityDocument != null && path.equals("/latest/dynamic/instance-identity/document")) {
133+
sendStringResponse(exchange, Strings.toString(instanceIdentityDocument));
134+
return;
110135
} else if (validCredentialsEndpoints.contains(path)) {
111136
final String accessKey = randomIdentifier();
112137
final String sessionToken = randomIdentifier();

‎test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsServiceBuilder.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
package fixture.aws.imds;
1111

1212
import org.elasticsearch.test.ESTestCase;
13+
import org.elasticsearch.xcontent.ToXContent;
1314

1415
import java.util.Collection;
1516
import java.util.HashMap;
@@ -24,6 +25,7 @@ public class Ec2ImdsServiceBuilder {
2425
private BiConsumer<String, String> newCredentialsConsumer = Ec2ImdsServiceBuilder::rejectNewCredentials;
2526
private Collection<String> alternativeCredentialsEndpoints = Set.of();
2627
private Supplier<String> availabilityZoneSupplier = Ec2ImdsServiceBuilder::rejectAvailabilityZone;
28+
private ToXContent instanceIdentityDocument = null;
2729
private final Map<String, String> instanceAddresses = new HashMap<>();
2830

2931
public Ec2ImdsServiceBuilder(Ec2ImdsVersion ec2ImdsVersion) {
@@ -64,8 +66,13 @@ public Ec2ImdsHttpHandler buildHandler() {
6466
newCredentialsConsumer,
6567
alternativeCredentialsEndpoints,
6668
availabilityZoneSupplier,
69+
instanceIdentityDocument,
6770
Map.copyOf(instanceAddresses)
6871
);
6972
}
7073

74+
public Ec2ImdsServiceBuilder instanceIdentityDocument(ToXContent instanceIdentityDocument) {
75+
this.instanceIdentityDocument = instanceIdentityDocument;
76+
return this;
77+
}
7178
}

‎test/fixtures/ec2-imds-fixture/src/test/java/fixture/aws/imds/Ec2ImdsHttpHandlerTests.java‎

Lines changed: 64 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -52,16 +52,13 @@ public void testImdsV1() throws IOException {
5252
assertTrue(Strings.hasText(profileName));
5353

5454
final var credentialsResponse = handleRequest(handler, "GET", SECURITY_CREDENTIALS_URI + profileName);
55-
assertEquals(RestStatus.OK, credentialsResponse.status());
5655

5756
assertThat(generatedCredentials, aMapWithSize(1));
58-
final var accessKey = generatedCredentials.keySet().iterator().next();
59-
final var sessionToken = generatedCredentials.values().iterator().next();
60-
61-
final var responseMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), credentialsResponse.body().streamInput(), false);
62-
assertEquals(Set.of("AccessKeyId", "Expiration", "RoleArn", "SecretAccessKey", "Token"), responseMap.keySet());
63-
assertEquals(accessKey, responseMap.get("AccessKeyId"));
64-
assertEquals(sessionToken, responseMap.get("Token"));
57+
assertValidCredentialsResponse(
58+
credentialsResponse,
59+
generatedCredentials.keySet().iterator().next(),
60+
generatedCredentials.values().iterator().next()
61+
);
6562
}
6663

6764
public void testImdsV2Disabled() {
@@ -78,6 +75,7 @@ public void testImdsV2() throws IOException {
7875

7976
final var tokenResponse = handleRequest(handler, "PUT", "/latest/api/token");
8077
assertEquals(RestStatus.OK, tokenResponse.status());
78+
assertEquals(List.of("86400" /* seconds in a day */), tokenResponse.responseHeaders().get("x-aws-ec2-metadata-token-ttl-seconds"));
8179
final var token = tokenResponse.body().utf8ToString();
8280

8381
final var roleResponse = checkImdsV2GetRequest(handler, SECURITY_CREDENTIALS_URI, token);
@@ -86,16 +84,13 @@ public void testImdsV2() throws IOException {
8684
assertTrue(Strings.hasText(profileName));
8785

8886
final var credentialsResponse = checkImdsV2GetRequest(handler, SECURITY_CREDENTIALS_URI + profileName, token);
89-
assertEquals(RestStatus.OK, credentialsResponse.status());
9087

9188
assertThat(generatedCredentials, aMapWithSize(1));
92-
final var accessKey = generatedCredentials.keySet().iterator().next();
93-
final var sessionToken = generatedCredentials.values().iterator().next();
94-
95-
final var responseMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), credentialsResponse.body().streamInput(), false);
96-
assertEquals(Set.of("AccessKeyId", "Expiration", "RoleArn", "SecretAccessKey", "Token"), responseMap.keySet());
97-
assertEquals(accessKey, responseMap.get("AccessKeyId"));
98-
assertEquals(sessionToken, responseMap.get("Token"));
89+
assertValidCredentialsResponse(
90+
credentialsResponse,
91+
generatedCredentials.keySet().iterator().next(),
92+
generatedCredentials.values().iterator().next()
93+
);
9994
}
10095

10196
public void testAvailabilityZone() {
@@ -113,7 +108,54 @@ public void testAvailabilityZone() {
113108
assertEquals(generatedAvailabilityZones, Set.of(availabilityZone));
114109
}
115110

116-
private record TestHttpResponse(RestStatus status, BytesReference body) {}
111+
public void testAlternativeCredentialsEndpoint() throws IOException {
112+
expectThrows(
113+
IllegalArgumentException.class,
114+
new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V2).alternativeCredentialsEndpoints(Set.of("/should-not-work"))::buildHandler
115+
);
116+
117+
final var alternativePaths = randomList(1, 5, () -> "/" + randomIdentifier());
118+
final Map<String, String> generatedCredentials = new HashMap<>();
119+
120+
final var handler = new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).alternativeCredentialsEndpoints(alternativePaths)
121+
.newCredentialsConsumer(generatedCredentials::put)
122+
.buildHandler();
123+
124+
final var credentialsResponse = handleRequest(handler, "GET", randomFrom(alternativePaths));
125+
126+
assertThat(generatedCredentials, aMapWithSize(1));
127+
assertValidCredentialsResponse(
128+
credentialsResponse,
129+
generatedCredentials.keySet().iterator().next(),
130+
generatedCredentials.values().iterator().next()
131+
);
132+
}
133+
134+
private static void assertValidCredentialsResponse(TestHttpResponse credentialsResponse, String accessKey, String sessionToken)
135+
throws IOException {
136+
assertEquals(RestStatus.OK, credentialsResponse.status());
137+
final var responseMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), credentialsResponse.body().streamInput(), false);
138+
assertEquals(Set.of("AccessKeyId", "Expiration", "RoleArn", "SecretAccessKey", "Token"), responseMap.keySet());
139+
assertEquals(accessKey, responseMap.get("AccessKeyId"));
140+
assertEquals(sessionToken, responseMap.get("Token"));
141+
}
142+
143+
public void testInstanceIdentityDocument() {
144+
final Set<String> generatedRegions = new HashSet<>();
145+
final var handler = new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).instanceIdentityDocument((builder, params) -> {
146+
final var newRegion = randomIdentifier();
147+
generatedRegions.add(newRegion);
148+
return builder.field("region", newRegion);
149+
}).buildHandler();
150+
151+
final var instanceIdentityResponse = handleRequest(handler, "GET", "/latest/dynamic/instance-identity/document");
152+
assertEquals(RestStatus.OK, instanceIdentityResponse.status());
153+
final var instanceIdentityString = instanceIdentityResponse.body().utf8ToString();
154+
155+
assertEquals(Strings.format("{\"region\":\"%s\"}", generatedRegions.iterator().next()), instanceIdentityString);
156+
}
157+
158+
private record TestHttpResponse(RestStatus status, Headers responseHeaders, BytesReference body) {}
117159

118160
private static TestHttpResponse checkImdsV2GetRequest(Ec2ImdsHttpHandler handler, String uri, String token) {
119161
final var unauthorizedResponse = handleRequest(handler, "GET", uri, null);
@@ -145,7 +187,11 @@ private static TestHttpResponse handleRequest(Ec2ImdsHttpHandler handler, String
145187
fail(e);
146188
}
147189
assertNotEquals(0, httpExchange.getResponseCode());
148-
return new TestHttpResponse(RestStatus.fromCode(httpExchange.getResponseCode()), httpExchange.getResponseBodyContents());
190+
return new TestHttpResponse(
191+
RestStatus.fromCode(httpExchange.getResponseCode()),
192+
httpExchange.getResponseHeaders(),
193+
httpExchange.getResponseBodyContents()
194+
);
149195
}
150196

151197
private static class TestHttpExchange extends HttpExchange {

0 commit comments

Comments
 (0)