Implemented GitHub feedback

Author: Robbie-Microsoft

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

@@ -288,18 +288,21 @@ public static async Task<bool> ProbeImdsEndpointAsync(
 
             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: requestContext.UserCancellationToken,
-                    retryPolicy: retryPolicy)
-                .ConfigureAwait(false);
+                using (var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(1)))
+                {
+                    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: timeoutCts.Token,
+                        retryPolicy: retryPolicy)
+                    .ConfigureAwait(false);
+                }
             }
             catch (Exception ex)
             {

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

@@ -106,10 +106,8 @@ internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(
             bool noImdsV2 = false)
         {
             // First check env vars to avoid the probe if possible
-            ManagedIdentitySource source = GetManagedIdentitySourceNoImdsV2(requestContext.Logger);
-#pragma warning disable CS0618 // Type or member is obsolete
-            if (source != ManagedIdentitySource.DefaultToImds)
-#pragma warning restore CS0618 // Type or member is obsolete
+            ManagedIdentitySource source = GetManagedIdentitySourceNoImds(requestContext.Logger);
+            if (source != ManagedIdentitySource.None)
             {
                 s_sourceName = source;
                 return source;
@@ -144,10 +142,18 @@ internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(
             return s_sourceName;
         }
 
-        // 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)
+        /// <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)
         {
             string identityEndpoint = EnvironmentVariables.IdentityEndpoint;
             string identityHeader = EnvironmentVariables.IdentityHeader;
@@ -189,9 +195,7 @@ internal static ManagedIdentitySource GetManagedIdentitySourceNoImdsV2(ILoggerAd
             }
             else
             {
-#pragma warning disable CS0618 // Type or member is obsolete
-                return ManagedIdentitySource.DefaultToImds;
-#pragma warning restore CS0618 // Type or member is obsolete
+                return ManagedIdentitySource.None;
             }
         }
 

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

@@ -48,7 +48,7 @@ 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("Use Imds instead. We are no longer defaulting to IMDS. An error will be thrown instead if no managed identity sources are available.")]
+        [Obsolete("In use only to support the now obsolete GetManagedIdentitySource API. Will be removed in a future version. Use GetManagedIdentitySourceAsync instead.")]
         DefaultToImds,
 
         /// <summary>

src/client/Microsoft.Identity.Client/ManagedIdentityApplication.cs

@@ -77,7 +77,16 @@ public async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync()
         [Obsolete("Use GetManagedIdentitySourceAsync() instead. \"ManagedIdentityApplication mi = miBuilder.Build() as ManagedIdentityApplication;\"")]
         public static ManagedIdentitySource GetManagedIdentitySource()
         {
-            return ManagedIdentityClient.GetManagedIdentitySourceNoImdsV2();
+            var source = ManagedIdentityClient.GetManagedIdentitySourceNoImds();
+            
+            return source == ManagedIdentitySource.None
+#pragma warning disable CS0618
+                // ManagedIdentitySource.DefaultToImds is marked obsolete, but is intentionally used here as a sentinel value to support legacy detection logic.
+                // This value signals that none of the environment-based managed identity sources were detected.
+                ? ManagedIdentitySource.DefaultToImds
+#pragma warning restore CS0618
+                : source;
+
         }
     }
 }

tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHttpMessageHandler.cs

@@ -41,8 +41,14 @@ internal class MockHttpMessageHandler : HttpClientHandler
 
         public X509Certificate2 ExpectedMtlsBindingCertificate { get; set; }
 
+        public Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> HandlerFunc { get; set; }
+
         protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
+            if (HandlerFunc != null)
+            {
+                return await HandlerFunc(request, cancellationToken).ConfigureAwait(false);
+            }
 
             ActualRequestMessage = request;
 

tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs

@@ -3,10 +3,12 @@
 
 using System;
 using System.Net;
+using System.Net.Http;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.AppConfig;
 using Microsoft.Identity.Client.ManagedIdentity;
+using Microsoft.Identity.Client.ManagedIdentity.V2;
 using Microsoft.Identity.Test.Common.Core.Helpers;
 using Microsoft.Identity.Test.Common.Core.Mocks;
 using Microsoft.Identity.Test.Unit.Helpers;
@@ -431,5 +433,44 @@ await mi.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource)
                 Assert.AreEqual(Num504Errors, requestsMade);
             }
         }
+
+        [TestMethod]
+        public async Task ProbeImdsEndpointAsync_TimesOutAfterOneSecond()
+        {
+            using (new EnvVariableContext())
+            using (var httpManager = new MockHttpManager())
+            {
+                var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned);
+
+                miBuilder
+                    .WithHttpManager(httpManager)
+                    .WithRetryPolicyFactory(_testRetryPolicyFactory);
+
+                var managedIdentityApp = miBuilder.Build();
+
+                httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2));
+
+                var imdsV1Handler = new MockHttpMessageHandler
+                {
+                    HandlerFunc = async (request, cancellationToken) =>
+                    {
+                        // IMDS (V1 and V2) probe has a timeout after 1 second
+                        await Task.Delay(TimeSpan.FromSeconds(2), cancellationToken).ConfigureAwait(false);
+                        return new HttpResponseMessage(HttpStatusCode.OK);
+                    }
+                };
+                httpManager.AddMockHandler(imdsV1Handler);
+
+                var miSource = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync().ConfigureAwait(false);
+                Assert.AreEqual(ManagedIdentitySource.None, miSource); // Probe timed out, no source available
+
+                var ex = await Assert.ThrowsExceptionAsync<MsalServiceException>(async () =>
+                    await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource)
+                    .ExecuteAsync().ConfigureAwait(false)
+                ).ConfigureAwait(false);
+
+                Assert.AreEqual(MsalError.ManagedIdentityAllSourcesUnavailable, ex.ErrorCode);
+            }
+        }
     }
 }