olai_http/azure/
credential.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::fmt::Debug;
19use std::ops::Deref;
20use std::process::Command;
21use std::str;
22use std::sync::Arc;
23use std::time::{Duration, Instant, SystemTime};
24
25use async_trait::async_trait;
26use chrono::DateTime;
27use reqwest::header::{ACCEPT, AUTHORIZATION, HeaderValue};
28use reqwest::{Client, Method, Request, RequestBuilder};
29use serde::Deserialize;
30
31use crate::retry::RetryExt;
32use crate::service::HttpService;
33use crate::token::{TemporaryToken, TokenCache};
34use crate::{CredentialProvider, RetryConfig, TokenProvider};
35
36const CONTENT_TYPE_JSON: &str = "application/json";
37const MSI_SECRET_ENV_KEY: &str = "IDENTITY_HEADER";
38const MSI_API_VERSION: &str = "2019-08-01";
39
40/// OIDC scope used when interacting with Azure Storage OAuth2 APIs
41pub(crate) const AZURE_STORAGE_SCOPE: &str = "https://storage.azure.com/.default";
42
43/// Resource ID used when obtaining an access token for Azure Storage from the metadata endpoint
44pub(crate) const AZURE_STORAGE_RESOURCE: &str = "https://storage.azure.com";
45
46/// Azure AD Application ID for Databricks; used as resource ID for Azure Databricks API tokens
47pub(crate) const AZURE_DATABRICKS_RESOURCE: &str = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d";
48
49/// OIDC scope for Azure Databricks OAuth2 APIs
50pub(crate) const AZURE_DATABRICKS_SCOPE: &str = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default";
51
52#[derive(Debug, thiserror::Error)]
53pub enum Error {
54    #[error("Error performing token request: {}", source)]
55    TokenRequest { source: crate::retry::Error },
56
57    #[error("Error getting token response body: {}", source)]
58    TokenResponseBody { source: reqwest::Error },
59
60    #[error("Error reading federated token file ")]
61    FederatedTokenFile,
62
63    #[error("Invalid Access Key: {}", source)]
64    InvalidAccessKey { source: base64::DecodeError },
65
66    #[error("'az account get-access-token' command failed: {message}")]
67    AzureCli { message: String },
68
69    #[error("Failed to parse azure cli response: {source}")]
70    AzureCliResponse { source: serde_json::Error },
71}
72
73pub(crate) type Result<T, E = Error> = std::result::Result<T, E>;
74
75impl From<Error> for crate::Error {
76    fn from(value: Error) -> Self {
77        Self::Generic {
78            source: Box::new(value),
79        }
80    }
81}
82
83/// An Azure storage credential.
84///
85/// Holds a long-lived secret used to authorize requests against Azure APIs, and
86/// must be handled as sensitive material — avoid logging it or otherwise
87/// exposing the inner value.
88#[derive(Eq, PartialEq)]
89pub enum AzureCredential {
90    /// An Azure AD OAuth 2.0 bearer token.
91    ///
92    /// The wrapped [`String`] is a secret access token; treat it as
93    /// confidential and do not log it.
94    BearerToken(String),
95}
96
97// Manual `Debug` impl: an `AzureCredential` holds a long-lived bearer token, so
98// we never expose its value. See the credential-redaction convention in the
99// README.
100impl std::fmt::Debug for AzureCredential {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        match self {
103            Self::BearerToken(_) => f.debug_tuple("BearerToken").field(&"<redacted>").finish(),
104        }
105    }
106}
107
108/// A list of known Azure authority hosts
109#[allow(dead_code)]
110pub mod authority_hosts {
111    /// China-based Azure Authority Host
112    pub const AZURE_CHINA: &str = "https://login.chinacloudapi.cn";
113    /// Germany-based Azure Authority Host
114    pub const AZURE_GERMANY: &str = "https://login.microsoftonline.de";
115    /// US Government Azure Authority Host
116    pub const AZURE_GOVERNMENT: &str = "https://login.microsoftonline.us";
117    /// Public Cloud Azure Authority Host
118    pub const AZURE_PUBLIC_CLOUD: &str = "https://login.microsoftonline.com";
119}
120
121/// Authorize a [`Request`] with an [`AzureAuthorizer`]
122#[derive(Debug)]
123pub struct AzureAuthorizer<'a> {
124    credential: &'a AzureCredential,
125}
126
127impl<'a> AzureAuthorizer<'a> {
128    /// Create a new [`AzureAuthorizer`] that authorizes requests with the given
129    /// [`AzureCredential`].
130    pub fn new(credential: &'a AzureCredential) -> Self {
131        AzureAuthorizer { credential }
132    }
133
134    /// Authorize `request` in place by appending the credential's
135    /// `Authorization` header.
136    ///
137    /// For an [`AzureCredential::BearerToken`] this appends an
138    /// `Authorization: Bearer <token>` header.
139    pub fn authorize(&self, request: &mut Request) {
140        match self.credential {
141            AzureCredential::BearerToken(token) => {
142                request.headers_mut().append(
143                    AUTHORIZATION,
144                    HeaderValue::from_str(format!("Bearer {token}").as_str()).unwrap(),
145                );
146            }
147        }
148    }
149}
150
151pub trait AzureCredentialExt {
152    /// Apply authorization to requests against azure storage accounts
153    fn with_azure_authorization(
154        self,
155        credential: &Option<impl Deref<Target = AzureCredential>>,
156    ) -> Self;
157}
158
159impl AzureCredentialExt for RequestBuilder {
160    fn with_azure_authorization(
161        self,
162        credential: &Option<impl Deref<Target = AzureCredential>>,
163    ) -> Self {
164        match credential.as_deref() {
165            Some(credential) => {
166                let (client, request) = self.build_split();
167                let mut request = request.expect("request valid");
168                AzureAuthorizer::new(credential).authorize(&mut request);
169                Self::from_parts(client, request)
170            }
171            None => self,
172        }
173    }
174}
175
176/// <https://learn.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-client-creds-grant-flow#successful-response-1>
177#[derive(Deserialize, Debug)]
178struct OAuthTokenResponse {
179    access_token: String,
180    expires_in: u64,
181}
182
183/// Encapsulates the logic to perform an OAuth token challenge using a client
184/// secret (shared-secret variant of the OAuth 2.0 client-credentials flow).
185///
186/// POSTs a `client_credentials` grant to the Azure AD token endpoint and
187/// returns a bearer token for the configured scope (defaults to
188/// `https://storage.azure.com/.default`).
189///
190/// # References
191/// - <https://learn.microsoft.com/en-us/entra/identity-platform/v2-oauth2-client-creds-grant-flow>
192/// - <https://learn.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-client-creds-grant-flow#first-case-access-token-request-with-a-shared-secret>
193#[derive(Debug)]
194pub(crate) struct ClientSecretOAuthProvider {
195    token_url: String,
196    client_id: String,
197    client_secret: String,
198    scope: String,
199}
200
201impl ClientSecretOAuthProvider {
202    /// Create a new provider with a custom OAuth scope (e.g. for Databricks).
203    pub(crate) fn new_with_scope(
204        client_id: String,
205        client_secret: String,
206        tenant_id: impl AsRef<str>,
207        authority_host: Option<String>,
208        scope: &str,
209    ) -> Self {
210        let authority_host =
211            authority_host.unwrap_or_else(|| authority_hosts::AZURE_PUBLIC_CLOUD.to_owned());
212
213        Self {
214            token_url: format!(
215                "{}/{}/oauth2/v2.0/token",
216                authority_host,
217                tenant_id.as_ref()
218            ),
219            client_id,
220            client_secret,
221            scope: scope.to_owned(),
222        }
223    }
224}
225
226#[async_trait::async_trait]
227impl TokenProvider for ClientSecretOAuthProvider {
228    type Credential = AzureCredential;
229
230    /// Fetch a bearer token via the OAuth 2.0 `client_credentials` grant.
231    ///
232    /// # Errors
233    ///
234    /// Returns [`Error::TokenRequest`] if the token request fails after the
235    /// configured retries are exhausted, or [`Error::TokenResponseBody`] if the
236    /// response body cannot be deserialized. The request is retried internally
237    /// per `retry`, so a returned error is terminal.
238    async fn fetch_token(
239        &self,
240        client: &Client,
241        service: &Arc<dyn HttpService>,
242        retry: &RetryConfig,
243    ) -> crate::Result<TemporaryToken<Arc<AzureCredential>>> {
244        let response: OAuthTokenResponse = client
245            .request(Method::POST, &self.token_url)
246            .header(ACCEPT, HeaderValue::from_static(CONTENT_TYPE_JSON))
247            .form(&[
248                ("client_id", self.client_id.as_str()),
249                ("client_secret", self.client_secret.as_str()),
250                ("scope", self.scope.as_str()),
251                ("grant_type", "client_credentials"),
252            ])
253            .retryable(retry, service.clone())
254            .idempotent(true)
255            .send()
256            .await
257            .map_err(|source| Error::TokenRequest { source })?
258            .json()
259            .await
260            .map_err(|source| Error::TokenResponseBody { source })?;
261
262        Ok(TemporaryToken {
263            token: Arc::new(AzureCredential::BearerToken(response.access_token)),
264            expiry: Some(Instant::now() + Duration::from_secs(response.expires_in)),
265        })
266    }
267}
268
269fn expires_on_string<'de, D>(deserializer: D) -> std::result::Result<Instant, D::Error>
270where
271    D: serde::de::Deserializer<'de>,
272{
273    let v = String::deserialize(deserializer)?;
274    let v = v.parse::<u64>().map_err(serde::de::Error::custom)?;
275    let now = SystemTime::now()
276        .duration_since(SystemTime::UNIX_EPOCH)
277        .map_err(serde::de::Error::custom)?;
278
279    Ok(Instant::now() + Duration::from_secs(v.saturating_sub(now.as_secs())))
280}
281
282#[derive(Debug, Clone, Deserialize)]
283struct ImdsTokenResponse {
284    pub access_token: String,
285    #[serde(deserialize_with = "expires_on_string")]
286    pub expires_on: Instant,
287}
288
289/// Attempts authentication using a managed identity that has been assigned to the deployment environment.
290///
291/// Calls the Azure Instance Metadata Service (IMDS) endpoint at
292/// `http://169.254.169.254/metadata/identity/oauth2/token` (or a custom
293/// endpoint when one is supplied) to obtain a bearer token for the configured
294/// Azure AD resource.  Supports user-assigned identities via `client_id`,
295/// `object_id`, or `msi_res_id` selectors.
296///
297/// # References
298/// - <https://learn.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token>
299/// - <https://learn.microsoft.com/en-gb/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http>
300#[derive(Debug)]
301pub(crate) struct ImdsManagedIdentityProvider {
302    msi_endpoint: String,
303    client_id: Option<String>,
304    object_id: Option<String>,
305    msi_res_id: Option<String>,
306    resource: String,
307}
308
309impl ImdsManagedIdentityProvider {
310    pub(crate) fn new(
311        client_id: Option<String>,
312        object_id: Option<String>,
313        msi_res_id: Option<String>,
314        msi_endpoint: Option<String>,
315    ) -> Self {
316        Self::new_with_resource(
317            client_id,
318            object_id,
319            msi_res_id,
320            msi_endpoint,
321            AZURE_STORAGE_RESOURCE,
322        )
323    }
324
325    /// Create a new provider with a custom Azure AD resource (e.g. for Databricks).
326    pub(crate) fn new_with_resource(
327        client_id: Option<String>,
328        object_id: Option<String>,
329        msi_res_id: Option<String>,
330        msi_endpoint: Option<String>,
331        resource: &str,
332    ) -> Self {
333        let msi_endpoint = msi_endpoint
334            .unwrap_or_else(|| "http://169.254.169.254/metadata/identity/oauth2/token".to_owned());
335
336        Self {
337            msi_endpoint,
338            client_id,
339            object_id,
340            msi_res_id,
341            resource: resource.to_owned(),
342        }
343    }
344}
345
346#[async_trait::async_trait]
347impl TokenProvider for ImdsManagedIdentityProvider {
348    type Credential = AzureCredential;
349
350    /// Fetch a bearer token from the Azure Instance Metadata Service.
351    ///
352    /// # Errors
353    ///
354    /// Returns [`Error::TokenRequest`] if the IMDS request fails after the
355    /// configured retries are exhausted, or [`Error::TokenResponseBody`] if the
356    /// response body cannot be deserialized. The request is retried internally
357    /// per `retry`, so a returned error is terminal.
358    async fn fetch_token(
359        &self,
360        client: &Client,
361        service: &Arc<dyn HttpService>,
362        retry: &RetryConfig,
363    ) -> crate::Result<TemporaryToken<Arc<AzureCredential>>> {
364        let mut query_items = vec![
365            ("api-version", MSI_API_VERSION),
366            ("resource", self.resource.as_str()),
367        ];
368
369        let mut identity = None;
370        if let Some(client_id) = &self.client_id {
371            identity = Some(("client_id", client_id));
372        }
373        if let Some(object_id) = &self.object_id {
374            identity = Some(("object_id", object_id));
375        }
376        if let Some(msi_res_id) = &self.msi_res_id {
377            identity = Some(("msi_res_id", msi_res_id));
378        }
379        if let Some((key, value)) = identity {
380            query_items.push((key, value));
381        }
382
383        let mut builder = client
384            .request(Method::GET, &self.msi_endpoint)
385            .header("metadata", "true")
386            .query(&query_items);
387
388        if let Ok(val) = std::env::var(MSI_SECRET_ENV_KEY) {
389            builder = builder.header("x-identity-header", val);
390        };
391
392        let response: ImdsTokenResponse = builder
393            .send_retry(retry, service.clone())
394            .await
395            .map_err(|source| Error::TokenRequest { source })?
396            .json()
397            .await
398            .map_err(|source| Error::TokenResponseBody { source })?;
399
400        Ok(TemporaryToken {
401            token: Arc::new(AzureCredential::BearerToken(response.access_token)),
402            expiry: Some(response.expires_on),
403        })
404    }
405}
406
407/// Credential for using workload identity federation.
408///
409/// Exchanges a federated OIDC token (read from a file, e.g. a Kubernetes
410/// projected service-account token) for an Azure AD bearer token using the
411/// `client_credentials` grant with a JWT client assertion.  Typically used in
412/// AKS workloads where the pod is annotated with a managed identity.
413///
414/// # References
415/// - <https://learn.microsoft.com/en-us/entra/workload-id/workload-identity-federation>
416/// - <https://learn.microsoft.com/en-us/azure/active-directory/develop/workload-identity-federation>
417#[derive(Debug)]
418pub(crate) struct WorkloadIdentityOAuthProvider {
419    token_url: String,
420    client_id: String,
421    federated_token_file: String,
422    scope: String,
423}
424
425impl WorkloadIdentityOAuthProvider {
426    pub(crate) fn new(
427        client_id: impl Into<String>,
428        federated_token_file: impl Into<String>,
429        tenant_id: impl AsRef<str>,
430        authority_host: Option<String>,
431    ) -> Self {
432        let authority_host =
433            authority_host.unwrap_or_else(|| authority_hosts::AZURE_PUBLIC_CLOUD.to_owned());
434
435        Self {
436            token_url: format!(
437                "{}/{}/oauth2/v2.0/token",
438                authority_host,
439                tenant_id.as_ref()
440            ),
441            client_id: client_id.into(),
442            federated_token_file: federated_token_file.into(),
443            scope: AZURE_STORAGE_SCOPE.to_owned(),
444        }
445    }
446
447    /// Request tokens for `scope` instead of Azure Storage.
448    pub(crate) fn with_scope(mut self, scope: impl Into<String>) -> Self {
449        self.scope = scope.into();
450        self
451    }
452}
453
454#[async_trait::async_trait]
455impl TokenProvider for WorkloadIdentityOAuthProvider {
456    type Credential = AzureCredential;
457
458    /// Exchange the federated OIDC token for a bearer token via the
459    /// `client_credentials` grant with a JWT client assertion.
460    ///
461    /// # Errors
462    ///
463    /// Returns [`Error::FederatedTokenFile`] if the federated token file cannot
464    /// be read, [`Error::TokenRequest`] if the token request fails after the
465    /// configured retries are exhausted, or [`Error::TokenResponseBody`] if the
466    /// response body cannot be deserialized. The request is retried internally
467    /// per `retry`, so a returned error is terminal.
468    async fn fetch_token(
469        &self,
470        client: &Client,
471        service: &Arc<dyn HttpService>,
472        retry: &RetryConfig,
473    ) -> crate::Result<TemporaryToken<Arc<AzureCredential>>> {
474        let token_str = std::fs::read_to_string(&self.federated_token_file)
475            .map_err(|_| Error::FederatedTokenFile)?;
476
477        let response: OAuthTokenResponse = client
478            .request(Method::POST, &self.token_url)
479            .header(ACCEPT, HeaderValue::from_static(CONTENT_TYPE_JSON))
480            .form(&[
481                ("client_id", self.client_id.as_str()),
482                (
483                    "client_assertion_type",
484                    "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
485                ),
486                ("client_assertion", token_str.as_str()),
487                ("scope", self.scope.as_str()),
488                ("grant_type", "client_credentials"),
489            ])
490            .retryable(retry, service.clone())
491            .idempotent(true)
492            .send()
493            .await
494            .map_err(|source| Error::TokenRequest { source })?
495            .json()
496            .await
497            .map_err(|source| Error::TokenResponseBody { source })?;
498
499        Ok(TemporaryToken {
500            token: Arc::new(AzureCredential::BearerToken(response.access_token)),
501            expiry: Some(Instant::now() + Duration::from_secs(response.expires_in)),
502        })
503    }
504}
505
506mod az_cli_date_format {
507    use chrono::{DateTime, TimeZone};
508    use serde::{self, Deserialize, Deserializer};
509
510    pub(crate) fn deserialize<'de, D>(deserializer: D) -> Result<DateTime<chrono::Local>, D::Error>
511    where
512        D: Deserializer<'de>,
513    {
514        let s = String::deserialize(deserializer)?;
515        let date = chrono::NaiveDateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S.%6f")
516            .map_err(serde::de::Error::custom)?;
517        chrono::Local
518            .from_local_datetime(&date)
519            .single()
520            .ok_or(serde::de::Error::custom(
521                "azure cli returned ambiguous expiry date",
522            ))
523    }
524}
525
526#[derive(Debug, Clone, Deserialize)]
527#[serde(rename_all = "camelCase")]
528struct AzureCliTokenResponse {
529    pub access_token: String,
530    #[serde(with = "az_cli_date_format")]
531    pub expires_on: DateTime<chrono::Local>,
532    pub token_type: String,
533}
534
535#[derive(Debug)]
536pub(crate) struct AzureCliCredential {
537    cache: TokenCache<Arc<AzureCredential>>,
538    scope: String,
539}
540
541impl AzureCliCredential {
542    /// A credential for Azure Storage.
543    pub(crate) fn new() -> Self {
544        Self::new_with_scope(AZURE_STORAGE_SCOPE)
545    }
546
547    /// A credential for a custom OAuth scope (e.g. Azure Databricks).
548    pub(crate) fn new_with_scope(scope: impl Into<String>) -> Self {
549        Self {
550            cache: TokenCache::default(),
551            scope: scope.into(),
552        }
553    }
554
555    async fn fetch_token(&self) -> Result<TemporaryToken<Arc<AzureCredential>>> {
556        let program = if cfg!(target_os = "windows") {
557            "cmd"
558        } else {
559            "az"
560        };
561        let mut args = Vec::new();
562        if cfg!(target_os = "windows") {
563            args.push("/C");
564            args.push("az");
565        }
566        args.push("account");
567        args.push("get-access-token");
568        args.push("--output");
569        args.push("json");
570        args.push("--scope");
571        args.push(&self.scope);
572
573        match Command::new(program).args(args).output() {
574            Ok(az_output) if az_output.status.success() => {
575                let output = str::from_utf8(&az_output.stdout).map_err(|_| Error::AzureCli {
576                    message: "az response is not a valid utf-8 string".to_string(),
577                })?;
578
579                let token_response = serde_json::from_str::<AzureCliTokenResponse>(output)
580                    .map_err(|source| Error::AzureCliResponse { source })?;
581
582                if !token_response.token_type.eq_ignore_ascii_case("bearer") {
583                    return Err(Error::AzureCli {
584                        message: format!(
585                            "got unexpected token type from azure cli: {0}",
586                            token_response.token_type
587                        ),
588                    });
589                }
590                let duration =
591                    token_response.expires_on.naive_local() - chrono::Local::now().naive_local();
592                Ok(TemporaryToken {
593                    token: Arc::new(AzureCredential::BearerToken(token_response.access_token)),
594                    expiry: Some(
595                        Instant::now()
596                            + duration.to_std().map_err(|_| Error::AzureCli {
597                                message: "az returned invalid lifetime".to_string(),
598                            })?,
599                    ),
600                })
601            }
602            Ok(az_output) => {
603                let message = String::from_utf8_lossy(&az_output.stderr);
604                Err(Error::AzureCli {
605                    message: message.into(),
606                })
607            }
608            Err(e) => match e.kind() {
609                std::io::ErrorKind::NotFound => Err(Error::AzureCli {
610                    message: "Azure Cli not installed".into(),
611                }),
612                error_kind => Err(Error::AzureCli {
613                    message: format!("io error: {error_kind:?}"),
614                }),
615            },
616        }
617    }
618}
619
620#[async_trait]
621impl CredentialProvider for AzureCliCredential {
622    type Credential = AzureCredential;
623
624    async fn get_credential(&self) -> crate::Result<Arc<Self::Credential>> {
625        Ok(self.cache.get_or_insert_with(|| self.fetch_token()).await?)
626    }
627}
628
629#[cfg(test)]
630mod tests {
631    use reqwest::Client;
632    use tempfile::NamedTempFile;
633
634    use super::*;
635    use crate::service::ReqwestService;
636    use mockito;
637
638    #[test]
639    fn test_debug_does_not_leak_secrets() {
640        // Deliberately non-secret-looking sentinel value so the secret scanner
641        // doesn't flag this test; we only care that the Debug output does not
642        // echo it back.
643        let credential = AzureCredential::BearerToken("fake-token-do-not-print".to_string());
644        let rendered = format!("{credential:?}");
645        assert!(
646            !rendered.contains("fake-token-do-not-print"),
647            "bearer token leaked: {rendered}"
648        );
649        assert!(rendered.contains("<redacted>"));
650    }
651
652    #[tokio::test]
653    async fn test_managed_identity() {
654        let mut server = mockito::Server::new_async().await;
655
656        unsafe { std::env::set_var(MSI_SECRET_ENV_KEY, "env-secret") };
657
658        let endpoint = server.url();
659        let client = Client::new();
660        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
661        let retry_config = RetryConfig::default();
662
663        let _mock = server
664            .mock("GET", "/metadata/identity/oauth2/token")
665            .match_query(mockito::Matcher::AllOf(vec![mockito::Matcher::UrlEncoded(
666                "client_id".into(),
667                "client_id".into(),
668            )]))
669            .match_header("x-identity-header", "env-secret")
670            .match_header("metadata", "true")
671            .with_status(200)
672            .with_body(
673                r#"
674            {
675                "access_token": "TOKEN",
676                "refresh_token": "",
677                "expires_in": "3599",
678                "expires_on": "1506484173",
679                "not_before": "1506480273",
680                "resource": "https://management.azure.com/",
681                "token_type": "Bearer"
682              }
683            "#,
684            )
685            .create_async()
686            .await;
687
688        let credential = ImdsManagedIdentityProvider::new(
689            Some("client_id".into()),
690            None,
691            None,
692            Some(format!("{endpoint}/metadata/identity/oauth2/token")),
693        );
694
695        let token = credential
696            .fetch_token(&client, &service, &retry_config)
697            .await
698            .unwrap();
699
700        assert_eq!(
701            token.token.as_ref(),
702            &AzureCredential::BearerToken("TOKEN".into())
703        );
704    }
705
706    #[tokio::test]
707    async fn test_workload_identity() {
708        let mut server = mockito::Server::new_async().await;
709        let tokenfile = NamedTempFile::new().unwrap();
710        let tenant = "tenant";
711        std::fs::write(tokenfile.path(), "federated-token").unwrap();
712
713        let endpoint = server.url();
714        let client = Client::new();
715        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
716        let retry_config = RetryConfig::default();
717
718        let _mock = server
719            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
720            .match_body(mockito::Matcher::Regex("federated-token".into()))
721            .with_status(200)
722            .with_body(
723                r#"
724            {
725                "access_token": "TOKEN",
726                "refresh_token": "",
727                "expires_in": 3599,
728                "expires_on": "1506484173",
729                "not_before": "1506480273",
730                "resource": "https://management.azure.com/",
731                "token_type": "Bearer"
732              }
733            "#,
734            )
735            .create_async()
736            .await;
737
738        let credential = WorkloadIdentityOAuthProvider::new(
739            "client_id",
740            tokenfile.path().to_str().unwrap(),
741            tenant,
742            Some(endpoint.to_string()),
743        );
744
745        let token = credential
746            .fetch_token(&client, &service, &retry_config)
747            .await
748            .unwrap();
749
750        assert_eq!(
751            token.token.as_ref(),
752            &AzureCredential::BearerToken("TOKEN".into())
753        );
754    }
755
756    #[tokio::test]
757    async fn test_client_secret_happy_path() {
758        let mut server = mockito::Server::new_async().await;
759        let tenant = "my-tenant";
760
761        let _mock = server
762            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
763            .match_body(mockito::Matcher::AllOf(vec![
764                mockito::Matcher::Regex("client_id=myclientid".into()),
765                mockito::Matcher::Regex("grant_type=client_credentials".into()),
766            ]))
767            .with_status(200)
768            .with_body(r#"{"access_token":"AZURE_TOKEN","expires_in":3599,"token_type":"Bearer"}"#)
769            .create_async()
770            .await;
771
772        let client = Client::new();
773        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
774        let retry_config = RetryConfig::default();
775
776        let provider = ClientSecretOAuthProvider::new_with_scope(
777            "myclientid".into(),
778            "myclientsecret".into(),
779            tenant,
780            Some(server.url()),
781            AZURE_STORAGE_SCOPE,
782        );
783
784        let token = provider
785            .fetch_token(&client, &service, &retry_config)
786            .await
787            .unwrap();
788
789        assert_eq!(
790            token.token.as_ref(),
791            &AzureCredential::BearerToken("AZURE_TOKEN".into())
792        );
793        assert!(token.expiry.is_some());
794
795        _mock.assert_async().await;
796    }
797
798    #[tokio::test]
799    async fn test_client_secret_invalid_client() {
800        let mut server = mockito::Server::new_async().await;
801        let tenant = "bad-tenant";
802
803        let _mock = server
804            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
805            .with_status(400)
806            .with_body(
807                r#"{"error":"invalid_client","error_description":"Client authentication failed"}"#,
808            )
809            .create_async()
810            .await;
811
812        let client = Client::new();
813        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
814        let retry_config = RetryConfig::default();
815
816        let provider = ClientSecretOAuthProvider::new_with_scope(
817            "badclientid".into(),
818            "badsecret".into(),
819            tenant,
820            Some(server.url()),
821            AZURE_STORAGE_SCOPE,
822        );
823
824        let result = provider.fetch_token(&client, &service, &retry_config).await;
825
826        assert!(result.is_err(), "Expected error for 400 response");
827
828        _mock.assert_async().await;
829    }
830
831    #[tokio::test]
832    async fn test_client_secret_token_refresh() {
833        let mut server = mockito::Server::new_async().await;
834        let tenant = "refresh-tenant";
835
836        // First call returns a token with very short expiry (1 second)
837        let _mock1 = server
838            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
839            .with_status(200)
840            .with_body(r#"{"access_token":"FIRST_TOKEN","expires_in":1,"token_type":"Bearer"}"#)
841            .create_async()
842            .await;
843
844        let client = Client::new();
845        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
846        let retry_config = RetryConfig::default();
847
848        let provider = ClientSecretOAuthProvider::new_with_scope(
849            "myclientid".into(),
850            "myclientsecret".into(),
851            tenant,
852            Some(server.url()),
853            AZURE_STORAGE_SCOPE,
854        );
855
856        let token1 = provider
857            .fetch_token(&client, &service, &retry_config)
858            .await
859            .unwrap();
860        assert_eq!(
861            token1.token.as_ref(),
862            &AzureCredential::BearerToken("FIRST_TOKEN".into())
863        );
864        // Token expiry should be very soon (1 second)
865        assert!(token1.expiry.is_some());
866
867        // Second call returns a new token
868        let _mock2 = server
869            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
870            .with_status(200)
871            .with_body(r#"{"access_token":"SECOND_TOKEN","expires_in":3600,"token_type":"Bearer"}"#)
872            .create_async()
873            .await;
874
875        let token2 = provider
876            .fetch_token(&client, &service, &retry_config)
877            .await
878            .unwrap();
879        assert_eq!(
880            token2.token.as_ref(),
881            &AzureCredential::BearerToken("SECOND_TOKEN".into())
882        );
883
884        _mock1.assert_async().await;
885        _mock2.assert_async().await;
886    }
887}
olai_http/azure/credential.rs

olai_http/azure/
credential.rs
