object_store: azure cli authorization (#3698)

roeap · tustvold · web-flow · commit e37e379f158c · 2023-02-13T14:40:16.000Z
* fix: pass bearer token credential as auth header

* feat: add azure cli credential

* fix: clippy

* Update object_store/src/azure/client.rs

Co-authored-by: Raphael Taylor-Davies &lt;1781103+tustvold@users.noreply.github.com&gt;

* chore: PR feedback

* docs: add azure cli link

---------

Co-authored-by: Raphael Taylor-Davies &lt;1781103+tustvold@users.noreply.github.com&gt;
diff --git a/object_store/src/azure/client.rs b/object_store/src/azure/client.rs
@@ -169,6 +169,18 @@ impl AzureClient {
             CredentialProvider::AccessKey(key) => {
                 Ok(AzureCredential::AccessKey(key.to_owned()))
             }
+            CredentialProvider::BearerToken(token) => {
+                Ok(AzureCredential::AuthorizationToken(
+                    // we do the conversion to a HeaderValue here, since it is fallible
+                    // and we want to use it in an infallible function
+                    HeaderValue::from_str(&format!("Bearer {token}")).map_err(|err| {
+                        crate::Error::Generic {
+                            store: "MicrosoftAzure",
+                            source: Box::new(err),
+                        }
+                    })?,
+                ))
+            }
             CredentialProvider::TokenCredential(cache, cred) => {
                 let token = cache
                     .get_or_insert_with(|| {
@@ -178,7 +190,7 @@ impl AzureClient {
                     .context(AuthorizationSnafu)?;
                 Ok(AzureCredential::AuthorizationToken(
                     // we do the conversion to a HeaderValue here, since it is fallible
-                    // and we wna to use it in an infallible function
+                    // and we want to use it in an infallible function
                     HeaderValue::from_str(&format!("Bearer {token}")).map_err(|err| {
                         crate::Error::Generic {
                             store: "MicrosoftAzure",
diff --git a/object_store/src/azure/credential.rs b/object_store/src/azure/credential.rs
@@ -21,7 +21,7 @@ use crate::util::hmac_sha256;
 use crate::RetryConfig;
 use base64::prelude::BASE64_STANDARD;
 use base64::Engine;
-use chrono::Utc;
+use chrono::{DateTime, Utc};
 use reqwest::header::ACCEPT;
 use reqwest::{
     header::{
@@ -34,6 +34,7 @@ use reqwest::{
 use serde::Deserialize;
 use snafu::{ResultExt, Snafu};
 use std::borrow::Cow;
+use std::process::Command;
 use std::str;
 use std::time::{Duration, Instant};
 use url::Url;
@@ -61,6 +62,12 @@ pub enum Error {
 
     #[snafu(display("Error reading federated token file "))]
     FederatedTokenFile,
+
+    #[snafu(display("'az account get-access-token' command failed: {message}"))]
+    AzureCli { message: String },
+
+    #[snafu(display("Failed to parse azure cli response: {source}"))]
+    AzureCliResponse { source: serde_json::Error },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -69,6 +76,7 @@ pub type Result<T, E = Error> = std::result::Result<T, E>;
 #[derive(Debug)]
 pub enum CredentialProvider {
     AccessKey(String),
+    BearerToken(String),
     SASToken(Vec<(String, String)>),
     TokenCredential(TokenCache<String>, Box<dyn TokenCredential>),
 }
@@ -540,6 +548,122 @@ impl TokenCredential for WorkloadIdentityOAuthProvider {
     }
 }
 
+mod az_cli_date_format {
+    use chrono::{DateTime, TimeZone};
+    use serde::{self, Deserialize, Deserializer};
+
+    pub fn deserialize<'de, D>(
+        deserializer: D,
+    ) -> Result<DateTime<chrono::Local>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let s = String::deserialize(deserializer)?;
+        // expiresOn from azure cli uses the local timezone
+        let date = chrono::NaiveDateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S.%6f")
+            .map_err(serde::de::Error::custom)?;
+        chrono::Local
+            .from_local_datetime(&date)
+            .single()
+            .ok_or(serde::de::Error::custom(
+                "azure cli returned ambiguous expiry date",
+            ))
+    }
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct AzureCliTokenResponse {
+    pub access_token: String,
+    #[serde(with = "az_cli_date_format")]
+    pub expires_on: DateTime<chrono::Local>,
+    pub token_type: String,
+}
+
+#[derive(Default, Debug)]
+pub struct AzureCliCredential {
+    _private: (),
+}
+
+impl AzureCliCredential {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
+#[async_trait::async_trait]
+impl TokenCredential for AzureCliCredential {
+    /// Fetch a token
+    async fn fetch_token(
+        &self,
+        _client: &Client,
+        _retry: &RetryConfig,
+    ) -> Result<TemporaryToken<String>> {
+        // on window az is a cmd and it should be called like this
+        // see https://doc.rust-lang.org/nightly/std/process/struct.Command.html
+        let program = if cfg!(target_os = "windows") {
+            "cmd"
+        } else {
+            "az"
+        };
+        let mut args = Vec::new();
+        if cfg!(target_os = "windows") {
+            args.push("/C");
+            args.push("az");
+        }
+        args.push("account");
+        args.push("get-access-token");
+        args.push("--output");
+        args.push("json");
+        args.push("--scope");
+        args.push(AZURE_STORAGE_SCOPE);
+
+        match Command::new(program).args(args).output() {
+            Ok(az_output) if az_output.status.success() => {
+                let output =
+                    str::from_utf8(&az_output.stdout).map_err(|_| Error::AzureCli {
+                        message: "az response is not a valid utf-8 string".to_string(),
+                    })?;
+
+                let token_response =
+                    serde_json::from_str::<AzureCliTokenResponse>(output)
+                        .context(AzureCliResponseSnafu)?;
+                if !token_response.token_type.eq_ignore_ascii_case("bearer") {
+                    return Err(Error::AzureCli {
+                        message: format!(
+                            "got unexpected token type from azure cli: {0}",
+                            token_response.token_type
+                        ),
+                    });
+                }
+                let duration = token_response.expires_on.naive_local()
+                    - chrono::Local::now().naive_local();
+                Ok(TemporaryToken {
+                    token: token_response.access_token,
+                    expiry: Instant::now()
+                        + duration.to_std().map_err(|_| Error::AzureCli {
+                            message: "az returned invalid lifetime".to_string(),
+                        })?,
+                })
+            }
+            Ok(az_output) => {
+                let message = String::from_utf8_lossy(&az_output.stderr);
+                Err(Error::AzureCli {
+                    message: message.into(),
+                })
+            }
+            Err(e) => match e.kind() {
+                std::io::ErrorKind::NotFound => Err(Error::AzureCli {
+                    message: "Azure Cli not installed".into(),
+                }),
+                error_kind => Err(Error::AzureCli {
+                    message: format!("io error: {error_kind:?}"),
+                }),
+            },
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/object_store/src/azure/mod.rs b/object_store/src/azure/mod.rs
@@ -400,6 +400,7 @@ pub struct MicrosoftAzureBuilder {
     object_id: Option<String>,
     msi_resource_id: Option<String>,
     federated_token_file: Option<String>,
+    use_azure_cli: bool,
     retry_config: RetryConfig,
     client_options: ClientOptions,
 }
@@ -533,6 +534,13 @@ pub enum AzureConfigKey {
     /// - `azure_federated_token_file`
     /// - `federated_token_file`
     FederatedTokenFile,
+
+    /// Use azure cli for acquiring access token
+    ///
+    /// Supported keys:
+    /// - `azure_use_azure_cli`
+    /// - `use_azure_cli`
+    UseAzureCli,
 }
 
 impl AsRef<str> for AzureConfigKey {
@@ -550,6 +558,7 @@ impl AsRef<str> for AzureConfigKey {
             Self::ObjectId => "azure_object_id",
             Self::MsiResourceId => "azure_msi_resource_id",
             Self::FederatedTokenFile => "azure_federated_token_file",
+            Self::UseAzureCli => "azure_use_azure_cli",
         }
     }
 }
@@ -593,6 +602,7 @@ impl FromStr for AzureConfigKey {
             "azure_federated_token_file" | "federated_token_file" => {
                 Ok(Self::FederatedTokenFile)
             }
+            "azure_use_azure_cli" | "use_azure_cli" => Ok(Self::UseAzureCli),
             _ => Err(Error::UnknownConfigurationKey { key: s.into() }.into()),
         }
     }
@@ -704,6 +714,9 @@ impl MicrosoftAzureBuilder {
             AzureConfigKey::FederatedTokenFile => {
                 self.federated_token_file = Some(value.into())
             }
+            AzureConfigKey::UseAzureCli => {
+                self.use_azure_cli = str_is_truthy(&value.into())
+            }
             AzureConfigKey::UseEmulator => {
                 self.use_emulator = str_is_truthy(&value.into())
             }
@@ -887,6 +900,13 @@ impl MicrosoftAzureBuilder {
         self
     }
 
+    /// Set if the Azure Cli should be used for acquiring access token
+    /// <https://learn.microsoft.com/en-us/cli/azure/account?view=azure-cli-latest#az-account-get-access-token>
+    pub fn with_use_azure_cli(mut self, use_azure_cli: bool) -> Self {
+        self.use_azure_cli = use_azure_cli;
+        self
+    }
+
     /// Configure a connection to container with given name on Microsoft Azure
     /// Blob store.
     pub fn build(mut self) -> Result<MicrosoftAzure> {
@@ -916,7 +936,7 @@ impl MicrosoftAzureBuilder {
             let url = Url::parse(&account_url)
                 .context(UnableToParseUrlSnafu { url: account_url })?;
             let credential = if let Some(bearer_token) = self.bearer_token {
-                credential::CredentialProvider::AccessKey(bearer_token)
+                credential::CredentialProvider::BearerToken(bearer_token)
             } else if let Some(access_key) = self.access_key {
                 credential::CredentialProvider::AccessKey(access_key)
             } else if let (Some(client_id), Some(tenant_id), Some(federated_token_file)) =
@@ -949,6 +969,11 @@ impl MicrosoftAzureBuilder {
                 credential::CredentialProvider::SASToken(query_pairs)
             } else if let Some(sas) = self.sas_key {
                 credential::CredentialProvider::SASToken(split_sas(&sas)?)
+            } else if self.use_azure_cli {
+                credential::CredentialProvider::TokenCredential(
+                    TokenCache::default(),
+                    Box::new(credential::AzureCliCredential::new()),
+                )
             } else {
                 let client =
                     self.client_options.clone().with_allow_http(true).client()?;
