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

ati/core/
auth_generator.rs
