From 788cac4aad1d56d48a25beafc9427f71c11a3013 Mon Sep 17 00:00:00 2001 From: David Turner Date: Tue, 17 Dec 2024 08:11:55 +0000 Subject: [PATCH] Add integ test for EC2 special network addresses (#118560) (#118616) Replaces the `Ec2NetworkTests` unit test suite with an integ test suite to cover the resolution process end-to-end. --- .../RepositoryS3ImdsV1CredentialsRestIT.java | 2 +- .../RepositoryS3ImdsV2CredentialsRestIT.java | 2 +- plugins/discovery-ec2/build.gradle | 5 +- .../DiscoveryEc2NetworkAddressesTestCase.java | 42 ++++ ...DiscoveryEc2RegularNetworkAddressesIT.java | 50 +++++ ...DiscoveryEc2SpecialNetworkAddressesIT.java | 77 ++++++++ ...yEc2AvailabilityZoneAttributeTestCase.java | 4 +- .../discovery/ec2/Ec2NetworkTests.java | 181 ------------------ .../fixture/aws/imds/Ec2ImdsHttpFixture.java | 37 ++++ .../fixture/aws/imds/Ec2ImdsHttpHandler.java | 26 ++- .../aws/imds/Ec2ImdsServiceBuilder.java | 16 +- 11 files changed, 247 insertions(+), 195 deletions(-) create mode 100644 plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2NetworkAddressesTestCase.java create mode 100644 plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2RegularNetworkAddressesIT.java create mode 100644 plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2SpecialNetworkAddressesIT.java delete mode 100644 plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/Ec2NetworkTests.java 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 dcdf52e963eef..bc41b9fd62ca9 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 @@ -44,7 +44,7 @@ public class RepositoryS3ImdsV1CredentialsRestIT extends AbstractRepositoryS3Res public static ElasticsearchCluster cluster = ElasticsearchCluster.local() .module("repository-s3") .setting("s3.client." + CLIENT + ".endpoint", s3Fixture::getAddress) - .systemProperty("com.amazonaws.sdk.ec2MetadataServiceEndpointOverride", ec2ImdsHttpFixture::getAddress) + .systemProperty(Ec2ImdsHttpFixture.ENDPOINT_OVERRIDE_SYSPROP_NAME, ec2ImdsHttpFixture::getAddress) .build(); @ClassRule 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 434fc9720fc29..34500ff5227f1 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 @@ -44,7 +44,7 @@ public class RepositoryS3ImdsV2CredentialsRestIT extends AbstractRepositoryS3Res public static ElasticsearchCluster cluster = ElasticsearchCluster.local() .module("repository-s3") .setting("s3.client." + CLIENT + ".endpoint", s3Fixture::getAddress) - .systemProperty("com.amazonaws.sdk.ec2MetadataServiceEndpointOverride", ec2ImdsHttpFixture::getAddress) + .systemProperty(Ec2ImdsHttpFixture.ENDPOINT_OVERRIDE_SYSPROP_NAME, ec2ImdsHttpFixture::getAddress) .build(); @ClassRule diff --git a/plugins/discovery-ec2/build.gradle b/plugins/discovery-ec2/build.gradle index e8390fc3b1f0f..0a35171309318 100644 --- a/plugins/discovery-ec2/build.gradle +++ b/plugins/discovery-ec2/build.gradle @@ -9,6 +9,7 @@ import org.elasticsearch.gradle.internal.info.BuildParams * License v3.0 only", or the "Server Side Public License, v 1". */ apply plugin: 'elasticsearch.internal-java-rest-test' +apply plugin: 'elasticsearch.internal-cluster-test' esplugin { description 'The EC2 discovery plugin allows to use AWS API for the unicast discovery mechanism.' @@ -31,6 +32,8 @@ dependencies { javaRestTestImplementation project(':plugins:discovery-ec2') javaRestTestImplementation project(':test:fixtures:ec2-imds-fixture') + + internalClusterTestImplementation project(':test:fixtures:ec2-imds-fixture') } tasks.named("dependencyLicenses").configure { @@ -84,7 +87,7 @@ tasks.register("writeTestJavaPolicy") { } } -tasks.named("test").configure { +tasks.withType(Test).configureEach { dependsOn "writeTestJavaPolicy" // this is needed for insecure plugins, remove if possible! systemProperty 'tests.artifact', project.name diff --git a/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2NetworkAddressesTestCase.java b/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2NetworkAddressesTestCase.java new file mode 100644 index 0000000000000..b6a4845977b09 --- /dev/null +++ b/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2NetworkAddressesTestCase.java @@ -0,0 +1,42 @@ +/* + * 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.common.settings.Settings; +import org.elasticsearch.common.util.CollectionUtils; +import org.elasticsearch.http.HttpServerTransport; +import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.test.ESIntegTestCase; + +import java.io.IOException; +import java.util.Collection; + +@ESIntegTestCase.ClusterScope(numDataNodes = 0) +public abstract class DiscoveryEc2NetworkAddressesTestCase extends ESIntegTestCase { + + @Override + protected Collection> nodePlugins() { + return CollectionUtils.appendToCopyNoNullElements(super.nodePlugins(), Ec2DiscoveryPlugin.class); + } + + @Override + protected boolean addMockHttpTransport() { + return false; + } + + void verifyPublishAddress(String publishAddressSetting, String expectedAddress) throws IOException { + final var node = internalCluster().startNode(Settings.builder().put("http.publish_host", publishAddressSetting)); + assertEquals( + expectedAddress, + internalCluster().getInstance(HttpServerTransport.class, node).boundAddress().publishAddress().getAddress() + ); + internalCluster().stopNode(node); + } +} diff --git a/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2RegularNetworkAddressesIT.java b/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2RegularNetworkAddressesIT.java new file mode 100644 index 0000000000000..491fa37a4c87d --- /dev/null +++ b/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2RegularNetworkAddressesIT.java @@ -0,0 +1,50 @@ +/* + * 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.common.settings.Settings; +import org.elasticsearch.transport.BindTransportException; + +import java.io.IOException; +import java.util.concurrent.ExecutionException; + +import static org.hamcrest.Matchers.containsString; + +public class DiscoveryEc2RegularNetworkAddressesIT extends DiscoveryEc2NetworkAddressesTestCase { + public void testLocalIgnoresImds() { + Ec2ImdsHttpFixture.runWithFixture(new Ec2ImdsServiceBuilder(randomFrom(Ec2ImdsVersion.values())), imdsFixture -> { + try (var ignored = Ec2ImdsHttpFixture.withEc2MetadataServiceEndpointOverride(imdsFixture.getAddress())) { + verifyPublishAddress("_local_", "127.0.0.1"); + } + }); + } + + public void testImdsNotAvailable() throws IOException { + try (var ignored = Ec2ImdsHttpFixture.withEc2MetadataServiceEndpointOverride("http://127.0.0.1")) { + // if IMDS is not running, regular values like `_local_` should still work + verifyPublishAddress("_local_", "127.0.0.1"); + + // but EC2 addresses will cause the node to fail to start + final var assertionError = expectThrows( + AssertionError.class, + () -> internalCluster().startNode(Settings.builder().put("http.publish_host", "_ec2_")) + ); + final var executionException = asInstanceOf(ExecutionException.class, assertionError.getCause()); + final var bindTransportException = asInstanceOf(BindTransportException.class, executionException.getCause()); + assertEquals("Failed to resolve publish address", bindTransportException.getMessage()); + final var ioException = asInstanceOf(IOException.class, bindTransportException.getCause()); + assertThat(ioException.getMessage(), containsString("/latest/meta-data/local-ipv4")); + } + } +} diff --git a/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2SpecialNetworkAddressesIT.java b/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2SpecialNetworkAddressesIT.java new file mode 100644 index 0000000000000..f541c4cdd979b --- /dev/null +++ b/plugins/discovery-ec2/src/internalClusterTest/java/org/elasticsearch/discovery/ec2/DiscoveryEc2SpecialNetworkAddressesIT.java @@ -0,0 +1,77 @@ +/* + * 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 com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import java.util.Map; +import java.util.stream.Stream; + +public class DiscoveryEc2SpecialNetworkAddressesIT extends DiscoveryEc2NetworkAddressesTestCase { + + private final String imdsAddressName; + private final String elasticsearchAddressName; + private final Ec2ImdsVersion imdsVersion; + + public DiscoveryEc2SpecialNetworkAddressesIT( + @Name("imdsAddressName") String imdsAddressName, + @Name("elasticsearchAddressName") String elasticsearchAddressName, + @Name("imdsVersion") Ec2ImdsVersion imdsVersion + ) { + this.imdsAddressName = imdsAddressName; + this.elasticsearchAddressName = elasticsearchAddressName; + this.imdsVersion = imdsVersion; + } + + @ParametersFactory + public static Iterable parameters() { + return Map.of( + "_ec2:privateIpv4_", + "local-ipv4", + "_ec2:privateDns_", + "local-hostname", + "_ec2:publicIpv4_", + "public-ipv4", + "_ec2:publicDns_", + "public-hostname", + "_ec2:publicIp_", + "public-ipv4", + "_ec2:privateIp_", + "local-ipv4", + "_ec2_", + "local-ipv4" + ) + .entrySet() + .stream() + .flatMap( + addresses -> Stream.of(Ec2ImdsVersion.values()) + .map(ec2ImdsVersion -> new Object[] { addresses.getValue(), addresses.getKey(), ec2ImdsVersion }) + ) + .toList(); + } + + public void testSpecialNetworkAddresses() { + final var publishAddress = "10.0." + between(0, 255) + "." + between(0, 255); + Ec2ImdsHttpFixture.runWithFixture( + new Ec2ImdsServiceBuilder(imdsVersion).addInstanceAddress(imdsAddressName, publishAddress), + imdsFixture -> { + try (var ignored = Ec2ImdsHttpFixture.withEc2MetadataServiceEndpointOverride(imdsFixture.getAddress())) { + verifyPublishAddress(elasticsearchAddressName, publishAddress); + } + } + ); + } + +} 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 index 7eb18eec5c0b9..178c5c3ad4cae 100644 --- 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 @@ -9,6 +9,8 @@ package org.elasticsearch.discovery.ec2; +import fixture.aws.imds.Ec2ImdsHttpFixture; + import org.elasticsearch.client.Request; import org.elasticsearch.common.util.concurrent.ConcurrentCollections; import org.elasticsearch.test.cluster.ElasticsearchCluster; @@ -34,7 +36,7 @@ protected static ElasticsearchCluster buildCluster(Supplier imdsFixtureA return ElasticsearchCluster.local() .plugin("discovery-ec2") .setting(AwsEc2Service.AUTO_ATTRIBUTE_SETTING.getKey(), "true") - .systemProperty("com.amazonaws.sdk.ec2MetadataServiceEndpointOverride", imdsFixtureAddressSupplier) + .systemProperty(Ec2ImdsHttpFixture.ENDPOINT_OVERRIDE_SYSPROP_NAME, imdsFixtureAddressSupplier) .build(); } diff --git a/plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/Ec2NetworkTests.java b/plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/Ec2NetworkTests.java deleted file mode 100644 index 82787f53c9f76..0000000000000 --- a/plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/Ec2NetworkTests.java +++ /dev/null @@ -1,181 +0,0 @@ -/* - * 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 com.sun.net.httpserver.HttpServer; - -import org.elasticsearch.common.Strings; -import org.elasticsearch.common.network.NetworkService; -import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.core.SuppressForbidden; -import org.elasticsearch.mocksocket.MockHttpServer; -import org.elasticsearch.rest.RestStatus; -import org.elasticsearch.test.ESTestCase; -import org.junit.AfterClass; -import org.junit.Before; -import org.junit.BeforeClass; - -import java.io.IOException; -import java.io.OutputStream; -import java.net.InetAddress; -import java.net.InetSocketAddress; -import java.security.AccessController; -import java.security.PrivilegedAction; -import java.util.Arrays; -import java.util.Collections; -import java.util.function.BiConsumer; - -import static com.amazonaws.SDKGlobalConfiguration.EC2_METADATA_SERVICE_OVERRIDE_SYSTEM_PROPERTY; -import static java.nio.charset.StandardCharsets.UTF_8; -import static org.hamcrest.Matchers.arrayContaining; -import static org.hamcrest.Matchers.equalTo; - -/** - * Test for EC2 network.host settings. - *

- * Warning: This test doesn't assert that the exceptions are thrown. - * They aren't. - */ -@SuppressForbidden(reason = "use http server") -public class Ec2NetworkTests extends ESTestCase { - - private static HttpServer httpServer; - - @BeforeClass - public static void startHttp() throws Exception { - httpServer = MockHttpServer.createHttp(new InetSocketAddress(InetAddress.getLoopbackAddress().getHostAddress(), 0), 0); - - BiConsumer registerContext = (path, v) -> { - final byte[] message = v.getBytes(UTF_8); - httpServer.createContext(path, (s) -> { - s.sendResponseHeaders(RestStatus.OK.getStatus(), message.length); - OutputStream responseBody = s.getResponseBody(); - responseBody.write(message); - responseBody.close(); - }); - }; - registerContext.accept("/latest/meta-data/local-ipv4", "127.0.0.1"); - registerContext.accept("/latest/meta-data/public-ipv4", "165.168.10.2"); - registerContext.accept("/latest/meta-data/public-hostname", "165.168.10.3"); - registerContext.accept("/latest/meta-data/local-hostname", "10.10.10.5"); - - httpServer.start(); - } - - @Before - public void setup() { - // redirect EC2 metadata service to httpServer - AccessController.doPrivileged( - (PrivilegedAction) () -> System.setProperty( - EC2_METADATA_SERVICE_OVERRIDE_SYSTEM_PROPERTY, - "http://" + httpServer.getAddress().getHostName() + ":" + httpServer.getAddress().getPort() - ) - ); - } - - @AfterClass - public static void stopHttp() { - httpServer.stop(0); - httpServer = null; - } - - /** - * Test for network.host: _ec2_ - */ - public void testNetworkHostEc2() throws IOException { - resolveEc2("_ec2_", InetAddress.getByName("127.0.0.1")); - } - - /** - * Test for network.host: _ec2_ - */ - public void testNetworkHostUnableToResolveEc2() { - // redirect EC2 metadata service to unknown location - AccessController.doPrivileged( - (PrivilegedAction) () -> System.setProperty(EC2_METADATA_SERVICE_OVERRIDE_SYSTEM_PROPERTY, "http://127.0.0.1/") - ); - - try { - resolveEc2("_ec2_", (InetAddress[]) null); - } catch (IOException e) { - assertThat( - e.getMessage(), - equalTo("IOException caught when fetching InetAddress from [http://127.0.0.1//latest/meta-data/local-ipv4]") - ); - } - } - - /** - * Test for network.host: _ec2:publicIp_ - */ - public void testNetworkHostEc2PublicIp() throws IOException { - resolveEc2("_ec2:publicIp_", InetAddress.getByName("165.168.10.2")); - } - - /** - * Test for network.host: _ec2:privateIp_ - */ - public void testNetworkHostEc2PrivateIp() throws IOException { - resolveEc2("_ec2:privateIp_", InetAddress.getByName("127.0.0.1")); - } - - /** - * Test for network.host: _ec2:privateIpv4_ - */ - public void testNetworkHostEc2PrivateIpv4() throws IOException { - resolveEc2("_ec2:privateIpv4_", InetAddress.getByName("127.0.0.1")); - } - - /** - * Test for network.host: _ec2:privateDns_ - */ - public void testNetworkHostEc2PrivateDns() throws IOException { - resolveEc2("_ec2:privateDns_", InetAddress.getByName("10.10.10.5")); - } - - /** - * Test for network.host: _ec2:publicIpv4_ - */ - public void testNetworkHostEc2PublicIpv4() throws IOException { - resolveEc2("_ec2:publicIpv4_", InetAddress.getByName("165.168.10.2")); - } - - /** - * Test for network.host: _ec2:publicDns_ - */ - public void testNetworkHostEc2PublicDns() throws IOException { - resolveEc2("_ec2:publicDns_", InetAddress.getByName("165.168.10.3")); - } - - private InetAddress[] resolveEc2(String host, InetAddress... expected) throws IOException { - Settings nodeSettings = Settings.builder().put("network.host", host).build(); - - NetworkService networkService = new NetworkService(Collections.singletonList(new Ec2NameResolver())); - - InetAddress[] addresses = networkService.resolveBindHostAddresses( - NetworkService.GLOBAL_NETWORK_BIND_HOST_SETTING.get(nodeSettings).toArray(Strings.EMPTY_ARRAY) - ); - if (expected == null) { - fail("We should get an IOException, resolved addressed:" + Arrays.toString(addresses)); - } - assertThat(addresses, arrayContaining(expected)); - return addresses; - } - - /** - * Test that we don't have any regression with network host core settings such as - * network.host: _local_ - */ - public void testNetworkHostCoreLocal() throws IOException { - NetworkService networkService = new NetworkService(Collections.singletonList(new Ec2NameResolver())); - InetAddress[] addresses = networkService.resolveBindHostAddresses(null); - assertThat(addresses, arrayContaining(networkService.resolveBindHostAddresses(new String[] { "_local_" }))); - } -} 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 cc268a6021cb3..e232d10fdddbd 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 @@ -10,15 +10,24 @@ import com.sun.net.httpserver.HttpServer; +import org.elasticsearch.core.CheckedConsumer; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.SuppressForbidden; import org.junit.rules.ExternalResource; +import org.junit.runner.Description; +import org.junit.runners.model.Statement; import java.net.InetAddress; import java.net.InetSocketAddress; import java.net.UnknownHostException; +import java.security.AccessController; +import java.security.PrivilegedAction; import java.util.Objects; public class Ec2ImdsHttpFixture extends ExternalResource { + public static final String ENDPOINT_OVERRIDE_SYSPROP_NAME = "com.amazonaws.sdk.ec2MetadataServiceEndpointOverride"; + private final Ec2ImdsServiceBuilder ec2ImdsServiceBuilder; private HttpServer server; @@ -52,4 +61,32 @@ private static InetSocketAddress resolveAddress() { throw new RuntimeException(e); } } + + @SuppressForbidden(reason = "deliberately adjusting system property for endpoint override for use in internal-cluster tests") + public static Releasable withEc2MetadataServiceEndpointOverride(String endpointOverride) { + final PrivilegedAction resetProperty = System.getProperty(ENDPOINT_OVERRIDE_SYSPROP_NAME) instanceof String originalValue + ? () -> System.setProperty(ENDPOINT_OVERRIDE_SYSPROP_NAME, originalValue) + : () -> System.clearProperty(ENDPOINT_OVERRIDE_SYSPROP_NAME); + doPrivileged(() -> System.setProperty(ENDPOINT_OVERRIDE_SYSPROP_NAME, endpointOverride)); + return () -> doPrivileged(resetProperty); + } + + private static void doPrivileged(PrivilegedAction privilegedAction) { + AccessController.doPrivileged(privilegedAction); + } + + public static void runWithFixture(Ec2ImdsServiceBuilder ec2ImdsServiceBuilder, CheckedConsumer action) { + final var imdsFixture = new Ec2ImdsHttpFixture(ec2ImdsServiceBuilder); + try { + imdsFixture.apply(new Statement() { + @Override + public void evaluate() throws Exception { + action.accept(imdsFixture); + } + }, Description.EMPTY).evaluate(); + } catch (Throwable e) { + throw new AssertionError(e); + } + } + } 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 fd2044357257b..0c58205ac8d60 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 @@ -23,6 +23,7 @@ import java.time.ZonedDateTime; import java.time.format.DateTimeFormatter; import java.util.Collection; +import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.function.BiConsumer; @@ -43,6 +44,7 @@ public class Ec2ImdsHttpHandler implements HttpHandler { private final Set validImdsTokens = ConcurrentCollections.newConcurrentSet(); private final BiConsumer newCredentialsConsumer; + private final Map instanceAddresses; private final Set validCredentialsEndpoints = ConcurrentCollections.newConcurrentSet(); private final Supplier availabilityZoneSupplier; @@ -50,10 +52,12 @@ public Ec2ImdsHttpHandler( Ec2ImdsVersion ec2ImdsVersion, BiConsumer newCredentialsConsumer, Collection alternativeCredentialsEndpoints, - Supplier availabilityZoneSupplier + Supplier availabilityZoneSupplier, + Map instanceAddresses ) { this.ec2ImdsVersion = Objects.requireNonNull(ec2ImdsVersion); this.newCredentialsConsumer = Objects.requireNonNull(newCredentialsConsumer); + this.instanceAddresses = instanceAddresses; this.validCredentialsEndpoints.addAll(alternativeCredentialsEndpoints); this.availabilityZoneSupplier = availabilityZoneSupplier; } @@ -97,17 +101,11 @@ public void handle(final HttpExchange exchange) throws IOException { if (path.equals(IMDS_SECURITY_CREDENTIALS_PATH)) { final var profileName = randomIdentifier(); validCredentialsEndpoints.add(IMDS_SECURITY_CREDENTIALS_PATH + profileName); - final byte[] response = profileName.getBytes(StandardCharsets.UTF_8); - exchange.getResponseHeaders().add("Content-Type", "text/plain"); - exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length); - exchange.getResponseBody().write(response); + sendStringResponse(exchange, profileName); 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); + sendStringResponse(exchange, availabilityZone); return; } else if (validCredentialsEndpoints.contains(path)) { final String accessKey = randomIdentifier(); @@ -132,10 +130,20 @@ public void handle(final HttpExchange exchange) throws IOException { exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length); exchange.getResponseBody().write(response); return; + } else if (instanceAddresses.get(path) instanceof String instanceAddress) { + sendStringResponse(exchange, instanceAddress); + return; } } ExceptionsHelper.maybeDieOnAnotherThread(new AssertionError("not supported: " + requestMethod + " " + path)); } } + + private void sendStringResponse(HttpExchange exchange, String value) throws IOException { + final byte[] response = value.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "text/plain"); + exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length); + exchange.getResponseBody().write(response); + } } 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 index bca43da8683b6..505c9978bc4fb 100644 --- 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 @@ -12,6 +12,8 @@ import org.elasticsearch.test.ESTestCase; import java.util.Collection; +import java.util.HashMap; +import java.util.Map; import java.util.Set; import java.util.function.BiConsumer; import java.util.function.Supplier; @@ -22,6 +24,7 @@ public class Ec2ImdsServiceBuilder { private BiConsumer newCredentialsConsumer = Ec2ImdsServiceBuilder::rejectNewCredentials; private Collection alternativeCredentialsEndpoints = Set.of(); private Supplier availabilityZoneSupplier = Ec2ImdsServiceBuilder::rejectAvailabilityZone; + private final Map instanceAddresses = new HashMap<>(); public Ec2ImdsServiceBuilder(Ec2ImdsVersion ec2ImdsVersion) { this.ec2ImdsVersion = ec2ImdsVersion; @@ -50,8 +53,19 @@ public Ec2ImdsServiceBuilder availabilityZoneSupplier(Supplier availabil return this; } + public Ec2ImdsServiceBuilder addInstanceAddress(String addressType, String addressValue) { + instanceAddresses.put("/latest/meta-data/" + addressType, addressValue); + return this; + } + public Ec2ImdsHttpHandler buildHandler() { - return new Ec2ImdsHttpHandler(ec2ImdsVersion, newCredentialsConsumer, alternativeCredentialsEndpoints, availabilityZoneSupplier); + return new Ec2ImdsHttpHandler( + ec2ImdsVersion, + newCredentialsConsumer, + alternativeCredentialsEndpoints, + availabilityZoneSupplier, + Map.copyOf(instanceAddresses) + ); } }