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

olai_http/azure/
credential.rs
