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.

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

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

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

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<Class<? extends Plugin>> 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);
+    }
+}

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"));
+        }
+    }
+}

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<Object[]> 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);
+                }
+            }
+        );
+    }
+
+}

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 @@ public abstract class DiscoveryEc2AvailabilityZoneAttributeTestCase extends ESRe
         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();
     }
 

plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/Ec2NetworkTests.java

@@ -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.
- * <p>
- * 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<String, String> 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<String>) () -> 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<String>) () -> 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_" })));
-    }
-}

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

@@ -10,15 +10,24 @@ package fixture.aws.imds;
 
 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 @@ public class Ec2ImdsHttpFixture extends ExternalResource {
             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<String> 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<Ec2ImdsHttpFixture, Exception> 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);
+        }
+    }
+
 }

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

@@ -23,6 +23,7 @@ import java.time.Clock;
 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<String> validImdsTokens = ConcurrentCollections.newConcurrentSet();
 
     private final BiConsumer<String, String> newCredentialsConsumer;
+    private final Map<String, String> instanceAddresses;
     private final Set<String> validCredentialsEndpoints = ConcurrentCollections.newConcurrentSet();
     private final Supplier<String> availabilityZoneSupplier;
 
@@ -50,10 +52,12 @@ public class Ec2ImdsHttpHandler implements HttpHandler {
         Ec2ImdsVersion ec2ImdsVersion,
         BiConsumer<String, String> newCredentialsConsumer,
         Collection<String> alternativeCredentialsEndpoints,
-        Supplier<String> availabilityZoneSupplier
+        Supplier<String> availabilityZoneSupplier,
+        Map<String, String> 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 class Ec2ImdsHttpHandler implements HttpHandler {
                 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 class Ec2ImdsHttpHandler implements HttpHandler {
                     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);
+    }
 }

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

@@ -12,6 +12,8 @@ package fixture.aws.imds;
 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<String, String> newCredentialsConsumer = Ec2ImdsServiceBuilder::rejectNewCredentials;
     private Collection<String> alternativeCredentialsEndpoints = Set.of();
     private Supplier<String> availabilityZoneSupplier = Ec2ImdsServiceBuilder::rejectAvailabilityZone;
+    private final Map<String, String> instanceAddresses = new HashMap<>();
 
     public Ec2ImdsServiceBuilder(Ec2ImdsVersion ec2ImdsVersion) {
         this.ec2ImdsVersion = ec2ImdsVersion;
@@ -50,8 +53,19 @@ public class Ec2ImdsServiceBuilder {
         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)
+        );
     }
 
 }