feat: support workload identity for azure storage

src/azure/storage/config.rs

@@ -1,3 +1,6 @@
+use std::env;
+use std::{collections::HashMap, fs};
+
 /// Config carries all the configuration for Azure Storage services.
 #[derive(Clone, Default)]
 #[cfg_attr(test, derive(Debug))]
@@ -23,6 +26,7 @@ pub struct Config {
  /// Specifies the application id (client id) associated with a user assigned managed service identity resource
  ///
  /// The values of object_id and msi_res_id are discarded
+ /// - cnv value: [`AZURE_CLIENT_ID`]
  ///
  /// This is part of use AAD(Azure Active Directory) authenticate on Azure VM
  pub client_id: Option<String>,
@@ -44,4 +48,61 @@ pub struct Config {
  ///
  /// This is part of use AAD(Azure Active Directory) authenticate on Azure VM
  pub endpoint: Option<String>,
+ /// `federated_token` value will be loaded from:
+ ///
+ /// - this field if it's `is_some`
+ /// - env value: [`AZURE_FEDERATED_TOKEN`]
+ /// - profile config: `federated_toen_file`
+ pub federated_token: Option<String>,
+ /// `tenant_id` value will be loaded from:
+ ///
+ /// - this field if it's `is_some`
+ /// - env value: [`AZURE_TENANT_ID`]
+ /// - profile config: `tenant_id`
+ pub tenant_id: Option<String>,
+ /// `authority_host` value will be loaded from:
+ ///
+ /// - this field if it's `is_some`
+ /// - env value: [`AZURE_AUTHORITY_HOST`]
+ /// - profile config: `authority_host`
+ pub authority_host: Option<String>,
+}
+
+pub const AZURE_FEDERATED_TOKEN: &str = "AZURE_FEDERATED_TOKEN";
+pub const AZURE_FEDERATED_TOKEN_FILE: &str = "AZURE_FEDERATED_TOKEN_FILE";
+pub const AZURE_TENANT_ID: &str = "AZURE_TENANT_ID";
+pub const AZURE_CLIENT_ID: &str = "AZURE_CLIENT_ID";
+pub const AZURE_AUTHORITY_HOST: &str = "AZURE_AUTHORITY_HOST";
+const AZURE_PUBLIC_CLOUD: &str = "https://login.microsoftonline.com";
+
+impl Config {
+ /// Load config from env.
+ pub fn from_env(mut self) -> Self {
+ let envs = env::vars().collect::<HashMap<_, _>>();
+
+ // federated_token can be loaded from both `AZURE_FEDERATED_TOKEN` and `AZURE_FEDERATED_TOKEN_FILE`.
+ if let Some(v) = envs.get(AZURE_FEDERATED_TOKEN_FILE) {
+ self.federated_token = Some(fs::read_to_string(v).unwrap_or_default());
+ }
+
+ if let Some(v) = envs.get(AZURE_FEDERATED_TOKEN) {
+ self.federated_token = Some(v.to_string());
+ }
+
+ if let Some(v) = envs.get(AZURE_TENANT_ID) {
+ self.tenant_id = Some(v.to_string());
+ }
+
+ if let Some(v) = envs.get(AZURE_CLIENT_ID) {
+ self.client_id = Some(v.to_string());
+ }
+
+ if let Some(v) = envs.get(AZURE_AUTHORITY_HOST) {
+ self.authority_host = Some(v.to_string());
+ } else {
+ self.authority_host = Some(AZURE_PUBLIC_CLOUD.to_string());
+ }
+
+ self
+ }
 }

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))]
@@ -31,7 +31,6 @@ impl Loader {
  if let Some(cred) = self.credential.lock().expect("lock poisoned").clone() {
  return Ok(Some(cred));
  }
-
  let cred = self.load_inner().await?;
 
  let mut lock = self.credential.lock().expect("lock poisoned");
@@ -45,6 +44,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 +75,13 @@ impl Loader {
 
  Ok(cred)
  }
+
+ async fn load_via_workload_identity(&self) -> Result<Option<Credential>> {
+ let workload_identity_token =
+ workload_identity_credential::get_workload_identity_token(&self.config).await?;
+ match workload_identity_token {
+ Some(token) => Ok(Some(Credential::BearerToken(token.access_token))),
+ None => Ok(None),
+ }
+ }
 }

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;

src/azure/storage/workload_identity_credential.rs

@@ -0,0 +1,79 @@
+use std::str;
+
+use http::HeaderValue;
+use http::Method;
+use http::Request;
+use reqwest::Client;
+use reqwest::Url;
+use serde::Deserialize;
+
+use super::config::Config;
+
+pub const API_VERSION: &str = "api-version";
+const STORAGE_TOKEN_SCOPE: &str = "https://storage.azure.com/.default";
+/// Gets an access token for the specified resource and configuration.
+///
+/// See <https://learn.microsoft.com/en-us/azure/app-service/overview-managed-identity?tabs=portal,http#using-the-rest-protocol>
+pub async fn get_workload_identity_token(config: &Config) -> anyhow::Result<Option<LoginResponse>> {
+ let (token, tenant_id, client_id, authority_host) = match (
+ &config.federated_token,
+ &config.tenant_id,
+ &config.client_id,
+ &config.authority_host,
+ ) {
+ (Some(token), Some(tenant_id), Some(client_id), Some(authority_host)) => {
+ (token, tenant_id, client_id, authority_host)
+ }
+ _ => return Ok(None),
+ };
+ let url = Url::parse(authority_host)?.join(&format!("/{tenant_id}/oauth2/v2.0/token"))?;
+ let scopes: &[&str] = &[STORAGE_TOKEN_SCOPE];
+ let encoded_body: String = form_urlencoded::Serializer::new(String::new())
+ .append_pair("client_id", client_id)
+ .append_pair("scope", &scopes.join(" "))
+ .append_pair(
+ "client_assertion_type",
+ "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+ )
+ .append_pair("client_assertion", token)
+ .append_pair("grant_type", "client_credentials")
+ .finish();
+
+ let mut req = Request::builder()
+ .method(Method::POST)
+ .uri(url.to_string())
+ .body(encoded_body)?;
+ req.headers_mut().insert(
+ http::header::CONTENT_TYPE.as_str(),
+ HeaderValue::from_static("application/x-www-form-urlencoded"),
+ );
+
+ req.headers_mut()
+ .insert(API_VERSION, HeaderValue::from_static("2019-06-01"));
+
+ 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, rsp_status = {}, rsp_body = {}",
+ rsp_status,
+ rsp_body
+ ));
+ }
+
+ let resp: LoginResponse = serde_json::from_str(&rsp_body)?;
+ Ok(Some(resp))
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct LoginResponse {
+ pub token_type: String,
+ pub expires_in: u64,
+ pub ext_expires_in: u64,
+ pub expires_on: Option<String>,
+ pub not_before: Option<String>,
+ pub resource: Option<String>,
+ pub access_token: String,
+}