Skip to content

Commit 8a61275

Browse files
DaveCTurnerDaveCTurner
authored
Add discovery-ec2 integration test for AZ attr (#118452) (#118541)
Verifies that the plugin sets the `aws_availability_zone` automatically by reading the AZ name from the IMDS at startup.
1 parent 96ec2f6 commit 8a61275

File tree

12 files changed

+265
-41
lines changed

12 files changed

+265
-41
lines changed

‎modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3EcsCredentialsRestIT.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
package org.elasticsearch.repositories.s3;
1111

1212
import fixture.aws.imds.Ec2ImdsHttpFixture;
13+
import fixture.aws.imds.Ec2ImdsServiceBuilder;
1314
import fixture.aws.imds.Ec2ImdsVersion;
1415
import fixture.s3.DynamicS3Credentials;
1516
import fixture.s3.S3HttpFixture;
@@ -37,9 +38,8 @@ public class RepositoryS3EcsCredentialsRestIT extends AbstractRepositoryS3RestTe
3738
private static final DynamicS3Credentials dynamicS3Credentials = new DynamicS3Credentials();
3839

3940
private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
40-
Ec2ImdsVersion.V1,
41-
dynamicS3Credentials::addValidCredentials,
42-
Set.of("/ecs_credentials_endpoint")
41+
new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).newCredentialsConsumer(dynamicS3Credentials::addValidCredentials)
42+
.alternativeCredentialsEndpoints(Set.of("/ecs_credentials_endpoint"))
4343
);
4444

4545
private static final S3HttpFixture s3Fixture = new S3HttpFixture(true, BUCKET, BASE_PATH, dynamicS3Credentials::isAuthorized);

‎modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV1CredentialsRestIT.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
package org.elasticsearch.repositories.s3;
1111

1212
import fixture.aws.imds.Ec2ImdsHttpFixture;
13+
import fixture.aws.imds.Ec2ImdsServiceBuilder;
1314
import fixture.aws.imds.Ec2ImdsVersion;
1415
import fixture.s3.DynamicS3Credentials;
1516
import fixture.s3.S3HttpFixture;
@@ -23,8 +24,6 @@
2324
import org.junit.rules.RuleChain;
2425
import org.junit.rules.TestRule;
2526

26-
import java.util.Set;
27-
2827
@ThreadLeakFilters(filters = { TestContainersThreadFilter.class })
2928
@ThreadLeakScope(ThreadLeakScope.Scope.NONE) // https://github.com/elastic/elasticsearch/issues/102482
3029
public class RepositoryS3ImdsV1CredentialsRestIT extends AbstractRepositoryS3RestTestCase {
@@ -37,9 +36,7 @@ public class RepositoryS3ImdsV1CredentialsRestIT extends AbstractRepositoryS3Res
3736
private static final DynamicS3Credentials dynamicS3Credentials = new DynamicS3Credentials();
3837

3938
private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
40-
Ec2ImdsVersion.V1,
41-
dynamicS3Credentials::addValidCredentials,
42-
Set.of()
39+
new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).newCredentialsConsumer(dynamicS3Credentials::addValidCredentials)
4340
);
4441

4542
private static final S3HttpFixture s3Fixture = new S3HttpFixture(true, BUCKET, BASE_PATH, dynamicS3Credentials::isAuthorized);

‎modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV2CredentialsRestIT.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
package org.elasticsearch.repositories.s3;
1111

1212
import fixture.aws.imds.Ec2ImdsHttpFixture;
13+
import fixture.aws.imds.Ec2ImdsServiceBuilder;
1314
import fixture.aws.imds.Ec2ImdsVersion;
1415
import fixture.s3.DynamicS3Credentials;
1516
import fixture.s3.S3HttpFixture;
@@ -23,8 +24,6 @@
2324
import org.junit.rules.RuleChain;
2425
import org.junit.rules.TestRule;
2526

26-
import java.util.Set;
27-
2827
@ThreadLeakFilters(filters = { TestContainersThreadFilter.class })
2928
@ThreadLeakScope(ThreadLeakScope.Scope.NONE) // https://github.com/elastic/elasticsearch/issues/102482
3029
public class RepositoryS3ImdsV2CredentialsRestIT extends AbstractRepositoryS3RestTestCase {
@@ -37,9 +36,7 @@ public class RepositoryS3ImdsV2CredentialsRestIT extends AbstractRepositoryS3Res
3736
private static final DynamicS3Credentials dynamicS3Credentials = new DynamicS3Credentials();
3837

3938
private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
40-
Ec2ImdsVersion.V2,
41-
dynamicS3Credentials::addValidCredentials,
42-
Set.of()
39+
new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V2).newCredentialsConsumer(dynamicS3Credentials::addValidCredentials)
4340
);
4441

4542
private static final S3HttpFixture s3Fixture = new S3HttpFixture(true, BUCKET, BASE_PATH, dynamicS3Credentials::isAuthorized);

‎plugins/discovery-ec2/build.gradle‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@ dependencies {
2828
api "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}"
2929
api "com.fasterxml.jackson.dataformat:jackson-dataformat-cbor:${versions.jackson}"
3030
api "joda-time:joda-time:2.10.10"
31+
32+
javaRestTestImplementation project(':plugins:discovery-ec2')
33+
javaRestTestImplementation project(':test:fixtures:ec2-imds-fixture')
3134
}
3235

3336
tasks.named("dependencyLicenses").configure {

‎plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV1IT.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.discovery.ec2;
11+
12+
import fixture.aws.imds.Ec2ImdsHttpFixture;
13+
import fixture.aws.imds.Ec2ImdsServiceBuilder;
14+
import fixture.aws.imds.Ec2ImdsVersion;
15+
16+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
17+
import org.junit.ClassRule;
18+
import org.junit.rules.RuleChain;
19+
import org.junit.rules.TestRule;
20+
21+
public class DiscoveryEc2AvailabilityZoneAttributeImdsV1IT extends DiscoveryEc2AvailabilityZoneAttributeTestCase {
22+
private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
23+
new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).availabilityZoneSupplier(
24+
DiscoveryEc2AvailabilityZoneAttributeTestCase::getAvailabilityZone
25+
)
26+
);
27+
28+
public static ElasticsearchCluster cluster = buildCluster(ec2ImdsHttpFixture::getAddress);
29+
30+
@ClassRule
31+
public static TestRule ruleChain = RuleChain.outerRule(ec2ImdsHttpFixture).around(cluster);
32+
33+
@Override
34+
protected String getTestRestCluster() {
35+
return cluster.getHttpAddresses();
36+
}
37+
}

‎plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV2IT.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.discovery.ec2;
11+
12+
import fixture.aws.imds.Ec2ImdsHttpFixture;
13+
import fixture.aws.imds.Ec2ImdsServiceBuilder;
14+
import fixture.aws.imds.Ec2ImdsVersion;
15+
16+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
17+
import org.junit.ClassRule;
18+
import org.junit.rules.RuleChain;
19+
import org.junit.rules.TestRule;
20+
21+
public class DiscoveryEc2AvailabilityZoneAttributeImdsV2IT extends DiscoveryEc2AvailabilityZoneAttributeTestCase {
22+
private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
23+
new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V2).availabilityZoneSupplier(
24+
DiscoveryEc2AvailabilityZoneAttributeTestCase::getAvailabilityZone
25+
)
26+
);
27+
28+
public static ElasticsearchCluster cluster = buildCluster(ec2ImdsHttpFixture::getAddress);
29+
30+
@ClassRule
31+
public static TestRule ruleChain = RuleChain.outerRule(ec2ImdsHttpFixture).around(cluster);
32+
33+
@Override
34+
protected String getTestRestCluster() {
35+
return cluster.getHttpAddresses();
36+
}
37+
}

‎plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeNoImdsIT.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.discovery.ec2;
11+
12+
import org.elasticsearch.client.Request;
13+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
14+
import org.elasticsearch.test.rest.ESRestTestCase;
15+
import org.junit.ClassRule;
16+
17+
import java.io.IOException;
18+
19+
public class DiscoveryEc2AvailabilityZoneAttributeNoImdsIT extends ESRestTestCase {
20+
@ClassRule
21+
public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
22+
.plugin("discovery-ec2")
23+
.setting(AwsEc2Service.AUTO_ATTRIBUTE_SETTING.getKey(), "true")
24+
.build();
25+
26+
@Override
27+
protected String getTestRestCluster() {
28+
return cluster.getHttpAddresses();
29+
}
30+
31+
public void testAvailabilityZoneAttribute() throws IOException {
32+
final var nodesInfoResponse = assertOKAndCreateObjectPath(client().performRequest(new Request("GET", "/_nodes/_all/_none")));
33+
for (final var nodeId : nodesInfoResponse.evaluateMapKeys("nodes")) {
34+
assertNull(nodesInfoResponse.evaluateExact("nodes", nodeId, "attributes", "aws_availability_zone"));
35+
}
36+
}
37+
}

‎plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeTestCase.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.discovery.ec2;
11+
12+
import org.elasticsearch.client.Request;
13+
import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
14+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
15+
import org.elasticsearch.test.rest.ESRestTestCase;
16+
import org.hamcrest.Matchers;
17+
18+
import java.io.IOException;
19+
import java.util.Objects;
20+
import java.util.Set;
21+
import java.util.function.Supplier;
22+
23+
public abstract class DiscoveryEc2AvailabilityZoneAttributeTestCase extends ESRestTestCase {
24+
25+
private static final Set<String> createdAvailabilityZones = ConcurrentCollections.newConcurrentSet();
26+
27+
protected static String getAvailabilityZone() {
28+
final var zoneName = randomIdentifier();
29+
createdAvailabilityZones.add(zoneName);
30+
return zoneName;
31+
}
32+
33+
protected static ElasticsearchCluster buildCluster(Supplier<String> imdsFixtureAddressSupplier) {
34+
return ElasticsearchCluster.local()
35+
.plugin("discovery-ec2")
36+
.setting(AwsEc2Service.AUTO_ATTRIBUTE_SETTING.getKey(), "true")
37+
.systemProperty("com.amazonaws.sdk.ec2MetadataServiceEndpointOverride", imdsFixtureAddressSupplier)
38+
.build();
39+
}
40+
41+
public void testAvailabilityZoneAttribute() throws IOException {
42+
final var nodesInfoResponse = assertOKAndCreateObjectPath(client().performRequest(new Request("GET", "/_nodes/_all/_none")));
43+
for (final var nodeId : nodesInfoResponse.evaluateMapKeys("nodes")) {
44+
assertThat(
45+
createdAvailabilityZones,
46+
Matchers.hasItem(
47+
Objects.requireNonNull(nodesInfoResponse.<String>evaluateExact("nodes", nodeId, "attributes", "aws_availability_zone"))
48+
)
49+
);
50+
}
51+
}
52+
}

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

Lines changed: 4 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
*/
99
package fixture.aws.imds;
1010

11-
import com.sun.net.httpserver.HttpHandler;
1211
import com.sun.net.httpserver.HttpServer;
1312

1413
import org.junit.rules.ExternalResource;
@@ -17,29 +16,14 @@
1716
import java.net.InetSocketAddress;
1817
import java.net.UnknownHostException;
1918
import java.util.Objects;
20-
import java.util.Set;
21-
import java.util.function.BiConsumer;
2219

2320
public class Ec2ImdsHttpFixture extends ExternalResource {
2421

22+
private final Ec2ImdsServiceBuilder ec2ImdsServiceBuilder;
2523
private HttpServer server;
2624

27-
private final Ec2ImdsVersion ec2ImdsVersion;
28-
private final BiConsumer<String, String> newCredentialsConsumer;
29-
private final Set<String> alternativeCredentialsEndpoints;
30-
31-
public Ec2ImdsHttpFixture(
32-
Ec2ImdsVersion ec2ImdsVersion,
33-
BiConsumer<String, String> newCredentialsConsumer,
34-
Set<String> alternativeCredentialsEndpoints
35-
) {
36-
this.ec2ImdsVersion = Objects.requireNonNull(ec2ImdsVersion);
37-
this.newCredentialsConsumer = Objects.requireNonNull(newCredentialsConsumer);
38-
this.alternativeCredentialsEndpoints = Objects.requireNonNull(alternativeCredentialsEndpoints);
39-
}
40-
41-
protected HttpHandler createHandler() {
42-
return new Ec2ImdsHttpHandler(ec2ImdsVersion, newCredentialsConsumer, alternativeCredentialsEndpoints);
25+
public Ec2ImdsHttpFixture(Ec2ImdsServiceBuilder ec2ImdsServiceBuilder) {
26+
this.ec2ImdsServiceBuilder = ec2ImdsServiceBuilder;
4327
}
4428

4529
public String getAddress() {
@@ -52,7 +36,7 @@ public void stop(int delay) {
5236

5337
protected void before() throws Throwable {
5438
server = HttpServer.create(resolveAddress(), 0);
55-
server.createContext("/", Objects.requireNonNull(createHandler()));
39+
server.createContext("/", Objects.requireNonNull(ec2ImdsServiceBuilder.buildHandler()));
5640
server.start();
5741
}
5842

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.util.Objects;
2727
import java.util.Set;
2828
import java.util.function.BiConsumer;
29+
import java.util.function.Supplier;
2930

3031
import static org.elasticsearch.test.ESTestCase.randomIdentifier;
3132
import static org.elasticsearch.test.ESTestCase.randomSecretKey;
@@ -43,15 +44,18 @@ public class Ec2ImdsHttpHandler implements HttpHandler {
4344

4445
private final BiConsumer<String, String> newCredentialsConsumer;
4546
private final Set<String> validCredentialsEndpoints = ConcurrentCollections.newConcurrentSet();
47+
private final Supplier<String> availabilityZoneSupplier;
4648

4749
public Ec2ImdsHttpHandler(
4850
Ec2ImdsVersion ec2ImdsVersion,
4951
BiConsumer<String, String> newCredentialsConsumer,
50-
Collection<String> alternativeCredentialsEndpoints
52+
Collection<String> alternativeCredentialsEndpoints,
53+
Supplier<String> availabilityZoneSupplier
5154
) {
5255
this.ec2ImdsVersion = Objects.requireNonNull(ec2ImdsVersion);
5356
this.newCredentialsConsumer = Objects.requireNonNull(newCredentialsConsumer);
5457
this.validCredentialsEndpoints.addAll(alternativeCredentialsEndpoints);
58+
this.availabilityZoneSupplier = availabilityZoneSupplier;
5559
}
5660

5761
@Override
@@ -98,6 +102,13 @@ public void handle(final HttpExchange exchange) throws IOException {
98102
exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
99103
exchange.getResponseBody().write(response);
100104
return;
105+
} else if (path.equals("/latest/meta-data/placement/availability-zone")) {
106+
final var availabilityZone = availabilityZoneSupplier.get();
107+
final byte[] response = availabilityZone.getBytes(StandardCharsets.UTF_8);
108+
exchange.getResponseHeaders().add("Content-Type", "text/plain");
109+
exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
110+
exchange.getResponseBody().write(response);
111+
return;
101112
} else if (validCredentialsEndpoints.contains(path)) {
102113
final String accessKey = randomIdentifier();
103114
final String sessionToken = randomIdentifier();

0 commit comments

Comments
 (0)