From f313a6afb7a46be49fc7c5dbd8fcb1054f33dc7d Mon Sep 17 00:00:00 2001 From: Brian Caswell Date: Tue, 16 Aug 2022 15:06:44 -0400 Subject: [PATCH 1/3] make checking IMDS for tokens happen last in DefaultAzureCredential The DefaultAzureCredential iterates through multiple methods for obtaining a credential until one succeeds (or all fail). Currently, the default is to check the Environment, IMDS, then the CLI. In many non-Azure workloads (such as MSFT corp computers), the IMDS token request takes an exceedingly long time to fail. This makes using DefaultAzureCredential without disabling IMDS on corp machines painfully slow. This PR changes the order to Environment, CLI, then IMDS, which is beneficial for developers using the examples. --- .../token_credentials/default_credentials.rs | 23 ++++++------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/sdk/identity/src/token_credentials/default_credentials.rs b/sdk/identity/src/token_credentials/default_credentials.rs index d01b871579..a034c04d28 100644 --- a/sdk/identity/src/token_credentials/default_credentials.rs +++ b/sdk/identity/src/token_credentials/default_credentials.rs @@ -56,16 +56,16 @@ impl DefaultAzureCredentialBuilder { super::EnvironmentCredential::default(), )); } - if self.include_managed_identity_credential { - sources.push(DefaultAzureCredentialEnum::ManagedIdentity( - ImdsManagedIdentityCredential::default(), - )) - } if self.include_azure_cli_credential { sources.push(DefaultAzureCredentialEnum::AzureCli( AzureCliCredential::new(), )); } + if self.include_managed_identity_credential { + sources.push(DefaultAzureCredentialEnum::ManagedIdentity( + ImdsManagedIdentityCredential::default(), + )) + } DefaultAzureCredential::with_sources(sources) } } @@ -113,8 +113,8 @@ impl TokenCredential for DefaultAzureCredentialEnum { /// /// The following credential types if enabled will be tried, in order: /// - EnvironmentCredential -/// - ManagedIdentityCredential /// - AzureCliCredential +/// - ManagedIdentityCredential /// Consult the documentation of these credential types for more information on how they attempt authentication. pub struct DefaultAzureCredential { sources: Vec, @@ -131,16 +131,7 @@ impl DefaultAzureCredential { impl Default for DefaultAzureCredential { fn default() -> Self { - DefaultAzureCredential { - sources: vec![ - #[cfg(feature = "enable_reqwest")] - DefaultAzureCredentialEnum::Environment(super::EnvironmentCredential::default()), - DefaultAzureCredentialEnum::ManagedIdentity( - ImdsManagedIdentityCredential::default(), - ), - DefaultAzureCredentialEnum::AzureCli(AzureCliCredential::new()), - ], - } + DefaultAzureCredentialBuilder::new().build() } } From e376633430c6ddf48d3bb814cbd0371c7c7cb051 Mon Sep 17 00:00:00 2001 From: Brian Caswell Date: Mon, 28 Nov 2022 11:42:09 -0500 Subject: [PATCH 2/3] address feedback --- sdk/identity/Cargo.toml | 1 + .../token_credentials/default_credentials.rs | 36 ++++++++++++------- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/sdk/identity/Cargo.toml b/sdk/identity/Cargo.toml index cb6246ad39..76cf2c24be 100644 --- a/sdk/identity/Cargo.toml +++ b/sdk/identity/Cargo.toml @@ -28,6 +28,7 @@ base64 = "0.13.0" uuid = { version = "1.0", features = ["v4"] } # work around https://github.com/rust-lang/rust/issues/63033 fix-hidden-lifetime-bug = "0.2" +futures-time = "3.0" [dev-dependencies] reqwest = { version = "0.11", features = ["json"], default-features = false } diff --git a/sdk/identity/src/token_credentials/default_credentials.rs b/sdk/identity/src/token_credentials/default_credentials.rs index 8154144cfd..3352195c92 100644 --- a/sdk/identity/src/token_credentials/default_credentials.rs +++ b/sdk/identity/src/token_credentials/default_credentials.rs @@ -1,6 +1,9 @@ use super::{AzureCliCredential, ImdsManagedIdentityCredential}; -use azure_core::auth::{TokenCredential, TokenResponse}; -use azure_core::error::{Error, ErrorKind, ResultExt}; +use azure_core::{ + auth::{TokenCredential, TokenResponse}, + error::{Error, ErrorKind, ResultExt}, +}; +use futures_time::{future::FutureExt, time::Duration}; #[derive(Debug)] /// Provides a mechanism of selectively disabling credentials used for a `DefaultAzureCredential` instance @@ -55,16 +58,16 @@ impl DefaultAzureCredentialBuilder { super::EnvironmentCredential::default(), )); } - if self.include_azure_cli_credential { - sources.push(DefaultAzureCredentialEnum::AzureCli( - AzureCliCredential::new(), - )); - } if self.include_managed_identity_credential { sources.push(DefaultAzureCredentialEnum::ManagedIdentity( ImdsManagedIdentityCredential::default(), )) } + if self.include_azure_cli_credential { + sources.push(DefaultAzureCredentialEnum::AzureCli( + AzureCliCredential::new(), + )); + } DefaultAzureCredential::with_sources(sources) } } @@ -91,10 +94,19 @@ impl TokenCredential for DefaultAzureCredentialEnum { ) } DefaultAzureCredentialEnum::ManagedIdentity(credential) => { - credential.get_token(resource).await.context( - ErrorKind::Credential, - "error getting managed identity credential", - ) + // IMSD timeout is only limited to 1 second when used in DefaultAzureCredential + credential + .get_token(resource) + .timeout(Duration::from_secs(1)) + .await + .context( + ErrorKind::Credential, + "getting managed identity credential timed out", + )? + .context( + ErrorKind::Credential, + "error getting managed identity credential", + ) } DefaultAzureCredentialEnum::AzureCli(credential) => { credential.get_token(resource).await.context( @@ -110,8 +122,8 @@ impl TokenCredential for DefaultAzureCredentialEnum { /// /// The following credential types if enabled will be tried, in order: /// - EnvironmentCredential -/// - AzureCliCredential /// - ManagedIdentityCredential +/// - AzureCliCredential /// Consult the documentation of these credential types for more information on how they attempt authentication. pub struct DefaultAzureCredential { sources: Vec, From b6793a2b93c16d342684e78264e11a4762396ebc Mon Sep 17 00:00:00 2001 From: Brian Caswell Date: Mon, 28 Nov 2022 17:34:51 -0500 Subject: [PATCH 3/3] use azure-core and a minimal future rather than futures-time Rather than expanding our dependencies, this uses azure-core's sleep to provide the future used for timing out and lifts the future from futures-time. --- sdk/identity/Cargo.toml | 2 +- sdk/identity/src/lib.rs | 1 + sdk/identity/src/timeout.rs | 70 +++++++++++++++++++ .../token_credentials/default_credentials.rs | 7 +- 4 files changed, 77 insertions(+), 3 deletions(-) create mode 100644 sdk/identity/src/timeout.rs diff --git a/sdk/identity/Cargo.toml b/sdk/identity/Cargo.toml index 76cf2c24be..a1bf4af43f 100644 --- a/sdk/identity/Cargo.toml +++ b/sdk/identity/Cargo.toml @@ -28,7 +28,7 @@ base64 = "0.13.0" uuid = { version = "1.0", features = ["v4"] } # work around https://github.com/rust-lang/rust/issues/63033 fix-hidden-lifetime-bug = "0.2" -futures-time = "3.0" +pin-project = "1.0" [dev-dependencies] reqwest = { version = "0.11", features = ["json"], default-features = false } diff --git a/sdk/identity/src/lib.rs b/sdk/identity/src/lib.rs index 1642b2371f..82427b4b34 100644 --- a/sdk/identity/src/lib.rs +++ b/sdk/identity/src/lib.rs @@ -49,6 +49,7 @@ pub mod development; pub mod device_code_flow; mod oauth2_http_client; pub mod refresh_token; +mod timeout; mod token_credentials; pub use crate::token_credentials::*; diff --git a/sdk/identity/src/timeout.rs b/sdk/identity/src/timeout.rs new file mode 100644 index 0000000000..20ec363301 --- /dev/null +++ b/sdk/identity/src/timeout.rs @@ -0,0 +1,70 @@ +// Copyright (c) 2020 Yoshua Wuyts +// +// based on https://crates.io/crates/futures-time +// Licensed under either of Apache License, Version 2.0 or MIT license at your option. + +use azure_core::sleep::{sleep, Sleep}; +use futures::Future; +use std::time::Duration; +use std::{ + pin::Pin, + task::{Context, Poll}, +}; + +#[pin_project::pin_project] +#[derive(Debug)] +pub(crate) struct Timeout { + #[pin] + future: F, + #[pin] + deadline: D, + completed: bool, +} + +impl Timeout { + pub(crate) fn new(future: F, deadline: D) -> Self { + Self { + future, + deadline, + completed: false, + } + } +} + +impl Future for Timeout { + type Output = azure_core::Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + + assert!(!*this.completed, "future polled after completing"); + + match this.future.poll(cx) { + Poll::Ready(v) => { + *this.completed = true; + Poll::Ready(Ok(v)) + } + Poll::Pending => match this.deadline.poll(cx) { + Poll::Ready(_) => { + *this.completed = true; + Poll::Ready(Err(azure_core::error::Error::with_message( + azure_core::error::ErrorKind::Other, + || String::from("operation timed out"), + ))) + } + Poll::Pending => Poll::Pending, + }, + } + } +} + +pub(crate) trait TimeoutExt: Future { + fn timeout(self, duration: Duration) -> Timeout + where + Self: Sized, + { + Timeout::new(self, sleep(duration)) + } +} + +impl TimeoutExt for T where T: Future {} diff --git a/sdk/identity/src/token_credentials/default_credentials.rs b/sdk/identity/src/token_credentials/default_credentials.rs index 3352195c92..75ad679263 100644 --- a/sdk/identity/src/token_credentials/default_credentials.rs +++ b/sdk/identity/src/token_credentials/default_credentials.rs @@ -1,9 +1,12 @@ -use super::{AzureCliCredential, ImdsManagedIdentityCredential}; +use crate::{ + timeout::TimeoutExt, + {AzureCliCredential, ImdsManagedIdentityCredential}, +}; use azure_core::{ auth::{TokenCredential, TokenResponse}, error::{Error, ErrorKind, ResultExt}, }; -use futures_time::{future::FutureExt, time::Duration}; +use std::time::Duration; #[derive(Debug)] /// Provides a mechanism of selectively disabling credentials used for a `DefaultAzureCredential` instance