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            openapi_spec: None,
566            openapi_include_tags: vec![],
567            openapi_exclude_tags: vec![],
568            openapi_include_operations: vec![],
569            openapi_exclude_operations: vec![],
570            openapi_max_operations: None,
571            openapi_overrides: HashMap::new(),
572            auth_generator: None,
573            category: None,
574            skills: vec![],
575        };
576
577        let gen = AuthGenerator {
578            gen_type: AuthGenType::Command,
579            command: Some("echo".into()),
580            args: vec!["hello-token".into()],
581            interpreter: None,
582            script: None,
583            cache_ttl_secs: 0,
584            output_format: AuthOutputFormat::Text,
585            env: HashMap::new(),
586            inject: HashMap::new(),
587            timeout_secs: 5,
588        };
589
590        let ctx = GenContext::default();
591        let keyring = Keyring::empty();
592        let cache = AuthCache::new();
593
594        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
595            .await
596            .unwrap();
597        assert_eq!(cred.value, "hello-token");
598        assert!(cred.extra_headers.is_empty());
599    }
600
601    #[tokio::test]
602    async fn test_generate_command_json() {
603        let provider = Provider {
604            name: "test".into(),
605            description: "test".into(),
606            base_url: String::new(),
607            auth_type: crate::core::manifest::AuthType::Bearer,
608            auth_key_name: None,
609            auth_header_name: None,
610            auth_query_name: None,
611            auth_value_prefix: None,
612            extra_headers: HashMap::new(),
613            oauth2_token_url: None,
614            auth_secret_name: None,
615            oauth2_basic_auth: false,
616            internal: false,
617            handler: "http".into(),
618            mcp_transport: None,
619            mcp_command: None,
620            mcp_args: vec![],
621            mcp_url: None,
622            mcp_env: HashMap::new(),
623            cli_command: None,
624            cli_default_args: vec![],
625            cli_env: HashMap::new(),
626            cli_timeout_secs: None,
627            openapi_spec: None,
628            openapi_include_tags: vec![],
629            openapi_exclude_tags: vec![],
630            openapi_include_operations: vec![],
631            openapi_exclude_operations: vec![],
632            openapi_max_operations: None,
633            openapi_overrides: HashMap::new(),
634            auth_generator: None,
635            category: None,
636            skills: vec![],
637        };
638
639        let mut inject = HashMap::new();
640        inject.insert(
641            "Credentials.AccessKeyId".into(),
642            crate::core::manifest::InjectTarget {
643                inject_type: "header".into(),
644                name: "X-Access-Key".into(),
645            },
646        );
647        inject.insert(
648            "Credentials.Secret".into(),
649            crate::core::manifest::InjectTarget {
650                inject_type: "env".into(),
651                name: "AWS_SECRET".into(),
652            },
653        );
654
655        let gen = AuthGenerator {
656            gen_type: AuthGenType::Command,
657            command: Some("echo".into()),
658            args: vec![
659                r#"{"Credentials":{"AccessKeyId":"AKIA123","Secret":"wJalr","SessionToken":"FwoG"}}"#.into(),
660            ],
661            interpreter: None,
662            script: None,
663            cache_ttl_secs: 0,
664            output_format: AuthOutputFormat::Json,
665            env: HashMap::new(),
666            inject,
667            timeout_secs: 5,
668        };
669
670        let ctx = GenContext::default();
671        let keyring = Keyring::empty();
672        let cache = AuthCache::new();
673
674        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
675            .await
676            .unwrap();
677        assert_eq!(cred.extra_headers.get("X-Access-Key").unwrap(), "AKIA123");
678        assert_eq!(cred.extra_env.get("AWS_SECRET").unwrap(), "wJalr");
679    }
680
681    #[tokio::test]
682    async fn test_generate_script() {
683        let provider = Provider {
684            name: "test".into(),
685            description: "test".into(),
686            base_url: String::new(),
687            auth_type: crate::core::manifest::AuthType::Bearer,
688            auth_key_name: None,
689            auth_header_name: None,
690            auth_query_name: None,
691            auth_value_prefix: None,
692            extra_headers: HashMap::new(),
693            oauth2_token_url: None,
694            auth_secret_name: None,
695            oauth2_basic_auth: false,
696            internal: false,
697            handler: "http".into(),
698            mcp_transport: None,
699            mcp_command: None,
700            mcp_args: vec![],
701            mcp_url: None,
702            mcp_env: HashMap::new(),
703            cli_command: None,
704            cli_default_args: vec![],
705            cli_env: HashMap::new(),
706            cli_timeout_secs: None,
707            openapi_spec: None,
708            openapi_include_tags: vec![],
709            openapi_exclude_tags: vec![],
710            openapi_include_operations: vec![],
711            openapi_exclude_operations: vec![],
712            openapi_max_operations: None,
713            openapi_overrides: HashMap::new(),
714            auth_generator: None,
715            category: None,
716            skills: vec![],
717        };
718
719        let gen = AuthGenerator {
720            gen_type: AuthGenType::Script,
721            command: None,
722            args: vec![],
723            interpreter: Some("bash".into()),
724            script: Some("echo script-token-42".into()),
725            cache_ttl_secs: 0,
726            output_format: AuthOutputFormat::Text,
727            env: HashMap::new(),
728            inject: HashMap::new(),
729            timeout_secs: 5,
730        };
731
732        let ctx = GenContext::default();
733        let keyring = Keyring::empty();
734        let cache = AuthCache::new();
735
736        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
737            .await
738            .unwrap();
739        assert_eq!(cred.value, "script-token-42");
740    }
741
742    #[tokio::test]
743    async fn test_generate_caches_result() {
744        let provider = Provider {
745            name: "cached_provider".into(),
746            description: "test".into(),
747            base_url: String::new(),
748            auth_type: crate::core::manifest::AuthType::Bearer,
749            auth_key_name: None,
750            auth_header_name: None,
751            auth_query_name: None,
752            auth_value_prefix: None,
753            extra_headers: HashMap::new(),
754            oauth2_token_url: None,
755            auth_secret_name: None,
756            oauth2_basic_auth: false,
757            internal: false,
758            handler: "http".into(),
759            mcp_transport: None,
760            mcp_command: None,
761            mcp_args: vec![],
762            mcp_url: None,
763            mcp_env: HashMap::new(),
764            cli_command: None,
765            cli_default_args: vec![],
766            cli_env: HashMap::new(),
767            cli_timeout_secs: None,
768            openapi_spec: None,
769            openapi_include_tags: vec![],
770            openapi_exclude_tags: vec![],
771            openapi_include_operations: vec![],
772            openapi_exclude_operations: vec![],
773            openapi_max_operations: None,
774            openapi_overrides: HashMap::new(),
775            auth_generator: None,
776            category: None,
777            skills: vec![],
778        };
779
780        let gen = AuthGenerator {
781            gen_type: AuthGenType::Command,
782            command: Some("date".into()),
783            args: vec!["+%s%N".into()],
784            interpreter: None,
785            script: None,
786            cache_ttl_secs: 300,
787            output_format: AuthOutputFormat::Text,
788            env: HashMap::new(),
789            inject: HashMap::new(),
790            timeout_secs: 5,
791        };
792
793        let ctx = GenContext {
794            jwt_sub: "test-agent".into(),
795            ..GenContext::default()
796        };
797        let keyring = Keyring::empty();
798        let cache = AuthCache::new();
799
800        let cred1 = generate(&provider, &gen, &ctx, &keyring, &cache)
801            .await
802            .unwrap();
803        let cred2 = generate(&provider, &gen, &ctx, &keyring, &cache)
804            .await
805            .unwrap();
806        // Second call should return cached value (same value)
807        assert_eq!(cred1.value, cred2.value);
808    }
809
810    #[tokio::test]
811    async fn test_generate_with_variable_expansion() {
812        let provider = Provider {
813            name: "test".into(),
814            description: "test".into(),
815            base_url: String::new(),
816            auth_type: crate::core::manifest::AuthType::Bearer,
817            auth_key_name: None,
818            auth_header_name: None,
819            auth_query_name: None,
820            auth_value_prefix: None,
821            extra_headers: HashMap::new(),
822            oauth2_token_url: None,
823            auth_secret_name: None,
824            oauth2_basic_auth: false,
825            internal: false,
826            handler: "http".into(),
827            mcp_transport: None,
828            mcp_command: None,
829            mcp_args: vec![],
830            mcp_url: None,
831            mcp_env: HashMap::new(),
832            cli_command: None,
833            cli_default_args: vec![],
834            cli_env: HashMap::new(),
835            cli_timeout_secs: None,
836            openapi_spec: None,
837            openapi_include_tags: vec![],
838            openapi_exclude_tags: vec![],
839            openapi_include_operations: vec![],
840            openapi_exclude_operations: vec![],
841            openapi_max_operations: None,
842            openapi_overrides: HashMap::new(),
843            auth_generator: None,
844            category: None,
845            skills: vec![],
846        };
847
848        let gen = AuthGenerator {
849            gen_type: AuthGenType::Command,
850            command: Some("echo".into()),
851            args: vec!["${JWT_SUB}".into()],
852            interpreter: None,
853            script: None,
854            cache_ttl_secs: 0,
855            output_format: AuthOutputFormat::Text,
856            env: HashMap::new(),
857            inject: HashMap::new(),
858            timeout_secs: 5,
859        };
860
861        let ctx = GenContext {
862            jwt_sub: "agent-42".into(),
863            jwt_scope: "*".into(),
864            tool_name: "brain:query".into(),
865            timestamp: 1234567890,
866        };
867        let keyring = Keyring::empty();
868        let cache = AuthCache::new();
869
870        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
871            .await
872            .unwrap();
873        assert_eq!(cred.value, "agent-42");
874    }
875
876    #[tokio::test]
877    async fn test_generate_timeout() {
878        let provider = Provider {
879            name: "test".into(),
880            description: "test".into(),
881            base_url: String::new(),
882            auth_type: crate::core::manifest::AuthType::Bearer,
883            auth_key_name: None,
884            auth_header_name: None,
885            auth_query_name: None,
886            auth_value_prefix: None,
887            extra_headers: HashMap::new(),
888            oauth2_token_url: None,
889            auth_secret_name: None,
890            oauth2_basic_auth: false,
891            internal: false,
892            handler: "http".into(),
893            mcp_transport: None,
894            mcp_command: None,
895            mcp_args: vec![],
896            mcp_url: None,
897            mcp_env: HashMap::new(),
898            cli_command: None,
899            cli_default_args: vec![],
900            cli_env: HashMap::new(),
901            cli_timeout_secs: None,
902            openapi_spec: None,
903            openapi_include_tags: vec![],
904            openapi_exclude_tags: vec![],
905            openapi_include_operations: vec![],
906            openapi_exclude_operations: vec![],
907            openapi_max_operations: None,
908            openapi_overrides: HashMap::new(),
909            auth_generator: None,
910            category: None,
911            skills: vec![],
912        };
913
914        let gen = AuthGenerator {
915            gen_type: AuthGenType::Command,
916            command: Some("sleep".into()),
917            args: vec!["10".into()],
918            interpreter: None,
919            script: None,
920            cache_ttl_secs: 0,
921            output_format: AuthOutputFormat::Text,
922            env: HashMap::new(),
923            inject: HashMap::new(),
924            timeout_secs: 1,
925        };
926
927        let ctx = GenContext::default();
928        let keyring = Keyring::empty();
929        let cache = AuthCache::new();
930
931        let err = generate(&provider, &gen, &ctx, &keyring, &cache)
932            .await
933            .unwrap_err();
934        assert!(matches!(err, AuthGenError::Timeout(1)));
935    }
936}
ati/core/auth_generator.rs

ati/core/
auth_generator.rs
