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/// Static bearer token from the environment (`RIGG_ACCESS_TOKEN`).
311/// Useful for CI systems that mint tokens out-of-band, and for tests.
312pub struct StaticTokenAuth {
313    token: String,
314}
315
316impl AuthProvider for StaticTokenAuth {
317    fn get_token(&self) -> Result<String, AuthError> {
318        Ok(self.token.clone())
319    }
320    fn method_name(&self) -> &'static str {
321        "Static token (RIGG_ACCESS_TOKEN)"
322    }
323}
324
325/// Get the best available authentication provider for a specific resource scope
326fn get_auth_provider_for_scope(scope: &'static str) -> Result<Box<dyn AuthProvider>, AuthError> {
327    // A pre-minted token wins over everything.
328    if let Ok(token) = std::env::var("RIGG_ACCESS_TOKEN") {
329        if !token.is_empty() {
330            return Ok(Box::new(StaticTokenAuth { token }));
331        }
332    }
333    // First try environment variables
334    if EnvAuth::is_configured() {
335        return Ok(Box::new(EnvAuth::from_env_for_scope(scope)?));
336    }
337
338    // Fall back to Azure CLI
339    AzCliAuth::check_status()?;
340    Ok(Box::new(AzCliAuth {
341        resource_scope: scope,
342    }))
343}
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348    use std::sync::Mutex;
349
350    // Env var tests must run serially since they share process-wide state.
351    static ENV_MUTEX: Mutex<()> = Mutex::new(());
352
353    /// # Safety
354    /// Must be called while holding ENV_MUTEX to avoid data races.
355    unsafe fn clear_azure_env_vars() {
356        unsafe {
357            std::env::remove_var("AZURE_CLIENT_ID");
358            std::env::remove_var("AZURE_CLIENT_SECRET");
359            std::env::remove_var("AZURE_TENANT_ID");
360        }
361    }
362
363    /// # Safety
364    /// Must be called while holding ENV_MUTEX to avoid data races.
365    unsafe fn set_azure_env_vars() {
366        unsafe {
367            std::env::set_var("AZURE_CLIENT_ID", "test-client-id");
368            std::env::set_var("AZURE_CLIENT_SECRET", "test-client-secret");
369            std::env::set_var("AZURE_TENANT_ID", "test-tenant-id");
370        }
371    }
372
373    #[test]
374    fn test_env_auth_from_env_success() {
375        let _lock = ENV_MUTEX.lock().unwrap();
376        unsafe { set_azure_env_vars() };
377
378        let result = EnvAuth::from_env();
379        assert!(result.is_ok());
380        let auth = result.unwrap();
381        assert_eq!(auth.client_id, "test-client-id");
382        assert_eq!(auth.client_secret, "test-client-secret");
383        assert_eq!(auth.tenant_id, "test-tenant-id");
384
385        unsafe { clear_azure_env_vars() };
386    }
387
388    #[test]
389    fn test_env_auth_from_env_missing_client_id() {
390        let _lock = ENV_MUTEX.lock().unwrap();
391        unsafe {
392            clear_azure_env_vars();
393            std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
394            std::env::set_var("AZURE_TENANT_ID", "test-tenant");
395        }
396
397        let result = EnvAuth::from_env();
398        assert!(result.is_err());
399        let err = result.unwrap_err();
400        assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_CLIENT_ID"));
401
402        unsafe { clear_azure_env_vars() };
403    }
404
405    #[test]
406    fn test_env_auth_from_env_missing_client_secret() {
407        let _lock = ENV_MUTEX.lock().unwrap();
408        unsafe {
409            clear_azure_env_vars();
410            std::env::set_var("AZURE_CLIENT_ID", "test-id");
411            std::env::set_var("AZURE_TENANT_ID", "test-tenant");
412        }
413
414        let result = EnvAuth::from_env();
415        assert!(result.is_err());
416        let err = result.unwrap_err();
417        assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_CLIENT_SECRET"));
418
419        unsafe { clear_azure_env_vars() };
420    }
421
422    #[test]
423    fn test_env_auth_from_env_missing_tenant_id() {
424        let _lock = ENV_MUTEX.lock().unwrap();
425        unsafe {
426            clear_azure_env_vars();
427            std::env::set_var("AZURE_CLIENT_ID", "test-id");
428            std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
429        }
430
431        let result = EnvAuth::from_env();
432        assert!(result.is_err());
433        let err = result.unwrap_err();
434        assert!(matches!(err, AuthError::MissingEnvVar(ref v) if v == "AZURE_TENANT_ID"));
435
436        unsafe { clear_azure_env_vars() };
437    }
438
439    #[test]
440    fn test_env_auth_is_configured_all_set() {
441        let _lock = ENV_MUTEX.lock().unwrap();
442        unsafe { set_azure_env_vars() };
443
444        assert!(EnvAuth::is_configured());
445
446        unsafe { clear_azure_env_vars() };
447    }
448
449    #[test]
450    fn test_env_auth_is_configured_none_set() {
451        let _lock = ENV_MUTEX.lock().unwrap();
452        unsafe { clear_azure_env_vars() };
453
454        assert!(!EnvAuth::is_configured());
455    }
456
457    #[test]
458    fn test_env_auth_is_configured_partial() {
459        let _lock = ENV_MUTEX.lock().unwrap();
460        unsafe {
461            clear_azure_env_vars();
462            std::env::set_var("AZURE_CLIENT_ID", "test-id");
463            std::env::set_var("AZURE_CLIENT_SECRET", "test-secret");
464        }
465        // AZURE_TENANT_ID intentionally missing
466
467        assert!(!EnvAuth::is_configured());
468
469        unsafe { clear_azure_env_vars() };
470    }
471
472    #[test]
473    fn test_env_auth_method_name() {
474        let _lock = ENV_MUTEX.lock().unwrap();
475        unsafe { set_azure_env_vars() };
476
477        let auth = EnvAuth::from_env().unwrap();
478        assert_eq!(
479            auth.method_name(),
480            "Environment Variables (Service Principal)"
481        );
482
483        unsafe { clear_azure_env_vars() };
484    }
485
486    #[test]
487    fn test_az_cli_auth_method_name() {
488        let auth = AzCliAuth::new();
489        assert_eq!(auth.method_name(), "Azure CLI");
490    }
491
492    #[test]
493    fn test_az_cli_auth_search_scope() {
494        let auth = AzCliAuth::for_search();
495        assert_eq!(auth.resource_scope, "https://search.azure.com");
496    }
497
498    #[test]
499    fn test_az_cli_auth_foundry_scope() {
500        let auth = AzCliAuth::for_foundry();
501        assert_eq!(auth.resource_scope, "https://ai.azure.com");
502    }
503
504    #[test]
505    fn test_az_cli_auth_cognitive_services_scope() {
506        let auth = AzCliAuth::for_cognitive_services();
507        assert_eq!(auth.resource_scope, "https://cognitiveservices.azure.com");
508    }
509
510    #[test]
511    fn test_az_cli_auth_new_defaults_to_search() {
512        let auth = AzCliAuth::new();
513        assert_eq!(auth.resource_scope, "https://search.azure.com");
514    }
515
516    #[test]
517    fn test_env_auth_from_env_scope_foundry() {
518        let _lock = ENV_MUTEX.lock().unwrap();
519        unsafe { set_azure_env_vars() };
520
521        let result = EnvAuth::from_env_for_scope("https://ai.azure.com");
522        assert!(result.is_ok());
523        let auth = result.unwrap();
524        assert_eq!(auth.resource_scope, "https://ai.azure.com");
525
526        unsafe { clear_azure_env_vars() };
527    }
528
529    #[test]
530    fn test_env_auth_from_env_default_scope_is_search() {
531        let _lock = ENV_MUTEX.lock().unwrap();
532        unsafe { set_azure_env_vars() };
533
534        let auth = EnvAuth::from_env().unwrap();
535        assert_eq!(auth.resource_scope, "https://search.azure.com");
536
537        unsafe { clear_azure_env_vars() };
538    }
539
540    #[test]
541    fn test_auth_status_fields() {
542        let status = AuthStatus {
543            logged_in: true,
544            user: Some("testuser@example.com".to_string()),
545            subscription: Some("My Subscription".to_string()),
546            subscription_id: Some("00000000-0000-0000-0000-000000000000".to_string()),
547        };
548
549        assert!(status.logged_in);
550        assert_eq!(status.user.as_deref(), Some("testuser@example.com"));
551        assert_eq!(status.subscription.as_deref(), Some("My Subscription"));
552        assert_eq!(
553            status.subscription_id.as_deref(),
554            Some("00000000-0000-0000-0000-000000000000")
555        );
556    }
557
558    #[test]
559    fn test_for_cosmos_uses_cosmos_scope() {
560        let auth = AzCliAuth::for_cosmos();
561        assert_eq!(auth.resource_scope, "https://cosmos.azure.com");
562    }
563}
rigg_client/auth.rs

rigg_client/
auth.rs
