Check status code for ManagedIdentityCredential (#511)

Co-authored-by: Joshua Gendein <jgendein@microsoft.com>

Author: JoshGendein

sdk/identity/src/token_credentials/default_credentials.rs

@@ -1,5 +1,5 @@
 use super::{
-    AzureCliCredential, EnvironmentCredential, ManagedIdentityCredential, TokenCredential,
+    AzureCliCredential, EnvironmentCredential, ImdsManagedIdentityCredential, TokenCredential,
 };
 use azure_core::TokenResponse;
 
@@ -47,7 +47,7 @@ impl DefaultAzureCredentialBuilder {
         }
         if self.include_managed_identity_credential {
             sources.push(DefaultAzureCredentialEnum::ManagedIdentity(
-                ManagedIdentityCredential {},
+                ImdsManagedIdentityCredential {},
             ))
         }
         if self.include_cli_credential {
@@ -76,7 +76,7 @@ pub enum DefaultAzureCredentialError {
 /// Types of TokenCredential supported by DefaultAzureCredential
 pub enum DefaultAzureCredentialEnum {
     Environment(EnvironmentCredential),
-    ManagedIdentity(ManagedIdentityCredential),
+    ManagedIdentity(ImdsManagedIdentityCredential),
     AzureCli(AzureCliCredential),
 }
 
@@ -124,7 +124,7 @@ impl Default for DefaultAzureCredential {
         DefaultAzureCredential {
             sources: vec![
                 DefaultAzureCredentialEnum::Environment(EnvironmentCredential::default()),
-                DefaultAzureCredentialEnum::ManagedIdentity(ManagedIdentityCredential {}),
+                DefaultAzureCredentialEnum::ManagedIdentity(ImdsManagedIdentityCredential {}),
                 DefaultAzureCredentialEnum::AzureCli(AzureCliCredential {}),
             ],
         }

sdk/identity/src/token_credentials/managed_identity_credentials.rs renamed to sdk/identity/src/token_credentials/imds_managed_identity_credentials.rs

@@ -18,7 +18,7 @@ const MSI_API_VERSION: &str = "2019-08-01";
 /// This authentication type works in Azure VMs, App Service and Azure Functions applications, as well as the Azure Cloud Shell
 ///
 /// Built up from docs at [https://docs.microsoft.com/azure/app-service/overview-managed-identity#using-the-rest-protocol](https://docs.microsoft.com/azure/app-service/overview-managed-identity#using-the-rest-protocol)
-pub struct ManagedIdentityCredential;
+pub struct ImdsManagedIdentityCredential;
 
 #[non_exhaustive]
 #[derive(Debug, thiserror::Error)]
@@ -32,14 +32,16 @@ pub enum ManagedIdentityCredentialError {
     MissingMsiSecret(std::env::VarError),
     #[error("Refresh token send error: {0}")]
     SendError(reqwest::Error),
-    #[error("Error getting text for refresh token: {0}")]
-    TextError(reqwest::Error),
     #[error("Error deserializing refresh token: {0}")]
-    DeserializeError(serde_json::Error),
+    DeserializeError(reqwest::Error),
+    #[error("The requested identity has not been assigned to this resource.")]
+    IdentityUnavailableError,
+    #[error("The request failed due to a gateway error.")]
+    GatewayError,
 }
 
 #[async_trait::async_trait]
-impl TokenCredential for ManagedIdentityCredential {
+impl TokenCredential for ImdsManagedIdentityCredential {
     type Error = ManagedIdentityCredentialError;
 
     async fn get_token(&self, resource: &str) -> Result<TokenResponse, Self::Error> {
@@ -55,29 +57,33 @@ impl TokenCredential for ManagedIdentityCredential {
             .map_err(ManagedIdentityCredentialError::MissingMsiSecret)?;
 
         let client = reqwest::Client::new();
-        let res_body = client
+        let response = client
             .get(msi_endpoint_url)
             .header("Metadata", "true")
             .header("X-IDENTITY-HEADER", msi_secret)
             .send()
             .await
-            .map_err(ManagedIdentityCredentialError::SendError)?
-            .text()
-            .await
-            .map_err(ManagedIdentityCredentialError::TextError)?;
-
-        let token_response = serde_json::from_str::<MsiTokenResponse>(&res_body)
-            .map_err(ManagedIdentityCredentialError::DeserializeError)?;
-
-        Ok(TokenResponse::new(
-            token_response.access_token,
-            token_response.expires_on,
-        ))
+            .map_err(ManagedIdentityCredentialError::SendError)?;
+
+        match response.status().as_u16() {
+            400 => Err(ManagedIdentityCredentialError::IdentityUnavailableError),
+            502 | 504 => Err(ManagedIdentityCredentialError::GatewayError),
+            _ => {
+                let token_response = response
+                    .json::<MsiTokenResponse>()
+                    .await
+                    .map_err(ManagedIdentityCredentialError::DeserializeError)?;
+                Ok(TokenResponse::new(
+                    token_response.access_token,
+                    token_response.expires_on,
+                ))
+            }
+        }
     }
 }
 
 #[async_trait::async_trait]
-impl azure_core::TokenCredential for ManagedIdentityCredential {
+impl azure_core::TokenCredential for ImdsManagedIdentityCredential {
     async fn get_token(
         &self,
         resource: &str,

sdk/identity/src/token_credentials/mod.rs

@@ -9,13 +9,13 @@ mod cli_credentials;
 mod client_secret_credentials;
 mod default_credentials;
 mod environment_credentials;
-mod managed_identity_credentials;
+mod imds_managed_identity_credentials;
 
 pub use cli_credentials::*;
 pub use client_secret_credentials::*;
 pub use default_credentials::*;
 pub use environment_credentials::*;
-pub use managed_identity_credentials::*;
+pub use imds_managed_identity_credentials::*;
 
 /// Represents a credential capable of providing an OAuth token.
 /// Same as [azure_core::TokenCredential](azure_core::TokenCredential), except a more specific error is returned.