From 06dd672e59d9bd6068e8ef282cb94ae17e6f75be Mon Sep 17 00:00:00 2001
From: Christopher Scott <chriscott@hotmail.com>
Date: Mon, 6 May 2024 17:16:54 -0500
Subject: [PATCH] Hotfix for DAC probe request

---
 sdk/identity/Azure.Identity/CHANGELOG.md      |  5 ++
 .../Azure.Identity/src/Azure.Identity.csproj  |  4 +-
 .../src/ImdsManagedIdentitySource.cs          |  4 +-
 .../tests/ImdsManagedIdentitySourceTests.cs   | 47 +++++++++++++++++++
 4 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/sdk/identity/Azure.Identity/CHANGELOG.md b/sdk/identity/Azure.Identity/CHANGELOG.md
index eec46bff925e..49f853c35971 100644
--- 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
index 2c56219422eb..858379b84b2b 100644
--- 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
index 4af0f42d9272..2fd0640f1368 100644
--- 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
index 0cbe779d0bac..8af14d281761 100644
--- 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()
         {