Implemented a Retry Policy for the IMDS Managed Identity Source (#7614)

The IMDS retry policy spec is located
[here](https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/blob/main/docs/imds_retry_based_on_errors.md).

The existing (linear) retry policy tests were moved out of Imds.spec.ts
and into a new file, DefaultManagedIdentityRetryPolicy.spec.ts.

Author: Robbie-Microsoft

change/@azure-msal-common-a198dabd-e74d-48c4-90fe-b9c1e4c81e93.json

@@ -0,0 +1,7 @@
+{
+  "type": "minor",
+  "comment": "Implemented a Retry Policy for the IMDS Managed Identity Source #7614",
+  "packageName": "@azure/msal-common",
+  "email": "[email protected]",
+  "dependentChangeType": "patch"
+}

change/@azure-msal-node-9aa1a5a1-daa9-4887-890e-6b873f5e0b8f.json

@@ -0,0 +1,7 @@
+{
+  "type": "minor",
+  "comment": "Implemented a Retry Policy for the IMDS Managed Identity Source #7614",
+  "packageName": "@azure/msal-node",
+  "email": "[email protected]",
+  "dependentChangeType": "patch"
+}

lib/msal-common/apiReview/msal-common.api.md

@@ -2412,6 +2412,7 @@ export const HttpStatus: {
     readonly UNAUTHORIZED: 401;
     readonly NOT_FOUND: 404;
     readonly REQUEST_TIMEOUT: 408;
+    readonly GONE: 410;
     readonly TOO_MANY_REQUESTS: 429;
     readonly CLIENT_ERROR_RANGE_END: 499;
     readonly SERVER_ERROR: 500;

lib/msal-common/src/utils/Constants.ts

@@ -70,6 +70,7 @@ export const HttpStatus = {
     UNAUTHORIZED: 401,
     NOT_FOUND: 404,
     REQUEST_TIMEOUT: 408,
+    GONE: 410,
     TOO_MANY_REQUESTS: 429,
     CLIENT_ERROR_RANGE_END: 499,
     SERVER_ERROR: 500,

lib/msal-node/src/client/ManagedIdentitySources/BaseManagedIdentitySource.ts

@@ -164,7 +164,8 @@ export abstract class BaseManagedIdentitySource {
             ? this.networkClient
             : new HttpClientWithRetries(
                   this.networkClient,
-                  networkRequest.retryPolicy
+                  networkRequest.retryPolicy,
+                  this.logger
               );
 
         const reqTimestamp = TimeUtils.nowSeconds();

lib/msal-node/src/client/ManagedIdentitySources/Imds.ts

@@ -18,6 +18,7 @@ import {
     RESOURCE_BODY_OR_QUERY_PARAMETER_NAME,
 } from "../../utils/Constants.js";
 import { NodeStorage } from "../../cache/NodeStorage.js";
+import { ImdsRetryPolicy } from "../../retry/ImdsRetryPolicy.js";
 
 // IMDS constants. Docs for IMDS are available here https://docs.microsoft.com/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http
 const IMDS_TOKEN_PATH: string = "/metadata/identity/oauth2/token";
@@ -131,6 +132,8 @@ export class Imds extends BaseManagedIdentitySource {
 
         // bodyParameters calculated in BaseManagedIdentity.acquireTokenWithManagedIdentity
 
+        request.retryPolicy = new ImdsRetryPolicy();
+
         return request;
     }
 }

lib/msal-node/src/config/ManagedIdentityRequestParameters.ts

@@ -8,14 +8,8 @@ import {
     UrlString,
     UrlUtils,
 } from "@azure/msal-common/node";
-import {
-    HttpMethod,
-    MANAGED_IDENTITY_HTTP_STATUS_CODES_TO_RETRY_ON,
-    MANAGED_IDENTITY_MAX_RETRIES,
-    MANAGED_IDENTITY_RETRY_DELAY,
-    RetryPolicies,
-} from "../utils/Constants.js";
-import { LinearRetryPolicy } from "../retry/LinearRetryPolicy.js";
+import { DefaultManagedIdentityRetryPolicy } from "../retry/DefaultManagedIdentityRetryPolicy.js";
+import { HttpMethod, RetryPolicies } from "../utils/Constants.js";
 
 export class ManagedIdentityRequestParameters {
     private _baseEndpoint: string;
@@ -36,12 +30,8 @@ export class ManagedIdentityRequestParameters {
         this.bodyParameters = {} as Record<string, string>;
         this.queryParameters = {} as Record<string, string>;
 
-        const defaultRetryPolicy: LinearRetryPolicy = new LinearRetryPolicy(
-            MANAGED_IDENTITY_MAX_RETRIES,
-            MANAGED_IDENTITY_RETRY_DELAY,
-            MANAGED_IDENTITY_HTTP_STATUS_CODES_TO_RETRY_ON
-        );
-        this.retryPolicy = retryPolicy || defaultRetryPolicy;
+        this.retryPolicy =
+            retryPolicy || new DefaultManagedIdentityRetryPolicy();
     }
 
     public computeUri(): string {

lib/msal-node/src/network/HttpClientWithRetries.ts

@@ -6,6 +6,7 @@
 import {
     HeaderNames,
     INetworkModule,
+    Logger,
     NetworkRequestOptions,
     NetworkResponse,
 } from "@azure/msal-common/node";
@@ -15,13 +16,16 @@ import { HttpMethod } from "../utils/Constants.js";
 export class HttpClientWithRetries implements INetworkModule {
     private httpClientNoRetries: INetworkModule;
     private retryPolicy: IHttpRetryPolicy;
+    private logger: Logger;
 
     constructor(
         httpClientNoRetries: INetworkModule,
-        retryPolicy: IHttpRetryPolicy
+        retryPolicy: IHttpRetryPolicy,
+        logger: Logger
     ) {
         this.httpClientNoRetries = httpClientNoRetries;
         this.retryPolicy = retryPolicy;
+        this.logger = logger;
     }
 
     private async sendNetworkRequestAsyncHelper<T>(
@@ -45,11 +49,16 @@ export class HttpClientWithRetries implements INetworkModule {
         let response: NetworkResponse<T> =
             await this.sendNetworkRequestAsyncHelper(httpMethod, url, options);
 
+        if ("isNewRequest" in this.retryPolicy) {
+            this.retryPolicy.isNewRequest = true;
+        }
+
         let currentRetry: number = 0;
         while (
             await this.retryPolicy.pauseForRetry(
                 response.status,
                 currentRetry,
+                this.logger,
                 response.headers[HeaderNames.RETRY_AFTER]
             )
         ) {

lib/msal-node/src/retry/DefaultManagedIdentityRetryPolicy.ts

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { IncomingHttpHeaders } from "http";
+import { HttpStatus, Logger } from "@azure/msal-common";
+import { IHttpRetryPolicy } from "./IHttpRetryPolicy.js";
+import { LinearRetryStrategy } from "./LinearRetryStrategy.js";
+
+const DEFAULT_MANAGED_IDENTITY_MAX_RETRIES: number = 3;
+const DEFAULT_MANAGED_IDENTITY_RETRY_DELAY_MS: number = 1000;
+const DEFAULT_MANAGED_IDENTITY_HTTP_STATUS_CODES_TO_RETRY_ON: Array<number> = [
+    HttpStatus.NOT_FOUND,
+    HttpStatus.REQUEST_TIMEOUT,
+    HttpStatus.TOO_MANY_REQUESTS,
+    HttpStatus.SERVER_ERROR,
+    HttpStatus.SERVICE_UNAVAILABLE,
+    HttpStatus.GATEWAY_TIMEOUT,
+];
+
+export class DefaultManagedIdentityRetryPolicy implements IHttpRetryPolicy {
+    /*
+     * this is defined here as a static variable despite being defined as a constant outside of the
+     * class because it needs to be overridden in the unit tests so that the unit tests run faster
+     */
+    static get DEFAULT_MANAGED_IDENTITY_RETRY_DELAY_MS(): number {
+        return DEFAULT_MANAGED_IDENTITY_RETRY_DELAY_MS;
+    }
+
+    private linearRetryStrategy: LinearRetryStrategy =
+        new LinearRetryStrategy();
+
+    async pauseForRetry(
+        httpStatusCode: number,
+        currentRetry: number,
+        logger: Logger,
+        retryAfterHeader: IncomingHttpHeaders["retry-after"]
+    ): Promise<boolean> {
+        if (
+            DEFAULT_MANAGED_IDENTITY_HTTP_STATUS_CODES_TO_RETRY_ON.includes(
+                httpStatusCode
+            ) &&
+            currentRetry < DEFAULT_MANAGED_IDENTITY_MAX_RETRIES
+        ) {
+            const retryAfterDelay: number =
+                this.linearRetryStrategy.calculateDelay(
+                    retryAfterHeader,
+                    DefaultManagedIdentityRetryPolicy.DEFAULT_MANAGED_IDENTITY_RETRY_DELAY_MS
+                );
+
+            logger.verbose(
+                `Retrying request in ${retryAfterDelay}ms (retry attempt: ${
+                    currentRetry + 1
+                })`
+            );
+
+            // pause execution for the calculated delay
+            await new Promise((resolve) => {
+                // retryAfterHeader value of 0 evaluates to false, and DEFAULT_MANAGED_IDENTITY_RETRY_DELAY_MS will be used
+                return setTimeout(resolve, retryAfterDelay);
+            });
+
+            return true;
+        }
+
+        // if the status code is not retriable or max retries have been reached, do not retry
+        return false;
+    }
+}

lib/msal-node/src/retry/ExponentialRetryStrategy.ts

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+export class ExponentialRetryStrategy {
+    // Minimum backoff time in milliseconds
+    private minExponentialBackoff: number;
+    // Maximum backoff time in milliseconds
+    private maxExponentialBackoff: number;
+    // Maximum backoff time in milliseconds
+    private exponentialDeltaBackoff: number;
+
+    constructor(
+        minExponentialBackoff: number,
+        maxExponentialBackoff: number,
+        exponentialDeltaBackoff: number
+    ) {
+        this.minExponentialBackoff = minExponentialBackoff;
+        this.maxExponentialBackoff = maxExponentialBackoff;
+        this.exponentialDeltaBackoff = exponentialDeltaBackoff;
+    }
+
+    /**
+     * Calculates the exponential delay based on the current retry attempt.
+     *
+     * @param {number} currentRetry - The current retry attempt number.
+     * @returns {number} - The calculated exponential delay in milliseconds.
+     *
+     * The delay is calculated using the formula:
+     * - If `currentRetry` is 0, it returns the minimum backoff time.
+     * - Otherwise, it calculates the delay as the minimum of:
+     *   - `(2^(currentRetry - 1)) * deltaBackoff`
+     *   - `maxBackoff`
+     *
+     * This ensures that the delay increases exponentially with each retry attempt,
+     * but does not exceed the maximum backoff time.
+     */
+    public calculateDelay(currentRetry: number): number {
+        // Attempt 1
+        if (currentRetry === 0) {
+            return this.minExponentialBackoff;
+        }
+
+        // Attempt 2+
+        const exponentialDelay = Math.min(
+            Math.pow(2, currentRetry - 1) * this.exponentialDeltaBackoff,
+            this.maxExponentialBackoff
+        );
+
+        return exponentialDelay;
+    }
+}

lib/msal-node/src/retry/IHttpRetryPolicy.ts

@@ -3,16 +3,25 @@
  * Licensed under the MIT License.
  */
 
-import http from "http";
+import { IncomingHttpHeaders } from "http";
+import { Logger } from "@azure/msal-common";
 
 export interface IHttpRetryPolicy {
-    /*
-     * if retry conditions occur, pauses and returns true
-     * otherwise return false
+    _isNewRequest?: boolean;
+    // set isNewRequest(value: boolean);
+
+    /**
+     * Pauses execution for a specified amount of time before retrying an HTTP request.
+     *
+     * @param httpStatusCode - The HTTP status code of the response.
+     * @param currentRetry - The current retry attempt number.
+     * @param retryAfterHeader - The value of the `retry-after` HTTP header, if present.
+     * @returns A promise that resolves to a boolean indicating whether to retry the request.
      */
     pauseForRetry(
         httpStatusCode: number,
         currentRetry: number,
-        retryAfterHeader: http.IncomingHttpHeaders["retry-after"]
+        logger: Logger,
+        retryAfterHeader?: IncomingHttpHeaders["retry-after"]
     ): Promise<boolean>;
 }