1use std::collections::HashMap;
4use std::process::Command;
5use std::sync::{Mutex, OnceLock};
6use std::time::{Duration, Instant};
7use thiserror::Error;
8
9struct TokenCache {
16 ttl: Duration,
17 entries: Mutex<HashMap<String, (String, Instant)>>,
18}
19
20impl TokenCache {
21 fn new(ttl: Duration) -> Self {
22 Self {
23 ttl,
24 entries: Mutex::new(HashMap::new()),
25 }
26 }
27
28 fn get_or_fetch(
29 &self,
30 scope: &str,
31 fetch: impl FnOnce() -> Result<String, AuthError>,
32 ) -> Result<String, AuthError> {
33 let mut entries = self.entries.lock().unwrap();
34 if let Some((token, acquired)) = entries.get(scope)
35 && acquired.elapsed() < self.ttl
36 {
37 return Ok(token.clone());
38 }
39 let token = fetch()?;
40 entries.insert(scope.to_string(), (token.clone(), Instant::now()));
41 Ok(token)
42 }
43}
44
45fn az_token_cache() -> &'static TokenCache {
46 static CACHE: OnceLock<TokenCache> = OnceLock::new();
47 CACHE.get_or_init(|| TokenCache::new(Duration::from_secs(300)))
48}
49
50#[derive(Debug, Error)]
52pub enum AuthError {
53 #[error(
54 "Azure CLI not found. Please install it: https://docs.microsoft.com/cli/azure/install-azure-cli"
55 )]
56 AzCliNotFound,
57 #[error("Not logged in to Azure CLI. Run: az login")]
58 NotLoggedIn,
59 #[error("Failed to get access token: {0}")]
60 TokenError(String),
61 #[error("Missing environment variable: {0}")]
62 MissingEnvVar(String),
63 #[error("Authentication failed: {0}")]
64 AuthFailed(String),
65}
66
67fn account_parse_error(e: serde_json::Error) -> AuthError {
74 AuthError::TokenError(format!(
75 "could not parse `az account show` output ({e}); this is usually a \
76 transient Azure CLI issue — try again, and if it persists run `az login`"
77 ))
78}
79
80fn token_error_detail(stderr: &str, status: std::process::ExitStatus) -> String {
84 if stderr.trim().is_empty() {
85 format!(
86 "az returned no error detail (exit {status}); usually transient — try again, or run `az login`"
87 )
88 } else {
89 stderr.trim().to_string()
90 }
91}
92
93pub trait AuthProvider: Send + Sync {
95 fn get_token(&self) -> Result<String, AuthError>;
97
98 fn method_name(&self) -> &'static str;
100}
101
102pub struct AzCliAuth {
104 resource_scope: &'static str,
105}
106
107impl AzCliAuth {
108 pub fn for_search() -> Self {
110 Self {
111 resource_scope: "https://search.azure.com",
112 }
113 }
114
115 pub fn for_foundry() -> Self {
117 Self {
118 resource_scope: "https://ai.azure.com",
119 }
120 }
121
122 pub fn for_cognitive_services() -> Self {
124 Self {
125 resource_scope: "https://cognitiveservices.azure.com",
126 }
127 }
128
129 pub fn for_cosmos() -> Self {
131 Self {
132 resource_scope: "https://cosmos.azure.com",
133 }
134 }
135
136 pub fn new() -> Self {
138 Self::for_search()
139 }
140
141 pub fn check_status() -> Result<AuthStatus, AuthError> {
143 let version_output = Command::new("az").arg("--version").output();
145
146 if version_output.is_err() {
147 return Err(AuthError::AzCliNotFound);
148 }
149
150 let account_output = Command::new("az")
152 .args(["account", "show", "--output", "json"])
153 .output()
154 .map_err(|e| AuthError::TokenError(e.to_string()))?;
155
156 if !account_output.status.success() {
157 return Err(AuthError::NotLoggedIn);
158 }
159
160 let account_json: serde_json::Value =
162 serde_json::from_slice(&account_output.stdout).map_err(account_parse_error)?;
163
164 Ok(AuthStatus {
165 logged_in: true,
166 user: account_json
167 .get("user")
168 .and_then(|u| u.get("name"))
169 .and_then(|n| n.as_str())
170 .map(String::from),
171 subscription: account_json
172 .get("name")
173 .and_then(|n| n.as_str())
174 .map(String::from),
175 subscription_id: account_json
176 .get("id")
177 .and_then(|i| i.as_str())
178 .map(String::from),
179 })
180 }
181
182 pub fn get_arm_token() -> Result<String, AuthError> {
184 az_token_cache().get_or_fetch("https://management.azure.com", Self::fetch_arm_token)
185 }
186
187 fn fetch_arm_token() -> Result<String, AuthError> {
188 let output = Command::new("az")
189 .args([
190 "account",
191 "get-access-token",
192 "--resource",
193 "https://management.azure.com",
194 "--query",
195 "accessToken",
196 "--output",
197 "tsv",
198 ])
199 .output()
200 .map_err(|e| AuthError::TokenError(e.to_string()))?;
201
202 if !output.status.success() {
203 let stderr = String::from_utf8_lossy(&output.stderr);
204 if stderr.contains("not logged in") || stderr.contains("AADSTS") {
205 return Err(AuthError::NotLoggedIn);
206 }
207 return Err(AuthError::TokenError(token_error_detail(
208 &stderr,
209 output.status,
210 )));
211 }
212
213 let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
214 if token.is_empty() {
215 return Err(AuthError::TokenError(
216 "Empty ARM token received".to_string(),
217 ));
218 }
219
220 Ok(token)
221 }
222}
223
224impl Default for AzCliAuth {
225 fn default() -> Self {
226 Self::new()
227 }
228}
229
230impl AuthProvider for AzCliAuth {
231 fn get_token(&self) -> Result<String, AuthError> {
232 az_token_cache().get_or_fetch(self.resource_scope, || self.fetch_token())
233 }
234
235 fn method_name(&self) -> &'static str {
236 "Azure CLI"
237 }
238}
239
240impl AzCliAuth {
241 fn fetch_token(&self) -> Result<String, AuthError> {
242 let output = Command::new("az")
243 .args([
244 "account",
245 "get-access-token",
246 "--resource",
247 self.resource_scope,
248 "--query",
249 "accessToken",
250 "--output",
251 "tsv",
252 ])
253 .output()
254 .map_err(|e| AuthError::TokenError(e.to_string()))?;
255
256 if !output.status.success() {
257 let stderr = String::from_utf8_lossy(&output.stderr);
258 if stderr.contains("not logged in") {
259 return Err(AuthError::NotLoggedIn);
260 }
261 if stderr.contains("AADSTS") {
262 let detail = stderr
264 .lines()
265 .find(|l| l.contains("AADSTS"))
266 .unwrap_or(&stderr)
267 .trim();
268 return Err(AuthError::TokenError(format!(
269 "Failed to get access token for {}: {}\n \
270 Debug: az account get-access-token --resource {}\n \
271 Fix: Ensure 'Cognitive Services User' role is assigned on the AI Services resource",
272 self.resource_scope, detail, self.resource_scope
273 )));
274 }
275 return Err(AuthError::TokenError(token_error_detail(
276 &stderr,
277 output.status,
278 )));
279 }
280
281 let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
282 if token.is_empty() {
283 return Err(AuthError::TokenError("Empty token received".to_string()));
284 }
285
286 Ok(token)
287 }
288}
289
290#[derive(Debug)]
292pub struct EnvAuth {
293 client_id: String,
294 client_secret: String,
295 tenant_id: String,
296 resource_scope: &'static str,
297}
298
299impl EnvAuth {
300 pub fn from_env() -> Result<Self, AuthError> {
302 Self::from_env_for_scope("https://search.azure.com")
303 }
304
305 pub fn from_env_for_scope(scope: &'static str) -> Result<Self, AuthError> {
307 let client_id = std::env::var("AZURE_CLIENT_ID")
308 .map_err(|_| AuthError::MissingEnvVar("AZURE_CLIENT_ID".to_string()))?;
309 let client_secret = std::env::var("AZURE_CLIENT_SECRET")
310 .map_err(|_| AuthError::MissingEnvVar("AZURE_CLIENT_SECRET".to_string()))?;
311 let tenant_id = std::env::var("AZURE_TENANT_ID")
312 .map_err(|_| AuthError::MissingEnvVar("AZURE_TENANT_ID".to_string()))?;
313
314 Ok(Self {
315 client_id,
316 client_secret,
317 tenant_id,
318 resource_scope: scope,
319 })
320 }
321
322 pub fn is_configured() -> bool {
324 std::env::var("AZURE_CLIENT_ID").is_ok()
325 && std::env::var("AZURE_CLIENT_SECRET").is_ok()
326 && std::env::var("AZURE_TENANT_ID").is_ok()
327 }
328}
329
330impl AuthProvider for EnvAuth {
331 fn get_token(&self) -> Result<String, AuthError> {
332 let output = Command::new("az")
334 .args([
335 "account",
336 "get-access-token",
337 "--resource",
338 self.resource_scope,
339 "--query",
340 "accessToken",
341 "--output",
342 "tsv",
343 "--tenant",
344 &self.tenant_id,
345 "--username",
346 &self.client_id,
347 ])
348 .env("AZURE_CLIENT_SECRET", &self.client_secret)
349 .output()
350 .map_err(|e| AuthError::TokenError(e.to_string()))?;
351
352 if !output.status.success() {
353 let stderr = String::from_utf8_lossy(&output.stderr);
354 return Err(AuthError::AuthFailed(stderr.to_string()));
355 }
356
357 let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
358 Ok(token)
359 }
360
361 fn method_name(&self) -> &'static str {
362 "Environment Variables (Service Principal)"
363 }
364}
365
366#[derive(Debug, Clone)]
368pub struct AuthStatus {
369 pub logged_in: bool,
370 pub user: Option<String>,
371 pub subscription: Option<String>,
372 pub subscription_id: Option<String>,
373}
374
375pub fn get_auth_provider() -> Result<Box<dyn AuthProvider>, AuthError> {
377 get_auth_provider_for_scope("https://search.azure.com")
378}
379
380pub fn get_auth_provider_for(
382 domain: rigg_core::ServiceDomain,
383) -> Result<Box<dyn AuthProvider>, AuthError> {
384 let scope = match domain {
385 rigg_core::ServiceDomain::Search => "https://search.azure.com",
386 rigg_core::ServiceDomain::Foundry => "https://ai.azure.com",
387 };
388 get_auth_provider_for_scope(scope)
389}
390
391pub fn get_cognitive_services_auth() -> Result<Box<dyn AuthProvider>, AuthError> {
393 get_auth_provider_for_scope("https://cognitiveservices.azure.com")
394}
395
396pub struct StaticTokenAuth {
399 token: String,
400}
401
402impl AuthProvider for StaticTokenAuth {
403 fn get_token(&self) -> Result<String, AuthError> {
404 Ok(self.token.clone())
405 }
406 fn method_name(&self) -> &'static str {
407 "Static token (RIGG_ACCESS_TOKEN)"
408 }
409}
410
411fn get_auth_provider_for_scope(scope: &'static str) -> Result<Box<dyn AuthProvider>, AuthError> {
413 if let Ok(token) = std::env::var("RIGG_ACCESS_TOKEN") {
415 if !token.is_empty() {
416 return Ok(Box::new(StaticTokenAuth { token }));
417 }
418 }
419 if EnvAuth::is_configured() {
421 return Ok(Box::new(EnvAuth::from_env_for_scope(scope)?));
422 }
423
424 AzCliAuth::check_status()?;
426 Ok(Box::new(AzCliAuth {
427 resource_scope: scope,
428 }))
429}
430
431#[cfg(test)]
432mod tests {
433 use super::*;
434 use std::sync::Mutex;
435
436 static ENV_MUTEX: Mutex<()> = Mutex::new(());
438
439 unsafe fn clear_azure_env_vars() {
442 unsafe {
443 std::env::remove_var("AZURE_CLIENT_ID");
444 std::env::remove_var("AZURE_CLIENT_SECRET");
445 std::env::remove_var("AZURE_TENANT_ID");
446 }
447 }
448
449 unsafe fn set_azure_env_vars() {
452 unsafe {
453 std::env::set_var("AZURE_CLIENT_ID", "test-client-id");
454 std::env::set_var("AZURE_CLIENT_SECRET", "test-client-secret");
455 std::env::set_var("AZURE_TENANT_ID", "test-tenant-id");
456 }
457 }
458
459 #[test]
460 fn test_env_auth_from_env_success() {
461 let _lock = ENV_MUTEX.lock().unwrap();
462 unsafe { set_azure_env_vars() };
463
464 let result = EnvAuth::from_env();
465 assert!(result.is_ok());
466 let auth = result.unwrap();
467 assert_eq!(auth.client_id, "test-client-id");
468 assert_eq!(auth.client_secret, "test-client-secret");
469 assert_eq!(auth.tenant_id, "test-tenant-id");
470
471 unsafe { clear_azure_env_vars() };
472 }
473
474 #[test]
475 fn test_env_auth_from_env_missing_client_id() {
476 let _lock = ENV_MUTEX.lock().unwrap();
477 unsafe {
478 clear_azure_env_vars();
479 std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
480 std::env::set_var("AZURE_TENANT_ID", "test-tenant");
481 }
482
483 let result = EnvAuth::from_env();
484 assert!(result.is_err());
485 let err = result.unwrap_err();
486 assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_CLIENT_ID"));
487
488 unsafe { clear_azure_env_vars() };
489 }
490
491 #[test]
492 fn test_env_auth_from_env_missing_client_secret() {
493 let _lock = ENV_MUTEX.lock().unwrap();
494 unsafe {
495 clear_azure_env_vars();
496 std::env::set_var("AZURE_CLIENT_ID", "test-id");
497 std::env::set_var("AZURE_TENANT_ID", "test-tenant");
498 }
499
500 let result = EnvAuth::from_env();
501 assert!(result.is_err());
502 let err = result.unwrap_err();
503 assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_CLIENT_SECRET"));
504
505 unsafe { clear_azure_env_vars() };
506 }
507
508 #[test]
509 fn test_env_auth_from_env_missing_tenant_id() {
510 let _lock = ENV_MUTEX.lock().unwrap();
511 unsafe {
512 clear_azure_env_vars();
513 std::env::set_var("AZURE_CLIENT_ID", "test-id");
514 std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
515 }
516
517 let result = EnvAuth::from_env();
518 assert!(result.is_err());
519 let err = result.unwrap_err();
520 assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_TENANT_ID"));
521
522 unsafe { clear_azure_env_vars() };
523 }
524
525 #[test]
526 fn test_env_auth_is_configured_all_set() {
527 let _lock = ENV_MUTEX.lock().unwrap();
528 unsafe { set_azure_env_vars() };
529
530 assert!(EnvAuth::is_configured());
531
532 unsafe { clear_azure_env_vars() };
533 }
534
535 #[test]
536 fn test_env_auth_is_configured_none_set() {
537 let _lock = ENV_MUTEX.lock().unwrap();
538 unsafe { clear_azure_env_vars() };
539
540 assert!(!EnvAuth::is_configured());
541 }
542
543 #[test]
544 fn test_env_auth_is_configured_partial() {
545 let _lock = ENV_MUTEX.lock().unwrap();
546 unsafe {
547 clear_azure_env_vars();
548 std::env::set_var("AZURE_CLIENT_ID", "test-id");
549 std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
550 }
551 assert!(!EnvAuth::is_configured());
554
555 unsafe { clear_azure_env_vars() };
556 }
557
558 #[test]
559 fn test_env_auth_method_name() {
560 let _lock = ENV_MUTEX.lock().unwrap();
561 unsafe { set_azure_env_vars() };
562
563 let auth = EnvAuth::from_env().unwrap();
564 assert_eq!(
565 auth.method_name(),
566 "Environment Variables (Service Principal)"
567 );
568
569 unsafe { clear_azure_env_vars() };
570 }
571
572 #[test]
573 fn test_az_cli_auth_method_name() {
574 let auth = AzCliAuth::new();
575 assert_eq!(auth.method_name(), "Azure CLI");
576 }
577
578 #[test]
579 fn test_az_cli_auth_search_scope() {
580 let auth = AzCliAuth::for_search();
581 assert_eq!(auth.resource_scope, "https://search.azure.com");
582 }
583
584 #[test]
585 fn test_az_cli_auth_foundry_scope() {
586 let auth = AzCliAuth::for_foundry();
587 assert_eq!(auth.resource_scope, "https://ai.azure.com");
588 }
589
590 #[test]
591 fn test_az_cli_auth_cognitive_services_scope() {
592 let auth = AzCliAuth::for_cognitive_services();
593 assert_eq!(auth.resource_scope, "https://cognitiveservices.azure.com");
594 }
595
596 #[test]
597 fn test_az_cli_auth_new_defaults_to_search() {
598 let auth = AzCliAuth::new();
599 assert_eq!(auth.resource_scope, "https://search.azure.com");
600 }
601
602 #[test]
603 fn test_env_auth_from_env_scope_foundry() {
604 let _lock = ENV_MUTEX.lock().unwrap();
605 unsafe { set_azure_env_vars() };
606
607 let result = EnvAuth::from_env_for_scope("https://ai.azure.com");
608 assert!(result.is_ok());
609 let auth = result.unwrap();
610 assert_eq!(auth.resource_scope, "https://ai.azure.com");
611
612 unsafe { clear_azure_env_vars() };
613 }
614
615 #[test]
616 fn test_env_auth_from_env_default_scope_is_search() {
617 let _lock = ENV_MUTEX.lock().unwrap();
618 unsafe { set_azure_env_vars() };
619
620 let auth = EnvAuth::from_env().unwrap();
621 assert_eq!(auth.resource_scope, "https://search.azure.com");
622
623 unsafe { clear_azure_env_vars() };
624 }
625
626 #[test]
627 fn test_auth_status_fields() {
628 let status = AuthStatus {
629 logged_in: true,
630 user: Some("testuser@example.com".to_string()),
631 subscription: Some("My Subscription".to_string()),
632 subscription_id: Some("00000000-0000-0000-0000-000000000000".to_string()),
633 };
634
635 assert!(status.logged_in);
636 assert_eq!(status.user.as_deref(), Some("testuser@example.com"));
637 assert_eq!(status.subscription.as_deref(), Some("My Subscription"));
638 assert_eq!(
639 status.subscription_id.as_deref(),
640 Some("00000000-0000-0000-0000-000000000000")
641 );
642 }
643
644 #[test]
645 fn test_for_cosmos_uses_cosmos_scope() {
646 let auth = AzCliAuth::for_cosmos();
647 assert_eq!(auth.resource_scope, "https://cosmos.azure.com");
648 }
649
650 #[test]
651 fn account_parse_error_is_actionable_not_raw_serde() {
652 let serde_err = serde_json::from_slice::<serde_json::Value>(b"").unwrap_err();
653 let err = account_parse_error(serde_err);
654 match err {
655 AuthError::TokenError(msg) => {
656 assert!(msg.contains("az login"), "{msg}");
657 assert!(msg.contains("transient"), "{msg}");
658 assert!(msg.contains("az account show"), "{msg}");
659 }
660 other => panic!("expected TokenError, got {other:?}"),
661 }
662 }
663
664 #[test]
665 fn token_error_detail_falls_back_when_stderr_empty() {
666 let status = std::process::Command::new("true")
667 .status()
668 .expect("failed to run `true`");
669 let detail = token_error_detail(" \n", status);
670 assert!(detail.contains("az login"), "{detail}");
671 assert!(detail.contains("transient"), "{detail}");
672 assert!(!detail.trim().is_empty());
673 }
674
675 #[test]
676 fn token_error_detail_preserves_nonempty_stderr() {
677 let status = std::process::Command::new("false")
678 .status()
679 .expect("failed to run `false`");
680 let detail = token_error_detail(" ERROR: something specific broke ", status);
681 assert_eq!(detail, "ERROR: something specific broke");
682 }
683}
684
685#[cfg(test)]
686mod token_cache_tests {
687 use super::*;
688 use std::cell::Cell;
689 use std::time::Duration;
690
691 #[test]
692 fn second_lookup_within_ttl_reuses_token_without_fetching() {
693 let cache = TokenCache::new(Duration::from_secs(300));
694 let fetches = Cell::new(0u32);
695 let fetch = || {
696 fetches.set(fetches.get() + 1);
697 Ok("tok-1".to_string())
698 };
699 assert_eq!(cache.get_or_fetch("scope-a", fetch).unwrap(), "tok-1");
700 assert_eq!(
701 cache
702 .get_or_fetch("scope-a", || panic!("must not fetch"))
703 .unwrap(),
704 "tok-1"
705 );
706 assert_eq!(fetches.get(), 1);
707 }
708
709 #[test]
710 fn expired_entry_is_refetched() {
711 let cache = TokenCache::new(Duration::ZERO);
712 cache
713 .get_or_fetch("scope-a", || Ok("old".to_string()))
714 .unwrap();
715 let got = cache
716 .get_or_fetch("scope-a", || Ok("new".to_string()))
717 .unwrap();
718 assert_eq!(got, "new");
719 }
720
721 #[test]
722 fn scopes_are_cached_independently() {
723 let cache = TokenCache::new(Duration::from_secs(300));
724 cache
725 .get_or_fetch("scope-a", || Ok("tok-a".to_string()))
726 .unwrap();
727 let got = cache
728 .get_or_fetch("scope-b", || Ok("tok-b".to_string()))
729 .unwrap();
730 assert_eq!(got, "tok-b");
731 }
732
733 #[test]
734 fn fetch_errors_are_not_cached() {
735 let cache = TokenCache::new(Duration::from_secs(300));
736 let err = cache.get_or_fetch("scope-a", || Err(AuthError::TokenError("boom".to_string())));
737 assert!(err.is_err());
738 let got = cache
739 .get_or_fetch("scope-a", || Ok("recovered".to_string()))
740 .unwrap();
741 assert_eq!(got, "recovered");
742 }
743}