Migrate from IMDSv1 to IMDSv2

src/Amazon.CloudWatch.EMF/Amazon.CloudWatch.EMF.csproj

@@ -2,7 +2,7 @@
 
     <PropertyGroup>
         <PackageId>Amazon.CloudWatch.EMF</PackageId>
-        <VersionPrefix>0.1.0</VersionPrefix>
+        <VersionPrefix>0.1.1</VersionPrefix>
         <Authors>Amazon Web Services</Authors>
         <Description>Amazon CloudWatch Embedded Metric Format Client Library</Description>
         <Language>en-US</Language>

src/Amazon.CloudWatch.EMF/Environment/EC2Environment.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using Amazon.CloudWatch.EMF.Config;
 using Amazon.CloudWatch.EMF.Model;
 using Microsoft.Extensions.Logging;
@@ -9,11 +10,18 @@ namespace Amazon.CloudWatch.EMF.Environment
 {
     public class EC2Environment : AgentBasedEnvironment
     {
+        // Documentation for configuring instance metadata can be found here:
+        // https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
         private const string INSTANCE_IDENTITY_URL = "http://169.254.169.254/latest/dynamic/instance-identity/document";
+        private const string TOKEN_URL = "http://169.254.169.254/latest/api/token";
+        private const string TOKEN_REQUEST_HEADER_KEY = "X-aws-ec2-metadata-token-ttl-seconds";
+        private const string TOKEN_REQUEST_HEADER_VALUE = "21600";
+        private const string METADATA_REQUEST_HEADER_KEY = "X-aws-ec2-metadata-token";
         private const string CFN_EC2_TYPE = "AWS::EC2::Instance";
 
         private readonly ILogger _logger;
         private readonly IResourceFetcher _resourceFetcher;
+        private string _token;
         private EC2Metadata _ec2Metadata;
 
         public EC2Environment(IConfiguration configuration, IResourceFetcher resourceFetcher)
@@ -32,10 +40,35 @@ public EC2Environment(IConfiguration configuration, IResourceFetcher resourceFet
 
         public override bool Probe()
         {
-            Uri uri = null;
+            Uri tokenUri = null;
+            var tokenRequestHeader = new Dictionary<string, string>();
+            tokenRequestHeader.Add(TOKEN_REQUEST_HEADER_KEY, TOKEN_REQUEST_HEADER_VALUE);
             try
             {
-                uri = new Uri(INSTANCE_IDENTITY_URL);
+                tokenUri = new Uri(TOKEN_URL);
+            }
+            catch (Exception)
+            {
+                _logger.LogDebug("Failed to construct url: " + TOKEN_URL);
+                return false;
+            }
+
+            try
+            {
+                _token = _resourceFetcher.FetchString(tokenUri, "PUT", tokenRequestHeader);
+            }
+            catch (EMFClientException ex)
+            {
+                _logger.LogDebug("Failed to get response from: " + tokenUri, ex);
+                return false;
+            }
+
+            Uri metadataUri = null;
+            var metadataRequestHeader = new Dictionary<string, string>();
+            metadataRequestHeader.Add(METADATA_REQUEST_HEADER_KEY, _token);
+            try
+            {
+                metadataUri = new Uri(INSTANCE_IDENTITY_URL);
             }
             catch (Exception)
             {
@@ -45,12 +78,12 @@ public override bool Probe()
 
             try
             {
-                _ec2Metadata = _resourceFetcher.Fetch<EC2Metadata>(uri);
+                _ec2Metadata = _resourceFetcher.FetchJson<EC2Metadata>(metadataUri, "GET", metadataRequestHeader);
                 return true;
             }
             catch (EMFClientException ex)
             {
-                _logger.LogDebug("Failed to get response from: " + uri, ex);
+                _logger.LogDebug("Failed to get response from: " + metadataUri, ex);
             }
 
             return false;

src/Amazon.CloudWatch.EMF/Environment/ECSEnvironment.cs

@@ -52,7 +52,7 @@ public override bool Probe()
             try
             {
                 var parsedUri = new Uri(uri);
-                _ecsMetadata = _resourceFetcher.Fetch<ECSMetadata>(parsedUri);
+                _ecsMetadata = _resourceFetcher.FetchJson<ECSMetadata>(parsedUri, "GET");
                 FormatImageName();
                 return true;
             }

src/Amazon.CloudWatch.EMF/Environment/IResourceFetcher.cs

@@ -1,9 +1,13 @@
 using System;
+using System.Collections.Generic;
+using System.Net.Http;
 
 namespace Amazon.CloudWatch.EMF.Environment
 {
     public interface IResourceFetcher
     {
-        public T Fetch<T>(Uri endpoint);
+        public T FetchJson<T>(Uri endpoint, string method, Dictionary<string, string> header = null);
+
+        public string FetchString(Uri endpoint, string method, Dictionary<string, string> header = null);
     }
 }

src/Amazon.CloudWatch.EMF/Environment/ResourceFetcher.cs

@@ -1,6 +1,8 @@
 using System;
-using System.IO;
+using System.Collections.Generic;
 using System.Net;
+using System.Net.Http;
+using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using Newtonsoft.Json;
@@ -25,33 +27,44 @@ public ResourceFetcher(ILoggerFactory loggerFactory)
         /// Fetch a json object from a given uri and deserialize it to the specified class: T
         /// </summary>
         /// <returns></returns>
-        public T Fetch<T>(Uri endpoint)
+        public T FetchJson<T>(Uri endpoint, string method, Dictionary<string, string> header = null)
         {
-            string response = ReadResource(endpoint, "GET");
+            string response = ReadResource(endpoint, method, header).Result;
 
             return JsonConvert.DeserializeObject<T>(response);
         }
 
-        private string ReadResource(Uri endpoint, string method)
+        /// <summary>
+        /// Fetch string from a given uri
+        /// </summary>
+        /// <returns></returns>
+        public string FetchString(Uri endpoint, string method, Dictionary<string, string> headers = null)
+        {
+            string response = ReadResource(endpoint, method, headers).Result;
+
+            return response;
+        }
+
+        private async Task<string> ReadResource(Uri endpoint, string method, Dictionary<string, string> headers)
         {
             try
             {
-                var httpWebRequest = GetHttpWebRequest(endpoint, method);
+                headers ??= new Dictionary<string, string>();
 
-                var httpWebResponse = (HttpWebResponse)httpWebRequest.GetResponse();
+                var httpResponse = GetResponse(endpoint, method, headers);
 
-                if (httpWebResponse.StatusCode == HttpStatusCode.OK)
-                {
-                    return GetResponse(httpWebResponse);
-                }
-                else if (httpWebResponse.StatusCode == HttpStatusCode.NotFound)
+                if (httpResponse.StatusCode == HttpStatusCode.OK)
                 {
-                    throw new EMFClientException("The requested metadata is not found at " + httpWebRequest.RequestUri.AbsolutePath);
+                    return await httpResponse.Content.ReadAsStringAsync();
                 }
-                else
+
+                if (httpResponse.StatusCode == HttpStatusCode.NotFound)
                 {
-                    HandleErrorResponse(httpWebResponse);
+                    throw new EMFClientException("The requested data is not found at " + endpoint.AbsolutePath);
                 }
+
+                throw new EMFClientException("Failed to get resource. Error code: " + httpResponse.StatusCode +
+                                             ", error message: " + httpResponse.ReasonPhrase);
             }
             catch (Exception e)
             {
@@ -61,52 +74,22 @@ private string ReadResource(Uri endpoint, string method)
                     + "\n Attempting to reconnect.");
                 throw new EMFClientException("Failed to connect to service endpoint: ", e);
             }
-
-            return string.Empty;
         }
 
-        private void HandleErrorResponse(HttpWebResponse httpWebResponse)
+        private HttpResponseMessage GetResponse(Uri endpoint, string method, Dictionary<string, string> headers)
         {
-            string errorResponse = GetResponse(httpWebResponse);
+            HttpClient client = new HttpClient();
 
-            try
-            {
-                /*JsonNode node = Jackson.jsonNodeOf(errorResponse);
-                JsonNode code = node.get("code");
-                JsonNode message = node.get("message");
-                if (code != null && message != null) {
-                    errorCode = code.asText();
-                    responseMessage = message.asText();
-                }
-
-                String exceptionMessage =
-                    String.format(
-                        "Failed to get resource. Error code: %s, error message: %s ",
-                        errorCode, responseMessage);
-                throw new EMFClientException(exceptionMessage);*/
-            }
-            catch (System.Exception exception)
+            var httpMethod = new HttpMethod(method.ToUpper());
+            HttpRequestMessage request = new HttpRequestMessage(httpMethod, endpoint);
+            foreach (KeyValuePair<string, string> header in headers)
             {
-                throw new EMFClientException("Unable to parse error stream: ", exception);
+                request.Headers.Add(header.Key, header.Value);
             }
-        }
-
-        private HttpWebRequest GetHttpWebRequest(Uri endpoint, string method)
-        {
-            var httpWebRequest = (HttpWebRequest)WebRequest.CreateHttp(endpoint);
-            httpWebRequest.Method = method;
-            httpWebRequest.Timeout = 1000;
-            httpWebRequest.ReadWriteTimeout = 1000;
-            return httpWebRequest;
-        }
-
-        private string GetResponse(HttpWebResponse response)
-        {
-            var inputStream = response.GetResponseStream();
 
-            // convert stream to string
-            using var reader = new StreamReader(inputStream);
-            return reader.ReadToEnd();
+            Task<HttpResponseMessage> response = client.SendAsync(request);
+            HttpResponseMessage result = response.Result;
+            return result;
         }
     }
 }

tests/Amazon.CloudWatch.EMF.Tests/Environment/EC2EnvironmentTests.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using Amazon.CloudWatch.EMF.Config;
 using Amazon.CloudWatch.EMF.Environment;
 using AutoFixture;
@@ -167,7 +168,9 @@ public void Probe_False()
             // Arrange
             var configuration = _fixture.Create<IConfiguration>();
             var resourceFetcher = _fixture.Create<IResourceFetcher>();
-            resourceFetcher.Fetch<EC2Metadata>(Arg.Any<Uri>()).Throws<EMFClientException>();
+            resourceFetcher.FetchString(
+                Arg.Any<Uri>(), Arg.Any<string>(), Arg.Any<Dictionary<string, string>>()
+                ).Throws<EMFClientException>();
             var environment = new EC2Environment(configuration, resourceFetcher);
 
             // Act
@@ -183,7 +186,12 @@ public void Type_WhenNoMetadata()
             // Arrange
             var configuration = _fixture.Create<IConfiguration>();
             var resourceFetcher = _fixture.Create<IResourceFetcher>();
-            resourceFetcher.Fetch<EC2Metadata>(Arg.Any<Uri>()).Throws<EMFClientException>();
+            resourceFetcher.FetchString(
+                Arg.Any<Uri>(), Arg.Any<string>(), Arg.Any<Dictionary<string, string>>()
+                ).Returns("fake_token");
+            resourceFetcher.FetchJson<EC2Metadata>(
+                Arg.Any<Uri>(), Arg.Any<string>(), Arg.Any<Dictionary<string, string>>()
+                ).Throws<EMFClientException>();
             var environment = new EC2Environment(configuration, resourceFetcher);
             environment.Probe();
 
@@ -199,7 +207,12 @@ public void Type_WithMetadata()
             // Arrange
             var configuration = _fixture.Create<IConfiguration>();
             var resourceFetcher = _fixture.Create<IResourceFetcher>();
-            resourceFetcher.Fetch<EC2Metadata>(Arg.Any<Uri>()).Returns(new EC2Metadata());
+            resourceFetcher.FetchString(
+                Arg.Any<Uri>(), Arg.Any<string>(), Arg.Any<Dictionary<string, string>>()
+                ).Returns("fake_token");
+            resourceFetcher.FetchJson<EC2Metadata>(
+                Arg.Any<Uri>(), Arg.Any<string>(), Arg.Any<Dictionary<string, string>>()
+                ).Returns(new EC2Metadata());
             var environment = new EC2Environment(configuration, resourceFetcher);
             environment.Probe();
 

tests/Amazon.CloudWatch.EMF.Tests/Environment/ECSEnvironmentTests.cs

@@ -152,7 +152,7 @@ public void Probe_False()
             // Arrange
             var configuration = _fixture.Create<IConfiguration>();
             var resourceFetcher = _fixture.Create<IResourceFetcher>();
-            resourceFetcher.Fetch<ECSMetadata>(Arg.Any<Uri>()).Throws<EMFClientException>();
+            resourceFetcher.FetchJson<ECSMetadata>(Arg.Any<Uri>(), Arg.Any<string>()).Throws<EMFClientException>();
             var environment = new ECSEnvironment(configuration, resourceFetcher);
 
             // Act