From 0b41a87afc2521f697c7351663bcadfc68290726 Mon Sep 17 00:00:00 2001 From: Charles Lowell <10964656+chlowell@users.noreply.github.com> Date: Fri, 17 Oct 2025 23:53:32 +0000 Subject: [PATCH 1/2] Implement retry logic for IMDS managed identity --- sdk/identity/azure_identity/CHANGELOG.md | 1 + ...app_service_managed_identity_credential.rs | 1 + .../src/imds_managed_identity_credential.rs | 6 ++- ...ual_machine_managed_identity_credential.rs | 38 +++++++++++++++++-- 4 files changed, 41 insertions(+), 5 deletions(-) diff --git a/sdk/identity/azure_identity/CHANGELOG.md b/sdk/identity/azure_identity/CHANGELOG.md index a96f377ca7..232db07900 100644 --- a/sdk/identity/azure_identity/CHANGELOG.md +++ b/sdk/identity/azure_identity/CHANGELOG.md @@ -15,6 +15,7 @@ ### Bugs Fixed - `ClientCertificateCredential::get_token()` returned an error when given multiple scopes +- `ManagedIdentityCredential` didn't follow IMDS retry guidance ### Other Changes diff --git a/sdk/identity/azure_identity/src/app_service_managed_identity_credential.rs b/sdk/identity/azure_identity/src/app_service_managed_identity_credential.rs index 5fb09e8d4a..edbb5eca64 100644 --- a/sdk/identity/azure_identity/src/app_service_managed_identity_credential.rs +++ b/sdk/identity/azure_identity/src/app_service_managed_identity_credential.rs @@ -48,6 +48,7 @@ impl AppServiceManagedIdentityCredential { SECRET_ENV, id, client_options, + None, env, ), })) diff --git a/sdk/identity/azure_identity/src/imds_managed_identity_credential.rs b/sdk/identity/azure_identity/src/imds_managed_identity_credential.rs index dcf38fe2c7..67a4444bcb 100644 --- a/sdk/identity/azure_identity/src/imds_managed_identity_credential.rs +++ b/sdk/identity/azure_identity/src/imds_managed_identity_credential.rs @@ -6,7 +6,7 @@ use azure_core::{ credentials::{AccessToken, Secret, TokenCredential, TokenRequestOptions}, error::{Error, ErrorKind}, http::{ - headers::HeaderName, request::Request, ClientOptions, Method, Pipeline, + headers::HeaderName, request::Request, ClientOptions, Method, Pipeline, PipelineOptions, PipelineSendOptions, StatusCode, Url, }, json::from_json, @@ -62,6 +62,7 @@ pub(crate) struct ImdsManagedIdentityCredential { } impl ImdsManagedIdentityCredential { + #[allow(clippy::too_many_arguments, reason = "private API")] pub fn new( endpoint: Url, api_version: &str, @@ -69,6 +70,7 @@ impl ImdsManagedIdentityCredential { secret_env: &str, id: ImdsId, client_options: ClientOptions, + pipeline_options: Option, env: Env, ) -> Self { let pipeline = Pipeline::new( @@ -77,7 +79,7 @@ impl ImdsManagedIdentityCredential { client_options, Vec::default(), Vec::default(), - None, + pipeline_options, ); Self { pipeline, diff --git a/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs b/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs index 988d40fa97..479c1fee54 100644 --- a/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs +++ b/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs @@ -3,10 +3,13 @@ use crate::env::Env; use crate::{ImdsId, ImdsManagedIdentityCredential}; -use azure_core::http::ClientOptions; use azure_core::{ credentials::{AccessToken, TokenCredential, TokenRequestOptions}, - http::{headers::HeaderName, Url}, + http::{ + headers::HeaderName, ClientOptions, ExponentialRetryOptions, PipelineOptions, RetryOptions, + StatusCode, Url, + }, + time::Duration, }; use std::sync::Arc; @@ -23,10 +26,38 @@ pub struct VirtualMachineManagedIdentityCredential { impl VirtualMachineManagedIdentityCredential { pub fn new( id: ImdsId, - client_options: ClientOptions, + mut client_options: ClientOptions, env: Env, ) -> azure_core::Result> { let endpoint = Url::parse(ENDPOINT).unwrap(); // valid url constant + let pipeline_options = Some(PipelineOptions { + // https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#error-handling + retry_status_codes: Vec::from([ + StatusCode::NotFound, + StatusCode::Gone, + StatusCode::TooManyRequests, + StatusCode::InternalServerError, + StatusCode::NotImplemented, + StatusCode::BadGateway, + StatusCode::ServiceUnavailable, + StatusCode::GatewayTimeout, + StatusCode::HttpVersionNotSupported, + StatusCode::VariantAlsoNegotiates, + StatusCode::InsufficientStorage, + StatusCode::LoopDetected, + StatusCode::NotExtended, + StatusCode::NetworkAuthenticationRequired, + ]), + ..Default::default() + }); + // these settings approximate the recommendations at + // https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#retry-guidance + client_options.retry = RetryOptions::exponential(ExponentialRetryOptions { + initial_delay: Duration::milliseconds(1340), + max_retries: 6, + max_total_elapsed: Duration::seconds(72), + ..Default::default() + }); Ok(Arc::new(Self { credential: ImdsManagedIdentityCredential::new( endpoint, @@ -35,6 +66,7 @@ impl VirtualMachineManagedIdentityCredential { SECRET_ENV, id, client_options, + pipeline_options, env, ), })) From bd23b9deaa78322b3e4b303540be61d002c15646 Mon Sep 17 00:00:00 2001 From: Charles Lowell <10964656+chlowell@users.noreply.github.com> Date: Mon, 10 Nov 2025 18:48:11 +0000 Subject: [PATCH 2/2] remove mut --- ...rtual_machine_managed_identity_credential.rs | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs b/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs index 479c1fee54..e4c5256d84 100644 --- a/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs +++ b/sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs @@ -26,7 +26,7 @@ pub struct VirtualMachineManagedIdentityCredential { impl VirtualMachineManagedIdentityCredential { pub fn new( id: ImdsId, - mut client_options: ClientOptions, + client_options: ClientOptions, env: Env, ) -> azure_core::Result> { let endpoint = Url::parse(ENDPOINT).unwrap(); // valid url constant @@ -52,12 +52,15 @@ impl VirtualMachineManagedIdentityCredential { }); // these settings approximate the recommendations at // https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#retry-guidance - client_options.retry = RetryOptions::exponential(ExponentialRetryOptions { - initial_delay: Duration::milliseconds(1340), - max_retries: 6, - max_total_elapsed: Duration::seconds(72), - ..Default::default() - }); + let client_options = ClientOptions { + retry: RetryOptions::exponential(ExponentialRetryOptions { + initial_delay: Duration::milliseconds(1340), + max_retries: 6, + max_total_elapsed: Duration::seconds(72), + ..Default::default() + }), + ..client_options + }; Ok(Arc::new(Self { credential: ImdsManagedIdentityCredential::new( endpoint,