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

rigg_client/
auth.rs
