Revert "IMDS Source Detection Logic Improvement (#5602)"

This reverts commit 964a074.

Author: bgavrilMS

src/client/Microsoft.Identity.Client/Http/Retry/ImdsProbeRetryPolicy.cs renamed to src/client/Microsoft.Identity.Client/Http/Retry/CsrMetadataProbeRetryPolicy.cs

@@ -5,11 +5,11 @@
 
 namespace Microsoft.Identity.Client.Http.Retry
 {
-    internal class ImdsProbeRetryPolicy : ImdsRetryPolicy
+    internal class CsrMetadataProbeRetryPolicy : ImdsRetryPolicy
     {
         protected override bool ShouldRetry(HttpResponse response, Exception exception)
         {
-            return HttpRetryConditions.ImdsProbe(response, exception);
+            return HttpRetryConditions.CsrMetadataProbe(response, exception);
         }
     }
 }

src/client/Microsoft.Identity.Client/Http/Retry/HttpRetryConditions.cs renamed to src/client/Microsoft.Identity.Client/Http/Retry/HttpRetryCondition.cs

@@ -27,21 +27,6 @@ public static bool DefaultManagedIdentity(HttpResponse response, Exception excep
             };
         }
 
-        /// <summary>
-        /// Retry policy specific to Imds v1 and v2 Probe.
-        /// Extends Imds retry policy but excludes 404 status code.
-        /// </summary>
-        public static bool ImdsProbe(HttpResponse response, Exception exception)
-        {
-            if (!Imds(response, exception))
-            {
-                return false;
-            }
-
-            // If Imds would retry but the status code is 404, don't retry
-            return (int)response.StatusCode is not 404;
-        }
-
         /// <summary>
         /// Retry policy specific to IMDS Managed Identity.
         /// </summary>
@@ -77,6 +62,21 @@ public static bool RegionDiscovery(HttpResponse response, Exception exception)
             return (int)response.StatusCode is not (404 or 408);
         }
 
+        /// <summary>
+        /// Retry policy specific to CSR Metadata Probe.
+        /// Extends Imds retry policy but excludes 404 status code.
+        /// </summary>
+        public static bool CsrMetadataProbe(HttpResponse response, Exception exception)
+        {
+            if (!Imds(response, exception))
+            {
+                return false;
+            }
+
+            // If Imds would retry but the status code is 404, don't retry
+            return (int)response.StatusCode is not 404;
+        }
+
         /// <summary>
         /// Retry condition for /token and /authorize endpoints
         /// </summary>

src/client/Microsoft.Identity.Client/Http/Retry/RetryPolicyFactory.cs

@@ -14,12 +14,12 @@ public virtual IRetryPolicy GetRetryPolicy(RequestType requestType)
                 case RequestType.STS:
                 case RequestType.ManagedIdentityDefault:
                     return new DefaultRetryPolicy(requestType);
-                case RequestType.ImdsProbe:
-                    return new ImdsProbeRetryPolicy();
                 case RequestType.Imds:
                     return new ImdsRetryPolicy();
                 case RequestType.RegionDiscovery:
                     return new RegionDiscoveryRetryPolicy();
+                case RequestType.CsrMetadataProbe:
+                    return new CsrMetadataProbeRetryPolicy();
                 default:
                     throw new ArgumentOutOfRangeException(nameof(requestType), requestType, "Unknown request type.");
             }

src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs

@@ -11,21 +11,17 @@
 using Microsoft.Identity.Client.ApiConfig.Parameters;
 using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.Http;
-using Microsoft.Identity.Client.Http.Retry;
 using Microsoft.Identity.Client.Internal;
-using Microsoft.Identity.Client.ManagedIdentity.V2;
-using Microsoft.Identity.Client.OAuth2;
 
 namespace Microsoft.Identity.Client.ManagedIdentity
 {
     internal class ImdsManagedIdentitySource : AbstractManagedIdentity
     {
         // 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
         // used in unit tests as well
-        public const string ApiVersionQueryParam = "api-version";
         public const string DefaultImdsBaseEndpoint= "http://169.254.169.254";
+        private const string ImdsTokenPath = "/metadata/identity/oauth2/token";
         public const string ImdsApiVersion = "2018-02-01";
-        public const string ImdsTokenPath = "/metadata/identity/oauth2/token";
 
         private const string DefaultMessage = "[Managed Identity] Service request failed.";
 
@@ -40,11 +36,6 @@ internal class ImdsManagedIdentitySource : AbstractManagedIdentity
 
         private static string s_cachedBaseEndpoint = null;
 
-        public static AbstractManagedIdentity Create(RequestContext requestContext)
-        {
-            return new ImdsManagedIdentitySource(requestContext);
-        }
-
         internal ImdsManagedIdentitySource(RequestContext requestContext) : 
             base(requestContext, ManagedIdentitySource.Imds)
         {
@@ -60,7 +51,7 @@ protected override Task<ManagedIdentityRequest> CreateRequestAsync(string resour
             ManagedIdentityRequest request = new(HttpMethod.Get, _imdsEndpoint);
 
             request.Headers.Add("Metadata", "true");
-            request.QueryParameters[ApiVersionQueryParam] = ImdsApiVersion;
+            request.QueryParameters["api-version"] = ImdsApiVersion;
             request.QueryParameters["resource"] = resource;
 
             switch (_requestContext.ServiceBundle.Config.ManagedIdentityId.IdType)
@@ -220,106 +211,5 @@ public static Uri GetValidatedEndpoint(
 
             return builder.Uri;
         }
-
-        public static string ImdsQueryParamsHelper(
-            RequestContext requestContext,
-            string apiVersionQueryParam,
-            string imdsApiVersion)
-        {
-            var queryParams = $"{apiVersionQueryParam}={imdsApiVersion}";
-
-            var userAssignedIdQueryParam = GetUserAssignedIdQueryParam(
-                requestContext.ServiceBundle.Config.ManagedIdentityId.IdType,
-                requestContext.ServiceBundle.Config.ManagedIdentityId.UserAssignedId,
-                requestContext.Logger);
-
-            if (userAssignedIdQueryParam != null)
-            {
-                queryParams += $"&{userAssignedIdQueryParam.Value.Key}={userAssignedIdQueryParam.Value.Value}";
-            }
-
-            return queryParams;
-        }
-
-        public static async Task<bool> ProbeImdsEndpointAsync(
-            RequestContext requestContext,
-            ImdsVersion imdsVersion,
-            CancellationToken cancellationToken)
-        {
-            string apiVersionQueryParam;
-            string imdsApiVersion;
-            string imdsEndpoint;
-            string imdsStringHelper;
-
-            switch (imdsVersion)
-            {
-                case ImdsVersion.V2:
-#if NET462
-                requestContext.Logger.Info("[Managed Identity] IMDSv2 flow is not supported on .NET Framework 4.6.2. Cryptographic operations required for managed identity authentication are unavailable on this platform. Skipping IMDSv2 probe.");
-                return false;
-#else
-                    apiVersionQueryParam = ImdsV2ManagedIdentitySource.ApiVersionQueryParam;
-                    imdsApiVersion = ImdsV2ManagedIdentitySource.ImdsV2ApiVersion;
-                    imdsEndpoint = ImdsV2ManagedIdentitySource.CsrMetadataPath;
-                    imdsStringHelper = "IMDSv2";
-                    break;
-#endif
-                case ImdsVersion.V1:
-                    apiVersionQueryParam = ApiVersionQueryParam;
-                    imdsApiVersion = ImdsApiVersion;
-                    imdsEndpoint = ImdsTokenPath;
-                    imdsStringHelper = "IMDSv1";
-                    break;
-
-                default:
-                    throw new ArgumentOutOfRangeException(nameof(imdsVersion), imdsVersion, null);
-            }
-
-            var queryParams = ImdsQueryParamsHelper(requestContext, apiVersionQueryParam, imdsApiVersion);
-
-            // probe omits the "Metadata: true" header and then treats 400 Bad Request as success
-            var headers = new Dictionary<string, string>
-            {
-                { OAuth2Header.XMsCorrelationId, requestContext.CorrelationId.ToString() }
-            };
-
-            IRetryPolicyFactory retryPolicyFactory = requestContext.ServiceBundle.Config.RetryPolicyFactory;
-            IRetryPolicy retryPolicy = retryPolicyFactory.GetRetryPolicy(RequestType.ImdsProbe);
-
-            HttpResponse response = null;
-
-            try
-            {
-                response = await requestContext.ServiceBundle.HttpManager.SendRequestAsync(
-                    GetValidatedEndpoint(requestContext.Logger, imdsEndpoint, queryParams),
-                    headers,
-                    body: null,
-                    method: HttpMethod.Get,
-                    logger: requestContext.Logger,
-                    doNotThrow: false,
-                    mtlsCertificate: null,
-                    validateServerCertificate: null,
-                    cancellationToken: cancellationToken,
-                    retryPolicy: retryPolicy)
-                .ConfigureAwait(false);
-            }
-            catch (Exception ex)
-            {
-                requestContext.Logger.Info($"[Managed Identity] {imdsStringHelper} probe endpoint failure. Exception occurred while sending request to probe endpoint: {ex}");
-                return false;
-            }
-
-            // probe omits the "Metadata: true" header and then treats 400 Bad Request as success
-            if (response.StatusCode == HttpStatusCode.BadRequest)
-            {
-                requestContext.Logger.Info(() => $"[Managed Identity] {imdsStringHelper} managed identity is available.");
-                return true;
-            }
-            else
-            {
-                requestContext.Logger.Info(() => $"[Managed Identity] {imdsStringHelper} managed identity is not available. Status code: {response.StatusCode}, Body: {response.Body}");
-                return false;
-            }
-        }
     }
 }

src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Concurrent;
 using System.IO;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
@@ -40,15 +41,12 @@ internal async Task<ManagedIdentityResponse> SendTokenRequestForManagedIdentityA
             AcquireTokenForManagedIdentityParameters parameters,
             CancellationToken cancellationToken)
         {
-            AbstractManagedIdentity msi = await GetOrSelectManagedIdentitySourceAsync(requestContext, parameters.IsMtlsPopRequested, cancellationToken).ConfigureAwait(false);
+            AbstractManagedIdentity msi = await GetOrSelectManagedIdentitySourceAsync(requestContext, parameters.IsMtlsPopRequested).ConfigureAwait(false);
             return await msi.AuthenticateAsync(parameters, cancellationToken).ConfigureAwait(false);
         }
 
         // This method tries to create managed identity source for different sources, if none is created then defaults to IMDS.
-        private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsync(
-            RequestContext requestContext,
-            bool isMtlsPopRequested,
-            CancellationToken cancellationToken)
+        private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsync(RequestContext requestContext, bool isMtlsPopRequested)
         {
             using (requestContext.Logger.LogMethodDuration())
             {
@@ -60,27 +58,28 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn
                 if (s_sourceName == ManagedIdentitySource.None)
                 {
                     // First invocation: detect and cache
-                    source = await GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested, cancellationToken).ConfigureAwait(false);
+                    source = await GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested).ConfigureAwait(false);
                 }
                 else
                 {
                     // Reuse cached value
                     source = s_sourceName;
                 }
 
-                // If the source has already been set to ImdsV2 (via this method, or GetManagedIdentitySourceAsync in ManagedIdentityApplication.cs)
-                // and mTLS PoP was NOT requested: fall back to ImdsV1, because ImdsV2 currently only supports mTLS PoP requests
+                // If the source has already been set to ImdsV2 (via this method,
+                // or GetManagedIdentitySourceAsync in ManagedIdentityApplication.cs) and mTLS PoP was NOT requested
+                // In this case, we need to fall back to ImdsV1, because ImdsV2 currently only supports mTLS PoP requests
                 if (source == ManagedIdentitySource.ImdsV2 && !isMtlsPopRequested)
                 {
                     requestContext.Logger.Info("[Managed Identity] ImdsV2 detected, but mTLS PoP was not requested. Falling back to ImdsV1 for this request only. Please use the \"WithMtlsProofOfPossession\" API to request a token via ImdsV2.");
                     // Do NOT modify s_sourceName; keep cached ImdsV2 so future PoP
                     // requests can leverage it.
-                    source = ManagedIdentitySource.Imds;
+                    source = ManagedIdentitySource.DefaultToImds;
                 }
 
                 // If the source is determined to be ImdsV1 and mTLS PoP was requested,
                 // throw an exception since ImdsV1 does not support mTLS PoP
-                if (source == ManagedIdentitySource.Imds && isMtlsPopRequested)
+                if (source == ManagedIdentitySource.DefaultToImds && isMtlsPopRequested)
                 {
                     throw new MsalClientException(
                         MsalError.MtlsPopTokenNotSupportedinImdsV1,
@@ -95,8 +94,7 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn
                     ManagedIdentitySource.CloudShell => CloudShellManagedIdentitySource.Create(requestContext),
                     ManagedIdentitySource.AzureArc => AzureArcManagedIdentitySource.Create(requestContext),
                     ManagedIdentitySource.ImdsV2 => ImdsV2ManagedIdentitySource.Create(requestContext),
-                    ManagedIdentitySource.Imds => ImdsManagedIdentitySource.Create(requestContext),
-                    _ => throw new MsalClientException(MsalError.ManagedIdentityAllSourcesUnavailable, MsalErrorMessage.ManagedIdentityAllSourcesUnavailable)
+                    _ => new ImdsManagedIdentitySource(requestContext)
                 };
             }
         }
@@ -105,58 +103,39 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn
         // This method is perf sensitive any changes should be benchmarked.
         internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(
             RequestContext requestContext,
-            bool isMtlsPopRequested,
-            CancellationToken cancellationToken)
+            bool isMtlsPopRequested)
         {
             // First check env vars to avoid the probe if possible
-            ManagedIdentitySource source = GetManagedIdentitySourceNoImds(requestContext.Logger);
-            if (source != ManagedIdentitySource.None)
+            ManagedIdentitySource source = GetManagedIdentitySourceNoImdsV2(requestContext.Logger);
+
+            // If a source is detected via env vars, or
+            // a source wasn't detected (it defaulted to ImdsV1) and MtlsPop was NOT requested,
+            // use the source.
+            // (don't trigger the ImdsV2 probe endpoint if MtlsPop was NOT requested)
+            if (source != ManagedIdentitySource.DefaultToImds || !isMtlsPopRequested)
             {
                 s_sourceName = source;
                 return source;
             }
 
-            // skip the ImdsV2 probe if MtlsPop was NOT requested
-            if (isMtlsPopRequested)
-            {
-                var imdsV2Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V2, cancellationToken).ConfigureAwait(false);
-                if (imdsV2Response)
-                {
-                    requestContext.Logger.Info("[Managed Identity] ImdsV2 detected.");
-                    s_sourceName = ManagedIdentitySource.ImdsV2;
-                    return s_sourceName;
-                }
-            }
-            else
-            {
-                requestContext.Logger.Info("[Managed Identity] Mtls Pop was not requested; skipping ImdsV2 probe.");
-            }
-
-            var imdsV1Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V1, cancellationToken).ConfigureAwait(false);
-            if (imdsV1Response)
+            // Otherwise, probe IMDSv2
+            var response = await ImdsV2ManagedIdentitySource.GetCsrMetadataAsync(requestContext, probeMode: true).ConfigureAwait(false);
+            if (response != null)
             {
-                requestContext.Logger.Info("[Managed Identity] ImdsV1 detected.");
-                s_sourceName = ManagedIdentitySource.Imds;
+                requestContext.Logger.Info("[Managed Identity] ImdsV2 detected.");
+                s_sourceName = ManagedIdentitySource.ImdsV2;
                 return s_sourceName;
             }
 
-            requestContext.Logger.Info($"[Managed Identity] {MsalErrorMessage.ManagedIdentityAllSourcesUnavailable}");
-            s_sourceName = ManagedIdentitySource.None;
+            requestContext.Logger.Info("[Managed Identity] IMDSv2 probe failed. Defaulting to IMDSv1.");
+            s_sourceName = ManagedIdentitySource.DefaultToImds;
             return s_sourceName;
         }
 
-        /// <summary>
-        /// Detects the managed identity source based on the availability of environment variables.
-        /// It does not probe IMDS, but it checks for all other sources.
-        /// This method does not cache its result, as reading environment variables is inexpensive.
-        /// It is performance sensitive; any changes should be benchmarked.
-        /// </summary>
-        /// <param name="logger">Optional logger for diagnostic output.</param>
-        /// <returns>
-        /// The detected <see cref="ManagedIdentitySource"/> based on environment variables.
-        /// Returns <c>ManagedIdentitySource.None</c> if no environment-based source is detected.
-        /// </returns>
-        internal static ManagedIdentitySource GetManagedIdentitySourceNoImds(ILoggerAdapter logger = null)
+        // Detect managed identity source based on the availability of environment variables.
+        // The result of this method is not cached because reading environment variables is cheap. 
+        // This method is perf sensitive any changes should be benchmarked.
+        internal static ManagedIdentitySource GetManagedIdentitySourceNoImdsV2(ILoggerAdapter logger = null)
         {
             string identityEndpoint = EnvironmentVariables.IdentityEndpoint;
             string identityHeader = EnvironmentVariables.IdentityHeader;
@@ -198,7 +177,7 @@ internal static ManagedIdentitySource GetManagedIdentitySourceNoImds(ILoggerAdap
             }
             else
             {
-                return ManagedIdentitySource.None;
+                return ManagedIdentitySource.DefaultToImds;
             }
         }
 

src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentitySource.cs

@@ -48,7 +48,6 @@ public enum ManagedIdentitySource
         /// Indicates that the source is defaulted to IMDS since no environment variables are set.
         /// This is used to detect the managed identity source.
         /// </summary>
-        [Obsolete("In use only to support the now obsolete GetManagedIdentitySource API. Will be removed in a future version. Use GetManagedIdentitySourceAsync instead.")]
         DefaultToImds,
 
         /// <summary>