Repository: awslabs/aws-embedded-metrics-node
Pull request #115 (2022-05-20 21:36:43)

Migrate from IMDSv1 to IMDSv2
*Issue #, if available:*
N/A

*Description of changes:*
IMDSv2 uses session authentication to retrieve EC2 instance metadata.

https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.


Patch:
diff --git a/src/environment/EC2Environment.ts b/src/environment/EC2Environment.ts
index 272644e..e869d9c 100644
--- a/src/environment/EC2Environment.ts
+++ b/src/environment/EC2Environment.ts
@@ -17,11 +17,19 @@ import config from '../config/Configuration';
 import { MetricsContext } from '../logger/MetricsContext';
 import { AgentSink } from '../sinks/AgentSink';
 import { ISink } from '../sinks/Sink';
-import { fetch } from '../utils/Fetch';
+import { fetchJSON, fetchString } from '../utils/Fetch';
 import { LOG } from '../utils/Logger';
 import { IEnvironment } from './IEnvironment';
+import { RequestOptions } from 'http';
 
-const endpoint = 'http://169.254.169.254/latest/dynamic/instance-identity/document';
+// Documentation for configuring instance metadata can be found here:
+// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
+const host = '169.254.169.254';
+const tokenPath = '/latest/api/token';
+const tokenRequestHeaderKey = 'X-aws-ec2-metadata-token-ttl-seconds';
+const tokenRequestHeaderValue = '21600';
+const metadataPath = '/latest/dynamic/instance-identity/document';
+const metadataRequestTokenHeaderKey = 'X-aws-ec2-metadata-token';
 
 interface IEC2MetadataResponse {
   imageId: string;
@@ -34,14 +42,31 @@ interface IEC2MetadataResponse {
 export class EC2Environment implements IEnvironment {
   private metadata: IEC2MetadataResponse | undefined;
   private sink: ISink | undefined;
+  private token: string | undefined;
 
   public async probe(): Promise<boolean> {
     try {
-      this.metadata = await fetch<IEC2MetadataResponse>(endpoint);
-      if (this.metadata) {
-        return true;
+      const options: RequestOptions = {
+        host,
+        path: tokenPath,
+        method: "PUT",
+        headers: {[tokenRequestHeaderKey]: tokenRequestHeaderValue}
       }
+      this.token = await fetchString(options);
+    } catch (e) {
+      LOG(e);
       return false;
+    }
+
+    try {
+      const metadataOptions: RequestOptions = {
+        host,
+        path: metadataPath,
+        method: "GET",
+        headers: {[metadataRequestTokenHeaderKey]: this.token}
+      }
+      this.metadata = await fetchJSON<IEC2MetadataResponse>(metadataOptions);
+      return !!this.metadata;
     } catch (e) {
       LOG(e);
       return false;
diff --git a/src/environment/ECSEnvironment.ts b/src/environment/ECSEnvironment.ts
index 8a55659..2e29ee7 100644
--- a/src/environment/ECSEnvironment.ts
+++ b/src/environment/ECSEnvironment.ts
@@ -18,7 +18,7 @@ import { MetricsContext } from '../logger/MetricsContext';
 import { AgentSink } from '../sinks/AgentSink';
 import { ISink } from '../sinks/Sink';
 import { IEnvironment } from './IEnvironment';
-import { fetch } from '../utils/Fetch';
+import { fetchJSON } from '../utils/Fetch';
 import { LOG } from '../utils/Logger';
 import * as os from 'os';
 import { Constants } from '../Constants';
@@ -77,7 +77,8 @@ export class ECSEnvironment implements IEnvironment {
     }
 
     try {
-      this.metadata = await fetch<IECSMetadataResponse>(process.env.ECS_CONTAINER_METADATA_URI);
+      const options = new URL(process.env.ECS_CONTAINER_METADATA_URI);
+      this.metadata = await fetchJSON<IECSMetadataResponse>(options);
       if (this.metadata) {
         this.metadata.FormattedImageName = formatImageName(this.metadata.Image);
         LOG(`Successfully collected ECS Container metadata.`);
diff --git a/src/logger/MetricsContext.ts b/src/logger/MetricsContext.ts
index 502f0b2..6193e93 100644
--- a/src/logger/MetricsContext.ts
+++ b/src/logger/MetricsContext.ts
@@ -89,7 +89,7 @@ export class MetricsContext {
     this.properties[key] = value;
   }
 
-  public setTimestamp(timestamp: Date | number) {
+  public setTimestamp(timestamp: Date | number): void {
     this.timestamp = timestamp;
     this.meta.Timestamp = MetricsContext.resolveMetaTimestamp(timestamp);
   }
diff --git a/src/utils/Fetch.ts b/src/utils/Fetch.ts
index cf775be..92ace33 100644
--- a/src/utils/Fetch.ts
+++ b/src/utils/Fetch.ts
@@ -13,24 +13,20 @@
  * limitations under the License.
  */
 
-import * as http from 'http';
+import { IncomingMessage, RequestOptions, request as httpRequest } from 'http';
 
 const SOCKET_TIMEOUT = 1000;
 
 /**
- * Fetch JSON data from an remote HTTP endpoint and de-serialize to the provided type.
- * There are no guarantees the response will conform to the contract defined by T.
- * It is up to the consumer to ensure the provided T captures all possible response types
- * from the provided endpoint.
+ * Fetch data from a remote HTTP endpoint with the provided headers.
  *
- * @param url - currently only supports HTTP
+ * @param options - HTTP request options
  */
-const fetch = <T>(url: string): Promise<T> => {
-  return new Promise<T>((resolve, reject) => {
-    const request = http
-      .get(url, { timeout: 2000 }, (response: http.IncomingMessage) => {
+const fetch = (options: RequestOptions): Promise<Buffer> => {
+  return new Promise<Buffer>((resolve, reject) => {
+    const request = httpRequest(options, (response: IncomingMessage) => {
         if (!response.statusCode) {
-          reject(`Received undefined response status code from '${url}'`);
+          reject(`Received undefined response status code from '${options.host}${options.path}'`);
           return;
         }
 
@@ -49,15 +45,8 @@ const fetch = <T>(url: string): Promise<T> => {
         });
 
         response.on('end', () => {
-          let responseString;
-          try {
-            const buffer: Buffer = Buffer.concat(body, bodyBytes);
-            responseString = buffer.toString();
-            const parsedJson = JSON.parse(responseString);
-            resolve(parsedJson as T);
-          } catch (e) {
-            reject(`Failed to parse response from '${url}' as JSON. Received: ${responseString}`);
-          }
+          const buffer: Buffer = Buffer.concat(body, bodyBytes);
+          resolve(buffer)
         });
       })
       .on('error', (err: unknown) => {
@@ -67,11 +56,36 @@ const fetch = <T>(url: string): Promise<T> => {
     request.on('socket', socket => {
       socket.on('timeout', () => {
         request.abort();
-        reject(`Socket timeout while connecting to '${url}'`);
+        reject(`Socket timeout while connecting to '${options.host}${options.path}'`);
       });
       socket.setTimeout(SOCKET_TIMEOUT);
     });
-  });
-};
 
-export { fetch };
+    request.end();
+  })
+}
+
+/**
+ * Fetch a string from a remote HTTP endpoint with the provided headers.
+ *
+ * @param options - HTTP request options
+ */
+const fetchString = async (options: RequestOptions): Promise<string> => {
+  const buffer = await fetch(options);
+  return buffer.toString();
+}
+
+/**
+ * Fetch JSON data from a remote HTTP endpoint with the provided headers and de-serialize to the provided type.
+ * There are no guarantees the response will conform to the contract defined by T.
+ * It is up to the consumer to ensure the provided T captures all possible response types
+ * from the provided endpoint.
+ *
+ * @param options - HTTP request options
+ */
+const fetchJSON = async <T>(options: RequestOptions): Promise<T> => {
+  const responseString = await fetchString(options)
+  return JSON.parse(responseString);
+}
+
+export { fetch, fetchJSON, fetchString };


Test patch:
diff --git a/src/environment/__tests__/ECSEnvironment.test.ts b/src/environment/__tests__/ECSEnvironment.test.ts
index 9cb13a4..c078b29 100644
--- a/src/environment/__tests__/ECSEnvironment.test.ts
+++ b/src/environment/__tests__/ECSEnvironment.test.ts
@@ -2,12 +2,12 @@ import * as faker from 'faker';
 import config from '../../config/Configuration';
 import { ECSEnvironment } from '../ECSEnvironment';
 import { MetricsContext } from '../../logger/MetricsContext';
-import { fetch } from '../../utils/Fetch';
+import { fetchJSON } from '../../utils/Fetch';
 
 import * as os from 'os';
 import { Constants } from '../../Constants';
 jest.mock('../../utils/Fetch', () => ({
-  fetch: jest.fn(),
+  fetchJSON: jest.fn(),
 }));
 
 beforeEach(() => {
@@ -15,7 +15,7 @@ beforeEach(() => {
   config.serviceName = undefined;
   config.logGroupName = undefined;
   process.env = {
-    ECS_CONTAINER_METADATA_URI: faker.internet.ip(),
+    ECS_CONTAINER_METADATA_URI: faker.internet.url(),
   };
 });
 
@@ -79,7 +79,7 @@ describe('configureContext()', () => {
     const expected = getRandomECSMetadata();
 
     // @ts-ignore
-    fetch.mockImplementation(() => expected);
+    fetchJSON.mockImplementation(() => expected);
     const env = new ECSEnvironment();
     await env.probe();
 
@@ -101,7 +101,7 @@ describe('configureContext()', () => {
     const expected = getRandomECSMetadata();
 
     // @ts-ignore
-    fetch.mockImplementation(() => expected);
+    fetchJSON.mockImplementation(() => expected);
     const env = new ECSEnvironment();
     await env.probe();
 
@@ -119,7 +119,7 @@ describe('configureContext()', () => {
     const expected = getRandomECSMetadata();
 
     // @ts-ignore
-    fetch.mockImplementation(() => expected);
+    fetchJSON.mockImplementation(() => expected);
     const env = new ECSEnvironment();
     await env.probe();
 
@@ -209,7 +209,7 @@ describe('getName()', () => {
     };
 
     // @ts-ignore
-    fetch.mockImplementation(() => metadata);
+    fetchJSON.mockImplementation(() => metadata);
     const env = new ECSEnvironment();
     await env.probe();
 
@@ -229,7 +229,7 @@ describe('getName()', () => {
     };
 
     // @ts-ignore
-    fetch.mockImplementation(() => metadata);
+    fetchJSON.mockImplementation(() => metadata);
     const env = new ECSEnvironment();
     await env.probe();
 
@@ -243,7 +243,7 @@ describe('getName()', () => {
   test('returns Unknown if image not available', async () => {
     // arrange
     // @ts-ignore
-    fetch.mockImplementation(() => {
+    fetchJSON.mockImplementation(() => {
       return {};
     });
     const env = new ECSEnvironment();
diff --git a/test/integ/agent/Dockerfile b/test/integ/agent/Dockerfile
index 675ae75..41bd10d 100644
--- a/test/integ/agent/Dockerfile
+++ b/test/integ/agent/Dockerfile
@@ -1,4 +1,5 @@
-FROM public.ecr.aws/lts/ubuntu:latest
+# TODO: Fix the NO_PUBKEY error for jammy and update to use the :latest tag
+FROM public.ecr.aws/lts/ubuntu:20.04
 
 RUN apt-get update &&  \
     apt-get install -y ca-certificates curl && \