483 lines
18 KiB
Rust
483 lines
18 KiB
Rust
use std::sync::{Arc, RwLock};
|
|
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
|
|
|
use anyhow::Result;
|
|
use base64::Engine;
|
|
use blocking::unblock;
|
|
use galaxy_core::channel::IapConfig;
|
|
use galaxyui_core::r#async::{BoxFuture, FutureExt as _, Timer};
|
|
use galaxyui_core::{AppContext, Entity, ModelContext, SingletonEntity};
|
|
use instant::Instant;
|
|
#[cfg(not(target_family = "wasm"))]
|
|
use websocket::connect_error_http_response;
|
|
|
|
const PROACTIVE_REFRESH_BUFFER: Duration = Duration::from_secs(5 * 60);
|
|
const INJECTED_TOKEN_ENV_VAR: &str = "WARP_IAP_TOKEN";
|
|
|
|
const BASE_FAILURE_RETRY_DELAY: Duration = Duration::from_secs(30);
|
|
const MAX_FAILURE_RETRY_DELAY: Duration = Duration::from_secs(5 * 60);
|
|
/// Maximum number of consecutive failed fetches to automatically retry
|
|
/// before giving up and waiting for a manual Refresh or an inbound
|
|
/// IAP challenge. i.e. so a persistently broken setup (no gcloud,
|
|
/// bad credentials) doesn't loop forever.
|
|
const MAX_FAILURE_RETRIES: u32 = 5;
|
|
|
|
pub type PathResolver = Box<dyn Fn(&mut AppContext) -> BoxFuture<'static, Option<String>>>;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct CachedToken {
|
|
pub token: String,
|
|
pub expires_at: Instant,
|
|
}
|
|
|
|
impl CachedToken {
|
|
fn valid_token(&self) -> Option<String> {
|
|
(self.expires_at > Instant::now()).then(|| self.token.clone())
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum IapCredentialsState {
|
|
Missing,
|
|
/// A credential fetch is in progress. `previous` carries the last
|
|
/// successfully-loaded token (if any). Allows us to attach it to
|
|
/// outbound requests while we're refreshing so that proactive refreshes
|
|
/// (i.e. refresh the token 5min before exp) don't prevent active requests.
|
|
Refreshing {
|
|
previous: Option<CachedToken>,
|
|
},
|
|
Loaded(CachedToken),
|
|
Failed {
|
|
message: String,
|
|
// in case the last token still works... we can try to use that for a couple more mins
|
|
previous: Option<CachedToken>,
|
|
},
|
|
/// Represents a terminal state in the iap creds state machine.
|
|
/// The gcloud refresh loop will never run, and an IAP challenge is logged
|
|
/// rather than triggering a refresh (we have no way to refresh a new token
|
|
/// from ambient agent context yet).
|
|
/// TODO(Isaiah/Jason): implement token refreshing scheme.
|
|
/// see: https://linear.app/warpdotdev/issue/REMOTE-1370/refresh-github-token
|
|
EnvInjected {
|
|
token: String,
|
|
},
|
|
}
|
|
|
|
impl IapCredentialsState {
|
|
fn previous_token(&self) -> Option<CachedToken> {
|
|
match self {
|
|
IapCredentialsState::Loaded(cached) => Some(cached.clone()),
|
|
IapCredentialsState::Refreshing { previous }
|
|
| IapCredentialsState::Failed { previous, .. } => previous.clone(),
|
|
IapCredentialsState::EnvInjected { .. } | IapCredentialsState::Missing => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct IapState {
|
|
audiences: String,
|
|
service_account_email: String,
|
|
inner: RwLock<IapCredentialsState>,
|
|
}
|
|
|
|
impl IapState {
|
|
pub fn new(config: &IapConfig) -> Self {
|
|
let initial = std::env::var(INJECTED_TOKEN_ENV_VAR)
|
|
.ok()
|
|
.filter(|s| !s.is_empty())
|
|
.map(|token| IapCredentialsState::EnvInjected { token })
|
|
.unwrap_or(IapCredentialsState::Missing);
|
|
Self {
|
|
audiences: config.audiences.to_string(),
|
|
service_account_email: config.service_account_email.to_string(),
|
|
inner: RwLock::new(initial),
|
|
}
|
|
}
|
|
|
|
pub fn get_cached(&self) -> Option<String> {
|
|
match &*self.inner.read().expect("IAP state lock poisoned") {
|
|
// Gate on expiry even while `Loaded`: if a proactive refresh is
|
|
// delayed (e.g. the machine slept across the refresh window), the
|
|
// token may already be expired, and attaching it would guarantee an
|
|
// IAP challenge. Returning `None` lets the caller proceed without a
|
|
// doomed token while the reactive refresh recovers.
|
|
IapCredentialsState::Loaded(cached) => cached.valid_token(),
|
|
IapCredentialsState::EnvInjected { token } => Some(token.clone()),
|
|
IapCredentialsState::Refreshing { previous }
|
|
| IapCredentialsState::Failed { previous, .. } => {
|
|
previous.as_ref().and_then(CachedToken::valid_token)
|
|
}
|
|
IapCredentialsState::Missing => None,
|
|
}
|
|
}
|
|
|
|
pub fn proxy_auth_header(&self) -> Option<(&'static str, String)> {
|
|
self.get_cached()
|
|
.map(|token| http_client::iap::proxy_auth_header(&token))
|
|
}
|
|
|
|
pub fn state(&self) -> IapCredentialsState {
|
|
self.inner.read().expect("IAP state lock poisoned").clone()
|
|
}
|
|
|
|
pub fn audiences(&self) -> &str {
|
|
&self.audiences
|
|
}
|
|
|
|
pub fn service_account_email(&self) -> &str {
|
|
&self.service_account_email
|
|
}
|
|
|
|
fn set_refreshing(&self) {
|
|
let mut state = self.inner.write().expect("IAP state lock poisoned");
|
|
*state = IapCredentialsState::Refreshing {
|
|
previous: state.previous_token(),
|
|
};
|
|
}
|
|
|
|
fn set_loaded(&self, cached: CachedToken) {
|
|
*self.inner.write().expect("IAP state lock poisoned") = IapCredentialsState::Loaded(cached);
|
|
}
|
|
|
|
fn set_failed(&self, message: String) {
|
|
let mut state = self.inner.write().expect("IAP state lock poisoned");
|
|
*state = IapCredentialsState::Failed {
|
|
message,
|
|
previous: state.previous_token(),
|
|
};
|
|
}
|
|
}
|
|
|
|
impl http_client::iap::IapTokenProvider for IapState {
|
|
fn cached_token(&self) -> Option<String> {
|
|
self.get_cached()
|
|
}
|
|
}
|
|
|
|
/// Owns the IAP refresh lifecycle: initial fetch, proactive time-based
|
|
/// refresh, and reactive refresh on challenge events.
|
|
pub struct IapManager {
|
|
state: Option<Arc<IapState>>,
|
|
path_resolver: PathResolver,
|
|
/// Number of consecutive failed fetches since the last success.
|
|
consecutive_failures: u32,
|
|
}
|
|
|
|
pub enum IapManagerEvent {
|
|
StateChanged,
|
|
RefreshFailed {
|
|
/// A human-readable error message describing why the refresh failed.
|
|
message: String,
|
|
/// Whether this is the first failure in a streak of failures.
|
|
is_first_failure_of_streak: bool,
|
|
},
|
|
}
|
|
|
|
impl IapManager {
|
|
pub fn new(
|
|
state: Option<Arc<IapState>>,
|
|
path_resolver: PathResolver,
|
|
ctx: &mut ModelContext<Self>,
|
|
) -> Self {
|
|
let mut manager = Self {
|
|
state,
|
|
path_resolver,
|
|
consecutive_failures: 0,
|
|
};
|
|
manager.start_refresh(ctx);
|
|
manager
|
|
}
|
|
|
|
/// Returns `true` if IAP is active for this build. When `false`, all
|
|
/// other methods on this type are no-ops.
|
|
pub fn is_enabled(&self) -> bool {
|
|
self.state.is_some()
|
|
}
|
|
|
|
pub fn state(&self) -> Option<IapCredentialsState> {
|
|
self.state.as_ref().map(|s| s.state())
|
|
}
|
|
|
|
/// Returns a handle to the shared IAP credential state, if IAP is active.
|
|
/// Mirrors how `AuthStateProvider` hands out the `Arc<AuthState>`, letting
|
|
/// callers read cached credentials (e.g. to build a proxy-auth header) off
|
|
/// a `ModelContext` without reaching through `ServerApi`.
|
|
pub fn iap_state(&self) -> Option<Arc<IapState>> {
|
|
self.state.clone()
|
|
}
|
|
|
|
pub fn handle_challenge(&mut self, ctx: &mut ModelContext<Self>) {
|
|
let Some(state) = self.state.as_ref() else {
|
|
return;
|
|
};
|
|
if matches!(state.state(), IapCredentialsState::EnvInjected { .. }) {
|
|
log::warn!(
|
|
"Env-injected IAP token ({INJECTED_TOKEN_ENV_VAR}) was rejected by IAP; \
|
|
token is likely stale — re-inject to recover"
|
|
);
|
|
return;
|
|
}
|
|
self.consecutive_failures = 0;
|
|
self.start_refresh(ctx);
|
|
}
|
|
|
|
pub fn start_refresh(&mut self, ctx: &mut ModelContext<Self>) {
|
|
let Some(state) = self.state.clone() else {
|
|
return;
|
|
};
|
|
// Don't touch state if a refresh is already running, or if we're
|
|
// in the terminal env-injected state (no refresh path exists).
|
|
if matches!(
|
|
state.state(),
|
|
IapCredentialsState::Refreshing { .. } | IapCredentialsState::EnvInjected { .. }
|
|
) {
|
|
return;
|
|
}
|
|
state.set_refreshing();
|
|
ctx.emit(IapManagerEvent::StateChanged);
|
|
ctx.notify();
|
|
|
|
let audiences = state.audiences().to_string();
|
|
let service_account_email = state.service_account_email().to_string();
|
|
|
|
// Make `gcloud` findable even when Warp is launched from the macOS GUI
|
|
// (i.e. in environments without something like `~/.zshrc && WarpDev` happening to init cli path)
|
|
let path_future = (self.path_resolver)(ctx);
|
|
|
|
ctx.spawn(
|
|
async move {
|
|
// Bound the interactive PATH capture. It spawns an interactive
|
|
// login shell (sourcing rc files), which can hang indefinitely
|
|
// on a misbehaving startup script. Without this bound the
|
|
// spawned task would never reach the `GCLOUD_TIMEOUT`-guarded
|
|
// fetch, stranding the state machine in `Refreshing` and
|
|
// silently disabling every future refresh and IAP challenge
|
|
// (both early-return while `Refreshing`). On timeout, fall back
|
|
// to the ambient PATH so the fetch still runs and the state
|
|
// machine can make progress (succeed or fail).
|
|
const PATH_CAPTURE_TIMEOUT: Duration = Duration::from_secs(10);
|
|
let path_env = match path_future.with_timeout(PATH_CAPTURE_TIMEOUT).await {
|
|
Ok(path_env) => path_env,
|
|
Err(_) => {
|
|
log::warn!(
|
|
"Interactive PATH capture timed out after {}s; \
|
|
falling back to ambient PATH for IAP token fetch",
|
|
PATH_CAPTURE_TIMEOUT.as_secs()
|
|
);
|
|
None
|
|
}
|
|
};
|
|
unblock(move || {
|
|
fetch_iap_token(&audiences, &service_account_email, path_env.as_deref())
|
|
})
|
|
.await
|
|
},
|
|
move |manager, result, ctx| {
|
|
let Some(state) = manager.state.as_ref() else {
|
|
return;
|
|
};
|
|
match result {
|
|
Ok(cached) => {
|
|
let expires_at = cached.expires_at;
|
|
state.set_loaded(cached);
|
|
manager.consecutive_failures = 0;
|
|
log::info!("Warp Staging IAP token refreshed");
|
|
ctx.emit(IapManagerEvent::StateChanged);
|
|
ctx.notify();
|
|
manager.schedule_next_refresh(expires_at, ctx);
|
|
}
|
|
Err(err) => {
|
|
let message = format!("{err:#}");
|
|
log::warn!("Warp Staging IAP token fetch failed: {message}");
|
|
let is_first_failure_of_streak = manager.consecutive_failures == 0;
|
|
state.set_failed(message.clone());
|
|
ctx.emit(IapManagerEvent::RefreshFailed {
|
|
message,
|
|
is_first_failure_of_streak,
|
|
});
|
|
ctx.emit(IapManagerEvent::StateChanged);
|
|
ctx.notify();
|
|
manager.schedule_failure_retry(ctx);
|
|
}
|
|
}
|
|
},
|
|
);
|
|
}
|
|
|
|
fn schedule_next_refresh(&mut self, expires_at: Instant, ctx: &mut ModelContext<Self>) {
|
|
let sleep_duration = expires_at
|
|
.saturating_duration_since(Instant::now())
|
|
.saturating_sub(PROACTIVE_REFRESH_BUFFER);
|
|
self.schedule_retry(sleep_duration, ctx);
|
|
}
|
|
|
|
fn schedule_failure_retry(&mut self, ctx: &mut ModelContext<Self>) {
|
|
if self.consecutive_failures >= MAX_FAILURE_RETRIES {
|
|
log::warn!(
|
|
"IAP token fetch failed {MAX_FAILURE_RETRIES} times in a row; giving up until \
|
|
manual refresh or server challenge"
|
|
);
|
|
return;
|
|
}
|
|
// Delay = BASE * 2^failures, capped at MAX. Using u32 shift is
|
|
// safe because we cap failures at MAX_FAILURE_RETRIES (< 32).
|
|
let delay = BASE_FAILURE_RETRY_DELAY
|
|
.saturating_mul(1u32 << self.consecutive_failures)
|
|
.min(MAX_FAILURE_RETRY_DELAY);
|
|
self.consecutive_failures += 1;
|
|
log::info!(
|
|
"Scheduling IAP refresh retry #{} in {}s",
|
|
self.consecutive_failures,
|
|
delay.as_secs()
|
|
);
|
|
self.schedule_retry(delay, ctx);
|
|
}
|
|
|
|
fn schedule_retry(&mut self, delay: Duration, ctx: &mut ModelContext<Self>) {
|
|
ctx.spawn(
|
|
async move {
|
|
Timer::after(delay).await;
|
|
},
|
|
|manager, _, ctx| {
|
|
manager.start_refresh(ctx);
|
|
},
|
|
);
|
|
}
|
|
|
|
/// Inspects a websocket *handshake* connect error for an IAP challenge.
|
|
/// If detected, triggers a refresh so the caller's retry loop can pick up
|
|
/// a fresh token on the next attempt.
|
|
#[cfg(not(target_family = "wasm"))]
|
|
pub fn check_ws_connect_error(&mut self, err: &anyhow::Error, ctx: &mut ModelContext<Self>) {
|
|
if ws_connect_is_iap_challenge(err) {
|
|
log::warn!("Received IAP challenge on websocket handshake; triggering refresh");
|
|
self.handle_challenge(ctx);
|
|
}
|
|
}
|
|
|
|
#[cfg(target_family = "wasm")]
|
|
pub fn check_ws_connect_error(&mut self, _err: &anyhow::Error, _ctx: &mut ModelContext<Self>) {}
|
|
}
|
|
|
|
#[cfg(not(target_family = "wasm"))]
|
|
pub fn ws_connect_is_iap_challenge(err: &anyhow::Error) -> bool {
|
|
connect_error_http_response(err).is_some_and(|response| {
|
|
http_client::iap::is_iap_challenge(response.status(), response.headers())
|
|
})
|
|
}
|
|
|
|
impl Entity for IapManager {
|
|
type Event = IapManagerEvent;
|
|
}
|
|
|
|
impl SingletonEntity for IapManager {}
|
|
|
|
/// How long to wait for `auth print-identity-token` command to respond before killing it.
|
|
const GCLOUD_TIMEOUT: Duration = Duration::from_secs(30);
|
|
|
|
// gcloud ships as `gcloud.cmd` on Windows
|
|
#[cfg(windows)]
|
|
const GCLOUD_PROGRAM: &str = "gcloud.cmd";
|
|
#[cfg(not(windows))]
|
|
const GCLOUD_PROGRAM: &str = "gcloud";
|
|
|
|
fn fetch_iap_token(
|
|
audiences: &str,
|
|
service_account_email: &str,
|
|
path_env: Option<&str>,
|
|
) -> Result<CachedToken> {
|
|
let args = [
|
|
"auth",
|
|
"print-identity-token",
|
|
"--audiences",
|
|
audiences,
|
|
"--impersonate-service-account",
|
|
service_account_email,
|
|
"--include-email",
|
|
];
|
|
let cmd_display = format!("{GCLOUD_PROGRAM} {}", args.join(" "));
|
|
|
|
let mut cmd = command::blocking::Command::new(GCLOUD_PROGRAM);
|
|
cmd
|
|
// Prevent gcloud from waiting for interactive input (fail fast instead of hanging)
|
|
.stdin(std::process::Stdio::null())
|
|
.stdout(std::process::Stdio::piped())
|
|
.stderr(std::process::Stdio::piped())
|
|
.args(args);
|
|
// allows warp to resolve `gcloud` cli path
|
|
if let Some(path_env) = path_env {
|
|
cmd.env("PATH", path_env);
|
|
}
|
|
let mut child = cmd
|
|
.spawn()
|
|
.map_err(|err| anyhow::anyhow!("Failed to spawn `{cmd_display}`: {err}"))?;
|
|
|
|
// Poll for completion, killing the child if it exceeds the timeout.
|
|
let start = Instant::now();
|
|
loop {
|
|
match child.try_wait() {
|
|
Ok(Some(_)) => break,
|
|
Ok(None) => {
|
|
if start.elapsed() > GCLOUD_TIMEOUT {
|
|
let _ = child.kill();
|
|
anyhow::bail!(
|
|
"`{cmd_display}` timed out after {}s",
|
|
GCLOUD_TIMEOUT.as_secs()
|
|
);
|
|
}
|
|
std::thread::sleep(Duration::from_millis(100));
|
|
}
|
|
Err(err) => anyhow::bail!("Failed to wait for `{cmd_display}`: {err}"),
|
|
}
|
|
}
|
|
|
|
let output = child
|
|
.wait_with_output()
|
|
.map_err(|err| anyhow::anyhow!("Failed to collect output from `{cmd_display}`: {err}"))?;
|
|
|
|
if !output.status.success() {
|
|
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
|
anyhow::bail!("`{cmd_display}` failed: {stderr}");
|
|
}
|
|
|
|
let token = String::from_utf8(output.stdout)
|
|
.map_err(|err| anyhow::anyhow!("gcloud output is not valid UTF-8: {err}"))?
|
|
.trim()
|
|
.to_string();
|
|
|
|
anyhow::ensure!(!token.is_empty(), "gcloud returned an empty token");
|
|
|
|
let expires_at = get_expires_at(&token)?;
|
|
Ok(CachedToken { token, expires_at })
|
|
}
|
|
|
|
fn get_expires_at(token: &str) -> Result<Instant> {
|
|
let exp = parse_exp_from_jwt(token).ok_or_else(|| {
|
|
anyhow::anyhow!("IAP token missing or unparseable `exp` claim; refusing to cache")
|
|
})?;
|
|
// `exp` is Unix wall-clock seconds; `Instant` is monotonic and
|
|
// has no Unix-time API, so bridge via `SystemTime::now()` to
|
|
// compute a delta, then add that to `Instant::now()`.
|
|
let now = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.map_err(|err| anyhow::anyhow!("system clock is before unix epoch: {err}"))?
|
|
.as_secs();
|
|
let secs_remaining = exp
|
|
.checked_sub(now)
|
|
.ok_or_else(|| anyhow::anyhow!("IAP token is already expired (exp={exp}, now={now})"))?;
|
|
Ok(Instant::now() + Duration::from_secs(secs_remaining))
|
|
}
|
|
|
|
fn parse_exp_from_jwt(token: &str) -> Option<u64> {
|
|
let payload_b64 = token.split('.').nth(1)?;
|
|
let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
|
.decode(payload_b64)
|
|
.ok()?;
|
|
let payload: serde_json::Value = serde_json::from_slice(&payload_bytes).ok()?;
|
|
payload.get("exp")?.as_u64()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "iap_tests.rs"]
|
|
mod tests;
|