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            mcp_url_env: None,
757            oauth2_basic_auth: false,
758            internal: false,
759            handler: "http".into(),
760            mcp_transport: None,
761            mcp_command: None,
762            mcp_args: vec![],
763            mcp_url: None,
764            mcp_env: HashMap::new(),
765            cli_command: None,
766            cli_default_args: vec![],
767            cli_env: HashMap::new(),
768            cli_timeout_secs: None,
769            cli_output_args: Vec::new(),
770            cli_output_positional: HashMap::new(),
771            upload_destinations: HashMap::new(),
772            upload_default_destination: None,
773            openapi_spec: None,
774            openapi_include_tags: vec![],
775            openapi_exclude_tags: vec![],
776            openapi_include_operations: vec![],
777            openapi_exclude_operations: vec![],
778            openapi_max_operations: None,
779            openapi_overrides: HashMap::new(),
780            auth_generator: None,
781            category: None,
782            skills: vec![],
783        };
784
785        let gen = AuthGenerator {
786            gen_type: AuthGenType::Command,
787            command: Some("echo".into()),
788            args: vec!["hello-token".into()],
789            interpreter: None,
790            script: None,
791            cache_ttl_secs: 0,
792            output_format: AuthOutputFormat::Text,
793            env: HashMap::new(),
794            inject: HashMap::new(),
795            timeout_secs: 5,
796        };
797
798        let ctx = GenContext::default();
799        let keyring = Keyring::empty();
800        let cache = AuthCache::new();
801
802        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
803            .await
804            .unwrap();
805        assert_eq!(cred.value, "hello-token");
806        assert!(cred.extra_headers.is_empty());
807    }
808
809    #[tokio::test]
810    async fn test_generate_command_json() {
811        let provider = Provider {
812            name: "test".into(),
813            description: "test".into(),
814            base_url: String::new(),
815            auth_type: crate::core::manifest::AuthType::Bearer,
816            auth_key_name: None,
817            auth_header_name: None,
818            auth_query_name: None,
819            auth_value_prefix: None,
820            extra_headers: HashMap::new(),
821            oauth2_token_url: None,
822            auth_secret_name: None,
823            auth_session_token_env: None,
824            mcp_url_env: None,
825            oauth2_basic_auth: false,
826            internal: false,
827            handler: "http".into(),
828            mcp_transport: None,
829            mcp_command: None,
830            mcp_args: vec![],
831            mcp_url: None,
832            mcp_env: HashMap::new(),
833            cli_command: None,
834            cli_default_args: vec![],
835            cli_env: HashMap::new(),
836            cli_timeout_secs: None,
837            cli_output_args: Vec::new(),
838            cli_output_positional: HashMap::new(),
839            upload_destinations: HashMap::new(),
840            upload_default_destination: None,
841            openapi_spec: None,
842            openapi_include_tags: vec![],
843            openapi_exclude_tags: vec![],
844            openapi_include_operations: vec![],
845            openapi_exclude_operations: vec![],
846            openapi_max_operations: None,
847            openapi_overrides: HashMap::new(),
848            auth_generator: None,
849            category: None,
850            skills: vec![],
851        };
852
853        let mut inject = HashMap::new();
854        inject.insert(
855            "Credentials.AccessKeyId".into(),
856            crate::core::manifest::InjectTarget {
857                inject_type: "header".into(),
858                name: "X-Access-Key".into(),
859            },
860        );
861        inject.insert(
862            "Credentials.Secret".into(),
863            crate::core::manifest::InjectTarget {
864                inject_type: "env".into(),
865                name: "AWS_SECRET".into(),
866            },
867        );
868
869        let gen = AuthGenerator {
870            gen_type: AuthGenType::Command,
871            command: Some("echo".into()),
872            args: vec![
873                r#"{"Credentials":{"AccessKeyId":"AKIA123","Secret":"wJalr","SessionToken":"FwoG"}}"#.into(),
874            ],
875            interpreter: None,
876            script: None,
877            cache_ttl_secs: 0,
878            output_format: AuthOutputFormat::Json,
879            env: HashMap::new(),
880            inject,
881            timeout_secs: 5,
882        };
883
884        let ctx = GenContext::default();
885        let keyring = Keyring::empty();
886        let cache = AuthCache::new();
887
888        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
889            .await
890            .unwrap();
891        assert_eq!(cred.extra_headers.get("X-Access-Key").unwrap(), "AKIA123");
892        assert_eq!(cred.extra_env.get("AWS_SECRET").unwrap(), "wJalr");
893    }
894
895    #[tokio::test]
896    async fn test_generate_script() {
897        let provider = Provider {
898            name: "test".into(),
899            description: "test".into(),
900            base_url: String::new(),
901            auth_type: crate::core::manifest::AuthType::Bearer,
902            auth_key_name: None,
903            auth_header_name: None,
904            auth_query_name: None,
905            auth_value_prefix: None,
906            extra_headers: HashMap::new(),
907            oauth2_token_url: None,
908            auth_secret_name: None,
909            auth_session_token_env: None,
910            mcp_url_env: None,
911            oauth2_basic_auth: false,
912            internal: false,
913            handler: "http".into(),
914            mcp_transport: None,
915            mcp_command: None,
916            mcp_args: vec![],
917            mcp_url: None,
918            mcp_env: HashMap::new(),
919            cli_command: None,
920            cli_default_args: vec![],
921            cli_env: HashMap::new(),
922            cli_timeout_secs: None,
923            cli_output_args: Vec::new(),
924            cli_output_positional: HashMap::new(),
925            upload_destinations: HashMap::new(),
926            upload_default_destination: None,
927            openapi_spec: None,
928            openapi_include_tags: vec![],
929            openapi_exclude_tags: vec![],
930            openapi_include_operations: vec![],
931            openapi_exclude_operations: vec![],
932            openapi_max_operations: None,
933            openapi_overrides: HashMap::new(),
934            auth_generator: None,
935            category: None,
936            skills: vec![],
937        };
938
939        let gen = AuthGenerator {
940            gen_type: AuthGenType::Script,
941            command: None,
942            args: vec![],
943            interpreter: Some("bash".into()),
944            script: Some("echo script-token-42".into()),
945            cache_ttl_secs: 0,
946            output_format: AuthOutputFormat::Text,
947            env: HashMap::new(),
948            inject: HashMap::new(),
949            timeout_secs: 5,
950        };
951
952        let ctx = GenContext::default();
953        let keyring = Keyring::empty();
954        let cache = AuthCache::new();
955
956        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
957            .await
958            .unwrap();
959        assert_eq!(cred.value, "script-token-42");
960    }
961
962    #[tokio::test]
963    async fn test_generate_caches_result() {
964        let provider = Provider {
965            name: "cached_provider".into(),
966            description: "test".into(),
967            base_url: String::new(),
968            auth_type: crate::core::manifest::AuthType::Bearer,
969            auth_key_name: None,
970            auth_header_name: None,
971            auth_query_name: None,
972            auth_value_prefix: None,
973            extra_headers: HashMap::new(),
974            oauth2_token_url: None,
975            auth_secret_name: None,
976            auth_session_token_env: None,
977            mcp_url_env: None,
978            oauth2_basic_auth: false,
979            internal: false,
980            handler: "http".into(),
981            mcp_transport: None,
982            mcp_command: None,
983            mcp_args: vec![],
984            mcp_url: None,
985            mcp_env: HashMap::new(),
986            cli_command: None,
987            cli_default_args: vec![],
988            cli_env: HashMap::new(),
989            cli_timeout_secs: None,
990            cli_output_args: Vec::new(),
991            cli_output_positional: HashMap::new(),
992            upload_destinations: HashMap::new(),
993            upload_default_destination: None,
994            openapi_spec: None,
995            openapi_include_tags: vec![],
996            openapi_exclude_tags: vec![],
997            openapi_include_operations: vec![],
998            openapi_exclude_operations: vec![],
999            openapi_max_operations: None,
1000            openapi_overrides: HashMap::new(),
1001            auth_generator: None,
1002            category: None,
1003            skills: vec![],
1004        };
1005
1006        let gen = AuthGenerator {
1007            gen_type: AuthGenType::Command,
1008            command: Some("date".into()),
1009            args: vec!["+%s%N".into()],
1010            interpreter: None,
1011            script: None,
1012            cache_ttl_secs: 300,
1013            output_format: AuthOutputFormat::Text,
1014            env: HashMap::new(),
1015            inject: HashMap::new(),
1016            timeout_secs: 5,
1017        };
1018
1019        let ctx = GenContext {
1020            jwt_sub: "test-agent".into(),
1021            ..GenContext::default()
1022        };
1023        let keyring = Keyring::empty();
1024        let cache = AuthCache::new();
1025
1026        let cred1 = generate(&provider, &gen, &ctx, &keyring, &cache)
1027            .await
1028            .unwrap();
1029        let cred2 = generate(&provider, &gen, &ctx, &keyring, &cache)
1030            .await
1031            .unwrap();
1032        // Second call should return cached value (same value)
1033        assert_eq!(cred1.value, cred2.value);
1034    }
1035
1036    #[tokio::test]
1037    async fn test_generate_with_variable_expansion() {
1038        let provider = Provider {
1039            name: "test".into(),
1040            description: "test".into(),
1041            base_url: String::new(),
1042            auth_type: crate::core::manifest::AuthType::Bearer,
1043            auth_key_name: None,
1044            auth_header_name: None,
1045            auth_query_name: None,
1046            auth_value_prefix: None,
1047            extra_headers: HashMap::new(),
1048            oauth2_token_url: None,
1049            auth_secret_name: None,
1050            auth_session_token_env: None,
1051            mcp_url_env: None,
1052            oauth2_basic_auth: false,
1053            internal: false,
1054            handler: "http".into(),
1055            mcp_transport: None,
1056            mcp_command: None,
1057            mcp_args: vec![],
1058            mcp_url: None,
1059            mcp_env: HashMap::new(),
1060            cli_command: None,
1061            cli_default_args: vec![],
1062            cli_env: HashMap::new(),
1063            cli_timeout_secs: None,
1064            cli_output_args: Vec::new(),
1065            cli_output_positional: HashMap::new(),
1066            upload_destinations: HashMap::new(),
1067            upload_default_destination: None,
1068            openapi_spec: None,
1069            openapi_include_tags: vec![],
1070            openapi_exclude_tags: vec![],
1071            openapi_include_operations: vec![],
1072            openapi_exclude_operations: vec![],
1073            openapi_max_operations: None,
1074            openapi_overrides: HashMap::new(),
1075            auth_generator: None,
1076            category: None,
1077            skills: vec![],
1078        };
1079
1080        let gen = AuthGenerator {
1081            gen_type: AuthGenType::Command,
1082            command: Some("echo".into()),
1083            args: vec!["${JWT_SUB}".into()],
1084            interpreter: None,
1085            script: None,
1086            cache_ttl_secs: 0,
1087            output_format: AuthOutputFormat::Text,
1088            env: HashMap::new(),
1089            inject: HashMap::new(),
1090            timeout_secs: 5,
1091        };
1092
1093        let ctx = GenContext {
1094            jwt_sub: "agent-42".into(),
1095            jwt_scope: "*".into(),
1096            tool_name: "brain:query".into(),
1097            timestamp: 1234567890,
1098            jwt_token: String::new(),
1099        };
1100        let keyring = Keyring::empty();
1101        let cache = AuthCache::new();
1102
1103        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
1104            .await
1105            .unwrap();
1106        assert_eq!(cred.value, "agent-42");
1107    }
1108
1109    #[tokio::test]
1110    async fn test_generate_timeout() {
1111        let provider = Provider {
1112            name: "test".into(),
1113            description: "test".into(),
1114            base_url: String::new(),
1115            auth_type: crate::core::manifest::AuthType::Bearer,
1116            auth_key_name: None,
1117            auth_header_name: None,
1118            auth_query_name: None,
1119            auth_value_prefix: None,
1120            extra_headers: HashMap::new(),
1121            oauth2_token_url: None,
1122            auth_secret_name: None,
1123            auth_session_token_env: None,
1124            mcp_url_env: None,
1125            oauth2_basic_auth: false,
1126            internal: false,
1127            handler: "http".into(),
1128            mcp_transport: None,
1129            mcp_command: None,
1130            mcp_args: vec![],
1131            mcp_url: None,
1132            mcp_env: HashMap::new(),
1133            cli_command: None,
1134            cli_default_args: vec![],
1135            cli_env: HashMap::new(),
1136            cli_timeout_secs: None,
1137            cli_output_args: Vec::new(),
1138            cli_output_positional: HashMap::new(),
1139            upload_destinations: HashMap::new(),
1140            upload_default_destination: None,
1141            openapi_spec: None,
1142            openapi_include_tags: vec![],
1143            openapi_exclude_tags: vec![],
1144            openapi_include_operations: vec![],
1145            openapi_exclude_operations: vec![],
1146            openapi_max_operations: None,
1147            openapi_overrides: HashMap::new(),
1148            auth_generator: None,
1149            category: None,
1150            skills: vec![],
1151        };
1152
1153        let gen = AuthGenerator {
1154            gen_type: AuthGenType::Command,
1155            command: Some("sleep".into()),
1156            args: vec!["10".into()],
1157            interpreter: None,
1158            script: None,
1159            cache_ttl_secs: 0,
1160            output_format: AuthOutputFormat::Text,
1161            env: HashMap::new(),
1162            inject: HashMap::new(),
1163            timeout_secs: 1,
1164        };
1165
1166        let ctx = GenContext::default();
1167        let keyring = Keyring::empty();
1168        let cache = AuthCache::new();
1169
1170        let err = generate(&provider, &gen, &ctx, &keyring, &cache)
1171            .await
1172            .unwrap_err();
1173        assert!(matches!(err, AuthGenError::Timeout(1)));
1174    }
1175}
ati/core/auth_generator.rs

ati/core/
auth_generator.rs
