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

M1/F1: Kimi Code OAuth device-code flow

Browse Source 
Replace the xAI OAuth stack with the Kimi device authorization grant:
- kimi_oauth.rs wire layer (device_authorization + token poll + refresh
  against kigi_env::oauth_host(); client_id per PRD; retryable statuses
  429/5xx with backoff; expired_token restarts authorization)
- X-Msh-Device-{Name,Model,Id} headers; device_id minted uuid4-hex at
  ~/.kigi/device_id (0600)
- Storage: system keyring service `kigi`, entry `oauth/kimi-code`
  (macOS/Windows native backends), atomic-file fallback under ~/.kigi;
  official client's keyring/~/.kimi never touched
- Refresh manager: 60s tick, threshold max(300, expires_in*0.5),
  401-tombstone keyed by rejected refresh token with 300s cooldown and
  rotation auto-clear, cross-process lock with sibling-adoption
  triple-check, sleep/wake forced refresh
- Deleted xAI machinery: enterprise OIDC (PKCE/JWKS/teams), devbox login,
  external auth provider, JWT tier gating + subscription paywall stack,
  X-XAI-Token-Auth marker headers, ZDR gates, /user enrichment
- kigi login / TUI /login both drive the device flow; login-host display
  now derives from kigi_env::oauth_host()
- 264 auth unit/wiremock tests; live contract probe of
  auth.kimi.com/api/oauth/device_authorization matches the wire shapes

Gates: check/clippy --all-targets clean, fmt, deny ok, kigi-shell lib
5131 tests green.
This commit is contained in:
ZacharyZhang-NY
2026-07-17 07:37:29 -04:00
parent d6c20fc13f
commit 021b82443d
117 changed files with 4052 additions and 19900 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/auth/attribution.rs
+10 -10
View File
@@ -370,7 +370,7 @@ pub(crate) fn record_auth_401(
/// 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`.
/// and the mint/expiry fields from the resulting `KimiAuth`.
///
/// `is_stale_snapshot` is `true` only when the live `current()` token
/// differs from the bearer the client sent. When `current()` returns
@@ -390,7 +390,7 @@ fn compute_attribution_payload(
// 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
// Single read-lock acquisition: pull the live `KimiAuth` (or
// `None`) once and derive every other field from it.
let current_auth = auth_manager.current();
let current_prefix_owned: Option<String> = current_auth
@@ -411,7 +411,7 @@ fn compute_attribution_payload(
//
// 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`
// `kimi_code_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
@@ -441,7 +441,7 @@ mod tests {
use chrono::{Duration, Utc};
use crate::auth::{AuthManager, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig};
use super::*;
@@ -449,17 +449,17 @@ mod tests {
/// 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 cfg = KimiCodeConfig::default();
let am = AuthManager::new(dir.path(), cfg);
(dir, am)
}
fn fresh_auth(key: &str) -> GrokAuth {
GrokAuth {
fn fresh_auth(key: &str) -> KimiAuth {
KimiAuth {
key: key.to_string(),
create_time: Utc::now(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
}
}
@@ -553,12 +553,12 @@ mod tests {
#[test]
fn legacy_token_uses_two_branch_fallback() {
let (_dir, am) = empty_auth_manager();
let auth = GrokAuth {
let auth = KimiAuth {
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()
..KimiAuth::test_default()
};
am.hot_swap(auth);
crates/codegen/kigi-shell/src/auth/config.rs
+30 -410
View File
@@ -1,423 +1,43 @@
use super::model::TEAM_PRINCIPAL_TYPE;
//! Kimi Code auth configuration.
//!
//! The wire endpoints come from [`kigi_env`] (`oauth_host()`, overridable via
//! `KIGI_OAUTH_HOST`) and the client id is fixed
//! ([`crate::auth::kimi_oauth::KIMI_CODE_CLIENT_ID`]), so this config carries
//! no per-deployment OAuth knobs. The struct is kept (deserialized from the
//! agent config TOML) as the extension point for future auth options.
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)]
/// Persisted-credential scope key for the Kimi Code OAuth session — both the
/// auth.json map key and the system-keyring entry name (service `kigi`).
pub const KIMI_CODE_OAUTH_SCOPE: &str = "oauth/kimi-code";
/// Auth configuration block (`[kimi_code_config]` in the agent config).
/// Currently empty: the OAuth host and client id are fixed by the
/// environment crate.
#[derive(Debug, Clone, Default, 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 struct KimiCodeConfig {}
impl KimiCodeConfig {
/// The persisted-credential scope key for this configuration.
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(),
})
KIMI_CODE_OAUTH_SCOPE.to_owned()
}
}
#[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");
fn auth_scope_is_the_kimi_code_key() {
assert_eq!(KimiCodeConfig::default().auth_scope(), "oauth/kimi-code");
}
#[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);
fn deserializes_from_empty_toml() {
let cfg: KimiCodeConfig = toml::from_str("").expect("empty config parses");
assert_eq!(cfg.auth_scope(), KIMI_CODE_OAUTH_SCOPE);
}
}
crates/codegen/kigi-shell/src/auth/credential_provider.rs
+26 -28
View File
@@ -1,5 +1,5 @@
use crate::auth::AuthManager;
use crate::util::grok_auth_credentials::GrokAuthCredentials;
use crate::util::kigi_auth_credentials::KigiAuthCredentials;
use kigi_auth::{
AuthCredentialProvider, CredentialSnapshot, HttpAuth, StaticAuthCredentialProvider,
};
@@ -7,7 +7,7 @@ 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> {
fn api_key_id_for(auth: Option<&crate::auth::KimiAuth>) -> Option<String> {
auth.filter(|a| matches!(a.auth_mode, crate::auth::AuthMode::ApiKey))
.map(|a| crate::agent::config::deployment_id_from_key(&a.key))
}
@@ -15,7 +15,7 @@ fn api_key_id_for(auth: Option<&crate::auth::GrokAuth>) -> Option<String> {
/// delegates to `AuthManager::unauthorized_recovery`.
pub struct ShellAuthCredentialProvider {
auth_manager: Arc<AuthManager>,
static_credentials: GrokAuthCredentials,
static_credentials: KigiAuthCredentials,
}
impl ShellAuthCredentialProvider {
pub(crate) fn new(
@@ -23,7 +23,7 @@ impl ShellAuthCredentialProvider {
deployment_key: Option<String>,
alpha_test_key: Option<String>,
) -> Self {
let mut static_credentials = GrokAuthCredentials::new(None);
let mut static_credentials = KigiAuthCredentials::new(None);
static_credentials.deployment_key = deployment_key;
static_credentials.alpha_test_key = alpha_test_key;
Self {
@@ -61,18 +61,19 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
};
}
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());
// The Kimi token response carries no account info; `user_id` stays
// empty until a later feature surfaces it.
let user_id = auth
.as_ref()
.map(|a| a.user_id.clone())
.filter(|id| !id.is_empty());
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 {
@@ -81,15 +82,12 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
}
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::KimiAuth;
use crate::auth::KimiCodeConfig;
use crate::auth::manager::AuthManager;
use chrono::{Duration as ChronoDuration, Utc};
use kigi_auth::AuthCredentialProvider;
@@ -128,19 +126,19 @@ mod tests {
}
}
}
fn make_auth(key: &str, expires_in: ChronoDuration) -> GrokAuth {
GrokAuth {
fn make_auth(key: &str, expires_in: ChronoDuration) -> KimiAuth {
KimiAuth {
key: key.to_string(),
user_id: "test-user".to_string(),
create_time: Utc::now(),
expires_at: Some(Utc::now() + expires_in),
..GrokAuth::test_default()
..KimiAuth::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());
fn make_manager(dir: &tempfile::TempDir, initial: Option<KimiAuth>) -> Arc<AuthManager> {
let mgr = AuthManager::new(dir.path(), KimiCodeConfig::default());
if let Some(auth) = initial {
mgr.hot_swap(auth);
}
@@ -210,16 +208,16 @@ mod tests {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(
dir.path(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
mgr.hot_swap(GrokAuth {
mgr.hot_swap(KimiAuth {
key: "stale".into(),
auth_mode: crate::auth::AuthMode::Oidc,
auth_mode: crate::auth::AuthMode::OAuth,
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()
..KimiAuth::test_default()
});
struct OkRefresher {
calls: Arc<std::sync::atomic::AtomicU32>,
@@ -231,14 +229,14 @@ mod tests {
_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 {
crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth {
key: "fresh".into(),
auth_mode: crate::auth::AuthMode::Oidc,
auth_mode: crate::auth::AuthMode::OAuth,
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()
..KimiAuth::test_default()
}))
}
}
@@ -282,11 +280,11 @@ mod tests {
Some(deployment_id_from_key("xai-token-EX").as_str())
);
assert!(dep.api_key_id.is_none());
let api_auth = GrokAuth {
let api_auth = KimiAuth {
key: "sk-apikey-xyz".into(),
auth_mode: crate::auth::AuthMode::ApiKey,
expires_at: Some(Utc::now() + ChronoDuration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
let api = ShellAuthCredentialProvider::new(make_manager(&dir, Some(api_auth)), None, None)
.snapshot();
crates/codegen/kigi-shell/src/auth/devbox_login_stub.rs
-37
View File
@@ -1,37 +0,0 @@
//! 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.rs
+297
View File
@@ -0,0 +1,297 @@
//! Device identity headers for the Kimi Code OAuth endpoints.
//!
//! Every OAuth call (device authorization, token poll, refresh) carries three
//! headers identifying this installation (PRD F1):
//!
//! - `X-Msh-Device-Name` — the local hostname
//! - `X-Msh-Device-Model` — an honest local OS/arch string (e.g.
//! "macOS 15.5 arm64"), ported from kimi-cli's `_device_model()`
//! - `X-Msh-Device-Id` — a uuid4 hex persisted at `~/.kigi/device_id`
//! (owner-only), created on first use
//!
//! All values are ASCII-sanitized (ported from kimi-cli's
//! `_ascii_header_value`) since HTTP header values must be ASCII.
use std::path::PathBuf;
use std::sync::OnceLock;
use anyhow::Context as _;
/// Sanitize a header value to ASCII: non-ASCII bytes are dropped; an empty
/// result falls back to `"unknown"`. Port of kimi-cli `_ascii_header_value`.
pub(crate) fn ascii_header_value(value: &str) -> String {
let sanitized: String = value.chars().filter(char::is_ascii).collect();
let trimmed = sanitized.trim();
if trimmed.is_empty() {
"unknown".to_owned()
} else {
trimmed.to_owned()
}
}
/// The three device-identity headers sent on every OAuth call.
///
/// Errors when the persistent device id cannot be created (e.g. read-only
/// `~/.kigi`): the OAuth endpoints require `X-Msh-Device-Id`, so login cannot
/// proceed without it.
pub(crate) fn device_headers() -> anyhow::Result<[(&'static str, String); 3]> {
Ok([
("X-Msh-Device-Name", ascii_header_value(&device_name())),
("X-Msh-Device-Model", ascii_header_value(device_model())),
("X-Msh-Device-Id", ascii_header_value(&device_id()?)),
])
}
/// Local hostname (kimi-cli: `platform.node() or socket.gethostname()`).
fn device_name() -> String {
#[cfg(unix)]
{
let mut buf = [0u8; 256];
// SAFETY: buf is a valid writable buffer of the passed length.
let rc = unsafe { libc::gethostname(buf.as_mut_ptr().cast(), buf.len()) };
if rc == 0 {
let end = buf.iter().position(|&b| b == 0).unwrap_or(buf.len());
let name = String::from_utf8_lossy(&buf[..end]).into_owned();
if !name.trim().is_empty() {
return name;
}
}
"unknown".to_owned()
}
#[cfg(windows)]
{
std::env::var("COMPUTERNAME").unwrap_or_else(|_| "unknown".to_owned())
}
#[cfg(not(any(unix, windows)))]
{
"unknown".to_owned()
}
}
/// Honest local device-model string, computed once per process. Port of
/// kimi-cli `_device_model()`:
/// - macOS → `macOS {product_version} {arch}` (e.g. "macOS 15.5 arm64")
/// - Windows → `Windows {10|11} {arch}` (build ≥ 22000 reports 11)
/// - other → `{sysname} {kernel_release} {machine}`
pub(crate) fn device_model() -> &'static str {
static MODEL: OnceLock<String> = OnceLock::new();
MODEL.get_or_init(compute_device_model)
}
fn compute_device_model() -> String {
#[cfg(target_os = "macos")]
{
// Match Python's platform.machine() spelling on macOS.
let arch = match std::env::consts::ARCH {
"aarch64" => "arm64",
other => other,
};
match macos_product_version() {
Some(version) => format!("macOS {version} {arch}"),
None => format!("macOS {arch}"),
}
}
#[cfg(windows)]
{
let arch = std::env::consts::ARCH;
match windows_release() {
Some(release) => format!("Windows {release} {arch}"),
None => format!("Windows {arch}"),
}
}
#[cfg(not(any(target_os = "macos", windows)))]
{
let (sysname, release, machine) = uname_fields();
match (release, machine) {
(Some(r), Some(m)) => format!("{sysname} {r} {m}"),
(Some(r), None) => format!("{sysname} {r}"),
(None, Some(m)) => format!("{sysname} {m}"),
(None, None) => sysname,
}
}
}
/// macOS product version (e.g. "15.5") from the SystemVersion plist — the
/// same source Python's `platform.mac_ver()` reads.
#[cfg(target_os = "macos")]
fn macos_product_version() -> Option<String> {
let plist = std::fs::read_to_string("/System/Library/CoreServices/SystemVersion.plist").ok()?;
plist_string_value(&plist, "ProductVersion")
}
/// Extract `<key>{key}</key><string>value</string>` from a plist XML body.
#[cfg(target_os = "macos")]
fn plist_string_value(plist: &str, key: &str) -> Option<String> {
let key_tag = format!("<key>{key}</key>");
let after_key = &plist[plist.find(&key_tag)? + key_tag.len()..];
let start = after_key.find("<string>")? + "<string>".len();
let end = after_key.find("</string>")?;
(start <= end).then(|| after_key[start..end].trim().to_owned())
}
/// Windows major release ("10" or "11"), from the build number reported by
/// `cmd /c ver` (kimi-cli: `sys.getwindowsversion().build >= 22000` → 11).
#[cfg(windows)]
fn windows_release() -> Option<String> {
let output = std::process::Command::new("cmd")
.args(["/c", "ver"])
.output()
.ok()?;
let text = String::from_utf8_lossy(&output.stdout);
// "Microsoft Windows [Version 10.0.22631.3155]"
let version = text.split("Version").nth(1)?.trim();
let mut parts = version.trim_end_matches(']').split('.');
let major = parts.next()?.trim().to_owned();
let _minor = parts.next()?;
let build: u32 = parts.next()?.trim().parse().ok()?;
if major == "10" && build >= 22000 {
Some("11".to_owned())
} else {
Some(major)
}
}
/// `uname(2)` sysname / release / machine for Linux and other Unix.
#[cfg(all(unix, not(target_os = "macos")))]
fn uname_fields() -> (String, Option<String>, Option<String>) {
// SAFETY: utsname is a plain-old-data struct; uname fills it in.
let mut uts: libc::utsname = unsafe { std::mem::zeroed() };
if unsafe { libc::uname(&mut uts) } != 0 {
return (std::env::consts::OS.to_owned(), None, None);
}
fn field(raw: &[libc::c_char]) -> Option<String> {
let bytes: Vec<u8> = raw
.iter()
.take_while(|&&c| c != 0)
.map(|&c| c as u8)
.collect();
let s = String::from_utf8_lossy(&bytes).trim().to_owned();
(!s.is_empty()).then_some(s)
}
(
field(&uts.sysname).unwrap_or_else(|| std::env::consts::OS.to_owned()),
field(&uts.release),
field(&uts.machine),
)
}
#[cfg(not(unix))]
#[cfg(not(windows))]
fn uname_fields() -> (String, Option<String>, Option<String>) {
(std::env::consts::OS.to_owned(), None, None)
}
/// Path of the persistent device id: `{kigi_home}/device_id`.
fn device_id_path() -> PathBuf {
kigi_config::kigi_home().join("device_id")
}
/// Persistent uuid4-hex device id, created (owner-only, 0o600) on first use
/// and cached for the process lifetime.
pub(crate) fn device_id() -> anyhow::Result<String> {
static DEVICE_ID: OnceLock<String> = OnceLock::new();
if let Some(id) = DEVICE_ID.get() {
return Ok(id.clone());
}
let id = load_or_create_device_id(&device_id_path())?;
Ok(DEVICE_ID.get_or_init(|| id).clone())
}
/// Read `path`, or mint a uuid4 hex and persist it owner-only.
fn load_or_create_device_id(path: &std::path::Path) -> anyhow::Result<String> {
if let Ok(existing) = std::fs::read_to_string(path) {
let trimmed = existing.trim();
if !trimmed.is_empty() {
return Ok(trimmed.to_owned());
}
}
let id = uuid::Uuid::new_v4().simple().to_string();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("creating {} for device_id", parent.display()))?;
}
std::fs::write(path, &id)
.with_context(|| format!("writing device id to {}", path.display()))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
.with_context(|| format!("chmod 600 {}", path.display()))?;
}
tracing::info!(path = %path.display(), "auth: created persistent device id");
Ok(id)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ascii_header_value_passes_ascii_through() {
assert_eq!(ascii_header_value("macOS 15.5 arm64"), "macOS 15.5 arm64");
assert_eq!(ascii_header_value(" padded "), "padded");
}
#[test]
fn ascii_header_value_strips_non_ascii() {
assert_eq!(ascii_header_value("café-host"), "caf-host");
assert_eq!(ascii_header_value("机器"), "unknown");
assert_eq!(ascii_header_value(" "), "unknown");
}
#[test]
fn device_model_is_nonempty_ascii() {
let model = device_model();
assert!(!model.is_empty());
assert!(model.is_ascii(), "device model must be ASCII: {model:?}");
// The honest local OS name must lead the string.
#[cfg(target_os = "macos")]
assert!(model.starts_with("macOS "), "got {model:?}");
#[cfg(windows)]
assert!(model.starts_with("Windows"), "got {model:?}");
}
#[test]
fn load_or_create_device_id_roundtrips_and_is_owner_only() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("device_id");
let created = load_or_create_device_id(&path).unwrap();
assert_eq!(created.len(), 32, "uuid4 hex is 32 chars: {created:?}");
assert!(created.chars().all(|c| c.is_ascii_hexdigit()));
// Second call reads the same id back.
let reread = load_or_create_device_id(&path).unwrap();
assert_eq!(created, reread);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o777, 0o600, "device_id must be owner-only");
}
}
#[test]
fn load_or_create_device_id_ignores_empty_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("device_id");
std::fs::write(&path, " \n").unwrap();
let created = load_or_create_device_id(&path).unwrap();
assert_eq!(created.len(), 32);
}
#[cfg(target_os = "macos")]
#[test]
fn plist_string_value_extracts_product_version() {
let plist = r#"<?xml version="1.0"?>
<dict>
<key>ProductBuildVersion</key>
<string>24F74</string>
<key>ProductVersion</key>
<string>15.5</string>
</dict>"#;
assert_eq!(
plist_string_value(plist, "ProductVersion").as_deref(),
Some("15.5")
);
assert_eq!(plist_string_value(plist, "Missing"), None);
}
}
crates/codegen/kigi-shell/src/auth/device_code.rs
+282 -812
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/error.rs
+20 -45
View File
@@ -3,30 +3,21 @@ use thiserror::Error;
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum AuthError {
#[error("Not logged in. Run `grok login`.")]
#[error("Not logged in. Run `kigi login`.")]
NotLoggedIn,
/// Token expired and no refresh authority available.
#[error("Token expired. Run `grok login` to re-authenticate.")]
#[error("Token expired. Run `kigi 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.")]
#[error("Authentication rejected by server. Run `kigi 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)]
@@ -38,7 +29,7 @@ pub enum AuthError {
/// 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.
/// The credential was rejected; the tombstone cooldown gates re-attempts.
#[error(transparent)]
Permanent(#[from] RefreshTokenFailedError),
/// Network / 5xx / unknown blip; safe to retry later. Carries the cause.
@@ -48,12 +39,10 @@ pub enum RefreshTokenError {
/// 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.
/// [`AuthError::transient_source`], so a stray `?` on some error can't
/// silently classify a permanent failure as retryable. Display frames the
/// cause as an auth-refresh failure so internal messages (lock timeout,
/// sleep defer) don't surface bare.
#[derive(Debug, Error)]
#[error("auth refresh failed: {0}")]
pub struct RefreshTransientError(#[source] Box<dyn std::error::Error + Send + Sync>);
@@ -74,45 +63,31 @@ impl From<RefreshTokenFailedReason> for RefreshTokenFailedError {
}
}
/// Why a token refresh terminally failed, grounded in the OAuth2 error codes
/// our IdP actually emits.
/// Why a token refresh terminally failed. Both reasons carry the same
/// tombstone semantics (PRD F1): a 300s cooldown scoped to the rejected
/// refresh token, auto-cleared when the persisted refresh token differs
/// (another process rotated) or a fresh login lands.
#[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).
/// The OAuth host answered 401/403 — the refresh token is no longer
/// valid (expired, reused, or revoked).
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`).
/// Non-retryable terminal failure that isn't an explicit rejection
/// (malformed payload, unexpected 4xx).
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.
/// User-facing copy for a terminal refresh failure; the raw wire detail
/// 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."
"Your session has expired. Run `kigi login` to sign in again."
}
Self::Other => {
"Authentication could not be refreshed. Run `grok login` to sign in again."
"Authentication could not be refreshed. Run `kigi login` to sign in again."
}
}
}
crates/codegen/kigi-shell/src/auth/external_auth.rs
-283
View File
@@ -1,283 +0,0 @@
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
+163 -1595
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/jwt.rs
-45
View File
@@ -1,45 +0,0 @@
//! 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/kimi_oauth.rs
+626
View File
@@ -0,0 +1,626 @@
//! Kimi Code OAuth wire protocol (PRD F1).
//!
//! Three calls against `{host}` (= `kigi_env::oauth_host()`), all
//! `application/x-www-form-urlencoded` POSTs carrying the device-identity
//! headers from [`super::device`]:
//!
//! - `POST /api/oauth/device_authorization` — form `client_id`
//! - `POST /api/oauth/token` (poll) — form `client_id` + `device_code` +
//! `grant_type=urn:ietf:params:oauth:grant-type:device_code`
//! - `POST /api/oauth/token` (refresh) — form `client_id` +
//! `grant_type=refresh_token` + `refresh_token`, with exponential backoff
//! over the retryable statuses {429, 500, 502, 503, 504} (3 tries) and
//! 401/403 mapped to [`RefreshError::Unauthorized`].
//!
//! Ported from kimi-cli `auth/oauth.py` (the authoritative reference).
use chrono::{Duration, Utc};
use serde::Deserialize;
use super::device::device_headers;
use super::model::{AuthMode, KimiAuth};
/// Kimi Code OAuth client id (fixed for the official device-flow client).
pub(crate) const KIMI_CODE_CLIENT_ID: &str = "17e5f671-d194-4dfb-9706-5516cb48c098";
const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code";
const REFRESH_GRANT_TYPE: &str = "refresh_token";
/// Refresh retry budget over the retryable statuses / network blips.
const MAX_REFRESH_RETRIES: u32 = 3;
/// HTTP statuses worth retrying a refresh for (kimi-cli parity).
const RETRYABLE_REFRESH_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
/// Result of `POST /api/oauth/device_authorization`.
#[derive(Debug, Clone)]
pub struct DeviceAuthorization {
pub user_code: String,
pub device_code: String,
/// Bare verification page (may be absent; the complete URI is required).
pub verification_uri: Option<String>,
/// Verification page with the user code pre-filled — what we display
/// and open in the browser.
pub verification_uri_complete: String,
/// Device-code lifetime; `None` when the server omits it.
pub expires_in: Option<i64>,
/// Poll interval in seconds (server default 5; floored at 1 by callers).
pub interval: i64,
}
#[derive(Deserialize)]
struct DeviceAuthorizationResponse {
user_code: String,
device_code: String,
#[serde(default)]
verification_uri: Option<String>,
verification_uri_complete: String,
#[serde(default)]
expires_in: Option<i64>,
#[serde(default)]
interval: Option<i64>,
}
/// Successful token payload (device grant and refresh grant share it).
#[derive(Debug, Deserialize)]
pub(crate) struct TokenResponse {
pub access_token: String,
pub refresh_token: String,
pub expires_in: i64,
#[serde(default)]
pub scope: Option<String>,
#[serde(default)]
pub token_type: Option<String>,
}
impl TokenResponse {
/// Materialize the credential: `expires_at = now + expires_in`.
pub(crate) fn into_auth(self) -> KimiAuth {
let now = Utc::now();
KimiAuth {
key: self.access_token,
auth_mode: AuthMode::OAuth,
create_time: now,
user_id: String::new(),
email: None,
refresh_token: Some(self.refresh_token),
expires_at: Some(now + Duration::seconds(self.expires_in)),
expires_in: Some(self.expires_in),
scope: self.scope,
token_type: self.token_type,
}
}
}
#[derive(Deserialize, Default)]
struct OAuthErrorBody {
#[serde(default)]
error: Option<String>,
#[serde(default)]
error_description: Option<String>,
}
/// One poll tick against the token endpoint.
#[derive(Debug)]
pub(crate) enum DevicePollResult {
/// 200 with an access token — login complete.
Success(Box<KimiAuth>),
/// `error == "expired_token"` — restart the whole device authorization.
Expired,
/// Any other non-200 outcome (`authorization_pending`, `slow_down`,
/// unknown errors) — wait and poll again. `slow_down` additionally bumps
/// the caller's interval.
Pending {
error: String,
description: Option<String>,
},
}
fn oauth_url(host: &str, path: &str) -> String {
format!("{}{path}", host.trim_end_matches('/'))
}
/// Attach the device-identity headers to a request.
fn with_device_headers(
mut builder: reqwest::RequestBuilder,
) -> anyhow::Result<reqwest::RequestBuilder> {
for (name, value) in device_headers()? {
builder = builder.header(name, value);
}
Ok(builder)
}
/// Defend against control characters / non-https redirects from a
/// compromised or mis-configured OAuth host.
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"),
}
}
/// `POST {host}/api/oauth/device_authorization` — start a device login.
pub(crate) async fn request_device_authorization(
host: &str,
) -> anyhow::Result<DeviceAuthorization> {
let url = oauth_url(host, "/api/oauth/device_authorization");
tracing::info!(url = %url, "auth: requesting device authorization");
let resp = with_device_headers(crate::http::shared_client().post(&url))?
.form(&[("client_id", KIMI_CODE_CLIENT_ID)])
.send()
.await?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
tracing::warn!(%status, "auth: device authorization failed");
anyhow::bail!("Device authorization failed (HTTP {status}): {body}");
}
let parsed: DeviceAuthorizationResponse = resp.json().await?;
if !parsed
.user_code
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '-')
{
anyhow::bail!("Server returned invalid user_code format (expected [A-Z0-9-])");
}
validate_verification_uri(&parsed.verification_uri_complete)?;
if let Some(ref uri) = parsed.verification_uri {
validate_verification_uri(uri)?;
}
tracing::info!(
user_code = %parsed.user_code,
interval = parsed.interval.unwrap_or(5),
expires_in = ?parsed.expires_in,
"auth: device authorization issued"
);
Ok(DeviceAuthorization {
user_code: parsed.user_code,
device_code: parsed.device_code,
verification_uri: parsed.verification_uri.filter(|u| !u.is_empty()),
verification_uri_complete: parsed.verification_uri_complete,
expires_in: parsed.expires_in.filter(|&e| e > 0),
interval: parsed.interval.unwrap_or(5),
})
}
/// One poll of `POST {host}/api/oauth/token` with the device grant.
///
/// 5xx and network/decode failures are errors (kimi-cli parity: the login
/// loop surfaces them); everything else maps onto [`DevicePollResult`].
pub(crate) async fn poll_device_token(
host: &str,
device_code: &str,
) -> anyhow::Result<DevicePollResult> {
let url = oauth_url(host, "/api/oauth/token");
let resp = with_device_headers(crate::http::shared_client().post(&url))?
.form(&[
("client_id", KIMI_CODE_CLIENT_ID),
("device_code", device_code),
("grant_type", DEVICE_GRANT_TYPE),
])
.send()
.await
.map_err(|e| anyhow::anyhow!("Token polling request failed: {e}"))?;
let status = resp.status();
if status.is_server_error() {
anyhow::bail!("Token polling server error: {status}");
}
let body = resp.bytes().await?;
if status.is_success() {
if let Ok(tokens) = serde_json::from_slice::<TokenResponse>(&body) {
tracing::info!("auth: device poll succeeded, access token issued");
return Ok(DevicePollResult::Success(Box::new(tokens.into_auth())));
}
// 200 without an access token: treat as still-pending (kimi-cli
// requires "access_token" in the payload before accepting).
tracing::warn!("auth: device poll returned 200 without access_token; continuing");
return Ok(DevicePollResult::Pending {
error: "missing_access_token".to_owned(),
description: None,
});
}
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
let error = err.error.unwrap_or_else(|| "unknown_error".to_owned());
if error == "expired_token" {
tracing::info!("auth: device code expired; restarting device authorization");
return Ok(DevicePollResult::Expired);
}
tracing::debug!(error = %error, "auth: device poll pending");
Ok(DevicePollResult::Pending {
error,
description: err.error_description,
})
}
/// Why a refresh call terminally or transiently failed.
#[derive(Debug, thiserror::Error)]
pub(crate) enum RefreshError {
/// 401/403 — the refresh token was rejected. Triggers the tombstone
/// cooldown in the manager.
#[error("token refresh unauthorized (HTTP {status}): {description}")]
Unauthorized { status: u16, description: String },
/// Non-retryable non-200 status.
#[error("token refresh failed (HTTP {status}): {description}")]
Fatal { status: u16, description: String },
/// Retry budget exhausted over retryable statuses / network blips.
#[error("token refresh failed after {MAX_REFRESH_RETRIES} attempts: {last_error}")]
Exhausted { last_error: String },
/// Local failure before the wire (e.g. device-id creation failed).
#[error(transparent)]
Local(#[from] anyhow::Error),
}
/// `POST {host}/api/oauth/token` with `grant_type=refresh_token`.
///
/// Retries the retryable statuses and network errors with exponential
/// backoff (`2^attempt` seconds); 401/403 returns immediately as
/// [`RefreshError::Unauthorized`].
pub(crate) async fn refresh_token(
host: &str,
refresh_token: &str,
) -> Result<KimiAuth, RefreshError> {
let url = oauth_url(host, "/api/oauth/token");
let mut last_error = String::from("no attempt made");
for attempt in 0..MAX_REFRESH_RETRIES {
if attempt > 0 {
let backoff = std::time::Duration::from_secs(1 << (attempt - 1));
tracing::warn!(
attempt,
backoff_secs = backoff.as_secs(),
last_error = %last_error,
"auth: retrying token refresh"
);
tokio::time::sleep(backoff).await;
}
tracing::info!(attempt, "auth: token refresh attempt");
let send_result = with_device_headers(crate::http::shared_client().post(&url))?
.form(&[
("client_id", KIMI_CODE_CLIENT_ID),
("grant_type", REFRESH_GRANT_TYPE),
("refresh_token", refresh_token),
])
.send()
.await;
let resp = match send_result {
Ok(resp) => resp,
Err(e) => {
last_error = format!("network error: {e}");
continue;
}
};
let status = resp.status().as_u16();
let body = resp.bytes().await.unwrap_or_default();
if status == 401 || status == 403 {
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
return Err(RefreshError::Unauthorized {
status,
description: err
.error_description
.unwrap_or_else(|| "Token refresh unauthorized.".to_owned()),
});
}
if status == 200 {
return match serde_json::from_slice::<TokenResponse>(&body) {
Ok(tokens) => Ok(tokens.into_auth()),
Err(e) => Err(RefreshError::Fatal {
status,
description: format!("malformed token payload: {e}"),
}),
};
}
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
let description = err
.error_description
.unwrap_or_else(|| format!("Token refresh failed (HTTP {status})."));
if RETRYABLE_REFRESH_STATUSES.contains(&status) {
last_error = description;
continue;
}
return Err(RefreshError::Fatal {
status,
description,
});
}
Err(RefreshError::Exhausted { last_error })
}
#[cfg(test)]
mod tests {
use super::*;
use wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
fn token_json(access: &str, refresh: &str) -> serde_json::Value {
serde_json::json!({
"access_token": access,
"refresh_token": refresh,
"expires_in": 3600,
"scope": "kimi-code",
"token_type": "bearer",
})
}
#[tokio::test]
async fn device_authorization_parses_wire_payload() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.and(body_string_contains(format!(
"client_id={KIMI_CODE_CLIENT_ID}"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "ABCD-1234",
"device_code": "dev-code-1",
"verification_uri": "https://auth.kimi.com/device",
"verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234",
"expires_in": 600,
"interval": 7,
})))
.expect(1)
.mount(&server)
.await;
let auth = request_device_authorization(&server.uri()).await.unwrap();
assert_eq!(auth.user_code, "ABCD-1234");
assert_eq!(auth.device_code, "dev-code-1");
assert_eq!(auth.interval, 7);
assert_eq!(auth.expires_in, Some(600));
assert_eq!(
auth.verification_uri_complete,
"https://auth.kimi.com/device?code=ABCD-1234"
);
}
#[tokio::test]
async fn device_authorization_sends_device_headers() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.and(wiremock::matchers::header_exists("X-Msh-Device-Name"))
.and(wiremock::matchers::header_exists("X-Msh-Device-Model"))
.and(wiremock::matchers::header_exists("X-Msh-Device-Id"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "https://auth.kimi.com/device?code=AAAA",
"interval": 5,
})))
.expect(1)
.mount(&server)
.await;
request_device_authorization(&server.uri()).await.unwrap();
}
#[tokio::test]
async fn device_authorization_defaults_interval_to_five() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "https://auth.kimi.com/device?code=AAAA",
})))
.mount(&server)
.await;
let auth = request_device_authorization(&server.uri()).await.unwrap();
assert_eq!(auth.interval, 5);
assert_eq!(auth.expires_in, None);
}
#[tokio::test]
async fn device_authorization_rejects_bad_verification_uri() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "javascript:alert(1)",
})))
.mount(&server)
.await;
let err = request_device_authorization(&server.uri())
.await
.unwrap_err();
assert!(err.to_string().contains("verification URI"), "{err}");
}
#[tokio::test]
async fn device_authorization_surfaces_http_errors() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.respond_with(ResponseTemplate::new(400).set_body_string("nope"))
.mount(&server)
.await;
let err = request_device_authorization(&server.uri())
.await
.unwrap_err();
assert!(err.to_string().contains("HTTP 400"), "{err}");
}
#[tokio::test]
async fn poll_success_builds_auth_with_expiry() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.and(body_string_contains("grant_type=urn"))
.and(body_string_contains("device_code=dev-1"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-1", "rt-1")))
.expect(1)
.mount(&server)
.await;
let result = poll_device_token(&server.uri(), "dev-1").await.unwrap();
let DevicePollResult::Success(auth) = result else {
panic!("expected success, got {result:?}");
};
assert_eq!(auth.key, "at-1");
assert_eq!(auth.refresh_token.as_deref(), Some("rt-1"));
assert_eq!(auth.expires_in, Some(3600));
let remaining = auth.expires_at.unwrap() - chrono::Utc::now();
assert!(
(3590..=3600).contains(&remaining.num_seconds()),
"expires_at must be ~now+expires_in, got {remaining:?}"
);
assert_eq!(auth.scope.as_deref(), Some("kimi-code"));
assert_eq!(auth.token_type.as_deref(), Some("bearer"));
}
#[tokio::test]
async fn poll_maps_expired_token_to_restart() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(400)
.set_body_json(serde_json::json!({ "error": "expired_token" })),
)
.mount(&server)
.await;
let result = poll_device_token(&server.uri(), "dev-1").await.unwrap();
assert!(matches!(result, DevicePollResult::Expired), "{result:?}");
}
#[tokio::test]
async fn poll_maps_pending_and_unknown_errors_to_pending() {
let server = MockServer::start().await;
for error in ["authorization_pending", "slow_down", "surprise_error"] {
server.reset().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(400).set_body_json(serde_json::json!({ "error": error })),
)
.mount(&server)
.await;
let result = poll_device_token(&server.uri(), "dev-1").await.unwrap();
match result {
DevicePollResult::Pending { error: got, .. } => assert_eq!(got, error),
other => panic!("expected pending for {error}, got {other:?}"),
}
}
}
#[tokio::test]
async fn poll_server_error_is_an_error() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(502))
.mount(&server)
.await;
let err = poll_device_token(&server.uri(), "dev-1").await.unwrap_err();
assert!(err.to_string().contains("server error"), "{err}");
}
#[tokio::test]
async fn refresh_success_round_trip() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.and(body_string_contains("grant_type=refresh_token"))
.and(body_string_contains("refresh_token=rt-old"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-new", "rt-new")))
.expect(1)
.mount(&server)
.await;
let auth = refresh_token(&server.uri(), "rt-old").await.unwrap();
assert_eq!(auth.key, "at-new");
assert_eq!(auth.refresh_token.as_deref(), Some("rt-new"));
}
#[tokio::test]
async fn refresh_401_maps_to_unauthorized_without_retry() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(401).set_body_json(
serde_json::json!({ "error_description": "refresh token revoked" }),
),
)
.expect(1)
.mount(&server)
.await;
let err = refresh_token(&server.uri(), "rt-dead").await.unwrap_err();
match err {
RefreshError::Unauthorized {
status,
description,
} => {
assert_eq!(status, 401);
assert_eq!(description, "refresh token revoked");
}
other => panic!("expected Unauthorized, got {other:?}"),
}
}
#[tokio::test]
async fn refresh_retries_retryable_status_then_succeeds() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(503))
.up_to_n_times(1)
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-2", "rt-2")))
.expect(1)
.mount(&server)
.await;
let auth = refresh_token(&server.uri(), "rt-old").await.unwrap();
assert_eq!(auth.key, "at-2");
}
#[tokio::test]
async fn refresh_exhausts_after_three_retryable_failures() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(500))
.expect(3)
.mount(&server)
.await;
let err = refresh_token(&server.uri(), "rt-old").await.unwrap_err();
assert!(matches!(err, RefreshError::Exhausted { .. }), "{err:?}");
}
#[tokio::test]
async fn refresh_non_retryable_status_is_fatal_without_retry() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(400)
.set_body_json(serde_json::json!({ "error_description": "bad request" })),
)
.expect(1)
.mount(&server)
.await;
let err = refresh_token(&server.uri(), "rt-old").await.unwrap_err();
match err {
RefreshError::Fatal {
status,
description,
} => {
assert_eq!(status, 400);
assert_eq!(description, "bad request");
}
other => panic!("expected Fatal, got {other:?}"),
}
}
}
crates/codegen/kigi-shell/src/auth/manager.rs
+383 -710
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/manager/enrichment.rs
-256
View File
@@ -1,256 +0,0 @@
//! 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_tests.rs
+463 -3802
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/meta.rs
+3 -19
View File
@@ -1,6 +1,8 @@
use serde::{Deserialize, Serialize};
/// Access gate from `grok_build_access_gate`.
/// Access-gate copy resolved from remote settings (message + optional CTA).
/// Auth no longer produces gates (tier gating was an xAI concept); the pager
/// still renders one when remote settings carry a gate message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GateInfo {
pub message: String,
@@ -17,24 +19,6 @@ pub struct AuthMeta {
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
+9 -21
View File
@@ -1,42 +1,30 @@
pub(crate) mod attribution;
mod config;
pub mod credential_provider;
#[path = "devbox_login_stub.rs"]
pub(crate) mod devbox_login;
pub(crate) mod device;
pub mod device_code;
pub mod error;
mod external_auth;
mod flow;
mod jwt;
pub(crate) mod kimi_oauth;
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 config::{KIMI_CODE_OAUTH_SCOPE, KimiCodeConfig};
pub(crate) use flow::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,
AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated,
ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow,
run_auth_flow_with_stderr_bridge, 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 model::{AuthMode, KimiAuth, lookup_auth};
pub(crate) use model::{TOKEN_TTL, 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
+171 -364
View File
@@ -1,108 +1,75 @@
//! Kimi Code auth data model: the persisted token set + expiry policy.
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use super::is_xai_oauth2_issuer;
/// Fallback TTL for credentials without a server-provided expiry
/// (plain API keys).
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";
/// Minimum refresh threshold (PRD F1): refresh when the remaining lifetime
/// drops below `max(300, expires_in × 0.5)` seconds.
const DEFAULT_EARLY_INVALIDATION_SECS: u64 = 300;
/// auth.json scope key for plain API key auth (desktop login, `grok login --api-key`).
pub const API_KEY_SCOPE: &str = "xai::api_key";
/// Fraction of `expires_in` that drives the dynamic refresh threshold.
const REFRESH_THRESHOLD_RATIO: f64 = 0.5;
const BLOCKED_REASON_NO_LOGS: &str = "BLOCKED_REASON_NO_LOGS";
const BLOCKED_REASON_NO_LOGS_MODERATED: &str = "BLOCKED_REASON_NO_LOGS_MODERATED";
/// auth.json scope key for plain API key auth (`kigi login --api-key`, F2).
pub const API_KEY_SCOPE: &str = "kigi::api_key";
/// Token provenance (debugging/auth.json only -- no code branches on this).
/// How this credential was obtained.
#[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`)
/// Kimi Code subscription OAuth (device-code flow).
#[serde(rename = "oauth")]
OAuth,
/// Plain 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";
/// The Kimi Code credential: the OAuth token set (or a bare API key) plus
/// local bookkeeping. The Kimi token response carries no user info; `user_id`
/// / `email` stay empty until a later feature surfaces account info.
#[derive(Clone, Serialize, Deserialize)]
pub struct GrokAuth {
pub struct KimiAuth {
/// The bearer sent on API calls (`Authorization: Bearer {key}`):
/// the OAuth access token, or the API key in `ApiKey` mode.
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>,
/// Account id — the Kimi token response has none; empty until a later
/// feature surfaces it.
#[serde(default)]
pub coding_data_retention_opt_out: bool,
/// Deprecated. Kept for deserializing existing auth.json files.
pub user_id: String,
/// Account email — the Kimi token response has none; `None` until a
/// later feature surfaces it.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub has_grok_code_access: Option<bool>,
/// Refresh token (OIDC/OAuth2 or external provider).
pub email: Option<String>,
/// OAuth refresh token; `None` for API keys.
#[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`.
/// `create_time + expires_in`, computed when the token was minted.
#[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`).
/// Server-reported token lifetime in seconds; drives the dynamic
/// refresh threshold `max(300, expires_in × 0.5)`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc_issuer: Option<String>,
/// OIDC client_id used to obtain this token (needed for refresh).
pub expires_in: Option<i64>,
/// OAuth scope string as returned by the token endpoint.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc_client_id: Option<String>,
pub scope: Option<String>,
/// Token type as returned by the token endpoint (e.g. "bearer").
#[serde(default, skip_serializing_if = "Option::is_none")]
pub token_type: Option<String>,
}
impl std::fmt::Debug for GrokAuth {
impl std::fmt::Debug for KimiAuth {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GrokAuth")
f.debug_struct("KimiAuth")
.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",
@@ -112,128 +79,44 @@ impl std::fmt::Debug for GrokAuth {
}
}
impl GrokAuth {
impl KimiAuth {
/// 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.
/// clock stepped back past `create_time` (NTP correction, VM restore).
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.
/// `true` for a refreshable subscription session (vs a bare API key).
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;
self.auth_mode == AuthMode::OAuth
}
}
impl Default for GrokAuth {
impl Default for KimiAuth {
fn default() -> Self {
Self {
key: String::new(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
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,
expires_in: None,
scope: None,
token_type: None,
}
}
}
#[cfg(test)]
impl GrokAuth {
/// Returns a `GrokAuth` with sensible defaults for tests. Override fields
/// with struct update syntax:
impl KimiAuth {
/// A `KimiAuth` with sensible defaults for tests. Override fields with
/// struct update syntax:
/// ```ignore
/// GrokAuth { key: "my-key".into(), ..GrokAuth::test_default() }
/// KimiAuth { key: "my-key".into(), ..KimiAuth::test_default() }
/// ```
pub fn test_default() -> Self {
Self {
@@ -244,82 +127,24 @@ impl GrokAuth {
}
}
pub(crate) type AuthStore = BTreeMap<String, GrokAuth>;
pub(crate) type AuthStore = BTreeMap<String, KimiAuth>;
/// 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.
/// Last 12 chars of a token string, safe for diagnostic logging. Uses the
/// tail because token prefixes are shared across a family; the tail is
/// unique per token and makes `key_changed` 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)
pub fn lookup_auth(map: &AuthStore, scope: &str) -> Option<KimiAuth> {
map.get(scope).cloned()
}
/// Early-invalidation buffer. Override with `KIGI_AUTH_EARLY_INVALIDATION_SECS`
/// for testing (e.g. `=5` to shrink the buffer to 5 seconds).
/// Minimum refresh-threshold component. 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()
@@ -328,14 +153,30 @@ pub(super) fn early_invalidation() -> Duration {
.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())
/// Dynamic refresh threshold (PRD F1): `max(min_threshold, expires_in × 0.5)`
/// where `min_threshold` defaults to 300s. Credentials without a positive
/// `expires_in` use the minimum alone.
pub(crate) fn refresh_threshold(auth: &KimiAuth) -> Duration {
let min = early_invalidation();
match auth.expires_in {
Some(expires_in) if expires_in > 0 => {
let ratio = Duration::seconds((expires_in as f64 * REFRESH_THRESHOLD_RATIO) as i64);
std::cmp::max(min, ratio)
}
_ => min,
}
}
/// Whether the credential is inside its refresh threshold (i.e. should be
/// treated as expiring-soon for refresh scheduling).
pub(crate) fn is_expired(auth: &KimiAuth) -> bool {
is_expired_with_buffer(auth, refresh_threshold(auth))
}
/// 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 {
/// really be rejected on the wire.
pub(crate) fn is_expired_with_buffer(auth: &KimiAuth, buffer: Duration) -> bool {
if let Some(expires_at) = auth.expires_at {
Utc::now() >= (expires_at - buffer)
} else {
@@ -348,142 +189,108 @@ pub(crate) fn is_expired_with_buffer(auth: &GrokAuth, buffer: Duration) -> bool
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,
fn auth_with_lifetime(expires_in: i64, remaining_secs: i64) -> KimiAuth {
KimiAuth {
expires_in: Some(expires_in),
expires_at: Some(Utc::now() + Duration::seconds(remaining_secs)),
refresh_token: Some("rt".into()),
..KimiAuth::test_default()
}
}
/// PRD threshold math: `max(300, expires_in × 0.5)`.
#[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)
fn refresh_threshold_is_max_of_min_and_half_life() {
// Short-lived token: the 300s floor wins (600 × 0.5 = 300 → tie; 400 × 0.5 = 200 < 300).
let short = auth_with_lifetime(400, 400);
assert_eq!(refresh_threshold(&short).num_seconds(), 300);
// Long-lived token: half the lifetime wins (7200 × 0.5 = 3600).
let long = auth_with_lifetime(7200, 7200);
assert_eq!(refresh_threshold(&long).num_seconds(), 3600);
// No expires_in: the floor alone.
let bare = KimiAuth::test_default();
assert_eq!(refresh_threshold(&bare).num_seconds(), 300);
// Non-positive expires_in must not produce a negative threshold.
let broken = KimiAuth {
expires_in: Some(-5),
..KimiAuth::test_default()
};
assert_eq!(refresh_threshold(&broken).num_seconds(), 300);
}
// 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());
/// A token past its dynamic threshold counts as expiring-soon while a
/// token comfortably before it does not.
#[test]
fn is_expired_uses_dynamic_threshold() {
// 7200s lifetime → threshold 3600s. 3000s remaining < 3600 → expiring.
assert!(is_expired(&auth_with_lifetime(7200, 3000)));
// 5000s remaining > 3600 → fresh.
assert!(!is_expired(&auth_with_lifetime(7200, 5000)));
// Hard expiry ignores the buffer entirely.
assert!(!is_expired_with_buffer(
&auth_with_lifetime(7200, 3000),
Duration::zero()
));
assert!(is_expired_with_buffer(
&auth_with_lifetime(7200, -1),
Duration::zero()
));
}
// 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());
/// Credentials without `expires_at` (API keys) age out via the 30-day TTL.
#[test]
fn no_expiry_falls_back_to_token_ttl() {
let fresh = KimiAuth::test_default();
assert!(!is_expired(&fresh));
let old = KimiAuth {
create_time: Utc::now() - Duration::days(31),
..KimiAuth::test_default()
};
assert!(is_expired(&old));
}
#[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)
fn lookup_auth_finds_scope_entry() {
let mut map = AuthStore::new();
map.insert("oauth/kimi-code".into(), KimiAuth::test_default());
assert!(lookup_auth(&map, "oauth/kimi-code").is_some());
assert!(lookup_auth(&map, "other").is_none());
}
#[test]
fn debug_redacts_tokens() {
let auth = KimiAuth {
key: "super-secret-access-token".into(),
refresh_token: Some("super-secret-refresh-token".into()),
..KimiAuth::test_default()
};
let debug = format!("{auth:?}");
assert!(!debug.contains("super-secret-access-token"));
assert!(!debug.contains("super-secret-refresh-token"));
}
// 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());
#[test]
fn serde_roundtrip_preserves_token_set() {
let auth = KimiAuth {
key: "at".into(),
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now()),
expires_in: Some(3600),
scope: Some("kimi-code".into()),
token_type: Some("bearer".into()),
..KimiAuth::test_default()
};
let json = serde_json::to_string(&auth).unwrap();
assert!(
!with_issuer(AuthMode::External, Some("https://idp.acme.example")).is_session_auth()
json.contains("\"oauth\""),
"wire spelling is \"oauth\": {json}"
);
// 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(""));
let back: KimiAuth = serde_json::from_str(&json).unwrap();
assert_eq!(back.key, "at");
assert_eq!(back.refresh_token.as_deref(), Some("rt"));
assert_eq!(back.expires_in, Some(3600));
assert_eq!(back.scope.as_deref(), Some("kimi-code"));
assert_eq!(back.token_type.as_deref(), Some("bearer"));
assert_eq!(back.auth_mode, AuthMode::OAuth);
}
}
crates/codegen/kigi-shell/src/auth/oidc/login.rs
-701
View File
@@ -1,701 +0,0 @@
//! 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
@@ -1,14 +0,0 @@
//! 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
@@ -1,249 +0,0 @@
//! 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
@@ -1,130 +0,0 @@
//! 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
+123 -311
View File
@@ -3,71 +3,54 @@
//! 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.
//! 1. **ReloadFromDisk** — re-read the persisted credential under a file
//! lock; if it differs from the rejected one, accept it (another process
//! may have refreshed).
//! 2. **RefreshFromAuthority** — run the refresh chain against the Kimi
//! OAuth host, unless the live token was minted moments ago (fresh-mint
//! guard).
//! 3. **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::model::KimiAuth;
use crate::auth::token_type::TokenType;
/// Whether a terminal `AuthError` forces a manual re-login (`None` cases
/// Whether a terminal `AuthError` forces a manual re-login (`false` 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,
// Self-healing via the tombstone cooldown, not a manual re-auth.
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,
| AuthError::TokenExpiredNoRefresh => true,
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.
/// 401s land ~20s after mint), while the refresh-threshold 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).
/// Re-read the persisted credential (file-locked).
ReloadFromDisk,
/// Refresh via the authority (OIDC, external binary, etc.).
/// Refresh via the Kimi OAuth host.
RefreshFromAuthority,
/// On devboxes: purge auth.json and mint fresh OIDC credentials.
DevboxRecovery,
/// All strategies exhausted.
Done,
}
@@ -79,8 +62,7 @@ pub struct UnauthorizedRecovery {
rejected_token: String,
/// Current step in the recovery sequence.
step: RecoveryStep,
/// Error from `RefreshFromAuthority`, propagated as fallback when
/// devbox recovery doesn't apply.
/// Error from `RefreshFromAuthority`, propagated on exhaustion.
authority_error: Option<AuthError>,
/// Whether the last authority failure was transient. Kept past the
/// `authority_error` handoff so exhaustion preserves the
@@ -90,7 +72,7 @@ pub struct UnauthorizedRecovery {
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 {
pub(crate) fn new(auth_manager: Arc<AuthManager>, rejected: Option<KimiAuth>) -> Self {
let rejected_token = rejected.as_ref().map(|a| a.key.clone()).unwrap_or_default();
Self {
auth_manager,
@@ -102,14 +84,14 @@ impl UnauthorizedRecovery {
}
/// Attempt the next recovery step. Walks
/// `ReloadFromDisk -> RefreshFromAuthority -> DevboxRecovery -> Done`.
/// `ReloadFromDisk -> RefreshFromAuthority -> 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> {
pub async fn next(&mut self) -> Result<KimiAuth, AuthError> {
let span = tracing::Span::current();
if !span.is_disabled() {
// Only acquire the inner-lock when tracing actually
@@ -122,23 +104,10 @@ impl UnauthorizedRecovery {
tracing::field::debug(self.auth_manager.token_type()),
);
}
self.resolve_next().await
self.next_step_loop().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> {
async fn next_step_loop(&mut self) -> Result<KimiAuth, AuthError> {
loop {
match self.step {
RecoveryStep::ReloadFromDisk => {
@@ -148,30 +117,20 @@ impl UnauthorizedRecovery {
}
}
RecoveryStep::RefreshFromAuthority => {
self.step = RecoveryStep::DevboxRecovery;
self.step = RecoveryStep::Done;
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);
return Err(self
.authority_error
.take()
.unwrap_or(AuthError::RecoveryExhausted));
}
}
}
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
@@ -187,9 +146,9 @@ impl UnauthorizedRecovery {
}
}
/// Re-read `auth.json` from disk. Accept the token only if it differs
/// Re-read the persisted credential. Accept the token only if it differs
/// from the one that was rejected.
async fn try_reload_from_disk(&self) -> Option<GrokAuth> {
async fn try_reload_from_disk(&self) -> Option<KimiAuth> {
let _lock = self
.auth_manager
.try_lock_auth_file_async(crate::auth::manager::AUTH_LOCK_TIMEOUT)
@@ -203,13 +162,13 @@ impl UnauthorizedRecovery {
// 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);
kigi_log::unified_log::debug("auth recovery: no persisted entry", None, None);
return None;
};
if crate::auth::is_expired(&disk_auth) {
tracing::debug!("auth recovery: disk token is expired, skipping");
tracing::debug!("auth recovery: persisted token is expired, skipping");
kigi_log::unified_log::debug(
"auth recovery: disk token expired",
"auth recovery: persisted token expired",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&disk_auth.key),
@@ -219,9 +178,9 @@ impl UnauthorizedRecovery {
return None;
}
if self.is_different_token(&disk_auth) {
tracing::info!("auth recovery: disk has a different token, accepting");
tracing::info!("auth recovery: persisted store has a different token, accepting");
kigi_log::unified_log::info(
"auth recovery: adopted disk token",
"auth recovery: adopted persisted token",
None,
Some(serde_json::json!({
"adopted_key_prefix": crate::auth::token_suffix(&disk_auth.key),
@@ -231,8 +190,12 @@ impl UnauthorizedRecovery {
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);
tracing::debug!("auth recovery: persisted token is same as rejected, skipping");
kigi_log::unified_log::debug(
"auth recovery: persisted token same as rejected",
None,
None,
);
None
}
}
@@ -242,14 +205,12 @@ impl UnauthorizedRecovery {
/// 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> {
/// the previous key) or validation lag on the new key — re-minting fixes
/// neither, and a crash between the token 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.
fn fresh_mint_guard(&self) -> Option<KimiAuth> {
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) {
@@ -272,14 +233,14 @@ impl UnauthorizedRecovery {
Some(auth)
}
/// Dispatch to the correct refresh chain based on the current `TokenType`.
/// Dispatch to the 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
/// - **OAuthSession**: full refresh chain via the OAuth host, unless the
/// live token is inside the fresh-mint guard window
/// ([`Self::fresh_mint_guard`]).
/// - **SessionNoRefresh / 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*
@@ -289,10 +250,10 @@ impl UnauthorizedRecovery {
/// 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> {
async fn try_refresh_from_authority(&self) -> Result<KimiAuth, AuthError> {
let tt = self.auth_manager.token_type();
match tt {
TokenType::OidcSession | TokenType::ExternalBinary => {
TokenType::OAuthSession => {
if let Some(auth) = self.fresh_mint_guard() {
return Ok(auth);
}
@@ -325,7 +286,7 @@ impl UnauthorizedRecovery {
}
result
}
TokenType::LegacySession | TokenType::ApiKey => {
TokenType::SessionNoRefresh | TokenType::ApiKey => {
kigi_log::unified_log::warn(
"auth recovery: no refresh authority for token type",
None,
@@ -338,7 +299,7 @@ impl UnauthorizedRecovery {
}
/// Check if a candidate token is different from the rejected one.
fn is_different_token(&self, candidate: &GrokAuth) -> bool {
fn is_different_token(&self, candidate: &KimiAuth) -> bool {
candidate.key != self.rejected_token
}
}
@@ -348,29 +309,28 @@ mod tests {
//! State-machine matrix tests for `UnauthorizedRecovery`.
//!
//! Coverage targets:
//! - All 5 `TokenType` variants x dispatch in `try_refresh_from_authority`.
//! - All 4 `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).
//! - Fresh-mint guard: ±window bounds, tombstone grace.
//!
//! 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::config::KimiCodeConfig;
use crate::auth::error::RefreshTokenError;
use crate::auth::model::{AuthMode, KimiAuth};
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 {
fn rejected_cred() -> Option<KimiAuth> {
Some(KimiAuth {
key: "rejected-tok".into(),
..GrokAuth::test_default()
..KimiAuth::test_default()
})
}
@@ -382,17 +342,17 @@ mod tests {
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 {
RefreshOutcome::Success(Box::new(KimiAuth {
key: "fresh-from-authority".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
}))
}
}
/// Refresher fake: returns PermanentFailure (invalid_grant).
/// Refresher fake: returns PermanentFailure (rejected refresh token).
struct FailRefresher {
calls: Arc<AtomicU32>,
}
@@ -406,19 +366,19 @@ mod tests {
fn mgr() -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let m = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let m = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
(dir, m)
}
fn seed(mgr: &AuthManager, mode: AuthMode, refresh_token: Option<&str>) {
let auth = GrokAuth {
let auth = KimiAuth {
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()
..KimiAuth::test_default()
};
mgr.hot_swap(auth);
}
@@ -426,9 +386,9 @@ mod tests {
// -- TokenType dispatch matrix ----------------------------------------
#[tokio::test]
async fn dispatch_oidc_session_uses_refresh_chain() {
async fn dispatch_oauth_session_uses_refresh_chain() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
@@ -441,39 +401,25 @@ mod tests {
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 {
fn seed_valid(mgr: &AuthManager, mint_age: Duration) {
mgr.hot_swap(KimiAuth {
key: "rejected-tok".into(),
auth_mode: mode,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
create_time: Utc::now() - mint_age,
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
expires_in: Some(3600),
..KimiAuth::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) {
async fn recover_with_ok_refresher(m: &Arc<AuthManager>) -> (Result<KimiAuth, AuthError>, u32) {
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
@@ -484,9 +430,9 @@ mod tests {
}
#[tokio::test]
async fn fresh_mint_guard_skips_idp_for_freshly_minted_token() {
async fn fresh_mint_guard_skips_wire_for_freshly_minted_token() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(10));
seed_valid(&m, Duration::seconds(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
@@ -495,23 +441,11 @@ mod tests {
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));
seed_valid(&m, Duration::seconds(-60));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
@@ -525,19 +459,19 @@ mod tests {
// 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));
seed_valid(&m, 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");
assert_eq!(calls, 1, "far-negative mint age must reach the wire");
}
#[tokio::test]
async fn fresh_mint_guard_lets_old_token_refresh() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::minutes(10));
seed_valid(&m, Duration::minutes(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("recovery should succeed").key,
@@ -545,25 +479,25 @@ mod tests {
);
assert_eq!(
calls, 1,
"outside the guard window ServerRejected must reach the IdP"
"outside the guard window ServerRejected must reach the wire"
);
}
#[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.
async fn fresh_mint_guard_wins_over_cached_tombstone() {
// A fresh *valid* token is served even when a tombstone is cached
// for its refresh token — mirrors `auth()`'s wire-valid grace arm;
// the tombstone re-applies once the guard window passes.
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(10));
seed_valid(&m, Duration::seconds(10));
m.record_permanent_failure(
"rejected-tok".into(),
"rt".into(),
RefreshTokenFailedReason::RefreshTokenRejected.into(),
);
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result
.expect("guard precedes the verdict short-circuit")
.expect("guard precedes the tombstone short-circuit")
.key,
"rejected-tok"
);
@@ -571,65 +505,10 @@ mod tests {
}
#[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() {
async fn dispatch_session_without_refresh_token_returns_server_rejected_no_recovery() {
// OAuth without refresh_token classifies as SessionNoRefresh.
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);
seed(&m, AuthMode::OAuth, None);
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
@@ -668,16 +547,17 @@ mod tests {
#[tokio::test]
async fn reload_from_disk_picks_up_different_token() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
// Sibling process wrote a different valid token to disk.
let scope = m.grok_com_config().auth_scope();
let fresh = GrokAuth {
let scope = m.kimi_code_config().auth_scope();
let fresh = KimiAuth {
key: "fresh-from-disk".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
expires_in: Some(3600),
..KimiAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, fresh);
@@ -694,16 +574,17 @@ mod tests {
#[tokio::test]
async fn reload_from_disk_skips_same_token_then_proceeds_to_authority() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
// Disk has the SAME token that was rejected -- skip, fall through.
let scope = m.grok_com_config().auth_scope();
let same = GrokAuth {
let scope = m.kimi_code_config().auth_scope();
let same = KimiAuth {
key: "rejected-tok".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
expires_in: Some(3600),
..KimiAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, same);
@@ -732,15 +613,11 @@ mod tests {
#[tokio::test]
async fn next_after_done_returns_recovery_exhausted() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, 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 {
@@ -773,9 +650,8 @@ mod tests {
}
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, 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.
@@ -799,21 +675,17 @@ mod tests {
!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) ------------
// -- Tombstone short-circuit (cross-check) ------------
#[tokio::test]
async fn refresh_authority_short_circuits_on_cached_permanent_failure() {
async fn refresh_authority_short_circuits_on_cached_tombstone() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
// Pre-record a permanent failure scoped to the seeded credential.
seed(&m, AuthMode::OAuth, Some("rt"));
// Pre-record a tombstone scoped to the seeded refresh token.
m.record_permanent_failure(
"rejected-tok".into(),
"rt".into(),
RefreshTokenFailedReason::RefreshTokenRejected.into(),
);
@@ -831,7 +703,7 @@ mod tests {
assert_eq!(
calls.load(Ordering::SeqCst),
0,
"refresher must not be invoked when permanent_failure is cached",
"refresher must not be invoked while the tombstone cooldown is live",
);
}
@@ -843,15 +715,15 @@ mod tests {
#[tokio::test]
async fn reload_from_disk_rejects_expired_different_token() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
let scope = m.grok_com_config().auth_scope();
let expired_different = GrokAuth {
let scope = m.kimi_code_config().auth_scope();
let expired_different = KimiAuth {
key: "different-but-expired".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, expired_different);
@@ -870,64 +742,4 @@ mod tests {
);
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
@@ -1,351 +0,0 @@
//! 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
@@ -1,176 +0,0 @@
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/kimi_refresher.rs
+356
View File
@@ -0,0 +1,356 @@
//! Kimi Code token refresher: drives `POST /api/oauth/token` with
//! `grant_type=refresh_token` through the [`TokenRefresher`] seam.
//!
//! Ports kimi-cli `OAuthManager._refresh_tokens`' sibling-safety behavior:
//! the persisted credential is re-read before the wire call (adopt a
//! rotation instead of refreshing), and after a 401/403 the persisted
//! credential is re-read once more (with a 1s grace) so a concurrent
//! process's freshly rotated token is adopted instead of tombstoning it.
use std::sync::Arc;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::kimi_oauth::{self, RefreshError};
use crate::auth::manager::RefreshReason;
use super::{AuthSnapshot, RefreshOutcome, TokenRefresher};
/// Grace period after a 401/403 before concluding the refresh token is dead:
/// a concurrent instance may still be persisting its rotated token
/// (kimi-cli parity: `await asyncio.sleep(1)`).
const POST_UNAUTHORIZED_GRACE: std::time::Duration = std::time::Duration::from_secs(1);
pub(crate) struct KimiRefresher {
auth: Arc<dyn AuthSnapshot>,
/// OAuth host; `kigi_env::oauth_host()` in production, injectable for
/// wiremock tests.
host: String,
}
impl KimiRefresher {
pub(crate) fn new(auth: Arc<dyn AuthSnapshot>, host: String) -> Self {
Self { auth, host }
}
/// Post-401 sibling check (kimi-cli parity): wait a beat, re-read the
/// persisted credential, and adopt it when its refresh token differs
/// from the one the server just rejected.
async fn adopt_rotation_after_unauthorized(&self, tried_rt: &str) -> Option<RefreshOutcome> {
tokio::time::sleep(POST_UNAUTHORIZED_GRACE).await;
let latest = self.auth.read_disk_auth()?;
let latest_rt = latest.refresh_token.as_deref()?;
if latest_rt == tried_rt {
return None;
}
kigi_log::unified_log::info(
"auth.refresh.adopted_rotation_after_401",
None,
Some(serde_json::json!({
"adopted_rt_prefix": crate::auth::token_suffix(latest_rt),
"rejected_rt_prefix": crate::auth::token_suffix(tried_rt),
})),
);
Some(RefreshOutcome::success(latest))
}
}
#[async_trait::async_trait]
impl TokenRefresher for KimiRefresher {
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
tracing::info!(?reason, "auth: kimi refresh attempt starting");
let disk_auth = self.auth.read_disk_auth();
// Sibling short-circuit: a valid persisted token whose key differs
// from in-memory means another process refreshed between the
// refresh_chain disk check (under lock) and here. Adopt directly.
if let Some(ref d) = disk_auth
&& !crate::auth::is_expired(d)
&& self.auth.current().map(|a| a.key).as_deref() != Some(&d.key)
{
kigi_log::unified_log::info(
"auth.refresh.adopted_sibling_token",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&d.key),
})),
);
return RefreshOutcome::success(d.clone());
}
let Some(auth) = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason)
else {
tracing::warn!(?reason, "auth: no credential available for refresh");
return RefreshOutcome::transient("no token with refresh_token available");
};
let Some(refresh_token) = auth.refresh_token.clone() else {
tracing::warn!(?reason, "auth: resolved credential has no refresh token");
return RefreshOutcome::transient("credential has no refresh token");
};
tracing::info!(
rt_prefix = crate::auth::token_suffix(&refresh_token),
expires_at = ?auth.expires_at,
"auth: sending refresh_token grant"
);
match kimi_oauth::refresh_token(&self.host, &refresh_token).await {
Ok(new_auth) => {
kigi_log::unified_log::info(
"auth.refresh.token_rotated",
None,
Some(serde_json::json!({
"new_key_prefix": crate::auth::token_suffix(&new_auth.key),
"expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
RefreshOutcome::success(new_auth)
}
Err(RefreshError::Unauthorized {
status,
description,
}) => {
tracing::warn!(status, %description, "auth: refresh token rejected");
if let Some(adopted) = self.adopt_rotation_after_unauthorized(&refresh_token).await
{
return adopted;
}
kigi_log::unified_log::warn(
"auth.refresh.unauthorized",
None,
Some(serde_json::json!({
"status": status,
"description": description,
"rt_prefix": crate::auth::token_suffix(&refresh_token),
})),
);
RefreshOutcome::permanent(
RefreshTokenFailedReason::RefreshTokenRejected,
Some(refresh_token),
)
}
// kimi-cli parity: non-401 failures never tombstone; the next
// 60s tick (or pre-request check) retries.
Err(
e @ (RefreshError::Exhausted { .. }
| RefreshError::Fatal { .. }
| RefreshError::Local(_)),
) => {
tracing::warn!(error = %e, "auth: refresh attempt failed (transient)");
kigi_log::unified_log::warn(
"auth.refresh.transient_wire_failure",
None,
Some(serde_json::json!({ "error": format!("{e}") })),
);
RefreshOutcome::transient(format!("token refresh failed: {e}"))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::model::KimiAuth;
use chrono::{Duration, Utc};
use parking_lot::Mutex;
use wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
/// Scriptable snapshot: `disk` can be swapped mid-test to simulate a
/// sibling process rotating the persisted credential.
struct FakeSnapshot {
current: Mutex<Option<KimiAuth>>,
disk: Mutex<Option<KimiAuth>>,
}
impl FakeSnapshot {
fn new(current: Option<KimiAuth>, disk: Option<KimiAuth>) -> Arc<Self> {
Arc::new(Self {
current: Mutex::new(current),
disk: Mutex::new(disk),
})
}
}
impl AuthSnapshot for FakeSnapshot {
fn current(&self) -> Option<KimiAuth> {
self.current
.lock()
.clone()
.filter(|a| !crate::auth::is_expired(a))
}
fn expired_auth(&self) -> Option<KimiAuth> {
self.current.lock().clone().filter(crate::auth::is_expired)
}
fn read_disk_auth(&self) -> Option<KimiAuth> {
self.disk.lock().clone()
}
fn is_expired(&self) -> bool {
self.current
.lock()
.as_ref()
.is_some_and(crate::auth::is_expired)
}
}
fn expired_session(key: &str, rt: &str) -> KimiAuth {
KimiAuth {
key: key.into(),
refresh_token: Some(rt.into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
expires_in: Some(3600),
..KimiAuth::test_default()
}
}
fn valid_session(key: &str, rt: &str) -> KimiAuth {
KimiAuth {
key: key.into(),
refresh_token: Some(rt.into()),
expires_at: Some(Utc::now() + Duration::hours(2)),
expires_in: Some(7200),
..KimiAuth::test_default()
}
}
fn token_json(access: &str, refresh: &str) -> serde_json::Value {
serde_json::json!({
"access_token": access,
"refresh_token": refresh,
"expires_in": 3600,
"scope": "kimi-code",
"token_type": "bearer",
})
}
#[tokio::test]
async fn refresh_success_returns_rotated_token() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.and(body_string_contains("refresh_token=rt-old"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-new", "rt-new")))
.expect(1)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-old");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
let RefreshOutcome::Success(new_auth) = outcome else {
panic!("expected success, got {outcome:?}");
};
assert_eq!(new_auth.key, "at-new");
assert_eq!(new_auth.refresh_token.as_deref(), Some("rt-new"));
}
#[tokio::test]
async fn adopts_valid_sibling_token_without_wire_call() {
// Disk has a fresh token with a different key: adopt, no HTTP.
let server = MockServer::start().await;
// No mock mounted: any request would 404 and fail the refresh.
let snap = FakeSnapshot::new(
Some(expired_session("at-old", "rt-old")),
Some(valid_session("at-sibling", "rt-sibling")),
);
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
let RefreshOutcome::Success(adopted) = outcome else {
panic!("expected sibling adoption, got {outcome:?}");
};
assert_eq!(adopted.key, "at-sibling");
assert!(
server.received_requests().await.unwrap().is_empty(),
"sibling adoption must not consume a refresh token on the wire"
);
}
#[tokio::test]
async fn unauthorized_tombstones_the_tried_refresh_token() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(401)
.set_body_json(serde_json::json!({ "error_description": "revoked" })),
)
.expect(1)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-dead");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
let RefreshOutcome::PermanentFailure {
error,
rejected_refresh_token,
} = outcome
else {
panic!("expected permanent failure, got {outcome:?}");
};
assert_eq!(error.reason, RefreshTokenFailedReason::RefreshTokenRejected);
assert_eq!(rejected_refresh_token.as_deref(), Some("rt-dead"));
}
#[tokio::test]
async fn unauthorized_adopts_sibling_rotation_instead_of_tombstoning() {
// 401 lands, but by the time we re-check, a sibling has persisted a
// rotated credential — adopt it (the mutual-logout race guard).
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-dead");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap.clone(), server.uri());
// Swap the persisted credential while the wire call is in flight.
let rotator = {
let snap = snap.clone();
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
*snap.disk.lock() = Some(valid_session("at-rotated", "rt-rotated"));
})
};
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
rotator.await.unwrap();
let RefreshOutcome::Success(adopted) = outcome else {
panic!("expected rotation adoption, got {outcome:?}");
};
assert_eq!(adopted.key, "at-rotated");
}
#[tokio::test]
async fn wire_exhaustion_is_transient_not_tombstoned() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(503))
.expect(3)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-old");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
assert!(
matches!(outcome, RefreshOutcome::TransientFailure { .. }),
"5xx exhaustion must stay transient: {outcome:?}"
);
}
#[tokio::test]
async fn no_credential_is_transient() {
let server = MockServer::start().await;
let snap = FakeSnapshot::new(None, None);
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
assert!(matches!(outcome, RefreshOutcome::TransientFailure { .. }));
}
}
crates/codegen/kigi-shell/src/auth/refresh/mod.rs
+39 -117
View File
@@ -1,42 +1,38 @@
mod external_refresher;
mod oidc_refresher;
mod kimi_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 crate::auth::model::KimiAuth;
use external_refresher::ExternalBinaryRefresher;
pub(crate) use oidc_refresher::OidcRefresher;
pub(crate) use kimi_refresher::KimiRefresher;
/// 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 current in-memory bearer outside the refresh threshold.
fn current(&self) -> Option<KimiAuth>;
/// 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>;
fn expired_auth(&self) -> Option<KimiAuth>;
/// Re-read the persisted credential (keyring → file) 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<KimiAuth>;
/// Whether the in-memory bearer is expired.
fn is_expired(&self) -> bool;
}
impl AuthSnapshot for AuthManager {
fn current(&self) -> Option<GrokAuth> {
fn current(&self) -> Option<KimiAuth> {
self.current()
}
fn expired_auth(&self) -> Option<GrokAuth> {
fn expired_auth(&self) -> Option<KimiAuth> {
self.expired_auth()
}
fn read_disk_auth(&self) -> Option<GrokAuth> {
fn read_disk_auth(&self) -> Option<KimiAuth> {
self.read_disk_auth()
}
fn is_expired(&self) -> bool {
@@ -44,31 +40,18 @@ impl AuthSnapshot for AuthManager {
}
}
/// 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.
/// The credential a refresh would send to the OAuth host: persisted
/// refresh-token first, then the expired in-mem bearer, then current (only on
/// `ServerRejected`). Single source of truth shared by
/// [`KimiRefresher::refresh`] (the attempt) and
/// `AuthManager::attempted_tombstone_key` (the tombstone scope), so the two
/// can't drift. The caller supplies the persisted read: the tombstone 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>,
disk_auth: Option<KimiAuth>,
reason: RefreshReason,
) -> Option<GrokAuth> {
) -> Option<KimiAuth> {
disk_auth
.filter(|a| a.refresh_token.is_some())
.or_else(|| snap.expired_auth())
@@ -84,18 +67,16 @@ pub(crate) fn resolve_refresh_credential(
#[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).
Success(Box<KimiAuth>),
/// Terminal failure (401/403 from the OAuth host). Caller records a
/// tombstone scoped to the rejected refresh token; the 300s cooldown (or
/// a rotated persisted refresh token) clears it.
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>,
/// The refresh-token value the refresher actually sent, so
/// `refresh_chain` scopes the tombstone to it. `None` when the
/// attempt never reached the wire.
rejected_refresh_token: Option<String>,
},
/// Transient / unknown failure. Caller may retry later. Message-only: the
/// underlying cause is logged structurally at the refresher, then flattened
@@ -105,19 +86,19 @@ pub(crate) enum RefreshOutcome {
impl RefreshOutcome {
/// A fresh credential from the authority (hides the `Box`).
pub(crate) fn success(auth: GrokAuth) -> Self {
pub(crate) fn success(auth: KimiAuth) -> 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).
/// Terminal failure for an already-classified reason against the
/// refresh token actually sent to the OAuth host.
pub(crate) fn permanent(
reason: crate::auth::error::RefreshTokenFailedReason,
tried_key: Option<String>,
rejected_refresh_token: Option<String>,
) -> Self {
Self::PermanentFailure {
error: reason.into(),
tried_key,
rejected_refresh_token,
}
}
@@ -139,67 +120,8 @@ pub(crate) trait TokenRefresher: Send + Sync {
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"
);
}
/// Build the production refresher against `kigi_env::oauth_host()`.
pub(crate) fn build_refresher(auth_manager: Arc<AuthManager>) -> Arc<dyn TokenRefresher> {
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host()))
}
crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher.rs
-260
View File
@@ -1,260 +0,0 @@
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
+257 -11
View File
@@ -2,7 +2,184 @@ 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};
use super::model::{API_KEY_SCOPE, AuthMode, AuthStore, KimiAuth, lookup_auth};
// ── System-keyring storage for the Kimi Code OAuth session ─────────────
//
// PRD F1: the OAuth token set lives in the system keyring (service `kigi`,
// entry `oauth/kimi-code`); when the keyring is unavailable we fall back to
// the file mechanism below (`auth.json`, owner-only, atomic writes). The
// official Kimi client's keyring entries (service `kimi-code`) and `~/.kimi`
// files are never touched.
/// Keyring service name — deliberately distinct from the official client's
/// `kimi-code` service.
#[cfg(any(target_os = "macos", windows))]
pub(crate) const KEYRING_SERVICE: &str = "kigi";
/// Outcome of a keyring read for the session scope.
#[derive(Debug)]
pub(crate) enum KeyringRead {
/// Backend reachable and the entry exists.
Found(Box<KimiAuth>),
/// Backend reachable, no entry stored.
Missing,
/// Keyring disabled, unsupported on this platform, or the backend
/// errored — callers fall back to the file store.
Unavailable,
}
/// Whether keyring storage participates for the session credential.
///
/// Disabled when:
/// - the platform has no supported backend (non-macOS/Windows builds),
/// - `KIGI_DISABLE_KEYRING` is set to a truthy value,
/// - a non-default credential location is in use (`KIGI_SHARE_DIR` /
/// `KIGI_AUTH_PATH`): the keyring entry belongs to the default user
/// install; alternate profiles (and tests) stay file-scoped, and
/// - in unit-test builds, unless a test explicitly opted into the mock
/// keyring via [`enable_mock_keyring_for_test`].
pub(crate) fn keyring_enabled() -> bool {
#[cfg(test)]
{
// Thread-local so keyring-specific tests (which opt in via
// `enable_mock_keyring_for_test`) can't leak the toggle into
// concurrently running persistence tests on other threads.
TEST_KEYRING_ENABLED.with(|flag| flag.get())
}
#[cfg(not(test))]
{
#[cfg(not(any(target_os = "macos", windows)))]
{
false
}
#[cfg(any(target_os = "macos", windows))]
{
let disabled = std::env::var("KIGI_DISABLE_KEYRING")
.is_ok_and(|v| !matches!(v.trim(), "" | "0" | "false" | "off" | "no"));
!disabled
&& std::env::var_os("KIGI_SHARE_DIR").is_none()
&& std::env::var_os("KIGI_AUTH_PATH").is_none()
}
}
}
#[cfg(test)]
thread_local! {
static TEST_KEYRING_ENABLED: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
/// Route keyring calls through an in-memory mock store for this process,
/// enable [`keyring_enabled`] on this thread, and clear any entry left by an
/// earlier test. Tests using this must serialize on the `kigi_keyring` key:
/// the mock entry is shared process-wide.
#[cfg(all(test, any(target_os = "macos", windows)))]
pub(crate) fn enable_mock_keyring_for_test() {
keyring::set_default_credential_builder(keyring::mock::default_credential_builder());
TEST_KEYRING_ENABLED.with(|flag| flag.set(true));
if let Err(e) = keyring_delete_session() {
panic!("mock keyring cleanup failed: {e}");
}
}
/// Disable the test keyring again (paired with
/// [`enable_mock_keyring_for_test`] in an RAII guard or test teardown).
#[cfg(test)]
pub(crate) fn disable_mock_keyring_for_test() {
TEST_KEYRING_ENABLED.with(|flag| flag.set(false));
}
/// The process-wide keyring entry handle. Cached so every reader/writer talks
/// to the same credential object: real backends read live state from the OS
/// store on each call, and the test mock keeps its state on the entry itself.
#[cfg(any(target_os = "macos", windows))]
fn keyring_entry() -> Result<&'static keyring::Entry, keyring::Error> {
static ENTRY: std::sync::OnceLock<Result<keyring::Entry, keyring::Error>> =
std::sync::OnceLock::new();
match ENTRY.get_or_init(|| {
keyring::Entry::new(KEYRING_SERVICE, crate::auth::config::KIMI_CODE_OAUTH_SCOPE)
}) {
Ok(entry) => Ok(entry),
// `keyring::Error` is not `Clone`; surface a stable equivalent.
Err(e) => {
tracing::warn!(error = %e, "auth: keyring entry construction failed");
Err(keyring::Error::Invalid(
"keyring entry".into(),
e.to_string(),
))
}
}
}
/// Read the session credential from the system keyring.
#[cfg(any(target_os = "macos", windows))]
pub(crate) fn keyring_read_session() -> KeyringRead {
if !keyring_enabled() {
return KeyringRead::Unavailable;
}
let entry = match keyring_entry() {
Ok(entry) => entry,
Err(e) => {
tracing::warn!(error = %e, "auth: keyring entry unavailable, falling back to file");
return KeyringRead::Unavailable;
}
};
match entry.get_password() {
Ok(raw) => match serde_json::from_str::<KimiAuth>(&raw) {
Ok(auth) => KeyringRead::Found(Box::new(auth)),
Err(e) => {
tracing::warn!(error = %e, "auth: keyring entry is not valid JSON, ignoring");
KeyringRead::Missing
}
},
Err(keyring::Error::NoEntry) => KeyringRead::Missing,
Err(e) => {
tracing::warn!(error = %e, "auth: keyring read failed, falling back to file");
KeyringRead::Unavailable
}
}
}
#[cfg(not(any(target_os = "macos", windows)))]
pub(crate) fn keyring_read_session() -> KeyringRead {
KeyringRead::Unavailable
}
/// Write the session credential to the system keyring.
#[cfg(any(target_os = "macos", windows))]
pub(crate) fn keyring_write_session(auth: &KimiAuth) -> anyhow::Result<()> {
anyhow::ensure!(keyring_enabled(), "keyring storage disabled");
let payload = serde_json::to_string(auth)?;
keyring_entry()?.set_password(&payload)?;
tracing::info!("auth: session credential written to system keyring");
Ok(())
}
#[cfg(not(any(target_os = "macos", windows)))]
pub(crate) fn keyring_write_session(_auth: &KimiAuth) -> anyhow::Result<()> {
anyhow::bail!("keyring storage is not supported on this platform")
}
/// Delete the session credential from the system keyring (Ok when absent).
#[cfg(any(target_os = "macos", windows))]
pub(crate) fn keyring_delete_session() -> anyhow::Result<()> {
if !keyring_enabled() {
return Ok(());
}
match keyring_entry()?.delete_credential() {
Ok(()) => {
tracing::info!("auth: session credential removed from system keyring");
Ok(())
}
Err(keyring::Error::NoEntry) => Ok(()),
Err(e) => Err(e.into()),
}
}
#[cfg(not(any(target_os = "macos", windows)))]
pub(crate) fn keyring_delete_session() -> anyhow::Result<()> {
Ok(())
}
/// RAII guard for an exclusive advisory lock on `auth.json.lock`.
/// The lock is released when the inner `File` is dropped (closing the FD).
@@ -325,25 +502,23 @@ fn restore_prior_bytes(auth_file: &Path, bytes: &[u8]) -> std::io::Result<()> {
}
/// 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`."))?;
read_auth_json(&path).map_err(|_| anyhow::anyhow!("Not logged in. Run `kigi 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.")
anyhow::anyhow!("Your auth token is invalid. Run `kigi login` to re-authenticate.")
})
}
/// Read the API key from the `xai::api_key` scope in auth.json.
/// Read the API key from the `kigi::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.
/// Store a plain API key in auth.json under the `kigi::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.
@@ -352,7 +527,7 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> {
let mut map = read_auth_json_or_empty_recovering_corrupt(&path)?;
map.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
KimiAuth {
key: api_key.to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
@@ -361,7 +536,7 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> {
write_auth_json(&path, &map)
}
/// Remove the `xai::api_key` scope from auth.json.
/// Remove the `kigi::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) {
@@ -383,7 +558,7 @@ mod write_fallback_tests {
let mut map = AuthStore::new();
map.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
KimiAuth {
key: "secret-key".to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
@@ -481,7 +656,7 @@ mod write_fallback_tests {
let mut replacement = AuthStore::new();
replacement.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
KimiAuth {
key: "replacement-key".to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
@@ -510,3 +685,74 @@ mod write_fallback_tests {
assert_eq!(mode & 0o777, 0o600, "restored file must stay 0o600");
}
}
#[cfg(all(test, any(target_os = "macos", windows)))]
mod keyring_tests {
use super::*;
use chrono::Utc;
/// RAII teardown so a panicking test doesn't leave the process-global
/// test-keyring toggle enabled for later tests.
struct MockKeyringGuard;
impl MockKeyringGuard {
fn enable() -> Self {
enable_mock_keyring_for_test();
Self
}
}
impl Drop for MockKeyringGuard {
fn drop(&mut self) {
disable_mock_keyring_for_test();
}
}
fn session_auth(key: &str, rt: &str) -> KimiAuth {
KimiAuth {
key: key.into(),
refresh_token: Some(rt.into()),
expires_at: Some(Utc::now() + chrono::Duration::seconds(3600)),
expires_in: Some(3600),
scope: Some("kimi-code".into()),
token_type: Some("bearer".into()),
..KimiAuth::test_default()
}
}
#[test]
#[serial_test::serial(kigi_keyring)]
fn keyring_session_roundtrip() {
let _guard = MockKeyringGuard::enable();
// Fresh mock store: nothing there yet.
assert!(matches!(keyring_read_session(), KeyringRead::Missing));
keyring_write_session(&session_auth("at-1", "rt-1")).unwrap();
let KeyringRead::Found(read) = keyring_read_session() else {
panic!("expected Found after write");
};
assert_eq!(read.key, "at-1");
assert_eq!(read.refresh_token.as_deref(), Some("rt-1"));
assert_eq!(read.expires_in, Some(3600));
// Overwrite rotates in place.
keyring_write_session(&session_auth("at-2", "rt-2")).unwrap();
let KeyringRead::Found(read) = keyring_read_session() else {
panic!("expected Found after rotate");
};
assert_eq!(read.key, "at-2");
// Delete is idempotent.
keyring_delete_session().unwrap();
assert!(matches!(keyring_read_session(), KeyringRead::Missing));
keyring_delete_session().unwrap();
}
#[test]
#[serial_test::serial(kigi_keyring)]
fn keyring_disabled_reads_unavailable() {
disable_mock_keyring_for_test();
assert!(matches!(keyring_read_session(), KeyringRead::Unavailable));
assert!(keyring_write_session(&session_auth("a", "r")).is_err());
// Delete when disabled is a no-op success (logout stays best-effort).
keyring_delete_session().unwrap();
}
}
crates/codegen/kigi-shell/src/auth/token_type.rs
+36 -20
View File
@@ -1,17 +1,13 @@
use crate::auth::model::{AuthMode, GrokAuth};
use crate::auth::model::{AuthMode, KimiAuth};
/// 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,
/// Kimi Code OAuth session with a refresh_token available.
OAuthSession,
/// OAuth session without a refresh_token (cannot be silently renewed).
SessionNoRefresh,
/// Plain API key (no refresh possible).
ApiKey,
/// No credentials loaded.
@@ -20,36 +16,56 @@ pub(crate) enum TokenType {
impl TokenType {
/// Classify the loaded credential (pure; no manager state).
pub(crate) fn from_auth(auth: Option<&GrokAuth>) -> Self {
pub(crate) fn from_auth(auth: Option<&KimiAuth>) -> 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::OAuth if a.refresh_token.is_some() => Self::OAuthSession,
AuthMode::OAuth => Self::SessionNoRefresh,
AuthMode::ApiKey => Self::ApiKey,
},
}
}
/// `true` for types that can be silently refreshed (OIDC, external binary).
/// `true` for types that can be silently refreshed.
pub(crate) fn is_refreshable(self) -> bool {
matches!(self, Self::OidcSession | Self::ExternalBinary)
matches!(self, Self::OAuthSession)
}
}
#[cfg(test)]
mod tests {
//! Per-variant matrix for `is_refreshable`.
//! Per-variant matrix for `is_refreshable` and classification.
use super::*;
#[test]
fn is_refreshable_matrix() {
assert!(TokenType::OidcSession.is_refreshable());
assert!(TokenType::ExternalBinary.is_refreshable());
assert!(!TokenType::LegacySession.is_refreshable());
assert!(TokenType::OAuthSession.is_refreshable());
assert!(!TokenType::SessionNoRefresh.is_refreshable());
assert!(!TokenType::ApiKey.is_refreshable());
assert!(!TokenType::None.is_refreshable());
}
#[test]
fn from_auth_classifies_by_mode_and_refresh_token() {
assert_eq!(TokenType::from_auth(None), TokenType::None);
let with_rt = KimiAuth {
refresh_token: Some("rt".into()),
..KimiAuth::test_default()
};
assert_eq!(
TokenType::from_auth(Some(&with_rt)),
TokenType::OAuthSession
);
let no_rt = KimiAuth::test_default();
assert_eq!(
TokenType::from_auth(Some(&no_rt)),
TokenType::SessionNoRefresh
);
let api = KimiAuth {
auth_mode: AuthMode::ApiKey,
..KimiAuth::test_default()
};
assert_eq!(TokenType::from_auth(Some(&api)), TokenType::ApiKey);
}
}