DAC only makes probe requeset on first GetToken (Azure#49159)

Author: christothes

sdk/identity/Azure.Identity/CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+- `DefaultAzureCredential` no longer sends a probe request on each call to `GetToken`. It now only happens on the first call.
 
 ### Other Changes
 

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

@@ -22,6 +22,7 @@ internal class ManagedIdentityClient
         private MsalManagedIdentityClient _msalManagedIdentityClient;
         private bool _isChainedCredential;
         private ManagedIdentityClientOptions _options;
+        private bool _probeRequestSent;
 
         protected ManagedIdentityClient()
         {
@@ -66,9 +67,11 @@ public async ValueTask<AccessToken> AuthenticateAsync(bool async, TokenRequestCo
             AzureIdentityEventSource.Singleton.ManagedIdentityCredentialSelected(availableSource.ToString(), _options.ManagedIdentityId.ToString());
 
             // If the source is DefaultToImds and the credential is chained, we should probe the IMDS endpoint first.
-            if (availableSource == MSAL.ManagedIdentitySource.DefaultToImds && _isChainedCredential)
+            if (availableSource == MSAL.ManagedIdentitySource.DefaultToImds && _isChainedCredential && !_probeRequestSent)
             {
-                return await AuthenticateCoreAsync(async, context, cancellationToken).ConfigureAwait(false);
+                var probedFlowTokenResult = await AuthenticateCoreAsync(async, context, cancellationToken).ConfigureAwait(false);
+                _probeRequestSent = true;
+                return probedFlowTokenResult;
             }
 
             // ServiceFabric does not support specifying user-assigned managed identity by client ID or resource ID. The managed identity selected is based on the resource configuration.

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

@@ -34,26 +34,6 @@ public ManagedIdentityCredentialTests(bool isAsync) : base(isAsync)
 
         private const string ExpectedToken = "mock-msi-access-token";
 
-        [Test]
-        public async Task VerifyExpiringTokenRefresh()
-        {
-            int callCount = 0;
-
-            var mockClient = new MockManagedIdentityClient(CredentialPipeline.GetInstance(null))
-            {
-                TokenFactory = () => { callCount++; return new AccessToken(Guid.NewGuid().ToString(), DateTimeOffset.UtcNow.AddMinutes(2)); }
-            };
-
-            var cred = InstrumentClient(new ManagedIdentityCredential(mockClient));
-
-            for (int i = 0; i < 5; i++)
-            {
-                await cred.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
-            }
-
-            Assert.AreEqual(5, callCount);
-        }
-
         [NonParallelizable]
         [Test]
         public async Task VerifyImdsRequestWithClientIdMockAsync()
@@ -84,6 +64,38 @@ public async Task VerifyImdsRequestWithClientIdMockAsync()
             Assert.IsTrue(query.Contains($"{Constants.ManagedIdentityClientId}=mock-client-id"));
         }
 
+        [NonParallelizable]
+        [Test]
+        public async Task VerifyImdsSendsProbeOnlyOnFirstRequest()
+        {
+            using var environment = new TestEnvVar(new() { { "MSI_ENDPOINT", null }, { "MSI_SECRET", null }, { "IDENTITY_ENDPOINT", null }, { "IDENTITY_HEADER", null }, { "AZURE_POD_IDENTITY_AUTHORITY_HOST", null } });
+
+            int probeCount = 0;
+            var mockTransport = new MockTransport(req =>
+            {
+                if (!req.Headers.TryGetValue("Metadata", out var _))
+                {
+                    probeCount++;
+                    return CreateErrorMockResponse(400, "mock error");
+                }
+                else
+                {
+                    return CreateMockResponse(200, ExpectedToken);
+                }
+            });
+            var options = new TokenCredentialOptions() { Transport = mockTransport, IsChainedCredential = true };
+            var pipeline = CredentialPipeline.GetInstance(options);
+            ManagedIdentityCredential credential = InstrumentClient(new ManagedIdentityCredential(
+                new ManagedIdentityClient(
+                    new ManagedIdentityClientOptions() { Pipeline = pipeline, ManagedIdentityId = ManagedIdentityId.FromUserAssignedClientId("mock-client-id"), IsForceRefreshEnabled = true, Options = options })));
+
+            await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
+            AccessToken actualToken = await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
+
+            Assert.AreEqual(ExpectedToken, actualToken.Token);
+            Assert.AreEqual(1, probeCount, "Probe was sent more than once.");
+        }
+
         [NonParallelizable]
         [Test]
         public async Task ImdsWithEmptyClientIdIsIgnoredMockAsync()