[8.x] Add discovery-ec2 integration test for AZ attr (#118452)

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

@@ -10,6 +10,7 @@
 package org.elasticsearch.repositories.s3;
 
 import fixture.aws.imds.Ec2ImdsHttpFixture;
+import fixture.aws.imds.Ec2ImdsServiceBuilder;
 import fixture.aws.imds.Ec2ImdsVersion;
 import fixture.s3.DynamicS3Credentials;
 import fixture.s3.S3HttpFixture;
@@ -37,9 +38,8 @@ public class RepositoryS3EcsCredentialsRestIT extends AbstractRepositoryS3RestTe
     private static final DynamicS3Credentials dynamicS3Credentials = new DynamicS3Credentials();
 
     private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
-        Ec2ImdsVersion.V1,
-        dynamicS3Credentials::addValidCredentials,
-        Set.of("/ecs_credentials_endpoint")
+        new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).newCredentialsConsumer(dynamicS3Credentials::addValidCredentials)
+            .alternativeCredentialsEndpoints(Set.of("/ecs_credentials_endpoint"))
     );
 
     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

@@ -10,6 +10,7 @@
 package org.elasticsearch.repositories.s3;
 
 import fixture.aws.imds.Ec2ImdsHttpFixture;
+import fixture.aws.imds.Ec2ImdsServiceBuilder;
 import fixture.aws.imds.Ec2ImdsVersion;
 import fixture.s3.DynamicS3Credentials;
 import fixture.s3.S3HttpFixture;
@@ -23,8 +24,6 @@
 import org.junit.rules.RuleChain;
 import org.junit.rules.TestRule;
 
-import java.util.Set;
-
 @ThreadLeakFilters(filters = { TestContainersThreadFilter.class })
 @ThreadLeakScope(ThreadLeakScope.Scope.NONE) // https://github.com/elastic/elasticsearch/issues/102482
 public class RepositoryS3ImdsV1CredentialsRestIT extends AbstractRepositoryS3RestTestCase {
@@ -37,9 +36,7 @@ public class RepositoryS3ImdsV1CredentialsRestIT extends AbstractRepositoryS3Res
     private static final DynamicS3Credentials dynamicS3Credentials = new DynamicS3Credentials();
 
     private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
-        Ec2ImdsVersion.V1,
-        dynamicS3Credentials::addValidCredentials,
-        Set.of()
+        new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).newCredentialsConsumer(dynamicS3Credentials::addValidCredentials)
     );
 
     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

@@ -10,6 +10,7 @@
 package org.elasticsearch.repositories.s3;
 
 import fixture.aws.imds.Ec2ImdsHttpFixture;
+import fixture.aws.imds.Ec2ImdsServiceBuilder;
 import fixture.aws.imds.Ec2ImdsVersion;
 import fixture.s3.DynamicS3Credentials;
 import fixture.s3.S3HttpFixture;
@@ -23,8 +24,6 @@
 import org.junit.rules.RuleChain;
 import org.junit.rules.TestRule;
 
-import java.util.Set;
-
 @ThreadLeakFilters(filters = { TestContainersThreadFilter.class })
 @ThreadLeakScope(ThreadLeakScope.Scope.NONE) // https://github.com/elastic/elasticsearch/issues/102482
 public class RepositoryS3ImdsV2CredentialsRestIT extends AbstractRepositoryS3RestTestCase {
@@ -37,9 +36,7 @@ public class RepositoryS3ImdsV2CredentialsRestIT extends AbstractRepositoryS3Res
     private static final DynamicS3Credentials dynamicS3Credentials = new DynamicS3Credentials();
 
     private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
-        Ec2ImdsVersion.V2,
-        dynamicS3Credentials::addValidCredentials,
-        Set.of()
+        new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V2).newCredentialsConsumer(dynamicS3Credentials::addValidCredentials)
     );
 
     private static final S3HttpFixture s3Fixture = new S3HttpFixture(true, BUCKET, BASE_PATH, dynamicS3Credentials::isAuthorized);

plugins/discovery-ec2/build.gradle

@@ -28,6 +28,9 @@ dependencies {
   api "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}"
   api "com.fasterxml.jackson.dataformat:jackson-dataformat-cbor:${versions.jackson}"
   api "joda-time:joda-time:2.10.10"
+
+  javaRestTestImplementation project(':plugins:discovery-ec2')
+  javaRestTestImplementation project(':test:fixtures:ec2-imds-fixture')
 }
 
 tasks.named("dependencyLicenses").configure {

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.discovery.ec2;
+
+import fixture.aws.imds.Ec2ImdsHttpFixture;
+import fixture.aws.imds.Ec2ImdsServiceBuilder;
+import fixture.aws.imds.Ec2ImdsVersion;
+
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.junit.ClassRule;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TestRule;
+
+public class DiscoveryEc2AvailabilityZoneAttributeImdsV1IT extends DiscoveryEc2AvailabilityZoneAttributeTestCase {
+    private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
+        new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).availabilityZoneSupplier(
+            DiscoveryEc2AvailabilityZoneAttributeTestCase::getAvailabilityZone
+        )
+    );
+
+    public static ElasticsearchCluster cluster = buildCluster(ec2ImdsHttpFixture::getAddress);
+
+    @ClassRule
+    public static TestRule ruleChain = RuleChain.outerRule(ec2ImdsHttpFixture).around(cluster);
+
+    @Override
+    protected String getTestRestCluster() {
+        return cluster.getHttpAddresses();
+    }
+}

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.discovery.ec2;
+
+import fixture.aws.imds.Ec2ImdsHttpFixture;
+import fixture.aws.imds.Ec2ImdsServiceBuilder;
+import fixture.aws.imds.Ec2ImdsVersion;
+
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.junit.ClassRule;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TestRule;
+
+public class DiscoveryEc2AvailabilityZoneAttributeImdsV2IT extends DiscoveryEc2AvailabilityZoneAttributeTestCase {
+    private static final Ec2ImdsHttpFixture ec2ImdsHttpFixture = new Ec2ImdsHttpFixture(
+        new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V2).availabilityZoneSupplier(
+            DiscoveryEc2AvailabilityZoneAttributeTestCase::getAvailabilityZone
+        )
+    );
+
+    public static ElasticsearchCluster cluster = buildCluster(ec2ImdsHttpFixture::getAddress);
+
+    @ClassRule
+    public static TestRule ruleChain = RuleChain.outerRule(ec2ImdsHttpFixture).around(cluster);
+
+    @Override
+    protected String getTestRestCluster() {
+        return cluster.getHttpAddresses();
+    }
+}

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.discovery.ec2;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.elasticsearch.test.rest.ESRestTestCase;
+import org.junit.ClassRule;
+
+import java.io.IOException;
+
+public class DiscoveryEc2AvailabilityZoneAttributeNoImdsIT extends ESRestTestCase {
+    @ClassRule
+    public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
+        .plugin("discovery-ec2")
+        .setting(AwsEc2Service.AUTO_ATTRIBUTE_SETTING.getKey(), "true")
+        .build();
+
+    @Override
+    protected String getTestRestCluster() {
+        return cluster.getHttpAddresses();
+    }
+
+    public void testAvailabilityZoneAttribute() throws IOException {
+        final var nodesInfoResponse = assertOKAndCreateObjectPath(client().performRequest(new Request("GET", "/_nodes/_all/_none")));
+        for (final var nodeId : nodesInfoResponse.evaluateMapKeys("nodes")) {
+            assertNull(nodesInfoResponse.evaluateExact("nodes", nodeId, "attributes", "aws_availability_zone"));
+        }
+    }
+}

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.discovery.ec2;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.elasticsearch.test.rest.ESRestTestCase;
+import org.hamcrest.Matchers;
+
+import java.io.IOException;
+import java.util.Objects;
+import java.util.Set;
+import java.util.function.Supplier;
+
+public abstract class DiscoveryEc2AvailabilityZoneAttributeTestCase extends ESRestTestCase {
+
+    private static final Set<String> createdAvailabilityZones = ConcurrentCollections.newConcurrentSet();
+
+    protected static String getAvailabilityZone() {
+        final var zoneName = randomIdentifier();
+        createdAvailabilityZones.add(zoneName);
+        return zoneName;
+    }
+
+    protected static ElasticsearchCluster buildCluster(Supplier<String> imdsFixtureAddressSupplier) {
+        return ElasticsearchCluster.local()
+            .plugin("discovery-ec2")
+            .setting(AwsEc2Service.AUTO_ATTRIBUTE_SETTING.getKey(), "true")
+            .systemProperty("com.amazonaws.sdk.ec2MetadataServiceEndpointOverride", imdsFixtureAddressSupplier)
+            .build();
+    }
+
+    public void testAvailabilityZoneAttribute() throws IOException {
+        final var nodesInfoResponse = assertOKAndCreateObjectPath(client().performRequest(new Request("GET", "/_nodes/_all/_none")));
+        for (final var nodeId : nodesInfoResponse.evaluateMapKeys("nodes")) {
+            assertThat(
+                createdAvailabilityZones,
+                Matchers.hasItem(
+                    Objects.requireNonNull(nodesInfoResponse.<String>evaluateExact("nodes", nodeId, "attributes", "aws_availability_zone"))
+                )
+            );
+        }
+    }
+}

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

@@ -8,7 +8,6 @@
  */
 package fixture.aws.imds;
 
-import com.sun.net.httpserver.HttpHandler;
 import com.sun.net.httpserver.HttpServer;
 
 import org.junit.rules.ExternalResource;
@@ -17,29 +16,14 @@
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
 import java.util.Objects;
-import java.util.Set;
-import java.util.function.BiConsumer;
 
 public class Ec2ImdsHttpFixture extends ExternalResource {
 
+    private final Ec2ImdsServiceBuilder ec2ImdsServiceBuilder;
     private HttpServer server;
 
-    private final Ec2ImdsVersion ec2ImdsVersion;
-    private final BiConsumer<String, String> newCredentialsConsumer;
-    private final Set<String> alternativeCredentialsEndpoints;
-
-    public Ec2ImdsHttpFixture(
-        Ec2ImdsVersion ec2ImdsVersion,
-        BiConsumer<String, String> newCredentialsConsumer,
-        Set<String> alternativeCredentialsEndpoints
-    ) {
-        this.ec2ImdsVersion = Objects.requireNonNull(ec2ImdsVersion);
-        this.newCredentialsConsumer = Objects.requireNonNull(newCredentialsConsumer);
-        this.alternativeCredentialsEndpoints = Objects.requireNonNull(alternativeCredentialsEndpoints);
-    }
-
-    protected HttpHandler createHandler() {
-        return new Ec2ImdsHttpHandler(ec2ImdsVersion, newCredentialsConsumer, alternativeCredentialsEndpoints);
+    public Ec2ImdsHttpFixture(Ec2ImdsServiceBuilder ec2ImdsServiceBuilder) {
+        this.ec2ImdsServiceBuilder = ec2ImdsServiceBuilder;
     }
 
     public String getAddress() {
@@ -52,7 +36,7 @@ public void stop(int delay) {
 
     protected void before() throws Throwable {
         server = HttpServer.create(resolveAddress(), 0);
-        server.createContext("/", Objects.requireNonNull(createHandler()));
+        server.createContext("/", Objects.requireNonNull(ec2ImdsServiceBuilder.buildHandler()));
         server.start();
     }
 

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

@@ -26,6 +26,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.function.Supplier;
 
 import static org.elasticsearch.test.ESTestCase.randomIdentifier;
 import static org.elasticsearch.test.ESTestCase.randomSecretKey;
@@ -43,15 +44,18 @@ public class Ec2ImdsHttpHandler implements HttpHandler {
 
     private final BiConsumer<String, String> newCredentialsConsumer;
     private final Set<String> validCredentialsEndpoints = ConcurrentCollections.newConcurrentSet();
+    private final Supplier<String> availabilityZoneSupplier;
 
     public Ec2ImdsHttpHandler(
         Ec2ImdsVersion ec2ImdsVersion,
         BiConsumer<String, String> newCredentialsConsumer,
-        Collection<String> alternativeCredentialsEndpoints
+        Collection<String> alternativeCredentialsEndpoints,
+        Supplier<String> availabilityZoneSupplier
     ) {
         this.ec2ImdsVersion = Objects.requireNonNull(ec2ImdsVersion);
         this.newCredentialsConsumer = Objects.requireNonNull(newCredentialsConsumer);
         this.validCredentialsEndpoints.addAll(alternativeCredentialsEndpoints);
+        this.availabilityZoneSupplier = availabilityZoneSupplier;
     }
 
     @Override
@@ -98,6 +102,13 @@ public void handle(final HttpExchange exchange) throws IOException {
                     exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
                     exchange.getResponseBody().write(response);
                     return;
+                } else if (path.equals("/latest/meta-data/placement/availability-zone")) {
+                    final var availabilityZone = availabilityZoneSupplier.get();
+                    final byte[] response = availabilityZone.getBytes(StandardCharsets.UTF_8);
+                    exchange.getResponseHeaders().add("Content-Type", "text/plain");
+                    exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
+                    exchange.getResponseBody().write(response);
+                    return;
                 } else if (validCredentialsEndpoints.contains(path)) {
                     final String accessKey = randomIdentifier();
                     final String sessionToken = randomIdentifier();