From de46128e22c3e1b345c952b0dae3a566a44d5165 Mon Sep 17 00:00:00 2001 From: congyi <15605187270@163.com> Date: Mon, 18 Mar 2024 15:35:32 +0800 Subject: [PATCH 1/2] support workload identity auth for Azure --- src/azure/storage/config.rs | 18 ++++ src/azure/storage/loader.rs | 18 +++- src/azure/storage/mod.rs | 2 + .../storage/workload_identity_credential.rs | 85 +++++++++++++++++++ 4 files changed, 122 insertions(+), 1 deletion(-) create mode 100644 src/azure/storage/workload_identity_credential.rs diff --git a/src/azure/storage/config.rs b/src/azure/storage/config.rs index fac6d7a1..667fc184 100644 --- a/src/azure/storage/config.rs +++ b/src/azure/storage/config.rs @@ -44,4 +44,22 @@ pub struct Config { /// /// This is part of use AAD(Azure Active Directory) authenticate on Azure VM pub endpoint: Option, + /// `azure_federated_token` value will be loaded from: + /// + /// - this field if it's `is_some` + /// - env value: [`AZURE_FEDERATED_TOKEN`] + /// - profile config: `azure_federated_token_file` + pub azure_federated_token: Option, + /// `azure_federated_token_file` value will be loaded from: + /// + /// - this field if it's `is_some` + /// - env value: [`AZURE_FEDERATED_TOKEN_FILE`] + /// - profile config: `azure_federated_token_file` + pub azure_federated_token_file: Option, + /// `azure_tenant_id_env_key` value will be loaded from: + /// + /// - this field if it's `is_some` + /// - env value: [`AZURE_TENANT_ID_ENV_KEY`] + /// - profile config: `azure_tenant_id_env_key` + pub azure_tenant_id_env_key: Option, } diff --git a/src/azure/storage/loader.rs b/src/azure/storage/loader.rs index 5a445bcc..b0b29d36 100644 --- a/src/azure/storage/loader.rs +++ b/src/azure/storage/loader.rs @@ -3,9 +3,9 @@ use std::sync::Mutex; use anyhow::Result; -use super::config::Config; use super::credential::Credential; use super::imds_credential; +use super::{config::Config, workload_identity_credential}; /// Loader will load credential from different methods. #[cfg_attr(test, derive(Debug))] @@ -45,6 +45,10 @@ impl Loader { return Ok(Some(cred)); } + if let Some(cred) = self.load_via_workload_identity().await? { + return Ok(Some(cred)); + } + // try to load credential using AAD(Azure Active Directory) authenticate on Azure VM // we may get an error if not running on Azure VM // see https://learn.microsoft.com/en-us/azure/app-service/overview-managed-identity?tabs=portal,http#using-the-rest-protocol @@ -72,4 +76,16 @@ impl Loader { Ok(cred) } + + async fn load_via_workload_identity(&self) -> Result> { + let workload_identity_token = workload_identity_credential::get_workload_identity_token( + "https://storage.azure.com/", + &self.config, + ) + .await?; + match workload_identity_token { + Some(token) => Ok(Some(Credential::BearerToken(token.access_token))), + None => Ok(None), + } + } } diff --git a/src/azure/storage/mod.rs b/src/azure/storage/mod.rs index a2549958..3fb2afa2 100644 --- a/src/azure/storage/mod.rs +++ b/src/azure/storage/mod.rs @@ -14,6 +14,8 @@ pub use credential::Credential as AzureStorageCredential; mod imds_credential; +mod workload_identity_credential; + mod loader; pub use loader::Loader as AzureStorageLoader; diff --git a/src/azure/storage/workload_identity_credential.rs b/src/azure/storage/workload_identity_credential.rs new file mode 100644 index 00000000..afa3467d --- /dev/null +++ b/src/azure/storage/workload_identity_credential.rs @@ -0,0 +1,85 @@ +use std::str; + +use http::HeaderValue; +use http::Method; +use http::Request; +use reqwest::Client; +use reqwest::Url; +use serde::Deserialize; +use std::fs; + +use super::config::Config; + +const MSI_API_VERSION: &str = "2019-08-01"; +const MSI_ENDPOINT: &str = "http://169.254.169.254/metadata/identity/oauth2/token"; + +/// Gets an access token for the specified resource and configuration. +/// +/// See +pub async fn get_workload_identity_token( + resource: &str, + config: &Config, +) -> anyhow::Result> { + let token = match ( + &config.azure_federated_token, + &config.azure_federated_token_file, + ) { + (Some(token), Some(_)) | (Some(token), None) => token.clone(), + (None, Some(token_file)) => { + let token = fs::read_to_string(token_file)?; + token + } + _ => return Ok(None), + }; + let tenant_id = if let Some(tenant_id) = &config.azure_tenant_id_env_key { + tenant_id + } else { + return Ok(None); + }; + let client_id = if let Some(client_id) = &config.client_id { + client_id + } else { + return Ok(None); + }; + + let endpoint = config.endpoint.as_deref().unwrap_or(MSI_ENDPOINT); + + let mut query_items = vec![("api-version", MSI_API_VERSION), ("resource", resource)]; + query_items.push(("token", &token)); + query_items.push(("tenant_id", &tenant_id)); + query_items.push(("client_id", &client_id)); + + let url = Url::parse_with_params(endpoint, &query_items)?; + let mut req = Request::builder() + .method(Method::GET) + .uri(url.to_string()) + .body("")?; + + req.headers_mut() + .insert("metadata", HeaderValue::from_static("true")); + + if let Some(secret) = &config.msi_secret { + req.headers_mut() + .insert("x-identity-header", HeaderValue::from_str(secret)?); + }; + + let res = Client::new().execute(req.try_into()?).await?; + let rsp_status = res.status(); + let rsp_body = res.text().await?; +; + + if !rsp_status.is_success() { + return Err(anyhow::anyhow!("Failed to get token from working identity credential")); + } + + let token: AccessToken = serde_json::from_str(&rsp_body)?; + Ok(Some(token)) +} + +#[derive(Debug, Clone, Deserialize)] +#[allow(unused)] +pub struct AccessToken { + pub access_token: String, + pub expires_on: String, + +} From e71ba53668e487b624f0f6e9f3e4c8f60a5a9fc7 Mon Sep 17 00:00:00 2001 From: congyi <15605187270@163.com> Date: Wed, 10 Apr 2024 22:28:42 +0800 Subject: [PATCH 2/2] resolve conflict --- src/azure/storage/loader.rs | 8 -- .../storage/workload_identity_credential.rs | 78 ------------------- 2 files changed, 86 deletions(-) diff --git a/src/azure/storage/loader.rs b/src/azure/storage/loader.rs index 4f23536c..b659b4aa 100644 --- a/src/azure/storage/loader.rs +++ b/src/azure/storage/loader.rs @@ -77,16 +77,8 @@ impl Loader { } async fn load_via_workload_identity(&self) -> Result> { -<<<<<<< HEAD - let workload_identity_token = workload_identity_credential::get_workload_identity_token( - "https://storage.azure.com/", - &self.config, - ) - .await?; -======= let workload_identity_token = workload_identity_credential::get_workload_identity_token(&self.config).await?; ->>>>>>> 31697b5ea627ffbecf328fe8d786127996865884 match workload_identity_token { Some(token) => Ok(Some(Credential::BearerToken(token.access_token))), None => Ok(None), diff --git a/src/azure/storage/workload_identity_credential.rs b/src/azure/storage/workload_identity_credential.rs index eaae3664..844690a6 100644 --- a/src/azure/storage/workload_identity_credential.rs +++ b/src/azure/storage/workload_identity_credential.rs @@ -6,64 +6,6 @@ use http::Request; use reqwest::Client; use reqwest::Url; use serde::Deserialize; -<<<<<<< HEAD -use std::fs; - -use super::config::Config; - -const MSI_API_VERSION: &str = "2019-08-01"; -const MSI_ENDPOINT: &str = "http://169.254.169.254/metadata/identity/oauth2/token"; - -/// Gets an access token for the specified resource and configuration. -/// -/// See -pub async fn get_workload_identity_token( - resource: &str, - config: &Config, -) -> anyhow::Result> { - let token = match ( - &config.azure_federated_token, - &config.azure_federated_token_file, - ) { - (Some(token), Some(_)) | (Some(token), None) => token.clone(), - (None, Some(token_file)) => { - let token = fs::read_to_string(token_file)?; - token - } - _ => return Ok(None), - }; - let tenant_id = if let Some(tenant_id) = &config.azure_tenant_id_env_key { - tenant_id - } else { - return Ok(None); - }; - let client_id = if let Some(client_id) = &config.client_id { - client_id - } else { - return Ok(None); - }; - - let endpoint = config.endpoint.as_deref().unwrap_or(MSI_ENDPOINT); - - let mut query_items = vec![("api-version", MSI_API_VERSION), ("resource", resource)]; - query_items.push(("token", &token)); - query_items.push(("tenant_id", &tenant_id)); - query_items.push(("client_id", &client_id)); - - let url = Url::parse_with_params(endpoint, &query_items)?; - let mut req = Request::builder() - .method(Method::GET) - .uri(url.to_string()) - .body("")?; - - req.headers_mut() - .insert("metadata", HeaderValue::from_static("true")); - - if let Some(secret) = &config.msi_secret { - req.headers_mut() - .insert("x-identity-header", HeaderValue::from_str(secret)?); - }; -======= use super::config::Config; @@ -108,29 +50,10 @@ pub async fn get_workload_identity_token(config: &Config) -> anyhow::Result>>>>>> 31697b5ea627ffbecf328fe8d786127996865884 let res = Client::new().execute(req.try_into()?).await?; let rsp_status = res.status(); let rsp_body = res.text().await?; -<<<<<<< HEAD -; - - if !rsp_status.is_success() { - return Err(anyhow::anyhow!("Failed to get token from working identity credential")); - } - - let token: AccessToken = serde_json::from_str(&rsp_body)?; - Ok(Some(token)) -} - -#[derive(Debug, Clone, Deserialize)] -#[allow(unused)] -pub struct AccessToken { - pub access_token: String, - pub expires_on: String, - -======= if !rsp_status.is_success() { return Err(anyhow::anyhow!( @@ -153,5 +76,4 @@ pub struct LoginResponse { pub not_before: Option, pub resource: Option, pub access_token: String, ->>>>>>> 31697b5ea627ffbecf328fe8d786127996865884 }