ati/core/
auth_generator.rs

1//! Dynamic credential generator — produces short-lived auth tokens at call time.
2//!
3//! Generators run where secrets live: on the proxy server in proxy mode,
4//! on the local machine in local mode. Signing keys never enter the sandbox.
5//!
6//! Two generator types:
7//! - `command`: runs an external command, captures stdout
8//! - `script`: writes an inline script to a temp file, runs via interpreter
9//!
10//! Results are cached per (provider, agent_sub) with configurable TTL.
11
12use std::collections::HashMap;
13use std::sync::Mutex;
14use std::time::{Duration, Instant};
15use thiserror::Error;
16
17use crate::core::keyring::Keyring;
18use crate::core::manifest::{AuthGenType, AuthGenerator, AuthOutputFormat, Provider};
19
20// ---------------------------------------------------------------------------
21// Errors
22// ---------------------------------------------------------------------------
23
24#[derive(Error, Debug)]
25pub enum AuthGenError {
26    #[error("Auth generator config error: {0}")]
27    Config(String),
28    #[error("Failed to spawn generator process: {0}")]
29    Spawn(String),
30    #[error("Generator timed out after {0}s")]
31    Timeout(u64),
32    #[error("Generator exited with code {code}: {stderr}")]
33    NonZeroExit { code: i32, stderr: String },
34    #[error("Failed to parse generator output: {0}")]
35    OutputParse(String),
36    #[error("Keyring key '{0}' not found (required by auth_generator)")]
37    KeyringMissing(String),
38    #[error("IO error: {0}")]
39    Io(#[from] std::io::Error),
40}
41
42// ---------------------------------------------------------------------------
43// Context for variable expansion
44// ---------------------------------------------------------------------------
45
46/// Context for expanding `${VAR}` placeholders in generator args/env.
47pub struct GenContext {
48    pub jwt_sub: String,
49    pub jwt_scope: String,
50    pub tool_name: String,
51    pub timestamp: u64,
52}
53
54impl Default for GenContext {
55    fn default() -> Self {
56        GenContext {
57            jwt_sub: "dev".into(),
58            jwt_scope: "*".into(),
59            tool_name: String::new(),
60            timestamp: std::time::SystemTime::now()
61                .duration_since(std::time::UNIX_EPOCH)
62                .unwrap_or_default()
63                .as_secs(),
64        }
65    }
66}
67
68// ---------------------------------------------------------------------------
69// Generated credential
70// ---------------------------------------------------------------------------
71
72/// Result of running a generator — primary token + optional extra injections.
73#[derive(Debug, Clone)]
74pub struct GeneratedCredential {
75    /// Primary token value (used for bearer/header/query auth).
76    pub value: String,
77    /// Extra headers from JSON inject targets with type="header".
78    pub extra_headers: HashMap<String, String>,
79    /// Extra env vars from JSON inject targets with type="env".
80    pub extra_env: HashMap<String, String>,
81}
82
83// ---------------------------------------------------------------------------
84// Cache
85// ---------------------------------------------------------------------------
86
87struct CachedCredential {
88    cred: GeneratedCredential,
89    expires_at: Instant,
90}
91
92/// TTL-based credential cache, keyed by (provider_name, agent_sub).
93pub struct AuthCache {
94    entries: Mutex<HashMap<(String, String), CachedCredential>>,
95}
96
97impl Default for AuthCache {
98    fn default() -> Self {
99        AuthCache {
100            entries: Mutex::new(HashMap::new()),
101        }
102    }
103}
104
105impl AuthCache {
106    pub fn new() -> Self {
107        Self::default()
108    }
109
110    pub fn get(&self, provider: &str, sub: &str) -> Option<GeneratedCredential> {
111        let cache = self.entries.lock().unwrap();
112        let key = (provider.to_string(), sub.to_string());
113        match cache.get(&key) {
114            Some(entry) if Instant::now() < entry.expires_at => Some(entry.cred.clone()),
115            _ => None,
116        }
117    }
118
119    pub fn insert(&self, provider: &str, sub: &str, cred: GeneratedCredential, ttl_secs: u64) {
120        if ttl_secs == 0 {
121            return; // No caching
122        }
123        let mut cache = self.entries.lock().unwrap();
124        let key = (provider.to_string(), sub.to_string());
125        cache.insert(
126            key,
127            CachedCredential {
128                cred,
129                expires_at: Instant::now() + Duration::from_secs(ttl_secs),
130            },
131        );
132    }
133}
134
135// ---------------------------------------------------------------------------
136// Main generate function
137// ---------------------------------------------------------------------------
138
139/// Generate a credential by running the provider's auth_generator.
140///
141/// 1. Check cache → return if hit
142/// 2. Expand variables in args and env
143/// 3. Spawn subprocess (command or script)
144/// 4. Parse output (text or JSON)
145/// 5. Cache and return
146pub async fn generate(
147    provider: &Provider,
148    gen: &AuthGenerator,
149    ctx: &GenContext,
150    keyring: &Keyring,
151    cache: &AuthCache,
152) -> Result<GeneratedCredential, AuthGenError> {
153    // 1. Check cache
154    if gen.cache_ttl_secs > 0 {
155        if let Some(cached) = cache.get(&provider.name, &ctx.jwt_sub) {
156            return Ok(cached);
157        }
158    }
159
160    // 2. Expand variables in args and env
161    let expanded_args: Vec<String> = gen
162        .args
163        .iter()
164        .map(|a| expand_variables(a, ctx, keyring))
165        .collect::<Result<Vec<_>, _>>()?;
166
167    let mut expanded_env: HashMap<String, String> = HashMap::new();
168    for (k, v) in &gen.env {
169        expanded_env.insert(k.clone(), expand_variables(v, ctx, keyring)?);
170    }
171
172    // 3. Build curated env (don't leak host secrets)
173    let mut final_env: HashMap<String, String> = HashMap::new();
174    for var in &["PATH", "HOME", "TMPDIR"] {
175        if let Ok(val) = std::env::var(var) {
176            final_env.insert(var.to_string(), val);
177        }
178    }
179    final_env.extend(expanded_env);
180
181    // 4. Spawn subprocess
182    let output =
183        match gen.gen_type {
184            AuthGenType::Command => {
185                let command = gen.command.as_deref().ok_or_else(|| {
186                    AuthGenError::Config("command required for type=command".into())
187                })?;
188
189                let child = tokio::process::Command::new(command)
190                    .args(&expanded_args)
191                    .env_clear()
192                    .envs(&final_env)
193                    .stdout(std::process::Stdio::piped())
194                    .stderr(std::process::Stdio::piped())
195                    .kill_on_drop(true)
196                    .spawn()
197                    .map_err(|e| AuthGenError::Spawn(format!("{command}: {e}")))?;
198
199                let timeout = Duration::from_secs(gen.timeout_secs);
200                tokio::time::timeout(timeout, child.wait_with_output())
201                    .await
202                    .map_err(|_| AuthGenError::Timeout(gen.timeout_secs))?
203                    .map_err(AuthGenError::Io)?
204            }
205            AuthGenType::Script => {
206                let interpreter = gen.interpreter.as_deref().ok_or_else(|| {
207                    AuthGenError::Config("interpreter required for type=script".into())
208                })?;
209                let script = gen.script.as_deref().ok_or_else(|| {
210                    AuthGenError::Config("script required for type=script".into())
211                })?;
212
213                // Write script to a temp file
214                let suffix: u32 = rand::random();
215                let tmp_path = std::env::temp_dir().join(format!("ati_gen_{suffix}.tmp"));
216                std::fs::write(&tmp_path, script).map_err(AuthGenError::Io)?;
217
218                let child = tokio::process::Command::new(interpreter)
219                    .arg(&tmp_path)
220                    .env_clear()
221                    .envs(&final_env)
222                    .stdout(std::process::Stdio::piped())
223                    .stderr(std::process::Stdio::piped())
224                    .kill_on_drop(true)
225                    .spawn()
226                    .map_err(|e| AuthGenError::Spawn(format!("{interpreter}: {e}")))?;
227
228                let timeout = Duration::from_secs(gen.timeout_secs);
229                let result = tokio::time::timeout(timeout, child.wait_with_output())
230                    .await
231                    .map_err(|_| AuthGenError::Timeout(gen.timeout_secs))?
232                    .map_err(AuthGenError::Io)?;
233
234                // Clean up temp file
235                let _ = std::fs::remove_file(&tmp_path);
236                result
237            }
238        };
239
240    if !output.status.success() {
241        let code = output.status.code().unwrap_or(-1);
242        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
243        return Err(AuthGenError::NonZeroExit { code, stderr });
244    }
245
246    let stdout = String::from_utf8_lossy(&output.stdout);
247
248    // 5. Parse output
249    let cred = match gen.output_format {
250        AuthOutputFormat::Text => GeneratedCredential {
251            value: stdout.trim().to_string(),
252            extra_headers: HashMap::new(),
253            extra_env: HashMap::new(),
254        },
255        AuthOutputFormat::Json => {
256            let json: serde_json::Value = serde_json::from_str(stdout.trim())
257                .map_err(|e| AuthGenError::OutputParse(format!("invalid JSON: {e}")))?;
258
259            let mut extra_headers = HashMap::new();
260            let mut extra_env = HashMap::new();
261            let mut primary_value = stdout.trim().to_string();
262
263            // If no inject map, use the whole output as the primary value
264            if gen.inject.is_empty() {
265                // Try to extract a "token" or "access_token" field as primary
266                if let Some(tok) = json.get("token").or(json.get("access_token")) {
267                    if let Some(s) = tok.as_str() {
268                        primary_value = s.to_string();
269                    }
270                }
271            } else {
272                // Extract fields per inject map
273                let mut found_primary = false;
274                for (json_path, target) in &gen.inject {
275                    let extracted = extract_json_path(&json, json_path).ok_or_else(|| {
276                        AuthGenError::OutputParse(format!(
277                            "JSON path '{}' not found in output",
278                            json_path
279                        ))
280                    })?;
281
282                    match target.inject_type.as_str() {
283                        "header" => {
284                            extra_headers.insert(target.name.clone(), extracted);
285                        }
286                        "env" => {
287                            extra_env.insert(target.name.clone(), extracted);
288                        }
289                        "query" => {
290                            // For query injection, use as primary value
291                            if !found_primary {
292                                primary_value = extracted;
293                                found_primary = true;
294                            }
295                        }
296                        _ => {
297                            // Default: treat as primary value
298                            if !found_primary {
299                                primary_value = extracted;
300                                found_primary = true;
301                            }
302                        }
303                    }
304                }
305            }
306
307            GeneratedCredential {
308                value: primary_value,
309                extra_headers,
310                extra_env,
311            }
312        }
313    };
314
315    // 6. Cache
316    cache.insert(
317        &provider.name,
318        &ctx.jwt_sub,
319        cred.clone(),
320        gen.cache_ttl_secs,
321    );
322
323    Ok(cred)
324}
325
326// ---------------------------------------------------------------------------
327// Variable expansion
328// ---------------------------------------------------------------------------
329
330/// Expand `${VAR}` placeholders in a string.
331///
332/// Recognized variables:
333/// - `${JWT_SUB}`, `${JWT_SCOPE}`, `${TOOL_NAME}`, `${TIMESTAMP}` — from GenContext
334/// - `${anything_else}` — looked up in the keyring
335fn expand_variables(
336    input: &str,
337    ctx: &GenContext,
338    keyring: &Keyring,
339) -> Result<String, AuthGenError> {
340    let mut result = input.to_string();
341    // Process all ${...} patterns
342    while let Some(start) = result.find("${") {
343        let rest = &result[start + 2..];
344        let end = match rest.find('}') {
345            Some(e) => e,
346            None => break,
347        };
348        let var_name = &rest[..end];
349
350        let replacement = match var_name {
351            "JWT_SUB" => ctx.jwt_sub.clone(),
352            "JWT_SCOPE" => ctx.jwt_scope.clone(),
353            "TOOL_NAME" => ctx.tool_name.clone(),
354            "TIMESTAMP" => ctx.timestamp.to_string(),
355            _ => {
356                // Keyring lookup
357                match keyring.get(var_name) {
358                    Some(val) => val.to_string(),
359                    None => return Err(AuthGenError::KeyringMissing(var_name.to_string())),
360                }
361            }
362        };
363
364        result = format!("{}{}{}", &result[..start], replacement, &rest[end + 1..]);
365    }
366    Ok(result)
367}
368
369// ---------------------------------------------------------------------------
370// JSON path extraction (dot-notation)
371// ---------------------------------------------------------------------------
372
373/// Extract a value from a JSON object using dot-notation path.
374///
375/// Example: `extract_json_path(json, "Credentials.AccessKeyId")`
376/// navigates `json["Credentials"]["AccessKeyId"]`.
377fn extract_json_path(value: &serde_json::Value, path: &str) -> Option<String> {
378    let mut current = value;
379    for segment in path.split('.') {
380        current = current.get(segment)?;
381    }
382    match current {
383        serde_json::Value::String(s) => Some(s.clone()),
384        serde_json::Value::Number(n) => Some(n.to_string()),
385        serde_json::Value::Bool(b) => Some(b.to_string()),
386        other => Some(other.to_string()),
387    }
388}
389
390// ---------------------------------------------------------------------------
391// Tests
392// ---------------------------------------------------------------------------
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397
398    #[test]
399    fn test_expand_variables_context() {
400        let ctx = GenContext {
401            jwt_sub: "agent-7".into(),
402            jwt_scope: "tool:brain:*".into(),
403            tool_name: "brain:query".into(),
404            timestamp: 1773096459,
405        };
406        let keyring = Keyring::empty();
407
408        assert_eq!(
409            expand_variables("${JWT_SUB}", &ctx, &keyring).unwrap(),
410            "agent-7"
411        );
412        assert_eq!(
413            expand_variables("${TOOL_NAME}", &ctx, &keyring).unwrap(),
414            "brain:query"
415        );
416        assert_eq!(
417            expand_variables("${TIMESTAMP}", &ctx, &keyring).unwrap(),
418            "1773096459"
419        );
420        assert_eq!(
421            expand_variables("sub=${JWT_SUB}&tool=${TOOL_NAME}", &ctx, &keyring).unwrap(),
422            "sub=agent-7&tool=brain:query"
423        );
424    }
425
426    #[test]
427    fn test_expand_variables_keyring() {
428        let dir = tempfile::TempDir::new().unwrap();
429        let path = dir.path().join("creds");
430        std::fs::write(&path, r#"{"my_secret":"s3cr3t"}"#).unwrap();
431        let keyring = Keyring::load_credentials(&path).unwrap();
432
433        let ctx = GenContext::default();
434        assert_eq!(
435            expand_variables("${my_secret}", &ctx, &keyring).unwrap(),
436            "s3cr3t"
437        );
438    }
439
440    #[test]
441    fn test_expand_variables_missing_key() {
442        let keyring = Keyring::empty();
443        let ctx = GenContext::default();
444        let err = expand_variables("${nonexistent}", &ctx, &keyring).unwrap_err();
445        assert!(matches!(err, AuthGenError::KeyringMissing(_)));
446    }
447
448    #[test]
449    fn test_expand_variables_no_placeholder() {
450        let keyring = Keyring::empty();
451        let ctx = GenContext::default();
452        assert_eq!(
453            expand_variables("plain text", &ctx, &keyring).unwrap(),
454            "plain text"
455        );
456    }
457
458    #[test]
459    fn test_extract_json_path_simple() {
460        let json: serde_json::Value = serde_json::json!({"token": "abc123", "expires_in": 3600});
461        assert_eq!(extract_json_path(&json, "token"), Some("abc123".into()));
462        assert_eq!(extract_json_path(&json, "expires_in"), Some("3600".into()));
463    }
464
465    #[test]
466    fn test_extract_json_path_nested() {
467        let json: serde_json::Value = serde_json::json!({
468            "Credentials": {
469                "AccessKeyId": "AKIA...",
470                "SecretAccessKey": "wJalrX...",
471                "SessionToken": "FwoGZ..."
472            }
473        });
474        assert_eq!(
475            extract_json_path(&json, "Credentials.AccessKeyId"),
476            Some("AKIA...".into())
477        );
478        assert_eq!(
479            extract_json_path(&json, "Credentials.SessionToken"),
480            Some("FwoGZ...".into())
481        );
482    }
483
484    #[test]
485    fn test_extract_json_path_missing() {
486        let json: serde_json::Value = serde_json::json!({"a": "b"});
487        assert_eq!(extract_json_path(&json, "nonexistent"), None);
488        assert_eq!(extract_json_path(&json, "a.b.c"), None);
489    }
490
491    #[test]
492    fn test_auth_cache_basic() {
493        let cache = AuthCache::new();
494        assert!(cache.get("provider", "sub").is_none());
495
496        let cred = GeneratedCredential {
497            value: "token123".into(),
498            extra_headers: HashMap::new(),
499            extra_env: HashMap::new(),
500        };
501        cache.insert("provider", "sub", cred.clone(), 300);
502
503        let cached = cache.get("provider", "sub").unwrap();
504        assert_eq!(cached.value, "token123");
505    }
506
507    #[test]
508    fn test_auth_cache_zero_ttl_no_cache() {
509        let cache = AuthCache::new();
510        let cred = GeneratedCredential {
511            value: "token".into(),
512            extra_headers: HashMap::new(),
513            extra_env: HashMap::new(),
514        };
515        cache.insert("provider", "sub", cred, 0);
516        assert!(cache.get("provider", "sub").is_none());
517    }
518
519    #[test]
520    fn test_auth_cache_different_keys() {
521        let cache = AuthCache::new();
522        let cred1 = GeneratedCredential {
523            value: "token-a".into(),
524            extra_headers: HashMap::new(),
525            extra_env: HashMap::new(),
526        };
527        let cred2 = GeneratedCredential {
528            value: "token-b".into(),
529            extra_headers: HashMap::new(),
530            extra_env: HashMap::new(),
531        };
532        cache.insert("provider", "agent-1", cred1, 300);
533        cache.insert("provider", "agent-2", cred2, 300);
534
535        assert_eq!(cache.get("provider", "agent-1").unwrap().value, "token-a");
536        assert_eq!(cache.get("provider", "agent-2").unwrap().value, "token-b");
537    }
538
539    #[tokio::test]
540    async fn test_generate_command_text() {
541        let provider = Provider {
542            name: "test".into(),
543            description: "test provider".into(),
544            base_url: String::new(),
545            auth_type: crate::core::manifest::AuthType::Bearer,
546            auth_key_name: None,
547            auth_header_name: None,
548            auth_query_name: None,
549            auth_value_prefix: None,
550            extra_headers: HashMap::new(),
551            oauth2_token_url: None,
552            auth_secret_name: None,
553            oauth2_basic_auth: false,
554            internal: false,
555            handler: "http".into(),
556            mcp_transport: None,
557            mcp_command: None,
558            mcp_args: vec![],
559            mcp_url: None,
560            mcp_env: HashMap::new(),
561            cli_command: None,
562            cli_default_args: vec![],
563            cli_env: HashMap::new(),
564            cli_timeout_secs: None,
565            cli_output_args: Vec::new(),
566            cli_output_positional: HashMap::new(),
567            upload_destinations: HashMap::new(),
568            upload_default_destination: None,
569            openapi_spec: None,
570            openapi_include_tags: vec![],
571            openapi_exclude_tags: vec![],
572            openapi_include_operations: vec![],
573            openapi_exclude_operations: vec![],
574            openapi_max_operations: None,
575            openapi_overrides: HashMap::new(),
576            auth_generator: None,
577            category: None,
578            skills: vec![],
579        };
580
581        let gen = AuthGenerator {
582            gen_type: AuthGenType::Command,
583            command: Some("echo".into()),
584            args: vec!["hello-token".into()],
585            interpreter: None,
586            script: None,
587            cache_ttl_secs: 0,
588            output_format: AuthOutputFormat::Text,
589            env: HashMap::new(),
590            inject: HashMap::new(),
591            timeout_secs: 5,
592        };
593
594        let ctx = GenContext::default();
595        let keyring = Keyring::empty();
596        let cache = AuthCache::new();
597
598        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
599            .await
600            .unwrap();
601        assert_eq!(cred.value, "hello-token");
602        assert!(cred.extra_headers.is_empty());
603    }
604
605    #[tokio::test]
606    async fn test_generate_command_json() {
607        let provider = Provider {
608            name: "test".into(),
609            description: "test".into(),
610            base_url: String::new(),
611            auth_type: crate::core::manifest::AuthType::Bearer,
612            auth_key_name: None,
613            auth_header_name: None,
614            auth_query_name: None,
615            auth_value_prefix: None,
616            extra_headers: HashMap::new(),
617            oauth2_token_url: None,
618            auth_secret_name: None,
619            oauth2_basic_auth: false,
620            internal: false,
621            handler: "http".into(),
622            mcp_transport: None,
623            mcp_command: None,
624            mcp_args: vec![],
625            mcp_url: None,
626            mcp_env: HashMap::new(),
627            cli_command: None,
628            cli_default_args: vec![],
629            cli_env: HashMap::new(),
630            cli_timeout_secs: None,
631            cli_output_args: Vec::new(),
632            cli_output_positional: HashMap::new(),
633            upload_destinations: HashMap::new(),
634            upload_default_destination: None,
635            openapi_spec: None,
636            openapi_include_tags: vec![],
637            openapi_exclude_tags: vec![],
638            openapi_include_operations: vec![],
639            openapi_exclude_operations: vec![],
640            openapi_max_operations: None,
641            openapi_overrides: HashMap::new(),
642            auth_generator: None,
643            category: None,
644            skills: vec![],
645        };
646
647        let mut inject = HashMap::new();
648        inject.insert(
649            "Credentials.AccessKeyId".into(),
650            crate::core::manifest::InjectTarget {
651                inject_type: "header".into(),
652                name: "X-Access-Key".into(),
653            },
654        );
655        inject.insert(
656            "Credentials.Secret".into(),
657            crate::core::manifest::InjectTarget {
658                inject_type: "env".into(),
659                name: "AWS_SECRET".into(),
660            },
661        );
662
663        let gen = AuthGenerator {
664            gen_type: AuthGenType::Command,
665            command: Some("echo".into()),
666            args: vec![
667                r#"{"Credentials":{"AccessKeyId":"AKIA123","Secret":"wJalr","SessionToken":"FwoG"}}"#.into(),
668            ],
669            interpreter: None,
670            script: None,
671            cache_ttl_secs: 0,
672            output_format: AuthOutputFormat::Json,
673            env: HashMap::new(),
674            inject,
675            timeout_secs: 5,
676        };
677
678        let ctx = GenContext::default();
679        let keyring = Keyring::empty();
680        let cache = AuthCache::new();
681
682        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
683            .await
684            .unwrap();
685        assert_eq!(cred.extra_headers.get("X-Access-Key").unwrap(), "AKIA123");
686        assert_eq!(cred.extra_env.get("AWS_SECRET").unwrap(), "wJalr");
687    }
688
689    #[tokio::test]
690    async fn test_generate_script() {
691        let provider = Provider {
692            name: "test".into(),
693            description: "test".into(),
694            base_url: String::new(),
695            auth_type: crate::core::manifest::AuthType::Bearer,
696            auth_key_name: None,
697            auth_header_name: None,
698            auth_query_name: None,
699            auth_value_prefix: None,
700            extra_headers: HashMap::new(),
701            oauth2_token_url: None,
702            auth_secret_name: None,
703            oauth2_basic_auth: false,
704            internal: false,
705            handler: "http".into(),
706            mcp_transport: None,
707            mcp_command: None,
708            mcp_args: vec![],
709            mcp_url: None,
710            mcp_env: HashMap::new(),
711            cli_command: None,
712            cli_default_args: vec![],
713            cli_env: HashMap::new(),
714            cli_timeout_secs: None,
715            cli_output_args: Vec::new(),
716            cli_output_positional: HashMap::new(),
717            upload_destinations: HashMap::new(),
718            upload_default_destination: None,
719            openapi_spec: None,
720            openapi_include_tags: vec![],
721            openapi_exclude_tags: vec![],
722            openapi_include_operations: vec![],
723            openapi_exclude_operations: vec![],
724            openapi_max_operations: None,
725            openapi_overrides: HashMap::new(),
726            auth_generator: None,
727            category: None,
728            skills: vec![],
729        };
730
731        let gen = AuthGenerator {
732            gen_type: AuthGenType::Script,
733            command: None,
734            args: vec![],
735            interpreter: Some("bash".into()),
736            script: Some("echo script-token-42".into()),
737            cache_ttl_secs: 0,
738            output_format: AuthOutputFormat::Text,
739            env: HashMap::new(),
740            inject: HashMap::new(),
741            timeout_secs: 5,
742        };
743
744        let ctx = GenContext::default();
745        let keyring = Keyring::empty();
746        let cache = AuthCache::new();
747
748        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
749            .await
750            .unwrap();
751        assert_eq!(cred.value, "script-token-42");
752    }
753
754    #[tokio::test]
755    async fn test_generate_caches_result() {
756        let provider = Provider {
757            name: "cached_provider".into(),
758            description: "test".into(),
759            base_url: String::new(),
760            auth_type: crate::core::manifest::AuthType::Bearer,
761            auth_key_name: None,
762            auth_header_name: None,
763            auth_query_name: None,
764            auth_value_prefix: None,
765            extra_headers: HashMap::new(),
766            oauth2_token_url: None,
767            auth_secret_name: None,
768            oauth2_basic_auth: false,
769            internal: false,
770            handler: "http".into(),
771            mcp_transport: None,
772            mcp_command: None,
773            mcp_args: vec![],
774            mcp_url: None,
775            mcp_env: HashMap::new(),
776            cli_command: None,
777            cli_default_args: vec![],
778            cli_env: HashMap::new(),
779            cli_timeout_secs: None,
780            cli_output_args: Vec::new(),
781            cli_output_positional: HashMap::new(),
782            upload_destinations: HashMap::new(),
783            upload_default_destination: None,
784            openapi_spec: None,
785            openapi_include_tags: vec![],
786            openapi_exclude_tags: vec![],
787            openapi_include_operations: vec![],
788            openapi_exclude_operations: vec![],
789            openapi_max_operations: None,
790            openapi_overrides: HashMap::new(),
791            auth_generator: None,
792            category: None,
793            skills: vec![],
794        };
795
796        let gen = AuthGenerator {
797            gen_type: AuthGenType::Command,
798            command: Some("date".into()),
799            args: vec!["+%s%N".into()],
800            interpreter: None,
801            script: None,
802            cache_ttl_secs: 300,
803            output_format: AuthOutputFormat::Text,
804            env: HashMap::new(),
805            inject: HashMap::new(),
806            timeout_secs: 5,
807        };
808
809        let ctx = GenContext {
810            jwt_sub: "test-agent".into(),
811            ..GenContext::default()
812        };
813        let keyring = Keyring::empty();
814        let cache = AuthCache::new();
815
816        let cred1 = generate(&provider, &gen, &ctx, &keyring, &cache)
817            .await
818            .unwrap();
819        let cred2 = generate(&provider, &gen, &ctx, &keyring, &cache)
820            .await
821            .unwrap();
822        // Second call should return cached value (same value)
823        assert_eq!(cred1.value, cred2.value);
824    }
825
826    #[tokio::test]
827    async fn test_generate_with_variable_expansion() {
828        let provider = Provider {
829            name: "test".into(),
830            description: "test".into(),
831            base_url: String::new(),
832            auth_type: crate::core::manifest::AuthType::Bearer,
833            auth_key_name: None,
834            auth_header_name: None,
835            auth_query_name: None,
836            auth_value_prefix: None,
837            extra_headers: HashMap::new(),
838            oauth2_token_url: None,
839            auth_secret_name: None,
840            oauth2_basic_auth: false,
841            internal: false,
842            handler: "http".into(),
843            mcp_transport: None,
844            mcp_command: None,
845            mcp_args: vec![],
846            mcp_url: None,
847            mcp_env: HashMap::new(),
848            cli_command: None,
849            cli_default_args: vec![],
850            cli_env: HashMap::new(),
851            cli_timeout_secs: None,
852            cli_output_args: Vec::new(),
853            cli_output_positional: HashMap::new(),
854            upload_destinations: HashMap::new(),
855            upload_default_destination: None,
856            openapi_spec: None,
857            openapi_include_tags: vec![],
858            openapi_exclude_tags: vec![],
859            openapi_include_operations: vec![],
860            openapi_exclude_operations: vec![],
861            openapi_max_operations: None,
862            openapi_overrides: HashMap::new(),
863            auth_generator: None,
864            category: None,
865            skills: vec![],
866        };
867
868        let gen = AuthGenerator {
869            gen_type: AuthGenType::Command,
870            command: Some("echo".into()),
871            args: vec!["${JWT_SUB}".into()],
872            interpreter: None,
873            script: None,
874            cache_ttl_secs: 0,
875            output_format: AuthOutputFormat::Text,
876            env: HashMap::new(),
877            inject: HashMap::new(),
878            timeout_secs: 5,
879        };
880
881        let ctx = GenContext {
882            jwt_sub: "agent-42".into(),
883            jwt_scope: "*".into(),
884            tool_name: "brain:query".into(),
885            timestamp: 1234567890,
886        };
887        let keyring = Keyring::empty();
888        let cache = AuthCache::new();
889
890        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
891            .await
892            .unwrap();
893        assert_eq!(cred.value, "agent-42");
894    }
895
896    #[tokio::test]
897    async fn test_generate_timeout() {
898        let provider = Provider {
899            name: "test".into(),
900            description: "test".into(),
901            base_url: String::new(),
902            auth_type: crate::core::manifest::AuthType::Bearer,
903            auth_key_name: None,
904            auth_header_name: None,
905            auth_query_name: None,
906            auth_value_prefix: None,
907            extra_headers: HashMap::new(),
908            oauth2_token_url: None,
909            auth_secret_name: None,
910            oauth2_basic_auth: false,
911            internal: false,
912            handler: "http".into(),
913            mcp_transport: None,
914            mcp_command: None,
915            mcp_args: vec![],
916            mcp_url: None,
917            mcp_env: HashMap::new(),
918            cli_command: None,
919            cli_default_args: vec![],
920            cli_env: HashMap::new(),
921            cli_timeout_secs: None,
922            cli_output_args: Vec::new(),
923            cli_output_positional: HashMap::new(),
924            upload_destinations: HashMap::new(),
925            upload_default_destination: None,
926            openapi_spec: None,
927            openapi_include_tags: vec![],
928            openapi_exclude_tags: vec![],
929            openapi_include_operations: vec![],
930            openapi_exclude_operations: vec![],
931            openapi_max_operations: None,
932            openapi_overrides: HashMap::new(),
933            auth_generator: None,
934            category: None,
935            skills: vec![],
936        };
937
938        let gen = AuthGenerator {
939            gen_type: AuthGenType::Command,
940            command: Some("sleep".into()),
941            args: vec!["10".into()],
942            interpreter: None,
943            script: None,
944            cache_ttl_secs: 0,
945            output_format: AuthOutputFormat::Text,
946            env: HashMap::new(),
947            inject: HashMap::new(),
948            timeout_secs: 1,
949        };
950
951        let ctx = GenContext::default();
952        let keyring = Keyring::empty();
953        let cache = AuthCache::new();
954
955        let err = generate(&provider, &gen, &ctx, &keyring, &cache)
956            .await
957            .unwrap_err();
958        assert!(matches!(err, AuthGenError::Timeout(1)));
959    }
960}
ati/core/auth_generator.rs

ati/core/
auth_generator.rs
