From b313ee9577d8ee50c79400847f3f5b0b6533343a Mon Sep 17 00:00:00 2001 From: YOMXXX Date: Fri, 29 May 2026 22:22:05 +0800 Subject: [PATCH] fix(socket): refresh token before reconnect, fast-fail on Invalid token (#2892) (#2905) --- src/openhuman/socket/manager.rs | 56 ++- src/openhuman/socket/mod.rs | 1 + src/openhuman/socket/schemas.rs | 24 +- src/openhuman/socket/token_provider.rs | 151 ++++++++ src/openhuman/socket/ws_loop.rs | 220 ++++++++++- src/openhuman/socket/ws_loop_tests.rs | 501 ++++++++++++++++++++++++- 6 files changed, 916 insertions(+), 37 deletions(-) create mode 100644 src/openhuman/socket/token_provider.rs diff --git a/src/openhuman/socket/manager.rs b/src/openhuman/socket/manager.rs index 1ad219963..5e80fc3d5 100644 --- a/src/openhuman/socket/manager.rs +++ b/src/openhuman/socket/manager.rs @@ -19,6 +19,7 @@ use tokio::time::Duration; use crate::api::models::socket::{ConnectionStatus, SocketState}; use crate::openhuman::webhooks::WebhookRouter; +use super::token_provider::{static_token_provider, TokenProvider}; use super::ws_loop::ws_loop; // --------------------------------------------------------------------------- @@ -139,7 +140,59 @@ impl SocketManager { log::error!("[socket] connect: refusing to start — empty session token"); return Err("empty session token — authenticate first".to_string()); } + // Wrap the static token in a provider closure. Existing callers that + // pass a concrete token value continue to work unchanged; the provider + // returns that same token on every call (static semantics). For + // live-session refresh, callers should use `connect_with_session` which + // builds a provider via `token_provider_from_config`. + let provider = static_token_provider(token.to_string()); + self.spawn_loop(url, provider).await + } + /// Connect using a **live-refresh token provider**. + /// + /// Unlike [`connect`] which wraps a single static token, this method + /// accepts a [`TokenProvider`] closure that is called before every + /// reconnect attempt. Use this when the token may change between retries + /// (e.g. after a session refresh or re-login) so the loop always sends the + /// freshest available credential. + /// + /// The provider is called immediately to validate that a token is available + /// before the background task is spawned — callers receive an actionable + /// `Err` if no token is stored rather than spawning a doomed retry loop. + pub async fn connect_with_provider( + &self, + url: &str, + token_provider: TokenProvider, + ) -> Result<(), String> { + // Validate that a token is available right now before spawning. This + // mirrors the empty-token guard in `connect()` and ensures callers + // see an immediate error if the session store is empty. + match token_provider() { + Ok(t) if !t.trim().is_empty() => {} + Ok(_) => { + log::error!( + "[socket] connect_with_provider: refusing to start — provider returned empty token" + ); + return Err("empty session token — authenticate first".to_string()); + } + Err(e) => { + log::error!( + "[socket] connect_with_provider: refusing to start — provider error: {e}" + ); + return Err(e); + } + } + self.spawn_loop(url, token_provider).await + } + + /// Shared spawn path used by both [`connect`] and [`connect_with_provider`]. + /// + /// Installs the rustls crypto provider, tears down any existing connection, + /// constructs the channel pair, and spawns the background `ws_loop` task. + /// Entry-point-specific validation (empty-token guard, provider pre-check) + /// is done by the callers before this is called. + async fn spawn_loop(&self, url: &str, provider: TokenProvider) -> Result<(), String> { // Ensure the rustls crypto provider is installed (needed for wss:// TLS). // This is a no-op if already installed. let _ = rustls::crypto::ring::default_provider().install_default(); @@ -160,11 +213,10 @@ impl SocketManager { *self.shutdown_tx.lock().await = Some(shutdown_tx); let url = url.to_string(); - let token = token.to_string(); let shared = Arc::clone(&self.shared); let handle = tokio::spawn(async move { - ws_loop(url, token, shared, emit_rx, shutdown_rx, internal_tx).await; + ws_loop(url, provider, shared, emit_rx, shutdown_rx, internal_tx).await; }); *self.loop_handle.lock().await = Some(handle); diff --git a/src/openhuman/socket/mod.rs b/src/openhuman/socket/mod.rs index f0ff5eebc..b20eacf1f 100644 --- a/src/openhuman/socket/mod.rs +++ b/src/openhuman/socket/mod.rs @@ -7,6 +7,7 @@ mod event_handlers; pub mod manager; mod schemas; +pub(crate) mod token_provider; pub mod types; pub(crate) mod ws_loop; diff --git a/src/openhuman/socket/schemas.rs b/src/openhuman/socket/schemas.rs index aa2620576..7f565d243 100644 --- a/src/openhuman/socket/schemas.rs +++ b/src/openhuman/socket/schemas.rs @@ -220,20 +220,34 @@ fn handle_connect_with_session(_params: Map) -> ControllerFuture log::info!("[socket:rpc] connect_with_session — resolving credentials"); - // Load config for API URL and session token. - let config = crate::openhuman::config::rpc::load_config_with_timeout().await?; + // Load config once to derive the API URL and perform the initial + // token validation. The config is then moved into a live-refresh + // `TokenProvider` so the reconnect loop can re-read the latest + // session token from the profile store on every subsequent attempt + // (fix for TAURI-RUST-9C / #2892 — stale-token retry storm). + let config = + std::sync::Arc::new(crate::openhuman::config::rpc::load_config_with_timeout().await?); let api_url = crate::api::config::effective_backend_api_url(&config.api_url); - let token = crate::api::jwt::get_session_token(&config) + + // Perform an eager check so the RPC caller gets an immediate error + // if no session is stored — the provider will do the same check on + // every reconnect attempt inside `ws_loop`. + let initial_token = crate::api::jwt::get_session_token(&config) .map_err(|e| format!("failed to read session token: {e}"))? .ok_or("no session token stored — user must log in first")?; log::info!( "[socket:rpc] connect_with_session url={} token_len={}", api_url, - token.len() + initial_token.len() ); - mgr.connect(&api_url, &token).await?; + // Build a live-refresh provider: on every reconnect attempt the loop + // calls this closure to obtain the most recently stored token, picking + // up any rotation or re-login that happened while sleeping. + let provider = super::token_provider::token_provider_from_config(config); + + mgr.connect_with_provider(&api_url, provider).await?; let state = mgr.get_state(); Ok(json!({ "status": format!("{:?}", state.status) })) diff --git a/src/openhuman/socket/token_provider.rs b/src/openhuman/socket/token_provider.rs new file mode 100644 index 000000000..37a02302c --- /dev/null +++ b/src/openhuman/socket/token_provider.rs @@ -0,0 +1,151 @@ +//! Token-provider abstraction for the WebSocket reconnect loop. +//! +//! Separating this into its own module isolates the token contract for unit +//! testing and keeps token-refresh logic out of the connection-loop hot path. +//! +//! ## Why a callback instead of `String`? +//! +//! The original `ws_loop` captured the session token once at spawn time and +//! reused it for every retry. When the backend invalidates a JWT mid-session +//! (rotation, explicit sign-out, server-side expiry), all subsequent reconnect +//! attempts carry the same dead token — exactly the "Invalid token" retry storm +//! tracked as TAURI-RUST-9C (#2892). A callback lets the loop re-read the +//! latest token from the profile store before each attempt. + +use std::sync::Arc; + +/// Callable that returns the current session token on demand. +/// +/// `Ok(token)` — a non-empty token was available; use it for the next attempt. +/// `Err(reason)` — no token is stored (user logged out, profile corrupt); the +/// caller should surface `reason` and exit the reconnect loop. +/// +/// The provider is intentionally **synchronous** — reading a token from the +/// profile store is a lock + disk read, not an async I/O round-trip. Wrapping +/// it in `Arc` lets it be cheaply cloned into the spawned task without +/// requiring `async_fn_in_trait` or boxing. +pub(super) type TokenProvider = Arc Result + Send + Sync>; + +/// Returns `true` iff the failure reason carries both the Socket.IO CONNECT +/// prefix AND the `"invalid token"` sentinel — a strict double anchor to avoid +/// misclassifying unrelated bare 401s (e.g. an upstream HTTP error message +/// that happens to contain `"invalid token"`). +/// +/// The upstream shape produced by `read_sio_connect_ack()` is: +/// ```text +/// Socket.IO connect error: Invalid token +/// ``` +/// Matching is case-insensitive so capitalisation variants are also caught. +/// +/// This function is `pub(super)` so `ws_loop.rs` and the tests module can call +/// it without exporting it beyond the `socket` domain. +pub(super) fn is_invalid_token_error(reason: &str) -> bool { + let lower = reason.to_ascii_lowercase(); + // Primary anchor: "invalid token" preceded by "socket.io connect error" + // produced by our own `read_sio_connect_ack()`. + lower.contains("socket.io connect error") && lower.contains("invalid token") +} + +/// Build a static provider that always returns the same token value. +/// +/// Used by `SocketManager::connect(url, token)` so the public API does not +/// change: existing callers that already have a token in hand (e.g. the CLI, +/// integration tests) can still pass it as a `&str` without touching the +/// provider layer. On each reconnect the loop will re-call the provider — in +/// this case it returns the same cloned string — which is equivalent to the +/// previous behaviour. +/// +/// For **live** session-token refresh, build a provider via +/// `token_provider_from_config` (used by `handle_connect_with_session`). +pub(super) fn static_token_provider(token: String) -> TokenProvider { + Arc::new(move || { + if token.trim().is_empty() { + Err("empty session token — authenticate first".to_string()) + } else { + Ok(token.clone()) + } + }) +} + +/// Build a provider that reads the latest session token from the profile store +/// on every call. +/// +/// This is the **live-refresh** path used by `handle_connect_with_session`: +/// when the loop retries after a disconnect it will see any token that was +/// refreshed or re-stored since the previous attempt. +pub(super) fn token_provider_from_config( + config: Arc, +) -> TokenProvider { + Arc::new(move || { + crate::api::jwt::get_session_token(&config) + .map_err(|e| format!("failed to read session token: {e}"))? + .ok_or_else(|| "no session token stored — user must log in first".to_string()) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn is_invalid_token_error_matches_exact_sio_wire_shape() { + assert!(is_invalid_token_error( + "Socket.IO connect error: Invalid token" + )); + } + + #[test] + fn is_invalid_token_error_is_case_insensitive() { + assert!(is_invalid_token_error( + "socket.io connect error: invalid token" + )); + assert!(is_invalid_token_error( + "SOCKET.IO CONNECT ERROR: INVALID TOKEN" + )); + } + + #[test] + fn is_invalid_token_error_requires_both_anchors() { + // "invalid token" without the "socket.io connect error" prefix must + // not fire — otherwise bare upstream 401s from other contexts would + // trigger the fast-fail session-expiry path. + assert!(!is_invalid_token_error("invalid token")); + assert!(!is_invalid_token_error("auth error: invalid token")); + // The SIO connect error prefix without the "invalid token" body must + // not fire either — a server-side config error, for instance. + assert!(!is_invalid_token_error( + "socket.io connect error: namespace not found" + )); + } + + #[test] + fn is_invalid_token_error_returns_false_for_unrelated_errors() { + assert!(!is_invalid_token_error( + "WebSocket connect: connection refused" + )); + assert!(!is_invalid_token_error("EIO OPEN: timeout")); + assert!(!is_invalid_token_error("")); + } + + #[test] + fn static_provider_returns_token() { + let provider = static_token_provider("my-token".to_string()); + assert_eq!(provider().unwrap(), "my-token"); + } + + #[test] + fn static_provider_rejects_empty_token() { + let provider = static_token_provider("".to_string()); + assert!(provider().is_err()); + let provider2 = static_token_provider(" ".to_string()); + assert!(provider2().is_err()); + } + + #[test] + fn static_provider_returns_same_token_on_repeated_calls() { + let provider = static_token_provider("tok-abc".to_string()); + // Simulates multiple reconnect attempts — must always return the same + // cloned token (static provider semantics). + assert_eq!(provider().unwrap(), provider().unwrap()); + } +} diff --git a/src/openhuman/socket/ws_loop.rs b/src/openhuman/socket/ws_loop.rs index c0e6a75c3..ba9828e47 100644 --- a/src/openhuman/socket/ws_loop.rs +++ b/src/openhuman/socket/ws_loop.rs @@ -16,6 +16,7 @@ use crate::openhuman::util::utf8_safe_prefix_at_byte_boundary; use super::event_handlers::{handle_sio_event, parse_sio_event}; use super::manager::{emit_state_change, SharedState}; +use super::token_provider::{is_invalid_token_error, TokenProvider}; use super::types::{ConnectionOutcome, WsStream}; /// Maximum HTTP redirect hops to follow during a single WebSocket connect attempt. @@ -49,32 +50,46 @@ const MAX_REDIRECT_HOPS: u8 = 3; const FAIL_ESCALATE_THRESHOLD: u32 = 5; /// Background loop that manages the WebSocket connection and reconnection. +/// +/// `token_provider` is called before **each** connection attempt, so a +/// token that was refreshed or re-stored on disk (e.g. after the user +/// re-logged-in while the loop was sleeping) is picked up automatically. +/// +/// On a `"Socket.IO connect error: Invalid token"` rejection the loop +/// performs one extra provider call to check whether a fresher token +/// became available. If the token is unchanged (or the second attempt also +/// fails with "Invalid token") the loop escalates immediately — it does +/// **not** waste the remaining back-off attempts on a provably dead token. +/// This is the fix for TAURI-RUST-9C (#2892). pub(super) async fn ws_loop( url: String, - token: String, + token_provider: TokenProvider, shared: Arc, mut emit_rx: mpsc::UnboundedReceiver, mut shutdown_rx: watch::Receiver, internal_tx: mpsc::UnboundedSender, ) { - // Defense in depth: `SocketManager::connect` is the primary guard and - // will reject an empty token before spawning this task. This check - // covers any future caller that invokes `ws_loop` directly (e.g. tests - // or internal plumbing bypassing the manager). Without it, every attempt - // would either 401 at the SIO CONNECT step or fail upstream at the - // gateway, producing the retry-storm Sentry noise this module is designed - // to suppress. - if token.trim().is_empty() { - log::error!("[socket] ws_loop: refusing to start — empty session token"); - *shared.status.write() = ConnectionStatus::Disconnected; - *shared.socket_id.write() = None; - emit_state_change(&shared); - return; - } - let mut backoff = Duration::from_millis(1000); let max_backoff = Duration::from_secs(30); let mut consecutive_failures: u32 = 0; + // How many `RetryImmediately` short-circuits we've taken in the **current** + // "fresh-token cycle" (i.e. since the last successful connection or + // non-token-related failure). Bounded to 1 so a buggy / non-deterministic + // provider that returns a *different* non-empty token on every call cannot + // hot-loop: connect → Invalid token → fresh token → connect → Invalid token + // → fresh token → ... with no sleep and no escalation. After one immediate + // shot we fall through to the normal backoff + escalation path. See the + // CodeRabbit Major on PR #2905. + let mut fresh_token_retries: u32 = 0; + // Fresh token carried forward from a `RetryImmediately` decision. When + // `decide_after_invalid_token` re-fetches the provider and finds a + // genuinely different value, we stash it here so the next loop iteration + // skips the redundant top-of-loop provider call and uses **exactly** the + // token that was validated by the decision step — avoiding a redundant + // lock + disk read and the case where the logged fresh-token length + // drifts from what's actually sent over the wire. See the CodeRabbit + // Minor on PR #2905. + let mut pending_token: Option = None; // `ws_url` is the *resolved* socket URL we're currently connecting to. // If the backend responds with an HTTP 3xx during the upgrade (typical when @@ -88,7 +103,45 @@ pub(super) async fn ws_loop( break; } - log::info!("[socket] Attempting connection..."); + // If a fresh token was carried forward from the previous iteration's + // `RetryImmediately` decision, consume it now and skip the redundant + // provider call — `decide_after_invalid_token` already validated that + // it is non-empty and distinct from the previously-rejected token, so + // re-reading the provider here would only re-acquire the same lock / + // re-hit the same disk file. Otherwise fetch the latest token + // afresh. If the provider returns an error (no token stored → user + // is logged out, profile corrupt), there is nothing useful to retry + // with — surface the error and exit cleanly rather than spamming + // the server. + let token = match pending_token.take() { + Some(t) => t, + None => match token_provider() { + Ok(t) if !t.trim().is_empty() => t, + Ok(_) => { + log::warn!("[socket] ws_loop: token provider returned empty token — stopping"); + *shared.error.write() = + Some("session expired — please sign in again".to_string()); + *shared.status.write() = ConnectionStatus::Disconnected; + *shared.socket_id.write() = None; + emit_state_change(&shared); + return; + } + Err(e) => { + log::warn!("[socket] ws_loop: token provider failed — stopping: {e}"); + *shared.error.write() = + Some("session expired — please sign in again".to_string()); + *shared.status.write() = ConnectionStatus::Disconnected; + *shared.socket_id.write() = None; + emit_state_change(&shared); + return; + } + }, + }; + + log::info!( + "[socket] Attempting connection (token_len={})...", + token.len() + ); *shared.status.write() = ConnectionStatus::Connecting; emit_state_change(&shared); @@ -111,7 +164,10 @@ pub(super) async fn ws_loop( // `Lost` is only returned after a successful SIO CONNECT ACK // (see `run_connection`), so reaching this arm proves the // backend is reachable and the token is valid. Reset both - // the backoff and the failure streak. + // the backoff and the failure streak — and clear the + // fresh-token retry counter so a long-lived session that + // accumulated a RetryImmediately on a prior reconnect cycle + // doesn't carry that dead state into the next one. if consecutive_failures > 0 { log::debug!( "[socket] Connection re-established; resetting failure streak ({} cleared)", @@ -119,9 +175,87 @@ pub(super) async fn ws_loop( ); } consecutive_failures = 0; + fresh_token_retries = 0; log::warn!("[socket] Connection lost: {}", reason); backoff = Duration::from_millis(1000); } + ConnectionOutcome::Failed(reason) if is_invalid_token_error(&reason) => { + // The server rejected our token explicitly. Try one more + // provider call — in case the token was refreshed on disk + // since we fetched it moments ago (e.g. another code path + // rotated it). If the provider returns a genuinely different + // token, we can give it one more shot without consuming the + // normal backoff budget. + log::warn!( + "[socket] Invalid token on attempt — checking for fresh token (current_len={})", + token.len() + ); + match decide_after_invalid_token(&token, &token_provider) { + InvalidTokenAction::RetryImmediately { token: fresh } => { + let fresh_len = fresh.len(); + fresh_token_retries = fresh_token_retries.saturating_add(1); + if fresh_token_retries > 1 { + // We already gave the fresh-token cycle one + // immediate shot. A provider that keeps returning + // *different* non-empty tokens (rapid server-side + // rotation, non-deterministic source, buggy impl) + // could otherwise hot-loop forever with no sleep + // and no escalation — arguably worse than the + // 5-retry storm this PR was originally fixing. Fall + // through to the normal failure path so backoff + // sleeps and `consecutive_failures` escalation + // converge on a definitive outcome. + log::warn!( + "[socket] Fresh token available (len={fresh_len}) but already \ + retried once this cycle — escalating to normal backoff path" + ); + consecutive_failures = consecutive_failures.saturating_add(1); + log_connection_failure(consecutive_failures, &reason); + // Fall through to the backoff sleep below. + // Intentionally drop `fresh` here: the bounded + // path now demands a backoff sleep, after which + // the next loop iteration will re-fetch the + // provider afresh (the token we just got may + // itself be stale by the time the sleep + // completes — this is the only knowable-correct + // policy for a rotating-source provider). + } else { + // We have a genuinely different token — try once + // immediately (no backoff sleep) with the **exact** + // token the decision step validated. Stash it for + // the next iteration so the top-of-loop provider + // call is skipped and we don't re-acquire the + // session-store lock or re-read from disk just to + // get the same value back. If this attempt also + // fails we will go through the normal escalation + // path on the next loop iteration (either + // same-token Escalate or the bounded fall-through + // above). + log::info!( + "[socket] Fresh token available (len={fresh_len}), retrying immediately" + ); + pending_token = Some(fresh); + // Don't increment consecutive_failures for an attempt we + // couldn't have avoided — the token we used was already + // stale at fetch time. + continue; + } + } + InvalidTokenAction::Escalate { reason } => { + // No fresh token — the session is definitively expired. + // Escalate immediately instead of wasting more attempts + // on what is provably a dead token. This is the core fix + // for TAURI-RUST-9C (#2892). + log::warn!("[socket] Session expired ({reason}) — stopping reconnect loop"); + *shared.error.write() = + Some("session expired — please sign in again".to_string()); + *shared.status.write() = ConnectionStatus::Disconnected; + *shared.socket_id.write() = None; + emit_state_change(&shared); + return; + } + } + } ConnectionOutcome::Failed(reason) => { consecutive_failures = consecutive_failures.saturating_add(1); log_connection_failure(consecutive_failures, &reason); @@ -207,6 +341,56 @@ fn log_connection_failure(consecutive: u32, reason: &str) { } } +// --------------------------------------------------------------------------- +// Invalid-token decision helper +// --------------------------------------------------------------------------- + +/// Action the reconnect loop should take after receiving an "Invalid token" +/// rejection from the server. +enum InvalidTokenAction { + /// A genuinely different token is available — retry the connection + /// immediately (no backoff sleep). The fresh token is carried forward so + /// the next `run_connection` uses **exactly** the validated value, not + /// whatever a subsequent re-read of the provider returns — avoiding a + /// redundant lock + disk I/O and the case where the logged `fresh_len` + /// drifts from the token actually sent over the wire. + RetryImmediately { token: String }, + /// No fresh token is available; the session is definitively expired. + /// `reason` is a short diagnostic string for the log line. + Escalate { reason: String }, +} + +/// Pure decision function: given the token that was just rejected and the +/// provider that may have a fresher one, decide what the reconnect loop +/// should do next. +/// +/// Calls `provider()` exactly once. No socket I/O. +/// +/// - Provider returns a **different, non-empty** token → `RetryImmediately` +/// carrying the fresh token for the caller to reuse on the next attempt. +/// - Provider returns the **same** token → `Escalate` (no point retrying). +/// - Provider returns an **empty** token → `Escalate` (treat as no session). +/// - Provider returns `Err` → `Escalate` with the provider error as reason. +fn decide_after_invalid_token( + previous_token: &str, + provider: &TokenProvider, +) -> InvalidTokenAction { + match provider() { + Ok(fresh) if !fresh.trim().is_empty() && fresh != previous_token => { + InvalidTokenAction::RetryImmediately { token: fresh } + } + Ok(same) if same == previous_token => InvalidTokenAction::Escalate { + reason: "token unchanged after provider re-fetch".to_string(), + }, + Ok(_) => InvalidTokenAction::Escalate { + reason: "provider returned empty token".to_string(), + }, + Err(e) => InvalidTokenAction::Escalate { + reason: format!("provider error: {e}"), + }, + } +} + // --------------------------------------------------------------------------- // Single connection attempt // --------------------------------------------------------------------------- diff --git a/src/openhuman/socket/ws_loop_tests.rs b/src/openhuman/socket/ws_loop_tests.rs index 7d8682f3c..6b9a3a563 100644 --- a/src/openhuman/socket/ws_loop_tests.rs +++ b/src/openhuman/socket/ws_loop_tests.rs @@ -1,7 +1,10 @@ use super::*; use parking_lot::RwLock; +use std::sync::atomic::{AtomicUsize, Ordering}; use tokio_tungstenite::tungstenite::http::{header::LOCATION, Response, StatusCode}; +use crate::openhuman::socket::token_provider::{is_invalid_token_error, static_token_provider}; + fn make_shared() -> Arc { Arc::new(SharedState { webhook_router: RwLock::new(None), @@ -519,7 +522,7 @@ async fn ws_loop_completes_handshake_and_shuts_down_cleanly() { let handle = tokio::spawn(async move { ws_loop( http_base_for(addr), - "test-token".into(), + static_token_provider("test-token".to_string()), loop_shared, emit_rx, shutdown_rx, @@ -577,7 +580,7 @@ async fn ws_loop_handles_connect_error_and_shutdown() { let handle = tokio::spawn(async move { ws_loop( http_base_for(addr), - "t".into(), + static_token_provider("t".to_string()), loop_shared, emit_rx, shutdown_rx, @@ -586,8 +589,9 @@ async fn ws_loop_handles_connect_error_and_shutdown() { .await; }); - // Give the loop a moment to observe the CONNECT_ERROR, then shut down - // before the reconnection backoff fires. + // Give the loop a moment to observe the CONNECT_ERROR (44{"message":"nope"} + // — not an "Invalid token", so it goes through the normal backoff path), + // then shut down. tokio::time::sleep(tokio::time::Duration::from_millis(400)).await; let _ = shutdown_tx.send(true); let _ = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; @@ -610,7 +614,7 @@ async fn ws_loop_handles_bad_eio_open_and_shutdown() { let handle = tokio::spawn(async move { ws_loop( http_base_for(addr), - "t".into(), + static_token_provider("t".to_string()), loop_shared, emit_rx, shutdown_rx, @@ -638,11 +642,10 @@ fn connect_behavior_variants_are_distinct() { } } -/// Empty-token guard: if a caller invokes the bare `connect(url, token)` RPC -/// with an empty string the loop must bail immediately rather than spin a -/// doomed reconnect cycle that fires Sentry events on every retry. The -/// status must end up `Disconnected` and the function must return without -/// completing the reconnect loop. +/// Empty-token guard: if the token provider returns an empty token the loop +/// must bail immediately rather than spin a doomed reconnect cycle that fires +/// Sentry events on every retry. The status must end up `Disconnected` and +/// the function must return without completing the reconnect loop. #[tokio::test] async fn ws_loop_refuses_to_start_with_empty_token() { let shared = make_shared(); @@ -665,7 +668,8 @@ async fn ws_loop_refuses_to_start_with_empty_token() { // URL is deliberately invalid — if the guard misfires, the // task would error on connect rather than return immediately. "http://invalid.example.invalid:1".into(), - " ".into(), // whitespace-only counts as empty + // static_token_provider(" ") returns Err for whitespace-only. + static_token_provider(" ".to_string()), loop_shared, emit_rx, shutdown_rx, @@ -684,6 +688,50 @@ async fn ws_loop_refuses_to_start_with_empty_token() { assert!(shared.socket_id.read().is_none()); } +/// Provider-error guard: if the token provider returns Err (e.g. no session +/// stored, profile corrupt) the loop exits immediately with an error set in +/// SharedState rather than attempting a connection. +#[tokio::test] +async fn ws_loop_exits_cleanly_when_provider_returns_error() { + let shared = make_shared(); + *shared.status.write() = ConnectionStatus::Connecting; + + let (_emit_tx, emit_rx) = mpsc::unbounded_channel::(); + let (_shutdown_tx, shutdown_rx) = watch::channel(false); + let (internal_tx, _internal_rx) = mpsc::unbounded_channel::(); + + let loop_shared = Arc::clone(&shared); + let handle = tokio::spawn(async move { + ws_loop( + "http://invalid.example.invalid:1".into(), + // Provider always returns Err — simulates logged-out state. + Arc::new(|| Err("no session token stored — user must log in first".to_string())), + loop_shared, + emit_rx, + shutdown_rx, + internal_tx, + ) + .await; + }); + + let res = tokio::time::timeout(tokio::time::Duration::from_secs(2), handle).await; + assert!( + matches!(res, Ok(Ok(()))), + "ws_loop must return cleanly when provider errors (no timeout, no panic)" + ); + + assert_eq!(*shared.status.read(), ConnectionStatus::Disconnected); + assert!(shared.socket_id.read().is_none()); + // The error slot must carry an actionable message. + let err = shared.error.read().clone(); + assert!( + err.as_deref() + .map(|e| e.contains("session expired")) + .unwrap_or(false), + "expected session-expired error in SharedState, got: {err:?}" + ); +} + // ── End-to-end redirect-follow (the real fix for the 301 noise) ── use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -740,7 +788,7 @@ async fn ws_loop_follows_301_to_working_backend() { let handle = tokio::spawn(async move { ws_loop( format!("http://{redirect_addr}"), - "redirect-test-token".into(), + static_token_provider("redirect-test-token".to_string()), loop_shared, emit_rx, shutdown_rx, @@ -818,3 +866,432 @@ async fn connect_with_redirects_fails_when_location_missing() { // No warning recorded because the redirect was never actually followed. assert!(shared.error.read().is_none()); } + +// ── Token-refresh and Invalid-token escalation (#2892) ──────────── + +/// Spawn a single-accept EIO v4 server that always rejects the SIO CONNECT +/// with `44{"message":"Invalid token"}`. Used to test the fast-fail path. +async fn spawn_mock_invalid_token_server() -> std::net::SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + tokio::spawn(async move { + // Accept connections in a loop so the server handles more than one + // attempt (the retry-on-fresh-token path triggers a second connection). + loop { + let Ok((stream, _)) = listener.accept().await else { + break; + }; + tokio::spawn(async move { + let ws = accept_async(stream).await.expect("ws accept"); + let (mut write, mut read) = ws.split(); + // 1. Send EIO OPEN. + let open = + r#"0{"sid":"mock-eio","upgrades":[],"pingInterval":1000,"pingTimeout":2000}"#; + let _ = write.send(WsMessage::Text(open.to_string())).await; + // 2. Drain the SIO CONNECT frame (don't care about its content). + let _ = read.next().await; + // 3. Reply with CONNECT_ERROR "Invalid token". + let _ = write + .send(WsMessage::Text( + r#"44{"message":"Invalid token"}"#.to_string(), + )) + .await; + let _ = write.close().await; + }); + } + }); + addr +} + +/// Provider called once per attempt: a counter-based provider proves the loop +/// re-fetches the token before each `run_connection` invocation. +#[tokio::test] +async fn ws_loop_calls_provider_before_each_attempt() { + // Use a server that always closes immediately (no EIO OPEN) so the loop + // cycles through failures quickly without hitting the backoff sleep. + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + // Accept and immediately close — simulates a connection refused / close. + tokio::spawn(async move { + let Ok((stream, _)) = listener.accept().await else { + return; + }; + let ws = accept_async(stream).await.expect("ws accept"); + let (mut write, _) = ws.split(); + let _ = write.close().await; + }); + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = Arc::clone(&call_count); + + let shared = make_shared(); + *shared.status.write() = ConnectionStatus::Disconnected; + let (_emit_tx, emit_rx) = mpsc::unbounded_channel::(); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let (internal_tx, _internal_rx) = mpsc::unbounded_channel::(); + + let loop_shared = Arc::clone(&shared); + let handle = tokio::spawn(async move { + ws_loop( + http_base_for(addr), + Arc::new(move || { + call_count_clone.fetch_add(1, Ordering::Relaxed); + Ok("counter-token".to_string()) + }), + loop_shared, + emit_rx, + shutdown_rx, + internal_tx, + ) + .await; + }); + + // Let the loop run for at least 2 attempts, then shut down. + // The first attempt triggers a provider call; after the connection closes, + // the loop sleeps (1s backoff) before attempt 2 — we just need to see ≥1 + // call to prove the provider is wired up, then shut down. + for _ in 0..50 { + if call_count.load(Ordering::Relaxed) >= 1 { + break; + } + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + } + let _ = shutdown_tx.send(true); + let _ = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; + + let calls = call_count.load(Ordering::Relaxed); + assert!( + calls >= 1, + "provider must be called at least once before each attempt; got {calls}" + ); +} + +/// "Invalid token" same-token escalation: when the server always rejects with +/// "Invalid token" and the provider keeps returning the same token, the loop +/// MUST exit within 2 attempts — not waste the remaining back-off retries. +/// This is the core regression fix for TAURI-RUST-9C (#2892). +#[tokio::test] +async fn ws_loop_escalates_immediately_on_invalid_token_no_refresh() { + let addr = spawn_mock_invalid_token_server().await; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = Arc::clone(&call_count); + + let shared = make_shared(); + *shared.status.write() = ConnectionStatus::Disconnected; + let (_emit_tx, emit_rx) = mpsc::unbounded_channel::(); + // Do NOT signal shutdown — the loop must exit on its own via fast-fail. + let (_shutdown_tx, shutdown_rx) = watch::channel(false); + let (internal_tx, _internal_rx) = mpsc::unbounded_channel::(); + + let loop_shared = Arc::clone(&shared); + let handle = tokio::spawn(async move { + ws_loop( + http_base_for(addr), + // Provider always returns the same (stale) token. + Arc::new(move || { + call_count_clone.fetch_add(1, Ordering::Relaxed); + Ok("stale-token-xyz".to_string()) + }), + loop_shared, + emit_rx, + shutdown_rx, + internal_tx, + ) + .await; + }); + + // The loop must exit by itself (fast-fail) well within 2 seconds. + // At FAIL_ESCALATE_THRESHOLD=5 the old code would still be sleeping through + // back-off at this point — finishing fast is what proves the fix. + let result = tokio::time::timeout(tokio::time::Duration::from_secs(4), handle).await; + assert!( + matches!(result, Ok(Ok(()))), + "ws_loop must exit cleanly after Invalid token (no timeout, no panic)" + ); + + // Loop must have called the provider at most 2 times (initial attempt + + // one re-fetch check). 3+ would mean it fell through to the old retry path. + let calls = call_count.load(Ordering::Relaxed); + assert!( + calls <= 2, + "provider must not be called more than 2 times on Invalid token fast-fail; got {calls}" + ); + + // Status must be Disconnected and error slot must carry session-expired. + assert_eq!(*shared.status.read(), ConnectionStatus::Disconnected); + let err = shared.error.read().clone(); + assert!( + err.as_deref() + .map(|e| e.contains("session expired")) + .unwrap_or(false), + "expected session-expired error in SharedState, got: {err:?}" + ); +} + +/// "Invalid token" with a fresh token available: provider returns token A on +/// the first call, token B on the second. The loop must NOT fast-fail — it +/// should detect the new token and retry once. Since the mock server also +/// rejects token B, the loop will fast-fail on the third call (same-token +/// case), but the important assertion is that it reached at least 2 actual +/// connection attempts (token A and token B) before stopping. +#[tokio::test] +async fn ws_loop_retries_with_fresh_token_on_invalid_token() { + // Server always replies with "Invalid token". + let addr = spawn_mock_invalid_token_server().await; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = Arc::clone(&call_count); + + let shared = make_shared(); + *shared.status.write() = ConnectionStatus::Disconnected; + let (_emit_tx, emit_rx) = mpsc::unbounded_channel::(); + let (_shutdown_tx, shutdown_rx) = watch::channel(false); + let (internal_tx, _internal_rx) = mpsc::unbounded_channel::(); + + let loop_shared = Arc::clone(&shared); + let handle = tokio::spawn(async move { + ws_loop( + http_base_for(addr), + // First call returns "token-a", second and beyond return "token-b". + Arc::new(move || { + let c = call_count_clone.fetch_add(1, Ordering::Relaxed); + if c == 0 { + Ok("token-a".to_string()) + } else { + Ok("token-b".to_string()) + } + }), + loop_shared, + emit_rx, + shutdown_rx, + internal_tx, + ) + .await; + }); + + // The loop should exit by itself (fast-fail after token-b also rejected). + let result = tokio::time::timeout(tokio::time::Duration::from_secs(6), handle).await; + assert!( + matches!(result, Ok(Ok(()))), + "ws_loop must exit cleanly after both tokens rejected" + ); + + // Provider must have been called at least 2 times. Per the Minor on + // PR #2905 the fresh token is now carried forward through + // `pending_token` so the second attempt does NOT re-call the provider — + // it uses the exact value the decision step validated. The expected + // sequence is therefore: + // + // call 1 → "token-a" (start of first attempt, no pending_token) + // call 2 → "token-b" (re-fetch check after Invalid token for + // "token-a") → RetryImmediately stashes "token-b" + // into pending_token + // (no extra call here: second attempt consumes pending_token = "token-b") + // call 3 → "token-b" (re-fetch check after Invalid token for + // "token-b") → same token → Escalate + // + // We assert ≥ 2 — i.e. the fresh-token check fired at least once. + let calls = call_count.load(Ordering::Relaxed); + assert!( + calls >= 2, + "provider must be called at least twice (initial + fresh-token check); got {calls}" + ); + + // End state must be session-expired. + assert_eq!(*shared.status.read(), ConnectionStatus::Disconnected); + let err = shared.error.read().clone(); + assert!( + err.as_deref() + .map(|e| e.contains("session expired")) + .unwrap_or(false), + "expected session-expired error in SharedState, got: {err:?}" + ); +} + +/// Regression guard for the CodeRabbit Major on PR #2905: a non-deterministic +/// provider that returns a **different** non-empty token on every call must +/// NOT cause the loop to hot-loop indefinitely down the `RetryImmediately` +/// path. The bound is one immediate retry per fresh-token cycle; subsequent +/// `RetryImmediately` outcomes must fall through to the normal +/// `consecutive_failures` + backoff sleep path, which converges on a +/// definitive outcome (escalation or session timeout) rather than hammering +/// the server in a tight loop. +/// +/// Setup: server always replies `Invalid token`; provider returns a brand-new +/// token on every call (token-0, token-1, token-2, …) — none of which the +/// server will accept. Before the bound, this would skip backoff on every +/// retry and produce tens-to-hundreds of attempts per second. After the +/// bound, total provider calls in a 2-second window must stay small (the +/// loop has to sleep through backoff between cycles). +#[tokio::test] +async fn ws_loop_bounds_fresh_token_retries_with_rotating_provider() { + // Server always replies with "Invalid token" — guaranteed rejection. + let addr = spawn_mock_invalid_token_server().await; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = Arc::clone(&call_count); + + let shared = make_shared(); + *shared.status.write() = ConnectionStatus::Disconnected; + let (_emit_tx, emit_rx) = mpsc::unbounded_channel::(); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let (internal_tx, _internal_rx) = mpsc::unbounded_channel::(); + + let loop_shared = Arc::clone(&shared); + let handle = tokio::spawn(async move { + ws_loop( + http_base_for(addr), + // Each provider call returns a fresh, distinct, non-empty token. + // This is the pathological provider the CodeRabbit Major calls + // out: rapid server-side rotation, non-deterministic source, or + // buggy implementation that never converges. + Arc::new(move || { + let n = call_count_clone.fetch_add(1, Ordering::Relaxed); + Ok(format!("rotating-token-{n}")) + }), + loop_shared, + emit_rx, + shutdown_rx, + internal_tx, + ) + .await; + }); + + // Give the loop ~2 seconds to misbehave. A hot-loop (no bound) would + // accumulate dozens-to-hundreds of provider calls in that window — each + // RetryImmediately is a `continue` with no sleep, and a roundtrip to the + // mock server is sub-100 ms on loopback. With the bound, every + // fresh-token cycle gets exactly one no-backoff retry; after that the + // loop must sleep through `backoff` (starts at 1 s, doubles each time) + // before re-attempting. The exponential backoff makes the upper bound + // here forgiving on slow CI while still being orders-of-magnitude below + // any plausible hot-loop count. + tokio::time::sleep(tokio::time::Duration::from_millis(2_000)).await; + + let calls_before_shutdown = call_count.load(Ordering::Relaxed); + let _ = shutdown_tx.send(true); + let _ = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; + + // The loop must NOT have hot-looped — a small bound proves the + // RetryImmediately path is now bounded. The exact number depends on CI + // scheduling, but 20 is comfortably above the steady-state expectation + // (~3–6 cycles in 2 s of backoff sleep, each cycle = 2 provider calls) + // and orders of magnitude below an unbounded loop. + assert!( + calls_before_shutdown <= 20, + "RetryImmediately must be bounded — rotating-token provider triggered \ + {calls_before_shutdown} provider calls in 2 s (suspected hot-loop; expected ≤ 20)" + ); + + // …and the loop must have made progress toward escalation, not stayed + // pinned on the no-backoff `continue` path. After the bound has fired + // once the loop has incremented `consecutive_failures` and started + // sleeping backoff. The clearest external proof of that is at least 3 + // provider calls (per-iteration flow with the Minor `pending_token` + // optimisation also applied): + // + // call 1 → initial connect attempt (token-0) + // call 2 → decide_after_invalid_token re-fetch (token-1) → + // RetryImmediately, stashes token-1 into pending_token + // (second attempt consumes pending_token; no new provider call) + // call 3 → decide_after_invalid_token after the second Invalid token + // (token-2) → RetryImmediately → BOUND HIT → falls through + // to backoff sleep + // + // We assert ≥ 3 (gives slack for a slow loopback) — i.e. we got past the + // single immediate retry into the bounded path. + assert!( + calls_before_shutdown >= 3, + "loop must have progressed past the initial connect + first immediate \ + retry — observed only {calls_before_shutdown} provider calls" + ); + + // End state must be Disconnected. The status field reflects the most + // recent state transition; whether the loop reached `session expired` + // depends on how many cycles fit into the 2-second window before + // shutdown — but Disconnected is the unconditional end state on both + // paths. + assert_eq!(*shared.status.read(), ConnectionStatus::Disconnected); +} + +/// `is_invalid_token_error` unit tests are in `token_provider::tests`. +/// This test pins the exact wire shape from `read_sio_connect_ack()` against +/// the classifier to guard against drift between the two modules. +#[test] +fn sio_connect_error_invalid_token_classifies_correctly() { + // This is the exact string produced by `read_sio_connect_ack()` when + // the server sends `44{"message":"Invalid token"}`. + assert!(is_invalid_token_error( + "Socket.IO connect error: Invalid token" + )); + // A different server error must not trigger the fast-fail path. + assert!(!is_invalid_token_error( + "Socket.IO connect error: namespace not found" + )); + // Internal errors (EIO, WS layer) must not match. + assert!(!is_invalid_token_error("EIO OPEN: timeout")); + assert!(!is_invalid_token_error( + "WebSocket connect: connection refused" + )); +} + +// ── decide_after_invalid_token ───────────────────────────────────────────── + +/// Provider returns a genuinely fresh token → loop should retry immediately, +/// carrying the validated fresh token forward so the next attempt sends +/// exactly that value instead of re-reading the provider. +#[test] +fn decide_after_invalid_token_fresh_token_returns_retry() { + let provider: TokenProvider = Arc::new(|| Ok("fresh-token".to_string())); + match decide_after_invalid_token("stale-token", &provider) { + InvalidTokenAction::RetryImmediately { token } => { + assert_eq!(token, "fresh-token"); + } + InvalidTokenAction::Escalate { reason } => { + panic!("expected RetryImmediately, got Escalate({reason})"); + } + } +} + +/// Provider returns the same token → session is definitively expired; escalate. +#[test] +fn decide_after_invalid_token_same_token_escalates() { + let provider: TokenProvider = Arc::new(|| Ok("same-token".to_string())); + match decide_after_invalid_token("same-token", &provider) { + InvalidTokenAction::Escalate { .. } => {} + InvalidTokenAction::RetryImmediately { .. } => { + panic!("expected Escalate when provider returns the same token"); + } + } +} + +/// Provider returns an error → escalate with the provider error as reason. +#[test] +fn decide_after_invalid_token_provider_error_escalates() { + let provider: TokenProvider = + Arc::new(|| Err("no session token stored — user must log in first".to_string())); + match decide_after_invalid_token("any-token", &provider) { + InvalidTokenAction::Escalate { reason } => { + assert!( + reason.contains("provider error"), + "expected 'provider error' in escalation reason, got: {reason}" + ); + } + InvalidTokenAction::RetryImmediately { .. } => { + panic!("expected Escalate when provider errors"); + } + } +} + +/// Provider returns an empty string → treat as no session; escalate. +#[test] +fn decide_after_invalid_token_empty_token_escalates() { + let provider: TokenProvider = Arc::new(|| Ok(String::new())); + match decide_after_invalid_token("prev-token", &provider) { + InvalidTokenAction::Escalate { .. } => {} + InvalidTokenAction::RetryImmediately { .. } => { + panic!("expected Escalate when provider returns empty token"); + } + } +}