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    /// Raw inbound bearer token (the `<token>` from `Authorization: Bearer
53    /// <token>`). Populated by the proxy from the request that triggered
54    /// this generator run, or by the CLI from `$ATI_SESSION_TOKEN` in
55    /// direct-CLI mode. Empty string when no inbound bearer is available
56    /// (dev mode, no JWT validation configured, no env var set). Exposed
57    /// to auth_generators via `${JWT_TOKEN}` so an MCP provider can forward
58    /// the calling sandbox's identity to its upstream MCP without the proxy
59    /// holding a long-lived static bearer that loses per-sandbox attribution.
60    /// See issue #115.
61    pub jwt_token: String,
62}
63
64impl Default for GenContext {
65    fn default() -> Self {
66        GenContext {
67            jwt_sub: "dev".into(),
68            jwt_scope: "*".into(),
69            tool_name: String::new(),
70            timestamp: std::time::SystemTime::now()
71                .duration_since(std::time::UNIX_EPOCH)
72                .unwrap_or_default()
73                .as_secs(),
74            jwt_token: String::new(),
75        }
76    }
77}
78
79// ---------------------------------------------------------------------------
80// Generated credential
81// ---------------------------------------------------------------------------
82
83/// Result of running a generator — primary token + optional extra injections.
84#[derive(Debug, Clone)]
85pub struct GeneratedCredential {
86    /// Primary token value (used for bearer/header/query auth).
87    pub value: String,
88    /// Extra headers from JSON inject targets with type="header".
89    pub extra_headers: HashMap<String, String>,
90    /// Extra env vars from JSON inject targets with type="env".
91    pub extra_env: HashMap<String, String>,
92}
93
94// ---------------------------------------------------------------------------
95// Cache
96// ---------------------------------------------------------------------------
97
98struct CachedCredential {
99    cred: GeneratedCredential,
100    expires_at: Instant,
101}
102
103/// TTL-based credential cache, keyed by `(provider_name, agent_sub,
104/// token_fingerprint)`.
105///
106/// The token fingerprint is the SHA-256 of the raw inbound bearer truncated
107/// to 16 hex chars (or the empty string when no bearer was presented). We
108/// added this third dimension so a per-sandbox JWT (which rotates per request
109/// in some deployments) doesn't either (a) cache one sandbox's generated
110/// credential and serve it to another sandbox with the same `sub`, or
111/// (b) force every consumer to set `cache_ttl_secs = 0`. With the fingerprint
112/// in the key, legitimate same-sandbox reuse hits the cache and cross-sandbox
113/// reuse misses — even when `sub` happens to collide. See issue #115.
114///
115/// The fingerprint is a one-way hash, so the cache map never carries raw
116/// bearer bytes. Truncating to 16 hex chars (64 bits) is plenty for collision
117/// avoidance at this cache's scale (provider × sub already partitions the
118/// keyspace; the token dimension just needs to distinguish concurrent
119/// per-request tokens for the same agent).
120pub struct AuthCache {
121    entries: Mutex<HashMap<(String, String, String), CachedCredential>>,
122}
123
124impl Default for AuthCache {
125    fn default() -> Self {
126        AuthCache {
127            entries: Mutex::new(HashMap::new()),
128        }
129    }
130}
131
132/// SHA-256 of `token`, truncated to 16 hex chars. Empty input → empty
133/// fingerprint (intentional — sentinel for "no inbound bearer" so all the
134/// no-bearer paths share one cache slot per (provider, sub)).
135pub fn token_fingerprint(token: &str) -> String {
136    if token.is_empty() {
137        return String::new();
138    }
139    use sha2::{Digest, Sha256};
140    let mut hasher = Sha256::new();
141    hasher.update(token.as_bytes());
142    let digest = hasher.finalize();
143    hex::encode(digest)[..16].to_string()
144}
145
146impl AuthCache {
147    pub fn new() -> Self {
148        Self::default()
149    }
150
151    pub fn get(&self, provider: &str, sub: &str, token: &str) -> Option<GeneratedCredential> {
152        let cache = self.entries.lock().unwrap();
153        let key = (
154            provider.to_string(),
155            sub.to_string(),
156            token_fingerprint(token),
157        );
158        match cache.get(&key) {
159            Some(entry) if Instant::now() < entry.expires_at => Some(entry.cred.clone()),
160            _ => None,
161        }
162    }
163
164    pub fn insert(
165        &self,
166        provider: &str,
167        sub: &str,
168        token: &str,
169        cred: GeneratedCredential,
170        ttl_secs: u64,
171    ) {
172        if ttl_secs == 0 {
173            return; // No caching
174        }
175        let mut cache = self.entries.lock().unwrap();
176        // Sweep expired entries on every insert. Without the token-fingerprint
177        // dimension the previous `(provider, sub)` key bounded the map to
178        // O(providers × subs), so stale entries were always overwritten and a
179        // background sweep wasn't needed. With per-token keys, rotating JWTs
180        // (the very production scenario this PR is targeting — see issue #115)
181        // create a new permanent map entry per request unless we prune. The
182        // sweep runs under the same lock we already hold for the insert, so
183        // it adds no contention; cost is O(n) over the live map, amortized
184        // across inserts. Greptile P1 on PR #117.
185        let now = Instant::now();
186        cache.retain(|_, v| now < v.expires_at);
187        let key = (
188            provider.to_string(),
189            sub.to_string(),
190            token_fingerprint(token),
191        );
192        cache.insert(
193            key,
194            CachedCredential {
195                cred,
196                expires_at: now + Duration::from_secs(ttl_secs),
197            },
198        );
199    }
200
201    /// Number of map entries (live + just-expired-not-yet-swept).
202    /// Test-only helper for asserting the sweep on insert actually evicts
203    /// expired rows. Not exposed publicly to avoid encouraging callers
204    /// to depend on a particular eviction cadence.
205    #[cfg(test)]
206    pub fn entry_count(&self) -> usize {
207        self.entries.lock().unwrap().len()
208    }
209}
210
211// ---------------------------------------------------------------------------
212// Main generate function
213// ---------------------------------------------------------------------------
214
215/// Generate a credential by running the provider's auth_generator.
216///
217/// 1. Check cache → return if hit
218/// 2. Expand variables in args and env
219/// 3. Spawn subprocess (command or script)
220/// 4. Parse output (text or JSON)
221/// 5. Cache and return
222pub async fn generate(
223    provider: &Provider,
224    gen: &AuthGenerator,
225    ctx: &GenContext,
226    keyring: &Keyring,
227    cache: &AuthCache,
228) -> Result<GeneratedCredential, AuthGenError> {
229    // 1. Check cache. Key includes a fingerprint of the inbound JWT so that
230    // per-sandbox bearers don't share cached credentials across sandboxes
231    // even when they happen to land on the same `sub` (see issue #115).
232    if gen.cache_ttl_secs > 0 {
233        if let Some(cached) = cache.get(&provider.name, &ctx.jwt_sub, &ctx.jwt_token) {
234            return Ok(cached);
235        }
236    }
237
238    // 2. Expand variables in args and env
239    let expanded_args: Vec<String> = gen
240        .args
241        .iter()
242        .map(|a| expand_variables(a, ctx, keyring))
243        .collect::<Result<Vec<_>, _>>()?;
244
245    let mut expanded_env: HashMap<String, String> = HashMap::new();
246    for (k, v) in &gen.env {
247        expanded_env.insert(k.clone(), expand_variables(v, ctx, keyring)?);
248    }
249
250    // 3. Build curated env (don't leak host secrets)
251    let mut final_env: HashMap<String, String> = HashMap::new();
252    for var in &["PATH", "HOME", "TMPDIR"] {
253        if let Ok(val) = std::env::var(var) {
254            final_env.insert(var.to_string(), val);
255        }
256    }
257    final_env.extend(expanded_env);
258
259    // 4. Spawn subprocess
260    let output =
261        match gen.gen_type {
262            AuthGenType::Command => {
263                let command = gen.command.as_deref().ok_or_else(|| {
264                    AuthGenError::Config("command required for type=command".into())
265                })?;
266
267                let child = tokio::process::Command::new(command)
268                    .args(&expanded_args)
269                    .env_clear()
270                    .envs(&final_env)
271                    .stdout(std::process::Stdio::piped())
272                    .stderr(std::process::Stdio::piped())
273                    .kill_on_drop(true)
274                    .spawn()
275                    .map_err(|e| AuthGenError::Spawn(format!("{command}: {e}")))?;
276
277                let timeout = Duration::from_secs(gen.timeout_secs);
278                tokio::time::timeout(timeout, child.wait_with_output())
279                    .await
280                    .map_err(|_| AuthGenError::Timeout(gen.timeout_secs))?
281                    .map_err(AuthGenError::Io)?
282            }
283            AuthGenType::Script => {
284                let interpreter = gen.interpreter.as_deref().ok_or_else(|| {
285                    AuthGenError::Config("interpreter required for type=script".into())
286                })?;
287                let script = gen.script.as_deref().ok_or_else(|| {
288                    AuthGenError::Config("script required for type=script".into())
289                })?;
290
291                // Write script to a temp file
292                let suffix: u32 = rand::random();
293                let tmp_path = std::env::temp_dir().join(format!("ati_gen_{suffix}.tmp"));
294                std::fs::write(&tmp_path, script).map_err(AuthGenError::Io)?;
295
296                let child = tokio::process::Command::new(interpreter)
297                    .arg(&tmp_path)
298                    .env_clear()
299                    .envs(&final_env)
300                    .stdout(std::process::Stdio::piped())
301                    .stderr(std::process::Stdio::piped())
302                    .kill_on_drop(true)
303                    .spawn()
304                    .map_err(|e| AuthGenError::Spawn(format!("{interpreter}: {e}")))?;
305
306                let timeout = Duration::from_secs(gen.timeout_secs);
307                let result = tokio::time::timeout(timeout, child.wait_with_output())
308                    .await
309                    .map_err(|_| AuthGenError::Timeout(gen.timeout_secs))?
310                    .map_err(AuthGenError::Io)?;
311
312                // Clean up temp file
313                let _ = std::fs::remove_file(&tmp_path);
314                result
315            }
316        };
317
318    if !output.status.success() {
319        let code = output.status.code().unwrap_or(-1);
320        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
321        return Err(AuthGenError::NonZeroExit { code, stderr });
322    }
323
324    let stdout = String::from_utf8_lossy(&output.stdout);
325
326    // 5. Parse output
327    let cred = match gen.output_format {
328        AuthOutputFormat::Text => GeneratedCredential {
329            value: stdout.trim().to_string(),
330            extra_headers: HashMap::new(),
331            extra_env: HashMap::new(),
332        },
333        AuthOutputFormat::Json => {
334            let json: serde_json::Value = serde_json::from_str(stdout.trim())
335                .map_err(|e| AuthGenError::OutputParse(format!("invalid JSON: {e}")))?;
336
337            let mut extra_headers = HashMap::new();
338            let mut extra_env = HashMap::new();
339            let mut primary_value = stdout.trim().to_string();
340
341            // If no inject map, use the whole output as the primary value
342            if gen.inject.is_empty() {
343                // Try to extract a "token" or "access_token" field as primary
344                if let Some(tok) = json.get("token").or(json.get("access_token")) {
345                    if let Some(s) = tok.as_str() {
346                        primary_value = s.to_string();
347                    }
348                }
349            } else {
350                // Extract fields per inject map
351                let mut found_primary = false;
352                for (json_path, target) in &gen.inject {
353                    let extracted = extract_json_path(&json, json_path).ok_or_else(|| {
354                        AuthGenError::OutputParse(format!(
355                            "JSON path '{}' not found in output",
356                            json_path
357                        ))
358                    })?;
359
360                    match target.inject_type.as_str() {
361                        "header" => {
362                            extra_headers.insert(target.name.clone(), extracted);
363                        }
364                        "env" => {
365                            extra_env.insert(target.name.clone(), extracted);
366                        }
367                        "query" => {
368                            // For query injection, use as primary value
369                            if !found_primary {
370                                primary_value = extracted;
371                                found_primary = true;
372                            }
373                        }
374                        _ => {
375                            // Default: treat as primary value
376                            if !found_primary {
377                                primary_value = extracted;
378                                found_primary = true;
379                            }
380                        }
381                    }
382                }
383            }
384
385            GeneratedCredential {
386                value: primary_value,
387                extra_headers,
388                extra_env,
389            }
390        }
391    };
392
393    // 6. Cache. See `cache.get` above for why the JWT fingerprint is part
394    // of the key.
395    cache.insert(
396        &provider.name,
397        &ctx.jwt_sub,
398        &ctx.jwt_token,
399        cred.clone(),
400        gen.cache_ttl_secs,
401    );
402
403    Ok(cred)
404}
405
406// ---------------------------------------------------------------------------
407// Variable expansion
408// ---------------------------------------------------------------------------
409
410/// Expand `${VAR}` placeholders in a string.
411///
412/// Recognized variables:
413/// - `${JWT_SUB}`, `${JWT_SCOPE}`, `${TOOL_NAME}`, `${TIMESTAMP}` — from GenContext
414/// - `${JWT_TOKEN}` — raw inbound bearer (proxy: from `Authorization: Bearer …`;
415///   CLI: from `$ATI_SESSION_TOKEN`). Empty when no inbound bearer is present.
416/// - `${anything_else}` — looked up in the keyring
417fn expand_variables(
418    input: &str,
419    ctx: &GenContext,
420    keyring: &Keyring,
421) -> Result<String, AuthGenError> {
422    let mut result = input.to_string();
423    // Process all ${...} patterns
424    while let Some(start) = result.find("${") {
425        let rest = &result[start + 2..];
426        let end = match rest.find('}') {
427            Some(e) => e,
428            None => break,
429        };
430        let var_name = &rest[..end];
431
432        let replacement = match var_name {
433            "JWT_SUB" => ctx.jwt_sub.clone(),
434            "JWT_SCOPE" => ctx.jwt_scope.clone(),
435            "TOOL_NAME" => ctx.tool_name.clone(),
436            "TIMESTAMP" => ctx.timestamp.to_string(),
437            // The raw inbound bearer. Empty in dev mode and in any context
438            // where no JWT was presented (CLI direct mode without
439            // `$ATI_SESSION_TOKEN`). Generators that *require* a non-empty
440            // token should fail explicitly inside their own script rather
441            // than relying on `expand_variables` to error — silently sending
442            // an empty bearer to the upstream is a debuggable 401, whereas
443            // synthesizing one would be a security hole.
444            "JWT_TOKEN" => ctx.jwt_token.clone(),
445            _ => {
446                // Keyring lookup
447                match keyring.get(var_name) {
448                    Some(val) => val.to_string(),
449                    None => return Err(AuthGenError::KeyringMissing(var_name.to_string())),
450                }
451            }
452        };
453
454        result = format!("{}{}{}", &result[..start], replacement, &rest[end + 1..]);
455    }
456    Ok(result)
457}
458
459// ---------------------------------------------------------------------------
460// JSON path extraction (dot-notation)
461// ---------------------------------------------------------------------------
462
463/// Extract a value from a JSON object using dot-notation path.
464///
465/// Example: `extract_json_path(json, "Credentials.AccessKeyId")`
466/// navigates `json["Credentials"]["AccessKeyId"]`.
467fn extract_json_path(value: &serde_json::Value, path: &str) -> Option<String> {
468    let mut current = value;
469    for segment in path.split('.') {
470        current = current.get(segment)?;
471    }
472    match current {
473        serde_json::Value::String(s) => Some(s.clone()),
474        serde_json::Value::Number(n) => Some(n.to_string()),
475        serde_json::Value::Bool(b) => Some(b.to_string()),
476        other => Some(other.to_string()),
477    }
478}
479
480// ---------------------------------------------------------------------------
481// Tests
482// ---------------------------------------------------------------------------
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487
488    #[test]
489    fn test_expand_variables_context() {
490        let ctx = GenContext {
491            jwt_sub: "agent-7".into(),
492            jwt_scope: "tool:brain:*".into(),
493            tool_name: "brain:query".into(),
494            timestamp: 1773096459,
495            jwt_token: "eyJhbGciOiJIUzI1NiJ9.payload.sig".into(),
496        };
497        let keyring = Keyring::empty();
498
499        assert_eq!(
500            expand_variables("${JWT_SUB}", &ctx, &keyring).unwrap(),
501            "agent-7"
502        );
503        assert_eq!(
504            expand_variables("${TOOL_NAME}", &ctx, &keyring).unwrap(),
505            "brain:query"
506        );
507        assert_eq!(
508            expand_variables("${TIMESTAMP}", &ctx, &keyring).unwrap(),
509            "1773096459"
510        );
511        assert_eq!(
512            expand_variables("${JWT_TOKEN}", &ctx, &keyring).unwrap(),
513            "eyJhbGciOiJIUzI1NiJ9.payload.sig"
514        );
515        assert_eq!(
516            expand_variables("sub=${JWT_SUB}&tool=${TOOL_NAME}", &ctx, &keyring).unwrap(),
517            "sub=agent-7&tool=brain:query"
518        );
519        // The Bearer-prefix form a passthrough/MCP auth_generator would use.
520        assert_eq!(
521            expand_variables("Bearer ${JWT_TOKEN}", &ctx, &keyring).unwrap(),
522            "Bearer eyJhbGciOiJIUzI1NiJ9.payload.sig"
523        );
524    }
525
526    #[test]
527    fn test_expand_variables_jwt_token_empty_when_unset() {
528        // Default context has no inbound bearer — expansion succeeds with
529        // an empty string rather than erroring. That keeps existing
530        // generators that never reference ${JWT_TOKEN} unaffected, and lets
531        // generators that DO reference it fail visibly downstream (401
532        // from the upstream) rather than silently substituting a static
533        // bootstrap token from the keyring.
534        let ctx = GenContext::default();
535        let keyring = Keyring::empty();
536        assert_eq!(
537            expand_variables("${JWT_TOKEN}", &ctx, &keyring).unwrap(),
538            ""
539        );
540    }
541
542    #[test]
543    fn test_expand_variables_keyring() {
544        let dir = tempfile::TempDir::new().unwrap();
545        let path = dir.path().join("creds");
546        std::fs::write(&path, r#"{"my_secret":"s3cr3t"}"#).unwrap();
547        let keyring = Keyring::load_credentials(&path).unwrap();
548
549        let ctx = GenContext::default();
550        assert_eq!(
551            expand_variables("${my_secret}", &ctx, &keyring).unwrap(),
552            "s3cr3t"
553        );
554    }
555
556    #[test]
557    fn test_expand_variables_missing_key() {
558        let keyring = Keyring::empty();
559        let ctx = GenContext::default();
560        let err = expand_variables("${nonexistent}", &ctx, &keyring).unwrap_err();
561        assert!(matches!(err, AuthGenError::KeyringMissing(_)));
562    }
563
564    #[test]
565    fn test_expand_variables_no_placeholder() {
566        let keyring = Keyring::empty();
567        let ctx = GenContext::default();
568        assert_eq!(
569            expand_variables("plain text", &ctx, &keyring).unwrap(),
570            "plain text"
571        );
572    }
573
574    #[test]
575    fn test_extract_json_path_simple() {
576        let json: serde_json::Value = serde_json::json!({"token": "abc123", "expires_in": 3600});
577        assert_eq!(extract_json_path(&json, "token"), Some("abc123".into()));
578        assert_eq!(extract_json_path(&json, "expires_in"), Some("3600".into()));
579    }
580
581    #[test]
582    fn test_extract_json_path_nested() {
583        let json: serde_json::Value = serde_json::json!({
584            "Credentials": {
585                "AccessKeyId": "AKIA...",
586                "SecretAccessKey": "wJalrX...",
587                "SessionToken": "FwoGZ..."
588            }
589        });
590        assert_eq!(
591            extract_json_path(&json, "Credentials.AccessKeyId"),
592            Some("AKIA...".into())
593        );
594        assert_eq!(
595            extract_json_path(&json, "Credentials.SessionToken"),
596            Some("FwoGZ...".into())
597        );
598    }
599
600    #[test]
601    fn test_extract_json_path_missing() {
602        let json: serde_json::Value = serde_json::json!({"a": "b"});
603        assert_eq!(extract_json_path(&json, "nonexistent"), None);
604        assert_eq!(extract_json_path(&json, "a.b.c"), None);
605    }
606
607    #[test]
608    fn test_auth_cache_basic() {
609        let cache = AuthCache::new();
610        assert!(cache.get("provider", "sub", "").is_none());
611
612        let cred = GeneratedCredential {
613            value: "token123".into(),
614            extra_headers: HashMap::new(),
615            extra_env: HashMap::new(),
616        };
617        cache.insert("provider", "sub", "", cred.clone(), 300);
618
619        let cached = cache.get("provider", "sub", "").unwrap();
620        assert_eq!(cached.value, "token123");
621    }
622
623    #[test]
624    fn test_auth_cache_zero_ttl_no_cache() {
625        let cache = AuthCache::new();
626        let cred = GeneratedCredential {
627            value: "token".into(),
628            extra_headers: HashMap::new(),
629            extra_env: HashMap::new(),
630        };
631        cache.insert("provider", "sub", "", cred, 0);
632        assert!(cache.get("provider", "sub", "").is_none());
633    }
634
635    #[test]
636    fn test_auth_cache_different_keys() {
637        let cache = AuthCache::new();
638        let cred1 = GeneratedCredential {
639            value: "token-a".into(),
640            extra_headers: HashMap::new(),
641            extra_env: HashMap::new(),
642        };
643        let cred2 = GeneratedCredential {
644            value: "token-b".into(),
645            extra_headers: HashMap::new(),
646            extra_env: HashMap::new(),
647        };
648        cache.insert("provider", "agent-1", "", cred1, 300);
649        cache.insert("provider", "agent-2", "", cred2, 300);
650
651        assert_eq!(
652            cache.get("provider", "agent-1", "").unwrap().value,
653            "token-a"
654        );
655        assert_eq!(
656            cache.get("provider", "agent-2", "").unwrap().value,
657            "token-b"
658        );
659    }
660
661    #[test]
662    fn test_auth_cache_per_token_isolation() {
663        // Issue #115: two sandboxes with the same `sub` but different
664        // inbound bearers must NOT share cached credentials. Otherwise the
665        // first sandbox's generated token gets served to the second.
666        let cache = AuthCache::new();
667        let cred_a = GeneratedCredential {
668            value: "for-sandbox-a".into(),
669            extra_headers: HashMap::new(),
670            extra_env: HashMap::new(),
671        };
672        let cred_b = GeneratedCredential {
673            value: "for-sandbox-b".into(),
674            extra_headers: HashMap::new(),
675            extra_env: HashMap::new(),
676        };
677        cache.insert("provider", "sandbox-svc", "jwt-A", cred_a, 300);
678        cache.insert("provider", "sandbox-svc", "jwt-B", cred_b, 300);
679
680        // Each token gets its own cached cred.
681        assert_eq!(
682            cache.get("provider", "sandbox-svc", "jwt-A").unwrap().value,
683            "for-sandbox-a"
684        );
685        assert_eq!(
686            cache.get("provider", "sandbox-svc", "jwt-B").unwrap().value,
687            "for-sandbox-b"
688        );
689        // A third unseen token misses (no leakage from either A or B).
690        assert!(cache.get("provider", "sandbox-svc", "jwt-C").is_none());
691    }
692
693    #[test]
694    fn test_token_fingerprint_stable_and_distinct() {
695        let f1 = token_fingerprint("token-one");
696        let f2 = token_fingerprint("token-two");
697        assert_eq!(token_fingerprint("token-one"), f1, "fingerprint is stable");
698        assert_ne!(f1, f2, "different tokens hash to different fingerprints");
699        assert_eq!(f1.len(), 16, "fingerprint is 16 hex chars");
700        assert!(f1.chars().all(|c| c.is_ascii_hexdigit()));
701        // Empty input gets the empty sentinel so all no-bearer paths share
702        // one cache slot per (provider, sub).
703        assert_eq!(token_fingerprint(""), "");
704    }
705
706    #[test]
707    fn test_auth_cache_insert_evicts_expired_entries() {
708        // Greptile P1 on PR #117: rotating per-request JWTs would create one
709        // permanent map entry per request because expired entries were never
710        // swept. The fix prunes on every `insert`. Verify the prune actually
711        // happens by inserting one entry with the smallest possible non-zero
712        // TTL (1 second), waiting for it to expire, then inserting a
713        // different key — the second insert MUST drop the first.
714        let cache = AuthCache::new();
715        let cred = GeneratedCredential {
716            value: "ephemeral".into(),
717            extra_headers: HashMap::new(),
718            extra_env: HashMap::new(),
719        };
720        cache.insert("p", "s", "old-jwt", cred.clone(), 1);
721        assert_eq!(cache.entry_count(), 1);
722
723        // Sleep past the 1s TTL. 1100ms is plenty of headroom on a busy
724        // runner without inflating test runtime.
725        std::thread::sleep(Duration::from_millis(1100));
726
727        // A second insert (different fingerprint) triggers the sweep.
728        cache.insert("p", "s", "new-jwt", cred, 60);
729        assert_eq!(
730            cache.entry_count(),
731            1,
732            "sweep on insert should have dropped the expired entry; \
733             only the new one should remain"
734        );
735
736        // Sanity-check which entry survived.
737        assert!(cache.get("p", "s", "new-jwt").is_some());
738        assert!(cache.get("p", "s", "old-jwt").is_none());
739    }
740
741    #[tokio::test]
742    async fn test_generate_command_text() {
743        let provider = Provider {
744            name: "test".into(),
745            description: "test provider".into(),
746            base_url: String::new(),
747            auth_type: crate::core::manifest::AuthType::Bearer,
748            auth_key_name: None,
749            auth_header_name: None,
750            auth_query_name: None,
751            auth_value_prefix: None,
752            extra_headers: HashMap::new(),
753            oauth2_token_url: None,
754            auth_secret_name: None,
755            auth_session_token_env: 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            cli_output_args: Vec::new(),
769            cli_output_positional: HashMap::new(),
770            upload_destinations: HashMap::new(),
771            upload_default_destination: None,
772            openapi_spec: None,
773            openapi_include_tags: vec![],
774            openapi_exclude_tags: vec![],
775            openapi_include_operations: vec![],
776            openapi_exclude_operations: vec![],
777            openapi_max_operations: None,
778            openapi_overrides: HashMap::new(),
779            auth_generator: None,
780            category: None,
781            skills: vec![],
782        };
783
784        let gen = AuthGenerator {
785            gen_type: AuthGenType::Command,
786            command: Some("echo".into()),
787            args: vec!["hello-token".into()],
788            interpreter: None,
789            script: None,
790            cache_ttl_secs: 0,
791            output_format: AuthOutputFormat::Text,
792            env: HashMap::new(),
793            inject: HashMap::new(),
794            timeout_secs: 5,
795        };
796
797        let ctx = GenContext::default();
798        let keyring = Keyring::empty();
799        let cache = AuthCache::new();
800
801        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
802            .await
803            .unwrap();
804        assert_eq!(cred.value, "hello-token");
805        assert!(cred.extra_headers.is_empty());
806    }
807
808    #[tokio::test]
809    async fn test_generate_command_json() {
810        let provider = Provider {
811            name: "test".into(),
812            description: "test".into(),
813            base_url: String::new(),
814            auth_type: crate::core::manifest::AuthType::Bearer,
815            auth_key_name: None,
816            auth_header_name: None,
817            auth_query_name: None,
818            auth_value_prefix: None,
819            extra_headers: HashMap::new(),
820            oauth2_token_url: None,
821            auth_secret_name: None,
822            auth_session_token_env: None,
823            oauth2_basic_auth: false,
824            internal: false,
825            handler: "http".into(),
826            mcp_transport: None,
827            mcp_command: None,
828            mcp_args: vec![],
829            mcp_url: None,
830            mcp_env: HashMap::new(),
831            cli_command: None,
832            cli_default_args: vec![],
833            cli_env: HashMap::new(),
834            cli_timeout_secs: None,
835            cli_output_args: Vec::new(),
836            cli_output_positional: HashMap::new(),
837            upload_destinations: HashMap::new(),
838            upload_default_destination: None,
839            openapi_spec: None,
840            openapi_include_tags: vec![],
841            openapi_exclude_tags: vec![],
842            openapi_include_operations: vec![],
843            openapi_exclude_operations: vec![],
844            openapi_max_operations: None,
845            openapi_overrides: HashMap::new(),
846            auth_generator: None,
847            category: None,
848            skills: vec![],
849        };
850
851        let mut inject = HashMap::new();
852        inject.insert(
853            "Credentials.AccessKeyId".into(),
854            crate::core::manifest::InjectTarget {
855                inject_type: "header".into(),
856                name: "X-Access-Key".into(),
857            },
858        );
859        inject.insert(
860            "Credentials.Secret".into(),
861            crate::core::manifest::InjectTarget {
862                inject_type: "env".into(),
863                name: "AWS_SECRET".into(),
864            },
865        );
866
867        let gen = AuthGenerator {
868            gen_type: AuthGenType::Command,
869            command: Some("echo".into()),
870            args: vec![
871                r#"{"Credentials":{"AccessKeyId":"AKIA123","Secret":"wJalr","SessionToken":"FwoG"}}"#.into(),
872            ],
873            interpreter: None,
874            script: None,
875            cache_ttl_secs: 0,
876            output_format: AuthOutputFormat::Json,
877            env: HashMap::new(),
878            inject,
879            timeout_secs: 5,
880        };
881
882        let ctx = GenContext::default();
883        let keyring = Keyring::empty();
884        let cache = AuthCache::new();
885
886        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
887            .await
888            .unwrap();
889        assert_eq!(cred.extra_headers.get("X-Access-Key").unwrap(), "AKIA123");
890        assert_eq!(cred.extra_env.get("AWS_SECRET").unwrap(), "wJalr");
891    }
892
893    #[tokio::test]
894    async fn test_generate_script() {
895        let provider = Provider {
896            name: "test".into(),
897            description: "test".into(),
898            base_url: String::new(),
899            auth_type: crate::core::manifest::AuthType::Bearer,
900            auth_key_name: None,
901            auth_header_name: None,
902            auth_query_name: None,
903            auth_value_prefix: None,
904            extra_headers: HashMap::new(),
905            oauth2_token_url: None,
906            auth_secret_name: None,
907            auth_session_token_env: None,
908            oauth2_basic_auth: false,
909            internal: false,
910            handler: "http".into(),
911            mcp_transport: None,
912            mcp_command: None,
913            mcp_args: vec![],
914            mcp_url: None,
915            mcp_env: HashMap::new(),
916            cli_command: None,
917            cli_default_args: vec![],
918            cli_env: HashMap::new(),
919            cli_timeout_secs: None,
920            cli_output_args: Vec::new(),
921            cli_output_positional: HashMap::new(),
922            upload_destinations: HashMap::new(),
923            upload_default_destination: None,
924            openapi_spec: None,
925            openapi_include_tags: vec![],
926            openapi_exclude_tags: vec![],
927            openapi_include_operations: vec![],
928            openapi_exclude_operations: vec![],
929            openapi_max_operations: None,
930            openapi_overrides: HashMap::new(),
931            auth_generator: None,
932            category: None,
933            skills: vec![],
934        };
935
936        let gen = AuthGenerator {
937            gen_type: AuthGenType::Script,
938            command: None,
939            args: vec![],
940            interpreter: Some("bash".into()),
941            script: Some("echo script-token-42".into()),
942            cache_ttl_secs: 0,
943            output_format: AuthOutputFormat::Text,
944            env: HashMap::new(),
945            inject: HashMap::new(),
946            timeout_secs: 5,
947        };
948
949        let ctx = GenContext::default();
950        let keyring = Keyring::empty();
951        let cache = AuthCache::new();
952
953        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
954            .await
955            .unwrap();
956        assert_eq!(cred.value, "script-token-42");
957    }
958
959    #[tokio::test]
960    async fn test_generate_caches_result() {
961        let provider = Provider {
962            name: "cached_provider".into(),
963            description: "test".into(),
964            base_url: String::new(),
965            auth_type: crate::core::manifest::AuthType::Bearer,
966            auth_key_name: None,
967            auth_header_name: None,
968            auth_query_name: None,
969            auth_value_prefix: None,
970            extra_headers: HashMap::new(),
971            oauth2_token_url: None,
972            auth_secret_name: None,
973            auth_session_token_env: None,
974            oauth2_basic_auth: false,
975            internal: false,
976            handler: "http".into(),
977            mcp_transport: None,
978            mcp_command: None,
979            mcp_args: vec![],
980            mcp_url: None,
981            mcp_env: HashMap::new(),
982            cli_command: None,
983            cli_default_args: vec![],
984            cli_env: HashMap::new(),
985            cli_timeout_secs: None,
986            cli_output_args: Vec::new(),
987            cli_output_positional: HashMap::new(),
988            upload_destinations: HashMap::new(),
989            upload_default_destination: None,
990            openapi_spec: None,
991            openapi_include_tags: vec![],
992            openapi_exclude_tags: vec![],
993            openapi_include_operations: vec![],
994            openapi_exclude_operations: vec![],
995            openapi_max_operations: None,
996            openapi_overrides: HashMap::new(),
997            auth_generator: None,
998            category: None,
999            skills: vec![],
1000        };
1001
1002        let gen = AuthGenerator {
1003            gen_type: AuthGenType::Command,
1004            command: Some("date".into()),
1005            args: vec!["+%s%N".into()],
1006            interpreter: None,
1007            script: None,
1008            cache_ttl_secs: 300,
1009            output_format: AuthOutputFormat::Text,
1010            env: HashMap::new(),
1011            inject: HashMap::new(),
1012            timeout_secs: 5,
1013        };
1014
1015        let ctx = GenContext {
1016            jwt_sub: "test-agent".into(),
1017            ..GenContext::default()
1018        };
1019        let keyring = Keyring::empty();
1020        let cache = AuthCache::new();
1021
1022        let cred1 = generate(&provider, &gen, &ctx, &keyring, &cache)
1023            .await
1024            .unwrap();
1025        let cred2 = generate(&provider, &gen, &ctx, &keyring, &cache)
1026            .await
1027            .unwrap();
1028        // Second call should return cached value (same value)
1029        assert_eq!(cred1.value, cred2.value);
1030    }
1031
1032    #[tokio::test]
1033    async fn test_generate_with_variable_expansion() {
1034        let provider = Provider {
1035            name: "test".into(),
1036            description: "test".into(),
1037            base_url: String::new(),
1038            auth_type: crate::core::manifest::AuthType::Bearer,
1039            auth_key_name: None,
1040            auth_header_name: None,
1041            auth_query_name: None,
1042            auth_value_prefix: None,
1043            extra_headers: HashMap::new(),
1044            oauth2_token_url: None,
1045            auth_secret_name: None,
1046            auth_session_token_env: None,
1047            oauth2_basic_auth: false,
1048            internal: false,
1049            handler: "http".into(),
1050            mcp_transport: None,
1051            mcp_command: None,
1052            mcp_args: vec![],
1053            mcp_url: None,
1054            mcp_env: HashMap::new(),
1055            cli_command: None,
1056            cli_default_args: vec![],
1057            cli_env: HashMap::new(),
1058            cli_timeout_secs: None,
1059            cli_output_args: Vec::new(),
1060            cli_output_positional: HashMap::new(),
1061            upload_destinations: HashMap::new(),
1062            upload_default_destination: None,
1063            openapi_spec: None,
1064            openapi_include_tags: vec![],
1065            openapi_exclude_tags: vec![],
1066            openapi_include_operations: vec![],
1067            openapi_exclude_operations: vec![],
1068            openapi_max_operations: None,
1069            openapi_overrides: HashMap::new(),
1070            auth_generator: None,
1071            category: None,
1072            skills: vec![],
1073        };
1074
1075        let gen = AuthGenerator {
1076            gen_type: AuthGenType::Command,
1077            command: Some("echo".into()),
1078            args: vec!["${JWT_SUB}".into()],
1079            interpreter: None,
1080            script: None,
1081            cache_ttl_secs: 0,
1082            output_format: AuthOutputFormat::Text,
1083            env: HashMap::new(),
1084            inject: HashMap::new(),
1085            timeout_secs: 5,
1086        };
1087
1088        let ctx = GenContext {
1089            jwt_sub: "agent-42".into(),
1090            jwt_scope: "*".into(),
1091            tool_name: "brain:query".into(),
1092            timestamp: 1234567890,
1093            jwt_token: String::new(),
1094        };
1095        let keyring = Keyring::empty();
1096        let cache = AuthCache::new();
1097
1098        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
1099            .await
1100            .unwrap();
1101        assert_eq!(cred.value, "agent-42");
1102    }
1103
1104    #[tokio::test]
1105    async fn test_generate_timeout() {
1106        let provider = Provider {
1107            name: "test".into(),
1108            description: "test".into(),
1109            base_url: String::new(),
1110            auth_type: crate::core::manifest::AuthType::Bearer,
1111            auth_key_name: None,
1112            auth_header_name: None,
1113            auth_query_name: None,
1114            auth_value_prefix: None,
1115            extra_headers: HashMap::new(),
1116            oauth2_token_url: None,
1117            auth_secret_name: None,
1118            auth_session_token_env: None,
1119            oauth2_basic_auth: false,
1120            internal: false,
1121            handler: "http".into(),
1122            mcp_transport: None,
1123            mcp_command: None,
1124            mcp_args: vec![],
1125            mcp_url: None,
1126            mcp_env: HashMap::new(),
1127            cli_command: None,
1128            cli_default_args: vec![],
1129            cli_env: HashMap::new(),
1130            cli_timeout_secs: None,
1131            cli_output_args: Vec::new(),
1132            cli_output_positional: HashMap::new(),
1133            upload_destinations: HashMap::new(),
1134            upload_default_destination: None,
1135            openapi_spec: None,
1136            openapi_include_tags: vec![],
1137            openapi_exclude_tags: vec![],
1138            openapi_include_operations: vec![],
1139            openapi_exclude_operations: vec![],
1140            openapi_max_operations: None,
1141            openapi_overrides: HashMap::new(),
1142            auth_generator: None,
1143            category: None,
1144            skills: vec![],
1145        };
1146
1147        let gen = AuthGenerator {
1148            gen_type: AuthGenType::Command,
1149            command: Some("sleep".into()),
1150            args: vec!["10".into()],
1151            interpreter: None,
1152            script: None,
1153            cache_ttl_secs: 0,
1154            output_format: AuthOutputFormat::Text,
1155            env: HashMap::new(),
1156            inject: HashMap::new(),
1157            timeout_secs: 1,
1158        };
1159
1160        let ctx = GenContext::default();
1161        let keyring = Keyring::empty();
1162        let cache = AuthCache::new();
1163
1164        let err = generate(&provider, &gen, &ctx, &keyring, &cache)
1165            .await
1166            .unwrap_err();
1167        assert!(matches!(err, AuthGenError::Timeout(1)));
1168    }
1169}
ati/core/auth_generator.rs

ati/core/
auth_generator.rs
