Skip to main content

olai_http/azure/
mod.rs

1use std::sync::Arc;
2
3use futures::future::BoxFuture;
4
5use crate::service::make_service;
6use crate::token::TemporaryToken;
7use crate::{ClientOptions, CredentialProvider, RequestSigner, Result, RetryConfig, TokenProvider};
8
9mod builder;
10pub(crate) mod credential;
11mod sas;
12
13pub(crate) use self::credential::*;
14pub use builder::*;
15pub use credential::AzureCredential;
16pub use sas::{
17    DEFAULT_TTL_SECS, SAS_READ, SAS_READ_WRITE, UserDelegationKey, UserDelegationKeyOptions,
18    build_storage_key_sas, build_user_delegation_sas, fetch_user_delegation_key,
19    fetch_user_delegation_key_with_retry, generate_storage_key_sas, generate_user_delegation_sas,
20    generate_user_delegation_sas_with_options,
21};
22
23pub type AzureCredentialProvider = Arc<dyn CredentialProvider<Credential = AzureCredential>>;
24
25/// Configuration for Azure authentication
26#[derive(Debug, Clone)]
27pub struct AzureConfig {
28    pub credentials: AzureCredentialProvider,
29    pub retry_config: RetryConfig,
30    pub skip_signature: bool,
31    pub client_options: ClientOptions,
32}
33
34impl AzureConfig {
35    pub(crate) async fn get_credential(&self) -> Result<Option<Arc<AzureCredential>>> {
36        if self.skip_signature {
37            Ok(None)
38        } else {
39            Some(self.credentials.get_credential().await).transpose()
40        }
41    }
42}
43
44impl RequestSigner for AzureConfig {
45    fn sign<'a>(
46        &'a self,
47        req: reqwest::RequestBuilder,
48    ) -> BoxFuture<'a, Result<reqwest::RequestBuilder>> {
49        Box::pin(async move {
50            let credential = self.get_credential().await?;
51            Ok(req.with_azure_authorization(&credential))
52        })
53    }
54}
55
56/// Fetch a short-lived Azure storage bearer token using an OAuth2 client-secret flow.
57///
58/// Calls the Azure AD token endpoint for the given tenant and returns a
59/// `TemporaryToken<Arc<AzureCredential>>` scoped to Azure Storage
60/// (`https://storage.azure.com/.default`).
61///
62/// This is commonly used by data-platform servers to vend temporary Azure credentials
63/// for external locations backed by an `AzureServicePrincipal` credential.
64pub async fn fetch_client_secret_token(
65    tenant_id: &str,
66    client_id: String,
67    client_secret: String,
68    authority_host: Option<String>,
69) -> Result<TemporaryToken<Arc<AzureCredential>>> {
70    let provider = ClientSecretOAuthProvider::new_with_scope(
71        client_id,
72        client_secret,
73        tenant_id,
74        authority_host,
75        credential::AZURE_STORAGE_SCOPE,
76    );
77    let client = ClientOptions::default().client()?;
78    let service = make_service(client.clone(), None);
79    provider
80        .fetch_token(&client, &service, &RetryConfig::default())
81        .await
82}
83
84/// Fetch a short-lived Azure storage bearer token using workload identity (federated token file).
85///
86/// Reads the OIDC token from the given file path and exchanges it for an Azure AD bearer
87/// token via the client-credentials grant with a JWT client assertion.
88pub async fn fetch_workload_identity_token(
89    tenant_id: &str,
90    client_id: String,
91    federated_token_file: String,
92    authority_host: Option<String>,
93) -> Result<TemporaryToken<Arc<AzureCredential>>> {
94    let provider = WorkloadIdentityOAuthProvider::new(
95        client_id,
96        federated_token_file,
97        tenant_id,
98        authority_host,
99    );
100    let client = ClientOptions::default().client()?;
101    let service = make_service(client.clone(), None);
102    provider
103        .fetch_token(&client, &service, &RetryConfig::default())
104        .await
105}
106
107/// Fetch a short-lived Azure storage bearer token using IMDS managed identity.
108///
109/// Calls the Azure Instance Metadata Service (IMDS) endpoint and returns a
110/// bearer token scoped to Azure Storage.  Supports user-assigned identities via
111/// `client_id`, `object_id`, or `msi_res_id`.
112pub async fn fetch_managed_identity_token(
113    client_id: Option<String>,
114    object_id: Option<String>,
115    msi_res_id: Option<String>,
116) -> Result<TemporaryToken<Arc<AzureCredential>>> {
117    let provider = ImdsManagedIdentityProvider::new(client_id, object_id, msi_res_id, None);
118    let client = ClientOptions::default().with_allow_http(true).client()?;
119    let service = make_service(client.clone(), None);
120    provider
121        .fetch_token(&client, &service, &RetryConfig::default())
122        .await
123}