rigg_client/
auth.rs

1//! Azure authentication providers
2
3use std::process::Command;
4use thiserror::Error;
5
6/// Authentication errors
7#[derive(Debug, Error)]
8pub enum AuthError {
9    #[error(
10        "Azure CLI not found. Please install it: https://docs.microsoft.com/cli/azure/install-azure-cli"
11    )]
12    AzCliNotFound,
13    #[error("Not logged in to Azure CLI. Run: az login")]
14    NotLoggedIn,
15    #[error("Failed to get access token: {0}")]
16    TokenError(String),
17    #[error("Missing environment variable: {0}")]
18    MissingEnvVar(String),
19    #[error("Authentication failed: {0}")]
20    AuthFailed(String),
21}
22
23/// Authentication provider trait
24pub trait AuthProvider: Send + Sync {
25    /// Get an access token for Azure Search
26    fn get_token(&self) -> Result<String, AuthError>;
27
28    /// Get the authentication method name
29    fn method_name(&self) -> &'static str;
30}
31
32/// Azure CLI authentication provider
33pub struct AzCliAuth {
34    resource_scope: &'static str,
35}
36
37impl AzCliAuth {
38    /// Create an auth provider for Azure Search
39    pub fn for_search() -> Self {
40        Self {
41            resource_scope: "https://search.azure.com",
42        }
43    }
44
45    /// Create an auth provider for Microsoft Foundry
46    pub fn for_foundry() -> Self {
47        Self {
48            resource_scope: "https://ai.azure.com",
49        }
50    }
51
52    /// Create an auth provider for Azure Cognitive Services (OpenAI)
53    pub fn for_cognitive_services() -> Self {
54        Self {
55            resource_scope: "https://cognitiveservices.azure.com",
56        }
57    }
58
59    /// Create a new auth provider (defaults to Search scope for backward compatibility)
60    pub fn new() -> Self {
61        Self::for_search()
62    }
63
64    /// Check if Azure CLI is available and logged in
65    pub fn check_status() -> Result<AuthStatus, AuthError> {
66        // Check if az CLI is installed
67        let version_output = Command::new("az").arg("--version").output();
68
69        if version_output.is_err() {
70            return Err(AuthError::AzCliNotFound);
71        }
72
73        // Check if logged in
74        let account_output = Command::new("az")
75            .args(["account", "show", "--output", "json"])
76            .output()
77            .map_err(|e| AuthError::TokenError(e.to_string()))?;
78
79        if !account_output.status.success() {
80            return Err(AuthError::NotLoggedIn);
81        }
82
83        // Parse account info
84        let account_json: serde_json::Value = serde_json::from_slice(&account_output.stdout)
85            .map_err(|e| AuthError::TokenError(e.to_string()))?;
86
87        Ok(AuthStatus {
88            logged_in: true,
89            user: account_json
90                .get("user")
91                .and_then(|u| u.get("name"))
92                .and_then(|n| n.as_str())
93                .map(String::from),
94            subscription: account_json
95                .get("name")
96                .and_then(|n| n.as_str())
97                .map(String::from),
98            subscription_id: account_json
99                .get("id")
100                .and_then(|i| i.as_str())
101                .map(String::from),
102        })
103    }
104
105    /// Get an access token for Azure Resource Manager (management.azure.com)
106    pub fn get_arm_token() -> Result<String, AuthError> {
107        let output = Command::new("az")
108            .args([
109                "account",
110                "get-access-token",
111                "--resource",
112                "https://management.azure.com",
113                "--query",
114                "accessToken",
115                "--output",
116                "tsv",
117            ])
118            .output()
119            .map_err(|e| AuthError::TokenError(e.to_string()))?;
120
121        if !output.status.success() {
122            let stderr = String::from_utf8_lossy(&output.stderr);
123            if stderr.contains("not logged in") || stderr.contains("AADSTS") {
124                return Err(AuthError::NotLoggedIn);
125            }
126            return Err(AuthError::TokenError(stderr.to_string()));
127        }
128
129        let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
130        if token.is_empty() {
131            return Err(AuthError::TokenError(
132                "Empty ARM token received".to_string(),
133            ));
134        }
135
136        Ok(token)
137    }
138}
139
140impl Default for AzCliAuth {
141    fn default() -> Self {
142        Self::new()
143    }
144}
145
146impl AuthProvider for AzCliAuth {
147    fn get_token(&self) -> Result<String, AuthError> {
148        let output = Command::new("az")
149            .args([
150                "account",
151                "get-access-token",
152                "--resource",
153                self.resource_scope,
154                "--query",
155                "accessToken",
156                "--output",
157                "tsv",
158            ])
159            .output()
160            .map_err(|e| AuthError::TokenError(e.to_string()))?;
161
162        if !output.status.success() {
163            let stderr = String::from_utf8_lossy(&output.stderr);
164            if stderr.contains("not logged in") {
165                return Err(AuthError::NotLoggedIn);
166            }
167            if stderr.contains("AADSTS") {
168                // Extract the first AADSTS error line for a concise message
169                let detail = stderr
170                    .lines()
171                    .find(|l| l.contains("AADSTS"))
172                    .unwrap_or(&stderr)
173                    .trim();
174                return Err(AuthError::TokenError(format!(
175                    "Failed to get access token for {}: {}\n  \
176                     Debug: az account get-access-token --resource {}\n  \
177                     Fix: Ensure 'Cognitive Services User' role is assigned on the AI Services resource",
178                    self.resource_scope, detail, self.resource_scope
179                )));
180            }
181            return Err(AuthError::TokenError(stderr.to_string()));
182        }
183
184        let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
185        if token.is_empty() {
186            return Err(AuthError::TokenError("Empty token received".to_string()));
187        }
188
189        Ok(token)
190    }
191
192    fn method_name(&self) -> &'static str {
193        "Azure CLI"
194    }
195}
196
197/// Environment variable authentication provider
198#[derive(Debug)]
199pub struct EnvAuth {
200    client_id: String,
201    client_secret: String,
202    tenant_id: String,
203    resource_scope: &'static str,
204}
205
206impl EnvAuth {
207    /// Create from environment variables (defaults to Search scope)
208    pub fn from_env() -> Result<Self, AuthError> {
209        Self::from_env_for_scope("https://search.azure.com")
210    }
211
212    /// Create from environment variables for a specific resource scope
213    pub fn from_env_for_scope(scope: &'static str) -> Result<Self, AuthError> {
214        let client_id = std::env::var("AZURE_CLIENT_ID")
215            .map_err(|_| AuthError::MissingEnvVar("AZURE_CLIENT_ID".to_string()))?;
216        let client_secret = std::env::var("AZURE_CLIENT_SECRET")
217            .map_err(|_| AuthError::MissingEnvVar("AZURE_CLIENT_SECRET".to_string()))?;
218        let tenant_id = std::env::var("AZURE_TENANT_ID")
219            .map_err(|_| AuthError::MissingEnvVar("AZURE_TENANT_ID".to_string()))?;
220
221        Ok(Self {
222            client_id,
223            client_secret,
224            tenant_id,
225            resource_scope: scope,
226        })
227    }
228
229    /// Check if environment variables are set
230    pub fn is_configured() -> bool {
231        std::env::var("AZURE_CLIENT_ID").is_ok()
232            && std::env::var("AZURE_CLIENT_SECRET").is_ok()
233            && std::env::var("AZURE_TENANT_ID").is_ok()
234    }
235}
236
237impl AuthProvider for EnvAuth {
238    fn get_token(&self) -> Result<String, AuthError> {
239        // Use Azure CLI to get token with service principal
240        let output = Command::new("az")
241            .args([
242                "account",
243                "get-access-token",
244                "--resource",
245                self.resource_scope,
246                "--query",
247                "accessToken",
248                "--output",
249                "tsv",
250                "--tenant",
251                &self.tenant_id,
252                "--username",
253                &self.client_id,
254            ])
255            .env("AZURE_CLIENT_SECRET", &self.client_secret)
256            .output()
257            .map_err(|e| AuthError::TokenError(e.to_string()))?;
258
259        if !output.status.success() {
260            let stderr = String::from_utf8_lossy(&output.stderr);
261            return Err(AuthError::AuthFailed(stderr.to_string()));
262        }
263
264        let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
265        Ok(token)
266    }
267
268    fn method_name(&self) -> &'static str {
269        "Environment Variables (Service Principal)"
270    }
271}
272
273/// Authentication status
274#[derive(Debug, Clone)]
275pub struct AuthStatus {
276    pub logged_in: bool,
277    pub user: Option<String>,
278    pub subscription: Option<String>,
279    pub subscription_id: Option<String>,
280}
281
282/// Get the best available authentication provider for Search (backward compat)
283pub fn get_auth_provider() -> Result<Box<dyn AuthProvider>, AuthError> {
284    get_auth_provider_for_scope("https://search.azure.com")
285}
286
287/// Get the best available authentication provider for a specific service domain
288pub fn get_auth_provider_for(
289    domain: rigg_core::ServiceDomain,
290) -> Result<Box<dyn AuthProvider>, AuthError> {
291    let scope = match domain {
292        rigg_core::ServiceDomain::Search => "https://search.azure.com",
293        rigg_core::ServiceDomain::Foundry => "https://ai.azure.com",
294    };
295    get_auth_provider_for_scope(scope)
296}
297
298/// Get the best available authentication provider for Azure Cognitive Services (OpenAI)
299pub fn get_cognitive_services_auth() -> Result<Box<dyn AuthProvider>, AuthError> {
300    get_auth_provider_for_scope("https://cognitiveservices.azure.com")
301}
302
303/// Get the best available authentication provider for a specific resource scope
304fn get_auth_provider_for_scope(scope: &'static str) -> Result<Box<dyn AuthProvider>, AuthError> {
305    // First try environment variables
306    if EnvAuth::is_configured() {
307        return Ok(Box::new(EnvAuth::from_env_for_scope(scope)?));
308    }
309
310    // Fall back to Azure CLI
311    AzCliAuth::check_status()?;
312    Ok(Box::new(AzCliAuth {
313        resource_scope: scope,
314    }))
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320    use std::sync::Mutex;
321
322    // Env var tests must run serially since they share process-wide state.
323    static ENV_MUTEX: Mutex<()> = Mutex::new(());
324
325    /// # Safety
326    /// Must be called while holding ENV_MUTEX to avoid data races.
327    unsafe fn clear_azure_env_vars() {
328        unsafe {
329            std::env::remove_var("AZURE_CLIENT_ID");
330            std::env::remove_var("AZURE_CLIENT_SECRET");
331            std::env::remove_var("AZURE_TENANT_ID");
332        }
333    }
334
335    /// # Safety
336    /// Must be called while holding ENV_MUTEX to avoid data races.
337    unsafe fn set_azure_env_vars() {
338        unsafe {
339            std::env::set_var("AZURE_CLIENT_ID", "test-client-id");
340            std::env::set_var("AZURE_CLIENT_SECRET", "test-client-secret");
341            std::env::set_var("AZURE_TENANT_ID", "test-tenant-id");
342        }
343    }
344
345    #[test]
346    fn test_env_auth_from_env_success() {
347        let _lock = ENV_MUTEX.lock().unwrap();
348        unsafe { set_azure_env_vars() };
349
350        let result = EnvAuth::from_env();
351        assert!(result.is_ok());
352        let auth = result.unwrap();
353        assert_eq!(auth.client_id, "test-client-id");
354        assert_eq!(auth.client_secret, "test-client-secret");
355        assert_eq!(auth.tenant_id, "test-tenant-id");
356
357        unsafe { clear_azure_env_vars() };
358    }
359
360    #[test]
361    fn test_env_auth_from_env_missing_client_id() {
362        let _lock = ENV_MUTEX.lock().unwrap();
363        unsafe {
364            clear_azure_env_vars();
365            std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
366            std::env::set_var("AZURE_TENANT_ID", "test-tenant");
367        }
368
369        let result = EnvAuth::from_env();
370        assert!(result.is_err());
371        let err = result.unwrap_err();
372        assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_CLIENT_ID"));
373
374        unsafe { clear_azure_env_vars() };
375    }
376
377    #[test]
378    fn test_env_auth_from_env_missing_client_secret() {
379        let _lock = ENV_MUTEX.lock().unwrap();
380        unsafe {
381            clear_azure_env_vars();
382            std::env::set_var("AZURE_CLIENT_ID", "test-id");
383            std::env::set_var("AZURE_TENANT_ID", "test-tenant");
384        }
385
386        let result = EnvAuth::from_env();
387        assert!(result.is_err());
388        let err = result.unwrap_err();
389        assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_CLIENT_SECRET"));
390
391        unsafe { clear_azure_env_vars() };
392    }
393
394    #[test]
395    fn test_env_auth_from_env_missing_tenant_id() {
396        let _lock = ENV_MUTEX.lock().unwrap();
397        unsafe {
398            clear_azure_env_vars();
399            std::env::set_var("AZURE_CLIENT_ID", "test-id");
400            std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
401        }
402
403        let result = EnvAuth::from_env();
404        assert!(result.is_err());
405        let err = result.unwrap_err();
406        assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_TENANT_ID"));
407
408        unsafe { clear_azure_env_vars() };
409    }
410
411    #[test]
412    fn test_env_auth_is_configured_all_set() {
413        let _lock = ENV_MUTEX.lock().unwrap();
414        unsafe { set_azure_env_vars() };
415
416        assert!(EnvAuth::is_configured());
417
418        unsafe { clear_azure_env_vars() };
419    }
420
421    #[test]
422    fn test_env_auth_is_configured_none_set() {
423        let _lock = ENV_MUTEX.lock().unwrap();
424        unsafe { clear_azure_env_vars() };
425
426        assert!(!EnvAuth::is_configured());
427    }
428
429    #[test]
430    fn test_env_auth_is_configured_partial() {
431        let _lock = ENV_MUTEX.lock().unwrap();
432        unsafe {
433            clear_azure_env_vars();
434            std::env::set_var("AZURE_CLIENT_ID", "test-id");
435            std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
436        }
437        // AZURE_TENANT_ID intentionally missing
438
439        assert!(!EnvAuth::is_configured());
440
441        unsafe { clear_azure_env_vars() };
442    }
443
444    #[test]
445    fn test_env_auth_method_name() {
446        let _lock = ENV_MUTEX.lock().unwrap();
447        unsafe { set_azure_env_vars() };
448
449        let auth = EnvAuth::from_env().unwrap();
450        assert_eq!(
451            auth.method_name(),
452            "Environment Variables (Service Principal)"
453        );
454
455        unsafe { clear_azure_env_vars() };
456    }
457
458    #[test]
459    fn test_az_cli_auth_method_name() {
460        let auth = AzCliAuth::new();
461        assert_eq!(auth.method_name(), "Azure CLI");
462    }
463
464    #[test]
465    fn test_az_cli_auth_search_scope() {
466        let auth = AzCliAuth::for_search();
467        assert_eq!(auth.resource_scope, "https://search.azure.com");
468    }
469
470    #[test]
471    fn test_az_cli_auth_foundry_scope() {
472        let auth = AzCliAuth::for_foundry();
473        assert_eq!(auth.resource_scope, "https://ai.azure.com");
474    }
475
476    #[test]
477    fn test_az_cli_auth_cognitive_services_scope() {
478        let auth = AzCliAuth::for_cognitive_services();
479        assert_eq!(auth.resource_scope, "https://cognitiveservices.azure.com");
480    }
481
482    #[test]
483    fn test_az_cli_auth_new_defaults_to_search() {
484        let auth = AzCliAuth::new();
485        assert_eq!(auth.resource_scope, "https://search.azure.com");
486    }
487
488    #[test]
489    fn test_env_auth_from_env_scope_foundry() {
490        let _lock = ENV_MUTEX.lock().unwrap();
491        unsafe { set_azure_env_vars() };
492
493        let result = EnvAuth::from_env_for_scope("https://ai.azure.com");
494        assert!(result.is_ok());
495        let auth = result.unwrap();
496        assert_eq!(auth.resource_scope, "https://ai.azure.com");
497
498        unsafe { clear_azure_env_vars() };
499    }
500
501    #[test]
502    fn test_env_auth_from_env_default_scope_is_search() {
503        let _lock = ENV_MUTEX.lock().unwrap();
504        unsafe { set_azure_env_vars() };
505
506        let auth = EnvAuth::from_env().unwrap();
507        assert_eq!(auth.resource_scope, "https://search.azure.com");
508
509        unsafe { clear_azure_env_vars() };
510    }
511
512    #[test]
513    fn test_auth_status_fields() {
514        let status = AuthStatus {
515            logged_in: true,
516            user: Some("testuser@example.com".to_string()),
517            subscription: Some("My Subscription".to_string()),
518            subscription_id: Some("00000000-0000-0000-0000-000000000000".to_string()),
519        };
520
521        assert!(status.logged_in);
522        assert_eq!(status.user.as_deref(), Some("testuser@example.com"));
523        assert_eq!(status.subscription.as_deref(), Some("My Subscription"));
524        assert_eq!(
525            status.subscription_id.as_deref(),
526            Some("00000000-0000-0000-0000-000000000000")
527        );
528    }
529}
rigg_client/auth.rs

rigg_client/
auth.rs
