ManagedIdentityCredential honors CancellationTokens

sdk/core/Azure.Core/CHANGELOG.md

@@ -9,6 +9,7 @@
 ### Bugs Fixed
 
 - Fixed an issue where `BearerTokenAuthenticationPolicy` throws `ArgumentOutOfRangeException` if the `ExpiresOn` property of the token is the default value. ([#47040](https://github.com/Azure/azure-sdk-for-net/pull/47040))
+- Fixed an issue where `ManagedIdentityCredential` does not honor the `CancellationToken` passed to `GetToken` and `GetTokenAsync`. ([#47156](https://github.com/Azure/azure-sdk-for-net/issues/47156))
 
 ### Other Changes
 

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

@@ -99,8 +99,8 @@ public virtual async ValueTask<AuthenticationResult> AcquireTokenForManagedIdent
             }
 #pragma warning disable AZC0102 // Do not use GetAwaiter().GetResult().
             return async ?
-                await builder.ExecuteAsync().ConfigureAwait(false) :
-                builder.ExecuteAsync().GetAwaiter().GetResult();
+                await builder.ExecuteAsync(cancellationToken).ConfigureAwait(false) :
+                builder.ExecuteAsync(cancellationToken).GetAwaiter().GetResult();
 #pragma warning restore AZC0102 // Do not use GetAwaiter().GetResult().
         }
     }

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

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.IO;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
 using Azure.Core.Pipeline;
@@ -174,6 +175,60 @@ public void ManagedIdentityCredentialUsesDefaultTimeoutAndRetries()
             CollectionAssert.AreEqual(expectedTimeouts, networkTimeouts);
         }
 
+        [Test]
+        public void ManagedIdentityCredentialRetryBehaviorIsOverriddenWithOptions()
+        {
+            int callCount = 0;
+            List<TimeSpan?> networkTimeouts = new();
+
+            var mockTransport = MockTransport.FromMessageCallback(msg =>
+            {
+                callCount++;
+                networkTimeouts.Add(msg.NetworkTimeout);
+                Assert.IsTrue(msg.Request.Headers.TryGetValue(ImdsManagedIdentitySource.metadataHeaderName, out _));
+                return CreateMockResponse(500, "Error").WithHeader("Content-Type", "application/json");
+            });
+
+            var options = new TokenCredentialOptions()
+            {
+                Transport = mockTransport,
+                RetryPolicy = new RetryPolicy(1, DelayStrategy.CreateFixedDelayStrategy(TimeSpan.Zero))
+            };
+            options.Retry.MaxDelay = TimeSpan.Zero;
+
+            var cred = new ManagedIdentityCredential(
+                "testCLientId", options);
+
+            Assert.ThrowsAsync<AuthenticationFailedException>(async () => await cred.GetTokenAsync(new(new[] { "test" })));
+
+            var expectedTimeouts = new TimeSpan?[] { null, null };
+            CollectionAssert.AreEqual(expectedTimeouts, networkTimeouts);
+        }
+
+        [Test]
+        public void ManagedIdentityCredentialRespectsCancellationToken()
+        {
+            int callCount = 0;
+
+            var mockTransport = MockTransport.FromMessageCallback(msg =>
+            {
+                callCount++;
+                return CreateMockResponse(500, "Error").WithHeader("Content-Type", "application/json");
+            });
+
+            var options = new TokenCredentialOptions() { Transport = mockTransport };
+            options.Retry.MaxDelay = TimeSpan.Zero;
+
+            var cred = new ManagedIdentityCredential(
+                "testCLientId", options);
+
+            var cts = new CancellationTokenSource();
+            cts.CancelAfter(TimeSpan.Zero);
+            Assert.ThrowsAsync<TaskCanceledException>(async () => await cred.GetTokenAsync(new(new[] { "test" }), cts.Token));
+
+            Assert.AreEqual(0, callCount);
+        }
+
         private MockResponse CreateMockResponse(int responseCode, string token)
         {
             var response = new MockResponse(responseCode);