ImdsManagedIdentitySource handles 502 response with CredentialUnavailableException

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

@@ -4,8 +4,6 @@
 using System;
 using System.Net;
 using System.Net.Sockets;
-using System.Runtime.CompilerServices;
-using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
@@ -104,7 +102,8 @@ protected override Request CreateRequest(string[] scopes)
 
         protected override async ValueTask<AccessToken> HandleResponseAsync(bool async, TokenRequestContext context, Response response, CancellationToken cancellationToken)
         {
-            if (response.Status == 400)
+            // 502 is typically due to the client dealing with a proxy configuration, which is not supported.
+            if (response.Status == 400 || response.Status == 502)
             {
                 string message = _identityUnavailableErrorMessage ?? await Pipeline.Diagnostics
                     .CreateRequestFailedMessageAsync(response, IdentityUnavailableError, null, null, async)

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

@@ -13,10 +13,12 @@ namespace Azure.Identity
     internal abstract class ManagedIdentitySource
     {
         internal const string AuthenticationResponseInvalidFormatError = "Invalid response, the authentication response was not in the expected format.";
+        private ManagedIdentityResponseClassifier _responseClassifier;
 
         protected ManagedIdentitySource(CredentialPipeline pipeline)
         {
             Pipeline = pipeline;
+            _responseClassifier = new ManagedIdentityResponseClassifier();
         }
 
         protected internal CredentialPipeline Pipeline { get; }
@@ -25,11 +27,17 @@ protected ManagedIdentitySource(CredentialPipeline pipeline)
         public virtual async ValueTask<AccessToken> AuthenticateAsync(bool async, TokenRequestContext context, CancellationToken cancellationToken)
         {
             using Request request = CreateRequest(context.Scopes);
-            Response response = async
-                ? await Pipeline.HttpPipeline.SendRequestAsync(request, cancellationToken).ConfigureAwait(false)
-                : Pipeline.HttpPipeline.SendRequest(request, cancellationToken);
+            using HttpMessage message = new HttpMessage(request, _responseClassifier);
+            if (async)
+            {
+                await Pipeline.HttpPipeline.SendAsync(message, cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                Pipeline.HttpPipeline.Send(message, cancellationToken);
+            }
 
-            return await HandleResponseAsync(async, context, response, cancellationToken).ConfigureAwait(false);
+            return await HandleResponseAsync(async, context, message.Response, cancellationToken).ConfigureAwait(false);
         }
 
         protected virtual async ValueTask<AccessToken> HandleResponseAsync(bool async, TokenRequestContext context, Response response, CancellationToken cancellationToken)
@@ -90,5 +98,18 @@ private static AccessToken GetTokenFromResponse(in JsonElement root)
 
             return null;
         }
+
+        private class ManagedIdentityResponseClassifier : ResponseClassifier
+        {
+            public override bool IsRetriableResponse(HttpMessage message)
+            {
+                return message.Response.Status switch
+                {
+                    404 => true,
+                    502 => false,
+                    _ => base.IsRetriableResponse(message)
+                };
+            }
+        }
     }
 }