mirror of
https://github.com/block/buzz.git
synced 2026-08-18 06:50:31 +02:00
fix(buzz-agent): rejected-aware auth, typed refresh, real process tests
Round-2 fixes to the DatabricksAuthCoordinator: - Rejected-aware acquisition: `acquire_with_intent(intent, rejected)` lets the saved-picker recovery path replace a locally-fresh bearer the server just 401'd, so Auto/UserInitiated escalate to a browser instead of re-returning the dead token. `refresh_now`'s hardcoded Headless could not. - Typed `RefreshOutcome`: only a token-endpoint grant rejection becomes RefreshRejected/browser-fallback; transport/timeout/5xx/decode failures stay NetworkUnavailable so a transient fault never pops a browser. - In-process shared-future joiner (INFLIGHT/InflightSlot/LeaderGuard): a pre-existing joiner receives the leader's SAME failure result rather than acquiring the lock afterward and launching a second browser. - Deadline holes closed: the HTTP client build error propagates (no untimed fallback), and every interactive timeout exits through the common outcome writer so TimedOut is recorded in the cooldown sidecar under the held lock. - Real cross-process tests: a `lock-holder` child bin takes the advisory lock so single-flight and crash-release are proven across processes, not simulated with same-process handles. buzz-agent added to the Justfile test-unit lane so CI executes these tests. - MSRV: use fs2::FileExt instead of std File::try_lock/unlock (1.89+), restoring the crate to the repo's declared 1.88 floor. The refresh-timeout test injects a short real-time HTTP timeout rather than pausing the clock: under start_paused tokio auto-advances into the timer while the real loopback discovery call is still in flight, tripping the timeout on the wrong request. Co-authored-by: Will Pfleger <pfleger.will@gmail.com> Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
This commit is contained in:
Generated
+11
@@ -863,6 +863,7 @@ dependencies = [
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"dirs",
|
||||
"fs2",
|
||||
"getrandom 0.4.3",
|
||||
"hex",
|
||||
"nix 0.31.3",
|
||||
@@ -2990,6 +2991,16 @@ dependencies = [
|
||||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs2"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs_extra"
|
||||
version = "1.3.0"
|
||||
|
||||
@@ -319,6 +319,13 @@ test-unit:
|
||||
# because nothing in CI runs `cargo test --workspace` — workspace
|
||||
# membership alone buys clippy/check, not a single executed test.
|
||||
cargo nextest run -p buzz-backend-kubernetes
|
||||
# buzz-agent: the OAuth auth coordinator's library concurrency matrix
|
||||
# plus the databricks integration tests (lock single-flight, cooldown,
|
||||
# cross-process crash recovery) are infra-free — a stub OIDC provider
|
||||
# and an injected browser opener, no network or Postgres. Enumerated
|
||||
# here because nothing in CI runs `cargo test --workspace`, so without
|
||||
# this stanza the crate's tests never execute remotely.
|
||||
cargo nextest run -p buzz-agent
|
||||
else
|
||||
./scripts/run-tests.sh unit
|
||||
fi
|
||||
|
||||
@@ -2,9 +2,7 @@
|
||||
name = "buzz-agent"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
# Above the 1.88 workspace floor: the auth coordinator's cross-process
|
||||
# single-flight uses `std::fs::File::try_lock`/`unlock`, stable since 1.89.
|
||||
rust-version = "1.89.0"
|
||||
rust-version.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Minimal, unbreakable ACP-compliant agent. Non-streaming. Tool-calls-as-output."
|
||||
@@ -26,6 +24,14 @@ path = "src/main.rs"
|
||||
name = "fake-mcp"
|
||||
path = "tests/bin/fake_mcp.rs"
|
||||
|
||||
# Test-only lock holder: a real second process that takes the coordinator's
|
||||
# cross-process advisory lock, so the auth tests can prove genuine
|
||||
# inter-process single-flight and crash-release rather than same-process
|
||||
# handles. Tiny; only used by the databricks auth integration tests.
|
||||
[[bin]]
|
||||
name = "lock-holder"
|
||||
path = "tests/bin/lock_holder.rs"
|
||||
|
||||
[dependencies]
|
||||
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time", "net"] }
|
||||
serde = { workspace = true }
|
||||
@@ -47,6 +53,11 @@ url = { workspace = true }
|
||||
urlencoding = "2"
|
||||
webbrowser = "1"
|
||||
dirs = "6"
|
||||
# Cross-process advisory file lock (flock on Unix, LockFileEx on Windows) for
|
||||
# the auth coordinator's single-flight. Kept off std's `File::try_lock` so the
|
||||
# crate stays buildable on the repo's declared 1.88 MSRV (those std APIs are
|
||||
# 1.89+).
|
||||
fs2 = "0.4"
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
|
||||
|
||||
+404
-55
@@ -15,19 +15,21 @@
|
||||
//! captures the redirect, and exchanges the code for a token. Subsequent
|
||||
//! calls hit the cache and silently refresh when expired.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::io::{self, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::Engine;
|
||||
use fs2::FileExt;
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::Digest;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::{watch, Mutex};
|
||||
|
||||
use crate::types::AgentError;
|
||||
|
||||
@@ -319,6 +321,24 @@ struct OidcEndpoints {
|
||||
token_endpoint: String,
|
||||
}
|
||||
|
||||
/// Typed result of a refresh-token grant, so the coordinator can separate an
|
||||
/// actual credential rejection from a transient fault.
|
||||
///
|
||||
/// - [`Refreshed`](Self::Refreshed): a fresh token — success.
|
||||
/// - [`Rejected`](Self::Rejected): the token endpoint rejected the *grant*
|
||||
/// (dead/rotated refresh token). This is the only outcome that becomes
|
||||
/// [`AuthError::RefreshRejected`] for `Headless` or drives a browser
|
||||
/// fallback for interactive intents.
|
||||
/// - [`Network`](Self::Network): transport error, timeout, 5xx, or an
|
||||
/// undecodable/malformed response — infrastructural, never a credential
|
||||
/// decision, so it surfaces as [`AuthError::NetworkUnavailable`] and never
|
||||
/// pops a browser.
|
||||
enum RefreshOutcome {
|
||||
Refreshed(CachedToken),
|
||||
Rejected,
|
||||
Network,
|
||||
}
|
||||
|
||||
/// PKCE OAuth token source with on-disk refresh cache.
|
||||
///
|
||||
/// First call:
|
||||
@@ -355,6 +375,24 @@ impl PkceOAuthTokenSource {
|
||||
pub fn new_with(
|
||||
cfg: PkceOAuthConfig,
|
||||
opener: Arc<dyn BrowserOpener>,
|
||||
) -> Result<Arc<Self>, AgentError> {
|
||||
Self::new_with_http_timeout(cfg, opener, HTTP_REQUEST_TIMEOUT)
|
||||
}
|
||||
|
||||
/// Construct with an injected opener *and* an explicit per-request HTTP
|
||||
/// timeout. Only the refresh-timeout integration test passes the timeout
|
||||
/// argument: it drives a hung token endpoint against a short bound so the
|
||||
/// per-request timeout classification (`NetworkUnavailable`, never
|
||||
/// `RefreshRejected`) is exercised in real time. A paused-clock test can't
|
||||
/// do this — tokio auto-advances into the timer while the real loopback
|
||||
/// discovery call is still in flight, tripping the timeout on the wrong
|
||||
/// request. Every production and other-test path goes through
|
||||
/// [`new`](Self::new) or [`new_with`](Self::new_with) at the default
|
||||
/// [`HTTP_REQUEST_TIMEOUT`].
|
||||
pub fn new_with_http_timeout(
|
||||
cfg: PkceOAuthConfig,
|
||||
opener: Arc<dyn BrowserOpener>,
|
||||
http_timeout: Duration,
|
||||
) -> Result<Arc<Self>, AgentError> {
|
||||
let cache_path = cache_path_for(&cfg)?;
|
||||
if let Some(parent) = cache_path.parent() {
|
||||
@@ -363,12 +401,14 @@ impl PkceOAuthTokenSource {
|
||||
}
|
||||
// Every OAuth HTTP call inherits this timeout so a hung provider can
|
||||
// never stall the caller — nor the same-key callers waiting on the
|
||||
// cross-process lock this holder owns. A build failure falls back to
|
||||
// the untimed default rather than making construction fallible.
|
||||
// cross-process lock this holder owns. Construction is fallible, so a
|
||||
// build failure propagates rather than silently falling back to an
|
||||
// untimed client — an untimed client would restore exactly the
|
||||
// unbounded-HTTP-under-lock failure the timeout exists to prevent.
|
||||
let http = Client::builder()
|
||||
.timeout(HTTP_REQUEST_TIMEOUT)
|
||||
.timeout(http_timeout)
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new());
|
||||
.map_err(|e| AgentError::Llm(format!("oauth http client: {e}")))?;
|
||||
let initial = read_cache(&cache_path);
|
||||
Ok(Arc::new(Self {
|
||||
cfg,
|
||||
@@ -439,32 +479,62 @@ impl PkceOAuthTokenSource {
|
||||
}
|
||||
|
||||
/// Exchange a refresh token for a fresh access token.
|
||||
async fn refresh(
|
||||
&self,
|
||||
endpoints: &OidcEndpoints,
|
||||
refresh_token: &str,
|
||||
) -> Result<CachedToken, AgentError> {
|
||||
///
|
||||
/// The outcome is typed so the caller can tell an actual credential
|
||||
/// rejection apart from a transient fault. Only a token-endpoint rejection
|
||||
/// of the grant itself (a 4xx `invalid_grant`-class response) is a dead
|
||||
/// refresh token; a transport failure, timeout, 5xx, or an
|
||||
/// undecodable/malformed response is infrastructural and must never be
|
||||
/// mistaken for a credential decision (it would otherwise pop a browser or
|
||||
/// return `RefreshRejected` when nothing was actually rejected).
|
||||
async fn refresh(&self, endpoints: &OidcEndpoints, refresh_token: &str) -> RefreshOutcome {
|
||||
let params = [
|
||||
("grant_type", "refresh_token"),
|
||||
("refresh_token", refresh_token),
|
||||
("client_id", &self.cfg.client_id),
|
||||
];
|
||||
let resp = self
|
||||
let resp = match self
|
||||
.http
|
||||
.post(&endpoints.token_endpoint)
|
||||
.form(¶ms)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth refresh: {e}")))?;
|
||||
if !resp.status().is_success() {
|
||||
{
|
||||
Ok(resp) => resp,
|
||||
// Transport error or the per-request timeout elapsed: no verdict
|
||||
// from the provider, so this is infrastructural, not a rejection.
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh transport failure");
|
||||
return RefreshOutcome::Network;
|
||||
}
|
||||
};
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(AgentError::Llm(format!("oauth refresh failed: {body}")));
|
||||
// A 4xx is the token endpoint rejecting the grant (dead/rotated
|
||||
// refresh token). A 5xx is a provider-side fault — transient, not a
|
||||
// credential decision — so it stays in the infrastructural bucket.
|
||||
if status.is_client_error() {
|
||||
tracing::warn!(status = %status, body = %body, "oauth refresh grant rejected");
|
||||
return RefreshOutcome::Rejected;
|
||||
}
|
||||
tracing::warn!(status = %status, body = %body, "oauth refresh server error");
|
||||
return RefreshOutcome::Network;
|
||||
}
|
||||
let v: Value = match resp.json().await {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh response decode failure");
|
||||
return RefreshOutcome::Network;
|
||||
}
|
||||
};
|
||||
match token_from_response(&v, Some(refresh_token)) {
|
||||
Ok(token) => RefreshOutcome::Refreshed(token),
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh response missing access_token");
|
||||
RefreshOutcome::Network
|
||||
}
|
||||
}
|
||||
let v: Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth refresh json: {e}")))?;
|
||||
token_from_response(&v, Some(refresh_token))
|
||||
}
|
||||
|
||||
/// Run the full browser-mediated Authorization Code + PKCE flow and cache
|
||||
@@ -473,19 +543,40 @@ impl PkceOAuthTokenSource {
|
||||
/// and single-flights with concurrent callers on the cross-process lock. A
|
||||
/// still-valid cached token short-circuits to success without re-prompting.
|
||||
///
|
||||
/// This is the no-rejected convenience: it trusts the local expiry clock,
|
||||
/// so a not-yet-expired cached token is accepted. When the caller already
|
||||
/// knows the cached bearer was rejected by the server (a 401), it must use
|
||||
/// [`acquire_with_intent`](Self::acquire_with_intent) with `rejected` set
|
||||
/// so the stale-but-fresh token can't short-circuit the sign-in.
|
||||
///
|
||||
/// [`UserInitiated`]: AuthIntent::UserInitiated
|
||||
pub async fn interactive_login(&self) -> Result<(), AgentError> {
|
||||
self.acquire(AuthIntent::UserInitiated, None).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Public entry for passive Desktop discovery (Phase 2): acquire a bearer
|
||||
/// under an explicit [`AuthIntent`], returning the typed [`AuthError`] so
|
||||
/// the caller can branch on a stable `code` rather than display text. The
|
||||
/// [`TokenSource`] trait methods wrap this and flatten the error into
|
||||
/// [`AgentError`].
|
||||
pub async fn acquire_with_intent(&self, intent: AuthIntent) -> Result<String, AuthError> {
|
||||
self.acquire(intent, None).await
|
||||
/// Public entry for passive Desktop discovery and the saved-model picker
|
||||
/// (Phase 2): acquire a bearer under an explicit [`AuthIntent`], returning
|
||||
/// the typed [`AuthError`] so the caller can branch on a stable `code`
|
||||
/// rather than display text. The [`TokenSource`] trait methods wrap this
|
||||
/// and flatten the error into [`AgentError`].
|
||||
///
|
||||
/// `rejected` carries the exact access token the provider just 401'd, if
|
||||
/// any. With `rejected = None` a locally-fresh cached token is a hit (the
|
||||
/// normal discovery path). With `rejected = Some(t)` the expiry clock is
|
||||
/// untrustworthy — the rejected token looked fresh — so a cached token
|
||||
/// equal to `t` is *not* a hit: the acquisition refreshes, and for `Auto`
|
||||
/// or `UserInitiated` falls through to a browser when the refresh grant is
|
||||
/// dead. This is what lets the saved-picker recovery path say "this
|
||||
/// locally-fresh bearer was just rejected — replace it" instead of
|
||||
/// re-returning the dead token, which `refresh_now`'s hardcoded
|
||||
/// [`Headless`](AuthIntent::Headless) can never escalate to a browser.
|
||||
pub async fn acquire_with_intent(
|
||||
&self,
|
||||
intent: AuthIntent,
|
||||
rejected: Option<&str>,
|
||||
) -> Result<String, AuthError> {
|
||||
self.acquire(intent, rejected).await
|
||||
}
|
||||
|
||||
/// Return a usable cached bearer, applying the identity rule for a
|
||||
@@ -555,14 +646,64 @@ impl PkceOAuthTokenSource {
|
||||
intent: AuthIntent,
|
||||
rejected: Option<&str>,
|
||||
) -> Result<String, AuthError> {
|
||||
// Fast path: no lock, no network.
|
||||
// Fast path: no lock, no network. `try_lock` rather than `lock().await`
|
||||
// so a caller arriving while a leader holds `state` across its browser
|
||||
// flow does not block here — it falls through to the in-process
|
||||
// registry below and joins the leader instead of waiting out the whole
|
||||
// flow and then racing in as a second leader. A cache hit is still
|
||||
// served without the file lock; a miss (or contention) coalesces.
|
||||
{
|
||||
let mut state = self.state.lock().await;
|
||||
if let Some(hit) = self.cached_hit(&mut state, rejected) {
|
||||
return Ok(hit);
|
||||
if let Ok(mut state) = self.state.try_lock() {
|
||||
if let Some(hit) = self.cached_hit(&mut state, rejected) {
|
||||
return Ok(hit);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// In-process single-flight (see [`INFLIGHT`]). Keyed by (lock path,
|
||||
// browser capability): browser-capable callers coalesce with each
|
||||
// other, so a caller that was already waiting when the leader's attempt
|
||||
// was in flight shares the leader's result instead of taking the lock
|
||||
// after it and launching a second browser. A `Headless` caller never
|
||||
// shares a browser-capable slot (and vice versa), so a racing inference
|
||||
// call is neither handed an interactive failure nor able to deny an
|
||||
// explicit sign-in its browser — those two intents still coordinate
|
||||
// only through the cross-process file lock.
|
||||
let key: InflightKey = (self.lock_path(), intent.may_open_browser());
|
||||
let (slot, is_leader) = {
|
||||
let mut reg = inflight_registry();
|
||||
match reg.get(&key) {
|
||||
Some(existing) => (existing.clone(), false),
|
||||
None => {
|
||||
let slot = Arc::new(InflightSlot::new());
|
||||
reg.insert(key.clone(), slot.clone());
|
||||
(slot, true)
|
||||
}
|
||||
}
|
||||
};
|
||||
if !is_leader {
|
||||
// Pre-existing joiner: observe the leader's outcome.
|
||||
return slot.wait().await;
|
||||
}
|
||||
|
||||
// Leader: run the real flow, then evict + publish. The guard makes
|
||||
// eviction and joiner wake-up happen even if this future is cancelled
|
||||
// or panics, so a dropped leader can never wedge its joiners or leave a
|
||||
// dead slot that turns later callers into joiners of nothing.
|
||||
let guard = LeaderGuard::new(key, slot);
|
||||
let result = self.acquire_leader(intent, rejected).await;
|
||||
guard.complete(result)
|
||||
}
|
||||
|
||||
/// The leader's slow-path body: take the cross-process lock, then run the
|
||||
/// bounded acquisition under it. Split out so [`acquire`] can wrap it in
|
||||
/// the in-process single-flight without the lock/deadline logic bleeding
|
||||
/// into the joiner path.
|
||||
async fn acquire_leader(
|
||||
&self,
|
||||
intent: AuthIntent,
|
||||
rejected: Option<&str>,
|
||||
) -> Result<String, AuthError> {
|
||||
// Slow path: one flow at a time per cache key. The waiter's deadline
|
||||
// exceeds a healthy holder's attempt deadline, so it never gives up on
|
||||
// a live holder.
|
||||
@@ -570,21 +711,31 @@ impl PkceOAuthTokenSource {
|
||||
let _guard = acquire_auth_lock(&self.lock_path(), deadline).await?;
|
||||
|
||||
// Bound the whole locked attempt so a wedged flow can't hold the lock
|
||||
// past the waiters' patience. On expiry the lock releases (guard drop)
|
||||
// and the attempt reports TimedOut.
|
||||
match tokio::time::timeout(AUTH_ATTEMPT_DEADLINE, self.acquire_locked(intent, rejected))
|
||||
// past the waiters' patience. The deadline is passed *into*
|
||||
// `acquire_locked` rather than wrapped around it in a cancelling
|
||||
// `tokio::time::timeout`: a cancel drops the future at an arbitrary
|
||||
// await point, which would skip the cooldown write for a timed-out
|
||||
// interactive attempt and let the next `Auto` caller re-pop a browser.
|
||||
// Threading the deadline lets every interactive timeout exit through
|
||||
// the common outcome writer while the lock is still held.
|
||||
let attempt_deadline = std::time::Instant::now() + AUTH_ATTEMPT_DEADLINE;
|
||||
self.acquire_locked(intent, rejected, attempt_deadline)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(_) => Err(AuthError::TimedOut),
|
||||
}
|
||||
}
|
||||
|
||||
/// Slow-path body, run while holding the cross-process auth lock.
|
||||
///
|
||||
/// `attempt_deadline` bounds the whole locked flow. Discovery and refresh
|
||||
/// are each bounded by the HTTP client's per-request timeout; the browser
|
||||
/// flow is wrapped in the *remaining* budget so a total-deadline expiry
|
||||
/// during the interactive step surfaces as [`AuthError::TimedOut`] through
|
||||
/// the same arm that records the cooldown — never as a cancellation that
|
||||
/// drops the guard without writing it.
|
||||
async fn acquire_locked(
|
||||
&self,
|
||||
intent: AuthIntent,
|
||||
rejected: Option<&str>,
|
||||
attempt_deadline: std::time::Instant,
|
||||
) -> Result<String, AuthError> {
|
||||
let mut state = self.state.lock().await;
|
||||
|
||||
@@ -602,10 +753,23 @@ impl PkceOAuthTokenSource {
|
||||
if let Some(rt) = state.as_ref().and_then(|t| t.refresh_token.clone()) {
|
||||
let eps = self.discover(&mut endpoints).await?;
|
||||
match self.refresh(eps, &rt).await {
|
||||
Ok(fresh) => return self.finish(&mut state, fresh),
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh failed; falling back");
|
||||
// A sibling may still have won the race while we ran.
|
||||
RefreshOutcome::Refreshed(fresh) => return self.finish(&mut state, fresh),
|
||||
// A transient fault (transport/timeout/5xx/decode) is not a
|
||||
// credential decision: never fall through to a browser or
|
||||
// report RefreshRejected. A sibling may have written a fresh
|
||||
// token while we ran, so honor that first; otherwise this is
|
||||
// infrastructural and surfaces as NetworkUnavailable.
|
||||
RefreshOutcome::Network => {
|
||||
if let Some(hit) = self.cached_hit(&mut state, rejected) {
|
||||
return Ok(hit);
|
||||
}
|
||||
return Err(AuthError::NetworkUnavailable);
|
||||
}
|
||||
// The token endpoint rejected the grant: a dead refresh token.
|
||||
// A sibling may still have won the race while we ran; if not,
|
||||
// fall through to a browser (interactive) or RefreshRejected
|
||||
// (headless).
|
||||
RefreshOutcome::Rejected => {
|
||||
if let Some(hit) = self.cached_hit(&mut state, rejected) {
|
||||
return Ok(hit);
|
||||
}
|
||||
@@ -636,7 +800,19 @@ impl PkceOAuthTokenSource {
|
||||
}
|
||||
|
||||
let eps = self.discover(&mut endpoints).await?;
|
||||
match browser_pkce_flow(&self.http, &self.cfg, eps, self.opener.as_ref()).await {
|
||||
// Wrap the browser flow in the *remaining* attempt budget so the total
|
||||
// locked time never exceeds `attempt_deadline` (and thus never
|
||||
// outlasts a waiter's `LOCK_WAIT_TIMEOUT`). A deadline expiry maps to
|
||||
// `TimedOut`, which is cooldown-worthy, so it flows through the same
|
||||
// writer arm below instead of being dropped by a cancel that would
|
||||
// release the lock without recording the cooldown.
|
||||
let remaining = attempt_deadline.saturating_duration_since(std::time::Instant::now());
|
||||
let flow = browser_pkce_flow(&self.http, &self.cfg, eps, self.opener.as_ref());
|
||||
let outcome = match tokio::time::timeout(remaining, flow).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => Err(AuthError::TimedOut),
|
||||
};
|
||||
match outcome {
|
||||
// `finish` clears the cooldown on success.
|
||||
Ok(fresh) => self.finish(&mut state, fresh),
|
||||
Err(e) => {
|
||||
@@ -824,20 +1000,24 @@ struct AuthLockGuard(fs::File);
|
||||
impl Drop for AuthLockGuard {
|
||||
fn drop(&mut self) {
|
||||
// Explicit for intent; closing the fd would release it regardless.
|
||||
let _ = self.0.unlock();
|
||||
let _ = FileExt::unlock(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Acquire the cross-process auth lock at `path`, polling until `deadline`.
|
||||
///
|
||||
/// `File::try_lock` is per–open-file-description, so a lock taken on one
|
||||
/// handle blocks every other handle — same process or not — which is exactly
|
||||
/// the single-flight guarantee we want without a separate in-memory registry.
|
||||
/// The lock is non-blocking, so we poll on [`LOCK_POLL_INTERVAL`] rather than
|
||||
/// parking a worker thread in a blocking `lock()`. A waiter whose `deadline`
|
||||
/// lapses returns [`AuthError::LockTimeout`]; because the caller sets that
|
||||
/// deadline longer than [`AUTH_ATTEMPT_DEADLINE`], a healthy holder always
|
||||
/// finishes first.
|
||||
/// `fs2::FileExt::try_lock_exclusive` maps to `flock(LOCK_EX | LOCK_NB)` on
|
||||
/// Unix and `LockFileEx` on Windows — advisory, per–open-file-description, so
|
||||
/// a lock taken on one handle blocks every other handle (same process or not),
|
||||
/// which is exactly the cross-process single-flight guarantee we want. The
|
||||
/// try-lock is non-blocking, so we poll on [`LOCK_POLL_INTERVAL`] rather than
|
||||
/// parking a worker thread in a blocking `lock_exclusive()`. Contention is
|
||||
/// reported as [`fs2::lock_contended_error`] (`EWOULDBLOCK`/`EACCES` on Unix,
|
||||
/// `ERROR_LOCK_VIOLATION` on Windows); we match its `raw_os_error` and retry.
|
||||
/// Any other error is a real fault and returns [`AuthError::LockTimeout`]. A
|
||||
/// waiter whose `deadline` lapses also returns [`AuthError::LockTimeout`];
|
||||
/// because the caller sets that deadline longer than [`AUTH_ATTEMPT_DEADLINE`],
|
||||
/// a healthy holder always finishes first.
|
||||
async fn acquire_auth_lock(
|
||||
path: &Path,
|
||||
deadline: std::time::Instant,
|
||||
@@ -851,20 +1031,151 @@ async fn acquire_auth_lock(
|
||||
.write(true)
|
||||
.open(path)
|
||||
.map_err(|_| AuthError::LockTimeout)?;
|
||||
let contended = fs2::lock_contended_error().raw_os_error();
|
||||
loop {
|
||||
match file.try_lock() {
|
||||
match file.try_lock_exclusive() {
|
||||
Ok(()) => return Ok(AuthLockGuard(file)),
|
||||
Err(fs::TryLockError::WouldBlock) => {
|
||||
Err(e) if e.raw_os_error() == contended => {
|
||||
if std::time::Instant::now() >= deadline {
|
||||
return Err(AuthError::LockTimeout);
|
||||
}
|
||||
tokio::time::sleep(LOCK_POLL_INTERVAL).await;
|
||||
}
|
||||
Err(fs::TryLockError::Error(_)) => return Err(AuthError::LockTimeout),
|
||||
Err(_) => return Err(AuthError::LockTimeout),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Key for the in-process single-flight registry: the cross-process lock path
|
||||
/// (one per cache key) paired with whether the caller may open a browser.
|
||||
/// Browser-capable callers (`Auto`/`UserInitiated`) coalesce with each other;
|
||||
/// a `Headless` caller keys separately so it neither inherits an interactive
|
||||
/// failure nor denies an explicit sign-in its browser — those two still
|
||||
/// coordinate through the cross-process file lock, not this registry.
|
||||
type InflightKey = (PathBuf, bool);
|
||||
|
||||
/// Process-global registry of in-flight auth attempts, the in-process
|
||||
/// counterpart to [`acquire_auth_lock`]'s cross-process file lock. The file
|
||||
/// lock serializes work across processes and shares *success* via a cache
|
||||
/// re-read, but a queued caller that acquires the lock after a browser denial
|
||||
/// would clear the sidecar and pop a second browser. This registry closes that
|
||||
/// gap: a caller that arrives while a leader's attempt is in flight joins the
|
||||
/// leader's [`InflightSlot`] and receives the *same* result — success or
|
||||
/// failure — instead of taking the lock afterward and launching again. Guarded
|
||||
/// by a `std::sync::Mutex` because every critical section is a cheap map lookup
|
||||
/// with no `.await` held.
|
||||
static INFLIGHT: LazyLock<std::sync::Mutex<HashMap<InflightKey, Arc<InflightSlot>>>> =
|
||||
LazyLock::new(|| std::sync::Mutex::new(HashMap::new()));
|
||||
|
||||
/// Lock the in-flight registry, recovering from a poisoned mutex rather than
|
||||
/// panicking: the only work done under this lock is map lookups that can't
|
||||
/// leave inconsistent state, so a poison from an unrelated panic must not wedge
|
||||
/// every future auth attempt.
|
||||
fn inflight_registry() -> std::sync::MutexGuard<'static, HashMap<InflightKey, Arc<InflightSlot>>> {
|
||||
INFLIGHT.lock().unwrap_or_else(|e| e.into_inner())
|
||||
}
|
||||
|
||||
/// The shared result of one leader's auth attempt, awaited by any joiner that
|
||||
/// arrived while the leader was in flight. A `watch` channel gives us
|
||||
/// publish-once plus wait-for-publish in one primitive: the leader publishes
|
||||
/// exactly once through [`LeaderGuard`]; joiners clone the published result.
|
||||
struct InflightSlot {
|
||||
tx: watch::Sender<Option<Result<String, AuthError>>>,
|
||||
rx: watch::Receiver<Option<Result<String, AuthError>>>,
|
||||
}
|
||||
|
||||
impl InflightSlot {
|
||||
fn new() -> Self {
|
||||
let (tx, rx) = watch::channel(None);
|
||||
Self { tx, rx }
|
||||
}
|
||||
|
||||
/// Block until the leader publishes, then clone out its result.
|
||||
///
|
||||
/// `borrow_and_update` marks the current value seen before awaiting, so a
|
||||
/// publish that lands between the read and the `changed()` await is not a
|
||||
/// lost wakeup — the version has advanced, so `changed()` returns at once.
|
||||
/// A closed channel (leader dropped without publishing — which
|
||||
/// [`LeaderGuard`]'s `Drop` prevents) surfaces as a transient so the caller
|
||||
/// retries rather than hangs.
|
||||
async fn wait(&self) -> Result<String, AuthError> {
|
||||
let mut rx = self.rx.clone();
|
||||
loop {
|
||||
if let Some(result) = rx.borrow_and_update().clone() {
|
||||
return result;
|
||||
}
|
||||
if rx.changed().await.is_err() {
|
||||
return Err(AuthError::NetworkUnavailable);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Publish `result` to every waiting joiner. A send error means no joiners
|
||||
/// remain, which is fine.
|
||||
fn publish(&self, result: Result<String, AuthError>) {
|
||||
let _ = self.tx.send(Some(result));
|
||||
}
|
||||
}
|
||||
|
||||
/// RAII owner of a leader's in-flight slot. Guarantees the slot is evicted from
|
||||
/// [`INFLIGHT`] and a result published to joiners even if the leader future is
|
||||
/// cancelled or panics: a leader that skipped this would leave a dead slot that
|
||||
/// turns every later caller into a joiner of an attempt that never publishes,
|
||||
/// wedging them until `LOCK_WAIT_TIMEOUT`.
|
||||
struct LeaderGuard {
|
||||
key: InflightKey,
|
||||
slot: Arc<InflightSlot>,
|
||||
done: bool,
|
||||
}
|
||||
|
||||
impl LeaderGuard {
|
||||
fn new(key: InflightKey, slot: Arc<InflightSlot>) -> Self {
|
||||
Self {
|
||||
key,
|
||||
slot,
|
||||
done: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Normal completion: evict the slot, publish `result` to joiners, and
|
||||
/// return it to the leader. Evicting *before* publishing means a caller
|
||||
/// arriving after this point starts a fresh attempt (a later explicit
|
||||
/// retry may launch), while joiners already holding the slot still receive
|
||||
/// the result. `Drop` covers the cancel/panic path.
|
||||
fn complete(mut self, result: Result<String, AuthError>) -> Result<String, AuthError> {
|
||||
self.done = true;
|
||||
Self::evict(&self.key, &self.slot);
|
||||
self.slot.publish(result.clone());
|
||||
result
|
||||
}
|
||||
|
||||
/// Remove this leader's slot from the registry, but only if it is still the
|
||||
/// same slot — defends against evicting a successor a later attempt may
|
||||
/// have installed under the same key.
|
||||
fn evict(key: &InflightKey, slot: &Arc<InflightSlot>) {
|
||||
let mut reg = inflight_registry();
|
||||
if reg
|
||||
.get(key)
|
||||
.is_some_and(|existing| Arc::ptr_eq(existing, slot))
|
||||
{
|
||||
reg.remove(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for LeaderGuard {
|
||||
fn drop(&mut self) {
|
||||
if self.done {
|
||||
return;
|
||||
}
|
||||
// Cancelled or panicked before `complete`: evict so later callers start
|
||||
// fresh, and wake joiners with a transient error so they retry rather
|
||||
// than hang on a leader that will never publish.
|
||||
Self::evict(&self.key, &self.slot);
|
||||
self.slot.publish(Err(AuthError::NetworkUnavailable));
|
||||
}
|
||||
}
|
||||
|
||||
/// Load a cached token, enforcing the owner-only invariant on load.
|
||||
///
|
||||
/// Owner-only permissions are a cache *lifecycle* invariant, not just a
|
||||
@@ -1810,6 +2121,44 @@ mod tests {
|
||||
drop(holder);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_lock_timeout_leaves_cooldown_sidecar_byte_for_byte_untouched() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let lock_path = dir.path().join("cache.json.lock");
|
||||
let cooldown_path = dir.path().join("cache.json.cooldown");
|
||||
|
||||
// A pre-existing cooldown sidecar written by an earlier interactive
|
||||
// failure. A waiter that can't take the lock must return before any
|
||||
// code that reads/clears/writes the cooldown, so these exact bytes
|
||||
// survive untouched — otherwise a lock-contended caller could clear a
|
||||
// live suppression and let the next Auto caller re-pop a browser.
|
||||
let original = br#"{"code":"denied","until":9999999999}"#;
|
||||
fs::write(&cooldown_path, original).unwrap();
|
||||
|
||||
// Holder owns the lock (RAII stand-in for another live process).
|
||||
let holder = acquire_auth_lock(&lock_path, Instant::now() + Duration::from_secs(30))
|
||||
.await
|
||||
.expect("holder should acquire the free lock");
|
||||
|
||||
// A waiter past its deadline gives up with LockTimeout — the `?` in
|
||||
// `acquire_leader` propagates this before `acquire_locked` (which owns
|
||||
// every sidecar mutation) is ever entered.
|
||||
let waiter = acquire_auth_lock(&lock_path, Instant::now()).await;
|
||||
assert!(
|
||||
matches!(waiter, Err(AuthError::LockTimeout)),
|
||||
"contended waiter past its deadline must return LockTimeout, got {waiter:?}"
|
||||
);
|
||||
|
||||
let after = fs::read(&cooldown_path).unwrap();
|
||||
assert_eq!(
|
||||
after.as_slice(),
|
||||
original.as_slice(),
|
||||
"a lock timeout must leave the cooldown sidecar byte-for-byte untouched"
|
||||
);
|
||||
|
||||
drop(holder);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_lock_released_on_holder_drop_lets_successor_proceed() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Test-only helper: a real second process that takes the coordinator's
|
||||
//! cross-process advisory lock and holds it until killed.
|
||||
//!
|
||||
//! The auth coordinator single-flights per cache key on an `fs2` advisory lock
|
||||
//! (`flock` on Unix, `LockFileEx` on Windows). To prove the *cross-process*
|
||||
//! contract — a genuine other process serializes the flow, and its death
|
||||
//! releases the lock with no PID files or lock-breaking — a test needs an
|
||||
//! actual separate process on the same lock file, not a second in-process
|
||||
//! handle. This binary is that process.
|
||||
//!
|
||||
//! Driven by two env vars:
|
||||
//! LOCK_HELPER_PATH — the lock file to acquire (the coordinator's
|
||||
//! `<hash>.json.lock`).
|
||||
//! LOCK_HELPER_READY — a marker file created *after* the lock is held, so
|
||||
//! the parent test can synchronize on ownership before
|
||||
//! racing the coordinator.
|
||||
//!
|
||||
//! After signaling readiness it blocks forever; the parent kills it to model a
|
||||
//! crash mid-flow.
|
||||
|
||||
use std::fs;
|
||||
|
||||
use fs2::FileExt;
|
||||
|
||||
fn main() {
|
||||
let lock_path = std::env::var("LOCK_HELPER_PATH").expect("LOCK_HELPER_PATH set");
|
||||
let ready_path = std::env::var("LOCK_HELPER_READY").expect("LOCK_HELPER_READY set");
|
||||
|
||||
if let Some(parent) = std::path::Path::new(&lock_path).parent() {
|
||||
fs::create_dir_all(parent).expect("create lock parent dir");
|
||||
}
|
||||
// Open exactly as the coordinator does so we contend on the same inode.
|
||||
let file = fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.write(true)
|
||||
.open(&lock_path)
|
||||
.expect("open lock file");
|
||||
file.lock_exclusive()
|
||||
.expect("hold the exclusive advisory lock");
|
||||
|
||||
// Signal ownership only once the lock is truly held.
|
||||
fs::write(&ready_path, b"held").expect("write ready marker");
|
||||
|
||||
// Hold the lock until the parent kills us (crash stand-in). The kernel
|
||||
// releases the advisory lock on process death.
|
||||
loop {
|
||||
std::thread::sleep(std::time::Duration::from_secs(3600));
|
||||
}
|
||||
}
|
||||
@@ -17,7 +17,7 @@ use std::io::Write;
|
||||
use std::net::{SocketAddr, TcpStream};
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use axum::extract::Form;
|
||||
use axum::{routing::get, routing::post, Json, Router};
|
||||
@@ -126,10 +126,41 @@ struct Stub {
|
||||
refresh_grants: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
/// How the stub's token endpoint answers a `refresh_token` grant. Lets a test
|
||||
/// distinguish the three ways a refresh can fail so it can assert the
|
||||
/// coordinator classifies each correctly: a `401` is a real credential
|
||||
/// rejection (dead refresh token), a `500` is a transient provider fault, and
|
||||
/// a hang models a slow/unreachable provider that must trip the per-request
|
||||
/// HTTP timeout. Authorization-code grants are never affected.
|
||||
#[derive(Clone, Copy)]
|
||||
enum RefreshMode {
|
||||
/// `200` with a fresh access token.
|
||||
Succeed,
|
||||
/// `401 invalid_grant` — the grant itself is rejected.
|
||||
Reject,
|
||||
/// `500` — a provider-side fault, transient rather than a credential
|
||||
/// decision.
|
||||
ServerError,
|
||||
/// Sleep `d` before answering, so the caller's per-request HTTP timeout
|
||||
/// elapses first (a transport timeout, not a verdict from the provider).
|
||||
Hang(Duration),
|
||||
}
|
||||
|
||||
/// Boot a stub provider. `reject_refresh` makes the token endpoint 401 every
|
||||
/// refresh-token grant (a dead refresh token); authorization-code grants
|
||||
/// always succeed with a fresh token.
|
||||
async fn spawn_stub(reject_refresh: bool) -> Stub {
|
||||
spawn_stub_with(if reject_refresh {
|
||||
RefreshMode::Reject
|
||||
} else {
|
||||
RefreshMode::Succeed
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Boot a stub provider whose refresh-token grant follows `mode`. Discovery and
|
||||
/// authorization-code grants always succeed instantly regardless of `mode`.
|
||||
async fn spawn_stub_with(mode: RefreshMode) -> Stub {
|
||||
let code_grants = Arc::new(AtomicU64::new(0));
|
||||
let refresh_grants = Arc::new(AtomicU64::new(0));
|
||||
|
||||
@@ -161,24 +192,34 @@ async fn spawn_stub(reject_refresh: bool) -> Stub {
|
||||
post(move |Form(form): Form<TokenForm>| {
|
||||
let code_grants = code_for_token.clone();
|
||||
let refresh_grants = refresh_for_token.clone();
|
||||
let reject_refresh = reject_refresh;
|
||||
let mode = mode;
|
||||
async move {
|
||||
if form.grant_type == "refresh_token" {
|
||||
let n = refresh_grants.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
if reject_refresh {
|
||||
return (
|
||||
// A hang delays the answer so the caller's per-request
|
||||
// HTTP timeout can elapse first (transport timeout, not
|
||||
// a credential decision).
|
||||
if let RefreshMode::Hang(d) = mode {
|
||||
tokio::time::sleep(d).await;
|
||||
}
|
||||
return match mode {
|
||||
RefreshMode::Reject => (
|
||||
axum::http::StatusCode::UNAUTHORIZED,
|
||||
Json(json!({ "error": "invalid_grant" })),
|
||||
);
|
||||
}
|
||||
return (
|
||||
axum::http::StatusCode::OK,
|
||||
Json(json!({
|
||||
"access_token": format!("refreshed-token-{n}"),
|
||||
"refresh_token": "rotated-refresh",
|
||||
"expires_in": 3600,
|
||||
})),
|
||||
);
|
||||
),
|
||||
RefreshMode::ServerError => (
|
||||
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({ "error": "temporarily_unavailable" })),
|
||||
),
|
||||
RefreshMode::Succeed | RefreshMode::Hang(_) => (
|
||||
axum::http::StatusCode::OK,
|
||||
Json(json!({
|
||||
"access_token": format!("refreshed-token-{n}"),
|
||||
"refresh_token": "rotated-refresh",
|
||||
"expires_in": 3600,
|
||||
})),
|
||||
),
|
||||
};
|
||||
}
|
||||
let n = code_grants.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
(
|
||||
@@ -222,7 +263,7 @@ fn future_secs() -> u64 {
|
||||
+ 3600
|
||||
}
|
||||
|
||||
fn seed_cache(cfg: &PkceOAuthConfig, cache_dir: &std::path::Path, body: serde_json::Value) {
|
||||
fn cache_file_path(cfg: &PkceOAuthConfig, cache_dir: &std::path::Path) -> std::path::PathBuf {
|
||||
use sha2::Digest;
|
||||
let mut h = sha2::Sha256::new();
|
||||
h.update(cfg.discovery_url.as_bytes());
|
||||
@@ -231,9 +272,23 @@ fn seed_cache(cfg: &PkceOAuthConfig, cache_dir: &std::path::Path, body: serde_js
|
||||
h.update(b"|");
|
||||
h.update(cfg.scopes.join(",").as_bytes());
|
||||
let hash = hex::encode(h.finalize());
|
||||
let path = cache_dir
|
||||
cache_dir
|
||||
.join(&cfg.cache_namespace)
|
||||
.join(format!("{hash}.json"));
|
||||
.join(format!("{hash}.json"))
|
||||
}
|
||||
|
||||
/// The cross-process advisory lock path for a config, matching the
|
||||
/// coordinator's `append_ext(cache_path, "lock")`. Used to point the
|
||||
/// out-of-process lock-holder helper at the exact file the coordinator
|
||||
/// contends on.
|
||||
fn lock_file_path(cfg: &PkceOAuthConfig, cache_dir: &std::path::Path) -> std::path::PathBuf {
|
||||
let mut p = cache_file_path(cfg, cache_dir).into_os_string();
|
||||
p.push(".lock");
|
||||
p.into()
|
||||
}
|
||||
|
||||
fn seed_cache(cfg: &PkceOAuthConfig, cache_dir: &std::path::Path, body: serde_json::Value) {
|
||||
let path = cache_file_path(cfg, cache_dir);
|
||||
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
|
||||
std::fs::write(&path, serde_json::to_vec(&body).unwrap()).unwrap();
|
||||
}
|
||||
@@ -259,8 +314,8 @@ async fn test_same_key_concurrent_callers_share_one_browser_attempt() {
|
||||
.unwrap();
|
||||
|
||||
let (ra, rb) = tokio::join!(
|
||||
a.acquire_with_intent(AuthIntent::Auto),
|
||||
b.acquire_with_intent(AuthIntent::Auto),
|
||||
a.acquire_with_intent(AuthIntent::Auto, None),
|
||||
b.acquire_with_intent(AuthIntent::Auto, None),
|
||||
);
|
||||
let ta = ra.expect("first caller authenticates");
|
||||
let tb = rb.expect("second caller authenticates");
|
||||
@@ -292,7 +347,7 @@ async fn test_denied_then_auto_reads_cooldown_without_second_launch() {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let first = src.acquire_with_intent(AuthIntent::Auto).await;
|
||||
let first = src.acquire_with_intent(AuthIntent::Auto, None).await;
|
||||
assert_eq!(
|
||||
first,
|
||||
Err(AuthError::Denied),
|
||||
@@ -302,7 +357,7 @@ async fn test_denied_then_auto_reads_cooldown_without_second_launch() {
|
||||
|
||||
// The denial wrote a cooldown; a subsequent Auto caller reads it and
|
||||
// returns the recorded outcome instead of popping a second browser.
|
||||
let second = src.acquire_with_intent(AuthIntent::Auto).await;
|
||||
let second = src.acquire_with_intent(AuthIntent::Auto, None).await;
|
||||
assert_eq!(
|
||||
second,
|
||||
Err(AuthError::Denied),
|
||||
@@ -327,7 +382,9 @@ async fn test_userinitiated_denial_is_visible_to_crossprocess_auto() {
|
||||
Arc::new(opener.clone()),
|
||||
)
|
||||
.unwrap();
|
||||
let denied = proc_a.acquire_with_intent(AuthIntent::UserInitiated).await;
|
||||
let denied = proc_a
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await;
|
||||
assert_eq!(denied, Err(AuthError::Denied));
|
||||
assert_eq!(opener.call_count(), 1);
|
||||
|
||||
@@ -340,7 +397,7 @@ async fn test_userinitiated_denial_is_visible_to_crossprocess_auto() {
|
||||
Arc::new(opener.clone()),
|
||||
)
|
||||
.unwrap();
|
||||
let auto = proc_b.acquire_with_intent(AuthIntent::Auto).await;
|
||||
let auto = proc_b.acquire_with_intent(AuthIntent::Auto, None).await;
|
||||
assert_eq!(
|
||||
auto,
|
||||
Err(AuthError::Denied),
|
||||
@@ -366,7 +423,9 @@ async fn test_userinitiated_retry_bypasses_cooldown_and_reopens() {
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
denier.acquire_with_intent(AuthIntent::UserInitiated).await,
|
||||
denier
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await,
|
||||
Err(AuthError::Denied)
|
||||
);
|
||||
|
||||
@@ -379,7 +438,7 @@ async fn test_userinitiated_retry_bypasses_cooldown_and_reopens() {
|
||||
)
|
||||
.unwrap();
|
||||
let token = retrier
|
||||
.acquire_with_intent(AuthIntent::UserInitiated)
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await
|
||||
.expect("explicit retry re-launches the browser and succeeds");
|
||||
assert_eq!(token, "browser-token-1");
|
||||
@@ -391,7 +450,7 @@ async fn test_userinitiated_retry_bypasses_cooldown_and_reopens() {
|
||||
|
||||
// Cooldown cleared on success: a follow-up Auto now sees a valid token,
|
||||
// never the stale denial.
|
||||
let auto = retrier.acquire_with_intent(AuthIntent::Auto).await;
|
||||
let auto = retrier.acquire_with_intent(AuthIntent::Auto, None).await;
|
||||
assert_eq!(auto, Ok("browser-token-1".to_string()));
|
||||
}
|
||||
|
||||
@@ -408,7 +467,7 @@ async fn test_distinct_hosts_do_not_inherit_cooldown() {
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
host_a.acquire_with_intent(AuthIntent::Auto).await,
|
||||
host_a.acquire_with_intent(AuthIntent::Auto, None).await,
|
||||
Err(AuthError::Denied)
|
||||
);
|
||||
|
||||
@@ -421,7 +480,7 @@ async fn test_distinct_hosts_do_not_inherit_cooldown() {
|
||||
)
|
||||
.unwrap();
|
||||
let token = host_b
|
||||
.acquire_with_intent(AuthIntent::Auto)
|
||||
.acquire_with_intent(AuthIntent::Auto, None)
|
||||
.await
|
||||
.expect("distinct host is unaffected by another key's cooldown");
|
||||
assert_eq!(token, "browser-token-1");
|
||||
@@ -441,7 +500,9 @@ async fn test_browser_open_failure_is_typed_and_retryable_by_user() {
|
||||
Arc::new(fail_opener.clone()),
|
||||
)
|
||||
.unwrap();
|
||||
let result = failing.acquire_with_intent(AuthIntent::UserInitiated).await;
|
||||
let result = failing
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await;
|
||||
assert_eq!(
|
||||
result,
|
||||
Err(AuthError::BrowserOpenFailed),
|
||||
@@ -459,7 +520,7 @@ async fn test_browser_open_failure_is_typed_and_retryable_by_user() {
|
||||
)
|
||||
.unwrap();
|
||||
let token = retrier
|
||||
.acquire_with_intent(AuthIntent::UserInitiated)
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await
|
||||
.expect("explicit retry reopens despite the prior launch failure");
|
||||
assert_eq!(token, "browser-token-1");
|
||||
@@ -487,7 +548,7 @@ async fn test_headless_dead_refresh_returns_refresh_rejected_without_browser() {
|
||||
);
|
||||
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let result = src.acquire_with_intent(AuthIntent::Headless).await;
|
||||
let result = src.acquire_with_intent(AuthIntent::Headless, None).await;
|
||||
assert_eq!(
|
||||
result,
|
||||
Err(AuthError::RefreshRejected),
|
||||
@@ -522,7 +583,7 @@ async fn test_interactive_dead_refresh_converts_to_browser() {
|
||||
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let token = src
|
||||
.acquire_with_intent(AuthIntent::UserInitiated)
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await
|
||||
.expect("interactive intent recovers via the browser");
|
||||
assert_eq!(token, "browser-token-1");
|
||||
@@ -550,7 +611,7 @@ async fn test_headless_expired_token_live_refresh_recovers_silently() {
|
||||
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let token = src
|
||||
.acquire_with_intent(AuthIntent::Headless)
|
||||
.acquire_with_intent(AuthIntent::Headless, None)
|
||||
.await
|
||||
.expect("live refresh recovers a Headless caller silently");
|
||||
assert_eq!(token, "refreshed-token-1");
|
||||
@@ -587,3 +648,321 @@ async fn test_interactive_login_reuses_valid_cache_without_browser() {
|
||||
"a valid cached token means no browser prompt"
|
||||
);
|
||||
}
|
||||
|
||||
// ---- locally-fresh rejected bearer (401) recovery ------------------------
|
||||
//
|
||||
// The saved-model picker's recovery path: model discovery 401s a bearer that
|
||||
// still looks locally fresh (its `expires_at` is in the future) and whose
|
||||
// refresh grant is dead. Passing that exact token as `rejected` makes the
|
||||
// clock untrustworthy, so the acquisition must not short-circuit on the fresh
|
||||
// cache. `Auto` and `UserInitiated` then convert to a browser; `Headless`
|
||||
// stays terminal with `RefreshRejected`. Seeding a *future*-expiry token is
|
||||
// what distinguishes this from the expired-token refresh path.
|
||||
|
||||
/// Seed a not-yet-expired access token with a (dead) refresh token and return
|
||||
/// the access token so the caller can pass it as `rejected`.
|
||||
fn seed_fresh_rejectable(cfg: &PkceOAuthConfig, cache_dir: &std::path::Path) -> String {
|
||||
let access = "fresh-but-rejected";
|
||||
seed_cache(
|
||||
cfg,
|
||||
cache_dir,
|
||||
json!({
|
||||
"access_token": access,
|
||||
"refresh_token": "dead-refresh",
|
||||
"expires_at": future_secs(),
|
||||
}),
|
||||
);
|
||||
access.to_string()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_auto_rejected_fresh_bearer_with_dead_refresh_launches_browser() {
|
||||
let stub = spawn_stub(true).await; // refresh grants 401
|
||||
let cache = TempDir::new().unwrap();
|
||||
let opener = ScriptedOpener::new(Script::Approve);
|
||||
let cfg = config(&stub, "/disco/a", cache.path());
|
||||
let rejected = seed_fresh_rejectable(&cfg, cache.path());
|
||||
|
||||
// The token is locally fresh, so without `rejected` it would be a cache
|
||||
// hit and never reach the browser. Passing it as rejected forces the
|
||||
// clock-based hit to fail, the dead refresh to be attempted, and an Auto
|
||||
// caller to fall through to the browser.
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let token = src
|
||||
.acquire_with_intent(AuthIntent::Auto, Some(&rejected))
|
||||
.await
|
||||
.expect("Auto recovers a rejected-but-fresh bearer via the browser");
|
||||
assert_eq!(token, "browser-token-1");
|
||||
assert_eq!(opener.call_count(), 1, "Auto launches a browser to recover");
|
||||
assert_eq!(stub.refresh_grants.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(stub.code_grants.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_userinitiated_rejected_fresh_bearer_with_dead_refresh_launches_browser() {
|
||||
let stub = spawn_stub(true).await; // refresh grants 401
|
||||
let cache = TempDir::new().unwrap();
|
||||
let opener = ScriptedOpener::new(Script::Approve);
|
||||
let cfg = config(&stub, "/disco/a", cache.path());
|
||||
let rejected = seed_fresh_rejectable(&cfg, cache.path());
|
||||
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let token = src
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, Some(&rejected))
|
||||
.await
|
||||
.expect("UserInitiated recovers a rejected-but-fresh bearer via the browser");
|
||||
assert_eq!(token, "browser-token-1");
|
||||
assert_eq!(
|
||||
opener.call_count(),
|
||||
1,
|
||||
"UserInitiated launches a browser to recover"
|
||||
);
|
||||
assert_eq!(stub.refresh_grants.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(stub.code_grants.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_headless_rejected_fresh_bearer_with_dead_refresh_returns_refresh_rejected() {
|
||||
let stub = spawn_stub(true).await; // refresh grants 401
|
||||
let cache = TempDir::new().unwrap();
|
||||
let opener = ScriptedOpener::new(Script::Approve);
|
||||
let cfg = config(&stub, "/disco/a", cache.path());
|
||||
let rejected = seed_fresh_rejectable(&cfg, cache.path());
|
||||
|
||||
// Same locally-fresh rejected seed, but a Headless caller cannot open a
|
||||
// browser: a dead refresh is terminal RefreshRejected, never a launch.
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let result = src
|
||||
.acquire_with_intent(AuthIntent::Headless, Some(&rejected))
|
||||
.await;
|
||||
assert_eq!(
|
||||
result,
|
||||
Err(AuthError::RefreshRejected),
|
||||
"Headless dead-refresh on a rejected fresh bearer is terminal"
|
||||
);
|
||||
assert_eq!(opener.call_count(), 0, "Headless never opens a browser");
|
||||
assert_eq!(stub.refresh_grants.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
// ---- refresh transport failures are not credential rejections ------------
|
||||
//
|
||||
// A refresh that never gets a verdict from the token endpoint — a per-request
|
||||
// timeout, or a 5xx — is infrastructural, not a dead credential. It must
|
||||
// surface as `NetworkUnavailable` and never pop a browser or return
|
||||
// `RefreshRejected`, which would misreport a transient fault as a rotated
|
||||
// token and (for interactive intents) prompt a needless sign-in.
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_timeout_is_network_unavailable_not_rejected() {
|
||||
// The token endpoint hangs far longer than the injected per-request HTTP
|
||||
// timeout, so the refresh call times out at the transport layer with no
|
||||
// verdict from the provider. A short real-time timeout is injected rather
|
||||
// than pausing the clock: under `start_paused` tokio auto-advances into
|
||||
// the timer while the real loopback discovery GET is still in flight, so
|
||||
// discovery — not the refresh — would trip the timeout, and the refresh
|
||||
// would never even be attempted. Real time keeps the timeout attached to
|
||||
// the request that actually hangs, which the `refresh_grants == 1` guard
|
||||
// below proves.
|
||||
let stub = spawn_stub_with(RefreshMode::Hang(Duration::from_secs(30))).await;
|
||||
let cache = TempDir::new().unwrap();
|
||||
let opener = ScriptedOpener::new(Script::Approve);
|
||||
let cfg = config(&stub, "/disco/a", cache.path());
|
||||
|
||||
// Expired token with a refresh token: the coordinator attempts the refresh,
|
||||
// which hangs past the HTTP timeout. A Headless caller must classify the
|
||||
// timeout as NetworkUnavailable, not RefreshRejected.
|
||||
seed_cache(
|
||||
&cfg,
|
||||
cache.path(),
|
||||
json!({
|
||||
"access_token": "stale",
|
||||
"refresh_token": "slow-refresh",
|
||||
"expires_at": 1u64,
|
||||
}),
|
||||
);
|
||||
|
||||
let src = PkceOAuthTokenSource::new_with_http_timeout(
|
||||
cfg,
|
||||
Arc::new(opener.clone()),
|
||||
Duration::from_millis(300),
|
||||
)
|
||||
.unwrap();
|
||||
let result = src.acquire_with_intent(AuthIntent::Headless, None).await;
|
||||
assert_eq!(
|
||||
result,
|
||||
Err(AuthError::NetworkUnavailable),
|
||||
"a refresh transport timeout is infrastructural, not a rejection"
|
||||
);
|
||||
assert_eq!(
|
||||
opener.call_count(),
|
||||
0,
|
||||
"a timed-out refresh never becomes a credential decision"
|
||||
);
|
||||
assert_eq!(
|
||||
stub.refresh_grants.load(Ordering::SeqCst),
|
||||
1,
|
||||
"the refresh was attempted exactly once before timing out"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_server_error_is_network_unavailable_not_rejected() {
|
||||
let stub = spawn_stub_with(RefreshMode::ServerError).await; // refresh 500s
|
||||
let cache = TempDir::new().unwrap();
|
||||
let opener = ScriptedOpener::new(Script::Approve);
|
||||
let cfg = config(&stub, "/disco/a", cache.path());
|
||||
|
||||
// A 5xx is a provider-side fault, not a grant rejection: an interactive
|
||||
// intent must NOT pop a browser off it, and it must surface as
|
||||
// NetworkUnavailable rather than RefreshRejected.
|
||||
seed_cache(
|
||||
&cfg,
|
||||
cache.path(),
|
||||
json!({
|
||||
"access_token": "stale",
|
||||
"refresh_token": "server-error-refresh",
|
||||
"expires_at": 1u64,
|
||||
}),
|
||||
);
|
||||
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let result = src
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await;
|
||||
assert_eq!(
|
||||
result,
|
||||
Err(AuthError::NetworkUnavailable),
|
||||
"a refresh 5xx is transient, not a credential rejection"
|
||||
);
|
||||
assert_eq!(
|
||||
opener.call_count(),
|
||||
0,
|
||||
"a provider 5xx must not trigger an interactive browser fallback"
|
||||
);
|
||||
assert_eq!(stub.refresh_grants.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
// ---- in-process joiner shares the leader's FAILURE result ----------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_two_concurrent_userinitiated_denials_share_one_browser() {
|
||||
// Two UserInitiated callers arrive together on one key. The first is the
|
||||
// leader and opens the browser; the second is a pre-existing joiner that
|
||||
// must receive the leader's SAME Denied result rather than acquire the
|
||||
// lock afterward, clear the cooldown, and pop a second browser. This is
|
||||
// the failure-sharing that a lock-alone protocol loses.
|
||||
let stub = spawn_stub(false).await;
|
||||
let cache = TempDir::new().unwrap();
|
||||
let opener = ScriptedOpener::new(Script::Deny);
|
||||
|
||||
let a = PkceOAuthTokenSource::new_with(
|
||||
config(&stub, "/disco/a", cache.path()),
|
||||
Arc::new(opener.clone()),
|
||||
)
|
||||
.unwrap();
|
||||
let b = PkceOAuthTokenSource::new_with(
|
||||
config(&stub, "/disco/a", cache.path()),
|
||||
Arc::new(opener.clone()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let (ra, rb) = tokio::join!(
|
||||
a.acquire_with_intent(AuthIntent::UserInitiated, None),
|
||||
b.acquire_with_intent(AuthIntent::UserInitiated, None),
|
||||
);
|
||||
assert_eq!(ra, Err(AuthError::Denied), "leader observes the denial");
|
||||
assert_eq!(
|
||||
rb,
|
||||
Err(AuthError::Denied),
|
||||
"the joiner shares the leader's denial, not a fresh attempt"
|
||||
);
|
||||
assert_eq!(
|
||||
opener.call_count(),
|
||||
1,
|
||||
"one browser launch shared across both concurrent UserInitiated callers"
|
||||
);
|
||||
}
|
||||
|
||||
// ---- genuine cross-process lock contention and crash release -------------
|
||||
//
|
||||
// The single-flight guarantee and its crash-release property are cross-process
|
||||
// claims, so they need a real second process — not a second in-process handle —
|
||||
// on the same lock file. The `lock-holder` helper binary takes the
|
||||
// coordinator's advisory lock and holds it until killed; killing it models a
|
||||
// crash mid-flow, and the kernel's release of the advisory lock is what lets
|
||||
// the coordinator's successor proceed with no PID files and no lock breaking.
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_crossprocess_lock_holder_blocks_then_crash_release_lets_successor_proceed() {
|
||||
let stub = spawn_stub(false).await; // refresh succeeds once the lock is free
|
||||
let cache = TempDir::new().unwrap();
|
||||
let opener = ScriptedOpener::new(Script::Approve);
|
||||
let cfg = config(&stub, "/disco/a", cache.path());
|
||||
|
||||
// Expired token with a LIVE refresh: a cache miss forces the coordinator
|
||||
// onto the slow path (it must take the lock), and once the lock is free the
|
||||
// refresh recovers a token without any browser — so success is a clean
|
||||
// signal that the successor proceeded.
|
||||
seed_cache(
|
||||
&cfg,
|
||||
cache.path(),
|
||||
json!({
|
||||
"access_token": "stale",
|
||||
"refresh_token": "live-refresh",
|
||||
"expires_at": 1u64,
|
||||
}),
|
||||
);
|
||||
|
||||
let lock_path = lock_file_path(&cfg, cache.path());
|
||||
let ready_marker = cache.path().join("holder.ready");
|
||||
|
||||
// A real second process grabs the lock and holds it.
|
||||
let mut holder = tokio::process::Command::new(env!("CARGO_BIN_EXE_lock-holder"))
|
||||
.env("LOCK_HELPER_PATH", &lock_path)
|
||||
.env("LOCK_HELPER_READY", &ready_marker)
|
||||
.kill_on_drop(true)
|
||||
.spawn()
|
||||
.expect("spawn the lock-holder helper process");
|
||||
|
||||
// Synchronize on real lock ownership before racing the coordinator.
|
||||
for _ in 0..600 {
|
||||
if ready_marker.exists() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
assert!(
|
||||
ready_marker.exists(),
|
||||
"lock-holder never signaled that it holds the lock"
|
||||
);
|
||||
|
||||
// The coordinator cannot make progress while another process holds the
|
||||
// lock: it polls the advisory lock rather than stealing it.
|
||||
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap();
|
||||
let task =
|
||||
tokio::spawn(async move { src.acquire_with_intent(AuthIntent::Headless, None).await });
|
||||
tokio::time::sleep(Duration::from_millis(400)).await;
|
||||
assert!(
|
||||
!task.is_finished(),
|
||||
"coordinator must block while a live process holds the cross-process lock"
|
||||
);
|
||||
|
||||
// Kill the holder: the kernel releases the advisory lock on process death,
|
||||
// with no PID file inspection or lock breaking on our side.
|
||||
holder.kill().await.expect("kill the lock holder");
|
||||
holder.wait().await.ok();
|
||||
|
||||
let token = task
|
||||
.await
|
||||
.expect("acquisition task joins")
|
||||
.expect("successor proceeds once the crashed holder's lock is released");
|
||||
assert_eq!(
|
||||
token, "refreshed-token-1",
|
||||
"successor completes the refresh after acquiring the freed lock"
|
||||
);
|
||||
assert_eq!(
|
||||
opener.call_count(),
|
||||
0,
|
||||
"Headless successor recovers via refresh without a browser"
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user