Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Increase IMDS retry count to 5 #43249

Merged
merged 3 commits into from
Apr 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions eng/Packages.Data.props
Original file line number Diff line number Diff line change
Expand Up @@ -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" />
Expand Down
24 changes: 22 additions & 2 deletions sdk/identity/Azure.Identity/src/CredentialPipeline.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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; }
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down