Skip to content

Commit

Permalink
Increase IMDS retry count to 5 (#43249)
Browse files Browse the repository at this point in the history
  • Loading branch information
christothes committed Apr 8, 2024
1 parent 0f3696b commit 1a1467c
Show file tree
Hide file tree
Showing 6 changed files with 31 additions and 14 deletions.
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

0 comments on commit 1a1467c

Please sign in to comment.