From 19854744b051ecb1fb83776146af17a720bf633e Mon Sep 17 00:00:00 2001 From: David Turner Date: Thu, 12 Dec 2024 10:57:57 +0000 Subject: [PATCH] Add `discovery-ec2` integration test for AZ attr (#118452) Verifies that the plugin sets the `aws_availability_zone` automatically by reading the AZ name from the IMDS at startup. --- .../s3/RepositoryS3EcsCredentialsRestIT.java | 6 +- .../RepositoryS3ImdsV1CredentialsRestIT.java | 7 +-- .../RepositoryS3ImdsV2CredentialsRestIT.java | 7 +-- plugins/discovery-ec2/build.gradle | 3 + ...yEc2AvailabilityZoneAttributeImdsV1IT.java | 37 ++++++++++++ ...yEc2AvailabilityZoneAttributeImdsV2IT.java | 37 ++++++++++++ ...yEc2AvailabilityZoneAttributeNoImdsIT.java | 37 ++++++++++++ ...yEc2AvailabilityZoneAttributeTestCase.java | 52 +++++++++++++++++ .../fixture/aws/imds/Ec2ImdsHttpFixture.java | 24 ++------ .../fixture/aws/imds/Ec2ImdsHttpHandler.java | 13 ++++- .../aws/imds/Ec2ImdsServiceBuilder.java | 57 +++++++++++++++++++ .../aws/imds/Ec2ImdsHttpHandlerTests.java | 26 ++++++--- 12 files changed, 265 insertions(+), 41 deletions(-) create mode 100644 plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV1IT.java create mode 100644 plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV2IT.java create mode 100644 plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeNoImdsIT.java create mode 100644 plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeTestCase.java create mode 100644 test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsServiceBuilder.java diff --git a/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3EcsCredentialsRestIT.java b/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3EcsCredentialsRestIT.java index a79ae4de7cc66..4f0bf83000642 100644 --- a/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3EcsCredentialsRestIT.java +++ b/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); diff --git a/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV1CredentialsRestIT.java b/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV1CredentialsRestIT.java index ead91981b3fa8..dcdf52e963eef 100644 --- a/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV1CredentialsRestIT.java +++ b/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); diff --git a/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV2CredentialsRestIT.java b/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV2CredentialsRestIT.java index 67adb096bd1ba..434fc9720fc29 100644 --- a/modules/repository-s3/src/javaRestTest/java/org/elasticsearch/repositories/s3/RepositoryS3ImdsV2CredentialsRestIT.java +++ b/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); diff --git a/plugins/discovery-ec2/build.gradle b/plugins/discovery-ec2/build.gradle index 591c42c0b803c..e8390fc3b1f0f 100644 --- a/plugins/discovery-ec2/build.gradle +++ b/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 { diff --git a/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV1IT.java b/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV1IT.java new file mode 100644 index 0000000000000..32291236ea158 --- /dev/null +++ b/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(); + } +} diff --git a/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV2IT.java b/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeImdsV2IT.java new file mode 100644 index 0000000000000..8b785d688e7c4 --- /dev/null +++ b/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(); + } +} diff --git a/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeNoImdsIT.java b/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeNoImdsIT.java new file mode 100644 index 0000000000000..602a98e17970d --- /dev/null +++ b/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")); + } + } +} diff --git a/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeTestCase.java b/plugins/discovery-ec2/src/javaRestTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2AvailabilityZoneAttributeTestCase.java new file mode 100644 index 0000000000000..7eb18eec5c0b9 --- /dev/null +++ b/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 createdAvailabilityZones = ConcurrentCollections.newConcurrentSet(); + + protected static String getAvailabilityZone() { + final var zoneName = randomIdentifier(); + createdAvailabilityZones.add(zoneName); + return zoneName; + } + + protected static ElasticsearchCluster buildCluster(Supplier 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.evaluateExact("nodes", nodeId, "attributes", "aws_availability_zone")) + ) + ); + } + } +} diff --git a/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsHttpFixture.java b/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsHttpFixture.java index c63c65a750d7c..cc268a6021cb3 100644 --- a/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsHttpFixture.java +++ b/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 newCredentialsConsumer; - private final Set alternativeCredentialsEndpoints; - - public Ec2ImdsHttpFixture( - Ec2ImdsVersion ec2ImdsVersion, - BiConsumer newCredentialsConsumer, - Set 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(); } diff --git a/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsHttpHandler.java b/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsHttpHandler.java index 281465b96de05..fd2044357257b 100644 --- a/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsHttpHandler.java +++ b/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 newCredentialsConsumer; private final Set validCredentialsEndpoints = ConcurrentCollections.newConcurrentSet(); + private final Supplier availabilityZoneSupplier; public Ec2ImdsHttpHandler( Ec2ImdsVersion ec2ImdsVersion, BiConsumer newCredentialsConsumer, - Collection alternativeCredentialsEndpoints + Collection alternativeCredentialsEndpoints, + Supplier 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(); diff --git a/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsServiceBuilder.java b/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsServiceBuilder.java new file mode 100644 index 0000000000000..bca43da8683b6 --- /dev/null +++ b/test/fixtures/ec2-imds-fixture/src/main/java/fixture/aws/imds/Ec2ImdsServiceBuilder.java @@ -0,0 +1,57 @@ +/* + * 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 fixture.aws.imds; + +import org.elasticsearch.test.ESTestCase; + +import java.util.Collection; +import java.util.Set; +import java.util.function.BiConsumer; +import java.util.function.Supplier; + +public class Ec2ImdsServiceBuilder { + + private final Ec2ImdsVersion ec2ImdsVersion; + private BiConsumer newCredentialsConsumer = Ec2ImdsServiceBuilder::rejectNewCredentials; + private Collection alternativeCredentialsEndpoints = Set.of(); + private Supplier availabilityZoneSupplier = Ec2ImdsServiceBuilder::rejectAvailabilityZone; + + public Ec2ImdsServiceBuilder(Ec2ImdsVersion ec2ImdsVersion) { + this.ec2ImdsVersion = ec2ImdsVersion; + } + + public Ec2ImdsServiceBuilder newCredentialsConsumer(BiConsumer newCredentialsConsumer) { + this.newCredentialsConsumer = newCredentialsConsumer; + return this; + } + + private static void rejectNewCredentials(String ignored1, String ignored2) { + ESTestCase.fail("credentials creation not supported"); + } + + public Ec2ImdsServiceBuilder alternativeCredentialsEndpoints(Collection alternativeCredentialsEndpoints) { + this.alternativeCredentialsEndpoints = alternativeCredentialsEndpoints; + return this; + } + + private static String rejectAvailabilityZone() { + return ESTestCase.fail(null, "availability zones not supported"); + } + + public Ec2ImdsServiceBuilder availabilityZoneSupplier(Supplier availabilityZoneSupplier) { + this.availabilityZoneSupplier = availabilityZoneSupplier; + return this; + } + + public Ec2ImdsHttpHandler buildHandler() { + return new Ec2ImdsHttpHandler(ec2ImdsVersion, newCredentialsConsumer, alternativeCredentialsEndpoints, availabilityZoneSupplier); + } + +} diff --git a/test/fixtures/ec2-imds-fixture/src/test/java/fixture/aws/imds/Ec2ImdsHttpHandlerTests.java b/test/fixtures/ec2-imds-fixture/src/test/java/fixture/aws/imds/Ec2ImdsHttpHandlerTests.java index bb613395a0fba..6d3eb3d14e9b2 100644 --- a/test/fixtures/ec2-imds-fixture/src/test/java/fixture/aws/imds/Ec2ImdsHttpHandlerTests.java +++ b/test/fixtures/ec2-imds-fixture/src/test/java/fixture/aws/imds/Ec2ImdsHttpHandlerTests.java @@ -30,6 +30,7 @@ import java.net.InetSocketAddress; import java.net.URI; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; @@ -43,7 +44,7 @@ public class Ec2ImdsHttpHandlerTests extends ESTestCase { public void testImdsV1() throws IOException { final Map generatedCredentials = new HashMap<>(); - final var handler = new Ec2ImdsHttpHandler(Ec2ImdsVersion.V1, generatedCredentials::put, Set.of()); + final var handler = new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).newCredentialsConsumer(generatedCredentials::put).buildHandler(); final var roleResponse = handleRequest(handler, "GET", SECURITY_CREDENTIALS_URI); assertEquals(RestStatus.OK, roleResponse.status()); @@ -66,18 +67,14 @@ public void testImdsV1() throws IOException { public void testImdsV2Disabled() { assertEquals( RestStatus.METHOD_NOT_ALLOWED, - handleRequest( - new Ec2ImdsHttpHandler(Ec2ImdsVersion.V1, (accessKey, sessionToken) -> fail(), Set.of()), - "PUT", - "/latest/api/token" - ).status() + handleRequest(new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).buildHandler(), "PUT", "/latest/api/token").status() ); } public void testImdsV2() throws IOException { final Map generatedCredentials = new HashMap<>(); - final var handler = new Ec2ImdsHttpHandler(Ec2ImdsVersion.V2, generatedCredentials::put, Set.of()); + final var handler = new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V2).newCredentialsConsumer(generatedCredentials::put).buildHandler(); final var tokenResponse = handleRequest(handler, "PUT", "/latest/api/token"); assertEquals(RestStatus.OK, tokenResponse.status()); @@ -101,6 +98,21 @@ public void testImdsV2() throws IOException { assertEquals(sessionToken, responseMap.get("Token")); } + public void testAvailabilityZone() { + final Set generatedAvailabilityZones = new HashSet<>(); + final var handler = new Ec2ImdsServiceBuilder(Ec2ImdsVersion.V1).availabilityZoneSupplier(() -> { + final var newAvailabilityZone = randomIdentifier(); + generatedAvailabilityZones.add(newAvailabilityZone); + return newAvailabilityZone; + }).buildHandler(); + + final var availabilityZoneResponse = handleRequest(handler, "GET", "/latest/meta-data/placement/availability-zone"); + assertEquals(RestStatus.OK, availabilityZoneResponse.status()); + final var availabilityZone = availabilityZoneResponse.body().utf8ToString(); + + assertEquals(generatedAvailabilityZones, Set.of(availabilityZone)); + } + private record TestHttpResponse(RestStatus status, BytesReference body) {} private static TestHttpResponse checkImdsV2GetRequest(Ec2ImdsHttpHandler handler, String uri, String token) {