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}
423
424impl WorkloadIdentityOAuthProvider {
425    pub(crate) fn new(
426        client_id: impl Into<String>,
427        federated_token_file: impl Into<String>,
428        tenant_id: impl AsRef<str>,
429        authority_host: Option<String>,
430    ) -> Self {
431        let authority_host =
432            authority_host.unwrap_or_else(|| authority_hosts::AZURE_PUBLIC_CLOUD.to_owned());
433
434        Self {
435            token_url: format!(
436                "{}/{}/oauth2/v2.0/token",
437                authority_host,
438                tenant_id.as_ref()
439            ),
440            client_id: client_id.into(),
441            federated_token_file: federated_token_file.into(),
442        }
443    }
444}
445
446#[async_trait::async_trait]
447impl TokenProvider for WorkloadIdentityOAuthProvider {
448    type Credential = AzureCredential;
449
450    /// Exchange the federated OIDC token for a bearer token via the
451    /// `client_credentials` grant with a JWT client assertion.
452    ///
453    /// # Errors
454    ///
455    /// Returns [`Error::FederatedTokenFile`] if the federated token file cannot
456    /// be read, [`Error::TokenRequest`] if the token request fails after the
457    /// configured retries are exhausted, or [`Error::TokenResponseBody`] if the
458    /// response body cannot be deserialized. The request is retried internally
459    /// per `retry`, so a returned error is terminal.
460    async fn fetch_token(
461        &self,
462        client: &Client,
463        service: &Arc<dyn HttpService>,
464        retry: &RetryConfig,
465    ) -> crate::Result<TemporaryToken<Arc<AzureCredential>>> {
466        let token_str = std::fs::read_to_string(&self.federated_token_file)
467            .map_err(|_| Error::FederatedTokenFile)?;
468
469        let response: OAuthTokenResponse = client
470            .request(Method::POST, &self.token_url)
471            .header(ACCEPT, HeaderValue::from_static(CONTENT_TYPE_JSON))
472            .form(&[
473                ("client_id", self.client_id.as_str()),
474                (
475                    "client_assertion_type",
476                    "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
477                ),
478                ("client_assertion", token_str.as_str()),
479                ("scope", AZURE_STORAGE_SCOPE),
480                ("grant_type", "client_credentials"),
481            ])
482            .retryable(retry, service.clone())
483            .idempotent(true)
484            .send()
485            .await
486            .map_err(|source| Error::TokenRequest { source })?
487            .json()
488            .await
489            .map_err(|source| Error::TokenResponseBody { source })?;
490
491        Ok(TemporaryToken {
492            token: Arc::new(AzureCredential::BearerToken(response.access_token)),
493            expiry: Some(Instant::now() + Duration::from_secs(response.expires_in)),
494        })
495    }
496}
497
498mod az_cli_date_format {
499    use chrono::{DateTime, TimeZone};
500    use serde::{self, Deserialize, Deserializer};
501
502    pub(crate) fn deserialize<'de, D>(deserializer: D) -> Result<DateTime<chrono::Local>, D::Error>
503    where
504        D: Deserializer<'de>,
505    {
506        let s = String::deserialize(deserializer)?;
507        let date = chrono::NaiveDateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S.%6f")
508            .map_err(serde::de::Error::custom)?;
509        chrono::Local
510            .from_local_datetime(&date)
511            .single()
512            .ok_or(serde::de::Error::custom(
513                "azure cli returned ambiguous expiry date",
514            ))
515    }
516}
517
518#[derive(Debug, Clone, Deserialize)]
519#[serde(rename_all = "camelCase")]
520struct AzureCliTokenResponse {
521    pub access_token: String,
522    #[serde(with = "az_cli_date_format")]
523    pub expires_on: DateTime<chrono::Local>,
524    pub token_type: String,
525}
526
527#[derive(Debug)]
528pub(crate) struct AzureCliCredential {
529    cache: TokenCache<Arc<AzureCredential>>,
530    scope: &'static str,
531}
532
533impl AzureCliCredential {
534    /// A credential for Azure Storage.
535    pub(crate) fn new() -> Self {
536        Self::new_with_scope(AZURE_STORAGE_SCOPE)
537    }
538
539    /// A credential for a custom OAuth scope (e.g. Azure Databricks).
540    pub(crate) fn new_with_scope(scope: &'static str) -> Self {
541        Self {
542            cache: TokenCache::default(),
543            scope,
544        }
545    }
546
547    async fn fetch_token(&self) -> Result<TemporaryToken<Arc<AzureCredential>>> {
548        let program = if cfg!(target_os = "windows") {
549            "cmd"
550        } else {
551            "az"
552        };
553        let mut args = Vec::new();
554        if cfg!(target_os = "windows") {
555            args.push("/C");
556            args.push("az");
557        }
558        args.push("account");
559        args.push("get-access-token");
560        args.push("--output");
561        args.push("json");
562        args.push("--scope");
563        args.push(self.scope);
564
565        match Command::new(program).args(args).output() {
566            Ok(az_output) if az_output.status.success() => {
567                let output = str::from_utf8(&az_output.stdout).map_err(|_| Error::AzureCli {
568                    message: "az response is not a valid utf-8 string".to_string(),
569                })?;
570
571                let token_response = serde_json::from_str::<AzureCliTokenResponse>(output)
572                    .map_err(|source| Error::AzureCliResponse { source })?;
573
574                if !token_response.token_type.eq_ignore_ascii_case("bearer") {
575                    return Err(Error::AzureCli {
576                        message: format!(
577                            "got unexpected token type from azure cli: {0}",
578                            token_response.token_type
579                        ),
580                    });
581                }
582                let duration =
583                    token_response.expires_on.naive_local() - chrono::Local::now().naive_local();
584                Ok(TemporaryToken {
585                    token: Arc::new(AzureCredential::BearerToken(token_response.access_token)),
586                    expiry: Some(
587                        Instant::now()
588                            + duration.to_std().map_err(|_| Error::AzureCli {
589                                message: "az returned invalid lifetime".to_string(),
590                            })?,
591                    ),
592                })
593            }
594            Ok(az_output) => {
595                let message = String::from_utf8_lossy(&az_output.stderr);
596                Err(Error::AzureCli {
597                    message: message.into(),
598                })
599            }
600            Err(e) => match e.kind() {
601                std::io::ErrorKind::NotFound => Err(Error::AzureCli {
602                    message: "Azure Cli not installed".into(),
603                }),
604                error_kind => Err(Error::AzureCli {
605                    message: format!("io error: {error_kind:?}"),
606                }),
607            },
608        }
609    }
610}
611
612#[async_trait]
613impl CredentialProvider for AzureCliCredential {
614    type Credential = AzureCredential;
615
616    async fn get_credential(&self) -> crate::Result<Arc<Self::Credential>> {
617        Ok(self.cache.get_or_insert_with(|| self.fetch_token()).await?)
618    }
619}
620
621#[cfg(test)]
622mod tests {
623    use reqwest::Client;
624    use tempfile::NamedTempFile;
625
626    use super::*;
627    use crate::service::ReqwestService;
628    use mockito;
629
630    #[test]
631    fn test_debug_does_not_leak_secrets() {
632        // Deliberately non-secret-looking sentinel value so the secret scanner
633        // doesn't flag this test; we only care that the Debug output does not
634        // echo it back.
635        let credential = AzureCredential::BearerToken("fake-token-do-not-print".to_string());
636        let rendered = format!("{credential:?}");
637        assert!(
638            !rendered.contains("fake-token-do-not-print"),
639            "bearer token leaked: {rendered}"
640        );
641        assert!(rendered.contains("<redacted>"));
642    }
643
644    #[tokio::test]
645    async fn test_managed_identity() {
646        let mut server = mockito::Server::new_async().await;
647
648        unsafe { std::env::set_var(MSI_SECRET_ENV_KEY, "env-secret") };
649
650        let endpoint = server.url();
651        let client = Client::new();
652        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
653        let retry_config = RetryConfig::default();
654
655        let _mock = server
656            .mock("GET", "/metadata/identity/oauth2/token")
657            .match_query(mockito::Matcher::AllOf(vec![mockito::Matcher::UrlEncoded(
658                "client_id".into(),
659                "client_id".into(),
660            )]))
661            .match_header("x-identity-header", "env-secret")
662            .match_header("metadata", "true")
663            .with_status(200)
664            .with_body(
665                r#"
666            {
667                "access_token": "TOKEN",
668                "refresh_token": "",
669                "expires_in": "3599",
670                "expires_on": "1506484173",
671                "not_before": "1506480273",
672                "resource": "https://management.azure.com/",
673                "token_type": "Bearer"
674              }
675            "#,
676            )
677            .create_async()
678            .await;
679
680        let credential = ImdsManagedIdentityProvider::new(
681            Some("client_id".into()),
682            None,
683            None,
684            Some(format!("{endpoint}/metadata/identity/oauth2/token")),
685        );
686
687        let token = credential
688            .fetch_token(&client, &service, &retry_config)
689            .await
690            .unwrap();
691
692        assert_eq!(
693            token.token.as_ref(),
694            &AzureCredential::BearerToken("TOKEN".into())
695        );
696    }
697
698    #[tokio::test]
699    async fn test_workload_identity() {
700        let mut server = mockito::Server::new_async().await;
701        let tokenfile = NamedTempFile::new().unwrap();
702        let tenant = "tenant";
703        std::fs::write(tokenfile.path(), "federated-token").unwrap();
704
705        let endpoint = server.url();
706        let client = Client::new();
707        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
708        let retry_config = RetryConfig::default();
709
710        let _mock = server
711            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
712            .match_body(mockito::Matcher::Regex("federated-token".into()))
713            .with_status(200)
714            .with_body(
715                r#"
716            {
717                "access_token": "TOKEN",
718                "refresh_token": "",
719                "expires_in": 3599,
720                "expires_on": "1506484173",
721                "not_before": "1506480273",
722                "resource": "https://management.azure.com/",
723                "token_type": "Bearer"
724              }
725            "#,
726            )
727            .create_async()
728            .await;
729
730        let credential = WorkloadIdentityOAuthProvider::new(
731            "client_id",
732            tokenfile.path().to_str().unwrap(),
733            tenant,
734            Some(endpoint.to_string()),
735        );
736
737        let token = credential
738            .fetch_token(&client, &service, &retry_config)
739            .await
740            .unwrap();
741
742        assert_eq!(
743            token.token.as_ref(),
744            &AzureCredential::BearerToken("TOKEN".into())
745        );
746    }
747
748    #[tokio::test]
749    async fn test_client_secret_happy_path() {
750        let mut server = mockito::Server::new_async().await;
751        let tenant = "my-tenant";
752
753        let _mock = server
754            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
755            .match_body(mockito::Matcher::AllOf(vec![
756                mockito::Matcher::Regex("client_id=myclientid".into()),
757                mockito::Matcher::Regex("grant_type=client_credentials".into()),
758            ]))
759            .with_status(200)
760            .with_body(r#"{"access_token":"AZURE_TOKEN","expires_in":3599,"token_type":"Bearer"}"#)
761            .create_async()
762            .await;
763
764        let client = Client::new();
765        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
766        let retry_config = RetryConfig::default();
767
768        let provider = ClientSecretOAuthProvider::new_with_scope(
769            "myclientid".into(),
770            "myclientsecret".into(),
771            tenant,
772            Some(server.url()),
773            AZURE_STORAGE_SCOPE,
774        );
775
776        let token = provider
777            .fetch_token(&client, &service, &retry_config)
778            .await
779            .unwrap();
780
781        assert_eq!(
782            token.token.as_ref(),
783            &AzureCredential::BearerToken("AZURE_TOKEN".into())
784        );
785        assert!(token.expiry.is_some());
786
787        _mock.assert_async().await;
788    }
789
790    #[tokio::test]
791    async fn test_client_secret_invalid_client() {
792        let mut server = mockito::Server::new_async().await;
793        let tenant = "bad-tenant";
794
795        let _mock = server
796            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
797            .with_status(400)
798            .with_body(
799                r#"{"error":"invalid_client","error_description":"Client authentication failed"}"#,
800            )
801            .create_async()
802            .await;
803
804        let client = Client::new();
805        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
806        let retry_config = RetryConfig::default();
807
808        let provider = ClientSecretOAuthProvider::new_with_scope(
809            "badclientid".into(),
810            "badsecret".into(),
811            tenant,
812            Some(server.url()),
813            AZURE_STORAGE_SCOPE,
814        );
815
816        let result = provider.fetch_token(&client, &service, &retry_config).await;
817
818        assert!(result.is_err(), "Expected error for 400 response");
819
820        _mock.assert_async().await;
821    }
822
823    #[tokio::test]
824    async fn test_client_secret_token_refresh() {
825        let mut server = mockito::Server::new_async().await;
826        let tenant = "refresh-tenant";
827
828        // First call returns a token with very short expiry (1 second)
829        let _mock1 = server
830            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
831            .with_status(200)
832            .with_body(r#"{"access_token":"FIRST_TOKEN","expires_in":1,"token_type":"Bearer"}"#)
833            .create_async()
834            .await;
835
836        let client = Client::new();
837        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(client.clone()));
838        let retry_config = RetryConfig::default();
839
840        let provider = ClientSecretOAuthProvider::new_with_scope(
841            "myclientid".into(),
842            "myclientsecret".into(),
843            tenant,
844            Some(server.url()),
845            AZURE_STORAGE_SCOPE,
846        );
847
848        let token1 = provider
849            .fetch_token(&client, &service, &retry_config)
850            .await
851            .unwrap();
852        assert_eq!(
853            token1.token.as_ref(),
854            &AzureCredential::BearerToken("FIRST_TOKEN".into())
855        );
856        // Token expiry should be very soon (1 second)
857        assert!(token1.expiry.is_some());
858
859        // Second call returns a new token
860        let _mock2 = server
861            .mock("POST", format!("/{tenant}/oauth2/v2.0/token").as_str())
862            .with_status(200)
863            .with_body(r#"{"access_token":"SECOND_TOKEN","expires_in":3600,"token_type":"Bearer"}"#)
864            .create_async()
865            .await;
866
867        let token2 = provider
868            .fetch_token(&client, &service, &retry_config)
869            .await
870            .unwrap();
871        assert_eq!(
872            token2.token.as_ref(),
873            &AzureCredential::BearerToken("SECOND_TOKEN".into())
874        );
875
876        _mock1.assert_async().await;
877        _mock2.assert_async().await;
878    }
879}
olai_http/azure/credential.rs

olai_http/azure/
credential.rs
