ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

M0: compilable skeleton — Kigi 0.1.0 fork surgery

Browse Source 
Hard fork of xai-org/grok-build (Apache-2.0) re-targeted as Kigi, an
unofficial Kimi Code CLI community build.

Rename & identity
- 72 xai-*/xai-grok-* crates -> kigi-* (explicit: xai-grok-pager-bin ->
  kigi-bin [binary `kigi`], xai-grok-pager -> kigi-tui; rest mechanical);
  ptyctl, ptyctl-cli, third_party/ unchanged; proto package
  xai.grok.tools.v1 -> kigi.tools.v1
- Config home ~/.kigi (KIGI_SHARE_DIR override), env prefix GROK_* ->
  KIGI_*, `kigi --version` carries the unofficial-community-build notice
- clap identity, help text, startup banner, prompt templates rebranded
  (templates re-encrypted)

Deletions (PRD removal list #5/#6/#7/#9/#10)
- voice input (xai-grok-voice) and all TUI wiring
- telemetry: Mixpanel client, external OTel stream, Sentry, OTLP layers,
  trace/GCS/S3 upload queues (kigi-file-utils halved), workspace upload
  module & dc_log, heap-profile uploader, auth-diagnostics uploader,
  session-analytics halves of feedback; local zero-egress observability
  preserved in new kigi-log crate (unified log, --debug firehose,
  subsystem file logs, opt-in instrumentation)
- announcements (crate, remote-settings fields, TUI surfaces)
- plugin marketplace (crate, sources/browse/CTA/extensions-modal tab);
  direct plugin install/uninstall/update via kigi-agent git_install kept
- relay/gateway/assets endpoints and features (agent relay, headless
  relay transport, gateway bridge, LeaderEnvUrls); leader IPC socket now
  ~/.kigi/leader.sock + KIGI_LEADER_SOCKET, no ws-url derivation
- functional types rehomed instead of deleted: PermissionMode ->
  kigi-config-types, McpInitStrategy -> kigi-mcp, PrCreationSource ->
  session signals, TerminalDiagnostics -> kigi-pager-render, agent_id ->
  shell util

Endpoints
- kigi-env rewritten: single production KigiEndpoints {coding_api_base_url
  https://api.kimi.com/coding/v1 (KIGI_CODE_BASE_URL), oauth_host
  https://auth.kimi.com (KIGI_OAUTH_HOST), update_base_url (GitHub
  Releases API), upgrade_page_url}; GrokBuildEnvironment enum deleted

Toolchain & workspace hygiene
- Rust 1.97.0 pinned; edition 2024; full cargo update; git2 hoisted to
  workspace at 0.21 (Option->Result API migration), quick-xml 0.41
- Root Cargo.toml hand-maintained (PRD §8.1): version 0.1.0 inherited by
  all members, members sorted, unused deps pruned
- cargo-deny advisories gate (deny.toml with documented transitive
  exceptions); CI workflow (check/clippy/fmt/deny/test, macOS+Linux)
- cross-crate test seams re-gated behind `test-support` cargo feature;
  insta snapshot baselines renamed to the kigi_tui prefix
- clippy --workspace --all-targets: zero warnings; fmt clean

Fixes surfaced by the port
- updater probe/installer divergence (bin/kigi vs bin/grok symlink set)
- idle model-metadata refresh dead under KIGI_CODE_BASE_URL override
  (new is_effective_coding_endpoint_url, loopback+override aware)
- macOS symlinked-TMPDIR fixture canonicalization (foreign_sessions,
  fast-worktree); RSS measurement tests serialized via serial_test

Docs & legal (Apache §4)
- NOTICE added (upstream attribution + change statement); THIRD-PARTY
  notices sustained; kigi-tools ported-code notices extended; README,
  CONTRIBUTING, SECURITY, AGENTS.md rewritten

Out of scope for M0 (tracked): Kimi auth/inference (M1), search/fetch,
command parity, config import (M2), Computer Hub excision & final
brand-token sweep (M2), distribution & self-update rewrite (M3).
This commit is contained in:
ZacharyZhang-NY
2026-07-17 05:31:01 -04:00
commit d6c20fc13f
2612 changed files with 1353757 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/auth/attribution.rs
+937
View File
@@ -0,0 +1,937 @@
//! Shell-side 401-attribution helpers.
//!
//! Every 401 emit site in the shell joins the bearer the client
//! actually sent on the wire (the `Authorization` value for OAI-compat
//! backends, `x-api-key` for Anthropic Messages, the API proxy
//! `Authorization` header for storage / feedback / registry /
//! idle-resume) with the live
//! [`AuthManager::current_api_key`] value. The two sinks are:
//!
//! 1. [`kigi_log::unified_log::warn`] for the local
//! `~/.kigi/logs/unified.jsonl` file (best-effort; ships to GCS
//! only on OIDC refresh failure via `auth/refresh.rs`).
//! 2. A discrete `tracing::warn_span!("auth_401_attribution", ...)`
//! captured by the OTel layer in `util/otel_layer.rs` and shipped
//! via OTLP export to the configured telemetry backend
//! (queryable by span name `auth_401_attribution`).
//!
//! # Schema (every emit)
//!
//! ```text
//! {
//! "sent_key_prefix": "<last 12 chars of bearer the client sent, or """>,
//! "current_key_prefix": "<last 12 chars of AuthManager::current_api_key()>",
//! "mint_age_seconds": <i64; current time minus auth.create_time, or -1>,
//! "expires_at_seconds_from_now": <i64; auth.expires_at minus now,
//! or 0 when no current token>,
//! "consumer": "OaiCompatClient.<endpoint>" | "FeedbackClient.<op>"
//! | "FeedbackClient.<op>" | "SessionRegistryClient.<op>"
//! | "IdleResumeModelRefresh",
//! "is_stale_snapshot": <bool; true iff sent_prefix differs from a *known* current_prefix>
//! }
//! ```
//!
//! # Cross-crate plumbing
//!
//! [`kigi_sampler`] is intentionally decoupled from this crate. It
//! invokes the trait [`kigi_sampler::Auth401AttributionCallback`] at
//! its six 401 arms; this module provides [`ShellAttribution`], the
//! concrete impl that the shell wires into
//! [`kigi_sampler::SamplerConfig::attribution_callback`] at every
//! sampler-construction site. Non-sampler sites (storage / feedback /
//! registry / idle-resume) call [`record_consumer_401`]
//! directly with their `(consumer_kind, op)` pair.
use std::sync::Arc;
use kigi_sampler::{Auth401AttributionCallback, SamplingConsumer};
use kigi_tools::{Auth401AttributionCallback as ToolAuth401AttributionCallback, ToolConsumer};
use serde_json::Value as JsonValue;
use crate::auth::{AuthManager, TOKEN_TTL, token_suffix};
/// `cfg(test)`-only process-global counter that bumps on every
/// successful `record_auth_401` invocation.
///
/// Because the counter is process-global, every test that observes it
/// MUST be annotated with `#[serial_test::serial(attribution_emit_count)]`.
#[cfg(test)]
static EMIT_COUNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// Read the test-only emit counter.
#[cfg(test)]
pub(crate) fn test_emit_count() -> u64 {
EMIT_COUNT.load(std::sync::atomic::Ordering::SeqCst)
}
/// Reset the test-only emit counter to zero. Tests that span multiple
/// instrumented call sites should call this at setup so leftover bumps
/// from earlier tests in the same process do not pollute the assertion.
#[cfg(test)]
pub(crate) fn reset_test_emit_count() {
EMIT_COUNT.store(0, std::sync::atomic::Ordering::SeqCst);
}
/// Concrete implementation of [`Auth401AttributionCallback`] for the
/// sampler crate's six 401 arms.
///
/// One instance is constructed per `SamplerConfig` and cloned cheaply
/// (the struct holds an `Arc` and an `Option<String>`). The
/// `session_id` is captured at construction time and used for the
/// `unified_log::warn` `sid` field; non-session callers may pass
/// `None`.
pub(crate) struct ShellAttribution {
auth_manager: Arc<AuthManager>,
session_id: Option<String>,
}
// `AuthManager` does not implement `Debug` (it carries a `RwLock` over
// auth state and would expose secrets if it did). Hand-roll a redacted
// `Debug` impl so the `Auth401AttributionCallback` trait's
// `Debug + Send + Sync` bound is satisfied without changing
// `AuthManager`'s API surface.
impl std::fmt::Debug for ShellAttribution {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ShellAttribution")
.field("auth_manager", &"<redacted>")
.field("session_id", &self.session_id)
.finish()
}
}
impl ShellAttribution {
/// Construct a shareable attribution callback wired to the given
/// [`AuthManager`]. Returns `Arc<dyn Trait>` for the sampler
/// trait so callers can drop the value directly into
/// [`kigi_sampler::SamplerConfig::attribution_callback`].
///
/// (Returns `Arc<dyn Trait>` rather than `Self` because the
/// `kigi_sampler::SamplerConfig` field expects exactly that;
/// keeping the boundary in one place avoids `as Arc<dyn _>`
/// coercions at every call site.)
#[allow(clippy::new_ret_no_self)]
pub fn new(
auth_manager: Arc<AuthManager>,
session_id: Option<String>,
) -> Arc<dyn Auth401AttributionCallback> {
Arc::new(Self {
auth_manager,
session_id,
})
}
/// Tool-side counterpart of [`Self::new`]: returns
/// `Arc<dyn kigi_tools::Auth401AttributionCallback>` for the
/// `with_attribution_callback(...)` builder on each tool HTTP
/// client (`ImageGenClient`, `VideoGenClient`, `WebSearchClient`).
/// The two callbacks share the same underlying impl and emit the
/// same `auth_401_attribution` event format -- only the trait
/// signature differs (`SamplingConsumer` vs. `ToolConsumer`).
pub fn new_tool_callback(
auth_manager: Arc<AuthManager>,
session_id: Option<String>,
) -> Arc<dyn ToolAuth401AttributionCallback> {
Arc::new(Self {
auth_manager,
session_id,
})
}
}
impl Auth401AttributionCallback for ShellAttribution {
fn record_401(&self, consumer: SamplingConsumer, sent_bearer_prefix: Option<&str>) {
// The sampler crate has already truncated `sent_bearer_prefix`
// to `kigi_sampler::SENT_BEARER_PREFIX_LEN` characters
// before this trait method fires (see
// `SamplingClient::extract_sent_bearer`); the truncation
// inside `compute_attribution_payload` (via `token_suffix`)
// is therefore idempotent for this code path. The doubled
// truncation is intentional belt-and-suspenders -- the
// sampler-side scrub keeps the full bearer from ever leaving
// that crate, and the shell-side scrub keeps the local-log
// and OTel-span sinks aligned with the existing 12-char
// convention used by every other auth log line.
record_consumer_401(
self.auth_manager.as_ref(),
self.session_id.as_deref(),
ConsumerKind::OaiCompatClient,
consumer.as_endpoint(),
sent_bearer_prefix,
);
}
}
/// Tool-side hook: each tool client (image_gen, video_gen, web_search)
/// in `kigi-tools` emits a 401 attribution event through this
/// trait when its HTTP request returns UNAUTHORIZED. Same shape as
/// the sampler-side impl above; routes to the same pair of sinks.
///
/// `ToolConsumer::VideoGenStart` and `VideoGenPoll` collapse to the
/// same [`ConsumerKind::VideoGen`] with different op strings so the
/// gate query can break down video-gen 401s by phase.
impl ToolAuth401AttributionCallback for ShellAttribution {
fn record_401(&self, consumer: ToolConsumer, sent_bearer_prefix: Option<&str>) {
let (kind, op) = match consumer {
ToolConsumer::ImageGen => (ConsumerKind::ImageGen, ""),
ToolConsumer::VideoGenStart => (ConsumerKind::VideoGen, "start"),
ToolConsumer::VideoGenPoll => (ConsumerKind::VideoGen, "poll"),
ToolConsumer::WebSearch => (ConsumerKind::WebSearch, ""),
};
record_consumer_401(
self.auth_manager.as_ref(),
self.session_id.as_deref(),
kind,
op,
sent_bearer_prefix,
);
}
}
/// Categories of 401-attribution emit sites. Each variant maps to a
/// fixed prefix in the rendered `consumer` field; the per-site `op`
/// string is appended after a `.` separator (omitted for variants that
/// have no per-operation discriminator, e.g.
/// [`ConsumerKind::IdleResumeModelRefresh`]).
#[derive(Debug, Clone, Copy)]
pub(crate) enum ConsumerKind {
/// Sampler-side OpenAI-compat / Anthropic Messages emit. The op
/// string is the [`SamplingConsumer::as_endpoint`] return value.
OaiCompatClient,
/// Feedback collection sites in `agent/feedback_client.rs`.
FeedbackClient,
/// Session registry register/update sites in
/// `agent/session_registry_client.rs`.
SessionRegistryClient,
/// Idle-resume model-metadata refresh in
/// `session/acp_session.rs::maybe_refresh_model_metadata_on_resume`.
/// No per-op discriminator -- the consumer string is just
/// `"IdleResumeModelRefresh"`.
IdleResumeModelRefresh,
/// `kigi_tools::ToolConsumer::ImageGen` -- Imagine API
/// (`POST /images/generations`). No per-op discriminator;
/// consumer string is just `"ImageGen"`.
ImageGen,
/// `kigi_tools::ToolConsumer::VideoGenStart` and
/// `VideoGenPoll` -- Video Generation API. The op string is
/// `"start"` (`POST /videos/generations`) or `"poll"`
/// (`GET /videos/{request_id}`).
VideoGen,
/// `kigi_tools::ToolConsumer::WebSearch` -- web search via
/// `POST /responses` with a `WebSearch` tool. No per-op
/// discriminator; consumer string is just `"WebSearch"`.
WebSearch,
}
impl ConsumerKind {
/// Fixed prefix for the rendered `consumer` field.
fn prefix(self) -> &'static str {
match self {
Self::OaiCompatClient => "OaiCompatClient",
Self::FeedbackClient => "FeedbackClient",
Self::SessionRegistryClient => "SessionRegistryClient",
Self::IdleResumeModelRefresh => "IdleResumeModelRefresh",
Self::ImageGen => "ImageGen",
Self::VideoGen => "VideoGen",
Self::WebSearch => "WebSearch",
}
}
/// `true` for variants that take a per-operation discriminator
/// appended as `<prefix>.<op>`. `false` for variants whose
/// `consumer` string is just the prefix
/// (`IdleResumeModelRefresh`, `ImageGen`, `WebSearch` -- each is
/// a single endpoint with no sub-operation).
fn takes_op(self) -> bool {
!matches!(
self,
Self::IdleResumeModelRefresh | Self::ImageGen | Self::WebSearch
)
}
}
/// Format a `(kind, op)` pair into the design-doc `consumer` string.
fn format_consumer(kind: ConsumerKind, op: &str) -> String {
if kind.takes_op() {
format!("{}.{}", kind.prefix(), op)
} else {
kind.prefix().to_string()
}
}
/// Emit a single `auth 401 attribution` event for a per-consumer 401.
///
/// Wraps [`record_auth_401`] with the design-doc `consumer` formatting
/// (e.g., `"FeedbackClient.submit"`, `"VideoGen.start"`).
/// All 401 emit sites in `kigi-shell` go through this helper -- the
/// per-client `record_401_attribution` wrappers in
/// `agent/feedback_client.rs` and `agent/session_registry_client.rs` each
/// resolve their bearer and call this with the right `(kind, op)`.
///
/// `sent_bearer` may be either a full bearer (passed by the
/// non-sampler call sites listed above, which read directly from the
/// client's `user_token` / `deployment_key` snapshot) or a 12-char
/// prefix (passed by the sampler-side
/// [`Auth401AttributionCallback`] boundary; the sampler scrubs to a
/// prefix before crossing the crate boundary). The truncation inside
/// [`record_auth_401`] / `compute_attribution_payload` is idempotent
/// for the prefix case.
pub(crate) fn record_consumer_401(
auth_manager: &AuthManager,
session_id: Option<&str>,
kind: ConsumerKind,
op: &str,
sent_bearer: Option<&str>,
) {
let consumer = format_consumer(kind, op);
record_auth_401(auth_manager, session_id, &consumer, sent_bearer);
}
/// Emit a single `auth 401 attribution` event to both sinks (local
/// unified log file + OTel span for OTLP export).
///
/// Schema:
/// `(sent_key_prefix, current_key_prefix, mint_age_seconds,
/// expires_at_seconds_from_now, consumer, is_stale_snapshot)`.
///
/// `sent_bearer` is the bearer that was sent on the wire (the
/// `Authorization` value with `"Bearer "` already stripped, or the
/// `x-api-key` value for Anthropic Messages backends), OR a 12-char
/// prefix of same -- the sampler boundary always passes a prefix
/// here, the non-sampler shell sites pass full bearers and rely on
/// the [`compute_attribution_payload`] truncation. `None` is fine;
/// the prefix becomes the empty string.
///
/// `consumer` should be one of the canonical strings used by the
/// per-client wrappers, e.g. `"OaiCompatClient.chat_completions_stream"`,
/// `"FeedbackClient.submit"`, `"IdleResumeModelRefresh"`. Most call
/// sites should go through [`record_consumer_401`] which formats the
/// consumer string from a [`ConsumerKind`] for them.
pub(crate) fn record_auth_401(
auth_manager: &AuthManager,
session_id: Option<&str>,
consumer: &str,
sent_bearer: Option<&str>,
) {
let payload = compute_attribution_payload(auth_manager, consumer, sent_bearer);
// Sink 1 -- local file (~/.kigi/logs/unified.jsonl) + scrubbed
// tracing event. The local file is reliable but only ships to GCS
// on OIDC refresh failure (auth/refresh.rs::spawn_diagnostic_upload),
// so by itself it does not give visibility into the steady-state
// 401 population. Sink 2 below provides that.
kigi_log::unified_log::warn("auth 401 attribution", session_id, Some(payload.clone()));
// Sink 2 -- discrete OTel span exported via OTLP
// (util/otel_layer.rs). Auth 401 attribution schema fields below
// become OTel span attributes under `attributes.custom.<name>`
// per the tracing-opentelemetry bridge; query by span name
// `auth_401_attribution` in the configured telemetry backend.
//
// Wrapping in a `warn_span!` (vs. plain `tracing::warn!`) ensures
// emission even when no parent span is active. The OTel layer
// attaches plain events to the currently-entered span only, so a
// `tracing::warn!` from a `spawn_blocking` closure (idle-resume
// model refresh) or a background sync task is silently dropped.
// A `warn_span!` itself is always emitted by the layer's
// `on_new_span`/`on_close` hooks regardless of parent context.
//
// The span carries no body and is dropped immediately at the end
// of this function, so its `duration` is a few microseconds and
// it is logically a one-shot record (not a wrapping context for
// any other work).
let _attribution_span = tracing::warn_span!(
"auth_401_attribution",
// String fields. tracing flattens Option<&str> via Display, so
// we pre-collapse `None` to "" for both prefix fields and for
// session_id; downstream queries should treat "" as absent.
sent_key_prefix = payload["sent_key_prefix"].as_str().unwrap_or(""),
current_key_prefix = payload["current_key_prefix"].as_str().unwrap_or(""),
consumer = consumer,
session_id = session_id.unwrap_or(""),
// Numeric fields. The sentinel values from
// `compute_attribution_payload` (-1, 0) carry through
// unchanged.
mint_age_seconds = payload["mint_age_seconds"].as_i64().unwrap_or(-1),
expires_at_seconds_from_now = payload["expires_at_seconds_from_now"].as_i64().unwrap_or(0),
// Boolean -- the load-bearing field for stale-vs-live splits.
is_stale_snapshot = payload["is_stale_snapshot"].as_bool().unwrap_or(false),
)
.entered();
#[cfg(test)]
EMIT_COUNT.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
/// Pure (no I/O) computation of the attribution payload. Extracted
/// from [`record_auth_401`] so unit tests can assert each field
/// directly without reaching into `unified_log`'s file writer or the
/// tracing layer.
///
/// This function performs **exactly one** read-side acquisition of
/// [`AuthManager`]'s internal `RwLock` -- it calls
/// [`AuthManager::current`] once and derives both `current_key_prefix`
/// and the mint/expiry fields from the resulting `GrokAuth`.
///
/// `is_stale_snapshot` is `true` only when the live `current()` token
/// differs from the bearer the client sent. When `current()` returns
/// `None` (the manager has no active token), the result is `false`:
/// absence of a live token is "no evidence of staleness," not stale.
fn compute_attribution_payload(
auth_manager: &AuthManager,
consumer: &str,
sent_bearer: Option<&str>,
) -> JsonValue {
let now = chrono::Utc::now();
// Last-12-char suffix of the bearer the wire actually carried
// (see [`token_suffix`]: JWT headers share a common base64 prefix).
// `""` when the request had no bearer at all (distinct case from
// "had a bearer that turned out to be stale" -- the gate-criteria
// query can break down on this).
let sent_prefix = sent_bearer.map(token_suffix).unwrap_or("");
// Single read-lock acquisition: pull the live `GrokAuth` (or
// `None`) once and derive every other field from it.
let current_auth = auth_manager.current();
let current_prefix_owned: Option<String> = current_auth
.as_ref()
.map(|a| token_suffix(&a.key).to_string());
// None current means "no evidence of staleness," not stale --
// the downstream stale-vs-live split should only count
// true-positive staleness (sent bearer differs from a known live
// bearer).
let is_stale_snapshot = match current_prefix_owned.as_deref() {
Some(c) => sent_prefix != c,
None => false,
};
// Mint-age + expiry come from the same `current_auth` we already
// read; sentinels `-1 / 0` when the manager has no current token.
//
// TODO: mirror the full External-with-ttl branch from
// `AuthManager::is_token_expired` (uses
// `grok_com_config.auth_token_ttl` when `expires_at` is `None`
// and `auth_mode == External`). The current 2-branch fallback
// (`expires_at` if Some else `create_time + TOKEN_TTL`) is good
// enough for diagnostic metadata; the External-ttl branch is
// worth wiring once a real consumer needs it.
let (mint_age_seconds, expires_at_seconds_from_now) = match current_auth {
Some(auth) => {
let mint_age = now.signed_duration_since(auth.create_time).num_seconds();
let expiry = auth.expires_at.unwrap_or(auth.create_time + TOKEN_TTL);
(mint_age, expiry.signed_duration_since(now).num_seconds())
}
None => (-1_i64, 0_i64),
};
serde_json::json!({
"sent_key_prefix": sent_prefix,
"current_key_prefix": current_prefix_owned,
"mint_age_seconds": mint_age_seconds,
"expires_at_seconds_from_now": expires_at_seconds_from_now,
"consumer": consumer,
"is_stale_snapshot": is_stale_snapshot,
})
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use chrono::{Duration, Utc};
use crate::auth::{AuthManager, GrokAuth, GrokComConfig};
use super::*;
/// Test helper: build a fresh `AuthManager` rooted at a tempdir so
/// nothing from a developer's actual `~/.kigi/auth.json` leaks in.
fn empty_auth_manager() -> (tempfile::TempDir, AuthManager) {
let dir = tempfile::tempdir().expect("tempdir");
let cfg = GrokComConfig::default();
let am = AuthManager::new(dir.path(), cfg);
(dir, am)
}
fn fresh_auth(key: &str) -> GrokAuth {
GrokAuth {
key: key.to_string(),
create_time: Utc::now(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
}
}
fn payload_field<'a>(payload: &'a JsonValue, key: &str) -> &'a JsonValue {
payload
.get(key)
.unwrap_or_else(|| panic!("payload missing field {key:?}: {payload:?}"))
}
/// Live token sent + 401 with matching `current()` ->
/// `is_stale_snapshot` must be `false`. Also assert the auxiliary
/// fields are set sensibly (prefix, mint age, expiry).
#[test]
fn live_token_sent_is_not_stale() {
let (_dir, am) = empty_auth_manager();
let sent = "live-token-1234567890abcdef";
am.hot_swap(fresh_auth(sent));
let payload = compute_attribution_payload(&am, "Test.live", Some(sent));
assert_eq!(payload_field(&payload, "is_stale_snapshot"), false);
assert_eq!(payload_field(&payload, "consumer"), "Test.live");
// Last 12 chars (tail prefix for JWT-friendly diagnostics).
assert_eq!(payload_field(&payload, "sent_key_prefix"), "567890abcdef");
assert_eq!(
payload_field(&payload, "current_key_prefix"),
"567890abcdef"
);
// mint_age_seconds: should be small and non-negative for a
// freshly-created auth.
let mint = payload_field(&payload, "mint_age_seconds")
.as_i64()
.unwrap();
assert!(
(0..5).contains(&mint),
"mint_age_seconds should be 0-5 sec for a freshly-created auth, got {mint}"
);
// expires_at_seconds_from_now: should be just under 1 hour
// (3600s), with a tolerance for elapsed time during the test.
let expires = payload_field(&payload, "expires_at_seconds_from_now")
.as_i64()
.unwrap();
assert!(
(3590..=3600).contains(&expires),
"expires_at_seconds_from_now should be ~3600 for a 1h-expiry token, got {expires}"
);
}
/// Stale snapshot sent + 401 with a different (newer) `current()`
/// -> `is_stale_snapshot` must be `true`.
#[test]
fn stale_snapshot_is_detected() {
let (_dir, am) = empty_auth_manager();
let stale = "stale-token-1234567890";
let live = "live-token-different";
am.hot_swap(fresh_auth(live));
let payload = compute_attribution_payload(&am, "Test.stale", Some(stale));
assert_eq!(payload_field(&payload, "is_stale_snapshot"), true);
assert_eq!(payload_field(&payload, "sent_key_prefix"), "n-1234567890");
assert_eq!(
payload_field(&payload, "current_key_prefix"),
"en-different"
);
assert_eq!(payload_field(&payload, "consumer"), "Test.stale");
}
/// Live token sent + 401 with `current() == None` ->
/// `is_stale_snapshot` must be `false` (no evidence of staleness).
/// Sentinel `mint_age_seconds = -1`,
/// `expires_at_seconds_from_now = 0`. `current_key_prefix` is JSON
/// `null`.
#[test]
fn absent_current_is_not_stale() {
let (_dir, am) = empty_auth_manager();
// Do NOT inject anything -- manager has no current token.
let payload = compute_attribution_payload(&am, "Test.absent", Some("any-token"));
assert_eq!(payload_field(&payload, "is_stale_snapshot"), false);
assert_eq!(payload_field(&payload, "sent_key_prefix"), "any-token");
assert!(payload_field(&payload, "current_key_prefix").is_null());
assert_eq!(payload_field(&payload, "mint_age_seconds"), -1);
assert_eq!(payload_field(&payload, "expires_at_seconds_from_now"), 0);
}
/// Two-branch fallback: legacy token (no `expires_at`) uses
/// `create_time + TOKEN_TTL` as the expiry source. We assert the
/// computed `expires_at_seconds_from_now` reflects that.
#[test]
fn legacy_token_uses_two_branch_fallback() {
let (_dir, am) = empty_auth_manager();
let auth = GrokAuth {
key: "k".into(),
create_time: Utc::now() - Duration::seconds(60),
// No expires_at => falls through to create_time + TOKEN_TTL
// (= 30 days).
..GrokAuth::test_default()
};
am.hot_swap(auth);
let payload = compute_attribution_payload(&am, "Test.legacy", Some("k"));
// mint_age_seconds: ~60.
let mint = payload_field(&payload, "mint_age_seconds")
.as_i64()
.unwrap();
assert!(
(60..=70).contains(&mint),
"mint_age_seconds should be ~60 for a 60s-old auth, got {mint}"
);
// expires_at_seconds_from_now: TOKEN_TTL minus 60s = roughly
// 30 * 86400 - 60 = 2_591_940. Tolerate ~10s drift.
let expires = payload_field(&payload, "expires_at_seconds_from_now")
.as_i64()
.unwrap();
let expected = TOKEN_TTL.num_seconds() - 60;
assert!(
(expected - 10..=expected + 10).contains(&expires),
"expires_at_seconds_from_now should be ~{expected}, got {expires}"
);
}
/// `format_consumer` matrix:
/// - generic ops append "." + op (`OaiCompatClient.foo`)
/// - IdleResumeModelRefresh and tool variants drop the op
/// (their consumer string has no sub-op axis).
#[test]
fn format_consumer_matrix() {
let cases: &[(ConsumerKind, &str, &str)] = &[
(
ConsumerKind::OaiCompatClient,
"chat_completions_stream",
"OaiCompatClient.chat_completions_stream",
),
(
ConsumerKind::IdleResumeModelRefresh,
"",
"IdleResumeModelRefresh",
),
(
ConsumerKind::IdleResumeModelRefresh,
"ignored",
"IdleResumeModelRefresh",
),
(ConsumerKind::ImageGen, "", "ImageGen"),
(ConsumerKind::ImageGen, "ignored", "ImageGen"),
(ConsumerKind::VideoGen, "start", "VideoGen.start"),
(ConsumerKind::VideoGen, "poll", "VideoGen.poll"),
(ConsumerKind::WebSearch, "", "WebSearch"),
(ConsumerKind::WebSearch, "ignored", "WebSearch"),
];
for (kind, op, expected) in cases {
assert_eq!(
format_consumer(*kind, op),
*expected,
"kind={kind:?} op={op:?}"
);
}
}
/// `format_consumer` formats `OaiCompatClient.<endpoint>`
/// correctly and omits the `.` separator for
/// `IdleResumeModelRefresh`.
#[test]
fn format_consumer_with_op_appends_dot() {
assert_eq!(
format_consumer(ConsumerKind::OaiCompatClient, "chat_completions_stream"),
"OaiCompatClient.chat_completions_stream"
);
}
/// `ShellAttribution` implements `kigi_tools::Auth401AttributionCallback`
/// by routing each `ToolConsumer` variant to the right
/// `(ConsumerKind, op)` pair, which formats to the expected
/// `consumer` string in the emitted payload.
#[test]
#[serial_test::serial(attribution_emit_count)]
fn shell_attribution_tool_impl_routes_to_correct_consumer_strings() {
reset_test_emit_count();
let (_dir, am) = empty_auth_manager();
am.hot_swap(fresh_auth("bearer-1234567890"));
let am_arc = Arc::new(am);
let cb: Arc<dyn ToolAuth401AttributionCallback> =
ShellAttribution::new_tool_callback(am_arc.clone(), Some("sid-tool".into()));
let cases = [
(ToolConsumer::ImageGen, "ImageGen"),
(ToolConsumer::VideoGenStart, "VideoGen.start"),
(ToolConsumer::VideoGenPoll, "VideoGen.poll"),
(ToolConsumer::WebSearch, "WebSearch"),
];
for (consumer, expected_consumer_str) in cases {
cb.record_401(consumer, Some("bearer-1234567890"));
let payload = compute_attribution_payload(
am_arc.as_ref(),
expected_consumer_str,
Some("bearer-1234567890"),
);
assert_eq!(
payload_field(&payload, "consumer"),
expected_consumer_str,
"ToolConsumer::{consumer:?} should render as {expected_consumer_str:?}",
);
}
// Each variant bumped the global counter exactly once.
assert_eq!(test_emit_count() as usize, cases.len());
}
/// Capture `tracing::Span` `on_new_span` callbacks into a
/// `Mutex<Vec<CapturedSpan>>` so tests can assert the
/// `warn_span!("auth_401_attribution", ...)` emit fired with the
/// expected name and field values.
///
/// We intentionally only need `on_new_span` (which the
/// tracing-opentelemetry layer uses as its `OTel span_started`
/// hook). `on_close` is not asserted because the test cares about
/// "did the span exist with these attributes," not its duration.
mod span_capture {
use std::sync::Mutex;
use tracing::Subscriber;
use tracing::field::{Field, Visit};
use tracing::span::Attributes;
use tracing_subscriber::layer::{Context, Layer};
use tracing_subscriber::registry::LookupSpan;
#[derive(Debug, Default, Clone)]
pub struct CapturedSpan {
pub name: String,
pub fields_str: std::collections::BTreeMap<String, String>,
pub fields_i64: std::collections::BTreeMap<String, i64>,
pub fields_bool: std::collections::BTreeMap<String, bool>,
}
pub struct SpanCollector {
pub spans: std::sync::Arc<Mutex<Vec<CapturedSpan>>>,
}
impl SpanCollector {
pub fn new() -> (Self, std::sync::Arc<Mutex<Vec<CapturedSpan>>>) {
let buf = std::sync::Arc::new(Mutex::new(Vec::new()));
(Self { spans: buf.clone() }, buf)
}
}
impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for SpanCollector {
fn on_new_span(&self, attrs: &Attributes<'_>, _id: &tracing::Id, _ctx: Context<'_, S>) {
let mut captured = CapturedSpan {
name: attrs.metadata().name().to_string(),
..Default::default()
};
let mut visitor = FieldVisitor {
captured: &mut captured,
};
attrs.record(&mut visitor);
self.spans.lock().unwrap().push(captured);
}
}
struct FieldVisitor<'a> {
captured: &'a mut CapturedSpan,
}
impl<'a> Visit for FieldVisitor<'a> {
fn record_str(&mut self, field: &Field, value: &str) {
self.captured
.fields_str
.insert(field.name().to_string(), value.to_string());
}
fn record_i64(&mut self, field: &Field, value: i64) {
self.captured
.fields_i64
.insert(field.name().to_string(), value);
}
fn record_u64(&mut self, field: &Field, value: u64) {
self.captured
.fields_i64
.insert(field.name().to_string(), value as i64);
}
fn record_bool(&mut self, field: &Field, value: bool) {
self.captured
.fields_bool
.insert(field.name().to_string(), value);
}
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
self.captured
.fields_str
.insert(field.name().to_string(), format!("{value:?}"));
}
}
}
/// `record_auth_401` emits a discrete `warn_span!` with name
/// `"auth_401_attribution"` and the attribution fields as span
/// attributes. This is the span the tracing-opentelemetry bridge
/// ships via OTLP export to the configured telemetry backend.
/// Verifies field names, types, and values match the schema
/// documented at the top of this module.
#[test]
#[serial_test::serial(attribution_emit_count)]
fn record_auth_401_emits_otel_span_with_attribution_fields() {
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
let (collector, captured) = span_capture::SpanCollector::new();
let subscriber = tracing_subscriber::registry().with(collector);
let _guard = subscriber.set_default();
reset_test_emit_count();
let (_dir, am) = empty_auth_manager();
am.hot_swap(fresh_auth("live-token-1234567890"));
record_auth_401(
&am,
Some("sid-otel-span"),
"OaiCompatClient.chat_completions_stream",
Some("stale-snapshot-aaaaaa"),
);
let spans = captured.lock().unwrap();
let attribution = spans
.iter()
.find(|s| s.name == "auth_401_attribution")
.expect("expected one auth_401_attribution span; got: {spans:?}");
// String fields: prefixes truncated to 12 chars, consumer +
// session_id passed verbatim.
assert_eq!(
attribution
.fields_str
.get("sent_key_prefix")
.map(String::as_str),
Some("pshot-aaaaaa"),
"sent_key_prefix should be last 12 chars",
);
assert_eq!(
attribution
.fields_str
.get("current_key_prefix")
.map(String::as_str),
Some("n-1234567890"),
);
assert_eq!(
attribution.fields_str.get("consumer").map(String::as_str),
Some("OaiCompatClient.chat_completions_stream"),
);
assert_eq!(
attribution.fields_str.get("session_id").map(String::as_str),
Some("sid-otel-span"),
);
// Boolean: the load-bearing field for stale-vs-live splits.
// `true` because `sent != current`.
assert_eq!(
attribution.fields_bool.get("is_stale_snapshot"),
Some(&true),
);
// Numeric: mint_age in [0, 5) for a freshly-injected auth;
// expires_at ~3600s away.
let mint = attribution
.fields_i64
.get("mint_age_seconds")
.copied()
.unwrap();
assert!(
(0..5).contains(&mint),
"mint_age_seconds should be 0-5, got {mint}",
);
let expires = attribution
.fields_i64
.get("expires_at_seconds_from_now")
.copied()
.unwrap();
assert!(
(3590..=3600).contains(&expires),
"expires_at_seconds_from_now should be ~3600, got {expires}",
);
}
/// `record_auth_401` (the I/O-bearing wrapper) bumps the
/// `cfg(test)` counter so cross-module tests can observe how many
/// times an attribution event was actually emitted.
///
/// `#[serial]` because `EMIT_COUNT` is process-global; concurrent
/// tests that exercise the counter would race each other.
#[test]
#[serial_test::serial(attribution_emit_count)]
fn record_auth_401_bumps_emit_counter() {
reset_test_emit_count();
let (_dir, am) = empty_auth_manager();
am.hot_swap(fresh_auth("k"));
record_auth_401(&am, None, "Test.counter", Some("k"));
assert_eq!(test_emit_count(), 1);
record_auth_401(&am, None, "Test.counter", Some("k"));
assert_eq!(test_emit_count(), 2);
}
/// The SubagentSpawnContext-borne callback flows through
/// `read_parent_sampling_config` into the inherited
/// `SamplerConfig.attribution_callback`. We can't drive the full
/// subagent path here (requires SessionActor + chat-state
/// scaffolding), but we can assert the structural property: the
/// callback the parent constructs is the one any later
/// `SamplerConfig` clone carries forward unchanged.
#[test]
#[serial_test::serial(attribution_emit_count)]
fn parent_callback_flows_through_arc_clone() {
reset_test_emit_count();
let (_dir, am) = empty_auth_manager();
let am_arc = Arc::new(am);
let parent_cb = ShellAttribution::new(am_arc.clone(), Some("parent-sid".into()));
// Simulate the inheritance hand-off: the parent callback flows
// through SessionHandle -> SubagentSpawnContext ->
// SamplerConfig.attribution_callback as plain Arc clones.
let inherited_cb = parent_cb.clone();
// Drive the inherited callback. The `record_401` should bump
// the same global counter the parent callback would, proving
// they refer to the same underlying impl.
inherited_cb.record_401(SamplingConsumer::ChatCompletionsStream, Some("bearer"));
assert_eq!(test_emit_count(), 1);
// Sanity: the parent_cb still works too (it's the same Arc).
parent_cb.record_401(SamplingConsumer::Messages, Some("bearer"));
assert_eq!(test_emit_count(), 2);
}
/// End-to-end: the trait impl wraps `consumer.as_endpoint()` in
/// `"OaiCompatClient.<endpoint>"` and delegates to
/// `record_consumer_401` for every variant of `SamplingConsumer`.
/// We assert one bump per variant via the test counter, plus the
/// rendered `consumer` string for one variant via a payload
/// recompute (the trait does not return the payload, so we
/// recompute directly from the same inputs).
#[test]
#[serial_test::serial(attribution_emit_count)]
fn shell_attribution_trait_impl_routes_through_helper() {
reset_test_emit_count();
let (_dir, am) = empty_auth_manager();
let am_arc = Arc::new(am);
let cb = ShellAttribution::new(am_arc.clone(), Some("sid-shell".into()));
let variants = [
SamplingConsumer::ChatCompletionsStream,
SamplingConsumer::ChatCompletions,
SamplingConsumer::ResponsesStream,
SamplingConsumer::Responses,
SamplingConsumer::MessagesStream,
SamplingConsumer::Messages,
];
for consumer in variants {
cb.record_401(consumer, Some("test-bearer"));
}
assert_eq!(test_emit_count() as usize, variants.len());
// Sanity-check the consumer-string formatting via direct
// payload computation.
let payload = compute_attribution_payload(
am_arc.as_ref(),
&format_consumer(
ConsumerKind::OaiCompatClient,
SamplingConsumer::MessagesStream.as_endpoint(),
),
Some("test-bearer"),
);
assert_eq!(
payload_field(&payload, "consumer"),
"OaiCompatClient.messages_stream"
);
}
}
crates/codegen/kigi-shell/src/auth/config.rs
+423
View File
@@ -0,0 +1,423 @@
use super::model::TEAM_PRINCIPAL_TYPE;
use serde::{Deserialize, Serialize};
// Transitional: the M1 auth rewrite (Kimi device flow) replaces this origin.
const AUTH_ORIGIN_DEFAULT: &str = "https://grok.com";
fn default_oidc_scopes() -> Vec<String> {
vec![
"openid".into(),
"profile".into(),
"email".into(),
"offline_access".into(),
"api:access".into(),
]
}
/// Default scopes for the xAI OAuth2 provider. Includes `grok-cli:access`
/// which authorizes the token for API proxy requests.
fn default_oauth2_scopes() -> Vec<String> {
vec![
"openid".into(),
"profile".into(),
"email".into(),
"offline_access".into(),
"grok-cli:access".into(),
"api:access".into(),
"conversations:read".into(),
"conversations:write".into(),
]
}
fn default_team_oauth2_scopes() -> Vec<String> {
vec![
"profile".into(),
"offline_access".into(),
"grok-cli:access".into(),
"api:access".into(),
"team:read".into(),
"conversations:read".into(),
"conversations:write".into(),
]
}
/// Pin automatic auth to one method (`[auth] preferred_method` in config.toml).
///
/// When set, only that method is used for automatic selection; if it is
/// unavailable, auth fails (no silent fallthrough to the other method).
/// Unset keeps today's multi-method fallthrough (session preferred when both
/// exist). Config-toml only — not remote settings, settings UI, or env.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PreferredAuthMethod {
/// `XAI_API_KEY` / auth.json `xai::api_key` / per-model BYOK (`xai.api_key`).
ApiKey,
/// OIDC / OAuth2 session (`cached_token`, interactive `grok.com` / `oidc`,
/// including devbox-minted OIDC).
Oidc,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct GrokComConfig {
/// Auth origin / login-host display (doubles as the legacy WS origin name).
pub grok_ws_origin: String,
pub token_header: String,
/// OIDC config for customer-provided IdPs. See [`OidcAuthConfig`].
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc: Option<OidcAuthConfig>,
/// OAuth2 provider config. When set, preferred over the legacy relay flow.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oauth2: Option<OAuth2ProviderConfig>,
/// External auth provider command (stdout = token, stderr = user UX, exit 0 = success).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_provider_command: Option<String>,
/// Login button label (env: `KIGI_AUTH_PROVIDER_LABEL`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_provider_label: Option<String>,
/// Token TTL in seconds for external auth providers that output bare
/// tokens without `expires_in`. Synthesizes `expires_at` so proactive
/// refresh works. Env: `KIGI_AUTH_TOKEN_TTL`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_token_ttl: Option<u64>,
/// Admin kill switch: when `Some(true)`, the `xai.api_key` auth method is
/// neither advertised nor accepted, so `XAI_API_KEY`/per-model credentials
/// can't bypass the deployment's IdP login. Env: `KIGI_DISABLE_API_KEY_AUTH`.
/// Parity with common force-login-method admin knobs.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub disable_api_key_auth: Option<bool>,
/// Restrict login to a specific team — the login token's team principal must
/// equal this. Put in `requirements.toml` to enforce as non-overridable policy.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub force_login_team_uuid: Option<ForceLoginTeam>,
/// Pin automatic auth to `api_key` or `oidc`. When set and the chosen
/// method is unavailable, auth fails (no fallthrough). Unset keeps
/// multi-method fallthrough. Config.toml only (`[auth] preferred_method`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub preferred_method: Option<PreferredAuthMethod>,
}
/// Team login restriction. TOML string or array; an empty array fails closed.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ForceLoginTeam {
/// The only allowed team.
Single(String),
/// Allowed teams; empty = fail closed.
AnyOf(Vec<String>),
}
/// Customer OIDC Identity Provider configuration (`[grok_com_config.oidc]`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OidcAuthConfig {
pub issuer: String,
pub client_id: String,
#[serde(default = "default_oidc_scopes")]
pub scopes: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub audience: Option<String>,
}
/// OAuth2 provider configuration (`KIGI_OAUTH2_ISSUER` / `KIGI_OAUTH2_CLIENT_ID`).
///
/// Uses the standard OAuth 2.1 Auth Code + PKCE flow via [`OidcAuthConfig`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OAuth2ProviderConfig {
pub issuer: String,
pub client_id: String,
#[serde(default = "default_oauth2_scopes")]
pub scopes: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_id: Option<String>,
/// Client-supplied referrer for OAuth usage-attribution analytics.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub referrer: Option<String>,
}
pub const XAI_OAUTH2_ISSUER: &str = "https://auth.x.ai";
/// Production accounts-app origin allowlist — the only origins builds without
/// non-production builds accept. Lives in its own const, referenced by both
/// profiles below, so the frozen-contract test (monorepo CI compiles with
/// that feature enabled) still pins this production-origin const.
const PROD_ACCOUNTS_APP_ORIGINS: &[&str] = &["https://accounts.x.ai"];
/// See the opt-in non-production feature variant above — builds without
/// the feature accept only the production accounts app.
pub fn allowed_accounts_app_origins() -> Vec<String> {
PROD_ACCOUNTS_APP_ORIGINS
.iter()
.map(|o| o.to_string())
.collect()
}
/// Build a CORS layer that accepts requests from the accounts-app deployments
/// listed in [`allowed_accounts_app_origins`] for the given HTTP method.
///
/// Callers can chain additional configuration (e.g. `.allow_headers(...)` or
/// `.allow_private_network(true)`) onto the returned layer.
pub fn accounts_app_cors_layer(method: axum::http::Method) -> tower_http::cors::CorsLayer {
tower_http::cors::CorsLayer::new()
.allow_origin(tower_http::cors::AllowOrigin::list(
allowed_accounts_app_origins()
.iter()
.filter_map(|origin| match origin.parse() {
Ok(value) => Some(value),
Err(_) => {
tracing::warn!(origin, "skipping malformed accounts-app CORS origin");
None
}
}),
))
.allow_methods([method])
}
/// Local-dev OAuth2 issuer (accounts-app running on localhost).
const XAI_OAUTH2_LOCAL_ISSUER: &str = "http://localhost:22255";
const DEFAULT_OAUTH2_REFERRER: &str = "grok-build";
/// Returns `true` when `KIGI_LOCAL_AUTH=1` is set,
/// indicating the local accounts-app should be used as the OAuth2 issuer.
pub fn use_local_auth() -> bool {
std::env::var("KIGI_LOCAL_AUTH")
.map(|v| !v.is_empty() && v != "0")
.unwrap_or(false)
}
/// Returns the active xAI OAuth2 issuer — the local-dev issuer when
/// `KIGI_LOCAL_AUTH=1` is set, otherwise the production issuer.
pub fn xai_oauth2_issuer() -> &'static str {
if use_local_auth() {
XAI_OAUTH2_LOCAL_ISSUER
} else {
XAI_OAUTH2_ISSUER
}
}
/// Returns `true` if `issuer` is a recognised xAI OAuth2 issuer
/// (production **or** local-dev). Use this instead of comparing against
/// [`XAI_OAUTH2_ISSUER`] directly so that local-dev sessions are still
/// treated as first-party xAI auth.
pub fn is_xai_oauth2_issuer(issuer: &str) -> bool {
issuer == XAI_OAUTH2_ISSUER || issuer == XAI_OAUTH2_LOCAL_ISSUER
}
/// auth.json scope key used by the pre-OIDC `grok login --legacy` flow.
/// Matches the key format produced by the original `accounts.x.ai` relay auth.
pub const LEGACY_AUTH_SCOPE: &str = "https://accounts.x.ai/sign-in";
impl GrokComConfig {
/// Whether `xai.api_key` auth is disabled. Pinning a team
/// (`force_login_team_uuid`) implies this — team membership can't be verified
/// from a bare API key, so it must go through IdP login. The
/// `KIGI_DISABLE_API_KEY_AUTH` env lockdown is sticky: because the env value
/// seeds `default()` (the merge base), a lower-trust user `config.toml` could
/// otherwise set `disable_api_key_auth = false` and override it — so the env
/// is OR-ed in here and cannot be turned back off by a user layer. Trusted
/// `requirements.toml` already wins over `config.toml` via layer precedence.
pub fn api_key_auth_disabled(&self) -> bool {
self.disable_api_key_auth == Some(true)
|| self.force_login_team_uuid.is_some()
|| env_lockdown_forced()
}
/// When `preferred_method = api_key`, automatic OIDC paths (devbox mint,
/// interactive browser login, external auth provider) must not run — the
/// pin is fail-closed. Explicit `grok login --devbox` / `--api-key` bypass
/// this by not consulting automatic flow helpers.
pub fn blocks_automatic_oidc(&self) -> bool {
matches!(self.preferred_method, Some(PreferredAuthMethod::ApiKey))
}
/// The auth.json scope key for this config.
pub fn auth_scope(&self) -> String {
if let Some(ref oidc) = self.oidc {
format!("{}::{}", oidc.issuer.trim_end_matches('/'), oidc.client_id)
} else if let Some(ref oauth2) = self.oauth2 {
oauth2.auth_scope()
} else {
unreachable!("oauth2 config is always present (xAI default or env override)")
}
}
}
impl OAuth2ProviderConfig {
pub fn is_team_principal(&self) -> bool {
self.principal_type.as_deref() == Some(TEAM_PRINCIPAL_TYPE)
}
pub fn from_env() -> Option<Self> {
let issuer = std::env::var("KIGI_OAUTH2_ISSUER").ok()?;
let client_id = std::env::var("KIGI_OAUTH2_CLIENT_ID").ok()?;
let principal_type = std::env::var("KIGI_OAUTH2_PRINCIPAL_TYPE").ok();
let principal_id = std::env::var("KIGI_OAUTH2_PRINCIPAL_ID").ok();
let default_scopes = match principal_type.as_deref() {
Some(TEAM_PRINCIPAL_TYPE) => default_team_oauth2_scopes(),
_ => default_oauth2_scopes(),
};
Some(Self {
issuer,
client_id,
scopes: std::env::var("KIGI_OAUTH2_SCOPES")
.map(|s| s.split(',').map(|s| s.trim().to_owned()).collect())
.unwrap_or(default_scopes),
principal_type,
principal_id,
referrer: Some(
std::env::var("KIGI_OAUTH2_REFERRER")
.unwrap_or_else(|_| DEFAULT_OAUTH2_REFERRER.to_owned()),
),
})
}
/// Convert to [`OidcAuthConfig`] to reuse the OIDC login flow.
pub fn as_oidc(&self) -> OidcAuthConfig {
OidcAuthConfig {
issuer: self.issuer.clone(),
client_id: self.client_id.clone(),
scopes: self.scopes.clone(),
audience: None,
}
}
pub fn base_auth_scope(&self) -> String {
format!("{}::{}", self.issuer.trim_end_matches('/'), self.client_id)
}
pub fn auth_scope(&self) -> String {
self.base_auth_scope()
}
}
impl Default for GrokComConfig {
fn default() -> Self {
let oidc = OidcAuthConfig::from_env();
let oauth2 = if oidc.is_some() {
None
} else {
Some(
OAuth2ProviderConfig::from_env().unwrap_or_else(|| OAuth2ProviderConfig {
issuer: xai_oauth2_issuer().to_owned(),
client_id: obfstr::obfstr!("b1a00492-073a-47ea-816f-4c329264a828").to_owned(),
scopes: default_oauth2_scopes(),
principal_type: None,
principal_id: None,
referrer: Some(DEFAULT_OAUTH2_REFERRER.to_owned()),
}),
)
};
Self {
grok_ws_origin: std::env::var("KIGI_WS_ORIGIN")
.unwrap_or_else(|_| AUTH_ORIGIN_DEFAULT.to_owned()),
token_header: "xai-grok-cli".to_owned(),
oidc,
oauth2,
auth_provider_command: std::env::var("KIGI_AUTH_PROVIDER_COMMAND").ok(),
auth_provider_label: std::env::var("KIGI_AUTH_PROVIDER_LABEL").ok(),
auth_token_ttl: std::env::var("KIGI_AUTH_TOKEN_TTL")
.ok()
.and_then(|v| v.parse().ok()),
disable_api_key_auth: std::env::var("KIGI_DISABLE_API_KEY_AUTH")
.ok()
.map(|v| env_flag_enabled(&v)),
force_login_team_uuid: None,
preferred_method: None,
}
}
}
/// Parse a boolean env-var value for grok's on/off flags. A bare presence
/// enables the flag, but the common falsy spellings (`0`, `false`, `off`,
/// `no`, empty) count as disabled — so e.g. `KIGI_DISABLE_API_KEY_AUTH=false`
/// does NOT turn the kill switch on.
fn env_flag_enabled(value: &str) -> bool {
!matches!(
value.trim().to_ascii_lowercase().as_str(),
"" | "0" | "false" | "off" | "no"
)
}
/// True when the admin has set `KIGI_DISABLE_API_KEY_AUTH` to a truthy value in
/// the process environment. Read live (call-time) and OR-ed into
/// `api_key_auth_disabled()` so the env lockdown is non-overridable by a
/// user-layer `config.toml`.
fn env_lockdown_forced() -> bool {
std::env::var("KIGI_DISABLE_API_KEY_AUTH")
.ok()
.is_some_and(|v| env_flag_enabled(&v))
}
impl OidcAuthConfig {
pub fn from_env() -> Option<Self> {
let issuer = std::env::var("KIGI_OIDC_ISSUER").ok()?;
let client_id = std::env::var("KIGI_OIDC_CLIENT_ID").ok()?;
Some(Self {
issuer,
client_id,
scopes: std::env::var("KIGI_OIDC_SCOPES")
.map(|s| s.split(',').map(|s| s.trim().to_owned()).collect())
.unwrap_or_else(|_| default_oidc_scopes()),
audience: std::env::var("KIGI_OIDC_AUDIENCE").ok(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn team_auth_scope_is_base_scope() {
let cfg = OAuth2ProviderConfig {
issuer: "https://auth.x.ai".into(),
client_id: "client-123".into(),
scopes: default_team_oauth2_scopes(),
principal_type: Some("Team".into()),
principal_id: Some("team-abc".into()),
referrer: Some("grok-build".into()),
};
assert_eq!(cfg.auth_scope(), "https://auth.x.ai::client-123");
}
#[test]
fn env_flag_enabled_treats_falsy_spellings_as_off() {
for off in ["", " ", "0", "false", "FALSE", "off", "No", " false "] {
assert!(!env_flag_enabled(off), "{off:?} should be off");
}
for on in ["1", "true", "yes", "on", "enabled"] {
assert!(env_flag_enabled(on), "{on:?} should be on");
}
}
#[test]
fn personal_auth_scope_is_base_scope() {
let cfg = OAuth2ProviderConfig {
issuer: "https://auth.x.ai".into(),
client_id: "client-123".into(),
scopes: default_oauth2_scopes(),
principal_type: None,
principal_id: None,
referrer: Some("grok-build".into()),
};
assert_eq!(cfg.auth_scope(), "https://auth.x.ai::client-123");
}
/// FROZEN loopback contract: the accounts-app origins the CLI's loopback
/// callback server accepts cross-origin requests from. The consent page
/// (served from accounts.x.ai) delivers the code via `fetch(..., cors)`, so
/// removing an origin breaks loopback delivery for already-installed CLIs.
/// Keep in sync with the oauth2-provider / accounts-app deployments.
/// Non-production / local-dev origins are opt-in only.
#[test]
fn allowed_accounts_app_origins_are_frozen() {
assert_eq!(PROD_ACCOUNTS_APP_ORIGINS, &["https://accounts.x.ai"]);
assert_eq!(allowed_accounts_app_origins(), PROD_ACCOUNTS_APP_ORIGINS);
}
/// FROZEN client contract: the 8 scopes the xAI OAuth2 client requests.
/// The server must keep accepting all of them; existing tokens carry
/// exactly this set. Frozen OAuth client scope contract.
#[test]
fn default_oauth2_scopes_are_frozen() {
let scopes = default_oauth2_scopes();
let scopes: Vec<&str> = scopes.iter().map(String::as_str).collect();
assert_eq!(
scopes,
[
"openid",
"profile",
"email",
"offline_access",
"grok-cli:access",
"api:access",
"conversations:read",
"conversations:write",
]
);
}
#[test]
fn preferred_method_deserializes_from_toml() {
let cfg: GrokComConfig = toml::from_str(
r#"
preferred_method = "api_key"
"#,
)
.expect("parse");
assert_eq!(cfg.preferred_method, Some(PreferredAuthMethod::ApiKey));
let cfg: GrokComConfig = toml::from_str(
r#"
preferred_method = "oidc"
"#,
)
.expect("parse");
assert_eq!(cfg.preferred_method, Some(PreferredAuthMethod::Oidc));
let cfg: GrokComConfig = toml::from_str("").expect("parse empty");
assert_eq!(cfg.preferred_method, None);
}
}
crates/codegen/kigi-shell/src/auth/credential_provider.rs
+325
View File
@@ -0,0 +1,325 @@
use crate::auth::AuthManager;
use crate::util::grok_auth_credentials::GrokAuthCredentials;
use kigi_auth::{
AuthCredentialProvider, CredentialSnapshot, HttpAuth, StaticAuthCredentialProvider,
};
use reqwest::RequestBuilder;
use std::sync::Arc;
/// `api_key.id` for the active credential: hash the stable API key, never the
/// OIDC bearer (which rotates). `None` for non-API-key auth.
fn api_key_id_for(auth: Option<&crate::auth::GrokAuth>) -> Option<String> {
auth.filter(|a| matches!(a.auth_mode, crate::auth::AuthMode::ApiKey))
.map(|a| crate::agent::config::deployment_id_from_key(&a.key))
}
/// Production impl: wraps the live `AuthManager`. 401 recovery
/// delegates to `AuthManager::unauthorized_recovery`.
pub struct ShellAuthCredentialProvider {
auth_manager: Arc<AuthManager>,
static_credentials: GrokAuthCredentials,
}
impl ShellAuthCredentialProvider {
pub(crate) fn new(
auth_manager: Arc<AuthManager>,
deployment_key: Option<String>,
alpha_test_key: Option<String>,
) -> Self {
let mut static_credentials = GrokAuthCredentials::new(None);
static_credentials.deployment_key = deployment_key;
static_credentials.alpha_test_key = alpha_test_key;
Self {
auth_manager,
static_credentials,
}
}
}
impl std::fmt::Debug for ShellAuthCredentialProvider {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ShellAuthCredentialProvider")
.field("auth_manager", &"<configured>")
.finish()
}
}
impl HttpAuth for ShellAuthCredentialProvider {
fn apply(&self, builder: RequestBuilder, base_url: &str) -> RequestBuilder {
let mut creds = self.static_credentials.clone();
if creds.deployment_key.is_none()
&& let Some(auth) = self.auth_manager.current_or_expired()
{
creds.user_token = Some(auth.key);
}
creds.apply(builder, base_url)
}
}
#[async_trait::async_trait]
impl AuthCredentialProvider for ShellAuthCredentialProvider {
fn snapshot(&self) -> CredentialSnapshot {
if let Some(ref dk) = self.static_credentials.deployment_key {
return CredentialSnapshot {
token: Some(dk.clone()),
deployment_id: crate::managed_config::resolve_deployment_id(Some(dk)),
..Default::default()
};
}
let auth = self.auth_manager.current_or_expired();
let user_id = auth.as_ref().map(|a| a.user_id.clone());
let team_id = auth.as_ref().and_then(|a| a.team_id.clone());
let organization_id = auth.as_ref().and_then(|a| a.organization_id.clone());
let api_key_id = api_key_id_for(auth.as_ref());
let token = auth.map(|a| a.key);
CredentialSnapshot {
token,
user_id,
team_id,
deployment_id: None,
api_key_id,
organization_id,
}
}
async fn refresh_after_unauthorized(&self) -> bool {
if self.static_credentials.deployment_key.is_some() {
return false;
}
self.auth_manager.try_recover_unauthorized().await
}
fn needs_token_auth_header(&self) -> bool {
self.static_credentials.deployment_key.is_none()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::GrokAuth;
use crate::auth::GrokComConfig;
use crate::auth::manager::AuthManager;
use chrono::{Duration as ChronoDuration, Utc};
use kigi_auth::AuthCredentialProvider;
use std::sync::Mutex;
/// Serializes tests that pin `KIGI_AUTH_EARLY_INVALIDATION_SECS`, since
/// env vars are process-global and parallel tests would race.
static EARLY_INVALIDATION_LOCK: Mutex<()> = Mutex::new(());
/// RAII guard: pins `KIGI_AUTH_EARLY_INVALIDATION_SECS` to the production
/// default (300s) while held, restoring the previous value on drop.
/// Acquires `EARLY_INVALIDATION_LOCK` so concurrent test runners can't
/// observe a half-mutated env.
struct EarlyInvalidationGuard {
_lock: std::sync::MutexGuard<'static, ()>,
previous: Option<String>,
}
impl EarlyInvalidationGuard {
fn pin_to_default() -> Self {
let lock = EARLY_INVALIDATION_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let previous = std::env::var("KIGI_AUTH_EARLY_INVALIDATION_SECS").ok();
unsafe { std::env::set_var("KIGI_AUTH_EARLY_INVALIDATION_SECS", "300") };
Self {
_lock: lock,
previous,
}
}
}
impl Drop for EarlyInvalidationGuard {
fn drop(&mut self) {
unsafe {
match self.previous.take() {
Some(prev) => std::env::set_var("KIGI_AUTH_EARLY_INVALIDATION_SECS", prev),
None => std::env::remove_var("KIGI_AUTH_EARLY_INVALIDATION_SECS"),
}
}
}
}
fn make_auth(key: &str, expires_in: ChronoDuration) -> GrokAuth {
GrokAuth {
key: key.to_string(),
user_id: "test-user".to_string(),
create_time: Utc::now(),
expires_at: Some(Utc::now() + expires_in),
..GrokAuth::test_default()
}
}
/// Build an `AuthManager` rooted at `dir`. Caller keeps `dir` alive for
/// the duration of the test so the `TempDir` `Drop` actually cleans up.
fn make_manager(dir: &tempfile::TempDir, initial: Option<GrokAuth>) -> Arc<AuthManager> {
let mgr = AuthManager::new(dir.path(), GrokComConfig::default());
if let Some(auth) = initial {
mgr.hot_swap(auth);
}
Arc::new(mgr)
}
/// `apply()` and `snapshot()` agree (snapshot==wire invariant) when the
/// in-memory token is fresh.
#[test]
fn apply_and_snapshot_agree_on_live_token() {
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let mgr = make_manager(
&dir,
Some(make_auth("live-token", ChronoDuration::hours(1))),
);
let provider = ShellAuthCredentialProvider::new(mgr, None, None);
let snap = provider.snapshot();
assert_eq!(snap.token.as_deref(), Some("live-token"));
assert_eq!(snap.user_id.as_deref(), Some("test-user"));
}
/// During the 5-minute pre-refresh buffer window, `auth_manager.current()`
/// returns `None` (the token is treated as expired-soon for refresh
/// scheduling), but the token is still valid at the proxy. The provider
/// must fall back to `expired_auth()` so the in-memory token gets sent
/// instead of nothing -- which is the fix for the bulk of the
/// `POST /v1/storage` 401s observed in production.
#[test]
fn falls_back_to_expired_auth_during_buffer_window() {
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let mgr = make_manager(
&dir,
Some(make_auth("buffer-token", ChronoDuration::minutes(4))),
);
assert!(mgr.current().is_none(), "buffer-window precondition");
assert!(mgr.expired_auth().is_some(), "buffer-window precondition");
let provider = ShellAuthCredentialProvider::new(mgr, None, None);
let snap = provider.snapshot();
assert_eq!(
snap.token.as_deref(),
Some("buffer-token"),
"snapshot should fall back to expired_auth instead of None"
);
assert_eq!(snap.user_id.as_deref(), Some("test-user"));
}
/// When `auth_manager` has nothing at all (no in-memory auth, expired
/// or otherwise), `snapshot()` returns `None` for the user-token branch.
/// `apply()` would then send no Authorization header.
#[test]
fn no_token_when_auth_manager_is_empty() {
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let mgr = make_manager(&dir, None);
let provider = ShellAuthCredentialProvider::new(mgr, None, None);
let snap = provider.snapshot();
assert!(
snap.token.is_none(),
"snapshot should be None when manager has no auth"
);
assert!(snap.user_id.is_none());
}
/// 401 recovery routes through `unauthorized_recovery` (pre-fix
/// it no-oped because the refresher arg was hardcoded `None`).
#[tokio::test]
async fn refresh_after_unauthorized_drives_recovery_state_machine() {
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(
dir.path(),
crate::auth::GrokComConfig::default(),
));
mgr.hot_swap(GrokAuth {
key: "stale".into(),
auth_mode: crate::auth::AuthMode::Oidc,
create_time: chrono::Utc::now() - ChronoDuration::hours(2),
user_id: "u".into(),
refresh_token: Some("rt-stale".into()),
expires_at: Some(chrono::Utc::now() - ChronoDuration::hours(1)),
..GrokAuth::test_default()
});
struct OkRefresher {
calls: Arc<std::sync::atomic::AtomicU32>,
}
#[async_trait::async_trait]
impl crate::auth::refresh::TokenRefresher for OkRefresher {
async fn refresh(
&self,
_r: crate::auth::manager::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
key: "fresh".into(),
auth_mode: crate::auth::AuthMode::Oidc,
create_time: chrono::Utc::now(),
user_id: "u".into(),
refresh_token: Some("rt-new".into()),
expires_at: Some(chrono::Utc::now() + ChronoDuration::hours(1)),
..GrokAuth::test_default()
}))
}
}
let calls = Arc::new(std::sync::atomic::AtomicU32::new(0));
mgr.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let provider = ShellAuthCredentialProvider::new(mgr.clone(), None, None);
assert!(provider.refresh_after_unauthorized().await);
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
assert_eq!(mgr.current().unwrap().key, "fresh");
assert_eq!(
provider.snapshot().token.as_deref(),
Some("fresh"),
"snapshot must reflect refreshed token for subsequent apply() calls"
);
}
/// Deployment-key path has no recovery (operator owns the bearer).
#[tokio::test]
async fn refresh_after_unauthorized_is_noop_for_deployment_key() {
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let mgr = make_manager(&dir, None);
let provider =
ShellAuthCredentialProvider::new(mgr, Some("deployment-key".to_string()), None);
assert!(!provider.refresh_after_unauthorized().await);
}
#[test]
fn snapshot_populates_tenant_id_per_auth_mode() {
use crate::agent::config::deployment_id_from_key;
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let dep = ShellAuthCredentialProvider::new(
make_manager(&dir, None),
Some("xai-token-EX".into()),
None,
)
.snapshot();
assert_eq!(
dep.deployment_id.as_deref(),
Some(deployment_id_from_key("xai-token-EX").as_str())
);
assert!(dep.api_key_id.is_none());
let api_auth = GrokAuth {
key: "sk-apikey-xyz".into(),
auth_mode: crate::auth::AuthMode::ApiKey,
expires_at: Some(Utc::now() + ChronoDuration::hours(1)),
..GrokAuth::test_default()
};
let api = ShellAuthCredentialProvider::new(make_manager(&dir, Some(api_auth)), None, None)
.snapshot();
assert_eq!(
api.api_key_id.as_deref(),
Some(deployment_id_from_key("sk-apikey-xyz").as_str())
);
assert!(api.deployment_id.is_none());
let oidc = ShellAuthCredentialProvider::new(
make_manager(
&dir,
Some(make_auth("oidc-token", ChronoDuration::hours(1))),
),
None,
None,
)
.snapshot();
assert!(oidc.deployment_id.is_none() && oidc.api_key_id.is_none());
}
/// Bootstrap mode: `snapshot()` re-reads disk so sibling-rotated
/// tokens are picked up without a live AuthManager.
#[test]
fn deployment_key_wins_over_resolved_user_token() {
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let mgr = make_manager(
&dir,
Some(make_auth("user-token", ChronoDuration::hours(1))),
);
let provider =
ShellAuthCredentialProvider::new(mgr, Some("deployment-key-12345".to_string()), None);
let snap = provider.snapshot();
assert_eq!(snap.token.as_deref(), Some("deployment-key-12345"));
assert!(snap.user_id.is_none());
}
}
crates/codegen/kigi-shell/src/auth/devbox_login_stub.rs
+37
View File
@@ -0,0 +1,37 @@
//! Stub for builds without the devbox auth feature.
//!
//! Compiled instead of `devbox_login.rs` when the devbox auth feature is
//! off, so the remote devbox login helper is not reached. The API
//! mirrors the real module: `is_devbox_environment()` is always `false`, which
//! short-circuits every auto-recovery/migration call site, and the entry
//! points that can still be reached directly (`grok login --devbox`) return a
//! descriptive error.
use super::manager::AuthManager;
use super::model::GrokAuth;
const UNAVAILABLE: &str =
"devbox login is not available in this build (compiled without the `devbox-login` feature)";
/// Always `false` without the devbox auth feature; callers treat the
/// process as running outside a devbox environment.
pub(crate) fn is_devbox_environment() -> bool {
false
}
/// Unreachable in practice (guarded by [`is_devbox_environment`]); errors
/// defensively if called.
pub(crate) async fn mint_devbox_auth(_auth_manager: &AuthManager) -> anyhow::Result<GrokAuth> {
anyhow::bail!(UNAVAILABLE)
}
/// Unreachable in practice (guarded by [`is_devbox_environment`]); errors
/// defensively if called.
pub(super) async fn mint_devbox_auth_raw() -> anyhow::Result<GrokAuth> {
anyhow::bail!(UNAVAILABLE)
}
/// `grok login --devbox` entry point: always errors in this build.
pub async fn run_devbox_login(_config: &crate::agent::config::Config) -> anyhow::Result<GrokAuth> {
anyhow::bail!(UNAVAILABLE)
}
crates/codegen/kigi-shell/src/auth/device_code.rs
+890
View File
@@ -0,0 +1,890 @@
//! RFC 8628 Device Authorization Grant -- CLI side.
//!
//! Two-phase API:
//! 1. `request_device_code()` -- POST to server, get code + URL
//! 2. `complete_device_code_login()` -- poll until approved, persist credentials
//!
//! Callers control what happens between the two phases (print to stderr,
//! show in TUI, display in IDE sidebar, etc.).
use std::sync::Arc;
use chrono::{Duration, Utc};
use serde::Deserialize;
use thiserror::Error;
use crate::auth::oidc::with_alpha_test_key;
use crate::auth::{AuthChannels, AuthManager, AuthMode, AuthUrlInfo, AuthUrlMode, GrokAuth};
const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code";
const DEFAULT_DEVICE_POLL_INTERVAL_SECS: i32 = 5;
const DEVICE_SLOW_DOWN_INCREMENT_SECS: u64 = 5;
const MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS: i64 = 10 * 60;
#[derive(Debug, Error)]
pub enum DeviceCodeError {
#[error(
"Device-code login is not available for this deployment. \
Try `grok login` or set XAI_API_KEY instead."
)]
NotEnabled,
#[error(transparent)]
Other(#[from] anyhow::Error),
}
impl From<reqwest::Error> for DeviceCodeError {
fn from(e: reqwest::Error) -> Self {
Self::Other(e.into())
}
}
// --- Public types ---
/// Low-cardinality client-surface hint sent to the OAuth2 provider as the
/// `x-grok-client-surface` header so device-flow metrics can separate logins a
/// human can actually finish (`Ui`, `Cli`) from headless automation
/// (`Headless`) that mints a device code but can never reach the browser
/// consent page — the traffic that otherwise pollutes the device-flow
/// conversion denominator.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientSurface {
/// An interactive front-end (TUI / IDE) renders the URL + code to a human.
Ui,
/// CLI attached to an interactive terminal (stderr is a TTY).
Cli,
/// No interactive surface (CI, container, script): no human can complete.
Headless,
}
impl ClientSurface {
fn as_str(self) -> &'static str {
match self {
Self::Ui => "ui",
Self::Cli => "cli",
Self::Headless => "headless",
}
}
}
/// Classify the CLI (non-TUI) surface: a TTY on stderr means a human is
/// watching the printed URL + code; otherwise we're headless (CI/container/
/// script) and no one will complete the flow.
fn detect_cli_surface() -> ClientSurface {
use std::io::IsTerminal as _;
if std::io::stderr().is_terminal() {
ClientSurface::Cli
} else {
ClientSurface::Headless
}
}
/// Result of requesting a device code from the server.
/// Callers display `verification_uri` + `user_code` to the user,
/// then pass this struct to `complete_device_code_login`.
#[derive(Debug, Clone)]
pub struct DeviceCode {
pub verification_uri: String,
pub verification_uri_complete: Option<String>,
pub user_code: String,
device_code: String,
interval: i32,
expires_in: i64,
}
// --- Wire types (serde) ---
#[derive(Deserialize)]
struct DeviceCodeResponse {
device_code: String,
user_code: String,
verification_uri: String,
verification_uri_complete: Option<String>,
expires_in: i64,
interval: Option<i32>,
}
#[derive(Deserialize)]
struct TokenOk {
access_token: String,
refresh_token: Option<String>,
expires_in: Option<i64>,
#[expect(dead_code, reason = "field retained for protocol compatibility")]
scope: Option<String>,
id_token: Option<String>,
}
#[derive(Deserialize)]
struct TokenErr {
error: String,
error_description: Option<String>,
}
#[derive(Deserialize)]
struct IdTokenClaims {
sub: Option<String>,
email: Option<String>,
}
// --- Phase 1: Request device code ---
/// Request a device code + user code from the OAuth2 provider.
///
/// This is a single HTTP POST. The caller is responsible for displaying
/// `DeviceCode::verification_uri` and `DeviceCode::user_code` to the user
/// before calling `complete_device_code_login`.
pub async fn request_device_code(
issuer: &str,
client_id: &str,
scopes: &[String],
surface: ClientSurface,
) -> Result<DeviceCode, DeviceCodeError> {
let client = crate::http::shared_client();
let url = format!("{}/oauth2/device/code", issuer.trim_end_matches('/'));
let scope_str = scopes.join(" ");
let resp = with_alpha_test_key(
client
.post(&url)
// Lets oauth2-provider segment device-flow success by client version.
.header("x-grok-client-version", kigi_version::VERSION)
// Lets oauth2-provider separate human-completable logins from
// headless automation in the device-flow funnel metrics.
.header("x-grok-client-surface", surface.as_str())
.form(&[
("client_id", client_id),
("scope", scope_str.as_str()),
("referrer", "grok-build"),
]),
&url,
)
.send()
.await?;
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
if status.as_u16() == 404 {
return Err(DeviceCodeError::NotEnabled);
}
return Err(anyhow::anyhow!("Device code request failed (HTTP {status}): {body}").into());
}
let server_resp: DeviceCodeResponse = resp.json().await?;
// Defend against control characters from a malicious issuer.
if !server_resp
.user_code
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '-')
{
return Err(anyhow::anyhow!(
"Server returned invalid user_code format (expected [A-Z0-9-])"
)
.into());
}
validate_verification_uri(&server_resp.verification_uri)?;
if let Some(ref verification_uri_complete) = server_resp.verification_uri_complete {
validate_verification_uri(verification_uri_complete)?;
}
Ok(DeviceCode {
verification_uri: server_resp.verification_uri,
verification_uri_complete: server_resp.verification_uri_complete,
user_code: server_resp.user_code,
device_code: server_resp.device_code,
interval: server_resp
.interval
.unwrap_or(DEFAULT_DEVICE_POLL_INTERVAL_SECS),
expires_in: server_resp.expires_in,
})
}
// --- Phase 2: Poll until approved ---
/// Poll the token endpoint until the user approves (or denies / expires).
///
/// On success, persists credentials to `~/.kigi/auth.json` and returns
/// the authenticated `GrokAuth`.
///
/// Callers should have already displayed `device_code.verification_uri`
/// and `device_code.user_code` to the user before calling this.
pub async fn complete_device_code_login(
issuer: &str,
client_id: &str,
device_code: DeviceCode,
auth_manager: &Arc<AuthManager>,
surface: ClientSurface,
) -> anyhow::Result<(GrokAuth, bool)> {
let client = crate::http::shared_client();
let token_url = format!("{}/oauth2/token", issuer.trim_end_matches('/'));
let mut poll_interval = std::time::Duration::from_secs(device_code.interval.max(1) as u64);
let deadline = tokio::time::Instant::now()
+ std::time::Duration::from_secs(
device_code
.expires_in
.max(MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS) as u64,
);
loop {
// Sleep first: an immediate poll on a fresh code only returns
// authorization_pending (and risks slow_down).
tokio::time::sleep(poll_interval).await;
if tokio::time::Instant::now() > deadline {
anyhow::bail!("Device code expired. Run `grok login --device-auth` again.");
}
let resp = with_alpha_test_key(
client
.post(&token_url)
.header("x-grok-client-version", kigi_version::VERSION)
.header("x-grok-client-surface", surface.as_str())
.form(&[
("grant_type", DEVICE_GRANT_TYPE),
("device_code", device_code.device_code.as_str()),
("client_id", client_id),
]),
&token_url,
)
.send()
.await?;
if resp.status().is_success() {
let tokens: TokenOk = resp.json().await?;
let auth = build_auth(&tokens, issuer, client_id, auth_manager).await?;
return Ok((auth, true));
}
let err: TokenErr = resp.json().await?;
let detail = err.error_description.as_deref().unwrap_or(&err.error);
match err.error.as_str() {
"authorization_pending" => {
// User hasn't acted yet -- keep polling.
continue;
}
"slow_down" => {
poll_interval += std::time::Duration::from_secs(DEVICE_SLOW_DOWN_INCREMENT_SECS);
continue;
}
"access_denied" => {
tracing::warn!(description = detail, "device auth authorization denied");
anyhow::bail!("Authorization denied. The user rejected the request.");
}
"expired_token" => {
tracing::warn!(description = detail, "device auth token expired");
anyhow::bail!("Device code expired. Run `grok login --device-auth` again.");
}
other => {
tracing::warn!(
error = other,
description = detail,
"device auth token exchange failed"
);
anyhow::bail!("Token exchange error: {detail}");
}
}
}
}
/// Device-code login shared by the TUI and CLI.
///
/// With `channels` (TUI) the verification URL goes to `url_tx` and the browser
/// opens automatically; on failure the copyable URL is the fallback. Without
/// `channels` (CLI) the URL + code are printed to stderr via `prompt_and_poll`.
/// `code_rx` is unused here. The caller reports success (`✓ Signed in`).
///
/// Takes `channels` by `&mut`, consuming it only after the device code is
/// obtained, so callers can reuse it for a loopback fallback on `NotEnabled`.
pub async fn run_device_code_login_channels(
issuer: &str,
client_id: &str,
scopes: &[String],
auth_manager: &Arc<AuthManager>,
channels: &mut Option<AuthChannels>,
) -> anyhow::Result<(GrokAuth, bool)> {
// A front-end (TUI/IDE) listening on `url_tx` renders the URL to a human, so
// it's `Ui`. Without one we're on the CLI: a TTY means a human can act
// (`Cli`), no TTY means headless automation (`Headless`) that will never
// complete. Computed before `take()` so the `request_device_code` call
// already carries the surface.
let surface = if channels.is_some() {
ClientSurface::Ui
} else {
detect_cli_surface()
};
let device_code = request_device_code(issuer, client_id, scopes, surface).await?;
let Some(channels) = channels.take() else {
// CLI: print the URL + code to stderr.
return prompt_and_poll(issuer, client_id, device_code, auth_manager, surface).await;
};
// TUI: push the URL through the channel BEFORE opening the browser, so
// `x.ai/auth/get_url` isn't blocked on a slow/hanging browser launch
// (e.g. SSH/headless). When the issuer omits `verification_uri_complete`,
// embed the code so the welcome screen can still show it (anti-phishing).
let display_uri = match device_code.verification_uri_complete.as_deref() {
Some(uri) => uri.to_owned(),
None => {
let sep = if device_code.verification_uri.contains('?') {
'&'
} else {
'?'
};
format!(
"{}{}user_code={}",
device_code.verification_uri, sep, device_code.user_code
)
}
};
if let Some(tx) = channels.url_tx {
let _ = tx.send(AuthUrlInfo {
url: display_uri.clone(),
mode: AuthUrlMode::Device,
});
}
open_browser_detached(&display_uri).await;
complete_device_code_login(issuer, client_id, device_code, auth_manager, surface).await
}
/// Display the device code to stderr and poll until approved.
async fn prompt_and_poll(
issuer: &str,
client_id: &str,
device_code: DeviceCode,
auth_manager: &Arc<AuthManager>,
surface: ClientSurface,
) -> anyhow::Result<(GrokAuth, bool)> {
let display_uri = device_code
.verification_uri_complete
.as_deref()
.unwrap_or(&device_code.verification_uri);
eprintln!();
eprintln!("To sign in, open this URL in your browser:");
eprintln!();
eprintln!(" {}", display_uri);
eprintln!();
if !open_browser_detached(display_uri).await {
eprintln!(" (Could not open browser automatically — open the URL above manually.)");
eprintln!();
}
// Show the code to confirm it matches the browser (anti-phishing): a complete
// URL pre-fills it (just confirm), otherwise the user types it.
if device_code.verification_uri_complete.is_some() {
eprintln!("Confirm this code in your browser:");
} else {
eprintln!("Then enter this code:");
}
eprintln!();
eprintln!(" {}", device_code.user_code);
eprintln!();
eprintln!(
"\x1b[90mOnly continue with a code you requested. \
Don't share it with anyone.\x1b[0m"
);
eprintln!();
eprintln!("Waiting for authorization...");
// The caller prints the `✓ Signed in` confirmation (it also owns the
// external-provider / devbox early-return paths that never reach here).
complete_device_code_login(issuer, client_id, device_code, auth_manager, surface).await
}
/// Open `url` in the browser off-thread: `webbrowser::open` is synchronous and
/// would stall the single-threaded TUI loop. Returns `true` on success so the
/// caller can decide how to notify the user (eprintln on CLI, nothing on TUI
/// where the URL is already rendered in the widget).
async fn open_browser_detached(url: &str) -> bool {
let url = url.to_owned();
match tokio::task::spawn_blocking(move || webbrowser::open(&url)).await {
Ok(Ok(())) => true,
Ok(Err(e)) => {
tracing::info!(error = %e, "device auth: could not open browser automatically");
false
}
Err(e) => {
tracing::info!(error = %e, "device auth: browser-open task failed");
false
}
}
}
// --- Internal helpers ---
/// No id_token signature verification -- token arrives over a direct HTTPS
/// channel (no browser redirect), and is only used for display info (email).
async fn build_auth(
tokens: &TokenOk,
issuer: &str,
client_id: &str,
auth_manager: &Arc<AuthManager>,
) -> anyhow::Result<GrokAuth> {
let (user_id, email) = if let Some(ref id_token) = tokens.id_token {
decode_jwt_claims(id_token)
} else {
(String::new(), None)
};
let (principal_type, principal_id, token_team_id) =
match crate::auth::oidc::peek_access_token_principal(&tokens.access_token) {
Some((pt, pid, tid)) => (Some(pt), Some(pid), tid),
None => (None, None, None),
};
// Device flow has no pre-selection; verify the token's principal here.
// Match the principal id even if `principal_type` is absent.
let principal_policy =
crate::auth::oidc::login_principal_policy(auth_manager.grok_com_config());
crate::auth::oidc::enforce_login_principal(
principal_policy.as_ref(),
crate::auth::oidc::peek_access_token_principal_id(&tokens.access_token).as_deref(),
)?;
let (user_id, email, team_id, organization_id) =
match (principal_type.as_deref(), principal_id.as_deref()) {
(Some(pt), Some(principal_id)) if pt == crate::auth::model::TEAM_PRINCIPAL_TYPE => (
principal_id.to_owned(),
None,
Some(principal_id.to_owned()),
None,
),
(Some("Organization"), Some(principal_id)) => (
principal_id.to_owned(),
None,
None,
Some(principal_id.to_owned()),
),
_ => (user_id, email, token_team_id, None),
};
let now = Utc::now();
let mut auth = GrokAuth {
key: tokens.access_token.clone(),
auth_mode: AuthMode::Oidc,
create_time: now,
user_id,
email,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type,
principal_id,
organization_id,
organization_name: None,
organization_role: None,
team_id,
team_name: None,
team_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: tokens.refresh_token.clone(),
expires_at: tokens.expires_in.map(|s| now + Duration::seconds(s)),
oidc_issuer: Some(issuer.to_owned()),
oidc_client_id: Some(client_id.to_owned()),
};
auth_manager.enrich_auth_inline(&mut auth).await;
auth_manager
.update(auth)
.await
.map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}"))
}
/// Decode JWT payload without signature verification.
/// Returns (sub, Option<email>).
fn decode_jwt_claims(jwt: &str) -> (String, Option<String>) {
use base64::Engine;
let parts: Vec<&str> = jwt.splitn(3, '.').collect();
if parts.len() < 2 {
return (String::new(), None);
}
let payload = match base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(parts[1]) {
Ok(bytes) => bytes,
Err(_) => return (String::new(), None),
};
let claims: IdTokenClaims = match serde_json::from_slice(&payload) {
Ok(claims) => claims,
Err(_) => return (String::new(), None),
};
(claims.sub.unwrap_or_default(), claims.email)
}
fn validate_verification_uri(uri: &str) -> anyhow::Result<()> {
if uri.chars().any(|c| c.is_ascii_control()) {
anyhow::bail!("Server returned invalid verification URI");
}
let parsed = url::Url::parse(uri)
.map_err(|_| anyhow::anyhow!("Server returned invalid verification URI"))?;
match parsed.scheme() {
"https" => Ok(()),
"http" if matches!(parsed.host_str(), Some("localhost") | Some("127.0.0.1")) => Ok(()),
_ => anyhow::bail!("Server returned unsupported verification URI scheme"),
}
}
#[cfg(test)]
pub(crate) mod tests {
use std::sync::Arc;
use super::{AuthManager, build_auth, validate_verification_uri};
use crate::auth::{AuthMode, GrokComConfig};
#[test]
fn validate_verification_uri_rejects_unsupported_scheme() {
let err = validate_verification_uri("javascript:alert(1)").unwrap_err();
assert_eq!(
"Server returned unsupported verification URI scheme",
err.to_string()
);
}
fn auth_manager_with_kigi_home(
kigi_home: &std::path::Path,
proxy_base_url: &str,
) -> Arc<AuthManager> {
Arc::new(
AuthManager::new(kigi_home, GrokComConfig::default())
.with_proxy_base_url(proxy_base_url),
)
}
#[test]
fn build_auth_persists_credentials_without_proxy_fetch() {
let temp_dir = tempfile::tempdir().unwrap();
let kigi_home = temp_dir.path().join(".kigi");
std::fs::create_dir_all(&kigi_home).unwrap();
let auth_manager = auth_manager_with_kigi_home(&kigi_home, "http://127.0.0.1:9");
let tokens = super::TokenOk {
access_token: "access-token".to_string(),
refresh_token: Some("refresh-token".to_string()),
expires_in: Some(900),
scope: Some("openid email offline_access grok-cli:access".to_string()),
id_token: Some(
"eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJzdWIiOiJ1c2VyLTEyMyIsImVtYWlsIjoiZGV2aWNlLWF1dGhAbG9jYWwudGVzdCJ9.sig".to_string(),
),
};
let auth = tokio::runtime::Runtime::new()
.unwrap()
.block_on(build_auth(
&tokens,
"http://localhost:22255",
"client-id",
&auth_manager,
))
.unwrap();
assert_eq!("access-token", auth.key);
assert_eq!(AuthMode::Oidc, auth.auth_mode);
assert_eq!("user-123", auth.user_id);
assert_eq!(Some("device-auth@local.test".to_string()), auth.email);
assert_eq!(Some("refresh-token".to_string()), auth.refresh_token);
assert_eq!(Some("http://localhost:22255".to_string()), auth.oidc_issuer);
assert_eq!(Some("client-id".to_string()), auth.oidc_client_id);
assert!(auth_manager.current().is_some());
}
/// jsonwebtoken needs a process-level CryptoProvider; tests that encode
/// JWTs can't rely on another test having installed it first.
fn ensure_crypto_provider() {
let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default();
}
#[test]
fn build_auth_seeds_team_metadata_from_access_token() {
ensure_crypto_provider();
let temp_dir = tempfile::tempdir().unwrap();
let kigi_home = temp_dir.path().join(".kigi");
std::fs::create_dir_all(&kigi_home).unwrap();
let auth_manager = auth_manager_with_kigi_home(&kigi_home, "http://127.0.0.1:9");
let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256);
let claims = serde_json::json!({
"sub": "user-42",
"iss": "https://auth.x.ai",
"aud": "client-id",
"exp": 9999999999u64,
"iat": 1000000000u64,
"scope": "offline_access grok-cli:access team:read",
"principal_type": "Team",
"principal_id": "team-123",
"client_id": "client-id",
"jti": "token-1",
});
let tokens = super::TokenOk {
access_token: jsonwebtoken::encode(
&header,
&claims,
&jsonwebtoken::EncodingKey::from_secret(b"test-secret"),
)
.unwrap(),
refresh_token: Some("refresh-token".to_owned()),
expires_in: Some(900),
scope: Some("offline_access grok-cli:access team:read".to_owned()),
id_token: None,
};
let auth = tokio::runtime::Runtime::new()
.unwrap()
.block_on(build_auth(
&tokens,
"http://localhost:22255",
"client-id",
&auth_manager,
))
.unwrap();
assert_eq!("team-123", auth.user_id);
assert_eq!(Some("Team".to_owned()), auth.principal_type);
assert_eq!(Some("team-123".to_owned()), auth.principal_id);
assert_eq!(Some("team-123".to_owned()), auth.team_id);
assert_eq!(None, auth.organization_id);
assert_eq!(None, auth.email);
}
/// Team access token carrying `principal_id` (signature irrelevant — only
/// the principal claims are peeked).
fn team_access_token(principal_id: &str) -> super::TokenOk {
let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256);
let claims = serde_json::json!({
"sub": "user-42",
"exp": 9999999999u64,
"principal_type": "Team",
"principal_id": principal_id,
});
super::TokenOk {
access_token: jsonwebtoken::encode(
&header,
&claims,
&jsonwebtoken::EncodingKey::from_secret(b"test-secret"),
)
.unwrap(),
refresh_token: Some("refresh-token".to_owned()),
expires_in: Some(900),
scope: None,
id_token: None,
}
}
/// `build_auth` with `token_principal` must fail with `expected_err` and
/// persist nothing.
fn assert_build_auth_rejected(cfg: GrokComConfig, token_principal: &str, expected_err: &str) {
ensure_crypto_provider();
let temp_dir = tempfile::tempdir().unwrap();
let kigi_home = temp_dir.path().join(".kigi");
std::fs::create_dir_all(&kigi_home).unwrap();
let auth_manager =
Arc::new(AuthManager::new(&kigi_home, cfg).with_proxy_base_url("http://127.0.0.1:9"));
let err = tokio::runtime::Runtime::new()
.unwrap()
.block_on(build_auth(
&team_access_token(token_principal),
"http://localhost:22255",
"client-id",
&auth_manager,
))
.unwrap_err();
assert_eq!(err.to_string(), expected_err);
assert!(
auth_manager.current().is_none(),
"rejected login must not persist credentials",
);
assert!(
!kigi_home.join("auth.json").exists(),
"rejected login must not write auth.json",
);
}
/// The legacy `oauth2.principal_id` only pre-selects a team; it must not
/// enforce a pin (only `force_login_team_uuid` does), so a different team's
/// token is accepted.
#[test]
fn build_auth_does_not_enforce_legacy_oauth2_principal_id() {
ensure_crypto_provider();
let cfg = GrokComConfig {
oauth2: Some(crate::auth::OAuth2ProviderConfig {
issuer: "http://localhost:22255".into(),
client_id: "client-id".into(),
scopes: vec!["offline_access".into()],
principal_type: Some("Team".into()),
principal_id: Some("team-required".into()),
referrer: None,
}),
..GrokComConfig::default()
};
let temp_dir = tempfile::tempdir().unwrap();
let kigi_home = temp_dir.path().join(".kigi");
std::fs::create_dir_all(&kigi_home).unwrap();
let auth_manager =
Arc::new(AuthManager::new(&kigi_home, cfg).with_proxy_base_url("http://127.0.0.1:9"));
let auth = tokio::runtime::Runtime::new()
.unwrap()
.block_on(build_auth(
&team_access_token("team-other"),
"http://localhost:22255",
"client-id",
&auth_manager,
))
.expect("legacy oauth2.principal_id must not enforce a pin");
assert_eq!(
auth.team_id.as_deref(),
Some("team-other"),
"the token's own team is used; the legacy pre-select id does not gate it",
);
}
/// Persistence-seam enforcement via a `force_login_team_uuid` list.
#[test]
fn build_auth_rejects_token_outside_force_login_team_list() {
let cfg = GrokComConfig {
force_login_team_uuid: Some(crate::auth::ForceLoginTeam::AnyOf(vec![
"team-a".into(),
"team-b".into(),
])),
..GrokComConfig::default()
};
assert_build_auth_rejected(
cfg,
"team-other",
"This deployment requires logging into one of teams: team-a, team-b; \
your login returned team-other",
);
}
// ── complete_device_code_login poll loop ────────────────────────────────
/// Spawn a mock `/oauth2/token` server that serves `responses` in order,
/// repeating the last entry. Returns the issuer base URL.
async fn spawn_token_server(
responses: Vec<(u16, serde_json::Value)>,
) -> (String, tokio::task::JoinHandle<()>) {
use std::sync::atomic::{AtomicUsize, Ordering};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let issuer = format!("http://{}", listener.local_addr().unwrap());
let counter = Arc::new(AtomicUsize::new(0));
let responses = Arc::new(responses);
let app = axum::Router::new().route(
"/oauth2/token",
axum::routing::post(move || {
let counter = counter.clone();
let responses = responses.clone();
async move {
let idx = counter
.fetch_add(1, Ordering::SeqCst)
.min(responses.len() - 1);
let (status, body) = &responses[idx];
(
axum::http::StatusCode::from_u16(*status).unwrap(),
axum::Json(body.clone()),
)
}
}),
);
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(issuer, handle)
}
fn device_code_for_test(interval: i32, expires_in: i64) -> super::DeviceCode {
super::DeviceCode {
verification_uri: "https://example.test/device".into(),
verification_uri_complete: Some(
"https://example.test/device?user_code=ABCD-EFGH".into(),
),
user_code: "ABCD-EFGH".into(),
device_code: "dev-code-123".into(),
interval,
expires_in,
}
}
// Real time (not `start_paused`: the shared client's 30s connect_timeout
// fires under auto-advance). Deadline-expiry isn't tested — the deadline is
// floored at 10 min (MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS).
async fn run_poll(
responses: Vec<(u16, serde_json::Value)>,
) -> anyhow::Result<(super::GrokAuth, bool)> {
let (issuer, server) = spawn_token_server(responses).await;
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager = auth_manager_with_kigi_home(temp_dir.path(), "http://127.0.0.1:9");
let device_code = device_code_for_test(1, 900);
let result = super::complete_device_code_login(
&issuer,
"client-id",
device_code,
&auth_manager,
super::ClientSurface::Cli,
)
.await;
server.abort();
result
}
fn success_body() -> serde_json::Value {
serde_json::json!({
"access_token": "mock-access-token",
"refresh_token": "mock-refresh-token",
"expires_in": 900,
"scope": "openid",
})
}
#[tokio::test]
async fn poll_succeeds_on_first_poll() {
let (auth, is_new) = run_poll(vec![(200, success_body())])
.await
.expect("should resolve to a token");
assert_eq!(auth.key, "mock-access-token");
assert!(is_new);
}
#[tokio::test]
async fn poll_succeeds_after_pending() {
let (auth, _) = run_poll(vec![
(400, serde_json::json!({ "error": "authorization_pending" })),
(200, success_body()),
])
.await
.expect("should resolve to a token after pending");
assert_eq!(auth.key, "mock-access-token");
}
#[tokio::test]
async fn poll_handles_slow_down_then_succeeds() {
// slow_down must be tolerated (interval bumped) without erroring.
let (auth, _) = run_poll(vec![
(400, serde_json::json!({ "error": "slow_down" })),
(200, success_body()),
])
.await
.expect("slow_down should be retried, not fatal");
assert_eq!(auth.key, "mock-access-token");
}
#[tokio::test]
async fn poll_maps_access_denied_to_error() {
let err = run_poll(vec![(400, serde_json::json!({ "error": "access_denied" }))])
.await
.expect_err("access_denied must be an error");
assert!(err.to_string().contains("denied"), "got: {err}");
}
#[tokio::test]
async fn poll_maps_expired_token_to_error() {
let err = run_poll(vec![(400, serde_json::json!({ "error": "expired_token" }))])
.await
.expect_err("expired_token must be an error");
assert!(err.to_string().contains("expired"), "got: {err}");
}
}
crates/codegen/kigi-shell/src/auth/error.rs
+144
View File
@@ -0,0 +1,144 @@
use thiserror::Error;
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum AuthError {
#[error("Not logged in. Run `grok login`.")]
NotLoggedIn,
/// Token expired and no refresh authority available.
#[error("Token expired. Run `grok login` to re-authenticate.")]
TokenExpiredNoRefresh,
/// Server rejected the token (401) with no recovery path.
#[error("Authentication rejected by server. Run `grok login` to re-authenticate.")]
ServerRejectedNoRecovery,
/// All recovery strategies exhausted.
#[error("Auth recovery exhausted; re-authentication required.")]
RecoveryExhausted,
/// A session's team principal violates the `force_login_team_uuid` pin.
/// `message` states which team is required vs. returned.
#[error("{message} Run `grok login` to sign in with the required team.")]
PinnedTeamMismatch { message: String },
/// Cached API-key session rejected because API-key auth is disabled.
#[error("API-key auth is disabled by your administrator. Run `grok login` to authenticate.")]
ApiKeyAuthDisabled,
/// Outcome of a refresh-authority attempt. Recoverability (and, for
/// permanent failures, the reason) lives in [`RefreshTokenError`].
#[error(transparent)]
Refresh(#[from] RefreshTokenError),
}
/// Recoverability axis of a token-refresh attempt. Deliberately total (no
/// `#[non_exhaustive]`): "permanent vs transient" is a closed decision every
/// caller must make, so a future third state should break consumers loudly.
#[derive(Debug, Error)]
pub enum RefreshTokenError {
/// The credential is dead; the user must re-authenticate.
#[error(transparent)]
Permanent(#[from] RefreshTokenFailedError),
/// Network / 5xx / unknown blip; safe to retry later. Carries the cause.
#[error(transparent)]
Transient(RefreshTransientError),
}
/// A retryable refresh failure, wrapping its cause. No public `From`:
/// construct only via [`AuthError::transient`] /
/// [`AuthError::transient_source`], so a stray `?` on some error can't silently
/// classify a permanent failure as retryable (mirrors the dedicated
/// [`RefreshTokenFailedError`] on the permanent arm). Display frames the cause
/// as an auth-refresh failure so internal messages (lock timeout, sleep defer)
/// don't surface bare; the permanent arm derives its copy from
/// [`RefreshTokenFailedReason::user_message`] and is not prefixed.
#[derive(Debug, Error)]
#[error("auth refresh failed: {0}")]
pub struct RefreshTransientError(#[source] Box<dyn std::error::Error + Send + Sync>);
/// A terminal refresh failure. `reason` is machine-readable; the user-facing
/// copy is derived from it via [`RefreshTokenFailedReason::user_message`], so
/// the two can never drift.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
#[error("{}", .reason.user_message())]
#[non_exhaustive]
pub struct RefreshTokenFailedError {
pub reason: RefreshTokenFailedReason,
}
impl From<RefreshTokenFailedReason> for RefreshTokenFailedError {
fn from(reason: RefreshTokenFailedReason) -> Self {
Self { reason }
}
}
/// Why a token refresh terminally failed, grounded in the OAuth2 error codes
/// our IdP actually emits.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum RefreshTokenFailedReason {
/// `invalid_grant` — the refresh token is no longer valid (expired, reused,
/// or revoked; the IdP does not distinguish these).
RefreshTokenRejected,
/// `invalid_client` — the client/app credential was rejected.
ClientRejected,
/// Escalation from repeated transient failures (OIDC) or a single
/// external-binary failure. Never a raw IdP code: an unrecognized terminal
/// code is classified transient, not `Other` (see `classify_terminal`).
Other,
}
impl RefreshTokenFailedReason {
/// Sticky until the credential changes (never ages out): a revoked refresh
/// token never self-heals, whereas client rotation / transient escalation
/// recover, so those age out past the TTL.
pub(crate) fn is_sticky(self) -> bool {
match self {
Self::RefreshTokenRejected => true,
Self::ClientRejected | Self::Other => false,
}
}
/// User-facing copy for a terminal refresh failure; the raw IdP code stays
/// in logs.
pub(crate) fn user_message(self) -> &'static str {
match self {
Self::RefreshTokenRejected => {
"Your session has expired. Run `grok login` to sign in again."
}
Self::ClientRejected => {
"Authentication is temporarily unavailable. Run `grok login` if this persists."
}
Self::Other => {
"Authentication could not be refreshed. Run `grok login` to sign in again."
}
}
}
}
impl AuthError {
/// A retryable refresh failure with a message-only cause, for the genuinely
/// message-only sites (lock timeout, sleep/dark-wake defer, no refresher);
/// use [`Self::transient_source`] when a real error is in hand.
pub(crate) fn transient(message: impl Into<String>) -> Self {
Self::transient_source(message.into())
}
/// A retryable refresh failure that preserves `source` in the error chain
/// (`Transient` carries the cause), so callers with a real error don't
/// flatten it to a string.
pub(crate) fn transient_source(
source: impl Into<Box<dyn std::error::Error + Send + Sync>>,
) -> Self {
AuthError::Refresh(RefreshTokenError::Transient(RefreshTransientError(
source.into(),
)))
}
/// A terminal refresh failure for an already-classified `reason`.
pub(crate) fn permanent(reason: RefreshTokenFailedReason) -> Self {
AuthError::Refresh(RefreshTokenError::Permanent(reason.into()))
}
}
crates/codegen/kigi-shell/src/auth/external_auth.rs
+283
View File
@@ -0,0 +1,283 @@
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.
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)
};
Ok(GrokAuth {
key: token,
auth_mode: AuthMode::External,
create_time: chrono::Utc::now(),
user_id: String::new(),
email: None,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token,
expires_at,
oidc_issuer: 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> {
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)");
let mut cmd = Command::new("sh");
cmd.args(["-c", command])
.stdin(Stdio::null())
.stdout(Stdio::piped())
// Pipe stderr — inherit would corrupt the TUI alternate screen.
.stderr(Stdio::piped());
if is_refresh {
cmd.env("KIGI_AUTH_EXPIRED", "1");
}
kigi_tools::util::detach_std_command(&mut cmd);
cmd.envs(kigi_tools::util::pager_env());
let mut child = cmd.spawn()
.map_err(|e| {
tracing::warn!(error = %e, cmd = %command, "auth: failed to start external auth provider");
e
})
.ok()?;
let timeout = std::time::Duration::from_secs(timeout_secs);
let start = std::time::Instant::now();
loop {
match child.try_wait() {
Ok(Some(_status)) => break,
Ok(None) => {
if start.elapsed() > timeout {
tracing::warn!(
cmd = %command,
timeout_secs,
"auth: external auth provider timed out (likely needs interactive auth), killing"
);
let _ = child.kill();
let _ = child.wait();
return None;
}
std::thread::sleep(std::time::Duration::from_millis(100));
}
Err(e) => {
tracing::warn!(error = %e, "auth: error waiting for external auth provider");
return None;
}
}
}
let output = child
.wait_with_output()
.map_err(|e| {
tracing::warn!(error = %e, "auth: failed to read external auth provider output");
e
})
.ok()?;
match parse_output(&output) {
Ok(auth) => {
tracing::info!("auth: external auth provider returned fresh token");
Some(auth)
}
Err(e) => {
tracing::warn!(error = %e, "auth: external auth provider failed");
None
}
}
}
/// Run external auth provider, carrying forward `/user`-derived fields from previous auth.
pub(crate) fn refresh_with_command(command: &str, prev_auth: &GrokAuth) -> Option<GrokAuth> {
let mut auth = run_external_auth_sync(command, true)?;
auth.carry_user_profile_from(prev_auth);
Some(auth)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_output_nonzero_exit_is_err() {
let output = std::process::Output {
status: std::process::Command::new("false").status().unwrap(),
stdout: b"token".to_vec(),
stderr: vec![],
};
assert!(parse_output(&output).is_err());
}
#[test]
fn parse_output_empty_stdout_is_err() {
let output = std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: b" \n".to_vec(),
stderr: vec![],
};
assert!(parse_output(&output).is_err());
}
#[test]
fn parse_output_issuer_claim_enables_xai_auth() {
let ok = |stdout: &str| std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: stdout.as_bytes().to_vec(),
stderr: vec![],
};
// x.ai issuer claim → first-party session (relay-eligible).
let auth = parse_output(&ok(
r#"{"access_token":"t","expires_in":900,"issuer":"https://auth.x.ai"}"#,
))
.unwrap();
assert_eq!(auth.oidc_issuer.as_deref(), Some("https://auth.x.ai"));
assert!(auth.is_xai_auth());
// Non-x.ai issuer is stored but stays third-party.
let auth = parse_output(&ok(
r#"{"access_token":"t","issuer":"https://idp.acme.example"}"#,
))
.unwrap();
assert_eq!(
auth.oidc_issuer.as_deref(),
Some("https://idp.acme.example")
);
assert!(!auth.is_xai_auth());
// Missing / empty / whitespace issuer → None.
let auth = parse_output(&ok(r#"{"access_token":"t"}"#)).unwrap();
assert_eq!(auth.oidc_issuer, None);
assert!(!auth.is_xai_auth());
let auth = parse_output(&ok(r#"{"access_token":"t","issuer":" "}"#)).unwrap();
assert_eq!(auth.oidc_issuer, None);
// Bare-token output never carries an issuer.
let auth = parse_output(&ok("bare-token")).unwrap();
assert_eq!(auth.oidc_issuer, None);
assert!(!auth.is_xai_auth());
}
#[test]
fn parse_output_malformed_json_falls_back_to_bare() {
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}");
}
#[test]
fn sync_spawn_failure_returns_none() {
assert!(run_external_auth_sync("/nonexistent/binary", false).is_none());
}
#[test]
fn sync_sets_grok_auth_expired_env_on_refresh() {
let auth = run_external_auth_sync("echo $KIGI_AUTH_EXPIRED", true).unwrap();
assert_eq!(auth.key, "1");
}
#[test]
fn refresh_carries_zdr_flags_forward() {
let prev = GrokAuth {
user_blocked_reason: Some("BLOCKED_REASON_OTHER".into()),
team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()],
coding_data_retention_opt_out: true,
organization_id: Some("org-1".into()),
..GrokAuth::test_default()
};
let auth = refresh_with_command("echo fresh-token", &prev).unwrap();
assert_eq!(auth.key, "fresh-token");
assert!(auth.is_zdr_team(), "ZDR flag must survive refresh");
assert!(auth.coding_data_retention_opt_out);
assert_eq!(
auth.user_blocked_reason.as_deref(),
Some("BLOCKED_REASON_OTHER")
);
assert_eq!(auth.user_id, "test-user", "profile must survive refresh");
assert_eq!(auth.organization_id.as_deref(), Some("org-1"));
}
#[test]
fn sync_refresh_interactive_times_out() {
// Binary writes link to stderr then blocks — 5s refresh timeout kills it.
let cmd = r#"echo 'Visit http://example.com/auth' >&2; sleep 20; echo token"#;
let start = std::time::Instant::now();
let result = run_external_auth_sync(cmd, true);
let elapsed = start.elapsed();
assert!(result.is_none(), "should timeout and return None");
assert!(
elapsed.as_secs() < 10,
"refresh should use 5s timeout, not 60s (took {}s)",
elapsed.as_secs()
);
}
}
crates/codegen/kigi-shell/src/auth/flow.rs
+1955
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/jwt.rs
+45
View File
@@ -0,0 +1,45 @@
//! JWT expiration detection. Returns `None`/`false` for non-JWT tokens.
use chrono::{DateTime, Duration, Utc};
use serde::Deserialize;
#[derive(Deserialize)]
struct Claims {
exp: Option<i64>,
}
pub fn parse_jwt_expiration(token: &str) -> Option<DateTime<Utc>> {
jsonwebtoken::dangerous::insecure_decode::<Claims>(token)
.ok()
.and_then(|data| data.claims.exp)
.and_then(|ts| DateTime::from_timestamp(ts, 0))
}
pub fn is_jwt_expired_or_near(token: &str, threshold: Duration) -> bool {
parse_jwt_expiration(token)
.map(|exp| exp <= Utc::now() + threshold)
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
/// Tokens with an `aud` claim must parse successfully.
/// `jsonwebtoken::Validation::default()` enables audience validation which
/// silently rejects these tokens unless `validate_aud = false` is set.
#[test]
fn parses_jwt_with_aud_claim() {
let token = build_test_jwt(r#"{"aud":["some-audience"],"exp":1772575524}"#);
let exp = parse_jwt_expiration(&token);
assert_eq!(exp.unwrap().timestamp(), 1772575524);
}
fn build_test_jwt(payload_json: &str) -> String {
use base64::Engine;
let enc = base64::engine::general_purpose::URL_SAFE_NO_PAD;
let header = enc.encode(r#"{"alg":"RS256","typ":"JWT"}"#);
let payload = enc.encode(payload_json);
format!("{header}.{payload}.fake-signature")
}
}
crates/codegen/kigi-shell/src/auth/manager.rs
+2208
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/manager/enrichment.rs
+256
View File
@@ -0,0 +1,256 @@
//! Background `/user` enrichment spawned by `AuthManager::update()`.
use std::sync::Arc;
use std::time::Duration as StdDuration;
use super::AuthManager;
use super::lock::try_lock_auth_file_async;
use crate::auth::manager::AUTH_LOCK_TIMEOUT;
use crate::auth::model::{GrokAuth, UserInfo, lookup_auth};
use crate::auth::storage::{read_auth_json, write_auth_json};
/// `/user` fetch budget, shared by the inline (login) and background paths.
const USER_FETCH_TIMEOUT: StdDuration = StdDuration::from_secs(10);
/// Logs `auth update enrichment dropped` if the task is cancelled
/// mid-flight. Disarmed on normal completion.
pub(super) struct EnrichmentExitGuard {
pub(super) started: std::time::Instant,
pub(super) armed: bool,
}
impl EnrichmentExitGuard {
pub(super) fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for EnrichmentExitGuard {
fn drop(&mut self) {
if !self.armed {
return;
}
kigi_log::unified_log::warn(
"auth update enrichment dropped",
None,
Some(serde_json::json!({
"elapsed_ms": self.started.elapsed().as_millis() as u64,
})),
);
}
}
pub(super) fn spawn(manager: Arc<AuthManager>, auth: GrokAuth) {
tokio::spawn(async move {
let mut exit_guard = EnrichmentExitGuard {
started: std::time::Instant::now(),
armed: true,
};
run_user_info_enrichment(&manager, auth).await;
exit_guard.disarm();
});
}
async fn fetch_user_info(manager: &AuthManager, key: &str, log_label: &str) -> Option<UserInfo> {
let user_url = format!("{}/user", manager.proxy_base_url);
let token_header = &manager.grok_com_config.token_header;
let started = std::time::Instant::now();
let http_client = crate::http::shared_client();
let response = http_client
.get(&user_url)
.timeout(USER_FETCH_TIMEOUT)
.header("Authorization", format!("Bearer {}", key))
.header("X-XAI-Token-Auth", token_header.as_str())
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.send()
.await;
match response {
Ok(resp) if resp.status().is_success() => match resp.json::<UserInfo>().await {
Ok(ui) if !ui.user_id.is_empty() => Some(ui),
Ok(_) => {
kigi_log::unified_log::warn(
&format!("{log_label} skipped"),
None,
Some(serde_json::json!({
"reason": "empty_user_id",
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
Err(e) => {
kigi_log::unified_log::warn(
&format!("{log_label} failed"),
None,
Some(serde_json::json!({
"reason": "parse",
"error": e.to_string(),
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
},
Ok(resp) => {
kigi_log::unified_log::warn(
&format!("{log_label} failed"),
None,
Some(serde_json::json!({
"reason": "http_status",
"http_status": resp.status().as_u16(),
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
Err(e) => {
kigi_log::unified_log::warn(
&format!("{log_label} failed"),
None,
Some(serde_json::json!({
"reason": if e.is_timeout() { "timeout" } else { "transport" },
"error": e.to_string(),
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
}
}
/// Blocking login-time enrichment: merge `/user` fields before the first save.
pub(super) async fn enrich_inline(manager: &AuthManager, auth: &mut GrokAuth) {
let Some(ui) = fetch_user_info(manager, &auth.key, "auth login enrichment").await else {
return;
};
apply_user_info_enrichment(auth, ui);
}
async fn run_user_info_enrichment(manager: &AuthManager, auth: GrokAuth) {
let started = std::time::Instant::now();
let Some(user_info) = fetch_user_info(manager, &auth.key, "auth update enrichment").await
else {
return;
};
let user_elapsed_ms = started.elapsed().as_millis() as u64;
// R-M-W file lock. On timeout, fall through to an unlocked write
// rather than drop the enrichment.
let lock_started = std::time::Instant::now();
let lock_guard = try_lock_auth_file_async(&manager.path, AUTH_LOCK_TIMEOUT).await;
let lock_wait_ms = lock_started.elapsed().as_millis() as u64;
if lock_guard.is_none() {
tracing::warn!("auth: enrichment proceeding without auth.json.lock");
}
let Ok(mut map) = read_auth_json(&manager.path) else {
kigi_log::unified_log::warn(
"auth update enrichment skipped",
None,
Some(serde_json::json!({ "reason": "read_disk_failed" })),
);
return;
};
let Some(mut disk) = lookup_auth(&map, &manager.scope) else {
kigi_log::unified_log::info(
"auth update enrichment skipped",
None,
Some(serde_json::json!({ "reason": "no_disk_auth" })),
);
return;
};
// Sibling-stomp guard. If either the access token or refresh
// token on disk differs from the one we wrote, a sibling process
// rotated tokens since our update(). Skip enrichment to avoid
// writing stale profile data over the sibling's fresher entry.
//
// OR logic (not AND): a single-field rotation (key changes, RT
// stays) is the common case during concurrent refresh. The old
// AND logic required ALL three fields to differ, letting
// single-field rotations through.
//
// Team-login transitions (placeholder→real user_id) don't rotate
// tokens, so OR correctly allows enrichment for that case.
if disk.key != auth.key || disk.refresh_token != auth.refresh_token {
kigi_log::unified_log::info(
"auth update enrichment skipped",
None,
Some(serde_json::json!({
"reason": "sibling_rotated",
"written_key_prefix": crate::auth::token_suffix(&auth.key),
"disk_key_prefix": crate::auth::token_suffix(&disk.key),
})),
);
return;
}
apply_user_info_enrichment(&mut disk, user_info);
map.insert(manager.scope.clone(), disk.clone());
let write_started = std::time::Instant::now();
if let Err(e) = write_auth_json(&manager.path, &map) {
kigi_log::unified_log::error(
"auth update enrichment write failed",
None,
Some(serde_json::json!({
"error": e.to_string(),
"user_ms": user_elapsed_ms,
"lock_wait_ms": lock_wait_ms,
"write_ms": write_started.elapsed().as_millis() as u64,
})),
);
return;
}
manager.with_inner_write(|inner| *inner = Some(disk));
kigi_log::unified_log::info(
"auth update enrichment done",
None,
Some(serde_json::json!({
"user_ms": user_elapsed_ms,
"lock_wait_ms": lock_wait_ms,
"write_ms": write_started.elapsed().as_millis() as u64,
"total_ms": started.elapsed().as_millis() as u64,
})),
);
}
/// Merge enrichment fields into disk auth. Does not touch token fields.
pub(super) fn apply_user_info_enrichment(disk: &mut GrokAuth, user_info: UserInfo) {
disk.user_id = user_info.user_id;
disk.first_name = user_info.first_name.or(disk.first_name.take());
disk.last_name = user_info.last_name.or(disk.last_name.take());
disk.profile_image_asset_id = user_info
.profile_image_asset_id
.or(disk.profile_image_asset_id.take());
disk.principal_type = user_info.principal_type.or(disk.principal_type.take());
disk.principal_id = user_info.principal_id.or(disk.principal_id.take());
disk.team_id = user_info.team_id.or(disk.team_id.take());
disk.team_name = user_info.team_name.or(disk.team_name.take());
disk.team_role = user_info.team_role.or(disk.team_role.take());
disk.organization_id = user_info.organization_id.or(disk.organization_id.take());
disk.organization_name = user_info
.organization_name
.or(disk.organization_name.take());
disk.organization_role = user_info
.organization_role
.or(disk.organization_role.take());
disk.user_blocked_reason = user_info
.user_blocked_reason
.or(disk.user_blocked_reason.take());
if let Some(reasons) = user_info.team_blocked_reasons {
disk.team_blocked_reasons = reasons;
}
if let Some(opt_out) = user_info.coding_data_retention_opt_out {
disk.coding_data_retention_opt_out = opt_out;
}
if let Some(ref email) = user_info.email
&& !email.is_empty()
{
disk.email = user_info.email;
}
}
crates/codegen/kigi-shell/src/auth/manager/lock.rs
+1187
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/manager/sleep_gate.rs
+433
View File
@@ -0,0 +1,433 @@
//! System-sleep refresh-straddle mitigation for [`AuthManager`].
//!
//! A refresh that straddles a suspend can lose its rotated successor token,
//! leaving a revoked refresh token on disk and forcing re-login. Two layers
//! guard against that straddle:
//!
//! 1. The gate `refresh_chain` consults *defers* a not-yet-started refresh; an
//! in-flight one is never aborted (dropping it could discard a rotated-token
//! response — the very revocation we guard against). See
//! [`AuthManager::refresh_chain`].
//! 2. When sleep becomes imminent and a refresh *is* already in flight,
//! [`AuthManager::set_system_sleep_imminent`] briefly **holds the OS sleep
//! acknowledgment** (macOS delays `IOAllowPowerChange`; Linux holds its
//! `delay` inhibitor — both via the blocking power-listener callback) until
//! the refresh drains or [`SLEEP_ACK_MAX_WAIT`] elapses, so the in-flight
//! exchange finishes *before* the machine suspends.
//!
//! Split out of `manager.rs` so the manager stays scannable: this is a
//! self-contained unit (the [`SleepGate`] type, the [`InFlightGuard`], and a
//! small `impl AuthManager` block driving them from OS power events).
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::{Duration as StdDuration, Instant, SystemTime};
use parking_lot::RwLock;
use super::AuthManager;
/// Max lifetime of the "system sleep imminent" gate. A wake event normally
/// clears it; this is the safety bound so a *missed* wake event can never
/// permanently block token refresh. Generous vs. the OS pre-sleep window
/// (macOS ~30 s, Linux ~5 s) — it only needs to outlast the sleep transition.
pub(super) const SLEEP_GATE_MAX: StdDuration = StdDuration::from_secs(120);
/// Max time a token refresh may stay deferred for **dark wake** before one is
/// forced through, mirroring [`SLEEP_GATE_MAX`]. A normal dark wake lasts
/// seconds and recurs interspersed with full wakes, so this rarely fires; it
/// rescues a machine that reports a *continuous* dark wake — e.g. an
/// interactive Mac with no display, whose system video capability is never set
/// — which would otherwise defer every refresh forever and reach the same
/// logged-out state this guard prevents. Bounded on two clocks (see
/// [`GateRaise`]) so it also survives the machine sleeping between dark wakes.
///
/// The straddle risk of one forced refresh is far smaller than a guaranteed
/// logout: requests only force through while the machine is busy enough to
/// issue them (so it is unlikely to re-sleep mid-exchange), and the idle
/// proactive loop reaches this at most once per [`BACKOFF_INTERVAL`].
///
/// [`BACKOFF_INTERVAL`]: super::BACKOFF_INTERVAL
pub(super) const DARK_WAKE_DEFER_MAX: StdDuration = StdDuration::from_secs(120);
/// Upper bound on how long a `WillSleep` transition will hold the OS sleep
/// acknowledgment waiting for in-flight IdP refreshes to drain (see
/// [`AuthManager::set_system_sleep_imminent`]). Must stay inside the OS
/// pre-sleep budgets — macOS allows ~30 s before `IOAllowPowerChange`; Linux
/// logind's `InhibitDelayMaxSec` defaults to 5 s — so we pick a value
/// comfortably under the smaller (Linux) budget; the inhibitor is released
/// before logind force-sleeps regardless. A healthy refresh round-trip is
/// ~1 s, so this is slack for a slow network, not the common path. Holding the
/// machine awake a few extra seconds is a negligible cost next to the forced
/// re-login a straddled refresh causes.
pub(super) const SLEEP_ACK_MAX_WAIT: StdDuration = StdDuration::from_secs(3);
/// When a gate was raised, captured on *two* clocks so the [`SLEEP_GATE_MAX`]
/// backstop survives a system sleep.
///
/// `Instant` is monotonic but, on macOS (`mach_absolute_time`) and Linux
/// (`CLOCK_MONOTONIC`), *pauses while the machine is asleep*. A gate raised just
/// before a long sleep would therefore never auto-expire on the monotonic clock
/// alone — the exact bug that let an expired token reach the server and 401.
/// The wall clock (`SystemTime`) keeps advancing through sleep, so we expire the
/// gate once *either* clock passes the bound:
/// - the monotonic clock bounds elapsed *awake* time (immune to wall-clock
/// jumps from NTP / manual changes), and
/// - the wall clock bounds elapsed *real* time (immune to the sleep pause).
#[derive(Clone, Copy)]
pub(super) struct GateRaise {
/// Monotonic; pauses during sleep. Bounds elapsed *awake* time.
pub(super) mono: Instant,
/// Wall clock; advances through sleep. Bounds elapsed *real* time.
pub(super) wall: SystemTime,
}
impl GateRaise {
pub(super) fn now() -> Self {
Self {
mono: Instant::now(),
wall: SystemTime::now(),
}
}
/// Elapsed on each clock as `(monotonic, wall)`. Wall-clock elapsed is
/// clamped to zero if the clock ran backwards (NTP step / manual change) so
/// a backward jump can never *extend* the gate — the monotonic clock still
/// bounds it in that case.
pub(super) fn elapsed(&self) -> (StdDuration, StdDuration) {
(
self.mono.elapsed(),
self.wall.elapsed().unwrap_or(StdDuration::ZERO),
)
}
}
/// A gate `refresh_chain` consults to avoid *starting* an IdP refresh just
/// before sleep. Only *defers* a not-yet-started refresh; an in-flight one is
/// left to finish (see [`AuthManager::refresh_chain`]).
#[derive(Default)]
pub(super) struct SleepGate {
pub(super) raised_at: RwLock<Option<GateRaise>>,
}
impl SleepGate {
pub(super) fn raise(&self) {
*self.raised_at.write() = Some(GateRaise::now());
kigi_log::unified_log::warn("auth.sleep.gate_set", None, None);
}
pub(super) fn lower(&self, reason: &str) {
let prev = self.raised_at.write().take();
let (mono_ms, wall_ms) = prev
.map(|r| {
let (mono, wall) = r.elapsed();
(mono.as_millis() as u64, wall.as_millis() as u64)
})
.unwrap_or((0, 0));
kigi_log::unified_log::info(
"auth.sleep.gate_cleared",
None,
Some(serde_json::json!({
"reason": reason,
"was_raised": prev.is_some(),
"mono_elapsed_ms": mono_ms,
"wall_elapsed_ms": wall_ms,
})),
);
}
/// A stale gate (a missed/late wake event) is lazily lowered here so it can
/// never permanently block refresh; this read can therefore have a side
/// effect. The gate expires once *either* clock passes [`SLEEP_GATE_MAX`]
/// (see [`GateRaise`]): without the wall-clock arm, a gate raised before a
/// long sleep would never auto-expire, because the monotonic clock pauses
/// while the machine is asleep.
pub(super) fn is_gated(&self) -> bool {
// Copy out so the read guard drops before the write lock below
// (parking_lot is not reentrant).
let raised_at = *self.raised_at.read();
let Some(raise) = raised_at else {
return false;
};
let (mono, wall) = raise.elapsed();
if mono < SLEEP_GATE_MAX && wall < SLEEP_GATE_MAX {
return true;
}
// Stale gate (missed/late wake). `sleep_straddle` = the monotonic clock
// is still under the bound but real (wall-clock) time is not: the
// machine slept through the gate without delivering a wake event. This
// is precisely the case the wall-clock arm was added to catch, so
// surface it explicitly to confirm the fix firing in the field.
let sleep_straddle = mono < SLEEP_GATE_MAX;
*self.raised_at.write() = None;
kigi_log::unified_log::info(
"auth.sleep.gate_cleared",
None,
Some(serde_json::json!({
"reason": "auto_expiry",
"sleep_straddle": sleep_straddle,
"mono_elapsed_ms": mono.as_millis() as u64,
"wall_elapsed_ms": wall.as_millis() as u64,
})),
);
false
}
}
/// RAII counter for in-flight IdP refreshes. Increments on construction and
/// decrements on drop so the count stays balanced even if the refresh future is
/// cancelled or panics. When the count returns to zero it wakes any
/// sleep-imminent waiter parked in
/// [`AuthManager::hold_sleep_ack_until_refresh_drains`].
pub(super) struct InFlightGuard<'a>(&'a AuthManager);
impl<'a> InFlightGuard<'a> {
pub(super) fn new(mgr: &'a AuthManager) -> Self {
mgr.begin_refresh_in_flight();
Self(mgr)
}
}
impl Drop for InFlightGuard<'_> {
fn drop(&mut self) {
self.0.end_refresh_in_flight();
}
}
impl AuthManager {
/// Report a system power transition (`true` = sleep imminent, `false` =
/// woke). Safe to call from any thread.
pub(crate) fn set_system_sleep_imminent(&self, imminent: bool) {
if imminent {
// Raise the gate first so a refresh that re-checks it right before
// its IdP call (see `refresh_chain`) backs out instead of starting
// into the suspend window. Then hold the OS sleep acknowledgment
// until any refresh already in flight drains, so it can finish
// before the machine suspends rather than straddling it.
self.sleep_gate.raise();
self.hold_sleep_ack_until_refresh_drains(SLEEP_ACK_MAX_WAIT);
} else {
self.sleep_gate.lower("wake");
// End any in-progress dark-wake deferral run on a *genuine* full
// wake so the next dark wake starts with a fresh budget — but only
// if we are not still in a dark wake. macOS delivers
// `SYSTEM_HAS_POWERED_ON` (→ `DidWake`) for dark wakes too;
// unconditionally clearing here would reset the
// `DARK_WAKE_DEFER_MAX` budget on every dark-wake cycle so it could
// never exhaust, and the forced refresh would never run on a machine
// stuck in continuous dark wake. (`should_defer_for_dark_wake` also
// clears lazily under the same `!is_dark_wake()` condition.)
if !self.is_dark_wake() {
*self.dark_wake_defer_since.write() = None;
}
}
}
/// Mark an IdP refresh as starting. Paired with [`Self::end_refresh_in_flight`]
/// via [`InFlightGuard`]; see [`Self::hold_sleep_ack_until_refresh_drains`].
fn begin_refresh_in_flight(&self) {
self.refresh_in_flight.fetch_add(1, Ordering::SeqCst);
}
/// Mark an IdP refresh as finished. When the count returns to zero, wake any
/// sleep-ack waiter under the same lock it parks on, so a held OS sleep ack
/// is released the moment the exchange finishes rather than after the full
/// timeout. `fetch_sub` returns the *previous* value, so `== 1` is the
/// drop-to-zero edge. Notifying with no waiter parked is cheap and harmless.
fn end_refresh_in_flight(&self) {
if self.refresh_in_flight.fetch_sub(1, Ordering::SeqCst) == 1 {
let _drain = self.refresh_drain_lock.lock();
self.refresh_drain_cv.notify_all();
}
}
/// Block the calling thread — the OS power-listener callback, so this
/// delays the macOS `IOAllowPowerChange` ack / Linux `delay`-inhibitor
/// release — until in-flight IdP refreshes drain or `max` elapses.
///
/// A refresh already on the wire when sleep is requested would otherwise
/// straddle the suspend and, on a long sleep, lose its rotated successor
/// token — revoking the refresh-token family and forcing re-login. We never
/// abort the refresh; we briefly delay the suspend so it can finish first.
///
/// Bounded by `max` (see [`SLEEP_ACK_MAX_WAIT`]) so a hung refresh can't
/// hold the machine awake past the OS pre-sleep budget: on timeout the
/// suspend proceeds and the in-flight refresh is left to finish (a resulting
/// straddle is surfaced by `auth.refresh.suspend_spanned`).
fn hold_sleep_ack_until_refresh_drains(&self, max: StdDuration) {
let in_flight = self.refresh_in_flight.load(Ordering::SeqCst);
if in_flight == 0 {
return;
}
kigi_log::unified_log::warn(
"auth.sleep.refresh_in_flight_at_suspend",
None,
Some(serde_json::json!({ "in_flight": in_flight })),
);
let started = Instant::now();
{
let mut drain = self.refresh_drain_lock.lock();
// Loop on the atomic (the authoritative predicate) under the lock so
// a notify that races the park — or a spurious wake — can neither
// lose the signal nor over-wait. `InFlightGuard::drop` notifies when
// the count hits zero.
while self.refresh_in_flight.load(Ordering::SeqCst) > 0 {
let Some(remaining) = max.checked_sub(started.elapsed()) else {
break;
};
if remaining.is_zero() {
break;
}
let _ = self.refresh_drain_cv.wait_for(&mut drain, remaining);
}
}
let remaining = self.refresh_in_flight.load(Ordering::SeqCst);
kigi_log::unified_log::info(
"auth.sleep.refresh_drain",
None,
Some(serde_json::json!({
"in_flight_at_start": in_flight,
"in_flight_remaining": remaining,
"drained": remaining == 0,
"waited_ms": started.elapsed().as_millis() as u64,
"max_wait_ms": max.as_millis() as u64,
})),
);
}
pub(crate) fn is_sleep_gated(&self) -> bool {
self.sleep_gate.is_gated()
}
/// Whether the system is currently in a **dark wake** (see
/// [`kigi_system_power::PowerState`] for the canonical explanation of what a
/// dark wake is and why an IdP refresh must avoid one). `refresh_chain`
/// gates on [`Self::should_defer_for_dark_wake`], which wraps this with a
/// deferral bound.
///
/// Scoped to processes that actively listen for power events (local /
/// interactive): if the OS power listener was never started
/// (headless / datacenter), we skip the query — both because dark wake is
/// not a concern there and because a screenless Mac can read as a permanent
/// dark wake (no video capability), which would otherwise wedge refresh.
pub(crate) fn is_dark_wake(&self) -> bool {
#[cfg(test)]
if let Some(forced) = *self.dark_wake_override.lock() {
return forced;
}
if !self.power_listener_started.load(Ordering::Acquire) {
return false;
}
matches!(
kigi_system_power::current_power_state(),
kigi_system_power::PowerState::DarkWake
)
}
/// Whether `refresh_chain` should defer this refresh because the system is
/// in a dark wake — bounded so deferral can never be indefinite.
///
/// Tracks when the current unbroken run of dark-wake deferrals began (on two
/// clocks; see [`GateRaise`]). While inside the [`DARK_WAKE_DEFER_MAX`]
/// budget it returns `true` (defer). Once either clock passes the bound it
/// forces one refresh through (`false`) and resets the clock, so a machine
/// stuck reporting a continuous dark wake refreshes periodically instead of
/// deferring forever and logging the user out. A full wake clears the run
/// (here, or eagerly in [`Self::set_system_sleep_imminent`]).
pub(crate) fn should_defer_for_dark_wake(&self) -> bool {
if !self.is_dark_wake() {
// Full wake (or no signal): end any deferral run in progress.
if self.dark_wake_defer_since.read().is_some() {
*self.dark_wake_defer_since.write() = None;
}
return false;
}
let Some(raise) = *self.dark_wake_defer_since.read() else {
// First deferral of this dark-wake run: start the budget clock.
*self.dark_wake_defer_since.write() = Some(GateRaise::now());
return true;
};
let (mono, wall) = raise.elapsed();
if mono < DARK_WAKE_DEFER_MAX && wall < DARK_WAKE_DEFER_MAX {
return true;
}
// Budget exhausted: force this refresh through and reset the clock so a
// still-continuous dark wake defers afresh (up to DARK_WAKE_DEFER_MAX)
// before the next forced refresh, rather than abandoning deferral
// entirely.
*self.dark_wake_defer_since.write() = None;
kigi_log::unified_log::warn(
"auth.dark_wake.defer_budget_exhausted",
None,
Some(serde_json::json!({
"mono_elapsed_ms": mono.as_millis() as u64,
"wall_elapsed_ms": wall.as_millis() as u64,
})),
);
false
}
/// Force the [`AuthManager::is_dark_wake`] result in tests.
#[cfg(test)]
pub(crate) fn set_dark_wake_for_test(&self, dark: bool) {
*self.dark_wake_override.lock() = Some(dark);
}
/// Test hook: simulate an IdP refresh entering flight (mirrors
/// [`InFlightGuard::new`]).
#[cfg(test)]
pub(crate) fn test_enter_refresh_in_flight(&self) {
self.begin_refresh_in_flight();
}
/// Test hook: simulate an in-flight IdP refresh finishing (mirrors
/// [`InFlightGuard`]'s drop), waking a sleep-ack waiter.
#[cfg(test)]
pub(crate) fn test_exit_refresh_in_flight(&self) {
self.end_refresh_in_flight();
}
/// Test hook: run the bounded sleep-ack hold directly so tests can pass a
/// short bound instead of [`SLEEP_ACK_MAX_WAIT`].
#[cfg(test)]
pub(crate) fn test_hold_sleep_ack(&self, max: StdDuration) {
self.hold_sleep_ack_until_refresh_drains(max);
}
/// Start the OS power listener so sleep/wake drives the gate. Idempotent and
/// a no-op where the listener is unavailable. Call only from local /
/// interactive entrypoints, never datacenter server/headless.
pub fn start_system_power_listener(self: &Arc<Self>) {
// Claim the one-time startup so concurrent/duplicate calls don't
// double-register.
if self
.power_listener_started
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
return;
}
// Weak ref to avoid a manager <-> listener Arc cycle.
let weak = Arc::downgrade(self);
let listener = kigi_system_power::SystemPowerListener::start(move |event| {
if let Some(this) = weak.upgrade() {
let imminent = matches!(event, kigi_system_power::PowerEvent::WillSleep);
this.set_system_sleep_imminent(imminent);
}
});
let available = listener.is_some();
if available {
*self.power_listener.lock() = listener;
} else {
// Unavailable (unsupported OS / no logind / registration failure):
// release the guard so a later call can retry rather than being
// permanently no-op'd for this manager.
self.power_listener_started.store(false, Ordering::Release);
}
kigi_log::unified_log::info(
"auth.sleep.power_listener_init",
None,
Some(serde_json::json!({ "available": available })),
);
}
}
crates/codegen/kigi-shell/src/auth/manager_tests.rs
+3912
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/meta.rs
+40
View File
@@ -0,0 +1,40 @@
use serde::{Deserialize, Serialize};
/// Access gate from `grok_build_access_gate`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GateInfo {
pub message: String,
#[serde(default)]
pub url: Option<String>,
#[serde(default)]
pub label: Option<String>,
}
/// Typed auth metadata passed from the shell to the pager via ACP.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AuthMeta {
#[serde(default)]
pub email: Option<String>,
#[serde(default)]
pub auth_mode: Option<String>,
/// Team principal UUID when the session is a team login (`None` for personal).
#[serde(default)]
pub team_id: Option<String>,
#[serde(default)]
pub team_name: Option<String>,
#[serde(default)]
pub is_zdr: bool,
#[serde(default)]
pub team_role: Option<String>,
#[serde(default)]
pub coding_data_retention_opt_out: bool,
#[serde(default)]
pub show_resolved_model: Option<bool>,
/// `Some` = user is blocked; `None` = user has access.
#[serde(default)]
pub gate: Option<GateInfo>,
/// User-friendly display name for the current subscription tier
/// (e.g. "SuperGrok Heavy", "X Premium", "Free"). From CCP `/settings`.
#[serde(default)]
pub subscription_tier: Option<String>,
}
crates/codegen/kigi-shell/src/auth/mod.rs
+42
View File
@@ -0,0 +1,42 @@
pub(crate) mod attribution;
mod config;
pub mod credential_provider;
#[path = "devbox_login_stub.rs"]
pub(crate) mod devbox_login;
pub mod device_code;
pub mod error;
mod external_auth;
mod flow;
mod jwt;
pub(crate) mod manager;
mod model;
pub mod oidc;
pub(crate) mod recovery;
pub(crate) mod refresh;
mod storage;
pub(crate) mod token_type;
pub(crate) use config::LEGACY_AUTH_SCOPE;
pub use config::{
ForceLoginTeam, GrokComConfig, OAuth2ProviderConfig, OidcAuthConfig, PreferredAuthMethod,
XAI_OAUTH2_ISSUER, is_xai_oauth2_issuer, xai_oauth2_issuer,
};
pub(crate) use external_auth::{parse_output, refresh_with_command};
pub(crate) use flow::{
AuthChannels, run_auth_flow, run_auth_flow_with_stderr_bridge,
try_ensure_session_noninteractive,
};
pub use flow::{
AuthUrlInfo, AuthUrlMode, LoginTransportOverride, LogoutResult, ensure_authenticated,
ensure_authenticated_or_noninteractive, ensure_authenticated_with_override, perform_logout,
run_cli_login, run_cli_logout, try_ensure_fresh_auth,
};
pub use jwt::{is_jwt_expired_or_near, parse_jwt_expiration};
mod meta;
pub use error::{AuthError, RefreshTokenError, RefreshTokenFailedReason};
pub use manager::{AuthManager, shared_api_key_provider};
pub use meta::{AuthMeta, GateInfo};
pub use model::{AuthMode, GrokAuth, lookup_auth};
pub(crate) use model::{TOKEN_TTL, UserInfo, is_expired, token_suffix};
pub use storage::{
clear_api_key, read_api_key, read_auth_json, read_token_by_scope, store_api_key,
};
crates/codegen/kigi-shell/src/auth/model.rs
+489
View File
@@ -0,0 +1,489 @@
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use super::is_xai_oauth2_issuer;
pub(crate) const TOKEN_TTL: Duration = Duration::days(30);
const DEFAULT_EARLY_INVALIDATION_SECS: u64 = 300; // 5 minutes
/// Legacy auth.json scope key. Fallback for old devbox auth files.
pub(super) const LEGACY_SCOPE: &str = "https://accounts.x.ai/sign-in";
/// auth.json scope key for plain API key auth (desktop login, `grok login --api-key`).
pub const API_KEY_SCOPE: &str = "xai::api_key";
const BLOCKED_REASON_NO_LOGS: &str = "BLOCKED_REASON_NO_LOGS";
const BLOCKED_REASON_NO_LOGS_MODERATED: &str = "BLOCKED_REASON_NO_LOGS_MODERATED";
/// Token provenance (debugging/auth.json only -- no code branches on this).
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AuthMode {
/// Deprecated. Kept for deserializing old auth.json files.
#[serde(alias = "grok")]
WebLogin,
/// OIDC or OAuth2 interactive login via customer IdP
#[serde(alias = "oidc")]
Oidc,
/// External auth provider binary
External,
/// Plain API key (e.g. from grok-desktop login or `grok login --api-key`)
ApiKey,
}
/// Wire value of `principal_type` for team OAuth principals (capitalized by
/// the auth service). Single source for every comparison site.
pub(crate) const TEAM_PRINCIPAL_TYPE: &str = "Team";
#[derive(Clone, Serialize, Deserialize)]
pub struct GrokAuth {
pub key: String,
pub auth_mode: AuthMode,
pub create_time: DateTime<Utc>,
pub user_id: String,
pub email: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub first_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub profile_image_asset_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub team_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub team_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub team_role: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub organization_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub organization_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub organization_role: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub user_blocked_reason: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub team_blocked_reasons: Vec<String>,
#[serde(default)]
pub coding_data_retention_opt_out: bool,
/// Deprecated. Kept for deserializing existing auth.json files.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub has_grok_code_access: Option<bool>,
/// Refresh token (OIDC/OAuth2 or external provider).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub refresh_token: Option<String>,
/// Server-provided expiration (from OIDC `expires_in`).
/// When present, takes precedence over the hardcoded `TOKEN_TTL`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub expires_at: Option<DateTime<Utc>>,
/// Issuer URL that issued this token. For OIDC credentials it drives
/// refresh via discovery; for external-provider credentials it is the
/// provider's `issuer` claim. In both modes an x.ai issuer marks the
/// credential first-party (`is_xai_auth`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc_issuer: Option<String>,
/// OIDC client_id used to obtain this token (needed for refresh).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc_client_id: Option<String>,
}
impl std::fmt::Debug for GrokAuth {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GrokAuth")
.field("key", &token_suffix(&self.key))
.field("auth_mode", &self.auth_mode)
.field("user_id", &self.user_id)
.field("expires_at", &self.expires_at)
.field(
"refresh_token",
&self.refresh_token.as_deref().map(token_suffix),
)
.finish_non_exhaustive()
}
}
impl GrokAuth {
/// Seconds since this credential was minted. Negative when the local
/// clock stepped back past `create_time` (NTP correction, VM restore, or
/// a sibling machine's clock via an adopted auth.json) — `create_time`
/// is always stamped from the minting machine's local clock.
pub(crate) fn mint_age_seconds(&self) -> i64 {
Utc::now()
.signed_duration_since(self.create_time)
.num_seconds()
}
/// `true` when the token comes from a first-party xAI account —
/// either an OIDC login against https://auth.x.ai (or the local-dev
/// equivalent), or an external auth provider that declared an xAI
/// issuer for its token.
///
/// The issuer is a client-side hint, not a trust assertion: everything
/// it unlocks still authenticates the actual token server-side, and it
/// never influences endpoints.
pub fn is_xai_auth(&self) -> bool {
match self.auth_mode {
AuthMode::Oidc | AuthMode::External => self
.oidc_issuer
.as_deref()
.is_some_and(is_xai_oauth2_issuer),
AuthMode::ApiKey | AuthMode::WebLogin => false,
}
}
/// `true` when this auth can access grok.com managed MCP connectors.
pub fn is_managed_mcp_eligible(&self) -> bool {
self.is_xai_auth() || self.auth_mode == AuthMode::WebLogin
}
/// Whether this credential can access `supported_in_api: false` models.
///
/// Session logins (WebLogin, OIDC — including enterprise issuers) always
/// qualify; external-provider credentials qualify only when first-party
/// (`is_xai_auth`), matching the built-in devbox login they replace.
/// Plain API keys never do.
pub fn is_session_auth(&self) -> bool {
match self.auth_mode {
AuthMode::WebLogin | AuthMode::Oidc => true,
AuthMode::External => self.is_xai_auth(),
AuthMode::ApiKey => false,
}
}
pub fn is_team_principal(&self) -> bool {
self.principal_type.as_deref() == Some(TEAM_PRINCIPAL_TYPE) && self.team_id.is_some()
}
/// `true` when the team has Zero Data Retention (ZDR) enabled.
pub fn is_zdr_team(&self) -> bool {
self.team_blocked_reasons
.iter()
.any(|r| r == BLOCKED_REASON_NO_LOGS || r == BLOCKED_REASON_NO_LOGS_MODERATED)
}
/// `true` when the team has ZDR or the user opted out of coding data
/// retention. Use this for trace-upload and research-data gates.
/// Product analytics (`telemetry_enabled`) and user-facing sync
/// features should use `is_zdr_team()` directly.
pub fn is_data_collection_disabled(&self) -> bool {
self.is_zdr_team() || self.coding_data_retention_opt_out
}
/// Carry `/user`-derived fields from a previous auth so refresh rebuilds don't drop them.
pub(crate) fn carry_user_profile_from(&mut self, prev: &GrokAuth) {
self.user_id = prev.user_id.clone();
self.email = prev.email.clone();
self.principal_type = prev.principal_type.clone();
self.principal_id = prev.principal_id.clone();
self.team_id = prev.team_id.clone();
self.team_name = prev.team_name.clone();
self.team_role = prev.team_role.clone();
self.organization_id = prev.organization_id.clone();
self.organization_name = prev.organization_name.clone();
self.organization_role = prev.organization_role.clone();
self.user_blocked_reason = prev.user_blocked_reason.clone();
self.team_blocked_reasons = prev.team_blocked_reasons.clone();
self.coding_data_retention_opt_out = prev.coding_data_retention_opt_out;
}
}
impl Default for GrokAuth {
fn default() -> Self {
Self {
key: String::new(),
auth_mode: AuthMode::Oidc,
create_time: Utc::now(),
user_id: String::new(),
email: None,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: None,
oidc_issuer: None,
oidc_client_id: None,
}
}
}
#[cfg(test)]
impl GrokAuth {
/// Returns a `GrokAuth` with sensible defaults for tests. Override fields
/// with struct update syntax:
/// ```ignore
/// GrokAuth { key: "my-key".into(), ..GrokAuth::test_default() }
/// ```
pub fn test_default() -> Self {
Self {
key: "test-key".into(),
user_id: "test-user".into(),
..Default::default()
}
}
}
pub(crate) type AuthStore = BTreeMap<String, GrokAuth>;
/// User information from the cli-chat-proxy `GET /v1/user` endpoint.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UserInfo {
pub(crate) user_id: String,
#[serde(default)]
pub(super) email: Option<String>,
#[serde(default)]
pub(super) first_name: Option<String>,
#[serde(default)]
pub(super) last_name: Option<String>,
#[serde(default)]
pub(super) profile_image_asset_id: Option<String>,
#[serde(default)]
pub(super) principal_type: Option<String>,
#[serde(default)]
pub(super) principal_id: Option<String>,
#[serde(default)]
pub(super) team_id: Option<String>,
#[serde(default)]
pub(super) team_name: Option<String>,
#[serde(default)]
pub(super) team_role: Option<String>,
#[serde(default)]
pub(super) organization_id: Option<String>,
#[serde(default)]
pub(super) organization_name: Option<String>,
#[serde(default)]
pub(super) organization_role: Option<String>,
#[serde(default)]
pub(super) user_blocked_reason: Option<String>,
#[serde(default)]
pub(super) team_blocked_reasons: Option<Vec<String>>,
#[serde(default)]
pub(super) coding_data_retention_opt_out: Option<bool>,
/// Live subscription tier from the backend (only present when
/// `?include=subscription` is passed to `/user`).
#[serde(default)]
pub(crate) subscription_tier: Option<String>,
}
/// Last 12 chars of a token string, safe for diagnostic logging.
/// Uses the tail because JWT access tokens all share the same base64
/// header prefix (`eyJ0eXAiOiJh…`); the tail (signature bytes) is
/// unique per token and makes `key_changed` / `is_stale_snapshot`
/// diagnostics meaningful.
pub(crate) fn token_suffix(t: &str) -> &str {
let len = t.len();
if len > 12 { &t[len - 12..] } else { t }
}
/// Look up auth from the store by scope key.
///
/// Legacy `WebLogin` tokens (from the pre-OIDC `grok login --legacy`
/// flow) are skipped — they are validated via a per-request DB lookup
/// server-side which fails at high volume. Skipping them here forces
/// affected users to re-authenticate via OIDC on next launch.
pub fn lookup_auth(map: &AuthStore, scope: &str) -> Option<GrokAuth> {
let auth = map.get(scope).cloned().or_else(|| {
if scope == LEGACY_SCOPE {
None
} else {
map.get(LEGACY_SCOPE).cloned()
}
})?;
if auth.auth_mode == AuthMode::WebLogin {
tracing::info!("auth: ignoring legacy WebLogin token — re-authentication required");
return None;
}
Some(auth)
}
/// Early-invalidation buffer. Override with `KIGI_AUTH_EARLY_INVALIDATION_SECS`
/// for testing (e.g. `=5` to shrink the buffer to 5 seconds).
pub(super) fn early_invalidation() -> Duration {
std::env::var("KIGI_AUTH_EARLY_INVALIDATION_SECS")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.map(|s| Duration::seconds(s as i64))
.unwrap_or_else(|| Duration::seconds(DEFAULT_EARLY_INVALIDATION_SECS as i64))
}
pub(crate) fn is_expired(auth: &GrokAuth) -> bool {
is_expired_with_buffer(auth, early_invalidation())
}
/// Like [`is_expired`] but with an explicit pre-expiry buffer. Pass
/// `Duration::zero()` for actual (hard) expiry — the instant the token would
/// really be rejected on the wire, with no early-invalidation margin.
pub(crate) fn is_expired_with_buffer(auth: &GrokAuth, buffer: Duration) -> bool {
if let Some(expires_at) = auth.expires_at {
Utc::now() >= (expires_at - buffer)
} else {
let age = Utc::now().signed_duration_since(auth.create_time);
age >= (TOKEN_TTL - buffer)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_auth(mode: AuthMode) -> GrokAuth {
GrokAuth {
key: "k".into(),
auth_mode: mode,
create_time: Utc::now(),
user_id: "u".into(),
email: None,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: None,
oidc_issuer: None,
oidc_client_id: None,
}
}
#[test]
fn is_xai_auth_matrix() {
use crate::auth::XAI_OAUTH2_ISSUER;
let with_issuer = |mode: AuthMode, issuer: Option<&str>| GrokAuth {
oidc_issuer: issuer.map(str::to_owned),
..make_auth(mode)
};
// Only Oidc/External qualify, and only with an x.ai issuer.
assert!(with_issuer(AuthMode::Oidc, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
assert!(with_issuer(AuthMode::External, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
assert!(!with_issuer(AuthMode::Oidc, None).is_xai_auth());
assert!(!with_issuer(AuthMode::External, None).is_xai_auth());
assert!(!with_issuer(AuthMode::Oidc, Some("https://idp.acme.example")).is_xai_auth());
assert!(!with_issuer(AuthMode::External, Some("https://idp.acme.example")).is_xai_auth());
// ApiKey / WebLogin stay false even with an x.ai issuer set.
assert!(!with_issuer(AuthMode::ApiKey, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
assert!(!with_issuer(AuthMode::WebLogin, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
}
#[test]
fn is_session_auth_requires_first_party_for_external() {
use crate::auth::XAI_OAUTH2_ISSUER;
let with_issuer = |mode: AuthMode, issuer: Option<&str>| GrokAuth {
oidc_issuer: issuer.map(str::to_owned),
..make_auth(mode)
};
// Session logins qualify regardless of issuer (incl. enterprise OIDC).
assert!(with_issuer(AuthMode::WebLogin, None).is_session_auth());
assert!(with_issuer(AuthMode::Oidc, None).is_session_auth());
assert!(with_issuer(AuthMode::Oidc, Some("https://idp.acme.example")).is_session_auth());
// External qualifies only when first-party (devbox-login parity).
assert!(with_issuer(AuthMode::External, Some(XAI_OAUTH2_ISSUER)).is_session_auth());
assert!(!with_issuer(AuthMode::External, None).is_session_auth());
assert!(
!with_issuer(AuthMode::External, Some("https://idp.acme.example")).is_session_auth()
);
// Plain API keys never do.
assert!(!with_issuer(AuthMode::ApiKey, Some(XAI_OAUTH2_ISSUER)).is_session_auth());
}
#[test]
fn lookup_auth_skips_weblogin_on_primary_scope() {
let mut map = AuthStore::new();
map.insert("scope".into(), make_auth(AuthMode::WebLogin));
assert!(lookup_auth(&map, "scope").is_none());
}
#[test]
fn lookup_auth_skips_weblogin_on_legacy_fallback() {
let mut map = AuthStore::new();
map.insert(LEGACY_SCOPE.into(), make_auth(AuthMode::WebLogin));
assert!(lookup_auth(&map, "other-scope").is_none());
}
#[test]
fn lookup_auth_returns_oidc_token() {
let mut map = AuthStore::new();
map.insert("scope".into(), make_auth(AuthMode::Oidc));
assert!(lookup_auth(&map, "scope").is_some());
}
#[test]
fn lookup_auth_returns_api_key_token() {
let mut map = AuthStore::new();
map.insert("scope".into(), make_auth(AuthMode::ApiKey));
assert!(lookup_auth(&map, "scope").is_some());
}
/// subscriptionTier present → deserializes to Some.
#[test]
fn user_info_subscription_tier_present() {
let json = r#"{
"userId": "u1",
"subscriptionTier": "SuperGrokPro"
}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert_eq!(info.subscription_tier.as_deref(), Some("SuperGrokPro"));
}
/// subscriptionTier absent → deserializes to None (backwards compat).
#[test]
fn user_info_subscription_tier_absent() {
let json = r#"{"userId": "u1"}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert!(info.subscription_tier.is_none());
}
/// subscriptionTier null → deserializes to None.
#[test]
fn user_info_subscription_tier_null() {
let json = r#"{"userId": "u1", "subscriptionTier": null}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert!(info.subscription_tier.is_none());
}
/// subscriptionTier empty string → deserializes to Some("").
/// The paywall poller treats this as "no subscription" (line 230:
/// `Some(tier) if !tier.is_empty()`) and keeps polling.
#[test]
fn user_info_subscription_tier_empty_string() {
let json = r#"{"userId": "u1", "subscriptionTier": ""}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert_eq!(info.subscription_tier.as_deref(), Some(""));
}
}
crates/codegen/kigi-shell/src/auth/oidc/login.rs
+701
View File
@@ -0,0 +1,701 @@
//! Interactive login orchestration: callback HTTP server, browser
//! handoff, stdin paste fallback, race between the two.
//!
//! Cross-references [`super::protocol`] for OIDC mechanics and
//! [`super::super::AuthManager`] for credential persistence.
use std::collections::HashMap;
use std::io::IsTerminal;
use std::sync::Arc;
use axum::{
Router,
extract::{Query, State},
http::{Method, StatusCode},
response::Html,
routing::get,
};
use tokio::net::TcpListener;
use super::super::config::{GrokComConfig, OidcAuthConfig};
use super::super::{AuthManager, GrokAuth};
use super::protocol::{
OidcError, build_authorize_url, build_grok_auth, discover, enforce_login_principal,
exchange_code, extract_user_info, generate_pkce, login_principal_policy,
peek_access_token_principal, peek_access_token_principal_id, validate_state,
};
/// Maximum time to wait for the browser OAuth callback (or manual paste of the code).
/// 10 minutes is long enough for users who step away briefly during login.
const AUTH_CALLBACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600);
/// Parse user-pasted input into `(code, state)`.
///
/// Accepts two formats:
/// 1. Full callback URL: `http://127.0.0.1:PORT/callback?code=XXX&state=YYY`
/// 2. Bare authorization code: `abc123`
fn parse_pasted_input(input: &str) -> Result<Callback, OidcError> {
let input = input.trim();
if input.is_empty() {
return Err(OidcError::InvalidPastedInput("empty input".into()));
}
if let Ok(url) = url::Url::parse(input) {
let params: HashMap<String, String> = url.query_pairs().into_owned().collect();
if let Some(code) = params.get("code") {
let state = params.get("state").cloned().unwrap_or_default();
return Ok(Callback {
code: code.clone(),
state,
});
}
if let Some(error) = params.get("error") {
let desc = params.get("error_description").cloned().unwrap_or_default();
return Err(OidcError::CallbackAuthFailed(if desc.is_empty() {
error.clone()
} else {
format!("{error}: {desc}")
}));
}
return Err(OidcError::InvalidPastedInput(
"URL has no 'code' query parameter".into(),
));
}
Ok(Callback {
code: input.to_owned(),
state: String::new(),
})
}
/// Render a styled callback page shown in the browser after the OAuth redirect.
pub(crate) fn callback_page(title: &str, message: &str, is_success: bool) -> String {
let icon = if is_success {
// Grok logo
r#"<svg xmlns="http://www.w3.org/2000/svg" width="48" height="48" fill="none" viewBox="0 0 33 33"><path fill="currentColor" d="m13.237 21.04 11.082-8.19c.543-.4 1.32-.244 1.578.38 1.363 3.288.754 7.241-1.957 9.955-2.71 2.714-6.482 3.31-9.93 1.954l-3.765 1.745c5.401 3.697 11.96 2.782 16.059-1.324 3.251-3.255 4.258-7.692 3.317-11.693l.008.009c-1.365-5.878.336-8.227 3.82-13.031q.123-.17.247-.345l-4.585 4.59v-.014L13.234 21.044M10.95 23.031c-3.877-3.707-3.208-9.446.1-12.755 2.446-2.449 6.454-3.448 9.952-1.979L24.76 6.56c-.677-.49-1.545-1.017-2.54-1.387A12.465 12.465 0 0 0 8.675 7.901c-3.519 3.523-4.625 8.94-2.725 13.561 1.42 3.454-.907 5.898-3.251 8.364-.83.874-1.664 1.749-2.335 2.674l10.583-9.466"/></svg>"#
} else {
// X circle
r#"<svg width="48" height="48" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" style="color:#ef4444"><circle cx="12" cy="12" r="10"/><line x1="15" y1="9" x2="9" y2="15"/><line x1="9" y1="9" x2="15" y2="15"/></svg>"#
};
format!(
r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width,initial-scale=1"/>
<meta name="color-scheme" content="light dark"/>
<title>{title}</title>
<style>
*{{margin:0;padding:0;box-sizing:border-box}}
body{{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,Helvetica,Arial,sans-serif;
display:flex;align-items:center;justify-content:center;min-height:100vh;
background:#0a0a0a;color:#e5e5e5}}
.card{{text-align:center;display:flex;flex-direction:column;align-items:center;gap:16px;padding:48px}}
h1{{font-size:18px;font-weight:600}}
p{{font-size:14px;color:#a3a3a3}}
@media(prefers-color-scheme:light){{
body{{background:#fafafa;color:#171717}}
p{{color:#525252}}
}}
</style>
</head>
<body>
<div class="card">
{icon}
<h1>{title}</h1>
<p>{message}</p>
</div>
</body>
</html>"#,
title = title,
icon = icon,
message = message,
)
}
/// Build the axum router for the OIDC loopback callback server.
fn build_callback_router(tx: tokio::sync::mpsc::Sender<CallbackResult>) -> Router {
let cors =
crate::auth::config::accounts_app_cors_layer(Method::GET).allow_private_network(true);
Router::new()
.route("/callback", get(handle_callback))
.layer(cors)
.with_state(tx)
}
async fn handle_callback(
State(tx): State<tokio::sync::mpsc::Sender<CallbackResult>>,
Query(params): Query<HashMap<String, String>>,
) -> (StatusCode, Html<String>) {
let result = parse_callback_params(&params);
let response = callback_response(&result);
if let Err(e) = tx.try_send(result) {
tracing::error!(?e, "OIDC: callback channel send failed; auth will time out");
}
response
}
fn parse_callback_params(params: &HashMap<String, String>) -> CallbackResult {
if let Some(code) = params.get("code") {
let state = params.get("state").cloned().unwrap_or_default();
tracing::debug!(state = %state, "OIDC: received code via loopback callback");
return Ok(Callback {
code: code.clone(),
state,
});
}
let error = params.get("error").cloned().unwrap_or_default();
let desc = params.get("error_description").cloned().unwrap_or_default();
tracing::error!(error = %error, desc = %desc, "OIDC: IdP returned error");
Err(if desc.is_empty() {
error
} else {
format!("{error}: {desc}")
})
}
fn callback_response(result: &CallbackResult) -> (StatusCode, Html<String>) {
let (title, message) = match result {
Ok(_) => (
"Signed in",
"You can close this window and return to Grok Build.",
),
Err(_) => ("Access denied", "Close this window and try again."),
};
(
StatusCode::OK,
Html(callback_page(title, message, result.is_ok())),
)
}
/// Wait until stdin has data or `tx` is closed. Returns `false` if closed.
#[cfg(unix)]
fn wait_for_stdin_or_closed(
stdin: &std::io::Stdin,
tx: &tokio::sync::mpsc::Sender<CallbackResult>,
) -> bool {
use std::os::unix::io::AsRawFd;
let fd = stdin.as_raw_fd();
loop {
if tx.is_closed() {
return false;
}
let ready = unsafe {
let mut fds = std::mem::zeroed::<libc::pollfd>();
fds.fd = fd;
fds.events = libc::POLLIN;
libc::poll(&mut fds, 1, 200)
};
if ready > 0 {
return true;
}
}
}
fn spawn_stdin_reader(tx: tokio::sync::mpsc::Sender<CallbackResult>) {
tokio::task::spawn_blocking(move || {
use std::io::BufRead;
let stdin = std::io::stdin();
let mut buf = String::new();
loop {
#[cfg(unix)]
if !wait_for_stdin_or_closed(&stdin, &tx) {
tracing::debug!("OIDC: stdin reader exiting, channel closed");
return;
}
#[cfg(not(unix))]
if tx.is_closed() {
tracing::debug!("OIDC: stdin reader exiting, channel closed");
return;
}
buf.clear();
let mut handle = stdin.lock();
match handle.read_line(&mut buf) {
Ok(0) => return,
Ok(_) => {}
Err(_) => return,
}
drop(handle);
let trimmed = buf.trim().to_owned();
if trimmed.is_empty() {
continue;
}
match parse_pasted_input(&trimmed) {
Ok(result) => {
tracing::debug!("OIDC: received code via stdin paste");
let _ = tx.blocking_send(Ok(result));
return;
}
Err(OidcError::InvalidPastedInput(msg)) => {
tracing::debug!(input = %msg, "OIDC: invalid stdin paste, retrying");
eprintln!(" Invalid input: {msg}. Try again:");
}
Err(e) => {
tracing::warn!(error = %e, "OIDC: stdin paste returned auth error");
let _ = tx.blocking_send(Err(e.to_string()));
return;
}
}
}
});
}
/// Race loopback callback against manual paste from `code_rx`.
async fn race_callback_and_client_ui(
listener: TcpListener,
code_rx: &mut tokio::sync::mpsc::Receiver<String>,
) -> anyhow::Result<Callback> {
tracing::debug!("OIDC: waiting for auth code (loopback + client paste)");
let (tx, mut rx) = tokio::sync::mpsc::channel::<CallbackResult>(1);
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let app = build_callback_router(tx.clone());
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await;
});
// Bridge client paste input into the callback channel.
let client_tx = tx.clone();
let client_bridge = async {
while let Some(code) = code_rx.recv().await {
match parse_pasted_input(&code) {
Ok(result) => {
tracing::debug!("OIDC: received code via client paste");
let _ = client_tx.send(Ok(result)).await;
return;
}
Err(e) => {
tracing::debug!(error = %e, "OIDC: invalid client paste input");
}
}
}
};
drop(tx);
let result = tokio::select! {
r = tokio::time::timeout(AUTH_CALLBACK_TIMEOUT, rx.recv()) => {
r.map_err(|_| anyhow::Error::new(OidcError::CallbackTimeout))?
.ok_or_else(|| anyhow::Error::new(OidcError::CallbackChannelClosed))?
}
_ = client_bridge => {
rx.recv().await
.ok_or_else(|| anyhow::Error::new(OidcError::CallbackChannelClosed))?
}
};
let _ = shutdown_tx.send(());
let _ = server.await;
result.map_err(|e| anyhow::Error::new(OidcError::CallbackAuthFailed(e)))
}
/// Race loopback callback against stdin paste.
async fn race_callback_and_stdin(
listener: TcpListener,
enable_stdin: bool,
) -> anyhow::Result<Callback> {
tracing::debug!(
enable_stdin = enable_stdin,
"OIDC: waiting for auth code (loopback + stdin)"
);
let (tx, mut rx) = tokio::sync::mpsc::channel::<CallbackResult>(1);
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let app = build_callback_router(tx.clone());
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await;
});
if enable_stdin {
spawn_stdin_reader(tx.clone());
}
drop(tx);
let result = tokio::time::timeout(AUTH_CALLBACK_TIMEOUT, rx.recv())
.await
.map_err(|_| {
// "10 minutes" must match AUTH_CALLBACK_TIMEOUT above
tracing::error!("auth: timed out after 10 minutes waiting for auth code");
anyhow::Error::new(OidcError::CallbackTimeout)
})?
.ok_or_else(|| {
tracing::error!(
"OIDC: callback channel closed, no code received from loopback or stdin"
);
anyhow::Error::new(OidcError::CallbackChannelClosed)
})?;
let _ = shutdown_tx.send(());
let _ = server.await;
result.map_err(|e| anyhow::Error::new(OidcError::CallbackAuthFailed(e)))
}
/// Run the full OIDC login flow: discovery → PKCE → browser → callback → token exchange → persist.
pub async fn run_login_flow(
config: &GrokComConfig,
auth_manager: &Arc<AuthManager>,
channels: Option<super::super::flow::AuthChannels>,
) -> anyhow::Result<(GrokAuth, bool)> {
let oidc = config
.oidc
.as_ref()
.ok_or_else(|| anyhow::Error::new(OidcError::NotConfigured))?;
run_login_flow_with_config(oidc, auth_manager, channels).await
}
/// Run the OIDC login flow with an explicit [`OidcAuthConfig`].
///
/// Also used by the OAuth2 provider path via [`OAuth2ProviderConfig::as_oidc`].
///
/// The flow races two input paths:
/// - **Path A**: A loopback HTTP server on `127.0.0.1` that receives the IdP redirect.
/// - **Path B**: Stdin paste — the user manually pastes the callback URL or bare auth code.
///
/// Path B is essential for remote VMs where the browser runs on a different machine
/// and the `127.0.0.1` redirect cannot reach the CLI process.
/// * `channels` — `Some`: pushes the auth URL to the TUI and receives pasted codes.
/// `None`: prints to stderr / reads stdin (CLI mode).
pub async fn run_login_flow_with_config(
oidc: &OidcAuthConfig,
auth_manager: &Arc<AuthManager>,
channels: Option<super::super::flow::AuthChannels>,
) -> anyhow::Result<(GrokAuth, bool)> {
tracing::info!(issuer = %oidc.issuer, client_id = %oidc.client_id, "OIDC: starting login flow");
// Ensure jsonwebtoken CryptoProvider is installed (required for JWT validation).
jsonwebtoken::crypto::CryptoProvider::install_default(
&jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER,
)
.ok();
let discovery = discover(&oidc.issuer).await?;
let pkce = generate_pkce();
let state = uuid::Uuid::now_v7().to_string();
let nonce = uuid::Uuid::now_v7().to_string();
// In local-dev mode, use a fixed callback port so the redirect_uri is stable
// and can be pre-registered with the local OAuth2 provider. In production the
// OS picks a random available port.
let callback_port: u16 = if super::super::config::use_local_auth() {
56121
} else {
0
};
let listener = TcpListener::bind(("127.0.0.1", callback_port))
.await
.map_err(|e| anyhow::Error::new(OidcError::BindLoopback(e.to_string())))?;
let port = listener.local_addr()?.port();
let redirect_uri = format!("http://127.0.0.1:{}/callback", port);
let oauth2 = auth_manager.grok_com_config().oauth2.as_ref();
let auth_url = build_authorize_url(
oidc,
oauth2,
&discovery,
&redirect_uri,
&pkce,
&state,
&nonce,
);
tracing::debug!(port = port, redirect_uri = %redirect_uri, "OIDC: callback server bound");
let (url_tx, code_rx) = match channels {
Some(ch) => (ch.url_tx, Some(ch.code_rx)),
None => (None, None),
};
let has_client_ui = code_rx.is_some();
if has_client_ui {
// Client provides its own auth UI; just open the browser.
if let Err(e) = webbrowser::open(&auth_url) {
tracing::debug!(error = %e, "OIDC: failed to open browser");
}
} else {
// No client UI — print to stderr.
eprintln!();
let provider_label = if oidc.issuer == super::super::config::XAI_OAUTH2_ISSUER {
"Grok".to_owned()
} else {
oidc.issuer.clone()
};
eprintln!("Signing in with {}...", provider_label);
eprintln!();
if let Err(e) = webbrowser::open(&auth_url) {
tracing::debug!(error = %e, "OIDC: failed to open browser");
}
eprintln!("Open this URL to sign in:");
eprintln!(" {}", auth_url);
}
let use_stdin = !has_client_ui && std::io::stdin().is_terminal();
if use_stdin {
eprintln!();
eprintln!("Paste the URL here if it doesn't connect:");
}
// Push auth URL to the TUI via oneshot.
if let Some(tx) = url_tx {
let _ = tx.send(super::super::flow::AuthUrlInfo {
url: auth_url.clone(),
mode: super::super::flow::AuthUrlMode::Loopback,
});
}
let Callback {
code,
state: received_state,
} = if let Some(mut rx) = code_rx {
// Client UI: race loopback against manual paste via code_rx.
race_callback_and_client_ui(listener, &mut rx).await?
} else {
// No client UI: race loopback against stdin paste.
race_callback_and_stdin(listener, use_stdin).await?
};
// Validate state (skip for bare code paste where state is empty)
if !received_state.is_empty() {
validate_state(&state, &received_state)?;
}
let tokens = exchange_code(
&discovery.token_endpoint,
&code,
&redirect_uri,
&oidc.client_id,
&pkce.code_verifier,
)
.await?;
tracing::info!(
has_refresh = tokens.refresh_token.is_some(),
expires_in = ?tokens.expires_in,
"OIDC: token exchange complete"
);
// Resolve the actual principal chosen on the consent screen.
//
// The shell's config may not have principal_type set (personal login),
// but the user might pick "Team" on the consent screen. The server
// encodes the chosen principal in the access token JWT. If the config
// doesn't specify a principal, peek at the token to discover it.
let token_principal = peek_access_token_principal(&tokens.access_token);
// The authorize URL only pre-selects; verify the token's principal here.
// Match the principal id even if `principal_type` is absent.
let principal_policy = login_principal_policy(auth_manager.grok_com_config());
enforce_login_principal(
principal_policy.as_ref(),
peek_access_token_principal_id(&tokens.access_token).as_deref(),
)?;
let (resolved_principal_type, resolved_principal_id, resolved_team_id) = {
let cfg_pt = oauth2.and_then(|cfg| cfg.principal_type.clone());
let cfg_pid = oauth2.and_then(|cfg| cfg.principal_id.clone());
if cfg_pt.is_some() {
(cfg_pt, cfg_pid, None)
} else if let Some((pt, pid, tid)) = token_principal {
tracing::info!(
principal_type = %pt,
principal_id = %pid,
team_id = ?tid,
"OIDC: resolved principal from access token"
);
(Some(pt), Some(pid), tid)
} else {
(cfg_pt, cfg_pid, None)
}
};
let user_info = extract_user_info(
tokens.id_token.as_deref(),
&discovery,
&oidc.issuer,
&oidc.client_id,
&nonce,
resolved_principal_type.as_deref(),
resolved_principal_id.as_deref(),
resolved_team_id,
)
.await?;
tracing::debug!(user_id = %user_info.user_id, "OIDC: extracted user info");
let mut auth = build_grok_auth(tokens, user_info, &oidc.issuer, &oidc.client_id);
auth_manager.enrich_auth_inline(&mut auth).await;
let auth = auth_manager
.update(auth)
.await
.map_err(|e| anyhow::Error::new(OidcError::SaveAuth(e.to_string())))?;
tracing::info!(user_id = %auth.user_id, "OIDC: login complete, credentials saved");
Ok((auth, true))
}
/// Successful OIDC callback payload.
#[derive(Debug, PartialEq, Eq)]
struct Callback {
code: String,
state: String,
}
/// Result from the OIDC callback: either a [`Callback`] or an IdP error message.
type CallbackResult = Result<Callback, String>;
#[cfg(test)]
mod tests {
use super::super::test_helpers::*;
use super::*;
/// End-to-end test: mock IdP + full login flow with code arriving via loopback.
/// Exercises discovery → PKCE → race_callback_and_stdin → token exchange → user info → persist.
#[tokio::test]
async fn full_login_flow_via_race() {
ensure_crypto_provider();
let (issuer, idp_server) = start_mock_idp().await;
let temp_dir = tempfile::tempdir().unwrap();
// Dead proxy port: inline `/user` enrichment fails fast in tests.
let dead_proxy = {
let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
format!("http://127.0.0.1:{}", l.local_addr().unwrap().port())
};
let auth_manager = Arc::new(
AuthManager::new(temp_dir.path(), GrokComConfig::default())
.with_proxy_base_url(&dead_proxy),
);
let oidc_cfg = OidcAuthConfig {
issuer: issuer.clone(),
client_id: TEST_CLIENT_ID.into(),
scopes: vec!["openid".into(), "email".into()],
audience: None,
};
let discovery = discover(&oidc_cfg.issuer).await.unwrap();
let pkce = generate_pkce();
let state = "test-state".to_string();
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let redirect_uri = format!("http://127.0.0.1:{port}/callback");
let _auth_url = build_authorize_url(
&oidc_cfg,
None,
&discovery,
&redirect_uri,
&pkce,
&state,
TEST_NONCE,
);
// Simulate browser callback via race_callback_and_stdin
let Callback {
code,
state: received_state,
} = tokio::join!(race_callback_and_stdin(listener, false), async {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
reqwest::get(format!(
"http://127.0.0.1:{port}/callback?code=mock-auth-code&state={state}"
))
.await
.unwrap();
})
.0
.unwrap();
assert_eq!(code, "mock-auth-code");
assert_eq!(received_state, state);
let tokens = exchange_code(
&discovery.token_endpoint,
&code,
&redirect_uri,
&oidc_cfg.client_id,
&pkce.code_verifier,
)
.await
.unwrap();
assert_eq!(tokens.access_token, "mock-access-token");
let user_info = extract_user_info(
tokens.id_token.as_deref(),
&discovery,
&oidc_cfg.issuer,
&oidc_cfg.client_id,
TEST_NONCE,
None,
None,
None,
)
.await
.unwrap();
let auth = build_grok_auth(tokens, user_info, &oidc_cfg.issuer, &oidc_cfg.client_id);
let auth = auth_manager.update(auth).await.unwrap();
assert_eq!(auth.key, "mock-access-token");
assert_eq!(auth.refresh_token.as_deref(), Some("mock-refresh-token"));
assert_eq!(auth.user_id, "user-42");
assert_eq!(auth.email.as_deref(), Some("test@corp.com"));
assert!(auth.principal_type.is_none());
assert!(auth.principal_id.is_none());
assert!(auth.expires_at.is_some());
assert_eq!(auth.oidc_issuer.as_deref(), Some(issuer.as_str()));
let auth_json = std::fs::read_to_string(temp_dir.path().join("auth.json")).unwrap();
assert!(auth_json.contains("mock-access-token"));
assert!(auth_json.contains("user-42"));
idp_server.abort();
}
/// Parser matrix: full callback URL, bare code, error URL, empty.
/// Each case is one bug class:
/// - full URL: regression in URL extraction
/// - bare code: paste-friendly fallback
/// - error URL: surfaces IdP error to user
/// - empty: input validation
#[test]
fn parse_pasted_input_matrix() {
// (input, expected: Ok((code, state)) | Err substring)
let ok_cases: &[(&str, &str, &str)] = &[
(
"http://127.0.0.1:54321/callback?code=abc123&state=xyz789",
"abc123",
"xyz789",
),
("abc123def456", "abc123def456", ""),
];
for (input, code, state) in ok_cases {
let cb =
parse_pasted_input(input).unwrap_or_else(|e| panic!("parse {input:?} failed: {e}"));
assert_eq!(cb.code, *code, "code for {input:?}");
assert_eq!(cb.state, *state, "state for {input:?}");
}
let err_cases: &[(&str, &str)] = &[
(
"http://127.0.0.1:54321/callback?error=access_denied&error_description=User+denied",
"access_denied",
),
("", ""),
(" ", ""),
];
for (input, expected_substr) in err_cases {
let err = parse_pasted_input(input).unwrap_err();
if !expected_substr.is_empty() {
assert!(
err.to_string().contains(expected_substr),
"input {input:?} -> unexpected err: {err}"
);
}
}
}
}
crates/codegen/kigi-shell/src/auth/oidc/mod.rs
+14
View File
@@ -0,0 +1,14 @@
//! OIDC authentication: protocol, login, and refresh submodules.
mod login;
pub(crate) mod protocol;
pub(crate) mod refresh;
#[cfg(test)]
mod test_helpers;
pub use login::{run_login_flow, run_login_flow_with_config};
pub(crate) use protocol::{
enforce_login_principal, is_configured, login_principal_policy, peek_access_token_principal,
peek_access_token_principal_id, with_alpha_test_key,
};
pub(crate) use refresh::{OidcRefreshResult, oidc_token_exchange};
crates/codegen/kigi-shell/src/auth/oidc/protocol.rs
+1283
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/oidc/refresh.rs
+249
View File
@@ -0,0 +1,249 @@
//! Pure-data OIDC refresh. Talks to the IdP and returns
//! [`OidcRefreshResult`] without touching [`AuthManager`].
use super::super::GrokAuth;
use super::protocol::{OidcError, OidcUserInfo, build_grok_auth, discover, refresh_tokens};
use crate::auth::error::RefreshTokenFailedReason;
/// Outcome of a pure OIDC token refresh (no AuthManager mutations).
pub(crate) enum OidcRefreshResult {
/// Fresh token obtained. Caller must persist.
Success(Box<GrokAuth>),
/// Terminal error from the IdP, already classified into a reason.
TerminalError { reason: RefreshTokenFailedReason },
/// Non-terminal failure (discovery failed, network error, etc.)
Failed,
}
/// Classify an OAuth2 `error` code as a terminal refresh failure. `None` means
/// non-terminal (retryable). Single source of truth for which codes are fatal;
/// the retry gate (`protocol::is_transient_refresh_error`) defers to this too.
pub(super) fn classify_terminal(error_code: &str) -> Option<RefreshTokenFailedReason> {
match error_code {
"invalid_grant" => Some(RefreshTokenFailedReason::RefreshTokenRejected),
"invalid_client" => Some(RefreshTokenFailedReason::ClientRejected),
_ => None,
}
}
/// `oauth2-provider` refresh-token rotation-grace window (ms). Only a clock
/// divergence past this bound is flagged as a suspected suspend-straddle, since
/// a longer suspend can turn a lost refresh response into a revoked RT.
const ROTATION_GRACE_MS: u64 = 60_000;
/// Exchange a refresh_token for fresh tokens at the IdP. Pure data return, no
/// `AuthManager` mutations; the caller (`OidcRefresher`) routes the result
/// through `refresh_chain`.
pub(crate) async fn oidc_token_exchange(auth: &GrokAuth) -> OidcRefreshResult {
let has_rt = auth.refresh_token.is_some();
let has_issuer = auth.oidc_issuer.is_some();
let has_client_id = auth.oidc_client_id.is_some();
tracing::debug!(
has_rt,
has_issuer,
has_client_id,
"oidc try_refresh_pure enter"
);
if !has_rt || !has_issuer || !has_client_id {
kigi_log::unified_log::warn(
"oidc try_refresh skipped: missing fields",
None,
Some(serde_json::json!({
"has_refresh_token": has_rt,
"has_issuer": has_issuer,
"has_client_id": has_client_id,
"auth_mode": format!("{:?}", auth.auth_mode),
})),
);
}
let Some(refresh_tok) = auth.refresh_token.as_ref() else {
return OidcRefreshResult::Failed;
};
let Some(issuer) = auth.oidc_issuer.as_ref() else {
return OidcRefreshResult::Failed;
};
let Some(client_id) = auth.oidc_client_id.as_ref() else {
return OidcRefreshResult::Failed;
};
crate::unified_log::info(
"oidc try_refresh_pure enter",
None,
Some(serde_json::json!({ "issuer": issuer, "client_id": client_id })),
);
// Suspend probe: the monotonic clock pauses while the machine is asleep
// but the wall clock does not, so a large divergence around the IdP call
// means the process was suspended mid-refresh — the exact condition that
// can revoke the refresh token (response lost across sleep).
let started_mono = std::time::Instant::now();
let started_wall = chrono::Utc::now();
let timing = || {
let mono_ms = started_mono.elapsed().as_millis() as u64;
let wall_ms = (chrono::Utc::now() - started_wall)
.num_milliseconds()
.max(0) as u64;
let suspended_ms = wall_ms.saturating_sub(mono_ms);
(
mono_ms,
wall_ms,
suspended_ms,
suspended_ms > ROTATION_GRACE_MS,
)
};
let discovery = match discover(issuer).await {
Ok(d) => d,
Err(e) => {
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
crate::unified_log::error(
"oidc try_refresh_pure discovery failed",
None,
Some(serde_json::json!({
"error": format!("{e:#}"),
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
})),
);
if suspected_suspend {
emit_suspend_spanned("discovery_failed", suspended_ms);
}
return OidcRefreshResult::Failed;
}
};
let tokens = match refresh_tokens(
&discovery.token_endpoint,
refresh_tok,
client_id,
auth.principal_type.as_deref(),
auth.principal_id.as_deref(),
)
.await
{
Ok(t) => t,
Err(e) => {
if let Some(OidcError::TokenRefreshHttp { body, .. }) = e.downcast_ref::<OidcError>()
&& let Some(error_code) = serde_json::from_str::<serde_json::Value>(body)
.ok()
.and_then(|v| v.get("error")?.as_str().map(str::to_owned))
&& let Some(reason) = classify_terminal(&error_code)
{
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
let cred_age_secs = auth.mint_age_seconds();
crate::unified_log::error(
"oidc try_refresh_pure terminal error",
None,
Some(serde_json::json!({
"error_code": error_code,
"client_id": client_id,
"tried_rt_prefix": auth.refresh_token.as_deref().map(crate::auth::token_suffix),
"error_description": serde_json::from_str::<serde_json::Value>(body)
.ok()
.and_then(|v| v.get("error_description").cloned()),
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
"cred_age_secs": cred_age_secs,
})),
);
if suspected_suspend {
emit_suspend_spanned(&error_code, suspended_ms);
}
return OidcRefreshResult::TerminalError { reason };
}
let http_status = e.downcast_ref::<OidcError>().and_then(|oe| match oe {
OidcError::TokenRefreshHttp { status, .. } => Some(*status),
_ => None,
});
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
crate::unified_log::error(
"oidc try_refresh_pure token exchange failed",
None,
Some(serde_json::json!({
"error": e.to_string(),
"client_id": client_id,
"http_status": http_status,
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
})),
);
tracing::warn!(
error = %e,
http_status = ?http_status,
client_id = %client_id,
issuer = %issuer,
"OIDC: token refresh failed"
);
if suspected_suspend {
emit_suspend_spanned("transient_failed", suspended_ms);
}
return OidcRefreshResult::Failed;
}
};
// Reuse identity from original login; new id_token from refresh is intentionally skipped.
let user_info = OidcUserInfo {
user_id: auth.user_id.clone(),
email: auth.email.clone(),
first_name: auth.first_name.clone(),
last_name: auth.last_name.clone(),
profile_image_asset_id: auth.profile_image_asset_id.clone(),
principal_type: auth.principal_type.clone(),
principal_id: auth.principal_id.clone(),
team_id: auth.team_id.clone(),
team_name: auth.team_name.clone(),
team_role: auth.team_role.clone(),
organization_id: auth.organization_id.clone(),
organization_name: auth.organization_name.clone(),
organization_role: auth.organization_role.clone(),
user_blocked_reason: auth.user_blocked_reason.clone(),
team_blocked_reasons: auth.team_blocked_reasons.clone(),
coding_data_retention_opt_out: auth.coding_data_retention_opt_out,
};
let mut new_auth = build_grok_auth(tokens, user_info, issuer, client_id);
let idp_rotated = new_auth.refresh_token.is_some();
// Keep old refresh token if IdP didn't rotate it
if new_auth.refresh_token.is_none() {
new_auth.refresh_token = auth.refresh_token.clone();
}
tracing::debug!(
idp_rotated,
key_prefix = crate::auth::token_suffix(&new_auth.key),
"oidc try_refresh_pure token obtained"
);
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
crate::unified_log::info(
"oidc try_refresh_pure succeeded",
None,
Some(serde_json::json!({
"expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()),
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
})),
);
if suspected_suspend {
emit_suspend_spanned("ok", suspended_ms);
}
OidcRefreshResult::Success(Box::new(new_auth))
}
/// Alertable event: an OIDC refresh's network call spanned a suspend (wall
/// clock ran far ahead of the monotonic clock) — the precondition for a
/// lost-response refresh-token revocation.
fn emit_suspend_spanned(outcome: &str, suspended_ms: u64) {
crate::unified_log::warn(
"auth.refresh.suspend_spanned",
None,
Some(serde_json::json!({
"outcome": outcome,
"suspended_ms": suspended_ms,
})),
);
}
crates/codegen/kigi-shell/src/auth/oidc/test_helpers.rs
+130
View File
@@ -0,0 +1,130 @@
//! Shared test helpers for `oidc::protocol::tests` and `oidc::login::tests`.
//! Both test modules need a mock IdP server (`start_mock_idp`), JWT
//! signing primitives (`generate_test_rsa_key`, `mock_idp_token`), and
//! the same constants. Extracted here so neither test mod has to
//! re-implement them.
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use super::protocol::{Discovery, discover};
pub(super) const TEST_KID: &str = "test-kid";
pub(super) const TEST_NONCE: &str = "test-nonce-value";
pub(super) const TEST_CLIENT_ID: &str = "test-client-id";
pub(super) fn ensure_crypto_provider() {
let _ = rustls::crypto::ring::default_provider().install_default();
let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default();
}
pub(super) fn generate_test_rsa_key() -> (String, String, String) {
use rsa::pkcs8::EncodePrivateKey;
use rsa::traits::PublicKeyParts;
let private_key = rsa::RsaPrivateKey::new(&mut rsa::rand_core::OsRng, 2048).unwrap();
let pem = private_key
.to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)
.unwrap()
.to_string();
let jwk_n = URL_SAFE_NO_PAD.encode(private_key.n().to_bytes_be());
let jwk_e = URL_SAFE_NO_PAD.encode(private_key.e().to_bytes_be());
(pem, jwk_n, jwk_e)
}
pub(super) async fn mock_idp_token() -> (String, String, Discovery, tokio::task::JoinHandle<()>) {
let (issuer, handle) = start_mock_idp().await;
let discovery = discover(&issuer).await.unwrap();
let resp: serde_json::Value = crate::http::shared_client()
.post(&discovery.token_endpoint)
.form(&[("grant_type", "authorization_code")])
.send()
.await
.unwrap()
.json()
.await
.unwrap();
let id_token = resp["id_token"]
.as_str()
.expect("mock missing id_token")
.to_string();
(issuer, id_token, discovery, handle)
}
pub(super) async fn start_mock_idp() -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let issuer = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
let issuer_for_discovery = issuer.clone();
let (rsa_pem, jwk_n, jwk_e) = generate_test_rsa_key();
#[derive(serde::Serialize)]
struct Claims {
sub: &'static str,
email: &'static str,
iss: String,
aud: &'static str,
nonce: &'static str,
exp: usize,
}
let id_token = {
let mut hdr = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);
hdr.kid = Some(TEST_KID.to_owned());
jsonwebtoken::encode(
&hdr,
&Claims {
sub: "user-42",
email: "test@corp.com",
iss: issuer.clone(),
aud: TEST_CLIENT_ID,
nonce: TEST_NONCE,
exp: (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp() as usize,
},
&jsonwebtoken::EncodingKey::from_rsa_pem(rsa_pem.as_bytes()).unwrap(),
)
.unwrap()
};
let app = axum::Router::new()
.route(
"/.well-known/openid-configuration",
axum::routing::get(move || {
let iss = issuer_for_discovery.clone();
async move {
axum::Json(serde_json::json!({
"authorization_endpoint": format!("{iss}/authorize"),
"token_endpoint": format!("{iss}/token"),
"jwks_uri": format!("{iss}/jwks"),
"id_token_signing_alg_values_supported": ["RS256"],
}))
}
}),
)
.route(
"/jwks",
axum::routing::get(move || {
let n = jwk_n.clone();
let e = jwk_e.clone();
async move {
axum::Json(serde_json::json!({
"keys": [{
"kty": "RSA", "alg": "RS256", "kid": TEST_KID,
"n": n, "e": e,
}]
}))
}
}),
)
.route(
"/token",
axum::routing::post(move || {
let tok = id_token.clone();
async move {
axum::Json(serde_json::json!({
"access_token": "mock-access-token",
"refresh_token": "mock-refresh-token",
"id_token": tok,
"expires_in": 3600,
}))
}
}),
);
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(issuer, handle)
}
crates/codegen/kigi-shell/src/auth/recovery.rs
+933
View File
@@ -0,0 +1,933 @@
//! Unauthorized (401) recovery state machine.
//!
//! When the server rejects a token, `UnauthorizedRecovery` walks through
//! a sequence of recovery steps before giving up:
//!
//! 1. **ReloadFromDisk** — re-read `auth.json` under a file lock; if the
//! on-disk token differs from the rejected one, accept it (another
//! process may have refreshed).
//! 2. **RefreshFromAuthority** — run the appropriate refresh chain
//! (OIDC token refresh, external binary, etc.) based on `TokenType`,
//! unless the live token was minted moments ago (fresh-mint guard).
//! 3. **DevboxRecovery** — on devboxes, purge `auth.json` and mint fresh
//! OIDC credentials.
//! 4. **Done** — all recovery strategies exhausted.
use std::sync::Arc;
use crate::auth::error::{AuthError, RefreshTokenError, RefreshTokenFailedReason};
use crate::auth::manager::AuthManager;
use crate::auth::model::GrokAuth;
use crate::auth::token_type::TokenType;
/// Whether a terminal `AuthError` forces a manual re-login (`None` cases
/// self-heal or are transient). Lives here (not on `AuthError`) so the error
/// model stays free of recovery policy.
pub(crate) fn forces_manual_reauth(err: &AuthError) -> bool {
match err {
AuthError::Refresh(RefreshTokenError::Permanent(e)) => match e.reason {
RefreshTokenFailedReason::RefreshTokenRejected => true,
// Self-healing via the TTL, not a manual re-auth.
RefreshTokenFailedReason::ClientRejected | RefreshTokenFailedReason::Other => false,
},
AuthError::ServerRejectedNoRecovery
| AuthError::RecoveryExhausted
| AuthError::TokenExpiredNoRefresh
| AuthError::PinnedTeamMismatch { .. } => true,
// API-key lockouts are out of scope: an admin disabling API-key auth
// means rotate the key, not `/login`.
AuthError::ApiKeyAuthDisabled
| AuthError::Refresh(RefreshTokenError::Transient(_))
| AuthError::NotLoggedIn => false,
}
}
/// Whether the relay should stop reconnecting on this recovery error. Its own
/// predicate rather than reusing `forces_manual_reauth`: the relay must give up
/// on any terminal auth failure, including `ApiKeyAuthDisabled` (a kill-switched
/// API key), which deliberately doesn't force a manual re-login.
pub(crate) fn relay_should_cancel(err: &AuthError) -> bool {
forces_manual_reauth(err) || matches!(err, AuthError::ApiKeyAuthDisabled)
}
/// Fresh-mint guard window (±) for `ServerRejected` refreshes
/// ([`UnauthorizedRecovery::fresh_mint_guard`]). 120s outlasts in-flight
/// requests sent with a previous key plus validation lag (observed stale
/// 401s land ~20s after mint), while `current()`'s 300s early-invalidation
/// buffer keeps any guard-returned token wire-valid. A genuinely-dead fresh
/// token waits at most this long to re-mint; the symmetric bound caps that
/// delay when the clock stepped back.
const FRESH_MINT_GUARD_SECS: i64 = 120;
/// Which recovery step to attempt next.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RecoveryStep {
/// Re-read auth.json from disk (file-locked).
ReloadFromDisk,
/// Refresh via the authority (OIDC, external binary, etc.).
RefreshFromAuthority,
/// On devboxes: purge auth.json and mint fresh OIDC credentials.
DevboxRecovery,
/// All strategies exhausted.
Done,
}
/// State machine that walks through recovery strategies after a 401.
pub struct UnauthorizedRecovery {
auth_manager: Arc<AuthManager>,
/// The token that was rejected by the server.
rejected_token: String,
/// Current step in the recovery sequence.
step: RecoveryStep,
/// Error from `RefreshFromAuthority`, propagated as fallback when
/// devbox recovery doesn't apply.
authority_error: Option<AuthError>,
/// Whether the last authority failure was transient. Kept past the
/// `authority_error` handoff so exhaustion preserves the
/// transient/permanent axis (see the `Done` arm).
authority_was_transient: bool,
}
impl UnauthorizedRecovery {
/// `rejected` is the credential the server rejected: its key drives recovery.
pub(crate) fn new(auth_manager: Arc<AuthManager>, rejected: Option<GrokAuth>) -> Self {
let rejected_token = rejected.as_ref().map(|a| a.key.clone()).unwrap_or_default();
Self {
auth_manager,
rejected_token,
step: RecoveryStep::ReloadFromDisk,
authority_error: None,
authority_was_transient: false,
}
}
/// Attempt the next recovery step. Walks
/// `ReloadFromDisk -> RefreshFromAuthority -> DevboxRecovery -> Done`.
/// `token_type` span field is recorded lazily via
/// `Span::is_disabled()` to avoid the lock when tracing is off.
#[tracing::instrument(
skip(self),
fields(step = ?self.step, token_type = tracing::field::Empty),
)]
pub async fn next(&mut self) -> Result<GrokAuth, AuthError> {
let span = tracing::Span::current();
if !span.is_disabled() {
// Only acquire the inner-lock when tracing actually
// collects the span. `token_type()` -> `inner.read()` is
// ~free but it's still a lock, and recovery is on the
// 401-recovery path; making the cost zero when tracing is
// off matches the no-trace-no-cost contract.
span.record(
"token_type",
tracing::field::debug(self.auth_manager.token_type()),
);
}
self.resolve_next().await
}
/// Walk the recovery steps and apply the team-pin policy gate.
async fn resolve_next(&mut self) -> Result<GrokAuth, AuthError> {
// Team-pin gate: 401 recovery must not resurrect a wrong-team session
// (disk adoption / refresh / devbox mint) for the relay to reconnect
// with. Clear + reject on mismatch.
let auth = self.next_step_loop().await?;
if let Some(e) = self.auth_manager.cached_token_policy_error(&auth) {
self.auth_manager.reject_and_clear(&e);
return Err(e);
}
Ok(auth)
}
async fn next_step_loop(&mut self) -> Result<GrokAuth, AuthError> {
loop {
match self.step {
RecoveryStep::ReloadFromDisk => {
self.step = RecoveryStep::RefreshFromAuthority;
if let Some(auth) = self.try_reload_from_disk().await {
return Ok(auth);
}
}
RecoveryStep::RefreshFromAuthority => {
self.step = RecoveryStep::DevboxRecovery;
match self.try_refresh_from_authority().await {
Ok(auth) => return Ok(auth),
Err(e) => {
self.authority_was_transient =
matches!(e, AuthError::Refresh(RefreshTokenError::Transient(_)));
self.authority_error = Some(e);
}
}
}
RecoveryStep::DevboxRecovery => {
self.step = RecoveryStep::Done;
// preferred_method=api_key forbids automatic OIDC mint.
if !self.auth_manager.grok_com_config().blocks_automatic_oidc()
&& self.auth_manager.is_devbox_environment()
&& let Ok(auth) = self.auth_manager.try_devbox_recovery().await
{
return Ok(auth);
}
return Err(self
.authority_error
.take()
.unwrap_or(AuthError::RecoveryExhausted));
}
RecoveryStep::Done => {
// Exhaustion after a *transient* authority failure stays
// transient: `RecoveryExhausted` here would count a network
// blip as a forced re-login and cancel the relay instead of
// letting it reconnect.
return Err(if self.authority_was_transient {
AuthError::transient("recovery exhausted after transient refresh failure")
} else {
AuthError::RecoveryExhausted
});
}
}
}
}
/// Re-read `auth.json` from disk. Accept the token only if it differs
/// from the one that was rejected.
async fn try_reload_from_disk(&self) -> Option<GrokAuth> {
let _lock = self
.auth_manager
.try_lock_auth_file_async(crate::auth::manager::AUTH_LOCK_TIMEOUT)
.await;
if _lock.is_none() {
tracing::warn!("auth recovery: proceeding without file lock");
}
let Some(disk_auth) = self.auth_manager.read_disk_auth() else {
// Every ReloadFromDisk outcome must log (adopted / expired /
// same-as-rejected / no entry): a silent arm hides which path
// a recovery loop is taking. Debug level — the disk-state
// *transition* is logged once by `read_disk_auth` itself.
kigi_log::unified_log::debug("auth recovery: no disk entry", None, None);
return None;
};
if crate::auth::is_expired(&disk_auth) {
tracing::debug!("auth recovery: disk token is expired, skipping");
kigi_log::unified_log::debug(
"auth recovery: disk token expired",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&disk_auth.key),
"expires_at": disk_auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
return None;
}
if self.is_different_token(&disk_auth) {
tracing::info!("auth recovery: disk has a different token, accepting");
kigi_log::unified_log::info(
"auth recovery: adopted disk token",
None,
Some(serde_json::json!({
"adopted_key_prefix": crate::auth::token_suffix(&disk_auth.key),
"expires_at": disk_auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
self.auth_manager.hot_swap(disk_auth.clone());
Some(disk_auth)
} else {
tracing::debug!("auth recovery: disk token is same as rejected, skipping");
kigi_log::unified_log::debug("auth recovery: disk token same as rejected", None, None);
None
}
}
/// Return the live token instead of refreshing when its mint age is
/// within ±[`FRESH_MINT_GUARD_SECS`]; anything outside (including a
/// clock that stepped far back) falls through to a normal refresh.
///
/// A 401 moments after a successful mint is a stale rejection (sent with
/// the previous key and mis-attributed — see `is_stale_snapshot`) or
/// validation lag on the new key — re-minting fixes neither, and a crash
/// between the IdP grant and persisting the response orphans the
/// replacement RT (forced re-login). Consumers retry with the returned
/// token; a genuinely-bad one refreshes once the window passes. Lives
/// here, not in `refresh_chain`, so paywall claims re-mints that call
/// `refresh_chain(ServerRejected)` directly are unaffected.
fn fresh_mint_guard(&self) -> Option<GrokAuth> {
let auth = self.auth_manager.current()?;
let mint_age_seconds = auth.mint_age_seconds();
if !(-FRESH_MINT_GUARD_SECS..FRESH_MINT_GUARD_SECS).contains(&mint_age_seconds) {
return None;
}
tracing::info!(
mint_age_seconds,
"auth recovery: current token freshly minted, skipping refresh"
);
kigi_log::unified_log::info(
"auth recovery: fresh mint, refresh skipped",
None,
Some(serde_json::json!({
"key_prefix": crate::auth::token_suffix(&auth.key),
"mint_age_seconds": mint_age_seconds,
"guard_seconds": FRESH_MINT_GUARD_SECS,
"expires_at": auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
Some(auth)
}
/// Dispatch to the correct refresh chain based on the current `TokenType`.
///
/// Per-variant outcome:
///
/// - **OidcSession / ExternalBinary**: full refresh chain via the
/// authority, unless the live token is inside the fresh-mint guard
/// window ([`Self::fresh_mint_guard`]).
/// - **LegacySession / ApiKey**: no refresh authority for these
/// types. We've already tried `ReloadFromDisk` (the previous
/// recovery step), so the server's 401 stands. Surface
/// [`AuthError::ServerRejectedNoRecovery`] -- *not*
/// `TokenExpiredNoRefresh`, because the trigger here is the
/// server rejecting the token (it may not have aged past any
/// local TTL; ApiKey in particular has no expiry). Consumers
/// reading the variant can distinguish "ran past local TTL" from
/// "server actively rejected".
/// - **None**: no credentials at all.
async fn try_refresh_from_authority(&self) -> Result<GrokAuth, AuthError> {
let tt = self.auth_manager.token_type();
match tt {
TokenType::OidcSession | TokenType::ExternalBinary => {
if let Some(auth) = self.fresh_mint_guard() {
return Ok(auth);
}
let result = self
.auth_manager
.refresh_chain(tt, crate::auth::manager::RefreshReason::ServerRejected)
.await;
match &result {
Ok(auth) => {
kigi_log::unified_log::info(
"auth recovery: refreshed from authority",
None,
Some(serde_json::json!({
"token_type": format!("{tt:?}"),
"new_key_prefix": crate::auth::token_suffix(&auth.key),
"expires_at": auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
}
Err(e) => {
kigi_log::unified_log::warn(
"auth recovery: refresh from authority failed",
None,
Some(serde_json::json!({
"token_type": format!("{tt:?}"),
"error": format!("{e}"),
})),
);
}
}
result
}
TokenType::LegacySession | TokenType::ApiKey => {
kigi_log::unified_log::warn(
"auth recovery: no refresh authority for token type",
None,
Some(serde_json::json!({ "token_type": format!("{tt:?}") })),
);
Err(AuthError::ServerRejectedNoRecovery)
}
TokenType::None => Err(AuthError::NotLoggedIn),
}
}
/// Check if a candidate token is different from the rejected one.
fn is_different_token(&self, candidate: &GrokAuth) -> bool {
candidate.key != self.rejected_token
}
}
#[cfg(test)]
mod tests {
//! State-machine matrix tests for `UnauthorizedRecovery`.
//!
//! Coverage targets:
//! - All 5 `TokenType` variants x dispatch in `try_refresh_from_authority`.
//! - `try_reload_from_disk`: same/different/no token on disk.
//! - `next()` exhaustion (Done -> RecoveryExhausted).
//! - Fresh-mint guard: ±window bounds, ExternalBinary, verdict grace,
//! policy-hidden fall-through (fail closed).
//!
//! These tests use the same in-process `AuthManager` that production
//! does and inject a counting refresher so we can observe whether the
//! authority was consulted.
use super::*;
use crate::auth::config::GrokComConfig;
use crate::auth::error::{RefreshTokenError, RefreshTokenFailedReason};
use crate::auth::model::{AuthMode, GrokAuth};
use crate::auth::refresh::{RefreshOutcome, TokenRefresher};
use crate::auth::storage::{read_auth_json, write_auth_json};
use chrono::{Duration, Utc};
use std::sync::atomic::{AtomicU32, Ordering};
/// The rejected wire bearer these tests seed into the manager.
fn rejected_cred() -> Option<GrokAuth> {
Some(GrokAuth {
key: "rejected-tok".into(),
..GrokAuth::test_default()
})
}
/// Refresher fake: returns Success with a fresh token on every call.
struct OkRefresher {
calls: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl TokenRefresher for OkRefresher {
async fn refresh(&self, _reason: crate::auth::manager::RefreshReason) -> RefreshOutcome {
self.calls.fetch_add(1, Ordering::SeqCst);
RefreshOutcome::Success(Box::new(GrokAuth {
key: "fresh-from-authority".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
}))
}
}
/// Refresher fake: returns PermanentFailure (invalid_grant).
struct FailRefresher {
calls: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl TokenRefresher for FailRefresher {
async fn refresh(&self, _reason: crate::auth::manager::RefreshReason) -> RefreshOutcome {
self.calls.fetch_add(1, Ordering::SeqCst);
RefreshOutcome::permanent(RefreshTokenFailedReason::RefreshTokenRejected, None)
}
}
fn mgr() -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let m = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
(dir, m)
}
fn seed(mgr: &AuthManager, mode: AuthMode, refresh_token: Option<&str>) {
let auth = GrokAuth {
key: "rejected-tok".into(),
auth_mode: mode,
refresh_token: refresh_token.map(str::to_string),
// Past expiry so `current()` returns None and the refresh
// chain actually has to do work.
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(auth);
}
// -- TokenType dispatch matrix ----------------------------------------
#[tokio::test]
async fn dispatch_oidc_session_uses_refresh_chain() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
// ReloadFromDisk fails (no disk auth), then RefreshFromAuthority succeeds.
let auth = rec.next().await.expect("recovery should succeed");
assert_eq!(auth.key, "fresh-from-authority");
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn dispatch_external_binary_uses_refresh_chain() {
let (_d, m) = mgr();
seed(&m, AuthMode::External, None);
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let auth = rec.next().await.expect("external-binary recovery succeeds");
assert_eq!(auth.key, "fresh-from-authority");
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
// -- Fresh-mint guard --------------------------------------------------
/// Seed a *valid* (unexpired) in-memory token whose `create_time` lies
/// `mint_age` in the past (negative = clock stepped back since mint).
fn seed_valid(mgr: &AuthManager, mode: AuthMode, mint_age: Duration) {
mgr.hot_swap(GrokAuth {
key: "rejected-tok".into(),
auth_mode: mode,
refresh_token: Some("rt".into()),
create_time: Utc::now() - mint_age,
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
});
}
/// Run one recovery against a counting refresher; return the outcome and
/// how many times the authority was consulted.
async fn recover_with_ok_refresher(m: &Arc<AuthManager>) -> (Result<GrokAuth, AuthError>, u32) {
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let result = rec.next().await;
(result, calls.load(Ordering::SeqCst))
}
#[tokio::test]
async fn fresh_mint_guard_skips_idp_for_freshly_minted_token() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
"rejected-tok"
);
assert_eq!(calls, 0, "a 10s-old token must not be re-minted");
}
#[tokio::test]
async fn fresh_mint_guard_applies_to_external_binary_tokens() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::External, Duration::seconds(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
"rejected-tok"
);
assert_eq!(calls, 0);
}
#[tokio::test]
async fn fresh_mint_guard_treats_small_negative_age_as_fresh() {
// Clock stepped back slightly since mint (NTP nudge).
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(-60));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
"rejected-tok"
);
assert_eq!(calls, 0);
}
#[tokio::test]
async fn fresh_mint_guard_refreshes_when_clock_stepped_far_back() {
// A large backwards clock step must not wedge recovery for the whole
// step: outside the ±window the guard stands down.
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::hours(-1));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("recovery should succeed").key,
"fresh-from-authority"
);
assert_eq!(calls, 1, "far-negative mint age must reach the IdP");
}
#[tokio::test]
async fn fresh_mint_guard_lets_old_token_refresh() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::minutes(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("recovery should succeed").key,
"fresh-from-authority"
);
assert_eq!(
calls, 1,
"outside the guard window ServerRejected must reach the IdP"
);
}
#[tokio::test]
async fn fresh_mint_guard_wins_over_cached_permanent_failure() {
// A fresh *valid* token is served even when a permanent-failure
// verdict is cached for it — mirrors `auth()`'s wire-valid grace arm;
// the verdict re-applies once the guard window passes.
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(10));
m.record_permanent_failure(
"rejected-tok".into(),
RefreshTokenFailedReason::RefreshTokenRejected.into(),
);
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result
.expect("guard precedes the verdict short-circuit")
.key,
"rejected-tok"
);
assert_eq!(calls, 0);
}
#[tokio::test]
async fn fresh_mint_guard_never_returns_policy_hidden_token() {
// Wrong-team fresh token: `current()` hides it (vet_cached), so the
// guard must fall through to a normal refresh — fail closed.
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig {
force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single(
"team-good".into(),
)),
..GrokComConfig::default()
};
let m = Arc::new(AuthManager::new(dir.path(), cfg));
m.hot_swap(GrokAuth {
key: team_jwt("team-wrong"),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
create_time: Utc::now(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
});
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let result = rec.next().await;
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"hidden token must not satisfy the guard"
);
if let Ok(auth) = result {
assert_ne!(
auth.key,
team_jwt("team-wrong"),
"wrong-team token must never be returned"
);
}
}
#[tokio::test]
async fn dispatch_legacy_session_returns_server_rejected_no_recovery() {
let (_d, m) = mgr();
// WebLogin (no refresh_token) -> LegacySession.
seed(&m, AuthMode::WebLogin, None);
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(
matches!(err, AuthError::ServerRejectedNoRecovery),
"LegacySession recovery should surface ServerRejectedNoRecovery, got {err:?}",
);
}
#[tokio::test]
async fn dispatch_oidc_without_refresh_token_returns_server_rejected_no_recovery() {
// Oidc without refresh_token classifies as LegacySession.
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, None);
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(matches!(err, AuthError::ServerRejectedNoRecovery));
}
#[tokio::test]
async fn dispatch_api_key_returns_server_rejected_no_recovery() {
let (_d, m) = mgr();
seed(&m, AuthMode::ApiKey, None);
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(
matches!(err, AuthError::ServerRejectedNoRecovery),
"ApiKey recovery should surface ServerRejectedNoRecovery (not \
TokenExpiredNoRefresh), got {err:?}",
);
}
#[tokio::test]
async fn dispatch_none_returns_not_logged_in() {
let (_d, m) = mgr();
// No seed — inner stays None → TokenType::None.
// Single next() falls through ReloadFromDisk → RefreshFromAuthority.
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(
matches!(err, AuthError::NotLoggedIn),
"None token type should surface NotLoggedIn, got {err:?}",
);
}
// -- ReloadFromDisk matrix --------------------------------------------
#[tokio::test]
async fn reload_from_disk_picks_up_different_token() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
// Sibling process wrote a different valid token to disk.
let scope = m.grok_com_config().auth_scope();
let fresh = GrokAuth {
key: "fresh-from-disk".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, fresh);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let mut rec = m.unauthorized_recovery(rejected_cred());
let auth = rec
.next()
.await
.expect("recovery should pick up the disk token");
assert_eq!(auth.key, "fresh-from-disk");
}
#[tokio::test]
async fn reload_from_disk_skips_same_token_then_proceeds_to_authority() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
// Disk has the SAME token that was rejected -- skip, fall through.
let scope = m.grok_com_config().auth_scope();
let same = GrokAuth {
key: "rejected-tok".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, same);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let auth = rec.next().await.expect("authority refresh succeeds");
assert_eq!(auth.key, "fresh-from-authority");
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"fall-through to authority must invoke the refresher exactly once",
);
}
// -- Done state -------------------------------------------------------
/// With no stored authority error (the first `next()` succeeded), driving
/// past `Done` surfaces `RecoveryExhausted`. The transient-failure case is
/// pinned by `exhaustion_after_transient_failure_stays_transient`.
#[tokio::test]
async fn next_after_done_returns_recovery_exhausted() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
m.set_refresher(Arc::new(OkRefresher {
calls: Arc::new(AtomicU32::new(0)),
}));
// Pin non-devbox so DevboxRecovery can't adopt the seeded token (CI runs
// in K8s pods where is_devbox_environment() is true).
m.set_devbox_env_for_test(false);
let mut rec = m.unauthorized_recovery(rejected_cred());
let _ = rec.next().await.unwrap();
let err = loop {
if let Err(e) = rec.next().await {
break e;
}
};
assert!(
matches!(err, AuthError::RecoveryExhausted),
"Done state must surface RecoveryExhausted, got {err:?}",
);
}
/// Exhaustion after a *transient* authority failure preserves the
/// transient axis: surfacing `RecoveryExhausted` would count a network
/// blip as a forced re-login (`manual_auth`) and make the relay cancel
/// instead of reconnect.
#[tokio::test]
async fn exhaustion_after_transient_failure_stays_transient() {
/// Refresher fake: transient failure on every call.
struct TransientFailRefresher;
#[async_trait::async_trait]
impl TokenRefresher for TransientFailRefresher {
async fn refresh(
&self,
_reason: crate::auth::manager::RefreshReason,
) -> RefreshOutcome {
RefreshOutcome::transient("network blip")
}
}
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
m.set_refresher(Arc::new(TransientFailRefresher));
m.set_devbox_env_for_test(false);
let mut rec = m.unauthorized_recovery(rejected_cred());
// First next(): the authority's transient error propagates as-is.
let first = rec.next().await.unwrap_err();
assert!(
matches!(first, AuthError::Refresh(RefreshTokenError::Transient(_))),
"authority transient must propagate, got {first:?}",
);
// Driving past exhaustion must stay transient too.
let err = loop {
if let Err(e) = rec.next().await {
break e;
}
};
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))),
"exhaustion after a transient failure must stay transient, got {err:?}",
);
assert!(
!forces_manual_reauth(&err),
"a transient exhaustion must not force a manual re-login",
);
assert!(
!relay_should_cancel(&err),
"the relay must reconnect (not cancel) on a transient exhaustion",
);
}
// -- Permanent failure short-circuit (cross-check) ------------
#[tokio::test]
async fn refresh_authority_short_circuits_on_cached_permanent_failure() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
// Pre-record a permanent failure scoped to the seeded credential.
m.record_permanent_failure(
"rejected-tok".into(),
RefreshTokenFailedReason::RefreshTokenRejected.into(),
);
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(FailRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(matches!(
err,
AuthError::Refresh(RefreshTokenError::Permanent(_))
));
assert_eq!(
calls.load(Ordering::SeqCst),
0,
"refresher must not be invoked when permanent_failure is cached",
);
}
// -- ReloadFromDisk rejects expired disk tokens -------------------------
/// Regression: disk holds a different but expired token. Recovery
/// must skip it and fall through to RefreshFromAuthority, not
/// return it for the caller to send on the wire (instant 401).
#[tokio::test]
async fn reload_from_disk_rejects_expired_different_token() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
let scope = m.grok_com_config().auth_scope();
let expired_different = GrokAuth {
key: "different-but-expired".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, expired_different);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let auth = rec.next().await.expect("should fall through to authority");
assert_eq!(
auth.key, "fresh-from-authority",
"must skip the expired disk token and use the refresher"
);
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
// -- force_login_team_uuid pin enforced on the 401-recovery path -------
fn ensure_crypto_provider() {
let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default();
}
fn team_jwt(principal_id: &str) -> String {
ensure_crypto_provider();
jsonwebtoken::encode(
&jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256),
&serde_json::json!({
"sub": "user-1",
"principal_type": "Team",
"principal_id": principal_id,
"exp": 9999999999u64,
}),
&jsonwebtoken::EncodingKey::from_secret(b"test-secret"),
)
.unwrap()
}
/// A sibling writes a wrong-team token to disk; 401 recovery (relay path)
/// must reject + clear it at `next()`, not hand it back as a bearer.
#[tokio::test]
async fn recovery_rejects_wrong_team_adopted_disk_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig {
force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single(
"team-good".into(),
)),
..GrokComConfig::default()
};
let scope = cfg.auth_scope();
let m = Arc::new(AuthManager::new(dir.path(), cfg));
// In-memory: the rejected (expired) session that triggered recovery.
seed(&m, AuthMode::Oidc, Some("rt"));
// Disk: a different, non-expired, *wrong-team* token a sibling wrote.
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(
scope,
GrokAuth {
key: team_jwt("team-wrong"),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-sibling".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
},
);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(
matches!(err, AuthError::PinnedTeamMismatch { .. }),
"recovery must reject a wrong-team disk token, got {err:?}"
);
}
}
crates/codegen/kigi-shell/src/auth/refresh/auth_backend_contract_tests.rs
+351
View File
@@ -0,0 +1,351 @@
//! End-to-end auth-backend contract tests: a mock IdP whose `/token` response
//! is forced per case, asserting the refresh outcome, the storm cap, and the
//! terminal-error classification on the live recovery path.
use super::*;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::{GrokAuth, GrokComConfig};
use chrono::{Duration, Utc};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
/// Mock IdP: OIDC discovery + a `/token` endpoint returning a fixed
/// `(status, body)` and counting every hit, plus the `/user` endpoint
/// `AuthManager::update` calls after a successful refresh. `delay_ms` widens
/// the in-lock window so concurrent callers queue on `refresh_lock`.
async fn start_idp(
token_status: u16,
token_body: String,
hits: Arc<AtomicU32>,
delay_ms: u64,
) -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
let disco = base.clone();
let app = axum::Router::new()
.route(
"/.well-known/openid-configuration",
axum::routing::get(move || {
let b = disco.clone();
async move {
axum::Json(serde_json::json!({
"authorization_endpoint": format!("{b}/authorize"),
"token_endpoint": format!("{b}/token"),
}))
}
}),
)
.route(
"/token",
axum::routing::post(move || {
let hits = hits.clone();
let body = token_body.clone();
async move {
hits.fetch_add(1, Ordering::SeqCst);
if delay_ms > 0 {
tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
}
(
axum::http::StatusCode::from_u16(token_status).unwrap(),
body,
)
}
}),
)
.route(
"/user",
axum::routing::get(|| async {
axum::Json(serde_json::json!({ "userId": "user-42", "email": "u@corp.com" }))
}),
);
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(base, handle)
}
fn expired_oidc(base_url: &str) -> GrokAuth {
GrokAuth {
key: "expired-at".into(),
create_time: Utc::now() - Duration::hours(2),
user_id: "user-42".into(),
auth_mode: crate::auth::model::AuthMode::Oidc,
refresh_token: Some("rt-under-test".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
oidc_issuer: Some(base_url.to_owned()),
oidc_client_id: Some("client-under-test".into()),
..GrokAuth::test_default()
}
}
#[derive(Debug)]
enum Expect {
Success,
Permanent(RefreshTokenFailedReason),
Transient,
}
/// The IdP token-endpoint contract: each response shape maps to one outcome.
/// `invalid_grant`/`invalid_client` are the only permanent verdicts; status
/// blips and unrecognized codes stay transient (never permanent-lock).
#[tokio::test]
async fn auth_backend_contract_token_responses_map_to_outcomes() {
use RefreshTokenFailedReason::{ClientRejected, RefreshTokenRejected};
let cases: &[(&str, u16, &str, Expect)] = &[
(
"success",
200,
r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#,
Expect::Success,
),
(
"invalid_grant",
400,
r#"{"error":"invalid_grant"}"#,
Expect::Permanent(RefreshTokenRejected),
),
(
"invalid_client",
401,
r#"{"error":"invalid_client"}"#,
Expect::Permanent(ClientRejected),
),
("server_error_5xx", 503, "{}", Expect::Transient),
("rate_limited_429", 429, "{}", Expect::Transient),
(
"temporarily_unavailable",
400,
r#"{"error":"temporarily_unavailable"}"#,
Expect::Transient,
),
("bare_4xx_no_body", 400, "", Expect::Transient),
("malformed_body", 400, "not json", Expect::Transient),
// Proxy/WAF-mangled bodies must degrade to retry, never a false permanent
// lock: a nested error object or a non-string `error` is not a recognized
// top-level code, so it stays transient.
(
"nested_error_object",
400,
r#"{"error":{"code":"invalid_grant"}}"#,
Expect::Transient,
),
(
"non_string_error",
400,
r#"{"error":123}"#,
Expect::Transient,
),
];
for (name, status, body, expect) in cases {
let hits = Arc::new(AtomicU32::new(0));
let (base_url, server) = start_idp(*status, body.to_string(), hits.clone(), 0).await;
let dir = tempfile::tempdir().unwrap();
let auth_manager = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url),
);
auth_manager.hot_swap(expired_oidc(&base_url));
let refresher = OidcRefresher::new(auth_manager.clone());
let result = refresher.refresh(RefreshReason::ServerRejected).await;
match (expect, &result) {
(Expect::Success, RefreshOutcome::Success(_)) => {}
(Expect::Permanent(want), RefreshOutcome::PermanentFailure { error, .. }) => {
assert_eq!(error.reason, *want, "{name}: wrong permanent reason");
}
(Expect::Transient, RefreshOutcome::TransientFailure { .. }) => {}
(exp, got) => panic!("{name}: expected {exp:?}, got {got:?}"),
}
server.abort();
}
}
/// A burst of concurrent 401s on the same revoked refresh token must hit the
/// IdP exactly once. The callers serialize on `refresh_lock`; the leader records
/// the verdict before releasing, so the in-lock re-check (`refresh_chain` step
/// 1b) short-circuits every follower. Delete step 1b and the count climbs to N.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auth_backend_contract_concurrent_401s_hit_idp_once() {
let hits = Arc::new(AtomicU32::new(0));
// 100ms /token delay so every caller passes the pre-lock check and queues
// on refresh_lock before the leader records the verdict, exercising step 1b.
let (base_url, server) = start_idp(
400,
r#"{"error":"invalid_grant"}"#.to_string(),
hits.clone(),
100,
)
.await;
let dir = tempfile::tempdir().unwrap();
let auth_manager = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url),
);
auth_manager.hot_swap(expired_oidc(&base_url));
auth_manager.set_refresher(Arc::new(OidcRefresher::new(auth_manager.clone())));
let mut tasks = Vec::new();
for _ in 0..6 {
let auth_manager = auth_manager.clone();
tasks.push(tokio::spawn(async move { auth_manager.auth().await }));
}
for t in tasks {
let outcome = t.await.unwrap();
assert!(
matches!(
outcome,
Err(crate::auth::AuthError::Refresh(
crate::auth::RefreshTokenError::Permanent(_)
))
),
"every concurrent caller must fail permanently on a revoked refresh token, got {outcome:?}",
);
}
assert_eq!(
hits.load(Ordering::SeqCst),
1,
"concurrent 401s on one dead credential must hit the IdP exactly once",
);
server.abort();
}
/// The classification loop through the live recovery state machine: a dead
/// refresh token terminates recovery with an error that forces a manual
/// re-login; a refreshable token auto-refreshes.
#[tokio::test]
async fn auth_backend_contract_dead_token_forces_manual_reauth() {
// A dead refresh token terminates recovery with a forced-relogin error.
let hits = Arc::new(AtomicU32::new(0));
let (url, server) = start_idp(400, r#"{"error":"invalid_grant"}"#.to_string(), hits, 0).await;
let dir = tempfile::tempdir().unwrap();
let auth_manager =
Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&url));
auth_manager.hot_swap(expired_oidc(&url));
auth_manager.set_refresher(Arc::new(OidcRefresher::new(auth_manager.clone())));
let err = auth_manager
.unauthorized_recovery(auth_manager.current_or_expired())
.next()
.await
.expect_err("a dead refresh token must fail recovery");
assert!(
crate::auth::recovery::forces_manual_reauth(&err),
"a dead refresh token must be a forced-relogin error, got {err:?}",
);
server.abort();
// Refreshable token: recovery auto-refreshes.
let ok_hits = Arc::new(AtomicU32::new(0));
let (ok_url, ok_server) = start_idp(
200,
r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#.to_string(),
ok_hits,
0,
)
.await;
let ok_dir = tempfile::tempdir().unwrap();
let ok_manager = Arc::new(
AuthManager::new(ok_dir.path(), GrokComConfig::default()).with_proxy_base_url(&ok_url),
);
ok_manager.hot_swap(expired_oidc(&ok_url));
ok_manager.set_refresher(Arc::new(OidcRefresher::new(ok_manager.clone())));
let refreshed = ok_manager
.unauthorized_recovery(ok_manager.current_or_expired())
.next()
.await
.expect("a refreshable token must auto-refresh");
assert_eq!(
refreshed.key, "fresh",
"recovery must return the fresh token"
);
ok_server.abort();
}
/// Consecutive transient failures self-heal up to a bound, then escalate to a
/// non-sticky `Other` permanent failure (which ages out via the TTL). A
/// regression here would turn recoverable blips into a permanent `/login`.
#[tokio::test]
async fn auth_backend_contract_transient_failures_escalate_to_non_sticky_permanent() {
let hits = Arc::new(AtomicU32::new(0));
// Persistent 503: every refresh attempt is transient.
let (base_url, server) = start_idp(503, "{}".to_string(), hits, 0).await;
let dir = tempfile::tempdir().unwrap();
let auth_manager = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url),
);
auth_manager.hot_swap(expired_oidc(&base_url));
// One refresher instance: it owns the consecutive-failure counter.
let refresher = OidcRefresher::new(auth_manager.clone());
let mut outcomes = Vec::new();
for _ in 0..3 {
outcomes.push(refresher.refresh(RefreshReason::ServerRejected).await);
}
assert!(
matches!(outcomes[0], RefreshOutcome::TransientFailure { .. }),
"first blip is transient, not a lockout: {:?}",
outcomes[0],
);
match &outcomes[2] {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(
error.reason,
RefreshTokenFailedReason::Other,
"escalation must use the generic Other reason",
);
assert!(
!error.reason.is_sticky(),
"an escalated transient must age out, not strand the user forever",
);
}
other => panic!("repeated transients must escalate to a permanent Other, got {other:?}"),
}
server.abort();
}
/// Two `AuthManager`s sharing one auth.json stand in for two CLI processes: the
/// auth.json flock must serialize their refreshes so the shared refresh token is
/// spent at the IdP exactly once. The loser adopts the rotated token from disk
/// instead of racing a second exchange (which the IdP could revoke as reuse).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auth_backend_contract_two_instances_share_one_idp_call() {
let hits = Arc::new(AtomicU32::new(0));
let (url, server) = start_idp(
200,
r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#.to_string(),
hits.clone(),
100,
)
.await;
let dir = tempfile::tempdir().unwrap();
// Distinct managers, same on-disk auth.json (separate flock OFDs => they
// genuinely contend, like two processes).
let new_instance = || {
let m = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&url),
);
m.hot_swap(expired_oidc(&url));
m.set_refresher(Arc::new(OidcRefresher::new(m.clone())));
m
};
let a = new_instance();
let b = new_instance();
let (ra, rb) = tokio::join!(a.auth(), b.auth());
assert_eq!(ra.expect("instance A must obtain a token").key, "fresh");
assert_eq!(rb.expect("instance B must obtain a token").key, "fresh");
assert_eq!(
hits.load(Ordering::SeqCst),
1,
"two instances sharing auth.json must spend the refresh token at the IdP only once",
);
server.abort();
}
crates/codegen/kigi-shell/src/auth/refresh/external_refresher.rs
+176
View File
@@ -0,0 +1,176 @@
use std::sync::Arc;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::manager::RefreshReason;
use super::{ExternalCommandRunner, RefreshOutcome, TokenRefresher};
/// Refreshes by re-running the operator's external auth binary via
/// `spawn_blocking`. Pure data return -- mutation lives in
/// `refresh_chain` (honors the [`TokenRefresher`] no-mutation contract).
pub(crate) struct ExternalBinaryRefresher {
runner: Arc<dyn ExternalCommandRunner>,
command: String,
timeout: std::time::Duration,
}
impl ExternalBinaryRefresher {
pub(crate) fn new(runner: Arc<dyn ExternalCommandRunner>, command: String) -> Self {
Self {
runner,
command,
timeout: EXTERNAL_REFRESH_TIMEOUT,
}
}
/// Override the binary timeout (tests use a short one to exercise the
/// timeout arm without a real 30s wait).
#[cfg(test)]
pub(crate) fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
self.timeout = timeout;
self
}
/// A failed binary run is a single-strike `Other` permanent failure; the
/// `PERMANENT_FAILURE_TTL` lets a flaky binary self-heal without `/login`.
/// No consecutive-blip tolerance like OIDC: a local binary failure is a
/// stronger signal than a network refresh blip.
fn record_failure(&self, message: String) -> RefreshOutcome {
tracing::warn!(%message, "auth: external binary refresh failed -> permanent");
// No token key in the binary flow; the caller scopes the verdict.
RefreshOutcome::permanent(RefreshTokenFailedReason::Other, None)
}
}
/// Timeout for the external auth binary. If the binary hangs, the
/// `spawn_blocking` thread is leaked (it cannot be interrupted), but this is
/// acceptable: the thread holds no locks and mutates no shared state.
const EXTERNAL_REFRESH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
#[async_trait::async_trait]
impl TokenRefresher for ExternalBinaryRefresher {
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
tracing::debug!(?reason, "auth: external binary refresh starting");
let runner = self.runner.clone();
let cmd = self.command.clone();
let timeout_ms = self.timeout.as_millis() as u64;
match tokio::time::timeout(
self.timeout,
tokio::task::spawn_blocking(move || runner.run_external_command(&cmd)),
)
.await
{
Err(_elapsed) => {
tracing::warn!(
timeout_ms,
"auth: external binary refresh timed out (thread leaked)"
);
crate::unified_log::warn(
"auth.refresh.external_timeout",
None,
Some(serde_json::json!({ "timeout_ms": timeout_ms })),
);
self.record_failure(format!("external binary timed out after {timeout_ms}ms"))
}
Ok(Ok(Some(auth))) => {
crate::unified_log::info("auth: external binary refresh succeeded", None, None);
RefreshOutcome::success(auth)
}
Ok(Ok(None)) => {
crate::unified_log::warn(
"auth: external binary refresh returned no token",
None,
None,
);
self.record_failure("external binary returned no token".into())
}
Ok(Err(e)) => {
tracing::warn!(error = %e, "auth: external binary refresh task failed");
self.record_failure(format!("external binary task failed: {e}"))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::GrokAuth;
/// Minimal runner whose external command returns a fixed result.
struct FakeRunner {
external_result: Option<GrokAuth>,
}
impl ExternalCommandRunner for FakeRunner {
fn run_external_command(&self, _command: &str) -> Option<GrokAuth> {
self.external_result.clone()
}
}
/// A failed binary run is a single-strike `Other` permanent failure that is
/// NON-sticky: it must age out via the TTL, never lock an external-binary
/// user out forever. (Flipping this to a sticky reason would be a silent
/// lockout regression.)
#[tokio::test]
async fn external_binary_failure_is_single_strike_non_sticky_permanent() {
let refresher = ExternalBinaryRefresher::new(
Arc::new(FakeRunner {
external_result: None,
}),
"auth-binary".into(),
);
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(error.reason, RefreshTokenFailedReason::Other);
assert!(
!error.reason.is_sticky(),
"external-binary failure must age out, not strand the user forever",
);
}
other => panic!("a failed binary run must be a permanent Other failure, got {other:?}"),
}
}
/// A binary that outlives the (test-shortened) timeout hits the `Elapsed`
/// arm and maps to the same non-sticky `Other` permanent failure.
#[tokio::test]
async fn external_binary_timeout_is_non_sticky_permanent() {
struct SlowRunner;
impl ExternalCommandRunner for SlowRunner {
fn run_external_command(&self, _command: &str) -> Option<GrokAuth> {
std::thread::sleep(std::time::Duration::from_millis(50));
Some(GrokAuth::test_default())
}
}
let refresher = ExternalBinaryRefresher::new(Arc::new(SlowRunner), "auth-binary".into())
.with_timeout(std::time::Duration::from_millis(5));
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(error.reason, RefreshTokenFailedReason::Other);
assert!(
!error.reason.is_sticky(),
"timeout must age out, not strand"
);
}
other => panic!("a timed-out binary must be a permanent Other failure, got {other:?}"),
}
}
#[tokio::test]
async fn external_binary_success_returns_fresh_token() {
let token = GrokAuth {
key: "ext-fresh".into(),
..GrokAuth::test_default()
};
let refresher = ExternalBinaryRefresher::new(
Arc::new(FakeRunner {
external_result: Some(token),
}),
"auth-binary".into(),
);
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::Success(auth) => assert_eq!(auth.key, "ext-fresh"),
other => panic!("a successful binary run must return Success, got {other:?}"),
}
}
}
crates/codegen/kigi-shell/src/auth/refresh/mod.rs
+205
View File
@@ -0,0 +1,205 @@
mod external_refresher;
mod oidc_refresher;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use crate::auth::manager::AuthManager;
pub(crate) use crate::auth::manager::RefreshReason;
use crate::auth::model::GrokAuth;
use external_refresher::ExternalBinaryRefresher;
pub(crate) use oidc_refresher::OidcRefresher;
/// Read-only view of `AuthManager` for refreshers. Enforces the
/// no-mutation contract on *credential* state at the type level: refreshers
/// hold `Arc<dyn AuthSnapshot>` and physically cannot call `update()`,
/// `clear()`, `hot_swap()`, or `refresh_chain()`.
pub(crate) trait AuthSnapshot: Send + Sync {
/// Read the current in-memory bearer outside the early-invalidation buffer.
fn current(&self) -> Option<GrokAuth>;
/// Read the expired in-memory bearer (for its `refresh_token`).
fn expired_auth(&self) -> Option<GrokAuth>;
/// Re-read auth.json from disk for the configured scope. Read-only w.r.t.
/// credentials, but may advance disk-observation state and emit transition
/// telemetry (not credential mutation).
fn read_disk_auth(&self) -> Option<GrokAuth>;
/// Whether the in-memory bearer is expired.
fn is_expired(&self) -> bool;
}
impl AuthSnapshot for AuthManager {
fn current(&self) -> Option<GrokAuth> {
self.current()
}
fn expired_auth(&self) -> Option<GrokAuth> {
self.expired_auth()
}
fn read_disk_auth(&self) -> Option<GrokAuth> {
self.read_disk_auth()
}
fn is_expired(&self) -> bool {
self.is_expired()
}
}
/// Capability to run the operator's external auth binary. Split out of
/// [`AuthSnapshot`] so OIDC refreshers (read-only) physically cannot reach it
/// (interface segregation); only [`ExternalBinaryRefresher`] depends on it.
pub(crate) trait ExternalCommandRunner: Send + Sync {
/// Run the external auth binary and return the parsed output.
fn run_external_command(&self, command: &str) -> Option<GrokAuth>;
}
impl ExternalCommandRunner for AuthManager {
fn run_external_command(&self, command: &str) -> Option<GrokAuth> {
self.run_external_refresh_command(command)
}
}
/// The credential a refresh would send to the IdP: disk refresh-token first,
/// then the expired in-mem bearer, then current (only on `ServerRejected`).
/// Single source of truth shared by [`OidcRefresher::refresh`] (the attempt) and
/// `AuthManager::attempted_verdict_key` (the verdict scope), so the two can't
/// drift. The caller supplies the disk read: the verdict path passes a
/// side-effect-free read, the refresher the observing one.
pub(crate) fn resolve_refresh_credential(
snap: &dyn AuthSnapshot,
disk_auth: Option<GrokAuth>,
reason: RefreshReason,
) -> Option<GrokAuth> {
disk_auth
.filter(|a| a.refresh_token.is_some())
.or_else(|| snap.expired_auth())
.or_else(|| {
(reason == RefreshReason::ServerRejected)
.then(|| snap.current())
.flatten()
})
}
/// Outcome of a refresh attempt. Data only -- `refresh_chain` handles mutations.
#[derive(Debug)]
#[must_use = "RefreshOutcome encodes a state transition; route it through refresh_chain"]
pub(crate) enum RefreshOutcome {
/// Authority returned a fresh token. Caller persists via `update()`.
Success(Box<GrokAuth>),
/// Terminal failure (e.g. invalid_grant), or a transient escalated to
/// `Other` after repeated blips. Caller records a verdict scoped to the
/// rejected credential and retains it (`RefreshTokenRejected` is sticky,
/// the rest age out past the TTL).
PermanentFailure {
error: crate::auth::error::RefreshTokenFailedError,
/// Key of the credential the refresher actually sent to the IdP, so
/// `refresh_chain` scopes the verdict to it. `None` when the authority
/// has no token key (external binary flow); the caller falls back to
/// its own resolution.
tried_key: Option<String>,
},
/// Transient / unknown failure. Caller may retry later. Message-only: the
/// underlying cause is logged structurally at the refresher, then flattened
/// here (the retry decision needs recoverability, not the source chain).
TransientFailure { message: String },
}
impl RefreshOutcome {
/// A fresh credential from the authority (hides the `Box`).
pub(crate) fn success(auth: GrokAuth) -> Self {
Self::Success(Box::new(auth))
}
/// Terminal failure for an already-classified reason against the credential
/// `tried_key` (the one actually sent to the IdP).
pub(crate) fn permanent(
reason: crate::auth::error::RefreshTokenFailedReason,
tried_key: Option<String>,
) -> Self {
Self::PermanentFailure {
error: reason.into(),
tried_key,
}
}
/// A retryable failure carrying a diagnostic message.
pub(crate) fn transient(message: impl Into<String>) -> Self {
Self::TransientFailure {
message: message.into(),
}
}
}
#[async_trait::async_trait]
pub(crate) trait TokenRefresher: Send + Sync {
/// Attempt to obtain a fresh token from the authority.
///
/// Implementations MUST NOT call auth_manager.update(), clear(),
/// hot_swap(), or any other state-mutating method. Return the
/// result and let refresh_chain handle all mutations.
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome;
}
pub(crate) fn build_refresher(
auth_manager: Arc<AuthManager>,
auth_provider_command: Option<String>,
) -> Arc<dyn TokenRefresher> {
match auth_provider_command {
Some(cmd) => {
let runner: Arc<dyn ExternalCommandRunner> = auth_manager;
Arc::new(ExternalBinaryRefresher::new(runner, cmd))
}
None => {
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
Arc::new(OidcRefresher::new(snapshot))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::{AuthMode, GrokAuth, GrokComConfig};
use chrono::{Duration, Utc};
/// auth_token_ttl makes is_token_expired use create_time + ttl for
/// External tokens without expires_at, instead of the 30-day fallback.
#[test]
fn token_ttl_expires_external_token_by_create_time() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig {
auth_token_ttl: Some(3600), // 1 hour
..GrokComConfig::default()
};
let mgr = AuthManager::new(dir.path(), cfg);
// Token created 2 hours ago, no expires_at. With auth_token_ttl=3600,
// is_token_expired should return true (age 2h > ttl 1h).
let old_token = GrokAuth {
key: "old-external-token".into(),
auth_mode: AuthMode::External,
create_time: Utc::now() - Duration::hours(2),
expires_at: None,
..GrokAuth::test_default()
};
mgr.hot_swap(old_token);
assert!(
mgr.current().is_none(),
"expired external token via auth_token_ttl"
);
assert!(mgr.is_expired());
// Fresh token created just now — should be valid.
let new_token = GrokAuth {
key: "new-external-token".into(),
auth_mode: AuthMode::External,
create_time: Utc::now(),
expires_at: None,
..GrokAuth::test_default()
};
mgr.hot_swap(new_token);
assert!(
mgr.current().is_some(),
"fresh external token should be valid"
);
}
}
crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher.rs
+260
View File
@@ -0,0 +1,260 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::manager::RefreshReason;
use crate::auth::oidc::OidcRefreshResult;
use super::{AuthSnapshot, RefreshOutcome, TokenRefresher};
#[cfg(test)]
use crate::auth::manager::AuthManager;
/// Escalate to `PermanentFailure` after this many consecutive transient
/// failures (then `PERMANENT_FAILURE_TTL` allows recovery). OIDC tolerates more
/// blips than `ExternalBinaryRefresher` (1) since network refreshes flake more
/// than a local binary.
const MAX_CONSECUTIVE_TRANSIENT_FAILURES: u32 = 3;
/// Consecutive transient-failure budget, scoped to the credential it accrued
/// against. Held under one lock so the credential check, reset, and increment
/// are a single atomic step.
#[derive(Default)]
struct TransientBudget {
/// Credential the count belongs to. A different credential (e.g. after
/// re-login on this long-lived refresher) re-arms the budget so a fresh,
/// valid token never inherits a dead one's escalation.
key: Option<String>,
count: u32,
}
pub(crate) struct OidcRefresher {
auth: Arc<dyn AuthSnapshot>,
transient_budget: parking_lot::Mutex<TransientBudget>,
}
impl OidcRefresher {
pub(crate) fn new(auth: Arc<dyn AuthSnapshot>) -> Self {
Self {
auth,
transient_budget: parking_lot::Mutex::new(TransientBudget::default()),
}
}
/// Clear the transient-blip budget on refresh progress (a fresh token or an
/// adopted sibling token), so later blips start from a full budget.
fn note_refresh_progress(&self) {
*self.transient_budget.lock() = TransientBudget::default();
}
fn record_transient_failure(
&self,
message: String,
tried_key: Option<String>,
) -> RefreshOutcome {
let escalate = {
let mut budget = self.transient_budget.lock();
// Re-arm when the credential changes so a fresh token never inherits
// a prior credential's accrued blips.
if budget.key != tried_key {
budget.key = tried_key.clone();
budget.count = 0;
}
budget.count += 1;
let escalate = budget.count >= MAX_CONSECUTIVE_TRANSIENT_FAILURES;
// On escalation reset the count so the next TTL window gets the full
// budget (the verdict gates refresh() meanwhile). The key is left in
// place; a same-key retry resumes from zero, a new key re-arms.
if escalate {
budget.count = 0;
}
escalate
};
if escalate {
tracing::warn!(%message, "auth: escalating consecutive transient failures to permanent");
RefreshOutcome::permanent(RefreshTokenFailedReason::Other, tried_key)
} else {
RefreshOutcome::transient(message)
}
}
/// One-shot retry with disk's RT after `invalid_grant`.
///
/// If disk already has a valid (unexpired) AT with a different key,
/// adopt it directly, without consuming the disk's RT in another IdP
/// call. This prevents cascading `invalid_grant` when a sibling
/// already refreshed and wrote a valid token.
async fn retry_with_fresh_disk_token(
&self,
tried: &crate::auth::GrokAuth,
) -> Option<RefreshOutcome> {
let disk_now = self.auth.read_disk_auth()?;
// If disk has a valid AT that differs from what we tried,
// a sibling already refreshed. Adopt directly — no IdP call.
if !crate::auth::is_expired(&disk_now) && disk_now.key != tried.key {
crate::unified_log::info(
"oidc refresh: disk has valid AT, adopting instead of consuming RT",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&disk_now.key),
"tried_key_prefix": crate::auth::token_suffix(&tried.key),
})),
);
self.note_refresh_progress();
return Some(RefreshOutcome::success(disk_now));
}
if disk_now.refresh_token.is_none()
|| disk_now.refresh_token.as_deref() == tried.refresh_token.as_deref()
{
return None;
}
crate::unified_log::info(
"oidc refresh retrying with disk token",
None,
Some(serde_json::json!({
"tried_rt_prefix": tried
.refresh_token
.as_deref()
.map(crate::auth::token_suffix),
"disk_rt_prefix": disk_now
.refresh_token
.as_deref()
.map(crate::auth::token_suffix),
})),
);
match crate::auth::oidc::oidc_token_exchange(&disk_now).await {
OidcRefreshResult::Success(new_auth) => {
self.note_refresh_progress();
Some(RefreshOutcome::Success(new_auth))
}
OidcRefreshResult::TerminalError { reason } => {
crate::unified_log::warn(
"oidc refresh disk retry exhausted",
None,
Some(serde_json::json!({ "reason": format!("{reason:?}") })),
);
Some(RefreshOutcome::permanent(
reason,
Some(disk_now.key.clone()),
))
}
OidcRefreshResult::Failed => {
Some(RefreshOutcome::transient("OIDC disk-retry refresh failed"))
}
}
}
}
#[async_trait::async_trait]
impl TokenRefresher for OidcRefresher {
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
crate::unified_log::debug(
"oidc refresh enter",
None,
Some(serde_json::json!({
"reason": format!("{reason:?}"),
"has_current": self.auth.current().is_some(),
"is_expired": self.auth.is_expired(),
})),
);
let disk_auth = self.auth.read_disk_auth();
// Short-circuit: if disk has a valid unexpired AT that differs
// from in-memory, a sibling refreshed between refresh_chain
// step 2 (disk check under lock) and here. Adopt it directly,
// no IdP call needed.
if let Some(ref d) = disk_auth
&& !crate::auth::is_expired(d)
&& self.auth.current().map(|a| a.key).as_deref() != Some(&d.key)
{
crate::unified_log::info(
"oidc refresh: sibling refreshed, adopting valid disk AT",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&d.key),
})),
);
self.note_refresh_progress();
return RefreshOutcome::success(d.clone());
}
let auth = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason);
let Some(auth) = auth else {
crate::unified_log::warn(
"oidc refresh no token available",
None,
Some(serde_json::json!({ "reason": format!("{reason:?}") })),
);
return RefreshOutcome::transient("no token with refresh_token available");
};
crate::unified_log::info(
"oidc refresh attempting idp",
None,
Some(serde_json::json!({
"has_rt": auth.refresh_token.is_some(),
"issuer": auth.oidc_issuer,
"client_id": auth.oidc_client_id,
"expires_at": auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
match crate::auth::oidc::oidc_token_exchange(&auth).await {
OidcRefreshResult::Success(new_auth) => {
self.note_refresh_progress();
RefreshOutcome::Success(new_auth)
}
OidcRefreshResult::TerminalError { reason } => {
// Sibling-rotation race: disk may hold a
// fresher RT than the one we tried. One-shot retry.
if reason == RefreshTokenFailedReason::RefreshTokenRejected
&& let Some(retry_outcome) = self.retry_with_fresh_disk_token(&auth).await
{
return retry_outcome;
}
RefreshOutcome::permanent(reason, Some(auth.key.clone()))
}
OidcRefreshResult::Failed => {
tracing::warn!(
refresh_reason = ?reason,
user_id = %auth.user_id,
has_refresh_token = auth.refresh_token.is_some(),
issuer = ?auth.oidc_issuer,
client_id = ?auth.oidc_client_id,
expires_at = ?auth.expires_at,
"auth: OIDC token refresh failed"
);
crate::unified_log::error(
"oidc refresh failed",
None,
Some(serde_json::json!({
"has_refresh_token": auth.refresh_token.is_some(),
"auth_mode": format!("{:?}", auth.auth_mode),
"issuer": auth.oidc_issuer,
"client_id": auth.oidc_client_id,
"expires_at": auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
self.record_transient_failure(
"OIDC token refresh failed".into(),
Some(auth.key.clone()),
)
}
}
}
}
#[cfg(test)]
#[path = "oidc_refresher_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "auth_backend_contract_tests.rs"]
mod auth_backend_contract_tests;
crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher_tests.rs
+1311
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/storage.rs
+512
View File
@@ -0,0 +1,512 @@
use std::fs::File;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use super::model::{API_KEY_SCOPE, AuthMode, AuthStore, GrokAuth, lookup_auth};
/// RAII guard for an exclusive advisory lock on `auth.json.lock`.
/// The lock is released when the inner `File` is dropped (closing the FD).
pub(crate) struct AuthFileLock {
pub(super) _file: File,
}
impl AuthFileLock {
/// Returns `true` while this guard still refers to the **live**
/// `auth.json.lock` inode.
///
/// A waiter that finds a holder stuck past the stale-lock timeout breaks
/// the lock by `unlink`ing the file and recreating it on a fresh inode
/// (see [`crate::auth::manager::lock`]). The usual cause of a "stuck"
/// holder is a process **suspended across system sleep** while holding the
/// lock: it stays alive (so the kernel never releases its flock) yet makes
/// no progress, so siblings break it. When such a holder resumes, its
/// flock lives on the now-deleted inode — it no longer holds the live lock
/// even though this `AuthFileLock` still exists.
///
/// Callers about to perform an irreversible, lock-protected action
/// (sending a refresh token to the IdP, writing `auth.json`) MUST
/// re-validate first; otherwise two processes can spend the same refresh
/// token and trip token-family revocation.
///
/// Non-Unix has no inode concept, so this conservatively returns `true`.
#[cfg(unix)]
pub(crate) fn still_live(&self, auth_json_path: &Path) -> bool {
use std::os::unix::fs::MetadataExt;
let lock_path = auth_json_path.with_file_name("auth.json.lock");
let (Ok(fd_meta), Ok(path_meta)) = (self._file.metadata(), std::fs::metadata(&lock_path))
else {
// Lock file gone or unreadable → we no longer hold the live lock.
return false;
};
fd_meta.ino() == path_meta.ino() && fd_meta.dev() == path_meta.dev()
}
#[cfg(not(unix))]
pub(crate) fn still_live(&self, _auth_json_path: &Path) -> bool {
true
}
}
pub fn read_auth_json(auth_file: &Path) -> std::io::Result<AuthStore> {
let mut file = File::open(auth_file)?;
let mut contents = String::new();
file.read_to_string(&mut contents)?;
// Empty files are valid (recover from prior crash/partial write).
let trimmed = contents.trim();
if trimmed.is_empty() {
return Ok(AuthStore::new());
}
let map = serde_json::from_str(trimmed)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
Ok(map)
}
/// Read auth.json, returning an empty map if the file does not exist.
///
/// Non-empty corrupt JSON, permission errors, etc. are returned as errors
/// so the caller can decide whether to skip the write (to avoid clobbering
/// sibling scopes).
///
/// Kept for the test-only `persist_and_swap` and as a strict reader.
#[cfg_attr(
not(test),
expect(
dead_code,
reason = "used from tests only; remove expect when wired in production"
)
)]
pub(crate) fn read_auth_json_or_empty(auth_file: &Path) -> std::io::Result<AuthStore> {
match read_auth_json(auth_file) {
Ok(map) => Ok(map),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(AuthStore::new()),
Err(e) => Err(e),
}
}
/// Best-effort backup of a corrupt (unparseable) auth.json.
///
/// If the file exists and `read_auth_json` fails with `InvalidData`,
/// it is renamed to `auth.json.corrupt.<millis>` (sibling in the same
/// directory) and the backup path is returned. Used before recovery
/// writes so the original bytes are never silently lost.
pub(crate) fn backup_corrupt_auth_file(path: &Path) -> Option<PathBuf> {
if !path.exists() {
return None;
}
if read_auth_json(path).is_ok() {
return None;
}
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
let file_name = path
.file_name()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_else(|| "auth.json".to_string());
let backup_name = format!("{}.corrupt.{}", file_name, ts);
let backup = path.with_file_name(backup_name);
match std::fs::rename(path, &backup) {
Ok(()) => {
tracing::warn!(
original = %path.display(),
backup = %backup.display(),
"auth: backed up corrupt auth.json before recovery write"
);
// Must reach unified.jsonl: the tracing line above is invisible
// in production captures, and this is the only record of both
// the corruption and where the original bytes went.
kigi_log::unified_log::error(
"auth: corrupt auth.json backed up",
None,
Some(serde_json::json!({
"original": path.display().to_string(),
"backup": backup.display().to_string(),
})),
);
Some(backup)
}
Err(e) => {
tracing::warn!(error = %e, "auth: failed to rename corrupt auth.json for backup");
kigi_log::unified_log::error(
"auth: corrupt auth.json backup failed",
None,
Some(serde_json::json!({
"original": path.display().to_string(),
"error": e.to_string(),
})),
);
None
}
}
}
/// Read auth.json for an upcoming write, with recovery for corrupt files.
///
/// - Missing/empty → empty map (safe to write fresh)
/// - Valid JSON → parsed map
/// - Non-empty corrupt JSON → backs up to `auth.json.corrupt.<millis>`,
/// then returns empty map so the caller can write the new credential.
///
/// Other I/O errors (PermissionDenied, etc.) are still returned as errors.
pub(crate) fn read_auth_json_or_empty_recovering_corrupt(
auth_file: &Path,
) -> std::io::Result<AuthStore> {
match read_auth_json(auth_file) {
Ok(map) => Ok(map),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(AuthStore::new()),
Err(e) if e.kind() == std::io::ErrorKind::InvalidData => {
let _ = backup_corrupt_auth_file(auth_file);
Ok(AuthStore::new())
}
Err(e) => Err(e),
}
}
/// Persist `auth.json`, preferring a crash-safe atomic write but falling
/// back to a non-atomic in-place write when the disk is full.
///
/// The atomic path (temp + rename) needs free space >= the file size,
/// because the old file and a full temp copy coexist until the rename. On a
/// nearly-full disk that temp copy can fail with `StorageFull` (ENOSPC)
/// even though the credentials themselves are tiny. When that happens we
/// retry with an in-place truncate+write, which only needs the freed blocks
/// of the old file — far less than the temp-copy approach.
///
/// The in-place path is non-atomic, with two accepted trade-offs:
/// - If the in-place write itself fails (e.g. a concurrent process grabs the
/// just-freed blocks, or a crash mid-write), the prior bytes are restored
/// best-effort so a torn/empty file never *replaces* the previous on-disk
/// credential — on-disk state ends up no worse than before the attempt.
/// - Unlocked concurrent readers can still observe a torn (partial) file
/// during the brief write window; a partial file is healed on the next
/// read via [`read_auth_json_or_empty_recovering_corrupt`] (backup +
/// relogin). This window is inherent to any sub-1×-free single-file
/// replace and is preferable to persisting nothing at all, which would
/// leave every concurrent process with a stale, already-revoked token.
pub(super) fn write_auth_json(auth_file: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
write_auth_json_with(auth_file, auth_store, write_auth_json_atomic)
}
/// Dispatch helper: run `atomic`, and on `StorageFull` fall back to an
/// in-place write. Split out (with `atomic` injectable) so the disk-full
/// fallback is unit-testable without an actually-full filesystem.
fn write_auth_json_with(
auth_file: &Path,
auth_store: &AuthStore,
atomic: fn(&Path, &AuthStore) -> std::io::Result<()>,
) -> std::io::Result<()> {
match atomic(auth_file, auth_store) {
Err(e) if e.kind() == std::io::ErrorKind::StorageFull => {
tracing::warn!(
path = %auth_file.display(),
"auth: disk full during atomic write, falling back to in-place write"
);
// Must reach unified.jsonl: a silent in-memory-only credential
// (the prior behavior) leaves sibling processes with a stale
// refresh token and no record of why. Surface it loudly.
kigi_log::unified_log::warn(
"auth: disk full, falling back to non-atomic in-place write",
None,
Some(serde_json::json!({
"path": auth_file.display().to_string(),
})),
);
write_auth_json_in_place(auth_file, auth_store)
}
other => other,
}
}
/// Serialize `auth_store` to `path` (truncate + rewrite), owner-only (0o600)
/// and `fsync`'d. Shared core of the atomic path (which targets the temp
/// file) and the in-place fallback (which targets `auth.json` directly).
///
/// Uses streaming `to_writer_pretty` through a `BufWriter` to avoid
/// allocating the entire JSON string in memory — eliminates OOM risk under
/// severe memory pressure.
fn write_store_to(path: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
use crate::util::secure_file::open_secure_file;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let file = open_secure_file(path)?;
let mut writer = std::io::BufWriter::new(file);
serde_json::to_writer_pretty(&mut writer, auth_store)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
writer.flush()?;
writer
.into_inner()
.map_err(|e| e.into_error())?
.sync_all()?;
#[cfg(windows)]
{
crate::util::secure_file::set_windows_secure_permissions(path)?;
}
Ok(())
}
/// Atomic write: tmp + rename. Unix `rename(2)` replaces atomically;
/// Windows `rename` requires removing the target first.
fn write_auth_json_atomic(auth_file: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
let tmp = auth_file.with_extension(format!("json.{}.tmp", std::process::id()));
write_store_to(&tmp, auth_store)?;
#[cfg(windows)]
{
let _ = std::fs::remove_file(auth_file);
}
std::fs::rename(&tmp, auth_file)?;
Ok(())
}
/// Non-atomic fallback: truncate and rewrite `auth.json` in place.
///
/// Used only when [`write_auth_json_atomic`] fails with `StorageFull`.
/// Opening with truncation first frees the old content's blocks before the
/// new bytes are written, so this needs only the file size in free space
/// rather than the temp-copy approach's file-size-of-headroom.
///
/// Truncation is destructive, so the prior bytes are snapshotted first and
/// restored best-effort if the rewrite fails partway — a failed fallback
/// must not leave an empty/torn file where a parseable (if stale) credential
/// used to be. A partial file that survives (because even the restore failed)
/// is healed on the next read via [`read_auth_json_or_empty_recovering_corrupt`].
fn write_auth_json_in_place(auth_file: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
write_auth_json_in_place_with(auth_file, auth_store, write_store_to)
}
/// Inner of [`write_auth_json_in_place`] with `write` injectable so the
/// rollback-on-failure path is unit-testable without an actually-full disk.
fn write_auth_json_in_place_with(
auth_file: &Path,
auth_store: &AuthStore,
write: fn(&Path, &AuthStore) -> std::io::Result<()>,
) -> std::io::Result<()> {
// Snapshot the prior bytes so a torn/empty write can be rolled back to
// the previous on-disk credential. `None` when the file is absent.
let prior = std::fs::read(auth_file).ok();
match write(auth_file, auth_store) {
Ok(()) => Ok(()),
Err(e) => {
if let Some(prior) = prior
&& let Err(restore_err) = restore_prior_bytes(auth_file, &prior)
{
tracing::warn!(
error = %restore_err,
"auth: failed to restore prior auth.json after in-place write failure"
);
}
Err(e)
}
}
}
/// Best-effort rollback: rewrite `bytes` (owner-only, `fsync`'d) after a
/// failed in-place write so a torn/empty file does not replace the prior
/// credential.
fn restore_prior_bytes(auth_file: &Path, bytes: &[u8]) -> std::io::Result<()> {
use crate::util::secure_file::open_secure_file;
let mut file = open_secure_file(auth_file)?;
file.write_all(bytes)?;
file.sync_all()?;
#[cfg(windows)]
{
crate::util::secure_file::set_windows_secure_permissions(auth_file)?;
}
Ok(())
}
/// Read a single auth token from `auth.json` by scope key.
/// Falls back to the legacy `https://accounts.x.ai/sign-in` scope key
/// when the requested scope is not found (devbox auth.json migration).
pub fn read_token_by_scope(kigi_home: &Path, scope: &str) -> anyhow::Result<String> {
let path = kigi_home.join("auth.json");
let store =
read_auth_json(&path).map_err(|_| anyhow::anyhow!("Not logged in. Run `grok login`."))?;
lookup_auth(&store, scope).map(|a| a.key).ok_or_else(|| {
anyhow::anyhow!("Your auth token is invalid. Run `grok login` to re-authenticate.")
})
}
/// Read the API key from the `xai::api_key` scope in auth.json.
pub fn read_api_key(kigi_home: &Path) -> Option<String> {
let path = kigi_home.join("auth.json");
let map = read_auth_json(&path).ok()?;
map.get(API_KEY_SCOPE).map(|a| a.key.clone())
}
/// Store a plain API key in auth.json under the `xai::api_key` scope.
///
/// Uses the corrupt-recovery reader so a malformed auth.json (e.g. from a
/// previous crash) can be healed when the user sets an API key.
pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> {
let path = kigi_home.join("auth.json");
let mut map = read_auth_json_or_empty_recovering_corrupt(&path)?;
map.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
key: api_key.to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
},
);
write_auth_json(&path, &map)
}
/// Remove the `xai::api_key` scope from auth.json.
pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> {
let path = kigi_home.join("auth.json");
if let Ok(mut map) = read_auth_json(&path) {
map.remove(API_KEY_SCOPE);
if map.is_empty() {
let _ = std::fs::remove_file(&path);
} else {
write_auth_json(&path, &map)?;
}
}
Ok(())
}
#[cfg(test)]
mod write_fallback_tests {
use super::*;
fn sample_store() -> AuthStore {
let mut map = AuthStore::new();
map.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
key: "secret-key".to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
},
);
map
}
fn read_key(path: &Path) -> Option<String> {
read_auth_json(path)
.ok()
.and_then(|m| m.get(API_KEY_SCOPE).map(|a| a.key.clone()))
}
fn fake_storage_full(_: &Path, _: &AuthStore) -> std::io::Result<()> {
Err(std::io::Error::from(std::io::ErrorKind::StorageFull))
}
fn fake_permission_denied(_: &Path, _: &AuthStore) -> std::io::Result<()> {
Err(std::io::Error::from(std::io::ErrorKind::PermissionDenied))
}
/// Simulates an in-place write that truncates the file (destroying the
/// old content, as `open_secure_file` does) and then fails partway — the
/// torn-write case the rollback must recover from.
fn fake_truncate_then_fail(path: &Path, _: &AuthStore) -> std::io::Result<()> {
crate::util::secure_file::open_secure_file(path)?; // truncates to 0 bytes
Err(std::io::Error::from(std::io::ErrorKind::StorageFull))
}
#[test]
fn in_place_write_roundtrips() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
write_auth_json_in_place(&path, &sample_store()).unwrap();
assert_eq!(read_key(&path).as_deref(), Some("secret-key"));
}
#[cfg(unix)]
#[test]
fn in_place_write_is_owner_only() {
use std::os::unix::fs::PermissionsExt;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
write_auth_json_in_place(&path, &sample_store()).unwrap();
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o777, 0o600, "in-place write must stay 0o600");
}
/// A `StorageFull` (ENOSPC) failure on the atomic path must fall back to
/// the in-place write so the credential still lands on disk.
#[test]
fn falls_back_to_in_place_on_storage_full() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
write_auth_json_with(&path, &sample_store(), fake_storage_full).unwrap();
assert_eq!(
read_key(&path).as_deref(),
Some("secret-key"),
"disk-full atomic write must fall back to a successful in-place write"
);
}
/// Non-ENOSPC errors must propagate unchanged and must NOT trigger the
/// in-place fallback (e.g. a permission error should not write the file).
#[test]
fn propagates_non_storage_full_errors() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
let err = write_auth_json_with(&path, &sample_store(), fake_permission_denied).unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::PermissionDenied);
assert!(!path.exists(), "non-ENOSPC failure must not write the file");
}
/// The normal (real atomic) path still works end to end.
#[test]
fn atomic_write_roundtrips() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
write_auth_json(&path, &sample_store()).unwrap();
assert_eq!(read_key(&path).as_deref(), Some("secret-key"));
}
/// A fallback write that truncates then fails must roll back to the prior
/// bytes instead of leaving an empty/torn file — otherwise a second
/// disk-full failure would destroy a previously-valid credential.
#[test]
fn in_place_restores_prior_bytes_on_failure() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
// Seed a valid prior credential.
write_auth_json_in_place(&path, &sample_store()).unwrap();
assert_eq!(read_key(&path).as_deref(), Some("secret-key"));
let mut replacement = AuthStore::new();
replacement.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
key: "replacement-key".to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
},
);
let err = write_auth_json_in_place_with(&path, &replacement, fake_truncate_then_fail)
.unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::StorageFull);
assert_eq!(
read_key(&path).as_deref(),
Some("secret-key"),
"a failed in-place write must restore the prior credential, not leave an empty file"
);
}
/// Rollback after a failed write must keep the file owner-only (0o600).
#[cfg(unix)]
#[test]
fn in_place_restore_is_owner_only() {
use std::os::unix::fs::PermissionsExt;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
write_auth_json_in_place(&path, &sample_store()).unwrap();
let _ = write_auth_json_in_place_with(&path, &sample_store(), fake_truncate_then_fail);
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o777, 0o600, "restored file must stay 0o600");
}
}
crates/codegen/kigi-shell/src/auth/token_type.rs
+55
View File
@@ -0,0 +1,55 @@
use crate::auth::model::{AuthMode, GrokAuth};
/// What kind of bearer is loaded right now. Dispatch key for
/// `auth()`, `unauthorized_recovery()`, and proactive refresh.
///
/// Not a session classifier — use `is_session_based_method` for that.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum TokenType {
/// OIDC/OAuth2 session with a refresh_token available.
OidcSession,
/// Legacy web-login session or OIDC without a refresh_token.
LegacySession,
/// External auth binary provides tokens.
ExternalBinary,
/// Plain API key (no refresh possible).
ApiKey,
/// No credentials loaded.
None,
}
impl TokenType {
/// Classify the loaded credential (pure; no manager state).
pub(crate) fn from_auth(auth: Option<&GrokAuth>) -> Self {
match auth {
None => Self::None,
// Oidc without a refresh_token degrades to the unrefreshable LegacySession shape.
Some(a) => match a.auth_mode {
AuthMode::Oidc if a.refresh_token.is_some() => Self::OidcSession,
AuthMode::Oidc | AuthMode::WebLogin => Self::LegacySession,
AuthMode::External => Self::ExternalBinary,
AuthMode::ApiKey => Self::ApiKey,
},
}
}
/// `true` for types that can be silently refreshed (OIDC, external binary).
pub(crate) fn is_refreshable(self) -> bool {
matches!(self, Self::OidcSession | Self::ExternalBinary)
}
}
#[cfg(test)]
mod tests {
//! Per-variant matrix for `is_refreshable`.
use super::*;
#[test]
fn is_refreshable_matrix() {
assert!(TokenType::OidcSession.is_refreshable());
assert!(TokenType::ExternalBinary.is_refreshable());
assert!(!TokenType::LegacySession.is_refreshable());
assert!(!TokenType::ApiKey.is_refreshable());
assert!(!TokenType::None.is_refreshable());
}
}