Migrate EC2 metadata lookup to JDK requests and support IMDSv2 (#155)

* Migrate EC2 metadata lookup to JDK requests and support IMDSv2

* More metadata

* idms -> imds

* Use identity document instead.

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java

@@ -0,0 +1,184 @@
+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;
+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.HashMap;
+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);
+
+    private static final JsonFactory JSON_FACTORY = new JsonFactory();
+
+    enum EC2Metadata {
+        INSTANCE_ID,
+        AVAILABILITY_ZONE,
+        INSTANCE_TYPE,
+        AMI_ID,
+    }
+
+    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;
+
+    EC2MetadataFetcher() {
+        this(System.getenv("IMDS_ENDPOINT") != null ? System.getenv("IMDS_ENDPOINT") : DEFAULT_IMDS_ENDPOINT);
+    }
+
+    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");
+        } catch (MalformedURLException e) {
+            throw new IllegalArgumentException("Illegal endpoint: " + endpoint);
+        }
+    }
+
+    Map<EC2Metadata, String> fetch() {
+        String token = fetchToken();
+
+        // If token is empty, either IMDSv2 isn't enabled or an unexpected failure happened. We can still get
+        // data if IMDSv1 is enabled.
+        String identity = fetchIdentity(token);
+        if (identity.isEmpty()) {
+            // If no identity document, assume we are not actually running on EC2.
+            return Collections.emptyMap();
+        }
+
+        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();
+        }
+
+        // Getting here means the document didn't have all the metadata fields we wanted.
+        logger.warn("Identity document missing metadata: " + identity);
+        return result;
+    }
+
+    private String fetchToken() {
+        return fetchString("PUT", tokenUrl, "", true);
+    }
+
+    private String fetchIdentity(String token) {
+        return fetchString("GET", identityDocumentUrl, token, false);
+    }
+
+    // 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 IMDS.", 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 IMDS.", e);
+            return "";
+        }
+
+        if (responseCode != 200) {
+            logger.warn("Error reponse from IMDS: 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);
+        }
+    }
+
+}

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2Plugin.java

@@ -1,24 +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.
  * 
@@ -28,8 +25,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 +37,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 final FileSystem fs;
 
-    private Set<AWSLogReference> logReferences;
+    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 +70,10 @@ 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));
+        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() {

aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/plugins/EC2MetadataFetcherTest.java

@@ -0,0 +1,108 @@
+package com.amazonaws.xray.plugins;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.any;
+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.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;
+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 {
+
+    // 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());
+
+    private EC2MetadataFetcher fetcher;
+
+    @Before
+    public void setUp() {
+        fetcher = new EC2MetadataFetcher("localhost:" + server.port());
+    }
+
+    @Test
+    public void imdsv2() {
+        stubFor(any(urlPathEqualTo("/latest/api/token")).willReturn(ok("token")));
+        stubFor(any(urlPathEqualTo("/latest/dynamic/instance-identity/document"))
+                    .willReturn(okJson(IDENTITY_DOCUMENT)));
+
+        Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
+        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/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/dynamic/instance-identity/document"))
+                    .willReturn(okJson(IDENTITY_DOCUMENT)));
+
+        Map<EC2MetadataFetcher.EC2Metadata, String> metadata = fetcher.fetch();
+        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/dynamic/instance-identity/document"))
+                   .withoutHeader("X-aws-ec2-metadata-token"));
+    }
+
+    @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"));
+    }
+}