add timeout to IMDS when used in DefaultAzureCredential (#1016)

* 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.

* address feedback

* 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.

Co-authored-by: Brian Caswell <[email protected]>

Author: bmc-msft

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"
+pin-project = "1.0"
 
 [dev-dependencies]
 reqwest = { version = "0.11", features = ["json"], default-features = false }

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::*;

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<F, D> {
+    #[pin]
+    future: F,
+    #[pin]
+    deadline: D,
+    completed: bool,
+}
+
+impl<F, D> Timeout<F, D> {
+    pub(crate) fn new(future: F, deadline: D) -> Self {
+        Self {
+            future,
+            deadline,
+            completed: false,
+        }
+    }
+}
+
+impl<F: Future, D: Future> Future for Timeout<F, D> {
+    type Output = azure_core::Result<F::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        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<Self, Sleep>
+    where
+        Self: Sized,
+    {
+        Timeout::new(self, sleep(duration))
+    }
+}
+
+impl<T> TimeoutExt for T where T: Future {}

sdk/identity/src/token_credentials/default_credentials.rs

@@ -1,6 +1,12 @@
-use super::{AzureCliCredential, ImdsManagedIdentityCredential};
-use azure_core::auth::{TokenCredential, TokenResponse};
-use azure_core::error::{Error, ErrorKind, ResultExt};
+use crate::{
+    timeout::TimeoutExt,
+    {AzureCliCredential, ImdsManagedIdentityCredential},
+};
+use azure_core::{
+    auth::{TokenCredential, TokenResponse},
+    error::{Error, ErrorKind, ResultExt},
+};
+use std::time::Duration;
 
 #[derive(Debug)]
 /// Provides a mechanism of selectively disabling credentials used for a `DefaultAzureCredential` instance
@@ -91,10 +97,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(
@@ -128,15 +143,7 @@ impl DefaultAzureCredential {
 
 impl Default for DefaultAzureCredential {
     fn default() -> Self {
-        DefaultAzureCredential {
-            sources: vec![
-                DefaultAzureCredentialEnum::Environment(super::EnvironmentCredential::default()),
-                DefaultAzureCredentialEnum::ManagedIdentity(
-                    ImdsManagedIdentityCredential::default(),
-                ),
-                DefaultAzureCredentialEnum::AzureCli(AzureCliCredential::new()),
-            ],
-        }
+        DefaultAzureCredentialBuilder::new().build()
     }
 }