Fix scenario where IMDS probe succeeds, but MI is unavailable

sdk/identity/Azure.Identity/CHANGELOG.md

@@ -7,8 +7,8 @@
 ### Breaking Changes
 
 ### Bugs Fixed
-- Fixed an issue that prevented `ManagedIdentityCredential` from attempting to detect if Workload Identity is enabled in the current environment. [#46653](https://github.com/Azure/azure-sdk-for-net/issues/46653)
-- Fixed an issue that prevented `DefaultAzureCredential` from progressing past `ManagedIdentityCredential` in some scenarios where the identity was not available. [#46709](https://github.com/Azure/azure-sdk-for-net/issues/46709)
+- Fixed a regression that prevented `ManagedIdentityCredential` from attempting to detect if Workload Identity is enabled in the current environment. [#46653](https://github.com/Azure/azure-sdk-for-net/issues/46653)
+- Fixed regression that prevented `DefaultAzureCredential` from progressing past `ManagedIdentityCredential` in some scenarios where the identity was not available. [#46709](https://github.com/Azure/azure-sdk-for-net/issues/46709)
 
 ### Other Changes
 

sdk/identity/Azure.Identity/src/ImdsManagedIdentityProbeSource.cs

@@ -6,6 +6,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
+using Microsoft.Identity.Client;
 
 namespace Azure.Identity
 {
@@ -22,6 +23,7 @@ internal class ImdsManagedIdentityProbeSource : ManagedIdentitySource
         internal const string TimeoutError = "ManagedIdentityCredential authentication unavailable. The request to the managed identity endpoint timed out.";
         internal const string GatewayError = "ManagedIdentityCredential authentication unavailable. The request failed due to a gateway error.";
         internal const string AggregateError = "ManagedIdentityCredential authentication unavailable. Multiple attempts failed to obtain a token from the managed identity endpoint.";
+        internal const string UnknownError = "ManagedIdentityCredential authentication unavailable. An unexpected error has occurred.";
 
         private readonly ManagedIdentityId _managedIdentityId;
         private readonly Uri _imdsEndpoint;
@@ -97,6 +99,7 @@ protected override HttpMessage CreateHttpMessage(Request request)
 
         public async override ValueTask<AccessToken> AuthenticateAsync(bool async, TokenRequestContext context, CancellationToken cancellationToken)
         {
+            bool continueIMDSRequestAfterProbe;
             try
             {
                 return await base.AuthenticateAsync(async, context, cancellationToken).ConfigureAwait(false);
@@ -127,12 +130,29 @@ public async override ValueTask<AccessToken> AuthenticateAsync(bool async, Token
             catch (ProbeRequestResponseException)
             {
                 // This was an expected response from the IMDS endpoint without the Metadata header set.
-                // Re-issue the request (CreateRequest will add the appropriate header).
+                // Fall-through to the code below to re-issue the request through MsalManagedIdentityClient.
+                continueIMDSRequestAfterProbe = true;
+            }
+
+            if (continueIMDSRequestAfterProbe)
+            {
+                try
+                {
 #pragma warning disable AZC0110 // DO NOT use await keyword in possibly synchronous scope.
-                var authResult = await _client.AcquireTokenForManagedIdentityAsync(context, cancellationToken).ConfigureAwait(false);
-                return authResult.ToAccessToken();
+                    var authResult = await _client.AcquireTokenForManagedIdentityAsync(context, cancellationToken).ConfigureAwait(false);
 #pragma warning restore AZC0110 // DO NOT use await keyword in possibly synchronous scope.
+                    return authResult.ToAccessToken();
+                }
+                catch (MsalServiceException e)
+                {
+                    if (e.Message.Contains("unavailable"))
+                    {
+                        throw new CredentialUnavailableException(IdentityUnavailableError, e);
+                    }
+                    throw new CredentialUnavailableException(UnknownError, e);
+                }
             }
+            throw new CredentialUnavailableException(UnknownError);
         }
 
         protected override async ValueTask<AccessToken> HandleResponseAsync(bool async, TokenRequestContext context, HttpMessage message, CancellationToken cancellationToken)

sdk/identity/Azure.Identity/tests/ImdsManagedIdentitySourceTests.cs

@@ -142,7 +142,7 @@ public void DefaultAzureCredentialRetryBehaviorIsOverriddenWithOptions()
 
             var cred = new DefaultAzureCredential(credOptions);
 
-            Assert.ThrowsAsync<AuthenticationFailedException>(async () => await cred.GetTokenAsync(new(new[] { "test" })));
+            Assert.ThrowsAsync<CredentialUnavailableException>(async () => await cred.GetTokenAsync(new(new[] { "test" })));
 
             var expectedTimeouts = new TimeSpan?[] { TimeSpan.FromSeconds(1), null, null, null, null, null, null, null, null };
             CollectionAssert.AreEqual(expectedTimeouts, networkTimeouts);

sdk/identity/Azure.Identity/tests/ManagedIdentityCredentialTests.cs

@@ -447,6 +447,37 @@ public void VerifyImdsRequestFailureWithValidJsonIdentityNotFoundErrorThrowsCUE(
             }
         }
 
+        [NonParallelizable]
+        [Test]
+        [TestCase(true)]
+        [TestCase(false)]
+        public void VerifyImdsProbeRequestSuccessWithIdentityNotFoundErrorThrowsCUE(bool isChained)
+        {
+            using var environment = new TestEnvVar(new() { { "MSI_ENDPOINT", null }, { "MSI_SECRET", null }, { "IDENTITY_ENDPOINT", null }, { "IDENTITY_HEADER", null }, { "AZURE_POD_IDENTITY_AUTHORITY_HOST", null } });
+
+            var mockTransport = new MockTransport(req =>
+            {
+                if (!req.Headers.TryGetValue(ImdsManagedIdentityProbeSource.metadataHeaderName, out _))
+                {
+                    return CreateResponse(400, """{"error":"invalid_request","error_description":"Required metadata header not specified"}""");
+                }
+                return CreateResponse(400, """{"error":"invalid_request","error_description":"Identity not found"}""");
+            });
+            var options = new TokenCredentialOptions() { Transport = mockTransport, IsChainedCredential = isChained };
+            var pipeline = CredentialPipeline.GetInstance(options);
+
+            ManagedIdentityCredential credential = InstrumentClient(new ManagedIdentityCredential("mock-client-id", pipeline, options));
+            if (isChained)
+            {
+                var ex = Assert.ThrowsAsync<CredentialUnavailableException>(async () => await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Alternate)));
+                Assert.That(ex.Message, Does.Contain(ImdsManagedIdentityProbeSource.IdentityUnavailableError));
+            }
+            else
+            {
+                var ex = Assert.ThrowsAsync<AuthenticationFailedException>(async () => await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Alternate)));
+            }
+        }
+
         [NonParallelizable]
         [Test]
         [TestCase(502, ImdsManagedIdentitySource.GatewayError)]
@@ -476,7 +507,10 @@ public async Task VerifyIMDSRequestWithPodIdentityEnvVarMockAsync(string clientI
             using var environment = new TestEnvVar(new() { { "MSI_ENDPOINT", null }, { "MSI_SECRET", null }, { "IDENTITY_ENDPOINT", null }, { "IDENTITY_HEADER", null }, { "AZURE_POD_IDENTITY_AUTHORITY_HOST", "https://mock.podid.endpoint/" } });
 
             var response = CreateMockResponse(200, ExpectedToken);
-            var mockTransport = new MockTransport(response);
+            var mockTransport = new MockTransport(req =>
+            {
+                return response;
+            });
             var options = new TokenCredentialOptions() { Transport = mockTransport };
             var pipeline = CredentialPipeline.GetInstance(options);