[Identity] Simplest possible exponential retry for IMDS's 404 (#14827)

* wip

* simplest approach

* CHANGELOG entry

* Suggestion from Harsha

* Jeffs feedback

* feedback from Vinay

sdk/identity/identity/CHANGELOG.md

@@ -5,6 +5,7 @@
 - Fixed issue with the logging of success messages on the `DefaultAzureCredential` and the `ChainedTokenCredential`. These messages will now mention the internal credential that succeeded.
 - The feature of persistence caching of credentials (introduced in 2.0.0-beta.1) is now supported on Node.js 15 as well.
 - `AuthenticationRequiredError` (introduced in 2.0.0-beta.1) now has the same impact on `ChainedTokenCredential` as the `CredentialUnavailableError` which is to allow the next credential in the chain to be tried.
+- `ManagedIdentityCredential` now retries with exponential back-off when a request for a token fails with a 404 status code on environments with available IMDS endpoints. 
 
 ## 2.0.0-beta.2 (2021-04-06)
 

sdk/identity/identity/src/credentials/managedIdentityCredential/imdsMsi.ts

@@ -1,8 +1,15 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { AccessToken, GetTokenOptions, RequestPrepareOptions, RestError } from "@azure/core-http";
+import {
+  AccessToken,
+  delay,
+  GetTokenOptions,
+  RequestPrepareOptions,
+  RestError
+} from "@azure/core-http";
 import { SpanStatusCode } from "@azure/core-tracing";
+import { AuthenticationError } from "../../client/errors";
 import { IdentityClient } from "../../client/identityClient";
 import { credentialLogger } from "../../util/logging";
 import { createSpan } from "../../util/tracing";
@@ -47,6 +54,13 @@ function prepareRequestOptions(resource?: string, clientId?: string): RequestPre
   };
 }
 
+// 800ms -> 1600ms -> 3200ms
+export const imdsMsiRetryConfig = {
+  maxRetries: 3,
+  startDelayInMs: 800,
+  intervalIncrement: 2
+};
+
 export const imdsMsi: MSI = {
   async isAvailable(
     identityClient: IdentityClient,
@@ -129,11 +143,28 @@ export const imdsMsi: MSI = {
       `Using the IMDS endpoint coming form the environment variable MSI_ENDPOINT=${process.env.MSI_ENDPOINT}, and using the cloud shell to proceed with the authentication.`
     );
 
-    return msiGenericGetToken(
-      identityClient,
-      prepareRequestOptions(resource, clientId),
-      expiresInParser,
-      getTokenOptions
+    let nextDelayInMs = imdsMsiRetryConfig.startDelayInMs;
+    for (let retries = 0; retries < imdsMsiRetryConfig.maxRetries; retries++) {
+      try {
+        return await msiGenericGetToken(
+          identityClient,
+          prepareRequestOptions(resource, clientId),
+          expiresInParser,
+          getTokenOptions
+        );
+      } catch (error) {
+        if (error.statusCode === 404) {
+          await delay(nextDelayInMs);
+          nextDelayInMs *= imdsMsiRetryConfig.intervalIncrement;
+          continue;
+        }
+        throw error;
+      }
+    }
+
+    throw new AuthenticationError(
+      404,
+      `Failed to retrieve IMDS token after ${imdsMsiRetryConfig.maxRetries} retries.`
     );
   }
 };

sdk/identity/identity/test/internal/node/managedIdentityCredential.spec.ts

@@ -11,6 +11,8 @@ import {
 } from "../../../src/credentials/managedIdentityCredential/constants";
 import { MockAuthHttpClient, MockAuthHttpClientOptions, assertRejects } from "../../authTestUtils";
 import { OAuthErrorResponse } from "../../../src/client/errors";
+import Sinon from "sinon";
+import { imdsMsiRetryConfig } from "../../../src/credentials/managedIdentityCredential/imdsMsi";
 
 interface AuthRequestDetails {
   requests: WebResource[];
@@ -21,15 +23,22 @@ describe("ManagedIdentityCredential", function() {
   // There are no types available for this dependency, at least at the time this test file was written.
   // eslint-disable-next-line @typescript-eslint/no-require-imports
   const mockFs = require("mock-fs");
-
   let envCopy: string = "";
+  let sandbox: Sinon.SinonSandbox;
+  let clock: Sinon.SinonFakeTimers;
+
   beforeEach(() => {
     envCopy = JSON.stringify(process.env);
     delete process.env.IDENTITY_ENDPOINT;
     delete process.env.IDENTITY_HEADER;
     delete process.env.MSI_ENDPOINT;
     delete process.env.MSI_SECRET;
     delete process.env.IDENTITY_SERVER_THUMBPRINT;
+    sandbox = Sinon.createSandbox();
+    clock = sandbox.useFakeTimers({
+      now: Date.now(),
+      shouldAdvanceTime: true
+    });
   });
   afterEach(() => {
     mockFs.restore();
@@ -39,6 +48,8 @@ describe("ManagedIdentityCredential", function() {
     process.env.MSI_ENDPOINT = env.MSI_ENDPOINT;
     process.env.MSI_SECRET = env.MSI_SECRET;
     process.env.IDENTITY_SERVER_THUMBPRINT = env.IDENTITY_SERVER_THUMBPRINT;
+    sandbox.restore();
+    clock.restore();
   });
 
   it("sends an authorization request with a modified resource name", async function() {
@@ -171,6 +182,68 @@ describe("ManagedIdentityCredential", function() {
     );
   });
 
+  it("IMDS MSI retries and succeeds on 404", async function() {
+    process.env.AZURE_CLIENT_ID = "errclient";
+
+    const mockHttpClient = new MockAuthHttpClient({
+      authResponse: [
+        // First response says the IMDS endpoint is available.
+        { status: 200 },
+        { status: 404 },
+        // Retries one time and fails
+        { status: 404 },
+        // Retries a second time and succeeds
+        {
+          status: 200,
+          parsedBody: {
+            access_token: "token"
+          }
+        }
+      ]
+    });
+
+    const credential = new ManagedIdentityCredential(process.env.AZURE_CLIENT_ID, {
+      ...mockHttpClient.tokenCredentialOptions
+    });
+
+    const response = await credential.getToken("scopes");
+    assert.equal(response.token, "token");
+  });
+
+  it("IMDS MSI retries up to a limit on 404", async function() {
+    process.env.AZURE_CLIENT_ID = "errclient";
+
+    const mockHttpClient = new MockAuthHttpClient({
+      // First response says the IMDS endpoint is available.
+      authResponse: [
+        { status: 200 },
+        { status: 404 },
+        { status: 404 },
+        { status: 404 },
+        { status: 404 }
+      ]
+    });
+
+    const credential = new ManagedIdentityCredential(process.env.AZURE_CLIENT_ID, {
+      ...mockHttpClient.tokenCredentialOptions
+    });
+
+    const promise = credential.getToken("scopes");
+
+    imdsMsiRetryConfig.startDelayInMs = 80; // 800ms / 10
+
+    // 800ms -> 1600ms -> 3200ms, results in 6400ms, / 10 = 640ms
+    clock.tick(640);
+
+    await assertRejects(
+      promise,
+      (error: AuthenticationError) =>
+        error.message.indexOf(
+          `Failed to retrieve IMDS token after ${imdsMsiRetryConfig.maxRetries} retries.`
+        ) > -1
+    );
+  });
+
   // Unavailable exception throws while IMDS endpoint is unavailable. This test not valid.
   // it("can extend timeout for IMDS endpoint", async function() {
   //   // Mock a timeout so that the endpoint ping fails