From 84778f090da1e175be331cfbb592766df2c8eb3e Mon Sep 17 00:00:00 2001
From: Anuraag Agrawal <aanuraag@amazon.co.jp>
Date: Thu, 28 May 2020 14:14:41 +0900
Subject: [PATCH 1/4] Migrate EC2 metadata lookup to JDK requests and support
 IMDSv2

---
 .../xray/plugins/EC2MetadataFetcher.java      | 151 ++++++++++++++++++
 .../com/amazonaws/xray/plugins/EC2Plugin.java |  30 ++--
 .../xray/plugins/EC2MetadataFetcherTest.java  |  74 +++++++++
 .../amazonaws/xray/plugins/EC2PluginTest.java |  56 +++++--
 4 files changed, 284 insertions(+), 27 deletions(-)
 create mode 100644 aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
 create mode 100644 aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java

diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
new file mode 100644
index 00000000..0e47fbd6
--- /dev/null
+++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
@@ -0,0 +1,151 @@
+package com.amazonaws.xray.plugins;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UnsupportedEncodingException;
+import java.net.HttpURLConnection;
+import java.net.MalformedURLException;
+import java.net.ProtocolException;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+class EC2MetadataFetcher {
+    private static final Log logger = LogFactory.getLog(EC2MetadataFetcher.class);
+
+    enum EC2Metadata {
+        INSTANCE_ID,
+        AVAILABILITY_ZONE,
+    }
+
+    private static final int TIMEOUT_MILLIS = 2000;
+    private static final String DEFAULT_IDMS_ENDPOINT = "169.254.169.254";
+
+    private final URL tokenUrl;
+    private final URL instanceIdUrl;
+    private final URL availabilityZoneUrl;
+
+    EC2MetadataFetcher() {
+        this(System.getenv("IMDS_ENDPOINT") != null ? System.getenv("IMDS_ENDPOINT") : DEFAULT_IDMS_ENDPOINT);
+    }
+
+    EC2MetadataFetcher(String endpoint) {
+        String urlBase = "http://" + endpoint;
+        try {
+            this.tokenUrl = new URL(urlBase + "/latest/api/token");
+            this.instanceIdUrl = new URL(urlBase + "/latest/meta-data/instance-id");
+            this.availabilityZoneUrl = new URL(urlBase + "/latest/meta-data/placement/availability-zone");
+        } catch (MalformedURLException e) {
+            throw new IllegalArgumentException("Illegal endpoint: " + endpoint);
+        }
+    }
+
+    Map<EC2Metadata, String> fetch() {
+        String token = fetchToken();
+
+        // If token is empty, either IDMSv2 isn't enabled or an unexpected failure happened. We can still get
+        // data if IDMSv1 is enabled.
+        String instanceId = fetchInstanceId(token);
+        if (instanceId.isEmpty()) {
+            // If no instance ID, assume we are not actually running on EC2.
+            return Collections.emptyMap();
+        }
+
+        String availabilityZone = fetchAvailabilityZone(token);
+
+        Map<EC2Metadata, String> result = new LinkedHashMap<>();
+        result.put(EC2Metadata.INSTANCE_ID, instanceId);
+        result.put(EC2Metadata.AVAILABILITY_ZONE, availabilityZone);
+        return result;
+    }
+
+    private String fetchToken() {
+        return fetchString("PUT", tokenUrl, "", true);
+    }
+
+    private String fetchInstanceId(String token) {
+        return fetchString("GET", instanceIdUrl, token, false);
+    }
+
+    private String fetchAvailabilityZone(String token) {
+        return fetchString("GET", availabilityZoneUrl, token, false);
+    }
+
+    // Generic HTTP fetch function for IDMS.
+    private static String fetchString(String httpMethod, URL url, String token, boolean includeTtl) {
+        final HttpURLConnection connection;
+        try {
+            connection = (HttpURLConnection) url.openConnection();
+        } catch (Exception e) {
+            logger.warn("Error connecting to IDMS.", e);
+            return "";
+        }
+
+        try {
+            connection.setRequestMethod(httpMethod);
+        } catch (ProtocolException e) {
+            logger.warn("Unknown HTTP method, this is a programming bug.", e);
+            return "";
+        }
+
+        connection.setConnectTimeout(TIMEOUT_MILLIS);
+        connection.setReadTimeout(TIMEOUT_MILLIS);
+
+        if (includeTtl) {
+            connection.setRequestProperty("X-aws-ec2-metadata-token-ttl-seconds", "60");
+        }
+        if (!token.isEmpty()) {
+            connection.setRequestProperty("X-aws-ec2-metadata-token", token);
+        }
+
+        final int responseCode;
+        try {
+            responseCode = connection.getResponseCode();
+        } catch (Exception e) {
+            logger.warn("Error connecting to IDMS.", e);
+            return "";
+        }
+
+        if (responseCode != 200) {
+            logger.warn("Error reponse from IDMS: code (" + responseCode + ") text " + readResponseString(connection));
+        }
+
+        return readResponseString(connection).trim();
+    }
+
+    private static String readResponseString(HttpURLConnection connection) {
+        ByteArrayOutputStream os = new ByteArrayOutputStream();
+        try (InputStream is = connection.getInputStream()) {
+            readTo(is, os);
+        } catch (IOException e) {
+            // Only best effort read if we can.
+        }
+        try (InputStream is = connection.getErrorStream()) {
+            readTo(is, os);
+        } catch (IOException e) {
+            // Only best effort read if we can.
+        }
+        try {
+            return os.toString(StandardCharsets.UTF_8.name());
+        } catch (UnsupportedEncodingException e) {
+            throw new IllegalStateException("UTF-8 not supported can't happen.");
+        }
+    }
+
+    private static void readTo(@Nullable InputStream is, ByteArrayOutputStream os) throws IOException {
+        if (is == null) {
+            return;
+        }
+        int b;
+        while ((b = is.read()) != -1) {
+            os.write(b);
+        }
+    }
+
+}
diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java
index ec556356..d4a8553d 100644
--- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java
+++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java
@@ -6,6 +6,7 @@
 import java.nio.file.Path;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -28,8 +29,6 @@
 public class EC2Plugin implements Plugin {
     private static final Log logger = LogFactory.getLog(EC2Plugin.class);
 
-    private static FileSystem fs;
-
     private static final String SERVICE_NAME = "ec2";
     public static final String ORIGIN = "AWS::EC2::Instance";
 
@@ -42,23 +41,28 @@ public class EC2Plugin implements Plugin {
     private static final String LINUX_ROOT = "/";
     private static final String LINUX_PATH = "opt/aws/amazon-cloudwatch-agent/etc/log-config.json";
 
-    private HashMap<String, Object> runtimeContext;
+    private final Map<String, Object> runtimeContext;
+
+    private final Set<AWSLogReference> logReferences;
 
-    private Set<AWSLogReference> logReferences;
+    private final FileSystem fs;
+
+    private final Map<EC2MetadataFetcher.EC2Metadata, String> metadata;
 
     public EC2Plugin() {
-        this(FileSystems.getDefault());
+        this(FileSystems.getDefault(), new EC2MetadataFetcher());
     }
 
-    public EC2Plugin(FileSystem fs) {
-        runtimeContext = new HashMap<>();
-        logReferences = new HashSet<>();
+    public EC2Plugin(FileSystem fs, EC2MetadataFetcher metadataFetcher) {
         this.fs = fs;
+        metadata = metadataFetcher.fetch();
+        runtimeContext = new LinkedHashMap<>();
+        logReferences = new HashSet<>();
     }
 
     @Override
     public boolean isEnabled() {
-        return EC2MetadataUtils.getInstanceId() != null;
+        return metadata.containsKey(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID);
     }
 
     @Override
@@ -70,12 +74,8 @@ public String getServiceName() {
      * Reads EC2 provided metadata to include it in trace document
      */
     public void populateRuntimeContext() {
-        if (null != EC2MetadataUtils.getInstanceId()) {
-            runtimeContext.put("instance_id", EC2MetadataUtils.getInstanceId());
-        }
-        if (null != EC2MetadataUtils.getAvailabilityZone()) {
-            runtimeContext.put("availability_zone", EC2MetadataUtils.getAvailabilityZone());
-        }
+        runtimeContext.put("instance_id", metadata.get(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID));
+        runtimeContext.put("availability_zone", metadata.get(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE));
     }
 
     public Map<String, Object> getRuntimeContext() {
diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
new file mode 100644
index 00000000..8d51e4e4
--- /dev/null
+++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
@@ -0,0 +1,74 @@
+package com.amazonaws.xray.plugins;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.any;
+import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl;
+import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
+import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.notFound;
+import static com.github.tomakehurst.wiremock.client.WireMock.ok;
+import static com.github.tomakehurst.wiremock.client.WireMock.putRequestedFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
+import static com.github.tomakehurst.wiremock.client.WireMock.verify;
+import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.entry;
+
+import com.github.tomakehurst.wiremock.junit.WireMockClassRule;
+import java.util.Map;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Test;
+
+public class EC2MetadataFetcherTest {
+
+    @ClassRule
+    public static WireMockClassRule server = new WireMockClassRule(wireMockConfig().dynamicPort());
+
+    private EC2MetadataFetcher fetcher;
+
+    @Before
+    public void setUp() {
+        fetcher = new EC2MetadataFetcher("localhost:" + server.port());
+    }
+
+    @Test
+    public void idmsv2() {
+        stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(ok("token")));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));
+
+        Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
+        assertThat(metadata).containsExactly(
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-123"),
+            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"));
+
+        verify(putRequestedFor(urlEqualTo("/latest/api/token"))
+                   .withHeader("X-aws-ec2-metadata-token-ttl-seconds", equalTo("60")));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-id"))
+                   .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/availability-zone"))
+                   .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
+    }
+
+    @Test
+    public void idmsv1() {
+        stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(notFound()));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));
+
+        Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
+        assertThat(metadata).containsExactly(
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-123"),
+            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"));
+
+        verify(putRequestedFor(urlEqualTo("/latest/api/token"))
+                   .withHeader("X-aws-ec2-metadata-token-ttl-seconds", equalTo("60")));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-id"))
+                   .withoutHeader("X-aws-ec2-metadata-token"));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/availability-zone"))
+                   .withoutHeader("X-aws-ec2-metadata-token"));
+    }
+}
diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java
index 190c4cec..8154d0e0 100644
--- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java
+++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java
@@ -1,14 +1,24 @@
 package com.amazonaws.xray.plugins;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.when;
+
 import com.amazonaws.util.EC2MetadataUtils;
 import com.amazonaws.xray.entities.AWSLogReference;
 import com.amazonaws.xray.utils.JsonUtils;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.BDDMockito;
+import org.mockito.Mock;
 import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 import org.powermock.api.mockito.PowerMockito;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
@@ -22,23 +32,46 @@
 import java.util.Set;
 
 @RunWith(PowerMockRunner.class)
-@PrepareForTest({JsonUtils.class, EC2MetadataUtils.class})
+@PrepareForTest({JsonUtils.class})
 public class EC2PluginTest {
+
+    @Rule
+    public MockitoRule mocks = MockitoJUnit.rule();
+
+    @Mock
+    private FileSystem fakeFs;
+
+    @Mock
+    private EC2MetadataFetcher metadataFetcher;
+
     private EC2Plugin ec2Plugin;
-    private FileSystem fakeFs = Mockito.mock(FileSystem.class);
 
     @Before
     public void setUpEC2Plugin() {
+        Map<EC2MetadataFetcher.EC2Metadata, String> metadata = new HashMap<>();
+        metadata.put(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-1234");
+        metadata.put(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "ap-northeast-1a");
         PowerMockito.mockStatic(JsonUtils.class);
-        PowerMockito.mockStatic(EC2MetadataUtils.class);
-        ec2Plugin = new EC2Plugin(fakeFs);
+        when(metadataFetcher.fetch()).thenReturn(metadata);
+        ec2Plugin = new EC2Plugin(fakeFs, metadataFetcher);
+    }
+
+    @Test
+    public void testMetadataPresent() {
+        assertThat(ec2Plugin.isEnabled()).isTrue();
+
+        ec2Plugin.populateRuntimeContext();
+        assertThat(ec2Plugin.getRuntimeContext())
+            .containsEntry("instance_id", "instance-1234")
+            .containsEntry("availability_zone", "ap-northeast-1a");
     }
 
     @Test
-    public void testInit() {
-        BDDMockito.given(EC2MetadataUtils.getInstanceId()).willReturn("12345");
+    public void testMetadataNotPresent() {
+        when(metadataFetcher.fetch()).thenReturn(Collections.emptyMap());
+        ec2Plugin = new EC2Plugin(fakeFs, metadataFetcher);
 
-        Assert.assertTrue(ec2Plugin.isEnabled());
+        assertThat(ec2Plugin.isEnabled()).isFalse();
     }
 
     @Test
@@ -49,7 +82,7 @@ public void testFilePathCreationFailure() {
 
         Set<AWSLogReference> logReferences = ec2Plugin.getLogReferences();
 
-        Assert.assertTrue(logReferences.isEmpty());
+        assertThat(logReferences).isEmpty();
     }
 
     @Test
@@ -64,10 +97,9 @@ public void testGenerationOfLogReference() throws IOException {
         BDDMockito.given(JsonUtils.getMatchingListFromJsonArrayNode(Mockito.any(), Mockito.any())).willReturn(groupList);
 
         Set<AWSLogReference> logReferences = ec2Plugin.getLogReferences();
-        AWSLogReference logReference = (AWSLogReference) logReferences.toArray()[0];
 
-        Assert.assertEquals(1, logReferences.size());
-        Assert.assertEquals("test_group", logReference.getLogGroup());
+        assertThat(logReferences).hasOnlyOneElementSatisfying(
+            reference -> assertThat(reference.getLogGroup()).isEqualTo("test_group"));
     }
 
     @Test
@@ -84,6 +116,6 @@ public void testGenerationOfMultipleLogReferences() throws IOException {
 
         Set<AWSLogReference> logReferences = ec2Plugin.getLogReferences();
 
-        Assert.assertEquals(2, logReferences.size());
+        assertThat(logReferences).hasSize(2);
     }
 }

From b8df2fdfe253a3912bd40795dc348ce03fa4b424 Mon Sep 17 00:00:00 2001
From: Anuraag Agrawal <aanuraag@amazon.co.jp>
Date: Fri, 29 May 2020 10:16:16 +0900
Subject: [PATCH 2/4] More metadata

---
 .../xray/plugins/EC2MetadataFetcher.java      | 18 +++++++++++++++
 .../com/amazonaws/xray/plugins/EC2Plugin.java | 14 +++++------
 .../xray/plugins/EC2MetadataFetcherTest.java  | 22 ++++++++++++++----
 .../amazonaws/xray/plugins/EC2PluginTest.java | 23 ++++++++++---------
 4 files changed, 54 insertions(+), 23 deletions(-)

diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
index 0e47fbd6..80b8ef95 100644
--- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
+++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
@@ -22,6 +22,8 @@ class EC2MetadataFetcher {
     enum EC2Metadata {
         INSTANCE_ID,
         AVAILABILITY_ZONE,
+        INSTANCE_TYPE,
+        AMI_ID,
     }
 
     private static final int TIMEOUT_MILLIS = 2000;
@@ -30,6 +32,8 @@ enum EC2Metadata {
     private final URL tokenUrl;
     private final URL instanceIdUrl;
     private final URL availabilityZoneUrl;
+    private final URL instanceTypeUrl;
+    private final URL amiIdUrl;
 
     EC2MetadataFetcher() {
         this(System.getenv("IMDS_ENDPOINT") != null ? System.getenv("IMDS_ENDPOINT") : DEFAULT_IDMS_ENDPOINT);
@@ -41,6 +45,8 @@ enum EC2Metadata {
             this.tokenUrl = new URL(urlBase + "/latest/api/token");
             this.instanceIdUrl = new URL(urlBase + "/latest/meta-data/instance-id");
             this.availabilityZoneUrl = new URL(urlBase + "/latest/meta-data/placement/availability-zone");
+            this.instanceTypeUrl = new URL(urlBase + "/latest/meta-data/instance-type");
+            this.amiIdUrl = new URL(urlBase + "/latest/meta-data/ami-id");
         } catch (MalformedURLException e) {
             throw new IllegalArgumentException("Illegal endpoint: " + endpoint);
         }
@@ -58,10 +64,14 @@ Map<EC2Metadata, String> fetch() {
         }
 
         String availabilityZone = fetchAvailabilityZone(token);
+        String instanceType = fetchInstanceType(token);
+        String amiId = fetchAmiId(token);
 
         Map<EC2Metadata, String> result = new LinkedHashMap<>();
         result.put(EC2Metadata.INSTANCE_ID, instanceId);
         result.put(EC2Metadata.AVAILABILITY_ZONE, availabilityZone);
+        result.put(EC2Metadata.INSTANCE_TYPE, instanceType);
+        result.put(EC2Metadata.AMI_ID, amiId);
         return result;
     }
 
@@ -77,6 +87,14 @@ private String fetchAvailabilityZone(String token) {
         return fetchString("GET", availabilityZoneUrl, token, false);
     }
 
+    private String fetchInstanceType(String token) {
+        return fetchString("GET", instanceTypeUrl, token, false);
+    }
+
+    private String fetchAmiId(String token) {
+        return fetchString("GET", amiIdUrl, token, false);
+    }
+
     // Generic HTTP fetch function for IDMS.
     private static String fetchString(String httpMethod, URL url, String token, boolean includeTtl) {
         final HttpURLConnection connection;
diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java
index d4a8553d..b2d71ec7 100644
--- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java
+++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java
@@ -1,25 +1,21 @@
 package com.amazonaws.xray.plugins;
 
+import com.amazonaws.xray.entities.AWSLogReference;
+import com.amazonaws.xray.entities.StringValidator;
+import com.amazonaws.xray.utils.JsonUtils;
+import com.fasterxml.jackson.databind.JsonNode;
 import java.io.IOException;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
 import java.nio.file.Path;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-
-import com.amazonaws.xray.entities.AWSLogReference;
-import com.amazonaws.xray.entities.StringValidator;
-import com.amazonaws.xray.utils.JsonUtils;
-import com.fasterxml.jackson.databind.JsonNode;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
-import com.amazonaws.util.EC2MetadataUtils;
-
 /**
  * A plugin, for use with the {@code AWSXRayRecorderBuilder} class, which will add EC2 instance information to segments generated by the built {@code AWSXRayRecorder} instance.
  * 
@@ -76,6 +72,8 @@ public String getServiceName() {
     public void populateRuntimeContext() {
         runtimeContext.put("instance_id", metadata.get(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID));
         runtimeContext.put("availability_zone", metadata.get(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE));
+        runtimeContext.put("instance_size", metadata.get(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE));
+        runtimeContext.put("ami_id", metadata.get(EC2MetadataFetcher.EC2Metadata.AMI_ID));
     }
 
     public Map<String, Object> getRuntimeContext() {
diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
index 8d51e4e4..f4babceb 100644
--- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
+++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
@@ -1,8 +1,6 @@
 package com.amazonaws.xray.plugins;
 
-import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
 import static com.github.tomakehurst.wiremock.client.WireMock.any;
-import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl;
 import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
 import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.notFound;
@@ -39,11 +37,15 @@ public void idmsv2() {
         stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(ok("token")));
         stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
         stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-type")).willReturn(ok("m4.xlarge")));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/ami-id")).willReturn(ok("ami-1234")));
 
         Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
         assertThat(metadata).containsExactly(
             entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-123"),
-            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"));
+            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"),
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE, "m4.xlarge"),
+            entry(EC2MetadataFetcher.EC2Metadata.AMI_ID, "ami-1234"));
 
         verify(putRequestedFor(urlEqualTo("/latest/api/token"))
                    .withHeader("X-aws-ec2-metadata-token-ttl-seconds", equalTo("60")));
@@ -51,6 +53,10 @@ public void idmsv2() {
                    .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
         verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/availability-zone"))
                    .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-type"))
+                   .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/ami-id"))
+                   .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
     }
 
     @Test
@@ -58,11 +64,15 @@ public void idmsv1() {
         stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(notFound()));
         stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
         stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-type")).willReturn(ok("m4.xlarge")));
+        stubFor(any(urlPathEqualTo("/latest/meta-data/ami-id")).willReturn(ok("ami-1234")));
 
         Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
         assertThat(metadata).containsExactly(
             entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-123"),
-            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"));
+            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"),
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE, "m4.xlarge"),
+            entry(EC2MetadataFetcher.EC2Metadata.AMI_ID, "ami-1234"));
 
         verify(putRequestedFor(urlEqualTo("/latest/api/token"))
                    .withHeader("X-aws-ec2-metadata-token-ttl-seconds", equalTo("60")));
@@ -70,5 +80,9 @@ public void idmsv1() {
                    .withoutHeader("X-aws-ec2-metadata-token"));
         verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/availability-zone"))
                    .withoutHeader("X-aws-ec2-metadata-token"));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-type"))
+                   .withoutHeader("X-aws-ec2-metadata-token"));
+        verify(getRequestedFor(urlEqualTo("/latest/meta-data/ami-id"))
+                   .withoutHeader("X-aws-ec2-metadata-token"));
     }
 }
diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java
index 8154d0e0..4281f290 100644
--- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java
+++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2PluginTest.java
@@ -3,13 +3,18 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.when;
 
-import com.amazonaws.util.EC2MetadataUtils;
 import com.amazonaws.xray.entities.AWSLogReference;
 import com.amazonaws.xray.utils.JsonUtils;
+import java.io.IOException;
+import java.nio.file.FileSystem;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
-import org.junit.Assert;
+import java.util.Set;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -23,14 +28,6 @@
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
 
-import java.io.IOException;
-import java.nio.file.FileSystem;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Set;
-
 @RunWith(PowerMockRunner.class)
 @PrepareForTest({JsonUtils.class})
 public class EC2PluginTest {
@@ -51,6 +48,8 @@ public void setUpEC2Plugin() {
         Map<EC2MetadataFetcher.EC2Metadata, String> metadata = new HashMap<>();
         metadata.put(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-1234");
         metadata.put(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "ap-northeast-1a");
+        metadata.put(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE, "m4.xlarge");
+        metadata.put(EC2MetadataFetcher.EC2Metadata.AMI_ID, "ami-1234");
         PowerMockito.mockStatic(JsonUtils.class);
         when(metadataFetcher.fetch()).thenReturn(metadata);
         ec2Plugin = new EC2Plugin(fakeFs, metadataFetcher);
@@ -63,7 +62,9 @@ public void testMetadataPresent() {
         ec2Plugin.populateRuntimeContext();
         assertThat(ec2Plugin.getRuntimeContext())
             .containsEntry("instance_id", "instance-1234")
-            .containsEntry("availability_zone", "ap-northeast-1a");
+            .containsEntry("availability_zone", "ap-northeast-1a")
+            .containsEntry("instance_size", "m4.xlarge")
+            .containsEntry("ami_id", "ami-1234");
     }
 
     @Test

From 9488444b8332dbf225a423ff7adf6ae123af8543 Mon Sep 17 00:00:00 2001
From: Anuraag Agrawal <aanuraag@amazon.co.jp>
Date: Fri, 29 May 2020 10:23:10 +0900
Subject: [PATCH 3/4] idms -> imds

---
 .../xray/plugins/EC2MetadataFetcher.java         | 16 ++++++++--------
 .../xray/plugins/EC2MetadataFetcherTest.java     |  4 ++--
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
index 80b8ef95..bd0164a3 100644
--- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
+++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
@@ -27,7 +27,7 @@ enum EC2Metadata {
     }
 
     private static final int TIMEOUT_MILLIS = 2000;
-    private static final String DEFAULT_IDMS_ENDPOINT = "169.254.169.254";
+    private static final String DEFAULT_IMDS_ENDPOINT = "169.254.169.254";
 
     private final URL tokenUrl;
     private final URL instanceIdUrl;
@@ -36,7 +36,7 @@ enum EC2Metadata {
     private final URL amiIdUrl;
 
     EC2MetadataFetcher() {
-        this(System.getenv("IMDS_ENDPOINT") != null ? System.getenv("IMDS_ENDPOINT") : DEFAULT_IDMS_ENDPOINT);
+        this(System.getenv("IMDS_ENDPOINT") != null ? System.getenv("IMDS_ENDPOINT") : DEFAULT_IMDS_ENDPOINT);
     }
 
     EC2MetadataFetcher(String endpoint) {
@@ -55,8 +55,8 @@ enum EC2Metadata {
     Map<EC2Metadata, String> fetch() {
         String token = fetchToken();
 
-        // If token is empty, either IDMSv2 isn't enabled or an unexpected failure happened. We can still get
-        // data if IDMSv1 is enabled.
+        // If token is empty, either IMDSv2 isn't enabled or an unexpected failure happened. We can still get
+        // data if IMDSv1 is enabled.
         String instanceId = fetchInstanceId(token);
         if (instanceId.isEmpty()) {
             // If no instance ID, assume we are not actually running on EC2.
@@ -95,13 +95,13 @@ private String fetchAmiId(String token) {
         return fetchString("GET", amiIdUrl, token, false);
     }
 
-    // Generic HTTP fetch function for IDMS.
+    // Generic HTTP fetch function for IMDS.
     private static String fetchString(String httpMethod, URL url, String token, boolean includeTtl) {
         final HttpURLConnection connection;
         try {
             connection = (HttpURLConnection) url.openConnection();
         } catch (Exception e) {
-            logger.warn("Error connecting to IDMS.", e);
+            logger.warn("Error connecting to IMDS.", e);
             return "";
         }
 
@@ -126,12 +126,12 @@ private static String fetchString(String httpMethod, URL url, String token, bool
         try {
             responseCode = connection.getResponseCode();
         } catch (Exception e) {
-            logger.warn("Error connecting to IDMS.", e);
+            logger.warn("Error connecting to IMDS.", e);
             return "";
         }
 
         if (responseCode != 200) {
-            logger.warn("Error reponse from IDMS: code (" + responseCode + ") text " + readResponseString(connection));
+            logger.warn("Error reponse from IMDS: code (" + responseCode + ") text " + readResponseString(connection));
         }
 
         return readResponseString(connection).trim();
diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
index f4babceb..ec6cab9f 100644
--- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
+++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
@@ -33,7 +33,7 @@ public void setUp() {
     }
 
     @Test
-    public void idmsv2() {
+    public void imdsv2() {
         stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(ok("token")));
         stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
         stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));
@@ -60,7 +60,7 @@ public void idmsv2() {
     }
 
     @Test
-    public void idmsv1() {
+    public void imdsv1() {
         stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(notFound()));
         stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
         stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));

From dfa6fa2331ca908458fbb14153121c6b7822c361 Mon Sep 17 00:00:00 2001
From: Anuraag Agrawal <aanuraag@amazon.co.jp>
Date: Fri, 29 May 2020 15:22:34 +0900
Subject: [PATCH 4/4] Use identity document instead.

---
 .../xray/plugins/EC2MetadataFetcher.java      | 83 +++++++++++--------
 .../xray/plugins/EC2MetadataFetcherTest.java  | 82 +++++++++++-------
 2 files changed, 100 insertions(+), 65 deletions(-)

diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
index bd0164a3..b526002f 100644
--- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
+++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java
@@ -1,5 +1,8 @@
 package com.amazonaws.xray.plugins;
 
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.JsonToken;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -10,7 +13,7 @@
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
-import java.util.LinkedHashMap;
+import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
 import org.apache.commons.logging.Log;
@@ -19,6 +22,8 @@
 class EC2MetadataFetcher {
     private static final Log logger = LogFactory.getLog(EC2MetadataFetcher.class);
 
+    private static final JsonFactory JSON_FACTORY = new JsonFactory();
+
     enum EC2Metadata {
         INSTANCE_ID,
         AVAILABILITY_ZONE,
@@ -29,11 +34,8 @@ enum EC2Metadata {
     private static final int TIMEOUT_MILLIS = 2000;
     private static final String DEFAULT_IMDS_ENDPOINT = "169.254.169.254";
 
+    private final URL identityDocumentUrl;
     private final URL tokenUrl;
-    private final URL instanceIdUrl;
-    private final URL availabilityZoneUrl;
-    private final URL instanceTypeUrl;
-    private final URL amiIdUrl;
 
     EC2MetadataFetcher() {
         this(System.getenv("IMDS_ENDPOINT") != null ? System.getenv("IMDS_ENDPOINT") : DEFAULT_IMDS_ENDPOINT);
@@ -42,11 +44,8 @@ enum EC2Metadata {
     EC2MetadataFetcher(String endpoint) {
         String urlBase = "http://" + endpoint;
         try {
+            this.identityDocumentUrl = new URL(urlBase + "/latest/dynamic/instance-identity/document");
             this.tokenUrl = new URL(urlBase + "/latest/api/token");
-            this.instanceIdUrl = new URL(urlBase + "/latest/meta-data/instance-id");
-            this.availabilityZoneUrl = new URL(urlBase + "/latest/meta-data/placement/availability-zone");
-            this.instanceTypeUrl = new URL(urlBase + "/latest/meta-data/instance-type");
-            this.amiIdUrl = new URL(urlBase + "/latest/meta-data/ami-id");
         } catch (MalformedURLException e) {
             throw new IllegalArgumentException("Illegal endpoint: " + endpoint);
         }
@@ -57,21 +56,49 @@ Map<EC2Metadata, String> fetch() {
 
         // If token is empty, either IMDSv2 isn't enabled or an unexpected failure happened. We can still get
         // data if IMDSv1 is enabled.
-        String instanceId = fetchInstanceId(token);
-        if (instanceId.isEmpty()) {
-            // If no instance ID, assume we are not actually running on EC2.
+        String identity = fetchIdentity(token);
+        if (identity.isEmpty()) {
+            // If no identity document, assume we are not actually running on EC2.
             return Collections.emptyMap();
         }
 
-        String availabilityZone = fetchAvailabilityZone(token);
-        String instanceType = fetchInstanceType(token);
-        String amiId = fetchAmiId(token);
+        Map<EC2Metadata, String> result = new HashMap<>();
+        try (JsonParser parser = JSON_FACTORY.createParser(identity)) {
+            parser.nextToken();
+
+            if (!parser.isExpectedStartObjectToken()) {
+                throw new IOException("Invalid JSON:" + identity);
+            }
+
+            while (parser.nextToken() != JsonToken.END_OBJECT) {
+                String value = parser.nextTextValue();
+                switch (parser.getCurrentName()) {
+                    case "instanceId":
+                        result.put(EC2Metadata.INSTANCE_ID, value);
+                        break;
+                    case "availabilityZone":
+                        result.put(EC2Metadata.AVAILABILITY_ZONE, value);
+                        break;
+                    case "instanceType":
+                        result.put(EC2Metadata.INSTANCE_TYPE, value);
+                        break;
+                    case "imageId":
+                        result.put(EC2Metadata.AMI_ID, value);
+                        break;
+                    default:
+                        parser.skipChildren();
+                }
+                if (result.size() == EC2Metadata.values().length) {
+                    return result;
+                }
+            }
+        } catch (IOException e) {
+            logger.warn("Could not parse identity document.", e);
+            return Collections.emptyMap();
+        }
 
-        Map<EC2Metadata, String> result = new LinkedHashMap<>();
-        result.put(EC2Metadata.INSTANCE_ID, instanceId);
-        result.put(EC2Metadata.AVAILABILITY_ZONE, availabilityZone);
-        result.put(EC2Metadata.INSTANCE_TYPE, instanceType);
-        result.put(EC2Metadata.AMI_ID, amiId);
+        // Getting here means the document didn't have all the metadata fields we wanted.
+        logger.warn("Identity document missing metadata: " + identity);
         return result;
     }
 
@@ -79,20 +106,8 @@ private String fetchToken() {
         return fetchString("PUT", tokenUrl, "", true);
     }
 
-    private String fetchInstanceId(String token) {
-        return fetchString("GET", instanceIdUrl, token, false);
-    }
-
-    private String fetchAvailabilityZone(String token) {
-        return fetchString("GET", availabilityZoneUrl, token, false);
-    }
-
-    private String fetchInstanceType(String token) {
-        return fetchString("GET", instanceTypeUrl, token, false);
-    }
-
-    private String fetchAmiId(String token) {
-        return fetchString("GET", amiIdUrl, token, false);
+    private String fetchIdentity(String token) {
+        return fetchString("GET", identityDocumentUrl, token, false);
     }
 
     // Generic HTTP fetch function for IMDS.
diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
index ec6cab9f..290ee49e 100644
--- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
+++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java
@@ -5,6 +5,7 @@
 import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.notFound;
 import static com.github.tomakehurst.wiremock.client.WireMock.ok;
+import static com.github.tomakehurst.wiremock.client.WireMock.okJson;
 import static com.github.tomakehurst.wiremock.client.WireMock.putRequestedFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
@@ -22,6 +23,26 @@
 
 public class EC2MetadataFetcherTest {
 
+    // From https://docs.amazonaws.cn/en_us/AWSEC2/latest/UserGuide/instance-identity-documents.html
+    private static final String IDENTITY_DOCUMENT =
+        "{\n"
+        + "    \"devpayProductCodes\" : null,\n"
+        + "    \"marketplaceProductCodes\" : [ \"1abc2defghijklm3nopqrs4tu\" ], \n"
+        + "    \"availabilityZone\" : \"us-west-2b\",\n"
+        + "    \"privateIp\" : \"10.158.112.84\",\n"
+        + "    \"version\" : \"2017-09-30\",\n"
+        + "    \"instanceId\" : \"i-1234567890abcdef0\",\n"
+        + "    \"billingProducts\" : null,\n"
+        + "    \"instanceType\" : \"t2.micro\",\n"
+        + "    \"accountId\" : \"123456789012\",\n"
+        + "    \"imageId\" : \"ami-5fb8c835\",\n"
+        + "    \"pendingTime\" : \"2016-11-19T16:32:11Z\",\n"
+        + "    \"architecture\" : \"x86_64\",\n"
+        + "    \"kernelId\" : null,\n"
+        + "    \"ramdiskId\" : null,\n"
+        + "    \"region\" : \"us-west-2\"\n"
+        + "}";
+
     @ClassRule
     public static WireMockClassRule server = new WireMockClassRule(wireMockConfig().dynamicPort());
 
@@ -35,54 +56,53 @@ public void setUp() {
     @Test
     public void imdsv2() {
         stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(ok("token")));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-type")).willReturn(ok("m4.xlarge")));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/ami-id")).willReturn(ok("ami-1234")));
+        stubFor(any(urlPathEqualTo("/latest/dynamic/instance-identity/document"))
+                    .willReturn(okJson(IDENTITY_DOCUMENT)));
 
         Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
-        assertThat(metadata).containsExactly(
-            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-123"),
-            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"),
-            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE, "m4.xlarge"),
-            entry(EC2MetadataFetcher.EC2Metadata.AMI_ID, "ami-1234"));
+        assertThat(metadata).containsOnly(
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "i-1234567890abcdef0"),
+            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "us-west-2b"),
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE, "t2.micro"),
+            entry(EC2MetadataFetcher.EC2Metadata.AMI_ID, "ami-5fb8c835"));
 
         verify(putRequestedFor(urlEqualTo("/latest/api/token"))
                    .withHeader("X-aws-ec2-metadata-token-ttl-seconds", equalTo("60")));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-id"))
-                   .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/availability-zone"))
-                   .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-type"))
-                   .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/ami-id"))
+        verify(getRequestedFor(urlEqualTo("/latest/dynamic/instance-identity/document"))
                    .withHeader("X-aws-ec2-metadata-token", equalTo("token")));
     }
 
     @Test
     public void imdsv1() {
         stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(notFound()));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-id")).willReturn(ok("instance-123")));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/placement/availability-zone")).willReturn(ok("asia-northeast-1a")));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/instance-type")).willReturn(ok("m4.xlarge")));
-        stubFor(any(urlPathEqualTo("/latest/meta-data/ami-id")).willReturn(ok("ami-1234")));
+        stubFor(any(urlPathEqualTo("/latest/dynamic/instance-identity/document"))
+                    .willReturn(okJson(IDENTITY_DOCUMENT)));
 
         Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
-        assertThat(metadata).containsExactly(
-            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "instance-123"),
-            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "asia-northeast-1a"),
-            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE, "m4.xlarge"),
-            entry(EC2MetadataFetcher.EC2Metadata.AMI_ID, "ami-1234"));
+        assertThat(metadata).containsOnly(
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_ID, "i-1234567890abcdef0"),
+            entry(EC2MetadataFetcher.EC2Metadata.AVAILABILITY_ZONE, "us-west-2b"),
+            entry(EC2MetadataFetcher.EC2Metadata.INSTANCE_TYPE, "t2.micro"),
+            entry(EC2MetadataFetcher.EC2Metadata.AMI_ID, "ami-5fb8c835"));
 
         verify(putRequestedFor(urlEqualTo("/latest/api/token"))
                    .withHeader("X-aws-ec2-metadata-token-ttl-seconds", equalTo("60")));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-id"))
+        verify(getRequestedFor(urlEqualTo("/latest/dynamic/instance-identity/document"))
                    .withoutHeader("X-aws-ec2-metadata-token"));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/placement/availability-zone"))
-                   .withoutHeader("X-aws-ec2-metadata-token"));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/instance-type"))
-                   .withoutHeader("X-aws-ec2-metadata-token"));
-        verify(getRequestedFor(urlEqualTo("/latest/meta-data/ami-id"))
+    }
+
+    @Test
+    public void badJson() {
+        stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(notFound()));
+        stubFor(any(urlPathEqualTo("/latest/dynamic/instance-identity/document"))
+                    .willReturn(okJson("I'm not JSON")));
+
+        Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
+        assertThat(metadata).isEmpty();
+
+        verify(putRequestedFor(urlEqualTo("/latest/api/token"))
+                   .withHeader("X-aws-ec2-metadata-token-ttl-seconds", equalTo("60")));
+        verify(getRequestedFor(urlEqualTo("/latest/dynamic/instance-identity/document"))
                    .withoutHeader("X-aws-ec2-metadata-token"));
     }
 }