From 17f139044134287a1cebe2bb2dcd7d5c7e342aad Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 02:38:35 +0000
Subject: [PATCH 01/13] chore: extract getProfile to new function

---
 .../src/fromInstanceMetadata.ts               | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index 45eabf71f924..b9ed621decff 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -23,15 +23,7 @@ export const fromInstanceMetadata = (
 ): CredentialProvider => {
   const { timeout, maxRetries } = providerConfigFromInit(init);
   return async () => {
-    const profile = (
-      await retry<string>(
-        async () =>
-          (
-            await httpRequest({ host: IMDS_IP, path: IMDS_PATH, timeout })
-          ).toString(),
-        maxRetries
-      )
-    ).trim();
+    const profile = await getProfile(timeout, maxRetries);
 
     return retry(async () => {
       const credsResponse = JSON.parse(
@@ -53,3 +45,14 @@ export const fromInstanceMetadata = (
     }, maxRetries);
   };
 };
+
+const getProfile = async (timeout: number, maxRetries: number) =>
+  (
+    await retry<string>(
+      async () =>
+        (
+          await httpRequest({ host: IMDS_IP, path: IMDS_PATH, timeout })
+        ).toString(),
+      maxRetries
+    )
+  ).trim();

From 87e00855f515c6066ad421b9c544a6bb2e79645f Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 02:43:26 +0000
Subject: [PATCH 02/13] chore: extract getCredentials to new function

---
 .../src/fromInstanceMetadata.ts               | 39 ++++++++++---------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index b9ed621decff..bbd36dfbf5b1 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -24,25 +24,7 @@ export const fromInstanceMetadata = (
   const { timeout, maxRetries } = providerConfigFromInit(init);
   return async () => {
     const profile = await getProfile(timeout, maxRetries);
-
-    return retry(async () => {
-      const credsResponse = JSON.parse(
-        (
-          await httpRequest({
-            host: IMDS_IP,
-            path: IMDS_PATH + profile,
-            timeout
-          })
-        ).toString()
-      );
-      if (!isImdsCredentials(credsResponse)) {
-        throw new ProviderError(
-          "Invalid response received from instance metadata service."
-        );
-      }
-
-      return fromImdsCredentials(credsResponse);
-    }, maxRetries);
+    return retry(async () => getCredentials(timeout, profile), maxRetries);
   };
 };
 
@@ -56,3 +38,22 @@ const getProfile = async (timeout: number, maxRetries: number) =>
       maxRetries
     )
   ).trim();
+
+const getCredentials = async (timeout: number, profile: string) => {
+  const credsResponse = JSON.parse(
+    (
+      await httpRequest({
+        host: IMDS_IP,
+        path: IMDS_PATH + profile,
+        timeout
+      })
+    ).toString()
+  );
+  if (!isImdsCredentials(credsResponse)) {
+    throw new ProviderError(
+      "Invalid response received from instance metadata service."
+    );
+  }
+
+  return fromImdsCredentials(credsResponse);
+};

From e84407e6819c8f62f1b6b699bd6dce4f076030db Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 02:49:57 +0000
Subject: [PATCH 03/13] chore: pass RequestOptions in helper functions

---
 .../src/fromInstanceMetadata.ts                   | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index bbd36dfbf5b1..3f59c4d96bf9 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -10,6 +10,7 @@ import {
 } from "./remoteProvider/ImdsCredentials";
 import { retry } from "./remoteProvider/retry";
 import { ProviderError } from "@aws-sdk/property-provider";
+import { RequestOptions } from "http";
 
 const IMDS_IP = "169.254.169.254";
 const IMDS_PATH = "/latest/meta-data/iam/security-credentials/";
@@ -23,29 +24,29 @@ export const fromInstanceMetadata = (
 ): CredentialProvider => {
   const { timeout, maxRetries } = providerConfigFromInit(init);
   return async () => {
-    const profile = await getProfile(timeout, maxRetries);
-    return retry(async () => getCredentials(timeout, profile), maxRetries);
+    const profile = await getProfile(maxRetries, { timeout });
+    return retry(async () => getCredentials(profile, { timeout }), maxRetries);
   };
 };
 
-const getProfile = async (timeout: number, maxRetries: number) =>
+const getProfile = async (maxRetries: number, options: RequestOptions) =>
   (
     await retry<string>(
       async () =>
         (
-          await httpRequest({ host: IMDS_IP, path: IMDS_PATH, timeout })
+          await httpRequest({ ...options, host: IMDS_IP, path: IMDS_PATH })
         ).toString(),
       maxRetries
     )
   ).trim();
 
-const getCredentials = async (timeout: number, profile: string) => {
+const getCredentials = async (profile: string, options: RequestOptions) => {
   const credsResponse = JSON.parse(
     (
       await httpRequest({
+        ...options,
         host: IMDS_IP,
-        path: IMDS_PATH + profile,
-        timeout
+        path: IMDS_PATH + profile
       })
     ).toString()
   );

From 8715c7ef277959f3bb703014179d13408f0f9b64 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 02:52:20 +0000
Subject: [PATCH 04/13] chore: move retry out of getProfile

---
 .../src/fromInstanceMetadata.ts                  | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index 3f59c4d96bf9..021199e900b1 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -24,21 +24,17 @@ export const fromInstanceMetadata = (
 ): CredentialProvider => {
   const { timeout, maxRetries } = providerConfigFromInit(init);
   return async () => {
-    const profile = await getProfile(maxRetries, { timeout });
+    const profile = (
+      await retry<string>(async () => getProfile({ timeout }), maxRetries)
+    ).trim();
     return retry(async () => getCredentials(profile, { timeout }), maxRetries);
   };
 };
 
-const getProfile = async (maxRetries: number, options: RequestOptions) =>
+const getProfile = async (options: RequestOptions) =>
   (
-    await retry<string>(
-      async () =>
-        (
-          await httpRequest({ ...options, host: IMDS_IP, path: IMDS_PATH })
-        ).toString(),
-      maxRetries
-    )
-  ).trim();
+    await httpRequest({ ...options, host: IMDS_IP, path: IMDS_PATH })
+  ).toString();
 
 const getCredentials = async (profile: string, options: RequestOptions) => {
   const credsResponse = JSON.parse(

From d43e3c3c0a3b4ea23c0f674a2dcd12f998f3b72d Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 02:56:04 +0000
Subject: [PATCH 05/13] chore: utility function getCredentials

---
 .../src/fromInstanceMetadata.ts               | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index 021199e900b1..66a518f83f6a 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -24,19 +24,30 @@ export const fromInstanceMetadata = (
 ): CredentialProvider => {
   const { timeout, maxRetries } = providerConfigFromInit(init);
   return async () => {
-    const profile = (
-      await retry<string>(async () => getProfile({ timeout }), maxRetries)
-    ).trim();
-    return retry(async () => getCredentials(profile, { timeout }), maxRetries);
+    return getCredentials(maxRetries, { timeout });
   };
 };
 
+const getCredentials = async (maxRetries: number, options: RequestOptions) => {
+  const profile = (
+    await retry<string>(async () => getProfile(options), maxRetries)
+  ).trim();
+
+  return retry(
+    async () => getCredentialsFromProfile(profile, options),
+    maxRetries
+  );
+};
+
 const getProfile = async (options: RequestOptions) =>
   (
     await httpRequest({ ...options, host: IMDS_IP, path: IMDS_PATH })
   ).toString();
 
-const getCredentials = async (profile: string, options: RequestOptions) => {
+const getCredentialsFromProfile = async (
+  profile: string,
+  options: RequestOptions
+) => {
   const credsResponse = JSON.parse(
     (
       await httpRequest({
@@ -46,6 +57,7 @@ const getCredentials = async (profile: string, options: RequestOptions) => {
       })
     ).toString()
   );
+
   if (!isImdsCredentials(credsResponse)) {
     throw new ProviderError(
       "Invalid response received from instance metadata service."

From 420bf56101455e0446079ca382a40f52bd45b872 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 04:06:22 +0000
Subject: [PATCH 06/13] chore: wip for secure IMDS workflow

---
 .../src/fromInstanceMetadata.ts               | 87 ++++++++++++++++---
 .../src/remoteProvider/httpRequest.ts         | 20 +++--
 2 files changed, 89 insertions(+), 18 deletions(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index 66a518f83f6a..cb38557ed3f3 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -1,4 +1,4 @@
-import { CredentialProvider } from "@aws-sdk/types";
+import { CredentialProvider, Credentials } from "@aws-sdk/types";
 import {
   RemoteProviderInit,
   providerConfigFromInit
@@ -14,6 +14,7 @@ import { RequestOptions } from "http";
 
 const IMDS_IP = "169.254.169.254";
 const IMDS_PATH = "/latest/meta-data/iam/security-credentials/";
+const IMDS_TOKEN_PATH = "/latest/api/token";
 
 /**
  * Creates a credential provider that will source credentials from the EC2
@@ -22,21 +23,81 @@ const IMDS_PATH = "/latest/meta-data/iam/security-credentials/";
 export const fromInstanceMetadata = (
   init: RemoteProviderInit = {}
 ): CredentialProvider => {
+  // when set to true, metadata service will not fetch token
+  let disableFetchToken = false;
   const { timeout, maxRetries } = providerConfigFromInit(init);
-  return async () => {
-    return getCredentials(maxRetries, { timeout });
-  };
-};
 
-const getCredentials = async (maxRetries: number, options: RequestOptions) => {
-  const profile = (
-    await retry<string>(async () => getProfile(options), maxRetries)
-  ).trim();
+  const getMetadataToken = async () =>
+    httpRequest({
+      host: IMDS_IP,
+      path: IMDS_TOKEN_PATH,
+      method: "PUT",
+      headers: {
+        "x-aws-ec2-metadata-token-ttl-seconds": "21600"
+      }
+    });
 
-  return retry(
-    async () => getCredentialsFromProfile(profile, options),
-    maxRetries
-  );
+  const getCredentials = async (
+    maxRetries: number,
+    options: RequestOptions
+  ) => {
+    const profile = (
+      await retry<string>(async () => {
+        let profile: string;
+        try {
+          profile = await getProfile(options);
+        } catch (err) {
+          if (err.statusCode === 401) {
+            disableFetchToken = false;
+          }
+          throw err;
+        }
+        return profile;
+      }, maxRetries)
+    ).trim();
+
+    return retry(async () => {
+      let creds: Credentials;
+      try {
+        creds = await getCredentialsFromProfile(profile, options);
+      } catch (err) {
+        if (err.statusCode === 401) {
+          disableFetchToken = false;
+        }
+        throw err;
+      }
+      return creds;
+    }, maxRetries);
+  };
+
+  return async () => {
+    if (disableFetchToken) {
+      return getCredentials(maxRetries, { timeout });
+    } else {
+      let token: string;
+      try {
+        token = (await getMetadataToken()).toString();
+      } catch (error) {
+        if (error.statusCode === 400) {
+          throw Object.assign(error, {
+            message: "EC2 Metadata token request returned error"
+          });
+        } else if (
+          error.message === "TimeoutError" ||
+          [403, 404, 405].includes(error.statusCode)
+        ) {
+          disableFetchToken = true;
+        }
+        return getCredentials(maxRetries, { timeout });
+      }
+      return getCredentials(maxRetries, {
+        timeout,
+        headers: {
+          "x-aws-ec2-metadata-token": token
+        }
+      });
+    }
+  };
 };
 
 const getProfile = async (options: RequestOptions) =>
diff --git a/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts b/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts
index 4eeed8ad17f0..51b33219360a 100644
--- a/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts
+++ b/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts
@@ -10,18 +10,28 @@ export function httpRequest(options: RequestOptions): Promise<Buffer> {
     const req = request({ method: "GET", ...options });
     req.on("error", err => {
       reject(
-        new ProviderError("Unable to connect to instance metadata service")
+        Object.assign(
+          new ProviderError("Unable to connect to instance metadata service"),
+          err
+        )
       );
     });
 
+    req.on("timeout", err => {
+      reject(Object.assign(err, { message: "TimeoutError" }));
+    });
+
     req.on("response", (res: IncomingMessage) => {
       const { statusCode = 400 } = res;
       if (statusCode < 200 || 300 <= statusCode) {
-        const error = new ProviderError(
-          "Error response received from instance metadata service"
+        reject(
+          Object.assign(
+            new ProviderError(
+              "Error response received from instance metadata service"
+            ),
+            { statusCode }
+          )
         );
-        (error as any).statusCode = statusCode;
-        reject(error);
       }
 
       const chunks: Array<Buffer> = [];

From 851b789ff707699e6af7e7aa1becec8c0dad438a Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 04:59:18 +0000
Subject: [PATCH 07/13] test: fix tests for fromInstanceMetadata

---
 .../src/fromInstanceMetadata.spec.ts          | 67 ++++++++++++-------
 .../src/fromInstanceMetadata.ts               | 23 ++++---
 2 files changed, 56 insertions(+), 34 deletions(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
index 1871a28c9402..e9bdce2824ff 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
@@ -14,14 +14,29 @@ jest.mock("./remoteProvider/retry");
 jest.mock("./remoteProvider/RemoteProviderInit");
 
 describe("fromInstanceMetadata", () => {
+  const host = "169.254.169.254";
   const mockTimeout = 1000;
   const mockMaxRetries = 3;
-  const mockProfile = "foo";
+  const mockToken = "fooToken";
+  const mockProfile = "fooProfile";
+
+  const mockTokenRequestOptions = {
+    host,
+    path: "/latest/api/token",
+    method: "PUT",
+    headers: {
+      "x-aws-ec2-metadata-token-ttl-seconds": "21600"
+    },
+    timeout: mockTimeout
+  };
 
-  const mockHttpRequestOptions = {
-    host: "169.254.169.254",
+  const mockProfileRequestOptions = {
+    host,
     path: "/latest/meta-data/iam/security-credentials/",
-    timeout: mockTimeout
+    timeout: mockTimeout,
+    headers: {
+      "x-aws-ec2-metadata-token": mockToken
+    }
   };
 
   const mockImdsCreds = Object.freeze({
@@ -50,8 +65,9 @@ describe("fromInstanceMetadata", () => {
     jest.resetAllMocks();
   });
 
-  it("gets profile name from IMDS, and passes profile name to fetch credentials", async () => {
+  it("gets token and profile name to fetch credentials", async () => {
     (httpRequest as jest.Mock)
+      .mockResolvedValueOnce(mockToken)
       .mockResolvedValueOnce(mockProfile)
       .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
 
@@ -59,16 +75,18 @@ describe("fromInstanceMetadata", () => {
     (fromImdsCredentials as jest.Mock).mockReturnValue(mockCreds);
 
     await expect(fromInstanceMetadata()()).resolves.toEqual(mockCreds);
-    expect(httpRequest).toHaveBeenCalledTimes(2);
-    expect(httpRequest).toHaveBeenNthCalledWith(1, mockHttpRequestOptions);
-    expect(httpRequest).toHaveBeenNthCalledWith(2, {
-      ...mockHttpRequestOptions,
-      path: `${mockHttpRequestOptions.path}${mockProfile}`
+    expect(httpRequest).toHaveBeenCalledTimes(3);
+    expect(httpRequest).toHaveBeenNthCalledWith(1, mockTokenRequestOptions);
+    expect(httpRequest).toHaveBeenNthCalledWith(2, mockProfileRequestOptions);
+    expect(httpRequest).toHaveBeenNthCalledWith(3, {
+      ...mockProfileRequestOptions,
+      path: `${mockProfileRequestOptions.path}${mockProfile}`
     });
   });
 
   it("trims profile returned name from IMDS", async () => {
     (httpRequest as jest.Mock)
+      .mockResolvedValueOnce(mockToken)
       .mockResolvedValueOnce("   " + mockProfile + "  ")
       .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
 
@@ -76,11 +94,9 @@ describe("fromInstanceMetadata", () => {
     (fromImdsCredentials as jest.Mock).mockReturnValue(mockCreds);
 
     await expect(fromInstanceMetadata()()).resolves.toEqual(mockCreds);
-    expect(httpRequest).toHaveBeenCalledTimes(2);
-    expect(httpRequest).toHaveBeenNthCalledWith(1, mockHttpRequestOptions);
-    expect(httpRequest).toHaveBeenNthCalledWith(2, {
-      ...mockHttpRequestOptions,
-      path: `${mockHttpRequestOptions.path}${mockProfile}`
+    expect(httpRequest).toHaveBeenNthCalledWith(3, {
+      ...mockProfileRequestOptions,
+      path: `${mockProfileRequestOptions.path}${mockProfile}`
     });
   });
 
@@ -118,6 +134,7 @@ describe("fromInstanceMetadata", () => {
 
   it("throws ProviderError if credentials returned are incorrect", async () => {
     (httpRequest as jest.Mock)
+      .mockResolvedValueOnce(mockToken)
       .mockResolvedValueOnce(mockProfile)
       .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
 
@@ -130,37 +147,41 @@ describe("fromInstanceMetadata", () => {
       )
     );
     expect(retry).toHaveBeenCalledTimes(2);
-    expect(httpRequest).toHaveBeenCalledTimes(2);
+    expect(httpRequest).toHaveBeenCalledTimes(3);
     expect(isImdsCredentials).toHaveBeenCalledTimes(1);
     expect(isImdsCredentials).toHaveBeenCalledWith(mockImdsCreds);
     expect(fromImdsCredentials).not.toHaveBeenCalled();
   });
 
-  it("throws Error if requestFromEc2Imds for profile fails", async () => {
+  it("throws Error if httpRequest for profile fails", async () => {
     const mockError = new Error("profile not found");
-    (httpRequest as jest.Mock).mockRejectedValueOnce(mockError);
+    (httpRequest as jest.Mock)
+      .mockResolvedValueOnce(mockToken)
+      .mockRejectedValueOnce(mockError);
     (retry as jest.Mock).mockImplementation((fn: any) => fn());
 
     await expect(fromInstanceMetadata()()).rejects.toEqual(mockError);
     expect(retry).toHaveBeenCalledTimes(1);
-    expect(httpRequest).toHaveBeenCalledTimes(1);
+    expect(httpRequest).toHaveBeenCalledTimes(2);
   });
 
-  it("throws Error if requestFromEc2Imds for credentials fails", async () => {
+  it("throws Error if httpRequest for credentials fails", async () => {
     const mockError = new Error("creds not found");
     (httpRequest as jest.Mock)
+      .mockResolvedValueOnce(mockToken)
       .mockResolvedValueOnce(mockProfile)
       .mockRejectedValueOnce(mockError);
     (retry as jest.Mock).mockImplementation((fn: any) => fn());
 
     await expect(fromInstanceMetadata()()).rejects.toEqual(mockError);
     expect(retry).toHaveBeenCalledTimes(2);
-    expect(httpRequest).toHaveBeenCalledTimes(2);
+    expect(httpRequest).toHaveBeenCalledTimes(3);
     expect(fromImdsCredentials).not.toHaveBeenCalled();
   });
 
-  it("throws SyntaxError if requestFromEc2Imds returns unparseable creds", async () => {
+  it("throws SyntaxError if httpRequest returns unparseable creds", async () => {
     (httpRequest as jest.Mock)
+      .mockResolvedValueOnce(mockToken)
       .mockResolvedValueOnce(mockProfile)
       .mockResolvedValueOnce(".");
     (retry as jest.Mock).mockImplementation((fn: any) => fn());
@@ -169,7 +190,7 @@ describe("fromInstanceMetadata", () => {
       new SyntaxError("Unexpected token . in JSON at position 0")
     );
     expect(retry).toHaveBeenCalledTimes(2);
-    expect(httpRequest).toHaveBeenCalledTimes(2);
+    expect(httpRequest).toHaveBeenCalledTimes(3);
     expect(fromImdsCredentials).not.toHaveBeenCalled();
   });
 });
diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index cb38557ed3f3..2a871b3da722 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -27,16 +27,6 @@ export const fromInstanceMetadata = (
   let disableFetchToken = false;
   const { timeout, maxRetries } = providerConfigFromInit(init);
 
-  const getMetadataToken = async () =>
-    httpRequest({
-      host: IMDS_IP,
-      path: IMDS_TOKEN_PATH,
-      method: "PUT",
-      headers: {
-        "x-aws-ec2-metadata-token-ttl-seconds": "21600"
-      }
-    });
-
   const getCredentials = async (
     maxRetries: number,
     options: RequestOptions
@@ -76,7 +66,7 @@ export const fromInstanceMetadata = (
     } else {
       let token: string;
       try {
-        token = (await getMetadataToken()).toString();
+        token = (await getMetadataToken({ timeout })).toString();
       } catch (error) {
         if (error.statusCode === 400) {
           throw Object.assign(error, {
@@ -100,6 +90,17 @@ export const fromInstanceMetadata = (
   };
 };
 
+const getMetadataToken = async (options: RequestOptions) =>
+  httpRequest({
+    ...options,
+    host: IMDS_IP,
+    path: IMDS_TOKEN_PATH,
+    method: "PUT",
+    headers: {
+      "x-aws-ec2-metadata-token-ttl-seconds": "21600"
+    }
+  });
+
 const getProfile = async (options: RequestOptions) =>
   (
     await httpRequest({ ...options, host: IMDS_IP, path: IMDS_PATH })

From 226f211e25d9f2be149da77162a043a89cd3250f Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 05:54:22 +0000
Subject: [PATCH 08/13] test: throws error if metadata token returns with
 statusCode 400

---
 .../src/fromInstanceMetadata.spec.ts                     | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
index e9bdce2824ff..74a111726727 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
@@ -193,4 +193,13 @@ describe("fromInstanceMetadata", () => {
     expect(httpRequest).toHaveBeenCalledTimes(3);
     expect(fromImdsCredentials).not.toHaveBeenCalled();
   });
+
+  it("throws error if metadata token errors with statusCode 400", async () => {
+    const tokenError = Object.assign(new Error("token not found"), {
+      statusCode: 400
+    });
+    (httpRequest as jest.Mock).mockRejectedValueOnce(tokenError);
+
+    await expect(fromInstanceMetadata()()).rejects.toEqual(tokenError);
+  });
 });

From e16216db8f4a3b0de5e89428a205913d0bd331fc Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 06:07:21 +0000
Subject: [PATCH 09/13] test: fromInstanceMetadata disables fetching of token

---
 .../src/fromInstanceMetadata.spec.ts          | 39 +++++++++++++++++++
 .../src/fromInstanceMetadata.ts               |  2 +-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
index 74a111726727..82d7bb9f698e 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
@@ -202,4 +202,43 @@ describe("fromInstanceMetadata", () => {
 
     await expect(fromInstanceMetadata()()).rejects.toEqual(tokenError);
   });
+
+  describe("disables fetching of token", () => {
+    beforeEach(() => {
+      (retry as jest.Mock).mockImplementation((fn: any) => fn());
+      (fromImdsCredentials as jest.Mock).mockReturnValue(mockCreds);
+    });
+
+    it("when token fetch returns with TimeoutError", async () => {
+      const tokenError = new Error("TimeoutError");
+
+      (httpRequest as jest.Mock)
+        .mockRejectedValueOnce(tokenError)
+        .mockResolvedValueOnce(mockProfile)
+        .mockResolvedValueOnce(JSON.stringify(mockImdsCreds))
+        .mockResolvedValueOnce(mockProfile)
+        .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
+
+      const fromInstanceMetadataFunc = fromInstanceMetadata();
+      await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+      await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+    });
+
+    [403, 404, 405].forEach(statusCode => {
+      it(`when token fetch errors with statusCode ${statusCode}`, async () => {
+        const tokenError = Object.assign(new Error(), { statusCode });
+
+        (httpRequest as jest.Mock)
+          .mockRejectedValueOnce(tokenError)
+          .mockResolvedValueOnce(mockProfile)
+          .mockResolvedValueOnce(JSON.stringify(mockImdsCreds))
+          .mockResolvedValueOnce(mockProfile)
+          .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
+
+        const fromInstanceMetadataFunc = fromInstanceMetadata();
+        await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+        await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+      });
+    });
+  });
 });
diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
index 2a871b3da722..b0b298123bc6 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.ts
@@ -68,7 +68,7 @@ export const fromInstanceMetadata = (
       try {
         token = (await getMetadataToken({ timeout })).toString();
       } catch (error) {
-        if (error.statusCode === 400) {
+        if (error?.statusCode === 400) {
           throw Object.assign(error, {
             message: "EC2 Metadata token request returned error"
           });

From a801ee9e7e80781c46d6feb357377e4d0ac07571 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 06:13:23 +0000
Subject: [PATCH 10/13] test: uses insecure data flow once, if error is not
 TimeoutError

---
 .../src/fromInstanceMetadata.spec.ts          | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
index 82d7bb9f698e..c5ac7a94484e 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
@@ -241,4 +241,23 @@ describe("fromInstanceMetadata", () => {
       });
     });
   });
+
+  it("uses insecure data flow once, if error is not TimeoutError", async () => {
+    const tokenError = new Error("Error");
+
+    (httpRequest as jest.Mock)
+      .mockRejectedValueOnce(tokenError)
+      .mockResolvedValueOnce(mockProfile)
+      .mockResolvedValueOnce(JSON.stringify(mockImdsCreds))
+      .mockResolvedValueOnce(mockToken)
+      .mockResolvedValueOnce(mockProfile)
+      .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
+
+    (retry as jest.Mock).mockImplementation((fn: any) => fn());
+    (fromImdsCredentials as jest.Mock).mockReturnValue(mockCreds);
+
+    const fromInstanceMetadataFunc = fromInstanceMetadata();
+    await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+    await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+  });
 });

From 1d7e13e26e0526473b59777208b602466c9e27b2 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 06:15:13 +0000
Subject: [PATCH 11/13] test: uses insecure data flow once, if error statusCode
 is not 400, 403, 404, 405

---
 .../src/fromInstanceMetadata.spec.ts          | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
index c5ac7a94484e..35aeb0fcf119 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
@@ -260,4 +260,23 @@ describe("fromInstanceMetadata", () => {
     await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
     await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
   });
+
+  it("uses insecure data flow once, if error statusCode is not 400, 403, 404, 405", async () => {
+    const tokenError = Object.assign(new Error("Error"), { statusCode: 406 });
+
+    (httpRequest as jest.Mock)
+      .mockRejectedValueOnce(tokenError)
+      .mockResolvedValueOnce(mockProfile)
+      .mockResolvedValueOnce(JSON.stringify(mockImdsCreds))
+      .mockResolvedValueOnce(mockToken)
+      .mockResolvedValueOnce(mockProfile)
+      .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
+
+    (retry as jest.Mock).mockImplementation((fn: any) => fn());
+    (fromImdsCredentials as jest.Mock).mockReturnValue(mockCreds);
+
+    const fromInstanceMetadataFunc = fromInstanceMetadata();
+    await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+    await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+  });
 });

From 30e17d65dbed70a13c9001b816c7ed8a9a92b243 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 06:23:05 +0000
Subject: [PATCH 12/13] test: re-enables fetching of token

---
 .../src/fromInstanceMetadata.spec.ts          | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
index 35aeb0fcf119..2196ba610b2f 100644
--- a/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
+++ b/packages/credential-provider-imds/src/fromInstanceMetadata.spec.ts
@@ -279,4 +279,47 @@ describe("fromInstanceMetadata", () => {
     await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
     await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
   });
+
+  describe("re-enables fetching of token", () => {
+    const error401 = Object.assign(new Error("error"), { statusCode: 401 });
+
+    beforeEach(() => {
+      const tokenError = new Error("TimeoutError");
+
+      (httpRequest as jest.Mock)
+        .mockRejectedValueOnce(tokenError)
+        .mockResolvedValueOnce(mockProfile)
+        .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
+
+      (retry as jest.Mock).mockImplementation((fn: any) => fn());
+      (fromImdsCredentials as jest.Mock).mockReturnValue(mockCreds);
+    });
+
+    it("when profile error with 401", async () => {
+      (httpRequest as jest.Mock)
+        .mockRejectedValueOnce(error401)
+        .mockResolvedValueOnce(mockToken)
+        .mockResolvedValueOnce(mockProfile)
+        .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
+
+      const fromInstanceMetadataFunc = fromInstanceMetadata();
+      await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+      await expect(fromInstanceMetadataFunc()).rejects.toEqual(error401);
+      await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+    });
+
+    it("when creds error with 401", async () => {
+      (httpRequest as jest.Mock)
+        .mockResolvedValueOnce(mockProfile)
+        .mockRejectedValueOnce(error401)
+        .mockResolvedValueOnce(mockToken)
+        .mockResolvedValueOnce(mockProfile)
+        .mockResolvedValueOnce(JSON.stringify(mockImdsCreds));
+
+      const fromInstanceMetadataFunc = fromInstanceMetadata();
+      await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+      await expect(fromInstanceMetadataFunc()).rejects.toEqual(error401);
+      await expect(fromInstanceMetadataFunc()).resolves.toEqual(mockCreds);
+    });
+  });
 });

From 94587626355c0e4ba56e8d13267fab8ebecea16e Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Fri, 10 Jul 2020 06:41:21 +0000
Subject: [PATCH 13/13] test: timeout for httpRequest

---
 .../src/remoteProvider/httpRequest.spec.ts         | 14 ++++++++++++++
 .../src/remoteProvider/httpRequest.ts              |  4 ++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/packages/credential-provider-imds/src/remoteProvider/httpRequest.spec.ts b/packages/credential-provider-imds/src/remoteProvider/httpRequest.spec.ts
index a05a8377d70a..a66bf63a1845 100644
--- a/packages/credential-provider-imds/src/remoteProvider/httpRequest.spec.ts
+++ b/packages/credential-provider-imds/src/remoteProvider/httpRequest.spec.ts
@@ -92,4 +92,18 @@ describe("httpRequest", () => {
       [300, 400, 500].forEach(errorOnStatusCode);
     });
   });
+
+  it("timeout", async () => {
+    const timeout = 1000;
+    const scope = nock(`http://${host}:${port}`)
+      .get(path)
+      .delay(timeout * 2)
+      .reply(200, "expectedResponse");
+
+    await expect(
+      httpRequest({ host, path, port, timeout })
+    ).rejects.toStrictEqual(new Error("TimeoutError"));
+
+    scope.done();
+  });
 });
diff --git a/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts b/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts
index 51b33219360a..25d8b0d4c105 100644
--- a/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts
+++ b/packages/credential-provider-imds/src/remoteProvider/httpRequest.ts
@@ -17,8 +17,8 @@ export function httpRequest(options: RequestOptions): Promise<Buffer> {
       );
     });
 
-    req.on("timeout", err => {
-      reject(Object.assign(err, { message: "TimeoutError" }));
+    req.on("timeout", () => {
+      reject(new Error("TimeoutError"));
     });
 
     req.on("response", (res: IncomingMessage) => {