Hotfix for DAC probe request

Azure · May 6, 2024 · 06dd672 · 06dd672
1 parent ae13ec2
commit 06dd672
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 3 deletions.
diff --git a/sdk/identity/Azure.Identity/CHANGELOG.md b/sdk/identity/Azure.Identity/CHANGELOG.md
@@ -1,5 +1,10 @@
 # Release History
 
+## 1.11.3 (2024-05-07)
+
+### Bugs Fixed
+- Fixed a regression in `DefaultAzureCredential` probe request behavior for IMDS managed identity environments. [#43796](https://github.com/Azure/azure-sdk-for-net/issues/43796)
+
 ## 1.11.2 (2024-04-19)
 
 ### Bugs Fixed

diff --git a/sdk/identity/Azure.Identity/src/Azure.Identity.csproj b/sdk/identity/Azure.Identity/src/Azure.Identity.csproj
@@ -2,9 +2,9 @@
   <PropertyGroup>
     <Description>This is the implementation of the Azure SDK Client Library for Azure Identity</Description>
     <AssemblyTitle>Microsoft Azure.Identity Component</AssemblyTitle>
-    <Version>1.11.2</Version>
+    <Version>1.11.3</Version>
     <!--The ApiCompatVersion is managed automatically and should not generally be modified manually.-->
-    <ApiCompatVersion>1.11.1</ApiCompatVersion>
+    <ApiCompatVersion>1.11.2</ApiCompatVersion>
     <PackageTags>Microsoft Azure Identity;$(PackageCommonTags)</PackageTags>
     <TargetFrameworks>$(RequiredTargetFrameworks)</TargetFrameworks>
     <NoWarn>$(NoWarn);3021;AZC0011</NoWarn>

diff --git a/sdk/identity/Azure.Identity/src/ImdsManagedIdentitySource.cs b/sdk/identity/Azure.Identity/src/ImdsManagedIdentitySource.cs
@@ -91,7 +91,6 @@ protected override HttpMessage CreateHttpMessage(Request request)
             if (_isFirstRequest && _isChainedCredential)
             {
                 message.NetworkTimeout = _imdsNetworkTimeout;
-                _isFirstRequest = false;
             }
 
             return message;
@@ -140,6 +139,9 @@ protected override async ValueTask<AccessToken> HandleResponseAsync(bool async,
             // if we got a response from IMDS we can stop limiting the network timeout
             _imdsNetworkTimeout = null;
 
+            // Mark that the first request has been made
+            _isFirstRequest = false;
+
             // handle error status codes indicating managed identity is not available
             string baseMessage = response.Status switch
             {

diff --git a/sdk/identity/Azure.Identity/tests/ImdsManagedIdentitySourceTests.cs b/sdk/identity/Azure.Identity/tests/ImdsManagedIdentitySourceTests.cs
@@ -62,6 +62,53 @@ public async Task DefaultAzureCredentialProbeUses1secTimeoutWithNoRetries()
             CollectionAssert.AreEqual(expectedTimeouts, networkTimeouts);
         }
 
+        [Test]
+        public async Task DefaultAzureCredentialUsesFirstRequestBehaviorUntilFirstResponse()
+        {
+            int callCount = 0;
+            List<TimeSpan?> networkTimeouts = new();
+
+            // the mock transport succeeds on the 2nd request to avoid long exponential back-offs,
+            // but is sufficient to validate the initial timeout and retry behavior
+            var mockTransport = MockTransport.FromMessageCallback(msg =>
+            {
+                callCount++;
+                networkTimeouts.Add(msg.NetworkTimeout);
+                return callCount switch
+                {
+                    1 => throw new TaskCanceledException(),
+                    2 => CreateMockResponse(400, "Error").WithHeader("Content-Type", "application/json"),
+                    _ => CreateMockResponse(200, "token").WithHeader("Content-Type", "application/json"),
+                };
+            });
+
+            var cred = new DefaultAzureCredential(new DefaultAzureCredentialOptions
+            {
+                ExcludeAzureCliCredential = true,
+                ExcludeAzureDeveloperCliCredential = true,
+                ExcludeAzurePowerShellCredential = true,
+                ExcludeEnvironmentCredential = true,
+                ExcludeSharedTokenCacheCredential = true,
+                ExcludeVisualStudioCodeCredential = true,
+                ExcludeVisualStudioCredential = true,
+                ExcludeWorkloadIdentityCredential = true,
+                Transport = mockTransport
+            });
+
+            //First request times out (throws TaskCancelledException) uses a 1 second timeout and no retries
+            Assert.ThrowsAsync<CredentialUnavailableException>(async () => await cred.GetTokenAsync(new(new[] { "test" })));
+
+            var expectedTimeouts = new TimeSpan?[] { TimeSpan.FromSeconds(1) };
+            CollectionAssert.AreEqual(expectedTimeouts, networkTimeouts);
+            networkTimeouts.Clear();
+
+            // Second request gets the expected probe response and should use the probe timeout on first request and default timeout on the retry
+            await cred.GetTokenAsync(new(new[] { "test" }));
+
+            expectedTimeouts = new TimeSpan?[] { TimeSpan.FromSeconds(1), null };
+            CollectionAssert.AreEqual(expectedTimeouts, networkTimeouts);
+        }
+
         [Test]
         public void DefaultAzureCredentialRetryBehaviorIsOverriddenWithOptions()
         {