1use std::collections::HashMap;
25use std::sync::{Arc, Mutex};
26use std::time::{Duration, SystemTime, UNIX_EPOCH};
27
28use agent_framework_core::error::{Error, Result};
29use async_trait::async_trait;
30use serde_json::Value;
31
32use crate::credential::TokenCredential;
33
34pub const REFRESH_SKEW: Duration = Duration::from_secs(120);
38
39pub const DEFAULT_IMDS_ENDPOINT: &str = "http://169.254.169.254/metadata/identity/oauth2/token";
42
43const IMDS_API_VERSION: &str = "2018-02-01";
45
46pub const DEFAULT_AUTHORITY: &str = "https://login.microsoftonline.com";
49
50const DEFAULT_TOKEN_TTL_SECS: u64 = 3600;
53
54const CLI_FALLBACK_TTL_SECS: u64 = 300;
60
61#[derive(Clone)]
66struct CachedToken {
67 token: String,
68 expires_at: SystemTime,
69}
70
71#[derive(Default)]
73struct TokenCache {
74 entries: Mutex<HashMap<String, CachedToken>>,
75}
76
77impl TokenCache {
78 fn get(&self, scope: &str) -> Option<String> {
81 let entries = self.entries.lock().unwrap();
82 let entry = entries.get(scope)?;
83 if SystemTime::now() + REFRESH_SKEW < entry.expires_at {
84 Some(entry.token.clone())
85 } else {
86 None
87 }
88 }
89
90 fn put(&self, scope: &str, token: String, expires_at: SystemTime) {
92 self.entries
93 .lock()
94 .unwrap()
95 .insert(scope.to_string(), CachedToken { token, expires_at });
96 }
97}
98
99fn json_u64(v: Option<&Value>) -> Option<u64> {
108 match v {
109 Some(Value::Number(n)) => n.as_u64().or_else(|| n.as_f64().map(|f| f as u64)),
110 Some(Value::String(s)) => s.trim().parse().ok(),
111 _ => None,
112 }
113}
114
115fn parse_oauth_token(body: &str) -> Result<(String, SystemTime)> {
119 let v: Value = serde_json::from_str(body)
120 .map_err(|e| Error::other(format!("invalid token response json: {e}")))?;
121 let token = v
122 .get("access_token")
123 .and_then(Value::as_str)
124 .ok_or_else(|| Error::other("token response missing 'access_token'"))?
125 .to_string();
126 let ttl = json_u64(v.get("expires_in")).unwrap_or(DEFAULT_TOKEN_TTL_SECS);
127 Ok((token, SystemTime::now() + Duration::from_secs(ttl)))
128}
129
130fn parse_cli_output(stdout: &[u8]) -> Result<(String, SystemTime)> {
136 let v: Value = serde_json::from_slice(stdout)
137 .map_err(|e| Error::other(format!("invalid Azure CLI token json: {e}")))?;
138 let token = v
139 .get("accessToken")
140 .and_then(Value::as_str)
141 .ok_or_else(|| Error::other("Azure CLI token json missing 'accessToken'"))?
142 .to_string();
143 let expires_at = match v.get("expires_on").and_then(Value::as_i64) {
144 Some(epoch) if epoch > 0 => UNIX_EPOCH + Duration::from_secs(epoch as u64),
145 _ => SystemTime::now() + Duration::from_secs(CLI_FALLBACK_TTL_SECS),
146 };
147 Ok((token, expires_at))
148}
149
150fn resource_from_scope(scope: &str) -> &str {
153 scope.strip_suffix("/.default").unwrap_or(scope)
154}
155
156fn missing_env_vars<'a>(vars: &[(&'a str, &Option<String>)]) -> Vec<&'a str> {
161 vars.iter()
162 .filter(|(_, v)| v.is_none())
163 .map(|(name, _)| *name)
164 .collect()
165}
166
167pub struct AzureCliCredential {
178 program: String,
179 default_scope: String,
180 cache: TokenCache,
181}
182
183impl AzureCliCredential {
184 pub fn new(scope: impl Into<String>) -> Self {
187 Self {
188 program: "az".to_string(),
189 default_scope: scope.into(),
190 cache: TokenCache::default(),
191 }
192 }
193
194 pub fn with_command(mut self, program: impl Into<String>) -> Self {
197 self.program = program.into();
198 self
199 }
200
201 async fn token(&self, scope: &str) -> Result<String> {
202 if let Some(t) = self.cache.get(scope) {
203 return Ok(t);
204 }
205 let (token, expires_at) = self.fetch(scope).await?;
206 self.cache.put(scope, token.clone(), expires_at);
207 Ok(token)
208 }
209
210 async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
211 let output = tokio::process::Command::new(&self.program)
212 .arg("account")
213 .arg("get-access-token")
214 .arg("--scope")
215 .arg(scope)
216 .arg("--output")
217 .arg("json")
218 .output()
219 .await;
220 match output {
221 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
222 Err(Error::Configuration(format!(
223 "Azure CLI ('{}') was not found on PATH; run `az login` after installing it, \
224 or use a different credential: {e}",
225 self.program
226 )))
227 }
228 Err(e) => Err(Error::service(format!("failed to run Azure CLI: {e}"))),
229 Ok(out) if !out.status.success() => {
230 let stderr = String::from_utf8_lossy(&out.stderr);
231 Err(Error::other(format!(
232 "Azure CLI token request failed ({}): {}",
233 out.status,
234 stderr.trim()
235 )))
236 }
237 Ok(out) => parse_cli_output(&out.stdout),
238 }
239 }
240}
241
242#[async_trait]
243impl TokenCredential for AzureCliCredential {
244 async fn get_token(&self) -> Result<String> {
245 self.token(&self.default_scope).await
246 }
247 async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
248 self.token(scope).await
249 }
250}
251
252pub struct ClientSecretCredential {
259 http: reqwest::Client,
260 authority: String,
261 tenant_id: String,
262 client_id: String,
263 client_secret: String,
264 default_scope: String,
265 cache: TokenCache,
266}
267
268impl ClientSecretCredential {
269 pub fn new(
272 tenant_id: impl Into<String>,
273 client_id: impl Into<String>,
274 client_secret: impl Into<String>,
275 scope: impl Into<String>,
276 ) -> Self {
277 Self {
278 http: reqwest::Client::new(),
279 authority: DEFAULT_AUTHORITY.to_string(),
280 tenant_id: tenant_id.into(),
281 client_id: client_id.into(),
282 client_secret: client_secret.into(),
283 default_scope: scope.into(),
284 cache: TokenCache::default(),
285 }
286 }
287
288 pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
291 self.authority = authority.into();
292 self
293 }
294
295 fn token_url(&self) -> String {
296 format!(
297 "{}/{}/oauth2/v2.0/token",
298 self.authority.trim_end_matches('/'),
299 self.tenant_id
300 )
301 }
302
303 async fn token(&self, scope: &str) -> Result<String> {
304 if let Some(t) = self.cache.get(scope) {
305 return Ok(t);
306 }
307 let (token, expires_at) = self.fetch(scope).await?;
308 self.cache.put(scope, token.clone(), expires_at);
309 Ok(token)
310 }
311
312 async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
313 let params = [
314 ("grant_type", "client_credentials"),
315 ("client_id", self.client_id.as_str()),
316 ("client_secret", self.client_secret.as_str()),
317 ("scope", scope),
318 ];
319 let resp = self
320 .http
321 .post(self.token_url())
322 .form(¶ms)
323 .send()
324 .await
325 .map_err(|e| Error::service(format!("client-secret token request failed: {e}")))?;
326 let status = resp.status();
327 let body = resp.text().await.unwrap_or_default();
328 if !status.is_success() {
329 return Err(Error::other(format!(
330 "client-secret token request rejected ({}): {}",
331 status,
332 body.trim()
333 )));
334 }
335 parse_oauth_token(&body)
336 }
337}
338
339#[async_trait]
340impl TokenCredential for ClientSecretCredential {
341 async fn get_token(&self) -> Result<String> {
342 self.token(&self.default_scope).await
343 }
344 async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
345 self.token(scope).await
346 }
347}
348
349pub struct EnvironmentCredential {
366 inner: ClientSecretCredential,
367}
368
369impl EnvironmentCredential {
370 pub fn new(scope: impl Into<String>) -> Result<Self> {
378 let tenant_id = std::env::var("AZURE_TENANT_ID").ok();
379 let client_id = std::env::var("AZURE_CLIENT_ID").ok();
380 let client_secret = std::env::var("AZURE_CLIENT_SECRET").ok();
381
382 let missing = missing_env_vars(&[
383 ("AZURE_TENANT_ID", &tenant_id),
384 ("AZURE_CLIENT_ID", &client_id),
385 ("AZURE_CLIENT_SECRET", &client_secret),
386 ]);
387 if !missing.is_empty() {
388 return Err(Error::Configuration(format!(
389 "EnvironmentCredential: missing required environment variable(s): {}",
390 missing.join(", ")
391 )));
392 }
393
394 Ok(Self {
395 inner: ClientSecretCredential::new(
396 tenant_id.unwrap(),
397 client_id.unwrap(),
398 client_secret.unwrap(),
399 scope,
400 ),
401 })
402 }
403
404 pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
407 self.inner = self.inner.with_authority(authority);
408 self
409 }
410}
411
412#[async_trait]
413impl TokenCredential for EnvironmentCredential {
414 async fn get_token(&self) -> Result<String> {
415 self.inner.get_token().await
416 }
417 async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
418 self.inner.get_token_for_scope(scope).await
419 }
420}
421
422pub struct ManagedIdentityCredential {
437 http: reqwest::Client,
438 endpoint: String,
439 client_id: Option<String>,
440 identity_header: Option<String>,
441 default_scope: String,
442 cache: TokenCache,
443}
444
445impl ManagedIdentityCredential {
446 pub fn new(scope: impl Into<String>) -> Self {
452 let endpoint = std::env::var("IDENTITY_ENDPOINT")
453 .or_else(|_| std::env::var("MSI_ENDPOINT"))
454 .unwrap_or_else(|_| DEFAULT_IMDS_ENDPOINT.to_string());
455 let identity_header = std::env::var("IDENTITY_HEADER").ok();
456 let http = reqwest::Client::builder()
462 .connect_timeout(std::time::Duration::from_secs(2))
463 .timeout(std::time::Duration::from_secs(10))
464 .build()
465 .unwrap_or_else(|_| reqwest::Client::new());
466 Self {
467 http,
468 endpoint,
469 client_id: None,
470 identity_header,
471 default_scope: scope.into(),
472 cache: TokenCache::default(),
473 }
474 }
475
476 pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
479 self.endpoint = endpoint.into();
480 self
481 }
482
483 pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
485 self.client_id = Some(client_id.into());
486 self
487 }
488
489 pub fn with_identity_header(mut self, header: impl Into<String>) -> Self {
491 self.identity_header = Some(header.into());
492 self
493 }
494
495 async fn token(&self, scope: &str) -> Result<String> {
496 if let Some(t) = self.cache.get(scope) {
497 return Ok(t);
498 }
499 let (token, expires_at) = self.fetch(scope).await?;
500 self.cache.put(scope, token.clone(), expires_at);
501 Ok(token)
502 }
503
504 async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
505 let resource = resource_from_scope(scope);
506 let mut params: Vec<(&str, &str)> =
507 vec![("api-version", IMDS_API_VERSION), ("resource", resource)];
508 if let Some(cid) = &self.client_id {
509 params.push(("client_id", cid.as_str()));
510 }
511 let mut req = self
512 .http
513 .get(&self.endpoint)
514 .query(¶ms)
515 .header("Metadata", "true");
516 if let Some(h) = &self.identity_header {
517 req = req.header("X-IDENTITY-HEADER", h);
518 }
519 let resp = req
520 .send()
521 .await
522 .map_err(|e| Error::service(format!("managed-identity token request failed: {e}")))?;
523 let status = resp.status();
524 let body = resp.text().await.unwrap_or_default();
525 if !status.is_success() {
526 return Err(Error::other(format!(
527 "managed-identity token request rejected ({}): {}",
528 status,
529 body.trim()
530 )));
531 }
532 parse_oauth_token(&body)
533 }
534}
535
536#[async_trait]
537impl TokenCredential for ManagedIdentityCredential {
538 async fn get_token(&self) -> Result<String> {
539 self.token(&self.default_scope).await
540 }
541 async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
542 self.token(scope).await
543 }
544}
545
546pub struct WorkloadIdentityCredential {
565 http: reqwest::Client,
566 authority: String,
567 tenant_id: String,
568 client_id: String,
569 token_file_path: String,
570 default_scope: String,
571 cache: TokenCache,
572}
573
574impl WorkloadIdentityCredential {
575 pub fn new(scope: impl Into<String>) -> Result<Self> {
583 Self::with_overrides(scope, None, None, None)
584 }
585
586 pub fn with_overrides(
595 scope: impl Into<String>,
596 tenant_id: Option<String>,
597 client_id: Option<String>,
598 token_file_path: Option<String>,
599 ) -> Result<Self> {
600 let tenant_id = tenant_id.or_else(|| std::env::var("AZURE_TENANT_ID").ok());
601 let client_id = client_id.or_else(|| std::env::var("AZURE_CLIENT_ID").ok());
602 let token_file_path =
603 token_file_path.or_else(|| std::env::var("AZURE_FEDERATED_TOKEN_FILE").ok());
604
605 let missing = missing_env_vars(&[
606 ("AZURE_TENANT_ID", &tenant_id),
607 ("AZURE_CLIENT_ID", &client_id),
608 ("AZURE_FEDERATED_TOKEN_FILE", &token_file_path),
609 ]);
610 if !missing.is_empty() {
611 return Err(Error::Configuration(format!(
612 "WorkloadIdentityCredential: missing required configuration: {}",
613 missing.join(", ")
614 )));
615 }
616
617 Ok(Self {
618 http: reqwest::Client::new(),
619 authority: DEFAULT_AUTHORITY.to_string(),
620 tenant_id: tenant_id.unwrap(),
621 client_id: client_id.unwrap(),
622 token_file_path: token_file_path.unwrap(),
623 default_scope: scope.into(),
624 cache: TokenCache::default(),
625 })
626 }
627
628 pub fn with_client_id(scope: impl Into<String>, client_id: Option<String>) -> Result<Self> {
633 Self::with_overrides(scope, None, client_id, None)
634 }
635
636 pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
639 self.authority = authority.into();
640 self
641 }
642
643 fn token_url(&self) -> String {
644 format!(
645 "{}/{}/oauth2/v2.0/token",
646 self.authority.trim_end_matches('/'),
647 self.tenant_id
648 )
649 }
650
651 async fn token(&self, scope: &str) -> Result<String> {
652 if let Some(t) = self.cache.get(scope) {
653 return Ok(t);
654 }
655 let (token, expires_at) = self.fetch(scope).await?;
656 self.cache.put(scope, token.clone(), expires_at);
657 Ok(token)
658 }
659
660 fn request_params(&self, scope: &str) -> Result<Vec<(String, String)>> {
665 let assertion = std::fs::read_to_string(&self.token_file_path).map_err(|e| {
666 Error::Configuration(format!(
667 "WorkloadIdentityCredential: failed to read federated token file '{}': {e}",
668 self.token_file_path
669 ))
670 })?;
671 Ok(vec![
672 ("grant_type".to_string(), "client_credentials".to_string()),
673 ("client_id".to_string(), self.client_id.clone()),
674 ("client_assertion".to_string(), assertion.trim().to_string()),
675 (
676 "client_assertion_type".to_string(),
677 "urn:ietf:params:oauth:client-assertion-type:jwt-bearer".to_string(),
678 ),
679 ("scope".to_string(), scope.to_string()),
680 ])
681 }
682
683 async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
684 let params = self.request_params(scope)?;
685 let resp = self
686 .http
687 .post(self.token_url())
688 .form(¶ms)
689 .send()
690 .await
691 .map_err(|e| Error::service(format!("workload-identity token request failed: {e}")))?;
692 let status = resp.status();
693 let body = resp.text().await.unwrap_or_default();
694 if !status.is_success() {
695 return Err(Error::other(format!(
696 "workload-identity token request rejected ({}): {}",
697 status,
698 body.trim()
699 )));
700 }
701 parse_oauth_token(&body)
702 }
703}
704
705#[async_trait]
706impl TokenCredential for WorkloadIdentityCredential {
707 async fn get_token(&self) -> Result<String> {
708 self.token(&self.default_scope).await
709 }
710 async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
711 self.token(scope).await
712 }
713}
714
715pub struct ChainedTokenCredential {
724 credentials: Vec<Arc<dyn TokenCredential>>,
725 remembered: Mutex<Option<usize>>,
726}
727
728impl ChainedTokenCredential {
729 pub fn new(credentials: Vec<Arc<dyn TokenCredential>>) -> Self {
731 Self {
732 credentials,
733 remembered: Mutex::new(None),
734 }
735 }
736
737 async fn acquire(&self, scope: Option<&str>) -> Result<String> {
738 let remembered = *self.remembered.lock().unwrap();
741
742 if let Some(idx) = remembered {
743 if let Some(cred) = self.credentials.get(idx) {
744 if let Ok(token) = Self::call(cred.as_ref(), scope).await {
745 return Ok(token);
746 }
747 }
748 }
749
750 let mut errors = Vec::new();
751 for (i, cred) in self.credentials.iter().enumerate() {
752 if Some(i) == remembered {
753 continue;
755 }
756 match Self::call(cred.as_ref(), scope).await {
757 Ok(token) => {
758 *self.remembered.lock().unwrap() = Some(i);
759 return Ok(token);
760 }
761 Err(e) => errors.push(format!("credential[{i}]: {e}")),
762 }
763 }
764
765 Err(Error::Configuration(format!(
766 "ChainedTokenCredential: no credential in the chain returned a token ({})",
767 if errors.is_empty() {
768 "the chain is empty".to_string()
769 } else {
770 errors.join("; ")
771 }
772 )))
773 }
774
775 async fn call(cred: &dyn TokenCredential, scope: Option<&str>) -> Result<String> {
776 match scope {
777 Some(s) => cred.get_token_for_scope(s).await,
778 None => cred.get_token().await,
779 }
780 }
781}
782
783#[async_trait]
784impl TokenCredential for ChainedTokenCredential {
785 async fn get_token(&self) -> Result<String> {
786 self.acquire(None).await
787 }
788 async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
789 self.acquire(Some(scope)).await
790 }
791}
792
793pub struct DefaultAzureCredential {
828 chain: ChainedTokenCredential,
829}
830
831impl DefaultAzureCredential {
832 pub fn new(scope: impl Into<String>) -> Self {
834 let scope = scope.into();
835 Self {
836 chain: ChainedTokenCredential::new(default_chain(&scope)),
837 }
838 }
839}
840
841fn default_chain(scope: &str) -> Vec<Arc<dyn TokenCredential>> {
845 let mut chain: Vec<Arc<dyn TokenCredential>> = Vec::new();
846 if let Ok(cred) = EnvironmentCredential::new(scope) {
847 chain.push(Arc::new(cred));
848 }
849 if let Ok(cred) = WorkloadIdentityCredential::new(scope) {
850 chain.push(Arc::new(cred));
851 }
852 chain.push(Arc::new(ManagedIdentityCredential::new(scope)));
853 chain.push(Arc::new(AzureCliCredential::new(scope)));
854 chain
855}
856
857#[async_trait]
858impl TokenCredential for DefaultAzureCredential {
859 async fn get_token(&self) -> Result<String> {
860 self.chain.get_token().await
861 }
862 async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
863 self.chain.get_token_for_scope(scope).await
864 }
865}
866
867#[cfg(test)]
868mod tests {
869 use super::*;
870 use std::sync::atomic::{AtomicUsize, Ordering};
871
872 #[test]
875 fn cache_returns_token_before_skew_window() {
876 let cache = TokenCache::default();
877 cache.put(
878 "scope",
879 "tok".into(),
880 SystemTime::now() + Duration::from_secs(3600),
881 );
882 assert_eq!(cache.get("scope").as_deref(), Some("tok"));
883 }
884
885 #[test]
886 fn cache_misses_within_refresh_skew() {
887 let cache = TokenCache::default();
888 cache.put(
890 "scope",
891 "tok".into(),
892 SystemTime::now() + Duration::from_secs(60),
893 );
894 assert!(cache.get("scope").is_none());
895 }
896
897 #[test]
898 fn cache_is_per_scope() {
899 let cache = TokenCache::default();
900 cache.put(
901 "a",
902 "tok-a".into(),
903 SystemTime::now() + Duration::from_secs(3600),
904 );
905 assert_eq!(cache.get("a").as_deref(), Some("tok-a"));
906 assert!(cache.get("b").is_none());
907 }
908
909 #[test]
914 fn parse_cli_output_reads_token_and_epoch_expiry() {
915 let json = br#"{"accessToken":"cli-token","expiresOn":"2099-01-01 00:00:00.000000","expires_on":4070908800,"tokenType":"Bearer"}"#;
916 let (token, expires_at) = parse_cli_output(json).unwrap();
917 assert_eq!(token, "cli-token");
918 assert_eq!(expires_at, UNIX_EPOCH + Duration::from_secs(4070908800));
919 }
920
921 #[test]
922 fn parse_cli_output_falls_back_to_ttl_without_epoch() {
923 let before = SystemTime::now();
924 let (token, expires_at) = parse_cli_output(br#"{"accessToken":"t"}"#).unwrap();
925 assert_eq!(token, "t");
926 assert!(expires_at >= before + Duration::from_secs(CLI_FALLBACK_TTL_SECS));
927 }
928
929 #[test]
930 fn parse_cli_output_errors_without_token() {
931 assert!(parse_cli_output(br#"{"tokenType":"Bearer"}"#).is_err());
932 }
933
934 #[test]
935 fn parse_oauth_token_number_and_string_expiry() {
936 let (t, _) = parse_oauth_token(r#"{"access_token":"a","expires_in":3600}"#).unwrap();
937 assert_eq!(t, "a");
938 let (t2, _) = parse_oauth_token(r#"{"access_token":"b","expires_in":"3600"}"#).unwrap();
940 assert_eq!(t2, "b");
941 }
942
943 #[test]
944 fn resource_strips_default_suffix() {
945 assert_eq!(
946 resource_from_scope("https://ai.azure.com/.default"),
947 "https://ai.azure.com"
948 );
949 assert_eq!(
950 resource_from_scope("https://vault.azure.net"),
951 "https://vault.azure.net"
952 );
953 }
954
955 #[tokio::test]
960 async fn azure_cli_missing_binary_gives_clear_error() {
961 let cred = AzureCliCredential::new("https://ai.azure.com/.default")
962 .with_command("definitely-not-a-real-az-binary-xyz");
963 let err = cred.get_token().await.unwrap_err();
964 let msg = err.to_string();
965 assert!(
966 msg.contains("was not found on PATH"),
967 "expected a clear missing-CLI message, got: {msg}"
968 );
969 }
970
971 struct FakeCredential {
976 result: std::result::Result<String, ()>,
977 calls: AtomicUsize,
978 }
979
980 impl FakeCredential {
981 fn ok(token: &str) -> Self {
982 Self {
983 result: Ok(token.to_string()),
984 calls: AtomicUsize::new(0),
985 }
986 }
987 fn fail() -> Self {
988 Self {
989 result: Err(()),
990 calls: AtomicUsize::new(0),
991 }
992 }
993 }
994
995 #[async_trait]
996 impl TokenCredential for FakeCredential {
997 async fn get_token(&self) -> Result<String> {
998 self.calls.fetch_add(1, Ordering::SeqCst);
999 self.result
1000 .clone()
1001 .map_err(|_| Error::other("fake credential failure"))
1002 }
1003 }
1004
1005 #[tokio::test]
1006 async fn chain_returns_first_success_and_remembers_it() {
1007 let first = Arc::new(FakeCredential::fail());
1008 let second = Arc::new(FakeCredential::ok("winner"));
1009 let third = Arc::new(FakeCredential::ok("unused"));
1010 let chain = ChainedTokenCredential::new(vec![first.clone(), second.clone(), third.clone()]);
1011
1012 assert_eq!(chain.get_token().await.unwrap(), "winner");
1013 assert_eq!(first.calls.load(Ordering::SeqCst), 1);
1015 assert_eq!(second.calls.load(Ordering::SeqCst), 1);
1016 assert_eq!(third.calls.load(Ordering::SeqCst), 0);
1017
1018 assert_eq!(chain.get_token().await.unwrap(), "winner");
1021 assert_eq!(first.calls.load(Ordering::SeqCst), 1);
1022 assert_eq!(second.calls.load(Ordering::SeqCst), 2);
1023 }
1024
1025 #[tokio::test]
1026 async fn chain_reports_all_failures_when_none_succeed() {
1027 let chain = ChainedTokenCredential::new(vec![
1028 Arc::new(FakeCredential::fail()),
1029 Arc::new(FakeCredential::fail()),
1030 ]);
1031 let err = chain.get_token().await.unwrap_err();
1032 assert!(err.to_string().contains("no credential in the chain"));
1033 }
1034
1035 #[tokio::test]
1036 async fn empty_chain_errors() {
1037 let chain = ChainedTokenCredential::new(vec![]);
1038 assert!(chain.get_token().await.is_err());
1039 }
1040
1041 static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
1051
1052 const ENTRA_ENV_VARS: &[&str] = &[
1053 "AZURE_TENANT_ID",
1054 "AZURE_CLIENT_ID",
1055 "AZURE_CLIENT_SECRET",
1056 "AZURE_FEDERATED_TOKEN_FILE",
1057 ];
1058
1059 fn clear_entra_env() {
1060 for key in ENTRA_ENV_VARS {
1061 unsafe { std::env::remove_var(key) };
1064 }
1065 }
1066
1067 const TEST_SCOPE: &str = "https://ai.azure.com/.default";
1068
1069 #[test]
1070 fn environment_credential_errors_listing_all_missing_vars() {
1071 let _guard = ENV_MUTEX.lock().unwrap();
1072 clear_entra_env();
1073 let err = EnvironmentCredential::new(TEST_SCOPE).err().unwrap();
1077 clear_entra_env();
1078 let msg = err.to_string();
1079 assert!(msg.contains("AZURE_TENANT_ID"), "{msg}");
1080 assert!(msg.contains("AZURE_CLIENT_ID"), "{msg}");
1081 assert!(msg.contains("AZURE_CLIENT_SECRET"), "{msg}");
1082 }
1083
1084 #[test]
1085 fn environment_credential_lists_only_the_vars_still_missing() {
1086 let _guard = ENV_MUTEX.lock().unwrap();
1087 clear_entra_env();
1088 unsafe {
1089 std::env::set_var("AZURE_TENANT_ID", "tenant");
1090 std::env::set_var("AZURE_CLIENT_ID", "client");
1091 }
1092 let err = EnvironmentCredential::new(TEST_SCOPE).err().unwrap();
1093 clear_entra_env();
1094 let msg = err.to_string();
1095 assert!(msg.contains("AZURE_CLIENT_SECRET"), "{msg}");
1096 assert!(!msg.contains("AZURE_TENANT_ID"), "{msg}");
1097 assert!(!msg.contains("AZURE_CLIENT_ID"), "{msg}");
1098 }
1099
1100 #[test]
1101 fn environment_credential_succeeds_with_all_three_vars() {
1102 let _guard = ENV_MUTEX.lock().unwrap();
1103 clear_entra_env();
1104 unsafe {
1105 std::env::set_var("AZURE_TENANT_ID", "tenant");
1106 std::env::set_var("AZURE_CLIENT_ID", "client");
1107 std::env::set_var("AZURE_CLIENT_SECRET", "secret");
1108 }
1109 let cred = EnvironmentCredential::new(TEST_SCOPE);
1110 clear_entra_env();
1111 assert!(cred.is_ok());
1112 }
1113
1114 #[test]
1115 fn workload_identity_errors_listing_all_missing_vars() {
1116 let _guard = ENV_MUTEX.lock().unwrap();
1117 clear_entra_env();
1118 let err = WorkloadIdentityCredential::new(TEST_SCOPE).err().unwrap();
1119 clear_entra_env();
1120 let msg = err.to_string();
1121 assert!(msg.contains("AZURE_TENANT_ID"), "{msg}");
1122 assert!(msg.contains("AZURE_CLIENT_ID"), "{msg}");
1123 assert!(msg.contains("AZURE_FEDERATED_TOKEN_FILE"), "{msg}");
1124 }
1125
1126 #[test]
1127 fn workload_identity_client_id_override_fills_in_for_missing_env_var() {
1128 let _guard = ENV_MUTEX.lock().unwrap();
1129 clear_entra_env();
1130 unsafe {
1131 std::env::set_var("AZURE_TENANT_ID", "tenant");
1132 std::env::set_var("AZURE_FEDERATED_TOKEN_FILE", "placeholder-token-file");
1133 }
1134 let cred = WorkloadIdentityCredential::with_client_id(
1135 TEST_SCOPE,
1136 Some("explicit-client".to_string()),
1137 );
1138 clear_entra_env();
1139 assert!(cred.is_ok());
1140 }
1141
1142 #[test]
1143 fn workload_identity_request_params_build_client_assertion_shape() {
1144 let _guard = ENV_MUTEX.lock().unwrap();
1145 let path = std::env::temp_dir().join(format!(
1147 "af-workload-identity-test-{}-{:?}.jwt",
1148 std::process::id(),
1149 std::thread::current().id()
1150 ));
1151 std::fs::write(&path, "fake.jwt.token").unwrap();
1152
1153 let cred = WorkloadIdentityCredential::with_overrides(
1154 TEST_SCOPE,
1155 Some("tenant-1".to_string()),
1156 Some("client-1".to_string()),
1157 Some(path.to_str().unwrap().to_string()),
1158 )
1159 .unwrap();
1160
1161 let params = cred.request_params(TEST_SCOPE).unwrap();
1162 let get = |key: &str| {
1163 params
1164 .iter()
1165 .find(|(k, _)| k == key)
1166 .map(|(_, v)| v.clone())
1167 };
1168 assert_eq!(get("grant_type").as_deref(), Some("client_credentials"));
1169 assert_eq!(get("client_id").as_deref(), Some("client-1"));
1170 assert_eq!(get("client_assertion").as_deref(), Some("fake.jwt.token"));
1171 assert_eq!(
1172 get("client_assertion_type").as_deref(),
1173 Some("urn:ietf:params:oauth:client-assertion-type:jwt-bearer")
1174 );
1175 assert_eq!(get("scope").as_deref(), Some(TEST_SCOPE));
1176 assert_eq!(
1177 cred.token_url(),
1178 "https://login.microsoftonline.com/tenant-1/oauth2/v2.0/token"
1179 );
1180
1181 std::fs::remove_file(&path).ok();
1182 }
1183
1184 #[test]
1185 fn workload_identity_rereads_token_file_on_each_request() {
1186 let path = std::env::temp_dir().join(format!(
1187 "af-workload-identity-rotate-{}-{:?}.jwt",
1188 std::process::id(),
1189 std::thread::current().id()
1190 ));
1191 std::fs::write(&path, "first-token").unwrap();
1192
1193 let cred = WorkloadIdentityCredential::with_overrides(
1194 TEST_SCOPE,
1195 Some("tenant-1".to_string()),
1196 Some("client-1".to_string()),
1197 Some(path.to_str().unwrap().to_string()),
1198 )
1199 .unwrap();
1200
1201 let first = cred.request_params("scope").unwrap();
1202 let assertion_of = |params: &[(String, String)]| {
1203 params
1204 .iter()
1205 .find(|(k, _)| k == "client_assertion")
1206 .map(|(_, v)| v.clone())
1207 };
1208 assert_eq!(assertion_of(&first).as_deref(), Some("first-token"));
1209
1210 std::fs::write(&path, "rotated-token").unwrap();
1211 let second = cred.request_params("scope").unwrap();
1212 assert_eq!(assertion_of(&second).as_deref(), Some("rotated-token"));
1213
1214 std::fs::remove_file(&path).ok();
1215 }
1216
1217 #[test]
1218 fn workload_identity_missing_token_file_is_a_clear_error() {
1219 let cred = WorkloadIdentityCredential::with_overrides(
1220 TEST_SCOPE,
1221 Some("tenant-1".to_string()),
1222 Some("client-1".to_string()),
1223 Some("/definitely/not/a/real/path/af-test.jwt".to_string()),
1224 )
1225 .unwrap();
1226 let err = cred.request_params(TEST_SCOPE).unwrap_err();
1227 assert!(err.to_string().contains("federated token file"), "{err}");
1228 }
1229
1230 #[test]
1235 fn default_chain_includes_only_always_on_credentials_without_env() {
1236 let _guard = ENV_MUTEX.lock().unwrap();
1237 clear_entra_env();
1238 let chain = default_chain(TEST_SCOPE);
1239 clear_entra_env();
1240 assert_eq!(chain.len(), 2);
1243 }
1244
1245 #[test]
1246 fn default_chain_adds_environment_credential_when_configured() {
1247 let _guard = ENV_MUTEX.lock().unwrap();
1248 clear_entra_env();
1249 unsafe {
1250 std::env::set_var("AZURE_TENANT_ID", "tenant");
1251 std::env::set_var("AZURE_CLIENT_ID", "client");
1252 std::env::set_var("AZURE_CLIENT_SECRET", "secret");
1253 }
1254 let chain = default_chain(TEST_SCOPE);
1255 clear_entra_env();
1256 assert_eq!(chain.len(), 3, "Environment + ManagedIdentity + AzureCli");
1257 }
1258
1259 #[test]
1260 fn default_chain_adds_workload_identity_when_configured() {
1261 let _guard = ENV_MUTEX.lock().unwrap();
1262 clear_entra_env();
1263 unsafe {
1264 std::env::set_var("AZURE_TENANT_ID", "tenant");
1265 std::env::set_var("AZURE_CLIENT_ID", "client");
1266 std::env::set_var("AZURE_FEDERATED_TOKEN_FILE", "placeholder-token-file");
1267 }
1268 let chain = default_chain(TEST_SCOPE);
1269 clear_entra_env();
1270 assert_eq!(
1271 chain.len(),
1272 3,
1273 "WorkloadIdentity + ManagedIdentity + AzureCli"
1274 );
1275 }
1276
1277 #[test]
1278 fn default_chain_adds_both_environment_and_workload_identity_in_order() {
1279 let _guard = ENV_MUTEX.lock().unwrap();
1280 clear_entra_env();
1281 unsafe {
1282 std::env::set_var("AZURE_TENANT_ID", "tenant");
1283 std::env::set_var("AZURE_CLIENT_ID", "client");
1284 std::env::set_var("AZURE_CLIENT_SECRET", "secret");
1285 std::env::set_var("AZURE_FEDERATED_TOKEN_FILE", "placeholder-token-file");
1286 }
1287 let chain = default_chain(TEST_SCOPE);
1288 clear_entra_env();
1289 assert_eq!(chain.len(), 4, "all four credentials configured");
1290 }
1291
1292 #[test]
1293 fn default_azure_credential_wraps_a_never_empty_chain() {
1294 let _guard = ENV_MUTEX.lock().unwrap();
1295 clear_entra_env();
1296 let cred = DefaultAzureCredential::new(TEST_SCOPE);
1297 clear_entra_env();
1298 assert_eq!(cred.chain.credentials.len(), 2);
1299 }
1300
1301 }