//! Token sources for the LLM transport layer. //! //! [`TokenSource`] decouples request auth from `Config::api_key`: providers //! can supply a static string ([`StaticTokenSource`]) or a refreshable OAuth //! 2.0 PKCE engine ([`PkceOAuthTokenSource`]). Engines own their own cache //! and refresh logic; the [`Llm`] just asks for a bearer per request. //! //! The PKCE engine implements RFC 6749 + RFC 7636 with on-disk token //! caching keyed by `sha256(discovery_url|client_id|scopes)`. It's the //! same shape goose uses for Databricks, but we own the wire format and //! cache directory so the two are independently upgradable. //! //! First-use (cache empty) requires a browser: the engine opens //! `authorization_endpoint` in `webbrowser`, listens on `127.0.0.1:0`, //! captures the redirect, and exchanges the code for a token. Subsequent //! calls hit the cache and silently refresh when expired. use std::fs; use std::io::{self, Write}; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use async_trait::async_trait; use base64::Engine; use reqwest::Client; use serde::{Deserialize, Serialize}; use serde_json::Value; use sha2::Digest; use tokio::sync::Mutex; use crate::types::AgentError; /// Buffer before `expires_at` to consider a cached token "still good". /// Keeps us off the cliff if the clock or the server's clock drifts. const TOKEN_REFRESH_LEEWAY: Duration = Duration::from_secs(60); /// Wall-clock budget for the interactive browser dance. Goose uses 60s. /// We match: any longer and the user has gone to lunch. const BROWSER_AUTH_TIMEOUT: Duration = Duration::from_secs(60); /// Asynchronous source of a bearer token. The [`Llm`] calls this per /// request, so impls are expected to be cheap on the cache-hit path. #[async_trait] pub trait TokenSource: Send + Sync { async fn bearer(&self) -> Result; /// Return a bearer token from cache or refresh, **never** opening a browser. /// /// The default delegates to [`bearer`](Self::bearer) — correct for token /// sources (e.g. static API keys) that can never trigger a browser flow. /// [`PkceOAuthTokenSource`] overrides this to stop before the browser step. async fn bearer_no_browser(&self) -> Result { self.bearer().await } /// Force a fresh bearer after the server rejected the current one (401). /// /// `rejected` is the exact access token that just got the 401. Unlike /// [`bearer`](Self::bearer), which trusts the local expiry clock, this is /// driven by the server's verdict: the cached token looked valid to us /// (well within its local expiry) but the provider rejected it — clock /// skew, server-side revocation, or a node that never saw it. The clock /// therefore can't decide whether to refresh; the caller passes the /// rejected token so the impl can refresh unless a concurrent caller has /// *already* replaced it. Implementations must obtain a new token without /// any interactive step, so a headless harness never hangs. The default /// returns the existing bearer — correct for sources that can't refresh /// (a static key); the caller's retry then fails terminally rather than /// looping. async fn refresh_now(&self, _rejected: &str) -> Result { self.bearer().await } } /// A token that never changes for the life of the process. pub struct StaticTokenSource(String); impl StaticTokenSource { pub fn new(token: impl Into) -> Self { Self(token.into()) } } #[async_trait] impl TokenSource for StaticTokenSource { async fn bearer(&self) -> Result { Ok(self.0.clone()) } } /// Static config for an OAuth 2.0 Authorization Code + PKCE provider. /// /// The `discovery_url` must return a JSON document with at least /// `authorization_endpoint` and `token_endpoint` (RFC 8414). The /// `cache_namespace` is the directory under `~/.config/buzz-agent/oauth/` /// the token JSON lives in — separates providers' caches cleanly. #[derive(Debug, Clone)] pub struct PkceOAuthConfig { pub discovery_url: String, pub client_id: String, pub scopes: Vec, pub cache_namespace: String, /// When `Some`, the engine writes tokens here instead of /// `~/.config/buzz-agent/oauth//`. Production code /// leaves this `None`. Integration tests use it to avoid stomping on /// a shared `$HOME` when running in parallel. pub cache_dir_override: Option, } #[derive(Debug, Serialize, Deserialize, Clone)] struct CachedToken { access_token: String, refresh_token: Option, /// Unix seconds. `None` means the server didn't advertise an expiry; /// we use it without checking and rely on refresh on 401. expires_at: Option, } #[derive(Debug, Clone)] struct OidcEndpoints { authorization_endpoint: String, token_endpoint: String, } /// PKCE OAuth token source with on-disk refresh cache. /// /// First call: /// 1. Loads from cache if present and unexpired. /// 2. Otherwise tries `refresh_token` if cached. /// 3. Otherwise runs the full browser flow. /// /// Subsequent calls hit an in-memory copy of the cached token and only /// touch disk/network if the access token is past `expires_at`. pub struct PkceOAuthTokenSource { cfg: PkceOAuthConfig, http: Client, cache_path: PathBuf, /// Single-flight guard: only one refresh/browser flow at a time, even /// if many tool calls land concurrently. state: Mutex>, } impl PkceOAuthTokenSource { pub fn new(cfg: PkceOAuthConfig) -> Result, AgentError> { let cache_path = cache_path_for(&cfg)?; if let Some(parent) = cache_path.parent() { fs::create_dir_all(parent) .map_err(|e| AgentError::Llm(format!("oauth cache dir {parent:?}: {e}")))?; } let initial = read_cache(&cache_path); Ok(Arc::new(Self { cfg, http: Client::new(), cache_path, state: Mutex::new(initial), })) } /// Discover authorization + token endpoints from the well-known URL. async fn endpoints(&self) -> Result { let v: Value = self .http .get(&self.cfg.discovery_url) .send() .await .map_err(|e| AgentError::Llm(format!("oauth discovery: {e}")))? .error_for_status() .map_err(|e| AgentError::Llm(format!("oauth discovery status: {e}")))? .json() .await .map_err(|e| AgentError::Llm(format!("oauth discovery json: {e}")))?; let auth = v .get("authorization_endpoint") .and_then(Value::as_str) .ok_or_else(|| { AgentError::Llm("oauth discovery: authorization_endpoint missing".into()) })? .to_string(); let token = v .get("token_endpoint") .and_then(Value::as_str) .ok_or_else(|| AgentError::Llm("oauth discovery: token_endpoint missing".into()))? .to_string(); Ok(OidcEndpoints { authorization_endpoint: auth, token_endpoint: token, }) } /// Persist a token to disk and the in-memory cell. /// /// The cache holds both the access and refresh tokens, so the on-disk /// file is written owner-only (`0o600` on Unix) via an atomic /// inode-swapping rename — see [`write_private_cache`]. fn save(&self, state: &mut Option, token: CachedToken) -> Result<(), AgentError> { let body = serde_json::to_vec_pretty(&token) .map_err(|e| AgentError::Llm(format!("oauth cache serialize: {e}")))?; write_private_cache(&self.cache_path, &body).map_err(|e| { AgentError::Llm(format!("oauth cache write {:?}: {e}", self.cache_path)) })?; *state = Some(token); Ok(()) } /// Exchange a refresh token for a fresh access token. async fn refresh( &self, endpoints: &OidcEndpoints, refresh_token: &str, ) -> Result { let params = [ ("grant_type", "refresh_token"), ("refresh_token", refresh_token), ("client_id", &self.cfg.client_id), ]; let resp = self .http .post(&endpoints.token_endpoint) .form(¶ms) .send() .await .map_err(|e| AgentError::Llm(format!("oauth refresh: {e}")))?; if !resp.status().is_success() { let body = resp.text().await.unwrap_or_default(); return Err(AgentError::Llm(format!("oauth refresh failed: {body}"))); } let v: Value = resp .json() .await .map_err(|e| AgentError::Llm(format!("oauth refresh json: {e}")))?; token_from_response(&v, Some(refresh_token)) } /// Run the full browser-mediated Authorization Code + PKCE flow. /// Caller must hold a TTY/browser: this opens a window and blocks. pub async fn interactive_login(&self) -> Result<(), AgentError> { let endpoints = self.endpoints().await?; let token = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?; let mut state = self.state.lock().await; self.save(&mut state, token)?; Ok(()) } } #[async_trait] impl TokenSource for PkceOAuthTokenSource { async fn bearer(&self) -> Result { let mut state = self.state.lock().await; // 1. In-memory cache hit, still fresh. if let Some(tok) = state.as_ref() { if !is_expired(tok) { return Ok(tok.access_token.clone()); } } // 2. Re-read disk — another process may have refreshed already. if let Some(disk_tok) = read_cache(&self.cache_path) { if !is_expired(&disk_tok) { let bearer = disk_tok.access_token.clone(); *state = Some(disk_tok); return Ok(bearer); } } // 3. Try refresh if we have a refresh token. Discover endpoints once // here — deliberately hoisted above the refresh-token check so the // browser flow at step 5 (which also needs them) reuses this call. let endpoints = self.endpoints().await?; let refresh = state.as_ref().and_then(|t| t.refresh_token.clone()); if let Some(rt) = refresh { match self.refresh(&endpoints, &rt).await { Ok(fresh) => { let bearer = fresh.access_token.clone(); self.save(&mut state, fresh)?; return Ok(bearer); } Err(e) => { tracing::warn!(error = %e, "oauth refresh failed; falling back to browser flow"); } } // 4. Re-read disk after refresh failure — another process may have won the race. if let Some(disk_tok) = read_cache(&self.cache_path) { if !is_expired(&disk_tok) { let bearer = disk_tok.access_token.clone(); *state = Some(disk_tok); return Ok(bearer); } } } // 5. No usable cache: full browser dance. let fresh = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?; let bearer = fresh.access_token.clone(); self.save(&mut state, fresh)?; Ok(bearer) } async fn bearer_no_browser(&self) -> Result { self.try_bearer_no_browser().await } /// Force-refresh after a 401, never touching the browser flow. /// /// `rejected` is the access token the server just 401'd. Coalescing keys /// off token *identity*, not the expiry clock: a 401 means the token was /// rejected while it still looked locally fresh, so `is_expired()` would /// say "keep it" and no grant would ever run. Instead, under the lock we /// compare the current cached token to `rejected` — if they differ, a /// concurrent caller (this process or a sibling) already refreshed, so we /// return the new token without burning a second grant. If they still /// match, this is the rejected token and we run the refresh-token grant /// unconditionally. The whole check→refresh→save runs under one lock hold /// so concurrent callers serialize. On any failure the refresh token is /// preserved (never nulled) and the error is terminal `LlmAuth` — no /// browser, no hang. async fn refresh_now(&self, rejected: &str) -> Result { let mut state = self.state.lock().await; // 1. Coalesce by identity: if the cached token (in-memory, then disk) // is no longer the one the server rejected, someone already // refreshed it. Return that instead of grabbing another grant. if let Some(tok) = state.as_ref() { if tok.access_token != rejected { return Ok(tok.access_token.clone()); } } if let Some(disk_tok) = read_cache(&self.cache_path) { if disk_tok.access_token != rejected { let bearer = disk_tok.access_token.clone(); *state = Some(disk_tok); return Ok(bearer); } } // 2. The cached token is still the rejected one. Run the refresh-token // grant unconditionally — the expiry clock can't be trusted here, a // locally-fresh token is exactly what got 401'd. let refresh = state.as_ref().and_then(|t| t.refresh_token.clone()); let Some(rt) = refresh else { return Err(AgentError::LlmAuth( "token rejected and no refresh token available".into(), )); }; let endpoints = self.endpoints().await?; match self.refresh(&endpoints, &rt).await { Ok(fresh) => { let bearer = fresh.access_token.clone(); self.save(&mut state, fresh)?; Ok(bearer) } // 3. Refresh token is itself dead. Terminal — surfacing LlmAuth // stops the retry loop instead of falling to the browser flow, // which would hang a headless harness. Err(e) => Err(AgentError::LlmAuth(format!("token refresh failed: {e}"))), } } } impl PkceOAuthTokenSource { /// Return a bearer token from cache or refresh, **never** opening a browser. /// /// Follows the same steps as [`bearer`](TokenSource::bearer) but stops at /// step 4 — if no usable token is available after cache + refresh attempts, /// returns `Err(LlmAuth(...))` instead of launching the browser PKCE flow. /// Used by model-discovery paths that must not block on user interaction. pub(crate) async fn try_bearer_no_browser(&self) -> Result { let mut state = self.state.lock().await; // 1. In-memory cache hit, still fresh. if let Some(tok) = state.as_ref() { if !is_expired(tok) { return Ok(tok.access_token.clone()); } } // 2. Re-read disk — another process may have refreshed already. if let Some(disk_tok) = read_cache(&self.cache_path) { if !is_expired(&disk_tok) { let bearer = disk_tok.access_token.clone(); *state = Some(disk_tok); return Ok(bearer); } } // 3. Try refresh if we have a refresh token. Endpoints are discovered // lazily here — only when a refresh token is actually present — so // that an unreachable OIDC discovery URL cannot prevent the // no-token/no-cache path from returning `LlmAuth` (graceful // fallback) instead of `Llm` (hard error). let refresh = state.as_ref().and_then(|t| t.refresh_token.clone()); if let Some(rt) = refresh { let endpoints = self.endpoints().await?; match self.refresh(&endpoints, &rt).await { Ok(fresh) => { let bearer = fresh.access_token.clone(); self.save(&mut state, fresh)?; return Ok(bearer); } Err(e) => { tracing::warn!(error = %e, "oauth refresh failed during model discovery"); } } // 4. Re-read disk after refresh failure. if let Some(disk_tok) = read_cache(&self.cache_path) { if !is_expired(&disk_tok) { let bearer = disk_tok.access_token.clone(); *state = Some(disk_tok); return Ok(bearer); } } } // No usable token — return error instead of opening a browser. Err(AgentError::LlmAuth( "no cached Databricks token; run `buzz-agent auth databricks` first".into(), )) } } // ---- helpers ------------------------------------------------------------- /// Aborts a spawned task when dropped. Used to guarantee the localhost /// callback server doesn't outlive a failed/abandoned PKCE attempt. struct AbortOnDrop(tokio::task::JoinHandle<()>); impl Drop for AbortOnDrop { fn drop(&mut self) { self.0.abort(); } } fn is_expired(t: &CachedToken) -> bool { let Some(exp) = t.expires_at else { return false; }; let now = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs()) .unwrap_or(0); now + TOKEN_REFRESH_LEEWAY.as_secs() >= exp } fn cache_path_for(cfg: &PkceOAuthConfig) -> Result { let mut h = sha2::Sha256::new(); h.update(cfg.discovery_url.as_bytes()); h.update(b"|"); h.update(cfg.client_id.as_bytes()); h.update(b"|"); h.update(cfg.scopes.join(",").as_bytes()); let hash = hex::encode(h.finalize()); let dir = match &cfg.cache_dir_override { Some(p) => p.join(&cfg.cache_namespace), None => dirs::home_dir() .ok_or_else(|| AgentError::Llm("oauth cache: home directory not found".into()))? .join(".config") .join("buzz-agent") .join("oauth") .join(&cfg.cache_namespace), }; Ok(dir.join(format!("{hash}.json"))) } /// Load a cached token, enforcing the owner-only invariant on load. /// /// Owner-only permissions are a cache *lifecycle* invariant, not just a /// write-path property: a world-readable cache left by an older buzz-agent /// (or any tampering) must be tightened the moment we touch it, before the /// tokens are used — otherwise a file that never expires stays exposed until /// some future refresh happens to rewrite it. Every load path (initial and /// cross-process re-reads) funnels through here, so the repair covers them /// all. Returns `None` when the cache is absent, unreadable, unparseable, or /// cannot be secured; the caller then falls through to refresh/browser. fn read_cache(path: &Path) -> Option { let body = read_private_cache(path).ok()?; serde_json::from_slice(&body).ok() } /// Open the cache, reject symlinks, tighten loose permissions to `0o600`, and /// return its bytes. /// /// On Unix `O_NOFOLLOW` rejects a symlinked cache path at the kernel level /// (no stat/open TOCTOU), and `fchmod` on the already-open handle repairs a /// loose mode against the pinned inode rather than re-resolving the path. /// A cache that exists but cannot be secured is an error, so the caller fails /// closed instead of using an exposed file. #[cfg(unix)] fn read_private_cache(path: &Path) -> io::Result> { use std::io::Read; use std::os::unix::fs::{OpenOptionsExt, PermissionsExt}; let mut file = fs::OpenOptions::new() .read(true) .custom_flags(nix::libc::O_NOFOLLOW) .open(path)?; let meta = file.metadata()?; if !meta.file_type().is_file() { return Err(io::Error::new( io::ErrorKind::InvalidData, "oauth cache is not a regular file", )); } // Tighten in place on the open fd if any group/other bit is set. fchmod // targets the inode we already hold, so no attacker can swap the path // between the check and the repair. if meta.permissions().mode() & 0o077 != 0 { file.set_permissions(fs::Permissions::from_mode(0o600))?; } let mut body = Vec::new(); file.read_to_end(&mut body)?; Ok(body) } /// Non-Unix fallback: read the cache as-is. Owner-only enforcement is the /// Windows DACL work deferred behind the [`create_private_temp_file`] seam. #[cfg(not(unix))] fn read_private_cache(path: &Path) -> io::Result> { fs::read(path) } /// Removes a temp file on drop unless it was already renamed away. Keeps a /// failed/partial write from leaving a stray token file behind. struct TmpFileGuard<'a>(&'a Path); impl Drop for TmpFileGuard<'_> { fn drop(&mut self) { let _ = fs::remove_file(self.0); } } /// A per-write-unique temp suffix so concurrent savers — sibling threads or /// separate processes sharing `$HOME` — never collide on one temp path. /// Falls back to a timestamp if the RNG is unavailable rather than panicking /// mid-auth. fn unique_suffix() -> String { let mut bytes = [0u8; 8]; if getrandom::fill(&mut bytes).is_ok() { return hex::encode(bytes); } let nanos = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_nanos()) .unwrap_or(0); format!("{nanos:x}") } /// Write `body` to `path` as an owner-only file via an atomic rename. /// /// The cache holds both the refresh and access tokens, so it must never be /// readable by other users. We create a uniquely-named temp file in the same /// directory with owner-only protection at creation time — mode `0o600` on /// Unix (see [`create_private_temp_file`]) — so it is never briefly /// world/other readable, write and fsync it, then rename over the /// destination. The rename swaps the inode/entry wholesale, so a pre-existing /// cache file with loose permissions is *replaced* by the new private one; /// its old mode never survives. `fs::rename` maps to /// `MOVEFILE_REPLACE_EXISTING` on Windows, so the atomic replace holds on /// both platforms; the Windows owner-only DACL is pending the unsafe-FFI /// decision noted at the seam. fn write_private_cache(path: &Path, body: &[u8]) -> io::Result<()> { let parent = path.parent().ok_or_else(|| { io::Error::new( io::ErrorKind::InvalidInput, "oauth cache path has no parent directory", ) })?; fs::create_dir_all(parent)?; let file_name = path .file_name() .and_then(|n| n.to_str()) .unwrap_or("oauth-cache"); let tmp = parent.join(format!(".{file_name}.{}.tmp", unique_suffix())); let guard = TmpFileGuard(&tmp); let mut f = create_private_temp_file(&tmp)?; f.write_all(body)?; f.sync_all()?; drop(f); fs::rename(&tmp, path)?; // The rename consumed the temp path; nothing left to clean up. std::mem::forget(guard); Ok(()) } /// Create `tmp` for writing with owner-only permissions from the moment it /// exists. Fails if the file already exists (`create_new`), which the /// per-write-unique suffix makes effectively impossible. #[cfg(unix)] fn create_private_temp_file(tmp: &Path) -> io::Result { use std::os::unix::fs::OpenOptionsExt; fs::OpenOptions::new() .write(true) .create_new(true) .mode(0o600) .open(tmp) } /// Non-Unix fallback: create the temp file if it does not already exist. /// /// On Windows the owner-only equivalent is an explicit DACL set at creation /// (`CreateFileW` with SDDL `D:P(A;;FA;;;OW)`, matching goose's /// `private_file.rs`), but that FFI needs `unsafe`, which this crate forbids. /// Reconciling the two — an isolated helper crate, a vetted safe dependency, /// or descoping Windows — is an open decision escalated to the maintainer, so /// this interim relies on the default per-user ACLs and drops the owner-only /// implementation in behind this seam once the decision lands. `create_new` /// fails if the file already exists. #[cfg(not(unix))] fn create_private_temp_file(tmp: &Path) -> io::Result { fs::OpenOptions::new() .write(true) .create_new(true) .open(tmp) } /// Parse a token-endpoint JSON response. Fails loudly when `access_token` /// is missing or empty — without this, a malformed server response would /// be cached and `bearer()` would silently return `""` until the entry /// expires or is deleted by hand. fn token_from_response( v: &Value, fallback_refresh: Option<&str>, ) -> Result { let access_token = v .get("access_token") .and_then(Value::as_str) .filter(|s| !s.is_empty()) .ok_or_else(|| AgentError::Llm("oauth: token response missing/empty access_token".into()))? .to_string(); let refresh_token = v .get("refresh_token") .and_then(Value::as_str) .map(str::to_string) .or_else(|| fallback_refresh.map(str::to_string)); let expires_at = v.get("expires_in").and_then(Value::as_u64).map(|secs| { SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs()) .unwrap_or(0) + secs }); Ok(CachedToken { access_token, refresh_token, expires_at, }) } /// PKCE pieces: URL-safe random verifier (~64 chars) and its SHA-256 /// challenge (RFC 7636 §4.2). fn pkce_pair() -> Result<(String, String), AgentError> { let mut bytes = [0u8; 48]; getrandom::fill(&mut bytes).map_err(|e| AgentError::Llm(format!("pkce rng: {e}")))?; let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes); let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD .encode(sha2::Sha256::digest(verifier.as_bytes())); Ok((verifier, challenge)) } fn random_state() -> Result { let mut bytes = [0u8; 16]; getrandom::fill(&mut bytes).map_err(|e| AgentError::Llm(format!("state rng: {e}")))?; Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)) } /// Decide the OAuth callback result and the HTML page to serve. /// /// Returns `(result, page)`: `result` carries the auth code (or a detail /// string on failure) to the waiting flow via the oneshot channel; `page` is /// the *static* HTML shown in the browser. The page never embeds any request /// parameter — the `error` query value is attacker-influenceable, so /// reflecting it would be an XSS sink on the localhost callback. Failure /// detail travels only through `result`, which surfaces in the process error /// and logs, never in the served markup. fn callback_outcome( params: &std::collections::HashMap, expected_state: &str, ) -> (Result, String) { let result = match (params.get("code"), params.get("state")) { (Some(code), Some(st)) if st == expected_state => Ok(code.clone()), (Some(_), Some(_)) => Err("state mismatch".to_string()), _ => Err(params .get("error") .map(|e| sanitize_callback_detail(e)) .unwrap_or_else(|| "missing code".into())), }; let page = match result { Ok(_) => "

Buzz: signed in

You can close this window.

", Err(_) => "

Buzz auth failed

You can close this window and try again.

", } .to_string(); (result, page) } /// Neutralize an attacker-controllable OAuth `error` value before it enters /// an error string that later reaches the logs. Control characters (CR/LF in /// particular) enable log-line injection, and an unbounded value could flood /// the logs — replace control chars with spaces and cap the length. fn sanitize_callback_detail(raw: &str) -> String { const MAX: usize = 200; raw.chars() .map(|c| if c.is_control() { ' ' } else { c }) .take(MAX) .collect() } /// Spin up a localhost callback server, open the authorize URL in a /// browser, wait up to [`BROWSER_AUTH_TIMEOUT`] for the redirect, then /// exchange the code for a token. async fn browser_pkce_flow( http: &Client, cfg: &PkceOAuthConfig, endpoints: &OidcEndpoints, ) -> Result { use axum::{extract::Query, response::Html, routing::get, Router}; use std::collections::HashMap; use std::net::SocketAddr; use tokio::sync::oneshot; let (verifier, challenge) = pkce_pair()?; let state = random_state()?; let (tx, rx) = oneshot::channel::>(); let tx = Arc::new(Mutex::new(Some(tx))); let expected_state = state.clone(); let app = Router::new().route( "/", get(move |Query(params): Query>| { let tx = Arc::clone(&tx); let expected = expected_state.clone(); async move { let (result, page) = callback_outcome(¶ms, &expected); if let Some(sender) = tx.lock().await.take() { let _ = sender.send(result); } Html(page) } }), ); let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))) .await .map_err(|e| AgentError::Llm(format!("oauth callback bind: {e}")))?; let port = listener .local_addr() .map_err(|e| AgentError::Llm(format!("oauth callback addr: {e}")))? .port(); let redirect_uri = format!("http://localhost:{port}"); // `_server` is held until this function returns; the drop guard aborts // the axum task on every exit path (timeout, callback error, token // exchange failure, or success), so we never leak a listener bound to // 127.0.0.1 past the auth attempt. let _server = AbortOnDrop(tokio::spawn(async move { let _ = axum::serve(listener, app).await; })); let auth_url = format!( "{}?response_type=code&client_id={}&redirect_uri={}&scope={}&state={}&code_challenge={}&code_challenge_method=S256", endpoints.authorization_endpoint, urlencoding::encode(&cfg.client_id), urlencoding::encode(&redirect_uri), urlencoding::encode(&cfg.scopes.join(" ")), urlencoding::encode(&state), urlencoding::encode(&challenge), ); eprintln!("Opening browser for authentication. If it doesn't open, visit:\n {auth_url}"); let _ = webbrowser::open(&auth_url); let code = tokio::time::timeout(BROWSER_AUTH_TIMEOUT, rx) .await .map_err(|_| AgentError::Llm("oauth: browser auth timed out".into()))? .map_err(|_| AgentError::Llm("oauth: callback sender dropped".into()))? .map_err(|e| AgentError::Llm(format!("oauth callback: {e}")))?; // Exchange code for token. let params = [ ("grant_type", "authorization_code"), ("code", &code), ("redirect_uri", &redirect_uri), ("code_verifier", &verifier), ("client_id", &cfg.client_id), ]; let resp = http .post(&endpoints.token_endpoint) .form(¶ms) .send() .await .map_err(|e| AgentError::Llm(format!("oauth exchange: {e}")))?; if !resp.status().is_success() { let body = resp.text().await.unwrap_or_default(); return Err(AgentError::Llm(format!("oauth exchange failed: {body}"))); } let v: Value = resp .json() .await .map_err(|e| AgentError::Llm(format!("oauth exchange json: {e}")))?; token_from_response(&v, None) } #[cfg(test)] mod tests { use super::*; #[test] fn pkce_pair_produces_valid_challenge() { let (verifier, challenge) = pkce_pair().unwrap(); assert!(verifier.len() >= 43); let expected = base64::engine::general_purpose::URL_SAFE_NO_PAD .encode(sha2::Sha256::digest(verifier.as_bytes())); assert_eq!(expected, challenge); } #[test] fn cached_token_no_expiry_is_not_expired() { let t = CachedToken { access_token: "x".into(), refresh_token: None, expires_at: None, }; assert!(!is_expired(&t)); } #[test] fn cached_token_far_future_is_not_expired() { let future = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs() + 3600; let t = CachedToken { access_token: "x".into(), refresh_token: None, expires_at: Some(future), }; assert!(!is_expired(&t)); } #[test] fn cached_token_within_leeway_is_expired() { let near = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs() + 10; // 10s away, leeway is 60s → counts as expired let t = CachedToken { access_token: "x".into(), refresh_token: None, expires_at: Some(near), }; assert!(is_expired(&t)); } #[test] fn cache_path_uses_platform_home_directory() { let cfg = PkceOAuthConfig { discovery_url: "https://example.com/.well-known".into(), client_id: "abc".into(), scopes: vec!["a".into(), "b".into()], cache_namespace: "demo".into(), cache_dir_override: None, }; let p = cache_path_for(&cfg).unwrap(); let expected_dir = dirs::home_dir() .unwrap() .join(".config") .join("buzz-agent") .join("oauth") .join("demo"); assert_eq!(p.parent(), Some(expected_dir.as_path())); assert_eq!(p.extension().and_then(|s| s.to_str()), Some("json")); } #[test] fn token_from_response_uses_fallback_refresh() { let v: Value = serde_json::from_str(r#"{"access_token":"abc","expires_in":3600}"#).unwrap(); let t = token_from_response(&v, Some("old-refresh")).unwrap(); assert_eq!(t.access_token, "abc"); assert_eq!(t.refresh_token.as_deref(), Some("old-refresh")); assert!(t.expires_at.is_some()); } #[test] fn token_from_response_rejects_missing_access_token() { let v: Value = serde_json::from_str(r#"{"expires_in":3600}"#).unwrap(); assert!(token_from_response(&v, None).is_err()); } #[test] fn token_from_response_rejects_empty_access_token() { let v: Value = serde_json::from_str(r#"{"access_token":""}"#).unwrap(); assert!(token_from_response(&v, None).is_err()); } #[tokio::test] async fn test_bearer_reuses_disk_token_after_expiry() { let dir = tempfile::tempdir().unwrap(); let cfg = PkceOAuthConfig { discovery_url: "https://example.com/.well-known".into(), client_id: "test-client".into(), scopes: vec!["offline_access".into()], cache_namespace: "test".into(), cache_dir_override: Some(dir.path().to_path_buf()), }; let source = PkceOAuthTokenSource::new(cfg).unwrap(); // Expire the in-memory state. { let mut state = source.state.lock().await; *state = Some(CachedToken { access_token: "stale".into(), refresh_token: None, expires_at: Some(0), // long expired }); } // Write a valid token to disk (simulating another process refreshing). let future_exp = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs() + 7200; let fresh_token = CachedToken { access_token: "fresh-from-disk".into(), refresh_token: Some("rt".into()), expires_at: Some(future_exp), }; let body = serde_json::to_vec_pretty(&fresh_token).unwrap(); fs::write(&source.cache_path, &body).unwrap(); // bearer() should pick up the disk token without any network call. let result = source.bearer().await.unwrap(); assert_eq!(result, "fresh-from-disk"); } #[tokio::test] async fn test_bearer_falls_through_to_browser_when_disk_also_expired() { let dir = tempfile::tempdir().unwrap(); let cfg = PkceOAuthConfig { discovery_url: "https://example.com/.well-known".into(), client_id: "test-client".into(), scopes: vec!["offline_access".into()], cache_namespace: "test".into(), cache_dir_override: Some(dir.path().to_path_buf()), }; let source = PkceOAuthTokenSource::new(cfg).unwrap(); // Expire the in-memory state. { let mut state = source.state.lock().await; *state = Some(CachedToken { access_token: "stale".into(), refresh_token: None, expires_at: Some(0), }); } // Write an expired token to disk too. let expired_token = CachedToken { access_token: "also-stale".into(), refresh_token: None, expires_at: Some(0), }; let body = serde_json::to_vec_pretty(&expired_token).unwrap(); fs::write(&source.cache_path, &body).unwrap(); // bearer() should fall through past the disk check. // It will fail at the endpoints() discovery call since there's no server, // which proves it didn't short-circuit on the expired disk token. let result = source.bearer().await; assert!(result.is_err()); let err_msg = format!("{}", result.unwrap_err()); assert!( err_msg.contains("oauth discovery"), "expected discovery error, got: {err_msg}" ); } /// `try_bearer_no_browser` with an empty cache and no refresh token must /// return `LlmAuth` immediately — it must NOT attempt OIDC discovery even /// when the `discovery_url` is unreachable/invalid. This guards the /// regression where `endpoints()` was called unconditionally before the /// refresh-token check, causing an `Llm` error (hard failure) instead of /// the intended graceful `LlmAuth` fallback. #[tokio::test] async fn test_try_bearer_no_browser_empty_cache_no_refresh_returns_llm_auth_without_discovery() { let dir = tempfile::tempdir().unwrap(); // Intentionally invalid/unreachable discovery URL — if endpoints() is // called, the test will get an `Llm` error and the assertion below fails. let cfg = PkceOAuthConfig { discovery_url: "https://invalid.example.test/.well-known/oauth-authorization-server" .into(), client_id: "test-client".into(), scopes: vec!["offline_access".into()], cache_namespace: "test".into(), cache_dir_override: Some(dir.path().to_path_buf()), }; let source = PkceOAuthTokenSource::new(cfg).unwrap(); // Empty in-memory state (no token, no refresh token). { let mut state = source.state.lock().await; *state = None; } // No disk cache file either — dir is empty. let result = source.try_bearer_no_browser().await; assert!(result.is_err(), "expected Err, got Ok"); match result.unwrap_err() { AgentError::LlmAuth(_) => {} // correct: graceful fallback other => panic!( "expected LlmAuth (no discovery attempted), got: {other:?}\n\ This means endpoints() was called before the refresh-token check." ), } } // ---- callback HTML must never reflect input -------------------------- #[test] fn test_callback_failure_page_omits_reflected_error_param() { // A hostile `error` query value carrying markup must not appear in // the served HTML — otherwise the localhost callback is an XSS sink. let payload = ""; let mut params = std::collections::HashMap::new(); params.insert("error".to_string(), payload.to_string()); let (result, page) = callback_outcome(¶ms, "expected-state"); // The failure detail still reaches the waiting flow via `result`... assert_eq!(result.as_ref().err().map(String::as_str), Some(payload)); // ...but the browser page is static and inert. assert!( !page.contains(payload), "callback page reflected the raw error param: {page}" ); assert!( !page.contains("