Synced from monorepo
Changes: - Persist submitter identity for /feedback - Let custom models use rotating tokens from named auth providers - Template stale tool/param name literals in server-native descriptions - Minimal mode commits thinking in full, lookups as one-liners - Tighten durable append internals - Nudge model to end turn on no-op bash commands - Per-fetch signing nonce in the managed-config envelope, with a server-side replay probe - Include working tree in startup status
This commit is contained in:
grokkybara[bot]
parent
ba76b0a683
commit
a881e6703f
140 changed files with 6746 additions and 2377 deletions
603
crates/codegen/xai-grok-shell/src/auth/auth_provider.rs
Normal file
603
crates/codegen/xai-grok-shell/src/auth/auth_provider.rs
Normal file
|
|
@ -0,0 +1,603 @@
|
|||
//! Model auth providers (`[auth_provider.<name>]`).
|
||||
//!
|
||||
//! A model opts in with `auth_provider = "<name>"`; the named table declares a
|
||||
//! command that prints a fresh bearer token, which this module mints, caches,
|
||||
//! and rotates for that model's requests.
|
||||
//!
|
||||
//! The minted token stays in memory only ([`AUTH_PROVIDER_SLOTS`] and chat
|
||||
//! state, never `auth.json`); the command is a credential helper that owns its
|
||||
//! own durable storage and OAuth2 refresh. See "Where model auth providers fit
|
||||
//! (and don't)"
|
||||
//! in `docs/internal/AUTH.md`.
|
||||
//!
|
||||
//! This is distinct from the `AuthCredentialProvider` HTTP consumers in
|
||||
//! [`crate::auth::credential_provider`].
|
||||
|
||||
use super::token_output::{expiry_after_seconds, parse_token_output};
|
||||
|
||||
/// One named `[auth_provider.<name>]` table, honored only from the trusted
|
||||
/// config layers (`parse_auth_providers`). A new field here needs a
|
||||
/// `parse_auth_providers` warning decision.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct AuthProviderConfig {
|
||||
/// Command that prints a bearer token on stdout, bare or as JSON
|
||||
/// `{access_token, expires_in}`. Without `args` it runs via `sh -c`.
|
||||
pub command: String,
|
||||
/// Arguments for `command`. When present (even empty), the command runs
|
||||
/// directly with no shell; `command` is a program name on `PATH`, or a path.
|
||||
pub args: Option<Vec<String>>,
|
||||
/// Fallback token lifetime in seconds, used when the command's output
|
||||
/// carries no `expires_in`. Takes precedence over a JWT `exp` claim.
|
||||
pub token_ttl_secs: Option<u64>,
|
||||
/// Maximum seconds to wait for the command (default 30, clamped to 1..=600).
|
||||
/// A turn waits up to this long on a mint, so keep helpers fast and
|
||||
/// non-interactive.
|
||||
pub timeout_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl AuthProviderConfig {
|
||||
pub(crate) fn is_usable(&self) -> bool {
|
||||
!self.command.trim().is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// A model's reference to a named auth provider, built by `resolve_model_list`.
|
||||
#[derive(Clone, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(from = "AuthProviderRefData", into = "AuthProviderRefData")]
|
||||
pub struct AuthProviderRef {
|
||||
pub(crate) name: String,
|
||||
pub(crate) config: AuthProviderConfig,
|
||||
slot: ProviderSlot,
|
||||
/// `true` once the trusted table is attached. A ref revived from bytes is
|
||||
/// `false` and never mints or reads until [`AuthProviderRef::attach_trusted_config`]
|
||||
/// joins the shared slot for its name.
|
||||
resolved: bool,
|
||||
}
|
||||
|
||||
/// Serialized form: the name only, so persisted bytes never carry a command.
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
struct AuthProviderRefData {
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl From<AuthProviderRefData> for AuthProviderRef {
|
||||
fn from(data: AuthProviderRefData) -> Self {
|
||||
AuthProviderRef::unresolved(data.name)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AuthProviderRef> for AuthProviderRefData {
|
||||
fn from(provider: AuthProviderRef) -> Self {
|
||||
Self {
|
||||
name: provider.name,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AuthProviderRef {
|
||||
/// Production uses `unresolved` + `attach_trusted_config`.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn new(name: String, config: AuthProviderConfig) -> Self {
|
||||
let slot = provider_slot(&name);
|
||||
Self {
|
||||
name,
|
||||
config,
|
||||
slot,
|
||||
resolved: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// The in-memory form of a ref revived from bytes;
|
||||
/// [`AuthProviderRef::attach_trusted_config`] resolves it.
|
||||
pub(crate) fn unresolved(name: String) -> Self {
|
||||
Self {
|
||||
name,
|
||||
config: AuthProviderConfig::default(),
|
||||
slot: ProviderSlot::default(),
|
||||
resolved: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Re-attach the trusted config for this name at model resolution
|
||||
/// (`None` = the table was removed, leaving an unusable config). The ref
|
||||
/// becomes authoritative, joins the shared slot for its name, and may mint.
|
||||
pub(crate) fn attach_trusted_config(&mut self, config: Option<&AuthProviderConfig>) {
|
||||
self.config = config.cloned().unwrap_or_default();
|
||||
self.slot = provider_slot(&self.name);
|
||||
self.resolved = true;
|
||||
}
|
||||
}
|
||||
|
||||
/// Ignores the slot; a deserialized ref compares unequal until resolution
|
||||
/// re-attaches its config.
|
||||
impl PartialEq for AuthProviderRef {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.name == other.name && self.config == other.config
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for AuthProviderRef {}
|
||||
|
||||
impl std::fmt::Debug for AuthProviderRef {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AuthProviderRef")
|
||||
.field("name", &self.name)
|
||||
.field("config", &self.config)
|
||||
.field("resolved", &self.resolved)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
struct MintedProviderToken {
|
||||
token: String,
|
||||
/// Handed back to the command on the next run; never sent on the wire.
|
||||
refresh_token: Option<String>,
|
||||
/// Drives the 401 fresh-mint guard.
|
||||
minted_at: std::time::Instant,
|
||||
expires_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
/// The table version that minted the token; a different version reads as
|
||||
/// stale (see [`token_identity`]), so edits re-mint.
|
||||
minted_with: AuthProviderConfig,
|
||||
}
|
||||
|
||||
/// The async lock is held across the command run, single-flighting mints
|
||||
/// per provider name (shared across sessions). This dedupes concurrent
|
||||
/// successes; a persistently failing helper is retried per waiter, each bounded
|
||||
/// by the timeout clamp.
|
||||
type ProviderSlot = std::sync::Arc<tokio::sync::Mutex<Option<MintedProviderToken>>>;
|
||||
|
||||
/// Shared token slots, one per resolved provider name. Bounded by the configured
|
||||
/// provider names (only `attach_trusted_config` and test `new` insert), so no
|
||||
/// eviction.
|
||||
static AUTH_PROVIDER_SLOTS: std::sync::OnceLock<
|
||||
std::sync::Mutex<std::collections::HashMap<String, ProviderSlot>>,
|
||||
> = std::sync::OnceLock::new();
|
||||
|
||||
fn provider_slot(name: &str) -> ProviderSlot {
|
||||
let map = AUTH_PROVIDER_SLOTS.get_or_init(Default::default);
|
||||
let mut map = map.lock().unwrap_or_else(|e| e.into_inner());
|
||||
map.entry(name.to_owned()).or_default().clone()
|
||||
}
|
||||
|
||||
/// Pre-refresh margin: re-mint when the token expires within this window.
|
||||
pub(crate) const PROVIDER_TOKEN_EXPIRY_SKEW_SECS: u64 = 60;
|
||||
const PROVIDER_TOKEN_EXPIRY_SKEW: chrono::Duration =
|
||||
chrono::Duration::seconds(PROVIDER_TOKEN_EXPIRY_SKEW_SECS as i64);
|
||||
/// 401 fresh-mint guard: a token minted this recently is never re-minted on
|
||||
/// rejection. Same idea as the guard in `unauthorized_recovery`, with a shorter
|
||||
/// window because a provider mint is local and cheap.
|
||||
const PROVIDER_TOKEN_FRESH_MINT_GUARD: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
const DEFAULT_PROVIDER_TIMEOUT_SECS: u64 = 30;
|
||||
/// The effective mint timeout is clamped to `[1, this]`. A configured value
|
||||
/// outside the range is honored up to the bound and draws a parse warning,
|
||||
/// since a turn waits on the mint.
|
||||
pub(crate) const PROVIDER_TIMEOUT_CEILING_SECS: u64 = 600;
|
||||
/// Caps on the helper's captured output so a runaway command can't exhaust
|
||||
/// memory before the timeout fires. A bearer (even a large JWT) is far under
|
||||
/// the stdout cap; stderr only ever appears truncated in the failure log.
|
||||
const PROVIDER_STDOUT_CAP_BYTES: u64 = 1 << 20; // 1 MiB
|
||||
const PROVIDER_STDERR_CAP_BYTES: u64 = 64 << 10; // 64 KiB
|
||||
|
||||
/// The table fields that shape the minted token; a cached token minted under a
|
||||
/// different set reads as stale, so a config edit re-mints. Destructured so a
|
||||
/// new `AuthProviderConfig` field is a compile error until it is classified as
|
||||
/// token-shaping (add it here) or an execution knob like `timeout_secs`
|
||||
/// (editing it never invalidates).
|
||||
fn token_identity(config: &AuthProviderConfig) -> (&str, Option<&[String]>, Option<u64>) {
|
||||
let AuthProviderConfig {
|
||||
command,
|
||||
args,
|
||||
token_ttl_secs,
|
||||
timeout_secs: _,
|
||||
} = config;
|
||||
(command, args.as_deref(), *token_ttl_secs)
|
||||
}
|
||||
|
||||
fn minted_token_is_stale(minted: &MintedProviderToken, config: &AuthProviderConfig) -> bool {
|
||||
token_identity(&minted.minted_with) != token_identity(config)
|
||||
|| minted
|
||||
.expires_at
|
||||
.is_some_and(|at| chrono::Utc::now() + PROVIDER_TOKEN_EXPIRY_SKEW >= at)
|
||||
}
|
||||
|
||||
/// Log the missing-command warning once per provider, then at debug, so a
|
||||
/// misconfigured model doesn't warn on every turn.
|
||||
fn warn_empty_command(name: &str) {
|
||||
static WARNED: std::sync::OnceLock<std::sync::Mutex<std::collections::HashSet<String>>> =
|
||||
std::sync::OnceLock::new();
|
||||
let first = WARNED
|
||||
.get_or_init(Default::default)
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.insert(name.to_owned());
|
||||
const MSG: &str = "auth provider has no usable command: the [auth_provider.*] table is \
|
||||
missing from the trusted config layers, or its `command` is empty";
|
||||
if first {
|
||||
tracing::warn!(provider = %name, "{MSG}");
|
||||
} else {
|
||||
tracing::debug!(provider = %name, "{MSG}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Read up to `keep` bytes into `buf`, then drain and discard any remainder so
|
||||
/// the child never blocks on a full pipe. Memory stays bounded by `keep`.
|
||||
async fn read_capped<R>(reader: R, keep: u64, buf: &mut Vec<u8>) -> std::io::Result<()>
|
||||
where
|
||||
R: tokio::io::AsyncRead + Unpin,
|
||||
{
|
||||
use tokio::io::AsyncReadExt;
|
||||
let mut limited = reader.take(keep);
|
||||
limited.read_to_end(buf).await?;
|
||||
tokio::io::copy(&mut limited.into_inner(), &mut tokio::io::sink()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove every first-party credential from the helper's environment. BYOK
|
||||
/// isolates these keys on the wire, so the helper (the agent puts them in its
|
||||
/// own env at startup) must not inherit them.
|
||||
fn scrub_first_party_credentials(cmd: &mut tokio::process::Command) {
|
||||
for var in crate::agent::config::FIRST_PARTY_CREDENTIAL_ENV_VARS {
|
||||
cmd.env_remove(var);
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn `cmd`, capture stdout/stderr with a byte cap (reading both
|
||||
/// concurrently so a full pipe on one can't deadlock the other; a runaway helper
|
||||
/// is drained to a sink past the cap so it can't wedge the wait), and bound the
|
||||
/// whole run by `timeout`. Exceeding the stdout cap is an error.
|
||||
///
|
||||
/// On timeout the child's entire process group is killed. The helper is a group
|
||||
/// leader (`detach_command`'s `setsid`), so a compound `sh -c` helper's
|
||||
/// grandchildren -- and the `GROK_AUTH_PROVIDER_*` credentials in their env --
|
||||
/// do not outlive the reported timeout; `kill_on_drop` alone would reap only the
|
||||
/// direct child.
|
||||
async fn run_capped(
|
||||
cmd: &mut tokio::process::Command,
|
||||
timeout: std::time::Duration,
|
||||
) -> anyhow::Result<std::process::Output> {
|
||||
let mut child = cmd
|
||||
.spawn()
|
||||
.map_err(|e| anyhow::anyhow!("command failed to start: {e}"))?;
|
||||
// Enroll the child's process group so the timeout path can tear down the
|
||||
// whole tree. Best-effort: if enrollment fails, `kill_on_drop` still reaps
|
||||
// the direct child.
|
||||
let mut group = xai_grok_tools::util::ProcessGroup::new()
|
||||
.map_err(|e| anyhow::anyhow!("process group setup failed: {e}"))?;
|
||||
if let Err(e) = group.attach(&child) {
|
||||
tracing::debug!(error = %e, "auth provider: could not enroll helper process group");
|
||||
}
|
||||
let stdout = child.stdout.take().expect("stdout is piped");
|
||||
let stderr = child.stderr.take().expect("stderr is piped");
|
||||
let mut out_buf = Vec::new();
|
||||
let mut err_buf = Vec::new();
|
||||
|
||||
// One extra stdout byte so an over-cap write is detectable, not truncated.
|
||||
// The stderr read is advisory (it only feeds the failure log), so only
|
||||
// stdout governs the mint.
|
||||
let capture = async {
|
||||
let (out_res, err_res) = tokio::join!(
|
||||
read_capped(stdout, PROVIDER_STDOUT_CAP_BYTES + 1, &mut out_buf),
|
||||
read_capped(stderr, PROVIDER_STDERR_CAP_BYTES, &mut err_buf),
|
||||
);
|
||||
if let Err(e) = err_res {
|
||||
tracing::debug!(error = %e, "auth provider: stderr capture failed (advisory)");
|
||||
}
|
||||
out_res.map_err(|e| anyhow::anyhow!("reading command stdout: {e}"))?;
|
||||
child
|
||||
.wait()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("waiting on command: {e}"))
|
||||
};
|
||||
|
||||
let status = match tokio::time::timeout(timeout, capture).await {
|
||||
Ok(res) => res?,
|
||||
Err(_elapsed) => {
|
||||
let _ = group.kill();
|
||||
anyhow::bail!("command timed out after {}s", timeout.as_secs());
|
||||
}
|
||||
};
|
||||
if out_buf.len() as u64 > PROVIDER_STDOUT_CAP_BYTES {
|
||||
anyhow::bail!("command wrote more than {PROVIDER_STDOUT_CAP_BYTES} bytes to stdout");
|
||||
}
|
||||
Ok(std::process::Output {
|
||||
status,
|
||||
stdout: out_buf,
|
||||
stderr: err_buf,
|
||||
})
|
||||
}
|
||||
|
||||
async fn mint_provider_token(
|
||||
provider: &AuthProviderRef,
|
||||
mark_expired: bool,
|
||||
previous: Option<&MintedProviderToken>,
|
||||
) -> anyhow::Result<MintedProviderToken> {
|
||||
use std::process::Stdio;
|
||||
|
||||
let name = &provider.name;
|
||||
let config = &provider.config;
|
||||
// Clamp to [1, ceiling]: the slot lock is held across the run, so an
|
||||
// unbounded timeout would let one hung helper stall every turn sharing this
|
||||
// provider name. The ceiling is a hard bound, not just a parse warning.
|
||||
let timeout_secs = config
|
||||
.timeout_secs
|
||||
.unwrap_or(DEFAULT_PROVIDER_TIMEOUT_SECS)
|
||||
.clamp(1, PROVIDER_TIMEOUT_CEILING_SECS);
|
||||
tracing::info!(
|
||||
provider = %name,
|
||||
mark_expired,
|
||||
timeout_secs,
|
||||
"auth provider: running helper command"
|
||||
);
|
||||
|
||||
let mut cmd = match config.args {
|
||||
Some(ref args) => {
|
||||
// Direct exec: the program name is a PATH lookup, so trim stray
|
||||
// whitespace that would otherwise fail to resolve.
|
||||
let mut cmd = tokio::process::Command::new(config.command.trim());
|
||||
cmd.args(args);
|
||||
cmd
|
||||
}
|
||||
None => {
|
||||
let mut cmd = tokio::process::Command::new("sh");
|
||||
cmd.args(["-c", &config.command]);
|
||||
cmd
|
||||
}
|
||||
};
|
||||
cmd.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
// Capture stderr for the failure log; inheriting corrupts the TUI.
|
||||
.stderr(Stdio::piped())
|
||||
// Reaps the direct child if the future is dropped; `run_capped`
|
||||
// additionally kills the whole process group on timeout.
|
||||
.kill_on_drop(true);
|
||||
if mark_expired {
|
||||
cmd.env("GROK_AUTH_EXPIRED", "1");
|
||||
}
|
||||
// Git-credential-helper handback: give the command the last stored
|
||||
// credential so it can refresh instead of re-authenticating.
|
||||
if let Some(prev) = previous {
|
||||
cmd.env("GROK_AUTH_PROVIDER_ACCESS_TOKEN", &prev.token);
|
||||
if let Some(refresh) = &prev.refresh_token {
|
||||
cmd.env("GROK_AUTH_PROVIDER_REFRESH_TOKEN", refresh);
|
||||
}
|
||||
if let Some(expires_at) = prev.expires_at {
|
||||
cmd.env("GROK_AUTH_PROVIDER_EXPIRES_AT", expires_at.to_rfc3339());
|
||||
}
|
||||
}
|
||||
xai_grok_tools::util::detach_command(&mut cmd);
|
||||
cmd.envs(xai_grok_tools::util::pager_env());
|
||||
// Scrub last so nothing above can reintroduce a first-party credential.
|
||||
scrub_first_party_credentials(&mut cmd);
|
||||
|
||||
let output = run_capped(&mut cmd, std::time::Duration::from_secs(timeout_secs)).await?;
|
||||
|
||||
let parsed = match parse_token_output(&output) {
|
||||
Ok(parsed) => parsed,
|
||||
Err(e) => {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
anyhow::bail!(
|
||||
"{e} (stderr: {})",
|
||||
crate::util::truncate(stderr.trim(), 300)
|
||||
);
|
||||
}
|
||||
};
|
||||
let expires_at = parsed
|
||||
.expires_at
|
||||
.or_else(|| config.token_ttl_secs.and_then(expiry_after_seconds))
|
||||
.or_else(|| crate::auth::parse_jwt_expiration(&parsed.access_token));
|
||||
tracing::info!(
|
||||
provider = %name,
|
||||
mark_expired,
|
||||
expires_at = ?expires_at,
|
||||
"auth provider minted token"
|
||||
);
|
||||
Ok(MintedProviderToken {
|
||||
token: parsed.access_token,
|
||||
refresh_token: parsed.refresh_token,
|
||||
minted_at: std::time::Instant::now(),
|
||||
expires_at,
|
||||
minted_with: config.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
#[must_use = "a rotated token must be written to chat-state, or the wire keeps the stale key"]
|
||||
pub(crate) enum ProviderRefreshOutcome {
|
||||
/// `current_key` is already the fresh cached token; nothing to write.
|
||||
Unchanged,
|
||||
/// A token that should replace `current_key` on the wire.
|
||||
Rotated(String),
|
||||
/// The provider is unusable (unresolved or removed); already warned.
|
||||
Unusable,
|
||||
/// The mint ran and failed (logged).
|
||||
MintFailed,
|
||||
}
|
||||
|
||||
impl ProviderRefreshOutcome {
|
||||
pub(crate) fn rotated(self) -> Option<String> {
|
||||
match self {
|
||||
Self::Rotated(token) => Some(token),
|
||||
Self::Unchanged | Self::Unusable | Self::MintFailed => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AuthProviderRef {
|
||||
/// The slot, locked for a mutating operation. A removed provider drops
|
||||
/// its cached token and yields `None`, failing closed. An unresolved ref
|
||||
/// (revived from bytes) fails closed without touching the shared slot.
|
||||
async fn locked_slot(
|
||||
&self,
|
||||
) -> Option<tokio::sync::OwnedMutexGuard<Option<MintedProviderToken>>> {
|
||||
if !self.resolved {
|
||||
return None;
|
||||
}
|
||||
let mut slot = self.slot.clone().lock_owned().await;
|
||||
if !self.config.is_usable() {
|
||||
if slot.take().is_some() {
|
||||
tracing::warn!(
|
||||
provider = %self.name,
|
||||
"auth provider removed from config: dropping its cached token"
|
||||
);
|
||||
}
|
||||
warn_empty_command(&self.name);
|
||||
return None;
|
||||
}
|
||||
Some(slot)
|
||||
}
|
||||
|
||||
/// Cache-only read for sync resolution: never runs the command, blocks, or
|
||||
/// mutates. `None` for an unresolved ref, a cold or stale cache, or a mint
|
||||
/// in progress; minting happens pre-turn via [`AuthProviderRef::ensure_fresh_token`].
|
||||
pub(crate) fn cached_token(&self) -> Option<String> {
|
||||
if !self.resolved {
|
||||
return None;
|
||||
}
|
||||
if !self.config.is_usable() {
|
||||
warn_empty_command(&self.name);
|
||||
return None;
|
||||
}
|
||||
// A mint in progress holds the lock; treat it as a miss rather than
|
||||
// block the sync path.
|
||||
let Ok(guard) = self.slot.try_lock() else {
|
||||
tracing::debug!(provider = %self.name, "cache read skipped: mint in progress");
|
||||
return None;
|
||||
};
|
||||
guard
|
||||
.as_ref()
|
||||
.filter(|m| !minted_token_is_stale(m, &self.config))
|
||||
.map(|m| m.token.clone())
|
||||
}
|
||||
|
||||
/// The token that should replace `current_key` on the wire: serves the
|
||||
/// fresh cached token when chat-state lags behind a rotation, mints when
|
||||
/// the cache is cold or stale. Mints or rotates a bearer; unrelated to an
|
||||
/// OAuth refresh token.
|
||||
pub(crate) async fn ensure_fresh_token(
|
||||
&self,
|
||||
current_key: Option<&str>,
|
||||
) -> ProviderRefreshOutcome {
|
||||
let Some(mut slot) = self.locked_slot().await else {
|
||||
return ProviderRefreshOutcome::Unusable;
|
||||
};
|
||||
if let Some(ref minted) = *slot
|
||||
&& !minted_token_is_stale(minted, &self.config)
|
||||
{
|
||||
return if current_key == Some(minted.token.as_str()) {
|
||||
ProviderRefreshOutcome::Unchanged
|
||||
} else {
|
||||
ProviderRefreshOutcome::Rotated(minted.token.clone())
|
||||
};
|
||||
}
|
||||
let mark_expired = slot.is_some();
|
||||
let minted = match mint_provider_token(self, mark_expired, slot.as_ref()).await {
|
||||
Ok(minted) => minted,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
provider = %self.name,
|
||||
error = %e,
|
||||
"auth provider pre-turn mint failed"
|
||||
);
|
||||
return ProviderRefreshOutcome::MintFailed;
|
||||
}
|
||||
};
|
||||
let token = minted.token.clone();
|
||||
*slot = Some(minted);
|
||||
ProviderRefreshOutcome::Rotated(token)
|
||||
}
|
||||
|
||||
/// The replacement for a server-rejected `rejected_key` (chat-state's
|
||||
/// current key): a fresher cached token is adopted without a re-run,
|
||||
/// otherwise the command runs once. `None` for a token minted moments ago
|
||||
/// under the current table (the fresh-mint guard, which an edited table
|
||||
/// bypasses).
|
||||
pub(crate) async fn recover_rejected_token(&self, rejected_key: &str) -> Option<String> {
|
||||
let mut slot = self.locked_slot().await?;
|
||||
if let Some(ref minted) = *slot {
|
||||
if minted.token != rejected_key && !minted_token_is_stale(minted, &self.config) {
|
||||
return Some(minted.token.clone());
|
||||
}
|
||||
if minted.token == rejected_key
|
||||
&& token_identity(&minted.minted_with) == token_identity(&self.config)
|
||||
&& minted.minted_at.elapsed() < PROVIDER_TOKEN_FRESH_MINT_GUARD
|
||||
{
|
||||
tracing::warn!(
|
||||
provider = %self.name,
|
||||
"auth provider token rejected moments after mint: not \
|
||||
re-running (fresh-mint guard); surfacing the 401"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
tracing::info!(provider = %self.name, "auth provider token rejected: re-minting");
|
||||
let minted = match mint_provider_token(self, true, slot.as_ref()).await {
|
||||
Ok(minted) => minted,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
provider = %self.name,
|
||||
error = %e,
|
||||
"auth provider 401 re-mint failed"
|
||||
);
|
||||
// The server rejected the cached token and the re-mint failed;
|
||||
// mark it stale so it is not re-served next turn (fail closed).
|
||||
// The entry stays so its refresh token still feeds the next
|
||||
// handback attempt.
|
||||
if let Some(minted) = slot.as_mut() {
|
||||
minted.expires_at = Some(chrono::Utc::now());
|
||||
}
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let token = minted.token.clone();
|
||||
*slot = Some(minted);
|
||||
Some(token)
|
||||
}
|
||||
}
|
||||
|
||||
/// Backdate a provider's mint time past the fresh-mint guard.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn test_backdate_provider_mint(name: &str, age: std::time::Duration) {
|
||||
let slot = provider_slot(name);
|
||||
let mut slot = slot
|
||||
.try_lock()
|
||||
.expect("no mint in flight during test mutation");
|
||||
if let Some(ref mut minted) = *slot {
|
||||
minted.minted_at = std::time::Instant::now()
|
||||
.checked_sub(age)
|
||||
.expect("backdate before the process epoch");
|
||||
}
|
||||
}
|
||||
|
||||
/// A counting provider that prints "tok-1", "tok-2", ... on successive runs.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn test_counting_provider(name: &str, dir: &std::path::Path) -> AuthProviderRef {
|
||||
let counter = dir.join("count");
|
||||
AuthProviderRef::new(
|
||||
name.to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: format!(
|
||||
"echo run >> {c}; printf 'tok-%s' \"$(wc -l < {c} | tr -d ' ')\"",
|
||||
c = counter.display()
|
||||
),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn test_expire_provider_token(name: &str) {
|
||||
let slot = provider_slot(name);
|
||||
let mut slot = slot
|
||||
.try_lock()
|
||||
.expect("no mint in flight during test mutation");
|
||||
if let Some(ref mut minted) = *slot {
|
||||
minted.expires_at = Some(chrono::Utc::now() - chrono::Duration::seconds(1));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "auth_provider_tests.rs"]
|
||||
mod tests;
|
||||
763
crates/codegen/xai-grok-shell/src/auth/auth_provider_tests.rs
Normal file
763
crates/codegen/xai-grok-shell/src/auth/auth_provider_tests.rs
Normal file
|
|
@ -0,0 +1,763 @@
|
|||
// Slot names are process-global, so every test uses a unique name (no #[serial]
|
||||
// needed). No test mutates the process env: the scrub test sets its leak values
|
||||
// on the child command instead.
|
||||
|
||||
use super::test_counting_provider as counting_provider;
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_token_is_cached_while_fresh() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let provider = counting_provider("test-cache", dir.path());
|
||||
assert_eq!(
|
||||
provider.cached_token(),
|
||||
None,
|
||||
"cache-only read must miss on a cold cache without running the command"
|
||||
);
|
||||
let first = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
let second = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
assert_eq!(first, "tok-1");
|
||||
assert_eq!(second, "tok-1", "fresh token must be served from cache");
|
||||
assert_eq!(
|
||||
provider.cached_token().as_deref(),
|
||||
Some("tok-1"),
|
||||
"sync cache-only read must serve the warm cache"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_token_reminted_when_expired() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let provider = counting_provider("test-expiry", dir.path());
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await.rotated().unwrap(),
|
||||
"tok-1"
|
||||
);
|
||||
test_expire_provider_token("test-expiry");
|
||||
assert_eq!(
|
||||
provider.cached_token(),
|
||||
None,
|
||||
"cache-only read must not serve a stale token"
|
||||
);
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await.rotated().unwrap(),
|
||||
"tok-2",
|
||||
"expired token must be re-minted"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_pre_turn_refresh_semantics() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let provider = counting_provider("test-stale", dir.path());
|
||||
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(Some(&token)).await,
|
||||
ProviderRefreshOutcome::Unchanged,
|
||||
"fresh matching token must not be re-minted pre-turn"
|
||||
);
|
||||
assert_eq!(
|
||||
provider
|
||||
.ensure_fresh_token(Some("lagging-chat-state-key"))
|
||||
.await
|
||||
.rotated()
|
||||
.as_deref(),
|
||||
Some("tok-1"),
|
||||
"chat-state lagging behind a rotation adopts the fresh cached token"
|
||||
);
|
||||
test_expire_provider_token("test-stale");
|
||||
assert_eq!(
|
||||
provider
|
||||
.ensure_fresh_token(Some(&token))
|
||||
.await
|
||||
.rotated()
|
||||
.as_deref(),
|
||||
Some("tok-2"),
|
||||
"stale token must be re-minted pre-turn"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_401_recovery_has_fresh_mint_guard() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let provider = counting_provider("test-401", dir.path());
|
||||
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
provider.recover_rejected_token(&token).await,
|
||||
None,
|
||||
"a token minted moments ago must not be re-minted on 401 (loop guard)"
|
||||
);
|
||||
|
||||
test_backdate_provider_mint("test-401", std::time::Duration::from_secs(60));
|
||||
assert_eq!(
|
||||
provider.recover_rejected_token(&token).await.as_deref(),
|
||||
Some("tok-2"),
|
||||
"an aged rejected token is re-minted once"
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
provider.recover_rejected_token(&token).await.as_deref(),
|
||||
Some("tok-2"),
|
||||
"a rejection of the already-replaced key adopts the fresh token without a re-run"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: a warm cache must not outlive the provider's config.
|
||||
#[tokio::test]
|
||||
async fn provider_removed_from_config_drops_cached_token() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let provider = counting_provider("test-removed", dir.path());
|
||||
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
|
||||
let removed = AuthProviderRef::new("test-removed".to_owned(), AuthProviderConfig::default());
|
||||
assert_eq!(
|
||||
removed.cached_token(),
|
||||
None,
|
||||
"empty command must fail closed even with a warm cache"
|
||||
);
|
||||
assert_eq!(
|
||||
removed.ensure_fresh_token(Some(&token)).await,
|
||||
ProviderRefreshOutcome::Unusable
|
||||
);
|
||||
let restored = counting_provider("test-removed", dir.path());
|
||||
assert_eq!(
|
||||
restored
|
||||
.ensure_fresh_token(Some(&token))
|
||||
.await
|
||||
.rotated()
|
||||
.as_deref(),
|
||||
Some("tok-2"),
|
||||
"the removed provider's token must not survive in the slot"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_config_edit_invalidates_cached_token() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let old = counting_provider("test-freshen", dir.path());
|
||||
assert_eq!(
|
||||
old.ensure_fresh_token(None).await.rotated().unwrap(),
|
||||
"tok-1"
|
||||
);
|
||||
|
||||
let edited = AuthProviderRef::new(
|
||||
"test-freshen".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "printf edited-token".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
edited.cached_token(),
|
||||
None,
|
||||
"the unexpired old token must not be served under the edited table"
|
||||
);
|
||||
assert_eq!(
|
||||
edited
|
||||
.ensure_fresh_token(Some("tok-1"))
|
||||
.await
|
||||
.rotated()
|
||||
.as_deref(),
|
||||
Some("edited-token"),
|
||||
"refresh must run the edited command without waiting for expiry"
|
||||
);
|
||||
}
|
||||
|
||||
/// The fresh-mint guard applies per table version.
|
||||
#[tokio::test]
|
||||
async fn provider_401_recovery_reminted_under_edited_config() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let old = counting_provider("test-401-edited", dir.path());
|
||||
let token = old.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
|
||||
let edited = AuthProviderRef::new(
|
||||
"test-401-edited".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "printf new-config-token".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
edited.recover_rejected_token(&token).await.as_deref(),
|
||||
Some("new-config-token"),
|
||||
"recovery must run the edited command, not adopt the old-table token"
|
||||
);
|
||||
}
|
||||
|
||||
/// Editing only `timeout_secs` keeps the token; it is not part of
|
||||
/// `token_identity`.
|
||||
#[tokio::test]
|
||||
async fn provider_timeout_edit_does_not_invalidate_token() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let provider = counting_provider("test-timeout-edit", dir.path());
|
||||
provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
|
||||
let retimed = AuthProviderRef::new(
|
||||
"test-timeout-edit".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: provider.config.command.clone(),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: Some(5),
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
retimed.cached_token().as_deref(),
|
||||
Some("tok-1"),
|
||||
"a timeout-only edit must not invalidate the cached token"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn attach_trusted_config_lets_a_revived_ref_mint() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let template = counting_provider("test-attach", dir.path());
|
||||
let mut revived: AuthProviderRef = serde_json::from_str(r#"{"name": "test-attach"}"#).unwrap();
|
||||
assert_eq!(
|
||||
revived.ensure_fresh_token(None).await,
|
||||
ProviderRefreshOutcome::Unusable
|
||||
);
|
||||
revived.attach_trusted_config(Some(&template.config));
|
||||
assert_eq!(
|
||||
revived.ensure_fresh_token(None).await.rotated().as_deref(),
|
||||
Some("tok-1"),
|
||||
"a re-attached ref must be able to mint"
|
||||
);
|
||||
}
|
||||
|
||||
/// A ref revived from bytes never mutates the shared slot: a mutating
|
||||
/// call fails closed and leaves a resolved ref's token intact.
|
||||
#[tokio::test]
|
||||
async fn deserialized_ref_never_drops_the_shared_token() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let resolved = counting_provider("test-unresolved", dir.path());
|
||||
resolved.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
|
||||
let revived: AuthProviderRef = serde_json::from_str(r#"{"name": "test-unresolved"}"#).unwrap();
|
||||
assert_eq!(
|
||||
revived.ensure_fresh_token(None).await,
|
||||
ProviderRefreshOutcome::Unusable
|
||||
);
|
||||
assert_eq!(revived.recover_rejected_token("tok-1").await, None);
|
||||
assert_eq!(
|
||||
resolved.cached_token().as_deref(),
|
||||
Some("tok-1"),
|
||||
"the resolved ref's token must survive a mutating call on the stub"
|
||||
);
|
||||
}
|
||||
|
||||
/// A ref serializes to its name only: the revived ref carries no command
|
||||
/// and fails closed until re-attached, while the shared slot still serves
|
||||
/// resolved refs of the same name.
|
||||
#[tokio::test]
|
||||
async fn provider_ref_serializes_name_only_and_drops_config() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let provider = counting_provider("test-serde", dir.path());
|
||||
provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
|
||||
let bytes = serde_json::to_string(&provider).unwrap();
|
||||
assert!(bytes.contains("test-serde"));
|
||||
assert!(
|
||||
!bytes.contains("tok-%s") && !bytes.contains("command"),
|
||||
"the serialized form must carry the name only: {bytes}"
|
||||
);
|
||||
let revived: AuthProviderRef = serde_json::from_str(&bytes).unwrap();
|
||||
assert_eq!(revived.name, "test-serde");
|
||||
assert_eq!(
|
||||
revived.config,
|
||||
AuthProviderConfig::default(),
|
||||
"a serialized command must not survive deserialization"
|
||||
);
|
||||
assert_eq!(
|
||||
revived.cached_token(),
|
||||
None,
|
||||
"an unresolved ref fails closed"
|
||||
);
|
||||
let same_name = counting_provider("test-serde", dir.path());
|
||||
assert_eq!(
|
||||
same_name.cached_token().as_deref(),
|
||||
Some("tok-1"),
|
||||
"the shared slot still serves refs constructed with the real config"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_refresh_sets_expired_env() {
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-expired-env".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "printf 'tok-%s' \"${GROK_AUTH_EXPIRED:-0}\"".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await.rotated().as_deref(),
|
||||
Some("tok-0"),
|
||||
"first mint runs without GROK_AUTH_EXPIRED"
|
||||
);
|
||||
test_expire_provider_token("test-expired-env");
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await.rotated().as_deref(),
|
||||
Some("tok-1"),
|
||||
"re-mints run with GROK_AUTH_EXPIRED=1"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_concurrent_mints_single_flight() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let counter = dir.path().join("count");
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-single-flight".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: format!(
|
||||
"sleep 0.3; echo run >> {c}; printf 'tok-%s' \"$(wc -l < {c} | tr -d ' ')\"",
|
||||
c = counter.display()
|
||||
),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
let (a, b) = tokio::join!(
|
||||
provider.ensure_fresh_token(None),
|
||||
provider.ensure_fresh_token(None)
|
||||
);
|
||||
assert_eq!(a.rotated().as_deref(), Some("tok-1"));
|
||||
assert_eq!(
|
||||
b.rotated().as_deref(),
|
||||
Some("tok-1"),
|
||||
"second caller adopts, never re-runs"
|
||||
);
|
||||
let runs = std::fs::read_to_string(&counter).unwrap().lines().count();
|
||||
assert_eq!(runs, 1, "the command must run exactly once");
|
||||
}
|
||||
|
||||
/// Proven by staleness: an expiry inside the 60s skew re-mints, a
|
||||
/// distant one serves from cache.
|
||||
#[tokio::test]
|
||||
async fn provider_expiry_source_precedence() {
|
||||
fn short_jwt() -> String {
|
||||
// exp within the skew window: stale immediately if consumed.
|
||||
jwt_with_exp(chrono::Utc::now().timestamp() + 30)
|
||||
}
|
||||
fn long_jwt() -> String {
|
||||
jwt_with_exp(chrono::Utc::now().timestamp() + 7200)
|
||||
}
|
||||
fn jwt_with_exp(exp: i64) -> String {
|
||||
jsonwebtoken::encode(
|
||||
&jsonwebtoken::Header::default(),
|
||||
&serde_json::json!({ "exp": exp }),
|
||||
&jsonwebtoken::EncodingKey::from_secret(b"test"),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
async fn mints_after_first(
|
||||
name: &str,
|
||||
command: String,
|
||||
token_ttl_secs: Option<u64>,
|
||||
counter: &std::path::Path,
|
||||
) -> usize {
|
||||
let provider = AuthProviderRef::new(
|
||||
name.to_owned(),
|
||||
AuthProviderConfig {
|
||||
command,
|
||||
args: None,
|
||||
token_ttl_secs,
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
let first = provider
|
||||
.ensure_fresh_token(None)
|
||||
.await
|
||||
.rotated()
|
||||
.expect("first mint");
|
||||
let _ = provider.ensure_fresh_token(Some(&first)).await;
|
||||
std::fs::read_to_string(counter).unwrap().lines().count()
|
||||
}
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
// expires_in=10 (stale) wins over token_ttl_secs=3600 (fresh): re-mints.
|
||||
let c1 = dir.path().join("c1");
|
||||
let cmd1 = format!(
|
||||
"echo run >> {}; printf '{{\"access_token\":\"t1\",\"expires_in\":10}}'",
|
||||
c1.display()
|
||||
);
|
||||
assert_eq!(
|
||||
mints_after_first("test-exp-expires-in", cmd1, Some(3600), &c1).await,
|
||||
2,
|
||||
"expires_in must win over token_ttl_secs"
|
||||
);
|
||||
|
||||
// token_ttl_secs=1 (stale) wins over a 2h JWT exp (fresh): re-mints.
|
||||
let c2 = dir.path().join("c2");
|
||||
let cmd2 = format!("echo run >> {}; printf '{}'", c2.display(), long_jwt());
|
||||
assert_eq!(
|
||||
mints_after_first("test-exp-ttl", cmd2, Some(1), &c2).await,
|
||||
2,
|
||||
"token_ttl_secs must win over the JWT exp claim"
|
||||
);
|
||||
|
||||
// JWT exp alone: a near-expiry claim (inside the skew) re-mints,
|
||||
// proving the claim is consumed when nothing else is configured.
|
||||
let c3 = dir.path().join("c3");
|
||||
let cmd3 = format!("echo run >> {}; printf '{}'", c3.display(), short_jwt());
|
||||
assert_eq!(
|
||||
mints_after_first("test-exp-jwt", cmd3, None, &c3).await,
|
||||
2,
|
||||
"the JWT exp claim must apply when expires_in and token_ttl_secs are absent"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_unusable_expiry_still_mints() {
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-overflow".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: format!(
|
||||
"printf '{{\"access_token\":\"t\",\"expires_in\":{}}}'",
|
||||
u64::MAX
|
||||
),
|
||||
args: None,
|
||||
token_ttl_secs: Some(u64::MAX),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await.rotated().as_deref(),
|
||||
Some("t"),
|
||||
"an unusable expiry still mints; the token just has no expiry"
|
||||
);
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(Some("t")).await,
|
||||
ProviderRefreshOutcome::Unchanged,
|
||||
"no expiry source: never proactively re-minted"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_args_run_without_a_shell() {
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-args".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "printf".to_owned(),
|
||||
// Shell metacharacters stay literal under direct exec.
|
||||
args: Some(vec!["tok-$HOME;42".to_owned()]),
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await.rotated().as_deref(),
|
||||
Some("tok-$HOME;42"),
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_command_times_out() {
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-timeout".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "sleep 20; printf never".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: None,
|
||||
timeout_secs: Some(1),
|
||||
},
|
||||
);
|
||||
let start = std::time::Instant::now();
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await,
|
||||
ProviderRefreshOutcome::MintFailed
|
||||
);
|
||||
assert!(
|
||||
start.elapsed().as_secs() < 5,
|
||||
"1s timeout_secs must bound the mint (took {}s)",
|
||||
start.elapsed().as_secs()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_zero_timeout_clamps_to_one_second() {
|
||||
// `timeout_secs = 0` clamps up to the 1s floor, so an instant helper mints
|
||||
// rather than failing immediately.
|
||||
let fast = AuthProviderRef::new(
|
||||
"test-zero-timeout-fast".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "printf tok".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: Some(0),
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
fast.ensure_fresh_token(None).await.rotated().as_deref(),
|
||||
Some("tok")
|
||||
);
|
||||
|
||||
// ...and clamps down from the 30s default: a helper that runs past 1s times
|
||||
// out, proving the effective bound is the clamp, not the default.
|
||||
let slow = AuthProviderRef::new(
|
||||
"test-zero-timeout-slow".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "sleep 5; printf tok".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: Some(0),
|
||||
},
|
||||
);
|
||||
assert!(
|
||||
matches!(
|
||||
slow.ensure_fresh_token(None).await,
|
||||
ProviderRefreshOutcome::MintFailed
|
||||
),
|
||||
"a >1s helper under timeout_secs=0 must time out at the 1s clamp"
|
||||
);
|
||||
}
|
||||
|
||||
/// The distinct mint-failure modes (timeout, spawn failure, ran-but-no-token)
|
||||
/// surface distinct, greppable error messages so operators can triage them.
|
||||
#[tokio::test]
|
||||
async fn mint_error_messages_distinguish_failure_modes() {
|
||||
let timed_out = AuthProviderRef::new(
|
||||
"test-classify-timeout".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "sleep 20".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: None,
|
||||
timeout_secs: Some(1),
|
||||
},
|
||||
);
|
||||
let err = mint_provider_token(&timed_out, false, None)
|
||||
.await
|
||||
.err()
|
||||
.expect("timeout must fail the mint");
|
||||
assert!(err.to_string().contains("timed out"), "got: {err}");
|
||||
|
||||
let missing = AuthProviderRef::new(
|
||||
"test-classify-spawn".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "/nonexistent/provider-binary".to_owned(),
|
||||
args: Some(vec![]),
|
||||
token_ttl_secs: None,
|
||||
timeout_secs: Some(5),
|
||||
},
|
||||
);
|
||||
let err = mint_provider_token(&missing, false, None)
|
||||
.await
|
||||
.err()
|
||||
.expect("spawn failure must fail the mint");
|
||||
assert!(err.to_string().contains("failed to start"), "got: {err}");
|
||||
|
||||
let empty_output = AuthProviderRef::new(
|
||||
"test-classify-permanent".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "printf ''".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: None,
|
||||
timeout_secs: Some(5),
|
||||
},
|
||||
);
|
||||
let err = mint_provider_token(&empty_output, false, None)
|
||||
.await
|
||||
.err()
|
||||
.expect("empty output must fail the mint");
|
||||
assert!(err.to_string().contains("no output"), "got: {err}");
|
||||
}
|
||||
|
||||
/// On an in-session re-mint, the prior credential is handed back to the command
|
||||
/// via `GROK_AUTH_PROVIDER_*`, so a refresh-grant command can refresh instead of
|
||||
/// re-authenticating. Nothing is written to disk.
|
||||
#[tokio::test]
|
||||
async fn re_mint_hands_the_prior_token_back_to_the_command() {
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-handback".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: "printf 'seen-%s' \"${GROK_AUTH_PROVIDER_ACCESS_TOKEN:-none}\"".to_owned(),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
|
||||
let first = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
assert_eq!(first, "seen-none", "the first mint has no prior credential");
|
||||
test_expire_provider_token("test-handback");
|
||||
assert_eq!(
|
||||
provider
|
||||
.ensure_fresh_token(Some(&first))
|
||||
.await
|
||||
.rotated()
|
||||
.as_deref(),
|
||||
Some("seen-seen-none"),
|
||||
"the re-mint must receive the prior access token via env"
|
||||
);
|
||||
}
|
||||
|
||||
/// A 401 whose re-mint fails invalidates the rejected token, so it is not
|
||||
/// re-served next turn (fail closed) even while still locally unexpired.
|
||||
#[tokio::test]
|
||||
async fn failed_401_remint_invalidates_the_cached_token() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let counter = dir.path().join("count");
|
||||
// Mints tok-1 on the first run, then exits non-zero on every later run.
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-401-invalidate".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: format!(
|
||||
"echo run >> {c}; n=$(wc -l < {c} | tr -d ' '); \
|
||||
[ \"$n\" = 1 ] && printf 'tok-1' || exit 1",
|
||||
c = counter.display()
|
||||
),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
|
||||
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
assert_eq!(token, "tok-1");
|
||||
// Age past the fresh-mint guard so recovery attempts a re-mint.
|
||||
test_backdate_provider_mint("test-401-invalidate", PROVIDER_TOKEN_FRESH_MINT_GUARD * 2);
|
||||
|
||||
assert_eq!(
|
||||
provider.recover_rejected_token(&token).await,
|
||||
None,
|
||||
"a failed re-mint surfaces the 401"
|
||||
);
|
||||
assert_eq!(
|
||||
provider.cached_token(),
|
||||
None,
|
||||
"a rejected token whose re-mint failed must not be re-served"
|
||||
);
|
||||
}
|
||||
|
||||
/// A pre-turn re-mint that fails over a now-stale cached token leaves nothing
|
||||
/// servable: the stale token is never handed to the wire (mirror of the 401
|
||||
/// path, for the pre-turn path).
|
||||
#[tokio::test]
|
||||
async fn failed_pre_turn_mint_does_not_serve_the_stale_token() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let counter = dir.path().join("count");
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-pre-turn-stale".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: format!(
|
||||
"echo run >> {c}; n=$(wc -l < {c} | tr -d ' '); \
|
||||
[ \"$n\" = 1 ] && printf 'tok-1' || exit 1",
|
||||
c = counter.display()
|
||||
),
|
||||
args: None,
|
||||
token_ttl_secs: Some(3600),
|
||||
timeout_secs: None,
|
||||
},
|
||||
);
|
||||
|
||||
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
|
||||
assert_eq!(token, "tok-1");
|
||||
// Make the cached token stale so the next pre-turn call re-mints (and fails).
|
||||
test_expire_provider_token("test-pre-turn-stale");
|
||||
|
||||
assert!(matches!(
|
||||
provider.ensure_fresh_token(Some(token.as_str())).await,
|
||||
ProviderRefreshOutcome::MintFailed
|
||||
));
|
||||
assert_eq!(
|
||||
provider.cached_token(),
|
||||
None,
|
||||
"a stale token whose pre-turn re-mint failed must not be served"
|
||||
);
|
||||
}
|
||||
|
||||
/// A helper that writes past the stdout cap fails closed (permanent), so a
|
||||
/// runaway command can't exhaust memory or put a huge token on the wire.
|
||||
#[tokio::test]
|
||||
async fn provider_output_over_cap_fails_closed() {
|
||||
let over = PROVIDER_STDOUT_CAP_BYTES + 4096;
|
||||
let provider = AuthProviderRef::new(
|
||||
"test-stdout-cap".to_owned(),
|
||||
AuthProviderConfig {
|
||||
command: format!("head -c {over} /dev/zero"),
|
||||
args: None,
|
||||
token_ttl_secs: None,
|
||||
timeout_secs: Some(5),
|
||||
},
|
||||
);
|
||||
let err = mint_provider_token(&provider, false, None)
|
||||
.await
|
||||
.err()
|
||||
.expect("over-cap output must fail the mint");
|
||||
assert!(
|
||||
err.to_string().contains("more than"),
|
||||
"an over-cap write must be reported as such, got: {err}"
|
||||
);
|
||||
assert_eq!(
|
||||
provider.ensure_fresh_token(None).await,
|
||||
ProviderRefreshOutcome::MintFailed
|
||||
);
|
||||
}
|
||||
|
||||
/// Every first-party credential env var is scrubbed from the helper, so a BYOK
|
||||
/// helper never inherits the keys BYOK isolates on the wire.
|
||||
///
|
||||
/// The test drives its set/echo from an independent audited `EXPECTED` list, not
|
||||
/// from the scrub const, so it is not tautological: dropping an entry from
|
||||
/// `FIRST_PARTY_CREDENTIAL_ENV_VARS` alone leaves that var set on the command and
|
||||
/// trips the assert below, and removing one from both requires deliberately
|
||||
/// editing this audited list.
|
||||
///
|
||||
/// The leak values are set on the child command, not the process env, so the
|
||||
/// test is hermetic: it needs no `#[serial]` and cannot race a sibling test that
|
||||
/// reads a first-party credential (e.g. the `auth::manager` session tests).
|
||||
#[tokio::test]
|
||||
async fn provider_helper_env_scrubs_first_party_credentials() {
|
||||
// The credentials a BYOK helper must never inherit. Editing this list is the
|
||||
// audit checkpoint: it must equal the production scrub const.
|
||||
const EXPECTED: &[&str] = &[
|
||||
"XAI_API_KEY",
|
||||
"GROK_CODE_XAI_API_KEY",
|
||||
"GROK_AUTH",
|
||||
"GROK_AUTH_PATH",
|
||||
"GROK_DEPLOYMENT_KEY",
|
||||
"GROK_EXTRA_AUTH_KEY",
|
||||
"GROK_TRACE_UPLOAD_CREDENTIALS_FILE",
|
||||
"OTEL_EXPORTER_OTLP_HEADERS",
|
||||
"GROK_INTERNAL_OTLP_HEADERS",
|
||||
];
|
||||
assert_eq!(
|
||||
crate::agent::config::FIRST_PARTY_CREDENTIAL_ENV_VARS,
|
||||
EXPECTED,
|
||||
"the scrub list changed: re-audit that every entry is a first-party \
|
||||
credential a BYOK helper must not inherit, then update EXPECTED"
|
||||
);
|
||||
|
||||
// Echo each expected var back; the scrub must leave every one empty. A
|
||||
// scrub-const entry that EXPECTED still lists but production stopped removing
|
||||
// stays at its leak value and surfaces here.
|
||||
let echo = EXPECTED
|
||||
.iter()
|
||||
.map(|v| format!("${{{v}-}}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
let mut cmd = tokio::process::Command::new("sh");
|
||||
cmd.args(["-c", &format!("printf 'tok[%s]' \"{echo}\"")]);
|
||||
for var in EXPECTED {
|
||||
cmd.env(var, "first-party-leak");
|
||||
}
|
||||
super::scrub_first_party_credentials(&mut cmd);
|
||||
|
||||
let output = cmd.output().await.expect("helper spawns");
|
||||
assert_eq!(
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
"tok[]",
|
||||
"no first-party credential may survive into the helper env"
|
||||
);
|
||||
}
|
||||
|
|
@ -1,60 +1,11 @@
|
|||
use crate::auth::token_output::parse_token_output;
|
||||
use crate::auth::{AuthMode, GrokAuth};
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
pub(crate) struct ExternalAuthOutput {
|
||||
pub access_token: String,
|
||||
#[serde(default)]
|
||||
pub refresh_token: Option<String>,
|
||||
#[serde(default)]
|
||||
pub expires_in: Option<u64>,
|
||||
/// Token issuer. An xAI issuer marks the credential as first-party;
|
||||
/// see [`GrokAuth::is_xai_auth`].
|
||||
#[serde(default)]
|
||||
pub issuer: Option<String>,
|
||||
}
|
||||
|
||||
/// Parse process output (stdout) into a `GrokAuth`. Accepts bare token or JSON.
|
||||
/// Parse stdout into a session-credential `GrokAuth`.
|
||||
pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result<GrokAuth> {
|
||||
if !output.status.success() {
|
||||
anyhow::bail!("exited with {}", output.status);
|
||||
}
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
|
||||
if stdout.is_empty() {
|
||||
anyhow::bail!("produced no output on stdout");
|
||||
}
|
||||
|
||||
let (token, refresh_token, expires_at, issuer) =
|
||||
if let Ok(parsed) = serde_json::from_str::<ExternalAuthOutput>(&stdout) {
|
||||
tracing::debug!(
|
||||
has_refresh_token = parsed.refresh_token.is_some(),
|
||||
expires_in = ?parsed.expires_in,
|
||||
issuer = ?parsed.issuer,
|
||||
"auth: parsed external provider output as JSON"
|
||||
);
|
||||
let expires_at = parsed
|
||||
.expires_in
|
||||
.map(|secs| chrono::Utc::now() + chrono::Duration::seconds(secs as i64));
|
||||
let issuer = parsed
|
||||
.issuer
|
||||
.map(|i| i.trim().to_owned())
|
||||
.filter(|i| !i.is_empty());
|
||||
(
|
||||
parsed.access_token,
|
||||
parsed.refresh_token,
|
||||
expires_at,
|
||||
issuer,
|
||||
)
|
||||
} else {
|
||||
tracing::debug!(
|
||||
stdout_len = stdout.len(),
|
||||
"auth: treating output as bare token"
|
||||
);
|
||||
(stdout, None, None, None)
|
||||
};
|
||||
|
||||
let parsed = parse_token_output(output)?;
|
||||
Ok(GrokAuth {
|
||||
key: token,
|
||||
key: parsed.access_token,
|
||||
auth_mode: AuthMode::External,
|
||||
create_time: chrono::Utc::now(),
|
||||
user_id: String::new(),
|
||||
|
|
@ -74,20 +25,25 @@ pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result<Grok
|
|||
team_blocked_reasons: vec![],
|
||||
coding_data_retention_opt_out: crate::auth::default_coding_data_retention_opt_out(),
|
||||
has_grok_code_access: None,
|
||||
refresh_token,
|
||||
expires_at,
|
||||
oidc_issuer: issuer,
|
||||
refresh_token: parsed.refresh_token,
|
||||
expires_at: parsed.expires_at,
|
||||
oidc_issuer: parsed.issuer,
|
||||
oidc_client_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Sync version for mid-session refresh. 5s timeout for refresh, 60s for initial.
|
||||
pub(crate) fn run_external_auth_sync(command: &str, is_refresh: bool) -> Option<GrokAuth> {
|
||||
let timeout_secs = if is_refresh { 5 } else { 60 };
|
||||
run_auth_command(command, timeout_secs, is_refresh)
|
||||
}
|
||||
|
||||
/// Runs `command` via `sh -c`; `mark_expired` sets `GROK_AUTH_EXPIRED=1` so the
|
||||
/// helper can distinguish re-mints from first runs.
|
||||
fn run_auth_command(command: &str, timeout_secs: u64, mark_expired: bool) -> Option<GrokAuth> {
|
||||
use std::process::{Command, Stdio};
|
||||
|
||||
let timeout_secs = if is_refresh { 5 } else { 60 };
|
||||
|
||||
tracing::info!(cmd = %command, is_refresh, timeout_secs, "auth: running external auth provider (sync)");
|
||||
tracing::info!(cmd = %command, mark_expired, timeout_secs, "auth: running external auth provider (sync)");
|
||||
|
||||
let mut cmd = Command::new("sh");
|
||||
cmd.args(["-c", command])
|
||||
|
|
@ -95,7 +51,7 @@ pub(crate) fn run_external_auth_sync(command: &str, is_refresh: bool) -> Option<
|
|||
.stdout(Stdio::piped())
|
||||
// Pipe stderr — inherit would corrupt the TUI alternate screen.
|
||||
.stderr(Stdio::piped());
|
||||
if is_refresh {
|
||||
if mark_expired {
|
||||
cmd.env("GROK_AUTH_EXPIRED", "1");
|
||||
}
|
||||
xai_grok_tools::util::detach_std_command(&mut cmd);
|
||||
|
|
@ -224,14 +180,13 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn parse_output_malformed_json_falls_back_to_bare() {
|
||||
fn parse_output_json_shaped_but_invalid_is_err() {
|
||||
let output = std::process::Output {
|
||||
status: std::process::Command::new("true").status().unwrap(),
|
||||
stdout: b"{not valid json}".to_vec(),
|
||||
stderr: vec![],
|
||||
};
|
||||
let auth = parse_output(&output).unwrap();
|
||||
assert_eq!(auth.key, "{not valid json}");
|
||||
assert!(parse_output(&output).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
pub(crate) mod attribution;
|
||||
mod auth_provider;
|
||||
mod config;
|
||||
pub mod credential_provider;
|
||||
#[path = "devbox_login_stub.rs"]
|
||||
|
|
@ -15,7 +16,14 @@ pub(crate) mod recovery;
|
|||
pub(crate) mod refresh;
|
||||
pub(crate) mod single_flight;
|
||||
mod storage;
|
||||
mod token_output;
|
||||
pub(crate) mod token_type;
|
||||
pub use auth_provider::{AuthProviderConfig, AuthProviderRef};
|
||||
pub(crate) use auth_provider::{
|
||||
PROVIDER_TIMEOUT_CEILING_SECS, PROVIDER_TOKEN_EXPIRY_SKEW_SECS, ProviderRefreshOutcome,
|
||||
};
|
||||
#[cfg(test)]
|
||||
pub(crate) use auth_provider::{test_backdate_provider_mint, test_counting_provider};
|
||||
pub(crate) use config::LEGACY_AUTH_SCOPE;
|
||||
pub use config::{
|
||||
ForceLoginTeam, GrokComConfig, OAuth2ProviderConfig, OidcAuthConfig, PreferredAuthMethod,
|
||||
|
|
|
|||
155
crates/codegen/xai-grok-shell/src/auth/token_output.rs
Normal file
155
crates/codegen/xai-grok-shell/src/auth/token_output.rs
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
//! Shared parser for an auth command's stdout.
|
||||
//!
|
||||
//! Both auth paths run a command that prints a bearer token and parse it here:
|
||||
//! the session external-auth path ([`super::external_auth`]) and the per-model
|
||||
//! provider mint ([`super::auth_provider`]).
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
pub(crate) struct ExternalAuthOutput {
|
||||
pub access_token: String,
|
||||
#[serde(default)]
|
||||
pub refresh_token: Option<String>,
|
||||
#[serde(default)]
|
||||
pub expires_in: Option<u64>,
|
||||
/// An xAI issuer marks the credential as first-party
|
||||
/// (see [`crate::auth::GrokAuth::is_xai_auth`]).
|
||||
#[serde(default)]
|
||||
pub issuer: Option<String>,
|
||||
}
|
||||
|
||||
/// A bearer must be a single line: reject control characters (including an
|
||||
/// interior newline) so a malformed token can never be smuggled onto an HTTP
|
||||
/// header, rather than relying on the HTTP layer to reject it later.
|
||||
fn reject_control_chars(token: &str) -> anyhow::Result<()> {
|
||||
if token.contains(char::is_control) {
|
||||
anyhow::bail!("token contains control characters");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `now + secs`, or `None` on overflow.
|
||||
pub(crate) fn expiry_after_seconds(secs: u64) -> Option<chrono::DateTime<chrono::Utc>> {
|
||||
let secs = i64::try_from(secs).ok()?;
|
||||
chrono::Utc::now().checked_add_signed(chrono::Duration::try_seconds(secs)?)
|
||||
}
|
||||
|
||||
pub(crate) struct ParsedTokenOutput {
|
||||
pub access_token: String,
|
||||
pub refresh_token: Option<String>,
|
||||
pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
pub issuer: Option<String>,
|
||||
}
|
||||
|
||||
/// Accepts a bare token or JSON `{access_token, expires_in, issuer, ...}`. A
|
||||
/// non-zero exit, non-UTF-8 or empty stdout, an empty `access_token`, or
|
||||
/// JSON-object output that is not a valid token payload are all errors, so a
|
||||
/// malformed mint fails closed rather than putting garbage on the wire.
|
||||
pub(crate) fn parse_token_output(
|
||||
output: &std::process::Output,
|
||||
) -> anyhow::Result<ParsedTokenOutput> {
|
||||
if !output.status.success() {
|
||||
anyhow::bail!("exited with {}", output.status);
|
||||
}
|
||||
let stdout = std::str::from_utf8(&output.stdout)
|
||||
.map_err(|_| anyhow::anyhow!("produced non-UTF-8 output on stdout"))?
|
||||
.trim();
|
||||
if stdout.is_empty() {
|
||||
anyhow::bail!("produced no output on stdout");
|
||||
}
|
||||
|
||||
// Output that starts with `{` is meant to be a token payload: require it to
|
||||
// parse and carry a non-empty access_token. Anything else is a bare token
|
||||
// (JWTs and opaque tokens never start with `{`), so an error object like
|
||||
// `{"error":"expired"}` can never be mistaken for a bearer.
|
||||
if stdout.starts_with('{') {
|
||||
let parsed: ExternalAuthOutput = serde_json::from_str(stdout)
|
||||
.map_err(|e| anyhow::anyhow!("produced JSON that is not a token payload: {e}"))?;
|
||||
let access_token = parsed.access_token.trim().to_owned();
|
||||
if access_token.is_empty() {
|
||||
anyhow::bail!("produced JSON with an empty access_token");
|
||||
}
|
||||
reject_control_chars(&access_token)?;
|
||||
tracing::debug!(
|
||||
has_refresh_token = parsed.refresh_token.is_some(),
|
||||
expires_in = ?parsed.expires_in,
|
||||
issuer = ?parsed.issuer,
|
||||
"auth: parsed external provider output as JSON"
|
||||
);
|
||||
return Ok(ParsedTokenOutput {
|
||||
access_token,
|
||||
refresh_token: parsed.refresh_token,
|
||||
expires_at: parsed.expires_in.and_then(expiry_after_seconds),
|
||||
issuer: parsed
|
||||
.issuer
|
||||
.map(|i| i.trim().to_owned())
|
||||
.filter(|i| !i.is_empty()),
|
||||
});
|
||||
}
|
||||
|
||||
reject_control_chars(stdout)?;
|
||||
tracing::debug!(
|
||||
stdout_len = stdout.len(),
|
||||
"auth: treating output as bare token"
|
||||
);
|
||||
Ok(ParsedTokenOutput {
|
||||
access_token: stdout.to_owned(),
|
||||
refresh_token: None,
|
||||
expires_at: None,
|
||||
issuer: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn expiry_after_seconds_returns_none_on_overflow() {
|
||||
assert_eq!(expiry_after_seconds(u64::MAX), None);
|
||||
assert_eq!(expiry_after_seconds(u64::try_from(i64::MAX).unwrap()), None);
|
||||
assert!(expiry_after_seconds(3600).is_some());
|
||||
}
|
||||
|
||||
/// The provider path reads `refresh_token`, which the bare-token fallback
|
||||
/// cannot carry; only JSON output does.
|
||||
#[test]
|
||||
fn parse_token_output_reads_refresh_token_from_json_only() {
|
||||
let ok = |stdout: &str| std::process::Output {
|
||||
status: std::process::Command::new("true").status().unwrap(),
|
||||
stdout: stdout.as_bytes().to_vec(),
|
||||
stderr: vec![],
|
||||
};
|
||||
|
||||
let parsed =
|
||||
parse_token_output(&ok(r#"{"access_token":"a","refresh_token":"r"}"#)).unwrap();
|
||||
assert_eq!(parsed.access_token, "a");
|
||||
assert_eq!(parsed.refresh_token.as_deref(), Some("r"));
|
||||
|
||||
assert_eq!(parse_token_output(&ok("bare")).unwrap().refresh_token, None);
|
||||
}
|
||||
|
||||
/// JSON-shaped output must be a valid, non-empty token payload; a botched or
|
||||
/// error payload fails closed instead of going on the wire as a bearer.
|
||||
#[test]
|
||||
fn parse_token_output_rejects_invalid_json_payloads() {
|
||||
let ok = |stdout: &str| std::process::Output {
|
||||
status: std::process::Command::new("true").status().unwrap(),
|
||||
stdout: stdout.as_bytes().to_vec(),
|
||||
stderr: vec![],
|
||||
};
|
||||
|
||||
assert!(parse_token_output(&ok(r#"{"access_token":""}"#)).is_err());
|
||||
assert!(parse_token_output(&ok(r#"{"access_token":" "}"#)).is_err());
|
||||
assert!(parse_token_output(&ok(r#"{"error":"expired"}"#)).is_err());
|
||||
assert!(parse_token_output(&ok("{not valid json}")).is_err());
|
||||
|
||||
// A JSON payload's access_token is trimmed of surrounding whitespace.
|
||||
let parsed = parse_token_output(&ok("{\"access_token\":\" tok \"}")).unwrap();
|
||||
assert_eq!(parsed.access_token, "tok");
|
||||
|
||||
// An interior control character is rejected on both paths, so a
|
||||
// malformed token can never reach an HTTP header.
|
||||
assert!(parse_token_output(&ok("{\"access_token\":\"tok\\ninjected\"}")).is_err());
|
||||
assert!(parse_token_output(&ok("tok\ninjected")).is_err());
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue