Increase IMDS retry count to 5 (#43249)

Azure · Apr 8, 2024 · 1a1467c · 1a1467c
1 parent 0f3696b
commit 1a1467c
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 14 deletions.
diff --git a/eng/Packages.Data.props b/eng/Packages.Data.props
@@ -153,13 +153,13 @@
     <!-- Other approved packages -->
     <PackageReference Update="Microsoft.Azure.Amqp" Version="2.6.5" />
     <PackageReference Update="Microsoft.Azure.WebPubSub.Common" Version="1.2.0" />
-    <PackageReference Update="Microsoft.Identity.Client" Version="4.59.0" />
-    <PackageReference Update="Microsoft.Identity.Client.Extensions.Msal" Version="4.59.0" />
+    <PackageReference Update="Microsoft.Identity.Client" Version="4.60.1" />
+    <PackageReference Update="Microsoft.Identity.Client.Extensions.Msal" Version="4.60.1" />
     <!--
       TODO: This package needs to be released as GA and arch-board approved before taking a dependency in any stable SDK library.
       Currently, it is referencd by Azure.Identity.Broker which is still in beta
     -->
-    <PackageReference Update="Microsoft.Identity.Client.Broker" Version="4.59.0" />
+    <PackageReference Update="Microsoft.Identity.Client.Broker" Version="4.60.1" />
 
     <!-- TODO: Make sure this package is arch-board approved -->
     <PackageReference Update="Microsoft.IdentityModel.Protocols.OpenIdConnect" Version="6.35.0" />

diff --git a/sdk/identity/Azure.Identity/src/CredentialPipeline.cs b/sdk/identity/Azure.Identity/src/CredentialPipeline.cs
@@ -26,9 +26,29 @@ public CredentialPipeline(HttpPipeline httpPipeline, ClientDiagnostics diagnosti
             Diagnostics = diagnostics;
         }
 
-        public static CredentialPipeline GetInstance(TokenCredentialOptions options)
+        public static CredentialPipeline GetInstance(TokenCredentialOptions options, bool IsManagedIdentityCredential = false)
         {
-            return options is null ? s_singleton.Value : new CredentialPipeline(options);
+            return options switch
+            {
+                _ when IsManagedIdentityCredential => configureOptionsForManagedIdentity(options),
+                not null => new CredentialPipeline(options),
+                _ => s_singleton.Value,
+
+            };
+        }
+
+        private static CredentialPipeline configureOptionsForManagedIdentity(TokenCredentialOptions options)
+        {
+            var clonedOptions = options switch
+            {
+                DefaultAzureCredentialOptions dac => dac.Clone<DefaultAzureCredentialOptions>(),
+                _ => options?.Clone<TokenCredentialOptions>() ?? new TokenCredentialOptions(),
+            };
+            // Set the custom retry policy
+            clonedOptions.Retry.MaxRetries = 5;
+            clonedOptions.RetryPolicy ??= new DefaultAzureCredentialImdsRetryPolicy(clonedOptions.Retry);
+            clonedOptions.IsChainedCredential = clonedOptions is DefaultAzureCredentialOptions;
+            return new CredentialPipeline(clonedOptions);
         }
 
         public HttpPipeline HttpPipeline { get; }

diff --git a/sdk/identity/Azure.Identity/src/Credentials/ManagedIdentityCredential.cs b/sdk/identity/Azure.Identity/src/Credentials/ManagedIdentityCredential.cs
@@ -41,7 +41,7 @@ protected ManagedIdentityCredential()
         /// </param>
         /// <param name="options">Options to configure the management of the requests sent to Microsoft Entra ID.</param>
         public ManagedIdentityCredential(string clientId = null, TokenCredentialOptions options = null)
-            : this(new ManagedIdentityClient(new ManagedIdentityClientOptions { ClientId = clientId, Pipeline = CredentialPipeline.GetInstance(options), Options = options }))
+            : this(new ManagedIdentityClient(new ManagedIdentityClientOptions { ClientId = clientId, Pipeline = CredentialPipeline.GetInstance(options, IsManagedIdentityCredential: true), Options = options }))
         {
             _logAccountDetails = options?.Diagnostics?.IsAccountIdentifierLoggingEnabled ?? false;
         }
@@ -55,7 +55,7 @@ public ManagedIdentityCredential(string clientId = null, TokenCredentialOptions
         /// </param>
         /// <param name="options">Options to configure the management of the requests sent to Microsoft Entra ID.</param>
         public ManagedIdentityCredential(ResourceIdentifier resourceId, TokenCredentialOptions options = null)
-            : this(new ManagedIdentityClient(new ManagedIdentityClientOptions { ResourceIdentifier = resourceId, Pipeline = CredentialPipeline.GetInstance(options), Options = options }))
+            : this(new ManagedIdentityClient(new ManagedIdentityClientOptions { ResourceIdentifier = resourceId, Pipeline = CredentialPipeline.GetInstance(options, IsManagedIdentityCredential: true), Options = options }))
         {
             _logAccountDetails = options?.Diagnostics?.IsAccountIdentifierLoggingEnabled ?? false;
             _clientId = resourceId.ToString();

diff --git a/sdk/identity/Azure.Identity/src/DefaultAzureCredentialFactory.cs b/sdk/identity/Azure.Identity/src/DefaultAzureCredentialFactory.cs
@@ -121,16 +121,13 @@ public virtual TokenCredential CreateWorkloadIdentityCredential()
         public virtual TokenCredential CreateManagedIdentityCredential()
         {
             var options = Options.Clone<DefaultAzureCredentialOptions>();
-            // Set the custom retry policy
-            options.Retry.MaxRetries = 4;
-            options.RetryPolicy ??= new DefaultAzureCredentialImdsRetryPolicy(options.Retry);
             options.IsChainedCredential = true;
 
             var miOptions = new ManagedIdentityClientOptions
             {
                 ResourceIdentifier = options.ManagedIdentityResourceId,
                 ClientId = options.ManagedIdentityClientId,
-                Pipeline = CredentialPipeline.GetInstance(options),
+                Pipeline = CredentialPipeline.GetInstance(options, IsManagedIdentityCredential: true),
                 Options = options,
                 InitialImdsConnectionTimeout = TimeSpan.FromSeconds(1),
                 ExcludeTokenExchangeManagedIdentitySource = options.ExcludeWorkloadIdentityCredential

diff --git a/sdk/identity/Azure.Identity/tests/ImdsManagedIdentitySourceTests.cs b/sdk/identity/Azure.Identity/tests/ImdsManagedIdentitySourceTests.cs
@@ -120,7 +120,7 @@ public void ManagedIdentityCredentialUsesDefaultTimeoutAndRetries()
 
             Assert.ThrowsAsync<AuthenticationFailedException>(async () => await cred.GetTokenAsync(new(new[] { "test" })));
 
-            var expectedTimeouts = new TimeSpan?[] { null, null, null, null };
+            var expectedTimeouts = new TimeSpan?[] { null, null, null, null, null, null };
             CollectionAssert.AreEqual(expectedTimeouts, networkTimeouts);
         }
 

diff --git a/sdk/identity/Azure.Identity/tests/ManagedIdentityCredentialTests.cs b/sdk/identity/Azure.Identity/tests/ManagedIdentityCredentialTests.cs
@@ -683,13 +683,13 @@ public async Task VerifyMsiUnavailableOnIMDSRequestFailedExcpetion()
         {
             using var environment = new TestEnvVar(new() { { "MSI_ENDPOINT", null }, { "MSI_SECRET", null }, { "IDENTITY_ENDPOINT", null }, { "IDENTITY_HEADER", null }, { "AZURE_POD_IDENTITY_AUTHORITY_HOST", "http://169.254.169.001/" } });
 
-            var options = new TokenCredentialOptions() { Retry = { MaxRetries = 0, NetworkTimeout = TimeSpan.FromMilliseconds(100) } };
+            var options = new TokenCredentialOptions() { Retry = { MaxRetries = 0, NetworkTimeout = TimeSpan.FromMilliseconds(100), MaxDelay = TimeSpan.Zero } };
 
             var credential = InstrumentClient(new ManagedIdentityCredential(options: options));
 
             var ex = Assert.ThrowsAsync<CredentialUnavailableException>(async () => await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default)));
 
-            Assert.That(ex.Message, Does.Contain(ImdsManagedIdentitySource.NoResponseError));
+            Assert.That(ex.Message, Does.Contain(ImdsManagedIdentitySource.AggregateError));
 
             await Task.CompletedTask;
         }