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..b526002f --- /dev/null +++ b/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 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 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); + } + } + +} 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..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,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 runtimeContext; + private final Map runtimeContext; + + private final Set logReferences; + + private final FileSystem fs; - private Set logReferences; + private final Map 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 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..290ee49e --- /dev/null +++ b/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 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 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 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")); + } +} 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..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 @@ -1,44 +1,78 @@ package com.amazonaws.xray.plugins; -import com.amazonaws.util.EC2MetadataUtils; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.when; + import com.amazonaws.xray.entities.AWSLogReference; import com.amazonaws.xray.utils.JsonUtils; -import org.junit.Assert; +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 java.util.Set; 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; -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, 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 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); - 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") + .containsEntry("instance_size", "m4.xlarge") + .containsEntry("ami_id", "ami-1234"); } @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 +83,7 @@ public void testFilePathCreationFailure() { Set logReferences = ec2Plugin.getLogReferences(); - Assert.assertTrue(logReferences.isEmpty()); + assertThat(logReferences).isEmpty(); } @Test @@ -64,10 +98,9 @@ public void testGenerationOfLogReference() throws IOException { BDDMockito.given(JsonUtils.getMatchingListFromJsonArrayNode(Mockito.any(), Mockito.any())).willReturn(groupList); Set 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 +117,6 @@ public void testGenerationOfMultipleLogReferences() throws IOException { Set logReferences = ec2Plugin.getLogReferences(); - Assert.assertEquals(2, logReferences.size()); + assertThat(logReferences).hasSize(2); } }