diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3ed6b77bc5..8a5ec82759 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -144,6 +144,19 @@ jobs: - name: Run core and desktop Rust tests run: cargo test --locked -p bitfun-core -p bitfun-desktop + # These crates own platform-sensitive behavior that is not exercised by + # testing bitfun-core/bitfun-desktop alone. Keep their focused contract + # suites in the OS matrix so Linux success cannot hide Windows/macOS + # regressions in worker interruption, relay storage, or OAuth handling. + - name: Run Page Functions runtime tests + run: cargo test --locked -p bitfun-page-function-runtime + + - name: Run Relay service tests + run: cargo test --locked -p bitfun-relay-service + + - name: Run subscription authentication tests + run: cargo test --locked -p bitfun-ai-adapters --features subscription-auth subscription_auth + # ── Frontend: build ──────────────────────────────────────────────── frontend-build: name: Frontend Build diff --git a/Cargo.toml b/Cargo.toml index 9a9a7b0044..2fb4321146 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -122,6 +122,10 @@ glob = "0.3" ignore = "0.4" notify = "8.2" dirs = "6.0" +keyring-core = "1.0.0" +apple-native-keyring-store = { version = "1.0.1", features = ["keychain"] } +windows-native-keyring-store = "1.1.0" +zbus-secret-service-keyring-store = { version = "1.0.0", features = ["crypto-rust"] } dark-light = "1.1" dunce = "1" filetime = "0.2" diff --git a/src/apps/cli/src/account.rs b/src/apps/cli/src/account.rs index c672b29a74..60c427bd94 100644 --- a/src/apps/cli/src/account.rs +++ b/src/apps/cli/src/account.rs @@ -9,24 +9,170 @@ //! //! The master key lives in memory only and is lost when the CLI exits. +use std::future::Future; use std::sync::{ - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, Arc, OnceLock, }; +use std::time::Duration; use anyhow::{anyhow, Result}; -use tokio::sync::RwLock; +use tokio::sync::{Notify, RwLock}; use bitfun_core::service::remote_connect::{ self, encryption, relay_client::RelayClient, relay_client::RelayEvent, session_store, - AccountClient, AccountSession, DeviceIdentity, RemoteServer, + validate_relay_base_url, AccountClient, AccountSession, DeviceIdentity, RemoteServer, }; -/// In-memory account session (token + master key). Lost on restart. -static ACCOUNT_SESSION: OnceLock>>> = OnceLock::new(); +#[derive(Clone)] +struct AccountContextState { + session: AccountSession, + relay_url: String, +} + +/// Session and relay URL are one atomic account context so concurrent login, +/// logout, routing and sync cannot observe a torn pair. +static ACCOUNT_CONTEXT: OnceLock>>> = OnceLock::new(); +static ACCOUNT_CONTEXT_GENERATION: AtomicU64 = AtomicU64::new(1); +static ACCOUNT_CONTEXT_TRANSITIONS: AtomicUsize = AtomicUsize::new(0); +static ACCOUNT_SYNC_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +/// Serializes candidate credential verification without hiding or stopping the +/// currently active account. Only a fully authenticated candidate may enter +/// the account transition that replaces it. +static ACCOUNT_LOGIN_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +static ACCOUNT_CONTEXT_TRANSITION_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +static ACCOUNT_SYNC_CANCEL: OnceLock = OnceLock::new(); +/// At most one delayed daemon-exit recovery poller may own a generation. +/// A newer generation supersedes an older poller without accumulating tasks. +static ROUTING_RECOVERY_GENERATION: AtomicU64 = AtomicU64::new(0); +/// Read leases cover one routing event through its side effects and response. +/// Account transitions and routing-client ownership changes take the write +/// lease, so a new owner cannot be published while an old handler is active. +static DEVICE_ROUTING_LIFECYCLE: RwLock<()> = RwLock::const_new(()); + +pub(crate) fn account_context_generation() -> u64 { + ACCOUNT_CONTEXT_GENERATION.load(Ordering::Acquire) +} + +pub(crate) fn account_context_is_current(generation: u64) -> bool { + ACCOUNT_CONTEXT_TRANSITIONS.load(Ordering::Acquire) == 0 + && account_context_generation() == generation +} + +struct AccountContextTransitionPermit; + +impl AccountContextTransitionPermit { + fn begin() -> Self { + ACCOUNT_CONTEXT_TRANSITIONS.fetch_add(1, Ordering::AcqRel); + ACCOUNT_CONTEXT_GENERATION.fetch_add(1, Ordering::AcqRel); + account_sync_cancel().notify_waiters(); + Self + } +} -/// The relay URL associated with the current account session. -static ACCOUNT_RELAY_URL: OnceLock>>> = OnceLock::new(); +impl Drop for AccountContextTransitionPermit { + fn drop(&mut self) { + // Reject work queued during the transition before exposing the newly + // installed (or cleared) account context. + ACCOUNT_CONTEXT_GENERATION.fetch_add(1, Ordering::AcqRel); + ACCOUNT_CONTEXT_TRANSITIONS.fetch_sub(1, Ordering::AcqRel); + } +} + +struct AccountContextTransitionGuard { + sync_guard: Option>, + transition: Option, + routing_guard: Option>, + transition_guard: Option>, +} + +impl AccountContextTransitionGuard { + fn finish(mut self) -> u64 { + drop(self.sync_guard.take()); + drop(self.transition.take()); + let generation = account_context_generation(); + drop(self.routing_guard.take()); + drop(self.transition_guard.take()); + generation + } +} + +impl Drop for AccountContextTransitionGuard { + fn drop(&mut self) { + drop(self.sync_guard.take()); + drop(self.transition.take()); + drop(self.routing_guard.take()); + drop(self.transition_guard.take()); + } +} + +pub(crate) async fn lock_account_sync( + generation: u64, +) -> Result> { + let guard = ACCOUNT_SYNC_LOCK.lock().await; + if !account_context_is_current(generation) { + return Err(anyhow!("account sync cancelled")); + } + Ok(guard) +} + +fn account_sync_cancel() -> &'static Notify { + ACCOUNT_SYNC_CANCEL.get_or_init(Notify::new) +} + +pub(crate) async fn await_account_sync_current(generation: u64, future: F) -> Result +where + F: Future, +{ + let mut cancelled = Box::pin(account_sync_cancel().notified()); + cancelled.as_mut().enable(); + if !account_context_is_current(generation) { + return Err(anyhow!("account sync cancelled")); + } + tokio::select! { + _ = &mut cancelled => Err(anyhow!("account sync cancelled")), + result = future => { + if !account_context_is_current(generation) { + Err(anyhow!("account sync cancelled")) + } else { + Ok(result) + } + } + } +} + +async fn invalidate_and_wait_for_account_sync() -> AccountContextTransitionGuard { + let transition_guard = ACCOUNT_CONTEXT_TRANSITION_LOCK.lock().await; + let transition = AccountContextTransitionPermit::begin(); + let sync_guard = ACCOUNT_SYNC_LOCK.lock().await; + bitfun_core::service::remote_connect::settings_sync::wait_for_sync_operations_idle().await; + let routing_guard = DEVICE_ROUTING_LIFECYCLE.write().await; + AccountContextTransitionGuard { + sync_guard: Some(sync_guard), + transition: Some(transition), + routing_guard: Some(routing_guard), + transition_guard: Some(transition_guard), + } +} + +async fn invalidate_and_wait_if_account_current( + expected_generation: u64, +) -> Option { + let transition_guard = ACCOUNT_CONTEXT_TRANSITION_LOCK.lock().await; + if !account_context_is_current(expected_generation) { + return None; + } + let transition = AccountContextTransitionPermit::begin(); + let sync_guard = ACCOUNT_SYNC_LOCK.lock().await; + bitfun_core::service::remote_connect::settings_sync::wait_for_sync_operations_idle().await; + let routing_guard = DEVICE_ROUTING_LIFECYCLE.write().await; + Some(AccountContextTransitionGuard { + sync_guard: Some(sync_guard), + transition: Some(transition), + routing_guard: Some(routing_guard), + transition_guard: Some(transition_guard), + }) +} /// The background device-routing relay client. Holding this keeps the WS /// connection alive (the internal read/write tasks own the socket). Dropping it @@ -43,12 +189,31 @@ static TOKEN_EXPIRED: AtomicBool = AtomicBool::new(false); /// `account_login` / `account_finalize_login`). static PENDING_SYNC_CHOICE: AtomicBool = AtomicBool::new(false); -fn account_session() -> &'static Arc>> { - ACCOUNT_SESSION.get_or_init(|| Arc::new(RwLock::new(None))) +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct AutomaticAccountSyncPolicy { + pub(crate) background_engine: bool, + pub(crate) management_push: bool, } -fn account_relay_url() -> &'static Arc>> { - ACCOUNT_RELAY_URL.get_or_init(|| Arc::new(RwLock::new(None))) +fn automatic_account_sync_policy_for_pending( + pending_sync_choice: bool, +) -> AutomaticAccountSyncPolicy { + let allowed = !pending_sync_choice; + AutomaticAccountSyncPolicy { + background_engine: allowed, + management_push: allowed, + } +} + +/// Automatic sync must remain idle while an authenticated account is waiting +/// for the user to choose whether cloud or local settings should win. Explicit +/// first-login sync is intentionally not governed by this policy. +pub(crate) fn automatic_account_sync_policy() -> AutomaticAccountSyncPolicy { + automatic_account_sync_policy_for_pending(PENDING_SYNC_CHOICE.load(Ordering::Acquire)) +} + +fn account_context() -> &'static Arc>> { + ACCOUNT_CONTEXT.get_or_init(|| Arc::new(RwLock::new(None))) } fn device_relay_client() -> &'static RwLock>> { @@ -58,29 +223,61 @@ fn device_relay_client() -> &'static RwLock>> { /// Read both the session and relay URL, returning owned clones to avoid holding /// locks across awaits. pub(crate) async fn read_account_context() -> Result<(AccountSession, String)> { - let session = account_session().read().await.clone(); - let relay_url = account_relay_url().read().await.clone(); - match (session, relay_url) { - (Some(s), Some(u)) => Ok((s, u)), - _ => Err(anyhow!("not logged in")), + let generation = account_context_generation(); + read_account_context_for_generation(generation).await +} + +async fn read_account_context_raw() -> Result<(AccountSession, String)> { + account_context() + .read() + .await + .clone() + .map(|context| (context.session, context.relay_url)) + .ok_or_else(|| anyhow!("not logged in")) +} + +pub(crate) async fn read_account_context_for_generation( + generation: u64, +) -> Result<(AccountSession, String)> { + if !account_context_is_current(generation) { + return Err(anyhow!("account context changed")); + } + let context = read_account_context_raw().await?; + if !account_context_is_current(generation) { + return Err(anyhow!("account context changed")); } + Ok(context) } /// Whether an account session is currently held and login is finalized. /// Matches desktop `account_status`: pending cloud/local sync choice is not /// treated as logged in. pub(crate) async fn is_logged_in() -> bool { - if PENDING_SYNC_CHOICE.load(Ordering::Relaxed) { + if PENDING_SYNC_CHOICE.load(Ordering::Acquire) { return false; } - account_session().read().await.is_some() + read_account_context().await.is_ok() +} + +fn normalize_relay_url(relay_url: &str) -> Result { + let parsed = validate_relay_base_url(relay_url.trim())?; + Ok(parsed.as_str().trim_end_matches('/').to_string()) } /// Attempt to restore a persisted session from disk. Called at startup. /// Returns `Some(user_id)` if a session was restored. pub(crate) async fn try_restore_session() -> Option { + let _sync_guard = invalidate_and_wait_for_account_sync().await; match session_store::load_session_detailed() { Ok(Some(loaded)) => { + let relay_url = match normalize_relay_url(&loaded.relay_url) { + Ok(url) => url, + Err(error) => { + tracing::warn!("Ignoring invalid persisted relay URL: {error}"); + session_store::clear_session(); + return None; + } + }; let user_id = loaded.user_id.clone(); if let Some(device_id) = loaded.device_id.as_deref() { if let Err(e) = DeviceIdentity::adopt_account_device_id(device_id) { @@ -92,8 +289,7 @@ pub(crate) async fn try_restore_session() -> Option { user_id: user_id.clone(), master_key: loaded.master_key, }; - *account_session().write().await = Some(session); - *account_relay_url().write().await = Some(loaded.relay_url); + *account_context().write().await = Some(AccountContextState { session, relay_url }); tracing::info!("Restored account session for user {user_id}"); Some(user_id) } @@ -142,9 +338,10 @@ pub(crate) async fn login_with_credentials( username: &str, password: &str, ) -> Result { - let relay_url = relay_url.trim(); + let _login_guard = ACCOUNT_LOGIN_LOCK.lock().await; + let relay_url_input = relay_url.trim(); let username = username.trim(); - if relay_url.is_empty() { + if relay_url_input.is_empty() { return Err(anyhow!("Auth Server is required")); } if username.is_empty() { @@ -153,69 +350,129 @@ pub(crate) async fn login_with_credentials( if password.is_empty() { return Err(anyhow!("Password is required")); } + let relay_url = normalize_relay_url(relay_url_input)?; + let expected_generation = account_context_generation(); + if !account_context_is_current(expected_generation) { + return Err(anyhow!("account context changed")); + } let device = current_device_identity()?; let client = AccountClient::new(); let session = client - .login(relay_url, username, password, &device) + .login(&relay_url, username, password, &device) .await .map_err(|e| anyhow!("login failed: {e}"))?; - let has_cloud_settings = client - .fetch_settings(relay_url, &session) - .await - .unwrap_or(None) - .is_some(); + let has_cloud_settings = + match resolve_cloud_settings_probe(client.fetch_settings(&relay_url, &session).await) { + Ok(has_cloud_settings) => has_cloud_settings, + Err(error) => { + revoke_rejected_login_candidate(&client, &relay_url, &session).await; + return Err(error); + } + }; + + // A daemon is a separate process with its own in-memory session and WebSocket. + // Retire it after candidate authentication but before beginning the local + // generation transition. If retirement fails, the old local owner keeps + // its original generation and remains usable. A clean daemon exit is not + // auto-restarted by the generated launchd/systemd service definitions. + // Snapshot the old owner before the guarded replacement. A generation race + // rejects the transition below, in which case this snapshot is never used. + let previous_account_context = account_context().read().await.clone(); + let (retired_daemon, transition_guard) = match begin_candidate_account_transition( + expected_generation, + retire_running_daemon_for_account_switch().await, + ) + .await + { + Ok(transition) => transition, + Err(CandidateAccountTransitionError::DaemonRetirement(failure)) => { + let recovery_message = if failure.daemon_may_exit { + schedule_routing_recovery_after_daemon_exit( + expected_generation, + device.device_name.clone(), + ); + "; this CLI will restore local routing if the daemon exits" + } else { + "" + }; + revoke_rejected_login_candidate(&client, &relay_url, &session).await; + return Err(anyhow!( + "{}; the old account context and generation were preserved{}", + failure.error, + recovery_message + )); + } + Err(CandidateAccountTransitionError::AccountContextChanged) => { + revoke_rejected_login_candidate(&client, &relay_url, &session).await; + return Err(anyhow!("account context changed")); + } + }; + + // The transition owns the routing lifecycle write lease. Retire any + // in-process owner before making the candidate context observable. + clear_replaced_persisted_session(); + stop_device_routing_locked().await; let user_id = session.user_id.clone(); let device_name = device.device_name.clone(); let token = session.token.clone(); let master_key = session.master_key; - *account_session().write().await = Some(session); - *account_relay_url().write().await = Some(relay_url.to_string()); - session_store::save_credential_hint(username, relay_url); + *account_context().write().await = Some(AccountContextState { + session, + relay_url: relay_url.clone(), + }); + session_store::save_credential_hint(username, &relay_url); TOKEN_EXPIRED.store(false, Ordering::Relaxed); if has_cloud_settings { // Defer disk persist until the sync choice is accepted. Killing the // process during the choice panel must not restore a logged-in session. - PENDING_SYNC_CHOICE.store(true, Ordering::Relaxed); + PENDING_SYNC_CHOICE.store(true, Ordering::Release); + transition_guard.finish(); + revoke_replaced_account_context(&client, previous_account_context, &relay_url, &token) + .await; return Ok(LoginResult { user_id: user_id.clone(), - relay_url: relay_url.to_string(), + relay_url: relay_url.clone(), has_cloud_settings, status_message: format!( - "Authenticated as user {} on {}. Choose cloud or local settings to finish login.", - user_id, relay_url + "Authenticated as user {} on {}. Choose cloud or local settings to finish login.{}", + user_id, + relay_url, + if retired_daemon { + " The previous CLI daemon was stopped; routing will resume after the sync choice." + } else { + "" + } ), }); } - PENDING_SYNC_CHOICE.store(false, Ordering::Relaxed); + PENDING_SYNC_CHOICE.store(false, Ordering::Release); if let Err(e) = session_store::save_session_with_device( &token, &user_id, &master_key, - relay_url, + &relay_url, Some(device.device_id.as_str()), ) { tracing::warn!("Failed to persist session: {e}"); } - let routing_msg = if crate::daemon::is_daemon_running() { - // The daemon already holds the relay connection; a second connection - // from this process would flap the shared device_id registration. - " Device routing is handled by the running CLI daemon.".to_string() - } else { - match spawn_device_routing(relay_url, &device_name).await { - Ok(()) => " Device routing connected (Peer Host ready). Tip: `bitfun daemon install` keeps this device reachable after exit or reboot.".to_string(), - Err(e) => format!(" (Warning: device routing failed: {e})"), - } + let generation = transition_guard.finish(); + let routing_msg = match spawn_device_routing(&relay_url, &device_name, generation).await { + Ok(()) if retired_daemon => " The previous CLI daemon was stopped and routing is connected in this CLI process. Restart `bitfun daemon run` to restore always-on routing.".to_string(), + Ok(()) => " Device routing connected (Peer Host ready). Tip: `bitfun daemon install` keeps this device reachable after exit or reboot.".to_string(), + Err(e) if retired_daemon => format!(" (Warning: the previous CLI daemon was stopped, but replacement routing failed: {e})"), + Err(e) => format!(" (Warning: device routing failed: {e})"), }; + revoke_replaced_account_context(&client, previous_account_context, &relay_url, &token).await; Ok(LoginResult { user_id: user_id.clone(), - relay_url: relay_url.to_string(), + relay_url: relay_url.clone(), has_cloud_settings, status_message: format!( "Logged in as user {} on {}.{}", @@ -224,11 +481,152 @@ pub(crate) async fn login_with_credentials( }) } +async fn revoke_rejected_login_candidate( + client: &AccountClient, + relay_url: &str, + session: &AccountSession, +) { + if let Err(error) = client.revoke_token(relay_url, session).await { + tracing::warn!("Failed to revoke rejected login candidate token: {error}"); + } +} + +fn clear_replaced_persisted_session() { + // Once this candidate has won the transition, the old account must never + // be restored after a crash. A finalized replacement is persisted below; + // a pending cloud-sync choice intentionally leaves no restorable session. + session_store::clear_session(); +} + +fn replaced_account_revocation_target( + previous: Option, + replacement_relay_url: &str, + replacement_token: &str, +) -> Option { + previous.filter(|context| { + context.relay_url != replacement_relay_url || context.session.token != replacement_token + }) +} + +async fn revoke_replaced_account_context( + client: &AccountClient, + previous: Option, + replacement_relay_url: &str, + replacement_token: &str, +) { + let Some(previous) = + replaced_account_revocation_target(previous, replacement_relay_url, replacement_token) + else { + return; + }; + if let Err(error) = client + .revoke_token(&previous.relay_url, &previous.session) + .await + { + // B is already the committed in-memory owner. Relay cleanup of A is + // best-effort and must never roll the replacement back. + tracing::warn!("Failed to revoke replaced account token: {error}"); + } +} + +fn resolve_cloud_settings_probe(result: Result>) -> Result { + result.map(|settings| settings.is_some()).map_err(|error| { + anyhow!("could not check cloud settings: {error}; the current account remains active") + }) +} + +struct DaemonRetirementFailure { + error: anyhow::Error, + daemon_may_exit: bool, +} + +enum CandidateAccountTransitionError { + DaemonRetirement(DaemonRetirementFailure), + AccountContextChanged, +} + +async fn begin_candidate_account_transition( + expected_generation: u64, + daemon_retirement: std::result::Result, +) -> std::result::Result<(bool, AccountContextTransitionGuard), CandidateAccountTransitionError> { + let retired_daemon = + daemon_retirement.map_err(CandidateAccountTransitionError::DaemonRetirement)?; + let transition_guard = invalidate_and_wait_if_account_current(expected_generation) + .await + .ok_or(CandidateAccountTransitionError::AccountContextChanged)?; + Ok((retired_daemon, transition_guard)) +} + +async fn retire_running_daemon_for_account_switch( +) -> std::result::Result { + if !crate::daemon::is_daemon_running() { + return Ok(false); + } + if !crate::daemon::request_daemon_shutdown() { + return Err(DaemonRetirementFailure { + error: anyhow!("could not stop the CLI daemon; the current account remains active"), + daemon_may_exit: false, + }); + } + + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + while crate::daemon::is_daemon_running() { + if tokio::time::Instant::now() >= deadline { + return Err(DaemonRetirementFailure { + error: anyhow!( + "CLI daemon did not stop in time; the current account remains active" + ), + daemon_may_exit: true, + }); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + Ok(true) +} + +fn schedule_routing_recovery_after_daemon_exit(expected_generation: u64, device_name: String) { + if !account_context_is_current(expected_generation) + || ROUTING_RECOVERY_GENERATION.swap(expected_generation, Ordering::AcqRel) + == expected_generation + { + return; + } + tokio::spawn(async move { + while ROUTING_RECOVERY_GENERATION.load(Ordering::Acquire) == expected_generation + && account_context_is_current(expected_generation) + && crate::daemon::is_daemon_running() + { + tokio::time::sleep(Duration::from_millis(100)).await; + } + if ROUTING_RECOVERY_GENERATION.load(Ordering::Acquire) == expected_generation + && account_context_is_current(expected_generation) + && !crate::daemon::is_daemon_running() + { + if let Err(error) = restore_device_routing(&device_name).await { + tracing::warn!( + "Failed to restore old account routing after delayed daemon exit: {error}" + ); + } + } + let _ = ROUTING_RECOVERY_GENERATION.compare_exchange( + expected_generation, + 0, + Ordering::AcqRel, + Ordering::Acquire, + ); + }); +} + /// Persist the in-memory session after the user accepts the sync choice, then /// start device routing (same as a first login with no cloud settings). pub(crate) async fn finalize_login_after_sync_choice() -> Result<()> { + let generation = account_context_generation(); + let sync_guard = lock_account_sync(generation).await?; let device = current_device_identity()?; let (session, relay_url) = read_account_context().await?; + let retired_daemon = retire_running_daemon_for_account_switch() + .await + .map_err(|failure| failure.error)?; session_store::save_session_with_device( &session.token, &session.user_id, @@ -237,12 +635,15 @@ pub(crate) async fn finalize_login_after_sync_choice() -> Result<()> { Some(device.device_id.as_str()), ) .map_err(|e| anyhow!("persist session: {e}"))?; - PENDING_SYNC_CHOICE.store(false, Ordering::Relaxed); + PENDING_SYNC_CHOICE.store(false, Ordering::Release); - if crate::daemon::is_daemon_running() { - return Ok(()); + if retired_daemon { + tracing::info!( + "Stopped the previous CLI daemon before finalizing replacement account routing" + ); } - spawn_device_routing(&relay_url, &device.device_name) + drop(sync_guard); + spawn_device_routing(&relay_url, &device.device_name, generation) .await .map_err(|e| anyhow!("device routing failed: {e}")) } @@ -269,26 +670,27 @@ pub(crate) async fn account_info() -> Result { /// Public wrapper for restoring device routing after session restore at startup. pub(crate) async fn restore_device_routing(device_name: &str) -> Result<()> { - let relay_url = account_relay_url() - .read() - .await - .clone() - .ok_or_else(|| anyhow!("not logged in"))?; - spawn_device_routing(&relay_url, device_name).await + let generation = account_context_generation(); + let (_, relay_url) = read_account_context().await?; + spawn_device_routing(&relay_url, device_name, generation).await } /// Connect to the account relay for device-to-device routing and spawn the /// background task that handles incoming RPC commands. -async fn spawn_device_routing(relay_url: &str, device_name: &str) -> Result<()> { +async fn spawn_device_routing( + relay_url: &str, + device_name: &str, + account_generation: u64, +) -> Result<()> { + let _sync_guard = lock_account_sync(account_generation).await?; + let relay_url = normalize_relay_url(relay_url)?; // Tear down any previous connection first. stop_device_routing().await; - let session_guard = account_session().read().await; - let session = session_guard - .as_ref() - .ok_or_else(|| anyhow!("not logged in"))? - .clone(); - drop(session_guard); + let (session, current_relay_url) = read_account_context().await?; + if current_relay_url != relay_url { + return Err(anyhow!("account context changed")); + } let ws_url = format!( "{}/ws", @@ -302,19 +704,43 @@ async fn spawn_device_routing(relay_url: &str, device_name: &str) -> Result<()> client .connect_authenticated(&session.token, device_name) .await?; - let client_arc = Arc::new(client); - *device_relay_client().write().await = Some(client_arc.clone()); + { + let _routing_guard = DEVICE_ROUTING_LIFECYCLE.write().await; + let mut current_client = device_relay_client().write().await; + if !account_context_is_current(account_generation) { + drop(current_client); + client_arc.disconnect().await; + return Err(anyhow!("account context changed")); + } + *current_client = Some(client_arc.clone()); + } - let session_arc = account_session().clone(); + let account_context = account_context().clone(); let relay_client_arc = client_arc.clone(); tokio::spawn(async move { - while let Some(event) = event_rx.recv().await { - handle_relay_event(event, &session_arc, &relay_client_arc).await; - } - if is_current_routing_client(&relay_client_arc).await { - crate::peer_host::update_controller_presence(Vec::new()).await; + loop { + if !routing_loop_is_current(account_generation, &relay_client_arc).await { + tracing::debug!("Stopping stale device routing event loop"); + break; + } + let Some(event) = event_rx.recv().await else { + break; + }; + if !routing_loop_is_current(account_generation, &relay_client_arc).await { + tracing::debug!("Stopping stale device routing event loop"); + break; + } + handle_relay_event( + event, + &account_context, + &relay_client_arc, + account_generation, + &session.token, + ) + .await; } + retire_routing_client_if_same(&relay_client_arc).await; tracing::info!("Device routing event loop exited"); }); @@ -323,6 +749,12 @@ async fn spawn_device_routing(relay_url: &str, device_name: &str) -> Result<()> /// Disconnect the device-routing connection (if any). pub(crate) async fn stop_device_routing() { + let _routing_guard = DEVICE_ROUTING_LIFECYCLE.write().await; + stop_device_routing_locked().await; +} + +/// Stop routing while the caller holds the lifecycle write lease. +async fn stop_device_routing_locked() { let client = { device_relay_client().write().await.take() }; if let Some(client) = client { client.disconnect().await; @@ -331,31 +763,64 @@ pub(crate) async fn stop_device_routing() { } async fn is_current_routing_client(client: &Arc) -> bool { - device_relay_client() - .read() - .await - .as_ref() - .is_some_and(|current| Arc::ptr_eq(current, client)) + same_routing_client(device_relay_client().read().await.as_ref(), client) +} + +fn same_routing_client(current: Option<&Arc>, expected: &Arc) -> bool { + current.is_some_and(|client| Arc::ptr_eq(client, expected)) +} + +fn take_routing_client_if_same(current: &mut Option>, expected: &Arc) -> bool { + if !same_routing_client(current.as_ref(), expected) { + return false; + } + current.take(); + true +} + +/// Validate both halves of a routing-loop lease. The generation is checked +/// again after awaiting the client slot so a concurrent account transition +/// cannot make the pre-lock snapshot look current. +async fn routing_loop_is_current(account_generation: u64, relay_client: &Arc) -> bool { + if !account_context_is_current(account_generation) { + return false; + } + let matches = is_current_routing_client(relay_client).await; + matches && account_context_is_current(account_generation) +} + +/// Retire only the client owned by this loop. Keep the lifecycle write lease +/// while clearing controller presence so a replacement cannot publish its +/// presence and then have it erased by the old loop's cleanup. +async fn retire_routing_client_if_same(relay_client: &Arc) -> bool { + let _routing_guard = DEVICE_ROUTING_LIFECYCLE.write().await; + let mut current = device_relay_client().write().await; + if !take_routing_client_if_same(&mut current, relay_client) { + return false; + } + drop(current); + crate::peer_host::update_controller_presence(Vec::new()).await; + true } /// Log out: tear down routing, revoke the token (best-effort), clear state. pub(crate) async fn logout() -> Result<()> { - stop_device_routing().await; + let _sync_guard = invalidate_and_wait_for_account_sync().await; + stop_device_routing_locked().await; // Take the always-on daemon down with the account: the token is revoked // below, so leaving the daemon connected would keep this device online // with a doomed token until its next reconnect fails. if crate::daemon::request_daemon_shutdown() { tracing::info!("Signalled the CLI daemon to shut down after logout"); } - let result = read_account_context().await; + let result = read_account_context_raw().await; if let Ok((session, relay_url)) = result { let _ = AccountClient::new() .revoke_token(&relay_url, &session) .await; } - *account_session().write().await = None; - *account_relay_url().write().await = None; - PENDING_SYNC_CHOICE.store(false, Ordering::Relaxed); + *account_context().write().await = None; + PENDING_SYNC_CHOICE.store(false, Ordering::Release); session_store::clear_session(); session_store::clear_credential_hint(); TOKEN_EXPIRED.store(false, Ordering::Relaxed); @@ -365,145 +830,346 @@ pub(crate) async fn logout() -> Result<()> { /// Handle a single relay event for the device-routing loop. async fn handle_relay_event( event: RelayEvent, - session_arc: &Arc>>, + account_context: &Arc>>, relay_client: &Arc, + account_generation: u64, + expected_token: &str, ) { - if !is_current_routing_client(relay_client).await { - tracing::debug!("Ignoring event from a stale device routing client"); - return; - } - match event { - RelayEvent::AuthOk { user_id, device_id } => { - tracing::info!("Device routing auth ok: user={user_id} device={device_id}"); - if let Err(e) = DeviceIdentity::adopt_account_device_id(&device_id) { - tracing::warn!("Failed to adopt AuthOk device_id: {e}"); - } else if let Some(session) = session_arc.read().await.clone() { - if let Some(relay_url) = account_relay_url().read().await.clone() { - if let Err(e) = session_store::save_session_with_device( - &session.token, - &session.user_id, - &session.master_key, - &relay_url, - Some(device_id.as_str()), - ) { - tracing::warn!("Failed to persist AuthOk device_id into session: {e}"); - } - } - } - } + let event = match event { RelayEvent::AuthError { message } => { - tracing::warn!("Device routing auth error: {message}"); - TOKEN_EXPIRED.store(true, Ordering::Relaxed); - // Keep CLI/daemon semantics aligned with Desktop: a relay-rejected - // token is no longer a usable local login and must not be restored - // again on the next process start. Preserve the non-secret hint so - // the re-login form can still be prefilled. - relay_client.disconnect().await; - *session_arc.write().await = None; - *account_relay_url().write().await = None; - PENDING_SYNC_CHOICE.store(false, Ordering::Relaxed); - session_store::clear_session(); - crate::peer_host::update_controller_presence(Vec::new()).await; - } - RelayEvent::DevicePresence { devices } => { - tracing::info!("Device presence updated: {} online", devices.len()); - crate::peer_host::update_controller_presence( - devices.into_iter().map(|device| device.device_id).collect(), + handle_relay_auth_error( + message, + account_context, + relay_client, + account_generation, + expected_token, ) .await; + return; } - RelayEvent::DeviceMessageReceived { - source_device_id, - correlation_id, - encrypted_data, - nonce, - } => { - let session_guard = session_arc.read().await.clone(); - let Some(session) = session_guard else { - return; - }; - let plaintext = - match encryption::decrypt_from_base64(&session.master_key, &encrypted_data, &nonce) - { - Ok(p) => p, - Err(e) => { - tracing::warn!("Failed to decrypt device message: {e}"); - return; - } - }; - use remote_connect::remote_server::{RemoteCommand, RemoteResponse}; - let cmd: RemoteCommand = match serde_json::from_str(&plaintext) { - Ok(c) => c, - Err(e) => { - tracing::warn!("Could not parse device command: {e}"); - return; - } - }; - tracing::info!("Device command from {source_device_id}: {cmd:?} corr={correlation_id}"); + event => event, + }; - let response = match &cmd { - RemoteCommand::HostInvoke { command, args } => { - crate::peer_host::handle_host_invoke(command, args.clone()).await + let _routing_lease = DEVICE_ROUTING_LIFECYCLE.read().await; + if !routing_loop_is_current(account_generation, relay_client).await { + tracing::debug!("Ignoring event from a stale device routing client"); + return; + } + let fanout_owner = PeerFanoutOwner { + account_generation, + account_token: expected_token.to_string(), + relay_client: Arc::clone(relay_client), + }; + ACTIVE_PEER_FANOUT_OWNER + .scope(fanout_owner, async { + match event { + RelayEvent::AuthOk { user_id, device_id } => { + tracing::info!("Device routing auth ok: user={user_id} device={device_id}"); + if let Err(e) = DeviceIdentity::adopt_account_device_id(&device_id) { + tracing::warn!("Failed to adopt AuthOk device_id: {e}"); + } else if let Some(context) = account_context.read().await.clone() { + if routing_loop_is_current(account_generation, relay_client).await + && context.session.token == expected_token + { + if let Err(e) = session_store::save_session_with_device( + &context.session.token, + &context.session.user_id, + &context.session.master_key, + &context.relay_url, + Some(device_id.as_str()), + ) { + tracing::warn!( + "Failed to persist AuthOk device_id into session: {e}" + ); + } + } + } } - RemoteCommand::DeviceEvent { .. } => { - crate::peer_host::handle_device_event_command() + RelayEvent::AuthError { .. } => { + unreachable!("AuthError handled before routing read lease") } - other => { - let server = RemoteServer::new(session.master_key); - server.dispatch(other).await + RelayEvent::DevicePresence { devices } => { + tracing::info!("Device presence updated: {} online", devices.len()); + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + crate::peer_host::update_controller_presence( + devices.into_iter().map(|device| device.device_id).collect(), + ) + .await; + if !routing_loop_is_current(account_generation, relay_client).await { + tracing::debug!("Account changed while applying device presence"); + } } - }; + RelayEvent::DeviceMessageReceived { + source_device_id, + correlation_id, + encrypted_data, + nonce, + } => { + let context = account_context.read().await.clone(); + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + let Some(context) = context else { + return; + }; + if context.session.token != expected_token { + return; + } + let plaintext = match encryption::decrypt_from_base64( + &context.session.master_key, + &encrypted_data, + &nonce, + ) { + Ok(p) => p, + Err(e) => { + tracing::warn!("Failed to decrypt device message: {e}"); + return; + } + }; + use remote_connect::remote_server::{RemoteCommand, RemoteResponse}; + let cmd: RemoteCommand = match serde_json::from_str(&plaintext) { + Ok(c) => c, + Err(e) => { + tracing::warn!("Could not parse device command: {e}"); + return; + } + }; + tracing::info!( + "Device command from {source_device_id}: {cmd:?} corr={correlation_id}" + ); - let resp_json = match serde_json::to_string(&response) { - Ok(s) => s, - Err(e) => { - tracing::warn!("Failed to serialize RPC response: {e}"); - serde_json::to_string(&RemoteResponse::Error { - message: format!("failed to serialize RPC response: {e}"), - }) - .unwrap_or_else(|_| { - r#"{"resp":"error","message":"serialize failed"}"#.to_string() - }) - } - }; + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + let response = match &cmd { + RemoteCommand::HostInvoke { command, args } => { + let response = + crate::peer_host::handle_host_invoke(command, args.clone()).await; + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + response + } + RemoteCommand::DeviceEvent { .. } => { + crate::peer_host::handle_device_event_command() + } + other => { + let server = RemoteServer::new(context.session.master_key); + let response = server.dispatch(other).await; + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + response + } + }; - match encryption::encrypt_to_base64(&session.master_key, &resp_json) { - Ok((enc_resp, resp_nonce)) => { - // HTTP RPC bridge expects replies targeted at "rpc". - let reply_target = if source_device_id == "rpc" { - "rpc" - } else { - source_device_id.as_str() + let resp_json = match serde_json::to_string(&response) { + Ok(s) => s, + Err(e) => { + tracing::warn!("Failed to serialize RPC response: {e}"); + serde_json::to_string(&RemoteResponse::Error { + message: format!("failed to serialize RPC response: {e}"), + }) + .unwrap_or_else(|_| { + r#"{"resp":"error","message":"serialize failed"}"#.to_string() + }) + } }; - if let Err(e) = relay_client - .send_device_message(reply_target, &correlation_id, &enc_resp, &resp_nonce) - .await - { - tracing::warn!("Failed to send RPC response: {e}"); + + match encryption::encrypt_to_base64(&context.session.master_key, &resp_json) { + Ok((enc_resp, resp_nonce)) => { + // HTTP RPC bridge expects replies targeted at "rpc". + let reply_target = if source_device_id == "rpc" { + "rpc" + } else { + source_device_id.as_str() + }; + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + let send_result = relay_client + .send_device_message( + reply_target, + &correlation_id, + &enc_resp, + &resp_nonce, + ) + .await; + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + if let Err(e) = send_result { + tracing::warn!("Failed to send RPC response: {e}"); + } + } + Err(e) => { + tracing::warn!("Failed to encrypt RPC response: {e}"); + } + } + } + RelayEvent::Disconnected => { + tracing::info!("Device routing disconnected"); + if !routing_loop_is_current(account_generation, relay_client).await { + return; + } + crate::peer_host::update_controller_presence(Vec::new()).await; + if !routing_loop_is_current(account_generation, relay_client).await { + tracing::debug!("Account changed while clearing device presence"); } } - Err(e) => { - tracing::warn!("Failed to encrypt RPC response: {e}"); + RelayEvent::Reconnected => { + tracing::info!("Device routing reconnected"); } + RelayEvent::Error { message } => { + tracing::warn!("Device routing error: {message}"); + } + _ => {} } + }) + .await; +} + +/// Auth failure starts an account transition, which owns the lifecycle write +/// lease. It cannot be handled under the ordinary event read lease because +/// upgrading a Tokio `RwLock` would deadlock. +async fn handle_relay_auth_error( + message: String, + account_context: &Arc>>, + relay_client: &Arc, + account_generation: u64, + expected_token: &str, +) { + tracing::warn!("Device routing auth error: {message}"); + let Some(_transition_guard) = invalidate_and_wait_if_account_current(account_generation).await + else { + tracing::debug!("Ignoring auth error from a stale account generation"); + return; + }; + if !is_current_routing_client(relay_client).await { + tracing::debug!("Ignoring auth error from a replaced routing client"); + return; + } + let token_matches = account_context + .read() + .await + .as_ref() + .is_some_and(|context| context.session.token == expected_token); + if !token_matches || !is_current_routing_client(relay_client).await { + tracing::debug!("Ignoring auth error from a replaced routing client"); + return; + } + + // Keep CLI/daemon semantics aligned with Desktop: a relay-rejected token + // is no longer a usable local login and must not be restored again on the + // next process start. Preserve the non-secret hint for the re-login form. + relay_client.disconnect().await; + if !is_current_routing_client(relay_client).await { + tracing::debug!("Ignoring auth error cleanup for a replaced routing client"); + return; + } + + let mut current_client = device_relay_client().write().await; + if !same_routing_client(current_client.as_ref(), relay_client) { + tracing::debug!("Ignoring auth error cleanup for a replaced routing client"); + return; + } + let mut current_context = account_context.write().await; + if !current_context + .as_ref() + .is_some_and(|context| context.session.token == expected_token) + { + tracing::debug!("Ignoring auth error cleanup for a replaced account"); + return; + } + take_routing_client_if_same(&mut current_client, relay_client); + *current_context = None; + drop(current_context); + drop(current_client); + + TOKEN_EXPIRED.store(true, Ordering::Relaxed); + PENDING_SYNC_CHOICE.store(false, Ordering::Release); + session_store::clear_session(); + crate::peer_host::update_controller_presence(Vec::new()).await; +} + +/// Immutable routing owner captured when a Peer DeviceEvent enters the bounded +/// delivery queue. It prevents an event from account A being encrypted or sent +/// through account B after waiting behind older events. +#[derive(Clone)] +pub(crate) struct PeerFanoutOwner { + account_generation: u64, + account_token: String, + relay_client: Arc, +} + +tokio::task_local! { + static ACTIVE_PEER_FANOUT_OWNER: PeerFanoutOwner; +} + +pub(crate) fn inherited_peer_fanout_owner() -> Option { + ACTIVE_PEER_FANOUT_OWNER + .try_with(PeerFanoutOwner::clone) + .ok() +} + +impl PeerFanoutOwner { + fn matches(&self, generation: u64, token: &str, relay_client: &Arc) -> bool { + self.account_generation == generation + && self.account_token == token + && Arc::ptr_eq(&self.relay_client, relay_client) + } + + #[cfg(test)] + pub(crate) fn for_test(account_generation: u64, account_token: &str) -> Self { + let (relay_client, _) = RelayClient::new(); + Self { + account_generation, + account_token: account_token.to_string(), + relay_client: Arc::new(relay_client), } - RelayEvent::Disconnected => { - tracing::info!("Device routing disconnected"); - crate::peer_host::update_controller_presence(Vec::new()).await; - } - RelayEvent::Reconnected => { - tracing::info!("Device routing reconnected"); - } - RelayEvent::Error { message } => { - tracing::warn!("Device routing error: {message}"); - } - _ => {} + } + + #[cfg(test)] + pub(crate) fn generation_for_test(&self) -> u64 { + self.account_generation } } -/// Context needed by Peer Host DeviceEvent fan-out. -pub(crate) async fn peer_fanout_context() -> Result<(AccountSession, Arc)> { - let session = account_session() +/// Stable fan-out context. The read lease is intentionally retained through +/// encryption and all target sends; account replacement takes the write lease. +pub(crate) struct PeerFanoutLease { + pub(crate) session: AccountSession, + pub(crate) relay_client: Arc, + _routing_lease: tokio::sync::RwLockReadGuard<'static, ()>, +} + +pub(crate) async fn capture_peer_fanout_owner() -> Result { + let generation = account_context_generation(); + let _routing_lease = DEVICE_ROUTING_LIFECYCLE.read().await; + let (session, _) = read_account_context_for_generation(generation).await?; + let client = device_relay_client() + .read() + .await + .clone() + .ok_or_else(|| anyhow!("device routing not connected"))?; + if !account_context_is_current(generation) || !is_current_routing_client(&client).await { + return Err(anyhow!("account context changed")); + } + Ok(PeerFanoutOwner { + account_generation: generation, + account_token: session.token, + relay_client: client, + }) +} + +pub(crate) async fn acquire_peer_fanout_lease(owner: &PeerFanoutOwner) -> Result { + let routing_lease = DEVICE_ROUTING_LIFECYCLE.read().await; + if !account_context_is_current(owner.account_generation) { + return Err(anyhow!("queued Peer event account changed")); + } + let context = account_context() .read() .await .clone() @@ -513,7 +1179,20 @@ pub(crate) async fn peer_fanout_context() -> Result<(AccountSession, Arc Result> { }) .collect()) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Duration; + + use super::{ + account_context_generation, automatic_account_sync_policy_for_pending, + begin_candidate_account_transition, clear_replaced_persisted_session, + inherited_peer_fanout_owner, login_with_credentials, replaced_account_revocation_target, + resolve_cloud_settings_probe, take_routing_client_if_same, AccountContextState, + CandidateAccountTransitionError, DaemonRetirementFailure, PeerFanoutOwner, + ACCOUNT_LOGIN_LOCK, ACTIVE_PEER_FANOUT_OWNER, DEVICE_ROUTING_LIFECYCLE, + }; + + #[test] + fn stale_routing_loop_cannot_clear_replacement_client() { + let stale = Arc::new("stale"); + let replacement = Arc::new("replacement"); + let mut current = Some(Arc::clone(&replacement)); + + assert!(!take_routing_client_if_same(&mut current, &stale)); + assert!(current + .as_ref() + .is_some_and(|client| Arc::ptr_eq(client, &replacement))); + } + + #[test] + fn routing_loop_can_clear_only_its_own_client() { + let owned = Arc::new("owned"); + let mut current = Some(Arc::clone(&owned)); + + assert!(take_routing_client_if_same(&mut current, &owned)); + assert!(current.is_none()); + } + + #[tokio::test] + async fn routing_replacement_waits_for_in_flight_event_lease() { + let event_lease = DEVICE_ROUTING_LIFECYCLE.read().await; + let (attempting_tx, attempting_rx) = tokio::sync::oneshot::channel(); + let replacement = tokio::spawn(async move { + let _ = attempting_tx.send(()); + let _replacement_lease = DEVICE_ROUTING_LIFECYCLE.write().await; + }); + + attempting_rx.await.expect("replacement task started"); + tokio::task::yield_now().await; + assert!(!replacement.is_finished()); + + drop(event_lease); + tokio::time::timeout(Duration::from_secs(1), replacement) + .await + .expect("replacement should acquire the lifecycle after event completion") + .expect("replacement task should finish"); + } + + #[tokio::test] + async fn inherited_fanout_owner_does_not_reacquire_routing_read_lease() { + let event_lease = DEVICE_ROUTING_LIFECYCLE.read().await; + let (attempting_tx, attempting_rx) = tokio::sync::oneshot::channel(); + let replacement = tokio::spawn(async move { + let _ = attempting_tx.send(()); + let _replacement_lease = DEVICE_ROUTING_LIFECYCLE.write().await; + }); + + attempting_rx.await.expect("replacement task started"); + tokio::task::yield_now().await; + assert!(!replacement.is_finished()); + + let owner = PeerFanoutOwner::for_test(21, "token-a"); + let inherited = tokio::time::timeout( + Duration::from_millis(100), + ACTIVE_PEER_FANOUT_OWNER.scope(owner, async { inherited_peer_fanout_owner() }), + ) + .await + .expect("inherited owner lookup must not wait behind the queued writer") + .expect("task-local owner should be visible"); + assert_eq!(inherited.generation_for_test(), 21); + assert!(!replacement.is_finished()); + + drop(event_lease); + tokio::time::timeout(Duration::from_secs(1), replacement) + .await + .expect("replacement should proceed after the outer event lease is released") + .expect("replacement task should finish"); + } + + #[tokio::test] + async fn invalid_login_does_not_invalidate_the_current_account() { + let generation = account_context_generation(); + + let error = login_with_credentials("", "user", "password") + .await + .expect_err("empty relay URL must be rejected"); + + assert!(error.to_string().contains("Auth Server is required")); + assert_eq!(account_context_generation(), generation); + } + + #[test] + fn cloud_settings_probe_errors_are_not_treated_as_missing_settings() { + assert!(!resolve_cloud_settings_probe(Ok(None)).expect("missing settings is valid")); + assert!( + resolve_cloud_settings_probe(Ok(Some("encrypted settings".to_string()))) + .expect("existing settings is valid") + ); + + let error = resolve_cloud_settings_probe(Err(anyhow::anyhow!("relay unavailable"))) + .expect_err("probe failure must reject the candidate login"); + assert!(error.to_string().contains("could not check cloud settings")); + assert!(error.to_string().contains("relay unavailable")); + } + + #[test] + fn pending_sync_choice_blocks_automatic_pull_and_push_until_finalized() { + let pending = automatic_account_sync_policy_for_pending(true); + assert!(!pending.background_engine); + assert!(!pending.management_push); + + let finalized = automatic_account_sync_policy_for_pending(false); + assert!(finalized.background_engine); + assert!(finalized.management_push); + } + + #[tokio::test] + async fn daemon_retirement_failure_does_not_begin_account_transition() { + let generation = account_context_generation(); + let result = begin_candidate_account_transition( + generation, + Err(DaemonRetirementFailure { + error: anyhow::anyhow!("daemon stayed alive"), + daemon_may_exit: true, + }), + ) + .await; + + assert!(matches!( + result, + Err(CandidateAccountTransitionError::DaemonRetirement(_)) + )); + assert_eq!(account_context_generation(), generation); + } + + #[test] + fn pending_replacement_cannot_restore_the_previous_persisted_account() { + let directory = std::env::temp_dir().join(format!( + "bitfun-cli-account-session-{}", + uuid::Uuid::new_v4() + )); + bitfun_core::service::remote_connect::session_store::set_session_store_directory_for_test( + directory, + ); + let old_master_key = [7_u8; 32]; + bitfun_core::service::remote_connect::session_store::save_session_with_device( + "account-a-token", + "account-a", + &old_master_key, + "https://relay-a.example", + Some("device-a"), + ) + .expect("persist account A"); + assert_eq!( + bitfun_core::service::remote_connect::session_store::load_session_detailed() + .expect("load account A") + .expect("account A should be persisted") + .token, + "account-a-token" + ); + + // This is the disk step used after candidate B wins the transition and + // before B is exposed as awaiting its cloud/local sync choice. + clear_replaced_persisted_session(); + + assert!( + bitfun_core::service::remote_connect::session_store::load_session_detailed() + .expect("load after candidate B becomes pending") + .is_none() + ); + } + + #[test] + fn replacement_revokes_only_the_previous_distinct_token() { + let previous = AccountContextState { + session: bitfun_core::service::remote_connect::AccountSession { + token: "old-token".to_string(), + user_id: "same-account".to_string(), + master_key: [3_u8; 32], + }, + relay_url: "https://relay.example".to_string(), + }; + + let target = replaced_account_revocation_target( + Some(previous.clone()), + "https://relay.example", + "new-token", + ) + .expect("a new token for the same account must retire the old bearer"); + assert_eq!(target.session.token, "old-token"); + assert_eq!(target.session.user_id, "same-account"); + assert_eq!(target.relay_url, "https://relay.example"); + + assert!(replaced_account_revocation_target( + Some(previous.clone()), + "https://relay.example", + "old-token" + ) + .is_none()); + assert!(replaced_account_revocation_target( + Some(previous), + "https://other-relay.example", + "old-token" + ) + .is_some()); + assert!( + replaced_account_revocation_target(None, "https://relay.example", "new-token") + .is_none() + ); + } + + #[tokio::test] + async fn candidate_login_attempts_are_serialized() { + let first_candidate = ACCOUNT_LOGIN_LOCK.lock().await; + let (attempting_tx, attempting_rx) = tokio::sync::oneshot::channel(); + let second_candidate = tokio::spawn(async move { + let _ = attempting_tx.send(()); + let _guard = ACCOUNT_LOGIN_LOCK.lock().await; + }); + + attempting_rx.await.expect("second candidate started"); + tokio::task::yield_now().await; + assert!(!second_candidate.is_finished()); + + drop(first_candidate); + tokio::time::timeout(Duration::from_secs(1), second_candidate) + .await + .expect("second candidate should proceed after the first") + .expect("second candidate task should finish"); + } + + #[test] + fn queued_fanout_owner_requires_generation_token_and_client_identity() { + let owner = PeerFanoutOwner::for_test(11, "token-a"); + let owned_client = Arc::clone(&owner.relay_client); + let replacement = PeerFanoutOwner::for_test(12, "token-b"); + + assert!(owner.matches(11, "token-a", &owned_client)); + assert!(!owner.matches(12, "token-a", &owned_client)); + assert!(!owner.matches(11, "token-b", &owned_client)); + assert!(!owner.matches(11, "token-a", &replacement.relay_client)); + } +} diff --git a/src/apps/cli/src/account_sync.rs b/src/apps/cli/src/account_sync.rs index b42e30a1c2..47be5869fc 100644 --- a/src/apps/cli/src/account_sync.rs +++ b/src/apps/cli/src/account_sync.rs @@ -15,7 +15,10 @@ use bitfun_core::service::config::get_global_config_service; use bitfun_core::service::remote_connect::settings_sync; use bitfun_core::service::remote_connect::{sync_state, AccountClient}; -use crate::account::read_account_context; +use crate::account::{ + account_context_generation, account_context_is_current, automatic_account_sync_policy, + await_account_sync_current, lock_account_sync, read_account_context, +}; const UPLOAD_CONCURRENCY_CHUNK: usize = 5; @@ -27,8 +30,24 @@ const UPLOAD_CONCURRENCY_CHUNK: usize = 5; pub(crate) fn start_settings_sync_loop() { let hooks = settings_sync::SettingsSyncHooks { account_context: Some(Arc::new(|| { - Box::pin(async { read_account_context().await }) + Box::pin(async { + if !automatic_account_sync_policy().background_engine { + return Err(anyhow!("account login is awaiting a sync choice")); + } + let generation = account_context_generation(); + if !account_context_is_current(generation) { + return Err(anyhow!("account context is transitioning")); + } + let (account, relay_url) = read_account_context().await?; + if !automatic_account_sync_policy().background_engine + || !account_context_is_current(generation) + { + return Err(anyhow!("account context changed while reading")); + } + Ok((account, relay_url, generation)) + }) })), + is_account_context_current: Some(Arc::new(account_context_is_current)), on_settings_applied: Some(Arc::new(|| { crate::peer_host::notify_controllers_settings_changed(); })), @@ -51,11 +70,21 @@ pub(crate) fn notify_local_settings_changed() { /// (e.g. `bitfun models set-default`) where the sync loop never starts. /// Silently no-ops when logged out; failures are logged, not fatal. pub(crate) async fn push_settings_after_local_change() { + if !automatic_account_sync_policy().management_push { + return; + } // Management commands never restore the persisted account session into // memory — do it on demand so the push can authenticate. if read_account_context().await.is_err() { crate::account::try_restore_session().await; } + let generation = account_context_generation(); + let Ok(_sync_guard) = lock_account_sync(generation).await else { + return; + }; + if !automatic_account_sync_policy().management_push { + return; + } let Ok((account, relay_url)) = read_account_context().await else { return; }; @@ -208,6 +237,8 @@ pub(crate) async fn run_auto_sync( is_first_login: bool, workspace_path: &Path, ) -> Result { + let generation = account_context_generation(); + let _sync_guard = lock_account_sync(generation).await?; set_progress(|p| { *p = SyncProgress { status: SyncStatus::Syncing, @@ -232,24 +263,32 @@ pub(crate) async fn run_auto_sync( .map_err(|e| anyhow!("export config: {e}"))?; let config_json = serde_json::to_string(&exported).map_err(|e| anyhow!("serialize config: {e}"))?; - settings_sync::upload_settings_payload(&acct_session, &relay_url, &config_json) - .await - .map_err(|e| anyhow!("upload settings: {e}"))?; + await_account_sync_current( + generation, + settings_sync::upload_settings_payload(&acct_session, &relay_url, &config_json), + ) + .await? + .map_err(|e| anyhow!("upload settings: {e}"))?; emit_progress("settings_done", 15, None, None, None).await; true } else { emit_progress("downloading_settings", 5, None, None, None).await; - let cloud = client - .fetch_settings_with_version(&relay_url, &acct_session) - .await - .map_err(|e| anyhow!("fetch settings: {e}"))?; + let cloud = await_account_sync_current( + generation, + client.fetch_settings_with_version(&relay_url, &acct_session), + ) + .await? + .map_err(|e| anyhow!("fetch settings: {e}"))?; if let Some(blob) = cloud { emit_progress("applying_settings", 10, None, None, None).await; // Explicit user choice ("use cloud") — always apply, even when the // cursor says this device already has this version. - settings_sync::apply_settings_blob(&acct_session, &blob, true) - .await - .map_err(|e| anyhow!("apply cloud config: {e}"))?; + await_account_sync_current( + generation, + settings_sync::apply_settings_blob(&acct_session, &blob, true), + ) + .await? + .map_err(|e| anyhow!("apply cloud config: {e}"))?; emit_progress("settings_done", 15, None, None, None).await; true } else { @@ -278,6 +317,9 @@ pub(crate) async fn run_auto_sync( let mut sync_state_local = sync_state::load(&acct_session.user_id); let mut pending_uploads: Vec<(String, String, String)> = Vec::new(); for meta in local_sessions.iter() { + if !account_context_is_current(generation) { + return Err(anyhow!("account sync cancelled")); + } let turns = compatibility .load_persisted_session_turns(&storage_path, &meta.session_id, None) .await @@ -318,16 +360,18 @@ pub(crate) async fn run_auto_sync( let bundle_json = bundle_json.clone(); let hash = hash.clone(); handles.push(tokio::spawn(async move { - let result = client - .upload_session(&relay_url, &acct_session, &session_id, &bundle_json) - .await; + let result = await_account_sync_current( + generation, + client.upload_session(&relay_url, &acct_session, &session_id, &bundle_json), + ) + .await; (session_id, hash, result) })); } for handle in handles { let done_base = chunk_idx * UPLOAD_CONCURRENCY_CHUNK; match handle.await { - Ok((session_id, hash, Ok(version))) => { + Ok((session_id, hash, Ok(Ok(version)))) => { uploaded.push((session_id.clone(), hash, version)); let done = uploaded.len(); let percent = if upload_total == 0 { @@ -344,15 +388,19 @@ pub(crate) async fn run_auto_sync( ) .await; } - Ok((session_id, _, Err(e))) => { + Ok((session_id, _, Ok(Err(e)))) => { tracing::warn!("Auto-sync upload {session_id} failed: {e}"); let _ = done_base; } + Ok((_, _, Err(e))) => return Err(e), Err(e) => { tracing::warn!("Auto-sync upload task join failed: {e}"); } } } + if !account_context_is_current(generation) { + return Err(anyhow!("account sync cancelled")); + } } let exported = uploaded.len(); @@ -368,6 +416,8 @@ pub(crate) async fn run_auto_sync( } let _ = sync_state::save(&acct_session.user_id, &sync_state_local); + ensure_session_backup_complete(upload_total, exported)?; + tracing::info!("Auto-sync: settings={settings_synced} exported={exported} imported=0"); emit_progress("done", 100, Some(exported), Some(0), None).await; @@ -378,6 +428,15 @@ pub(crate) async fn run_auto_sync( }) } +fn ensure_session_backup_complete(total: usize, uploaded: usize) -> Result<()> { + if uploaded == total { + return Ok(()); + } + Err(anyhow!( + "session backup incomplete: uploaded {uploaded} of {total}; retry will resume remaining sessions" + )) +} + pub(crate) fn sync_phase_label(progress: &SyncProgress) -> String { match progress.phase.as_str() { "uploading_settings" => "Uploading settings…".into(), @@ -398,3 +457,17 @@ pub(crate) fn sync_phase_label(progress: &SyncProgress) -> String { other => other.to_string(), } } + +#[cfg(test)] +mod tests { + use super::ensure_session_backup_complete; + + #[test] + fn partial_session_backup_is_not_reported_as_success() { + assert!(ensure_session_backup_complete(4, 4).is_ok()); + assert!(ensure_session_backup_complete(4, 1) + .unwrap_err() + .to_string() + .contains("uploaded 1 of 4")); + } +} diff --git a/src/apps/cli/src/peer_host/fanout.rs b/src/apps/cli/src/peer_host/fanout.rs index eec31f4f8b..679e0730a9 100644 --- a/src/apps/cli/src/peer_host/fanout.rs +++ b/src/apps/cli/src/peer_host/fanout.rs @@ -10,12 +10,15 @@ use bitfun_core::service::remote_connect::remote_server::RemoteCommand; use bitfun_events::{project_agentic_frontend_event, AgenticEvent, ToolEventData}; use tokio::sync::{broadcast, mpsc}; +use crate::account::PeerFanoutOwner; + use super::control::{attached_controllers, controller_delivery_lease}; use super::state::{PeerHostState, PeerTurnKey}; const PEER_EVENT_DELIVERY_CAPACITY: usize = 512; struct QueuedPeerDeviceEvent { + owner: PeerFanoutOwner, targets: Vec, event: String, payload: serde_json::Value, @@ -24,8 +27,14 @@ struct QueuedPeerDeviceEvent { } impl QueuedPeerDeviceEvent { - fn new(targets: Vec, event: String, payload: serde_json::Value) -> Self { + fn new( + owner: PeerFanoutOwner, + targets: Vec, + event: String, + payload: serde_json::Value, + ) -> Self { Self { + owner, targets, event, payload, @@ -35,6 +44,7 @@ impl QueuedPeerDeviceEvent { } fn for_agent_event( + owner: PeerFanoutOwner, targets: Vec, event: String, payload: serde_json::Value, @@ -44,6 +54,7 @@ impl QueuedPeerDeviceEvent { ) -> Self { let terminal = terminal_turn.map(|turn| (turns.clone(), generation, turn)); Self { + owner, targets, event, payload, @@ -398,9 +409,13 @@ async fn handle_agentic_event(state: &PeerHostState, event: AgenticEvent) -> Res return Err("no attached Peer controller can receive Agent events".to_string()); } let generation = state.turns.current_event_stream_generation()?; + let owner = crate::account::capture_peer_fanout_owner() + .await + .map_err(|error| format!("Peer event routing owner unavailable: {error}"))?; enqueue_peer_device_event( peer_event_sender(), QueuedPeerDeviceEvent::for_agent_event( + owner, targets, projected.event_name, projected.payload, @@ -528,7 +543,27 @@ pub(crate) async fn fanout_peer_device_event(event: String, payload: serde_json: if targets.is_empty() { return; } - let queued = QueuedPeerDeviceEvent::new(targets, event, payload); + let inherited_owner = crate::account::inherited_peer_fanout_owner(); + let inherits_routing_lease = inherited_owner.is_some(); + let owner = match inherited_owner { + Some(owner) => owner, + None => match crate::account::capture_peer_fanout_owner().await { + Ok(owner) => owner, + Err(error) => { + tracing::debug!("Peer event fanout skipped before enqueue: {error}"); + return; + } + }, + }; + let queued = QueuedPeerDeviceEvent::new(owner, targets, event, payload); + if inherits_routing_lease { + // HostInvoke already holds the lifecycle read lease. Never await queue + // capacity or acquire a nested read here: a queued transition writer + // would otherwise create a writer-priority self-deadlock. A detached + // task preserves backpressure and validates the captured owner later. + enqueue_inherited_peer_device_event(peer_event_sender().clone(), queued); + return; + } if let Err(queued) = enqueue_peer_device_event(peer_event_sender(), queued).await { tracing::warn!( "Peer event delivery queue closed before accepting command event; using direct delivery" @@ -537,6 +572,33 @@ pub(crate) async fn fanout_peer_device_event(event: String, payload: serde_json: } } +fn enqueue_inherited_peer_device_event( + sender: mpsc::Sender, + queued: QueuedPeerDeviceEvent, +) { + match sender.try_send(queued) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(queued)) => { + tokio::spawn(async move { + if let Err(queued) = enqueue_peer_device_event(&sender, queued).await { + tracing::warn!( + "Peer event delivery queue closed while draining inherited routing event" + ); + fanout_peer_device_event_once(queued).await; + } + }); + } + Err(mpsc::error::TrySendError::Closed(queued)) => { + tokio::spawn(async move { + tracing::warn!( + "Peer event delivery queue closed for inherited routing event; using direct delivery" + ); + fanout_peer_device_event_once(queued).await; + }); + } + } +} + async fn enqueue_peer_device_event( sender: &mpsc::Sender, queued: QueuedPeerDeviceEvent, @@ -546,6 +608,7 @@ async fn enqueue_peer_device_event( async fn fanout_peer_device_event_once(queued: QueuedPeerDeviceEvent) { let QueuedPeerDeviceEvent { + owner, targets, event, payload, @@ -560,13 +623,15 @@ async fn fanout_peer_device_event_once(queued: QueuedPeerDeviceEvent) { return; } - let (session, relay_client) = match crate::account::peer_fanout_context().await { - Ok(ctx) => ctx, + let routing_lease = match crate::account::acquire_peer_fanout_lease(&owner).await { + Ok(lease) => lease, Err(error) => { - tracing::debug!("Peer event fanout skipped: {error}"); + tracing::debug!("Queued Peer event dropped after owner change: {error}"); return; } }; + let session = &routing_lease.session; + let relay_client = &routing_lease.relay_client; let envelope = match serde_json::to_string(&RemoteCommand::DeviceEvent { event, payload }) { Ok(envelope) => envelope, @@ -655,16 +720,21 @@ mod tests { use bitfun_events::{AgenticEvent, AgenticEventEnvelope, AgenticEventPriority}; use super::{ - continuity_is_current, drain_broadcast_receiver, enqueue_peer_device_event, event_turn_key, - interrupted_turn_failure_projection, retained_delivery_targets, QueuedPeerDeviceEvent, - TerminalDeliveryGuard, + continuity_is_current, drain_broadcast_receiver, enqueue_inherited_peer_device_event, + enqueue_peer_device_event, event_turn_key, interrupted_turn_failure_projection, + retained_delivery_targets, QueuedPeerDeviceEvent, TerminalDeliveryGuard, }; use crate::peer_host::state::{PeerTurnKey, PeerTurnTracker}; + fn test_owner(generation: u64) -> crate::account::PeerFanoutOwner { + crate::account::PeerFanoutOwner::for_test(generation, "test-token") + } + #[test] fn queued_events_keep_the_target_snapshot_from_enqueue_time() { let mut current_targets = vec!["controller-1".to_string()]; let queued = QueuedPeerDeviceEvent::new( + test_owner(7), current_targets.clone(), "dialog_turn_started".to_string(), serde_json::json!({}), @@ -714,6 +784,7 @@ mod tests { let (tx, rx) = tokio::sync::mpsc::channel(1); drop(rx); let queued = QueuedPeerDeviceEvent::new( + test_owner(7), vec!["controller-1".to_string()], "agentic://dialog-turn-failed".to_string(), serde_json::json!({ "turnId": "turn-1" }), @@ -725,6 +796,36 @@ mod tests { assert_eq!(recovered.event, "agentic://dialog-turn-failed"); } + #[tokio::test] + async fn inherited_enqueue_does_not_wait_for_full_delivery_queue() { + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + tx.send(QueuedPeerDeviceEvent::new( + test_owner(7), + vec!["controller-1".to_string()], + "first".to_string(), + serde_json::json!({}), + )) + .await + .expect("seed queue"); + + enqueue_inherited_peer_device_event( + tx, + QueuedPeerDeviceEvent::new( + test_owner(7), + vec!["controller-1".to_string()], + "second".to_string(), + serde_json::json!({}), + ), + ); + + assert_eq!(rx.recv().await.expect("first queued event").event, "first"); + let second = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv()) + .await + .expect("detached enqueue should complete after capacity is available") + .expect("second queued event"); + assert_eq!(second.event, "second"); + } + #[test] fn terminal_turn_stays_owned_until_delivery_completion() { let tracker = PeerTurnTracker::new(); @@ -780,6 +881,7 @@ mod tests { .current_event_stream_generation() .expect("ready generation"); let queued = QueuedPeerDeviceEvent::for_agent_event( + test_owner(7), vec!["controller-1".to_string()], "dialog_turn_started".to_string(), serde_json::json!({}), @@ -789,6 +891,7 @@ mod tests { ); assert!(continuity_is_current(&queued.continuity)); + assert_eq!(queued.owner.generation_for_test(), 7); turns.interrupt_event_stream(false); turns.mark_event_stream_ready(); assert!(!continuity_is_current(&queued.continuity)); diff --git a/src/apps/desktop/src/api/commands.rs b/src/apps/desktop/src/api/commands.rs index 8070eb70ce..6d0bfb5fbc 100644 --- a/src/apps/desktop/src/api/commands.rs +++ b/src/apps/desktop/src/api/commands.rs @@ -4367,6 +4367,13 @@ pub struct SubscriptionProviderRequest { pub provider: bitfun_core::infrastructure::subscription_auth::SubscriptionProvider, } +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SubscriptionLoginRequest { + pub provider: bitfun_core::infrastructure::subscription_auth::SubscriptionProvider, + pub session_id: String, +} + #[tauri::command] pub async fn list_subscription_accounts( ) -> Result, String> { @@ -4375,30 +4382,42 @@ pub async fn list_subscription_accounts( #[tauri::command] pub async fn start_subscription_login( - request: SubscriptionProviderRequest, + request: SubscriptionLoginRequest, ) -> Result { - bitfun_core::infrastructure::subscription_auth::start_login(request.provider) - .await - .map_err(|e| format!("Failed to start subscription login: {e:#}")) + bitfun_core::infrastructure::subscription_auth::start_login( + request.provider, + request.session_id, + ) + .await + .map_err(|e| format!("Failed to start subscription login: {e:#}")) } #[tauri::command] pub async fn get_subscription_login_status( - request: SubscriptionProviderRequest, + request: SubscriptionLoginRequest, ) -> Result { - Ok(bitfun_core::infrastructure::subscription_auth::login_status(request.provider).await) + bitfun_core::infrastructure::subscription_auth::login_status( + request.provider, + &request.session_id, + ) + .await + .map_err(|e| format!("Failed to get subscription login status: {e:#}")) } #[tauri::command] -pub async fn cancel_subscription_login(request: SubscriptionProviderRequest) -> Result<(), String> { - bitfun_core::infrastructure::subscription_auth::cancel_login(request.provider).await; - Ok(()) +pub async fn cancel_subscription_login(request: SubscriptionLoginRequest) -> Result<(), String> { + bitfun_core::infrastructure::subscription_auth::cancel_login( + request.provider, + &request.session_id, + ) + .await + .map_err(|e| format!("Failed to cancel subscription login: {e:#}")) } #[tauri::command] pub async fn logout_subscription_account( request: SubscriptionProviderRequest, -) -> Result<(), String> { +) -> Result { bitfun_core::infrastructure::subscription_auth::logout(request.provider) .await .map_err(|e| format!("Failed to logout subscription account: {e:#}")) diff --git a/src/apps/desktop/src/api/pages_api.rs b/src/apps/desktop/src/api/pages_api.rs index 98d4a6f135..63a7707462 100644 --- a/src/apps/desktop/src/api/pages_api.rs +++ b/src/apps/desktop/src/api/pages_api.rs @@ -1,14 +1,26 @@ //! BitFun Page Tauri commands (Save Version → Deploy). use bitfun_services_integrations::remote_connect::{ - delete_page_version_on_relay, deploy_page_version_on_relay, list_page_versions_from_relay, - list_pages_from_relay, publish_page_to_relay, save_page_version_to_relay, - unpublish_page_from_relay, update_page_on_relay, PageInfo, PagePublishResult, - PageSaveVersionResult, PageVersionInfo, + create_page_open_link_on_relay, delete_page_from_relay, delete_page_version_on_relay, + deploy_page_version_on_relay, list_page_versions_from_relay, list_pages_from_relay, + publish_page_to_relay, save_page_version_to_relay, unpublish_page_from_relay, + update_page_on_relay, PageInfo, PageOpenLink, PagePublishResult, PageSaveVersionResult, + PageVersionInfo, }; use serde::Deserialize; -use super::remote_connect_api::read_account_context; +use super::remote_connect_api::{ + account_context_generation, account_context_is_current, read_account_context_for_generation, +}; + +fn ensure_page_account_is_current(generation: u64) -> Result<(), String> { + if account_context_is_current(generation) { + Ok(()) + } else { + Err("Account changed while the Page operation was in progress; retry for the current account" + .to_string()) + } +} #[derive(Debug, Deserialize)] pub struct PagePublishRequest { @@ -22,11 +34,13 @@ pub struct PagePublishRequest { #[derive(Debug, Deserialize)] pub struct PageSlugRequest { pub slug: String, + pub generation: String, } #[derive(Debug, Deserialize)] pub struct PageUpdateRequest { pub slug: String, + pub generation: String, pub visibility: Option, pub title: Option, } @@ -34,22 +48,36 @@ pub struct PageUpdateRequest { #[derive(Debug, Deserialize)] pub struct PageDeployRequest { pub slug: String, + pub generation: String, pub version_id: String, } #[derive(Debug, Deserialize)] pub struct PageDeleteVersionRequest { pub slug: String, + pub generation: String, pub version_id: String, } +#[derive(Debug, Deserialize)] +pub struct PageOpenRequest { + pub slug: String, + pub generation: String, + pub version_id: Option, +} + /// Save a new immutable version (does not change production). #[tauri::command] pub async fn page_save_version( request: PagePublishRequest, ) -> Result { - let (session, relay_url) = read_account_context().await?; - save_page_version_to_relay( + // `directory` is explicitly a desktop-host path. Remote PagePublish tool + // calls are rejected from their typed workspace context before reaching + // this adapter; inferring ownership from the global remote-root registry + // would misclassify every local absolute path when an SSH root is `/`. + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = save_page_version_to_relay( &relay_url, &session.token, &request.directory, @@ -58,15 +86,17 @@ pub async fn page_save_version( request.title.as_deref(), request.note.as_deref(), ) - .await - .map_err(|e| e.to_string()) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) } /// Legacy alias for [`page_save_version`] (save only, no deploy). #[tauri::command] pub async fn page_publish(request: PagePublishRequest) -> Result { - let (session, relay_url) = read_account_context().await?; - publish_page_to_relay( + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = publish_page_to_relay( &relay_url, &session.token, &request.directory, @@ -74,70 +104,126 @@ pub async fn page_publish(request: PagePublishRequest) -> Result Result, String> { - let (session, relay_url) = read_account_context().await?; - list_pages_from_relay(&relay_url, &session.token) - .await - .map_err(|e| e.to_string()) + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = list_pages_from_relay(&relay_url, &session.token).await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) } #[tauri::command] pub async fn page_list_versions(request: PageSlugRequest) -> Result, String> { - let (session, relay_url) = read_account_context().await?; - list_page_versions_from_relay(&relay_url, &session.token, &request.slug) - .await - .map_err(|e| e.to_string()) + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = list_page_versions_from_relay( + &relay_url, + &session.token, + &request.slug, + &request.generation, + ) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn page_create_open_link(request: PageOpenRequest) -> Result { + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = create_page_open_link_on_relay( + &relay_url, + &session.token, + &request.slug, + request.version_id.as_deref(), + &request.generation, + ) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) } #[tauri::command] pub async fn page_deploy(request: PageDeployRequest) -> Result { - let (session, relay_url) = read_account_context().await?; - deploy_page_version_on_relay( + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = deploy_page_version_on_relay( &relay_url, &session.token, &request.slug, &request.version_id, + &request.generation, ) - .await - .map_err(|e| e.to_string()) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) } #[tauri::command] pub async fn page_delete_version(request: PageDeleteVersionRequest) -> Result<(), String> { - let (session, relay_url) = read_account_context().await?; - delete_page_version_on_relay( + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = delete_page_version_on_relay( &relay_url, &session.token, &request.slug, &request.version_id, + &request.generation, ) - .await - .map_err(|e| e.to_string()) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) } #[tauri::command] pub async fn page_update(request: PageUpdateRequest) -> Result { - let (session, relay_url) = read_account_context().await?; - update_page_on_relay( + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = update_page_on_relay( &relay_url, &session.token, &request.slug, + &request.generation, request.visibility.as_deref(), request.title.as_deref(), ) - .await - .map_err(|e| e.to_string()) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) } #[tauri::command] pub async fn page_unpublish(request: PageSlugRequest) -> Result<(), String> { - let (session, relay_url) = read_account_context().await?; - unpublish_page_from_relay(&relay_url, &session.token, &request.slug) - .await - .map_err(|e| e.to_string()) + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = unpublish_page_from_relay( + &relay_url, + &session.token, + &request.slug, + &request.generation, + ) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn page_delete(request: PageSlugRequest) -> Result<(), String> { + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let result = delete_page_from_relay( + &relay_url, + &session.token, + &request.slug, + &request.generation, + ) + .await; + ensure_page_account_is_current(generation)?; + result.map_err(|e| e.to_string()) } diff --git a/src/apps/desktop/src/api/peer_host_invoke.rs b/src/apps/desktop/src/api/peer_host_invoke.rs index a9cbfb40ff..95c1d4f3e7 100644 --- a/src/apps/desktop/src/api/peer_host_invoke.rs +++ b/src/apps/desktop/src/api/peer_host_invoke.rs @@ -43,6 +43,7 @@ static LOCAL_ONLY_COMMANDS: &[&str] = &[ // Account identity / peer mode control (stay on controller) "account_login", "account_finalize_login", + "account_cancel_pending_login", "account_logout", "account_status", "account_get_credential_hint", diff --git a/src/apps/desktop/src/api/remote_connect_api.rs b/src/apps/desktop/src/api/remote_connect_api.rs index 92cbc241e8..ec5bca3bfc 100644 --- a/src/apps/desktop/src/api/remote_connect_api.rs +++ b/src/apps/desktop/src/api/remote_connect_api.rs @@ -11,33 +11,300 @@ use bitfun_core::service::remote_connect::session_store::{ }; use bitfun_core::service::remote_connect::{ bot::{self, weixin, BotConfig}, - lan, session_store, sync_state, AccountClient, AccountSession, ConnectionMethod, - ConnectionResult, DeviceIdentity, PairingState, RemoteConnectConfig, RemoteConnectService, + lan, session_store, sync_state, AccountClient, AccountPairingVerification, AccountSession, + ConnectionMethod, ConnectionResult, DelegatedIdentityAuthorization, DeviceIdentity, + PairingState, RemoteConnectConfig, RemoteConnectService, }; use bitfun_core::service::session::{DialogTurnData, SessionMetadata}; -use bitfun_services_integrations::remote_connect::account::error_indicates_expired_token; +use bitfun_services_integrations::remote_connect::account::{ + error_indicates_expired_token, validate_relay_base_url, +}; use bitfun_services_integrations::remote_connect::{ - deploy_page_version_on_relay, join_relay_url, publish_page_content_on_relay, + deploy_page_version_on_relay, join_relay_url, list_pages_from_relay, + publish_page_content_on_relay, }; use futures::stream::{self, StreamExt}; use regex::Regex; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::future::Future; use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; use tauri::{AppHandle, Emitter, State}; -use tokio::sync::RwLock; +use tokio::sync::{Notify, RwLock}; static REMOTE_CONNECT_SERVICE: OnceLock>>> = OnceLock::new(); -/// In-memory account session (token + master key). The master key is never -/// persisted to disk; it is lost on restart and re-derived on next login. -static ACCOUNT_SESSION: OnceLock>>> = OnceLock::new(); +/// Session and relay URL must move together. Keeping them behind one lock +/// prevents a request from observing a token from one login and the URL from +/// another while an account transition is in progress. +#[derive(Clone)] +struct AccountContextState { + session: AccountSession, + relay_url: String, +} + +static ACCOUNT_CONTEXT: OnceLock>>> = OnceLock::new(); + +/// Serializes explicit login-time syncs and lets logout/new login invalidate +/// the active operation before account state is changed. +static ACCOUNT_AUTO_SYNC_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +/// Serializes credential verification attempts without hiding or disconnecting +/// the currently active account. A successful candidate acquires the account +/// transition guard only after all login-time network requests complete. +static ACCOUNT_LOGIN_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +static ACCOUNT_CONTEXT_TRANSITION_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +static ACCOUNT_AUTO_SYNC_CANCEL: OnceLock = OnceLock::new(); +static ACTIVE_ACCOUNT_AUTO_SYNC_OPERATION_ID: AtomicU64 = AtomicU64::new(0); +static ACCOUNT_CONTEXT_GENERATION: AtomicU64 = AtomicU64::new(1); +static ACCOUNT_CONTEXT_TRANSITIONS: AtomicUsize = AtomicUsize::new(0); + +/// Device-routing effects take a read lease; connection replacement and +/// teardown take the write lease. Together with `DeviceRoutingOwner`, this +/// prevents a retiring event loop from dispatching through a newer socket. +static DEVICE_ROUTING_LIFECYCLE_LOCK: tokio::sync::RwLock<()> = tokio::sync::RwLock::const_new(()); +static DEVICE_ROUTING_CONNECTION_ID: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone, Debug, Eq, PartialEq)] +struct DeviceRoutingOwner { + account_generation: u64, + account_token: String, + connection_id: u64, + service_connection_id: u64, +} + +#[derive(Default)] +struct DeviceRoutingState { + owner: Option, + online_devices: Vec, +} + +static DEVICE_ROUTING_STATE: OnceLock> = OnceLock::new(); + +pub(crate) fn account_context_generation() -> u64 { + ACCOUNT_CONTEXT_GENERATION.load(Ordering::Acquire) +} + +pub(crate) fn account_context_is_current(generation: u64) -> bool { + ACCOUNT_CONTEXT_TRANSITIONS.load(Ordering::Acquire) == 0 + && account_context_generation() == generation +} + +struct AccountContextTransitionPermit; + +impl AccountContextTransitionPermit { + fn begin() -> Self { + ACCOUNT_CONTEXT_TRANSITIONS.fetch_add(1, Ordering::AcqRel); + ACCOUNT_CONTEXT_GENERATION.fetch_add(1, Ordering::AcqRel); + ACTIVE_ACCOUNT_AUTO_SYNC_OPERATION_ID.store(0, Ordering::Release); + clear_last_finalized_pending_login(); + account_auto_sync_cancel().notify_waiters(); + Self + } +} -/// The relay URL associated with the current account session (needed for sync -/// and device-routing calls). -static ACCOUNT_RELAY_URL: OnceLock>>> = OnceLock::new(); +impl Drop for AccountContextTransitionPermit { + fn drop(&mut self) { + // Invalidate work queued during the transition before making the newly + // installed (or cleared) account context discoverable. + ACCOUNT_CONTEXT_GENERATION.fetch_add(1, Ordering::AcqRel); + ACCOUNT_CONTEXT_TRANSITIONS.fetch_sub(1, Ordering::AcqRel); + } +} + +struct AccountContextTransitionGuard { + sync_guard: Option>, + transition: Option, + transition_guard: Option>, +} + +impl AccountContextTransitionGuard { + /// Make the committed context observable while retaining the transition + /// mutex. Login-state listeners can now probe `account_status`, while a + /// competing logout or replacement remains blocked until publication ends. + fn make_context_observable(&mut self) { + drop(self.sync_guard.take()); + drop(self.transition.take()); + } +} + +struct PendingLoginFinalizeGuard { + sync_guard: Option>, + transition_guard: Option>, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct FinalizedPendingLoginOwner { + pending_login_id: String, + account_generation: u64, + account_token: String, +} + +impl Drop for PendingLoginFinalizeGuard { + fn drop(&mut self) { + drop(self.sync_guard.take()); + drop(self.transition_guard.take()); + } +} + +impl Drop for AccountContextTransitionGuard { + fn drop(&mut self) { + // Release the operation lock before reopening context discovery. A + // queued operation that wins this handoff still fails on the gate. + drop(self.sync_guard.take()); + drop(self.transition.take()); + drop(self.transition_guard.take()); + } +} + +async fn lock_account_sync( + generation: u64, +) -> Result, String> { + let guard = ACCOUNT_AUTO_SYNC_LOCK.lock().await; + if !account_context_is_current(generation) { + return Err("account sync cancelled".to_string()); + } + Ok(guard) +} + +fn ensure_account_auto_sync_current(operation_id: u64) -> Result<(), String> { + if operation_id != 0 + && ACTIVE_ACCOUNT_AUTO_SYNC_OPERATION_ID.load(Ordering::Acquire) == operation_id + { + Ok(()) + } else { + Err("account sync cancelled".to_string()) + } +} + +fn account_auto_sync_cancel() -> &'static Notify { + ACCOUNT_AUTO_SYNC_CANCEL.get_or_init(Notify::new) +} + +async fn await_account_auto_sync(operation_id: u64, future: F) -> Result +where + F: Future, +{ + let mut cancelled = Box::pin(account_auto_sync_cancel().notified()); + cancelled.as_mut().enable(); + ensure_account_auto_sync_current(operation_id)?; + tokio::select! { + _ = &mut cancelled => Err("account sync cancelled".to_string()), + result = future => { + ensure_account_auto_sync_current(operation_id)?; + Ok(result) + } + } +} + +async fn cancel_and_wait_for_account_auto_sync() -> AccountContextTransitionGuard { + // Serialize transition creation so a stale invalidation can re-check its + // generation before it makes the current account undiscoverable. + let transition_guard = ACCOUNT_CONTEXT_TRANSITION_LOCK.lock().await; + let transition = AccountContextTransitionPermit::begin(); + let sync_guard = ACCOUNT_AUTO_SYNC_LOCK.lock().await; + bitfun_core::service::remote_connect::settings_sync::wait_for_sync_operations_idle().await; + AccountContextTransitionGuard { + sync_guard: Some(sync_guard), + transition: Some(transition), + transition_guard: Some(transition_guard), + } +} + +async fn cancel_and_wait_if_account_current( + expected_generation: u64, +) -> Option { + let transition_guard = ACCOUNT_CONTEXT_TRANSITION_LOCK.lock().await; + if !account_context_is_current(expected_generation) { + return None; + } + let transition = AccountContextTransitionPermit::begin(); + let sync_guard = ACCOUNT_AUTO_SYNC_LOCK.lock().await; + bitfun_core::service::remote_connect::settings_sync::wait_for_sync_operations_idle().await; + Some(AccountContextTransitionGuard { + sync_guard: Some(sync_guard), + transition: Some(transition), + transition_guard: Some(transition_guard), + }) +} + +fn pending_login_is_owned_by(expected_pending_login_id: &str) -> bool { + if expected_pending_login_id.is_empty() + || !PENDING_SYNC_CHOICE.load(std::sync::atomic::Ordering::Acquire) + { + return false; + } + PENDING_LOGIN_ID + .get_or_init(|| std::sync::Mutex::new(None)) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_deref() + == Some(expected_pending_login_id) +} + +fn set_pending_login_id(pending_login_id: Option) { + let mut current = PENDING_LOGIN_ID + .get_or_init(|| std::sync::Mutex::new(None)) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *current = pending_login_id; + PENDING_SYNC_CHOICE.store(current.is_some(), std::sync::atomic::Ordering::Release); +} + +fn background_account_sync_is_allowed() -> bool { + !PENDING_SYNC_CHOICE.load(std::sync::atomic::Ordering::Acquire) +} + +fn clear_last_finalized_pending_login() { + *LAST_FINALIZED_PENDING_LOGIN + .get_or_init(|| std::sync::Mutex::new(None)) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = None; +} + +fn record_finalized_pending_login(owner: FinalizedPendingLoginOwner) { + *LAST_FINALIZED_PENDING_LOGIN + .get_or_init(|| std::sync::Mutex::new(None)) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(owner); +} + +async fn finalized_pending_login_is_current(pending_login_id: &str) -> bool { + let owner = LAST_FINALIZED_PENDING_LOGIN + .get_or_init(|| std::sync::Mutex::new(None)) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone(); + let Some(owner) = owner else { + return false; + }; + owner.pending_login_id == pending_login_id + && account_context_matches(owner.account_generation, &owner.account_token).await +} + +async fn lock_pending_login_for_finalize( + expected_pending_login_id: &str, +) -> Result { + let transition_guard = ACCOUNT_CONTEXT_TRANSITION_LOCK.lock().await; + let generation = account_context_generation(); + if !account_context_is_current(generation) + || !pending_login_is_owned_by(expected_pending_login_id) + { + return Err("pending login changed".to_string()); + } + let sync_guard = ACCOUNT_AUTO_SYNC_LOCK.lock().await; + if !account_context_is_current(generation) + || !pending_login_is_owned_by(expected_pending_login_id) + { + return Err("pending login changed".to_string()); + } + Ok(PendingLoginFinalizeGuard { + sync_guard: Some(sync_guard), + transition_guard: Some(transition_guard), + }) +} /// Global handle to the DialogScheduler, set during app startup. Used by the /// device-routing background task to execute commands received from peer @@ -97,8 +364,39 @@ fn emit_settings_applied() { ); } +async fn disconnect_peer_controllers(reason: &'static str) { + let request_ids = crate::api::peer_host_invoke::disconnect_controllers(); + if let Err(error) = + crate::api::peer_host_invoke::fail_closed_permission_requests(request_ids, reason).await + { + log::warn!("Peer permission requests were not fully cancelled: {error}"); + } + emit_device_presence(&[]); +} + +/// Retire the active routing owner before touching the shared service. The +/// lifecycle write lease ensures no retiring event handler can cross this +/// boundary and dispatch through a subsequently installed connection. +async fn stop_and_clear_device_routing(reason: &'static str) { + let _lifecycle = DEVICE_ROUTING_LIFECYCLE_LOCK.write().await; + clear_device_routing_state(); + if let Some(service) = get_service_holder().read().await.as_ref() { + service.stop_device_connection().await; + } + disconnect_peer_controllers(reason).await; +} + +async fn finish_device_routing_event_loop(owner: &DeviceRoutingOwner) { + let _lifecycle = DEVICE_ROUTING_LIFECYCLE_LOCK.write().await; + if !clear_device_routing_if_owner(owner) { + return; + } + disconnect_peer_controllers("Peer device-routing stream closed").await; +} + /// Emit granular auto-sync progress for the account login / devices UI. fn emit_sync_progress( + operation_id: u64, phase: &str, percent: u8, current: Option, @@ -108,6 +406,7 @@ fn emit_sync_progress( emit_account_event( "account://sync-progress", serde_json::json!({ + "operation_id": operation_id, "phase": phase, "percent": percent.min(100), "current": current, @@ -126,45 +425,62 @@ pub fn fanout_peer_device_event(event: String, payload: serde_json::Value) { if crate::api::peer_host_invoke::attached_controllers().is_empty() { return; } + let Some(routing_owner) = current_device_routing_owner_snapshot() else { + return; + }; let tx = PEER_EVENT_FANOUT_TX.get_or_init(|| { - let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<(String, serde_json::Value)>(); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); tokio::spawn(async move { - while let Some((event, payload)) = rx.recv().await { - fanout_peer_device_event_once(event, payload).await; + while let Some(item) = rx.recv().await { + fanout_peer_device_event_once(item).await; } }); tx }); - if let Err(e) = tx.send((event, payload)) { + if let Err(e) = tx.send(PeerEventFanoutItem { + routing_owner, + event, + payload, + }) { log::debug!("peer event fanout queue closed: {e}"); } } -static PEER_EVENT_FANOUT_TX: OnceLock< - tokio::sync::mpsc::UnboundedSender<(String, serde_json::Value)>, -> = OnceLock::new(); +struct PeerEventFanoutItem { + routing_owner: DeviceRoutingOwner, + event: String, + payload: serde_json::Value, +} -async fn fanout_peer_device_event_once(event: String, payload: serde_json::Value) { +static PEER_EVENT_FANOUT_TX: OnceLock> = + OnceLock::new(); + +async fn fanout_peer_device_event_once(item: PeerEventFanoutItem) { + let Some(_routing_effect) = lock_current_device_routing(&item.routing_owner).await else { + return; + }; let targets = crate::api::peer_host_invoke::attached_controllers(); if targets.is_empty() { return; } - let (session, _) = match read_account_context().await { - Ok(ctx) => ctx, - Err(e) => { - log::debug!("peer event fanout skipped (no account): {e}"); - return; - } - }; - let holder = get_service_holder().read().await; - let Some(ref service) = *holder else { + let (session, _) = + match read_account_context_for_generation(item.routing_owner.account_generation).await { + Ok(ctx) => ctx, + Err(e) => { + log::debug!("peer event fanout skipped (no account): {e}"); + return; + } + }; + if session.token != item.routing_owner.account_token + || !device_routing_owner_is_current(&item.routing_owner).await + { return; - }; + } use bitfun_core::service::remote_connect::encryption::encrypt_to_base64; use bitfun_core::service::remote_connect::remote_server::RemoteCommand; let envelope = match serde_json::to_string(&RemoteCommand::DeviceEvent { - event: event.clone(), - payload, + event: item.event.clone(), + payload: item.payload, }) { Ok(s) => s, Err(e) => { @@ -181,9 +497,14 @@ async fn fanout_peer_device_event_once(event: String, payload: serde_json::Value }; for target in targets { let correlation_id = uuid::Uuid::new_v4().to_string(); - if let Err(e) = service - .send_device_message(&target, &correlation_id, &encrypted_data, &nonce) - .await + if let Err(e) = send_device_message_with_routing_lease( + &item.routing_owner, + &target, + &correlation_id, + &encrypted_data, + &nonce, + ) + .await { log::debug!("peer event fanout to {target} failed: {e}"); } @@ -247,12 +568,49 @@ pub fn wrap_peer_aware_emitter( Arc::new(PeerAwareEmitter::new(inner)) } +async fn send_device_message_with_routing_lease( + owner: &DeviceRoutingOwner, + target_device_id: &str, + correlation_id: &str, + encrypted_data: &str, + nonce: &str, +) -> Result<(), String> { + if !device_routing_owner_is_current(owner).await { + return Err("device routing changed".to_string()); + } + let holder = get_service_holder().read().await; + if !device_routing_owner_is_current(owner).await { + return Err("device routing changed".to_string()); + } + let service = holder + .as_ref() + .ok_or_else(|| "remote connect service not initialized".to_string())?; + let sent = service + .send_device_message_if_connection( + owner.service_connection_id, + target_device_id, + correlation_id, + encrypted_data, + nonce, + ) + .await + .map_err(|error| error.to_string())?; + if !sent || !device_routing_owner_is_current(owner).await { + return Err("device routing changed".to_string()); + } + Ok(()) +} + /// Encrypt and send an RPC response (or error) back to the HTTP caller via relay. async fn send_rpc_envelope( + owner: &DeviceRoutingOwner, session: &AccountSession, correlation_id: &str, resp_value: serde_json::Value, ) { + if !device_routing_owner_is_current(owner).await { + return; + } let resp_json = match serde_json::to_string(&resp_value) { Ok(s) => s, Err(e) => { @@ -267,14 +625,16 @@ async fn send_rpc_envelope( use bitfun_core::service::remote_connect::encryption::encrypt_to_base64; match encrypt_to_base64(&session.master_key, &resp_json) { Ok((enc_resp, resp_nonce)) => { - let holder = get_service_holder().read().await; - if let Some(ref svc) = *holder { - if let Err(e) = svc - .send_device_message("rpc", correlation_id, &enc_resp, &resp_nonce) - .await - { - log::warn!("RPC: send response failed: {e}"); - } + if let Err(e) = send_device_message_with_routing_lease( + owner, + "rpc", + correlation_id, + &enc_resp, + &resp_nonce, + ) + .await + { + log::warn!("RPC: send response failed: {e}"); } } Err(e) => { @@ -284,11 +644,13 @@ async fn send_rpc_envelope( } async fn send_rpc_error( + owner: &DeviceRoutingOwner, session: &AccountSession, correlation_id: &str, message: impl Into, ) { send_rpc_envelope( + owner, session, correlation_id, serde_json::json!({ @@ -314,24 +676,45 @@ fn is_token_expired_error(e: &anyhow::Error) -> bool { error_indicates_expired_token(&e.to_string()) } -/// Drop the local account session after the relay rejects the token. +/// Drop the local account session after the relay rejects the token, but only +/// if the response still belongs to the same account generation and token. /// Keeps the username/relay hint so the login form can be prefilled. -async fn invalidate_local_account_session(reason: &str) { +async fn invalidate_local_account_session_if_current( + expected_generation: u64, + expected_token: &str, + reason: &str, +) -> bool { + let Some(_transition_guard) = cancel_and_wait_if_account_current(expected_generation).await + else { + log::info!("Ignored auth failure from a stale account generation"); + return false; + }; + let token_matches = get_account_context() + .read() + .await + .as_ref() + .is_some_and(|context| context.session.token == expected_token); + if !token_matches { + log::info!("Ignored auth failure from a replaced account token"); + return false; + } + sync_account_login_capability(false); TOKEN_EXPIRED.store(true, std::sync::atomic::Ordering::Relaxed); + stop_and_clear_device_routing("Account session expired").await; if let Some(service) = get_service_holder().read().await.as_ref() { - service.stop_device_connection().await; service.clear_account_pairing_context().await; service.clear_trusted_mobile_identity().await; + service.clear_bot_delegated_identities().await; } - *get_account_session().write().await = None; - *get_account_relay_url().write().await = None; + *get_account_context().write().await = None; + set_pending_login_id(None); session_store::clear_session(); - sync_account_login_capability(false); emit_account_event( "account://login-state", serde_json::json!({ "logged_in": false, "reason": "session_expired" }), ); log::warn!("Invalidated local account session after relay auth failure: {reason}"); + true } fn sync_account_login_capability(logged_in: bool) { @@ -341,10 +724,28 @@ fn sync_account_login_capability(logged_in: bool) { fn register_page_deploy_host() { set_page_deploy_handler(Arc::new(|slug, version_id| { Box::pin(async move { - let (session, relay_url) = read_account_context().await?; - let info = deploy_page_version_on_relay(&relay_url, &session.token, &slug, &version_id) + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; + let page = list_pages_from_relay(&relay_url, &session.token) .await - .map_err(|e| e.to_string())?; + .map_err(|e| e.to_string())? + .into_iter() + .find(|page| page.slug == slug) + .ok_or_else(|| { + "Page not found; refresh the Page list before deploying".to_string() + })?; + let info = deploy_page_version_on_relay( + &relay_url, + &session.token, + &slug, + &version_id, + &page.generation, + ) + .await + .map_err(|e| e.to_string())?; + if !account_context_is_current(generation) { + return Err("account context changed".to_string()); + } let mut value = serde_json::to_value(info).map_err(|e| e.to_string())?; if let Some(obj) = value.as_object_mut() { let path = obj @@ -376,7 +777,8 @@ fn register_page_deploy_host() { fn register_page_publish_host() { set_page_publish_handler(Arc::new(|request| { Box::pin(async move { - let (session, relay_url) = read_account_context().await?; + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; let result = publish_page_content_on_relay( &relay_url, &session.token, @@ -390,27 +792,217 @@ fn register_page_publish_host() { ) .await .map_err(|e| e.to_string())?; + if !account_context_is_current(generation) { + return Err("account context changed".to_string()); + } serde_json::to_value(result).map_err(|e| e.to_string()) }) })); } -fn get_account_session() -> &'static Arc>> { - ACCOUNT_SESSION.get_or_init(|| Arc::new(RwLock::new(None))) -} - -fn get_account_relay_url() -> &'static Arc>> { - ACCOUNT_RELAY_URL.get_or_init(|| Arc::new(RwLock::new(None))) +fn get_account_context() -> &'static Arc>> { + ACCOUNT_CONTEXT.get_or_init(|| Arc::new(RwLock::new(None))) } /// Read both the session and relay URL, returning owned clones to avoid /// holding locks across awaits. pub(crate) async fn read_account_context() -> Result<(AccountSession, String), String> { - let session = get_account_session().read().await.clone(); - let relay_url = get_account_relay_url().read().await.clone(); - match (session, relay_url) { - (Some(s), Some(u)) => Ok((s, u)), - _ => Err("not logged in".to_string()), + let generation = account_context_generation(); + read_account_context_for_generation(generation).await +} + +async fn read_account_context_raw() -> Result<(AccountSession, String), String> { + get_account_context() + .read() + .await + .clone() + .map(|context| (context.session, context.relay_url)) + .ok_or_else(|| "not logged in".to_string()) +} + +pub(crate) async fn read_account_context_for_generation( + generation: u64, +) -> Result<(AccountSession, String), String> { + if !account_context_is_current(generation) { + return Err("account context changed".to_string()); + } + let context = read_account_context_raw().await?; + if !account_context_is_current(generation) { + return Err("account context changed".to_string()); + } + Ok(context) +} + +async fn account_context_matches(generation: u64, token: &str) -> bool { + if !account_context_is_current(generation) { + return false; + } + let token_matches = get_account_context() + .read() + .await + .as_ref() + .is_some_and(|context| context.session.token == token); + token_matches && account_context_is_current(generation) +} + +fn get_device_routing_state() -> &'static std::sync::Mutex { + DEVICE_ROUTING_STATE.get_or_init(|| std::sync::Mutex::new(DeviceRoutingState::default())) +} + +fn with_device_routing_state(update: impl FnOnce(&mut DeviceRoutingState) -> T) -> T { + let mut state = get_device_routing_state() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + update(&mut state) +} + +fn new_device_routing_owner( + account_generation: u64, + account_token: &str, + service_connection_id: u64, +) -> DeviceRoutingOwner { + DeviceRoutingOwner { + account_generation, + account_token: account_token.to_string(), + connection_id: DEVICE_ROUTING_CONNECTION_ID.fetch_add(1, Ordering::AcqRel) + 1, + service_connection_id, + } +} + +fn install_device_routing_owner(owner: DeviceRoutingOwner) { + with_device_routing_state(|state| { + state.owner = Some(owner); + state.online_devices.clear(); + }); +} + +fn device_routing_owner_is_registered(owner: &DeviceRoutingOwner) -> bool { + with_device_routing_state(|state| state.owner.as_ref() == Some(owner)) +} + +fn device_routing_owner_for_account( + account_generation: u64, + account_token: &str, +) -> Option { + with_device_routing_state(|state| { + let owner = state.owner.as_ref()?; + if owner.account_generation != account_generation || owner.account_token != account_token { + return None; + } + Some(owner.clone()) + }) +} + +fn current_device_routing_owner_snapshot() -> Option { + with_device_routing_state(|state| state.owner.clone()).filter(|owner| { + account_context_is_current(owner.account_generation) + && device_routing_owner_is_registered(owner) + }) +} + +async fn device_routing_owner_is_current(owner: &DeviceRoutingOwner) -> bool { + device_routing_owner_is_registered(owner) + && account_context_matches(owner.account_generation, &owner.account_token).await + && device_routing_owner_is_registered(owner) +} + +async fn lock_current_device_routing( + owner: &DeviceRoutingOwner, +) -> Option> { + let guard = DEVICE_ROUTING_LIFECYCLE_LOCK.read().await; + if device_routing_owner_is_current(owner).await { + Some(guard) + } else { + None + } +} + +fn replace_device_presence_if_owner( + owner: &DeviceRoutingOwner, + devices: Vec, +) -> bool { + with_device_routing_state(|state| { + if state.owner.as_ref() != Some(owner) { + return false; + } + state.online_devices = devices; + true + }) +} + +fn device_presence_for_account( + account_generation: u64, + account_token: &str, +) -> Option> { + with_device_routing_state(|state| { + let owner = state.owner.as_ref()?; + if owner.account_generation != account_generation || owner.account_token != account_token { + return None; + } + Some(state.online_devices.clone()) + }) +} + +fn clear_device_routing_if_owner(owner: &DeviceRoutingOwner) -> bool { + with_device_routing_state(|state| { + if state.owner.as_ref() != Some(owner) { + return false; + } + state.owner = None; + state.online_devices.clear(); + true + }) +} + +fn clear_device_routing_state() -> bool { + with_device_routing_state(|state| { + let had_owner = state.owner.take().is_some(); + state.online_devices.clear(); + had_owner + }) +} + +fn normalize_relay_url(relay_url: &str) -> Result { + let parsed = validate_relay_base_url(relay_url.trim()).map_err(|error| error.to_string())?; + Ok(parsed.as_str().trim_end_matches('/').to_string()) +} + +fn cloud_settings_exist_from_probe(result: Result, E>) -> Result { + result.map(|settings| settings.is_some()) +} + +async fn revoke_login_candidate( + client: &AccountClient, + relay_url: &str, + session: &AccountSession, + reason: &str, +) { + if let Err(error) = client.revoke_token(relay_url, session).await { + log::warn!("Failed to revoke rejected login candidate after {reason}: {error}"); + } +} + +fn select_replaced_account_for_revocation( + previous: Option, + replacement_relay_url: &str, + replacement_token: &str, +) -> Option { + previous.filter(|account| { + account.relay_url != replacement_relay_url || account.session.token != replacement_token + }) +} + +async fn revoke_replaced_account(client: &AccountClient, account: Option) { + let Some(account) = account else { + return; + }; + if let Err(error) = client + .revoke_token(&account.relay_url, &account.session) + .await + { + // The replacement is already fully published. Revocation is + // best-effort and must never roll the new account back. + log::warn!("Failed to revoke replaced account token: {error}"); } } @@ -445,21 +1037,46 @@ pub fn set_mobile_web_resource_path(path: PathBuf) { /// after fresh login. async fn register_delegated_identity_providers() { // Room-channel provider for mobile-web. - let session_clone = get_account_session().clone(); - let relay_url_clone = get_account_relay_url().clone(); + let account_context = get_account_context().clone(); if let Some(service) = get_service_holder().read().await.as_ref() { service .set_delegated_identity_provider(move || { - let session_arc = session_clone.clone(); - let relay_url_arc = relay_url_clone.clone(); + let account_context = account_context.clone(); Box::pin(async move { - let session = session_arc.read().await.clone()?; - let relay_url = relay_url_arc.read().await.clone()?; + let generation = account_context_generation(); + if !account_context_is_current(generation) { + return None; + } + // Core calls this provider while holding the room lifecycle + // lease. Acquire the account lease second and return it with + // the credentials so account replacement cannot begin until + // Core has encrypted and sent the response. + let account_lease = lock_account_sync(generation).await.ok()?; + let context = account_context.read().await.clone()?; + if !account_context_matches(generation, &context.session.token).await { + return None; + } match AccountClient::new() - .delegate_token(&relay_url, &session) + .delegate_token(&context.relay_url, &context.session) .await { - Ok(delegated) => Some((delegated.token, session.master_key, relay_url)), + Ok(delegated) => { + if delegated.user_id != context.session.user_id { + log::warn!( + "Delegated identity user did not match the desktop account" + ); + return None; + } + if !account_context_matches(generation, &context.session.token).await { + return None; + } + Some(DelegatedIdentityAuthorization::with_host_lease( + delegated.token, + delegated.user_id, + context.session.master_key, + account_lease, + )) + } Err(e) => { log::warn!("Delegate token failed: {e}"); None @@ -475,22 +1092,32 @@ async fn register_delegated_identity_providers() { // Login/restore may switch accounts; drop any prior URL-bound mobile // identity so the next pair can bind to the current account user id. service.clear_trusted_mobile_identity().await; + service.clear_bot_delegated_identities().await; } // Global provider for IM bots. - let session_arc = get_account_session().clone(); - let relay_url_arc = get_account_relay_url().clone(); + let account_context = get_account_context().clone(); bitfun_core::service::remote_connect::bot::set_delegated_identity_provider(move || { - let session_arc = session_arc.clone(); - let relay_url_arc = relay_url_arc.clone(); + let account_context = account_context.clone(); Box::pin(async move { - let session = session_arc.read().await.clone()?; - let relay_url = relay_url_arc.read().await.clone()?; + let generation = account_context_generation(); + if !account_context_is_current(generation) { + return None; + } + let context = account_context.read().await.clone()?; + if !account_context_is_current(generation) { + return None; + } match AccountClient::new() - .delegate_token(&relay_url, &session) + .delegate_token(&context.relay_url, &context.session) .await { - Ok(delegated) => Some((relay_url, delegated.token, session.master_key.to_vec())), + Ok(delegated) if account_context_is_current(generation) => Some(( + context.relay_url, + delegated.token, + context.session.master_key.to_vec(), + )), + Ok(_) => None, Err(e) => { log::warn!("Bot delegate token failed: {e}"); None @@ -508,15 +1135,17 @@ async fn register_account_pairing_context(service: &RemoteConnectService) { .unwrap_or_default(); service.set_account_pairing_username(Some(username)).await; - let session_arc = get_account_session().clone(); - let relay_url_arc = get_account_relay_url().clone(); + let account_context = get_account_context().clone(); let pairing_attempts = Arc::new(tokio::sync::Mutex::new((0_u32, None::))); service .set_account_pairing_verifier(move |username, password| { - let session_arc = session_arc.clone(); - let relay_url_arc = relay_url_arc.clone(); + let account_context = account_context.clone(); let pairing_attempts = pairing_attempts.clone(); async move { + let generation = account_context_generation(); + if !account_context_is_current(generation) { + return Err("Desktop account is changing; scan again".to_string()); + } { let mut attempts = pairing_attempts.lock().await; if let Some(locked_until) = attempts.1 { @@ -529,24 +1158,31 @@ async fn register_account_pairing_context(service: &RemoteConnectService) { *attempts = (0, None); } } - let session = session_arc + let context = account_context .read() .await .clone() .ok_or_else(|| "Desktop is not logged into a BitFun account".to_string())?; - let relay_url = relay_url_arc - .read() + if !account_context_is_current(generation) { + return Err("Desktop account is changing; scan again".to_string()); + } + let account_lease = lock_account_sync(generation) .await - .clone() - .ok_or_else(|| "Desktop is not logged into a BitFun account".to_string())?; + .map_err(|_| "Desktop account is changing; scan again".to_string())?; + if !account_context_matches(generation, &context.session.token).await { + return Err("Desktop account is changing; scan again".to_string()); + } let verification = AccountClient::new() .verify_password_for_master_key( - &relay_url, + &context.relay_url, &username, &password, - &session.master_key, + &context.session.master_key, ) .await; + if !account_context_matches(generation, &context.session.token).await { + return Err("Desktop account changed; scan again".to_string()); + } if let Err(error) = verification { // Keep the real cause in desktop logs (network vs bad // credentials); the mobile only gets the unified message. @@ -562,7 +1198,10 @@ async fn register_account_pairing_context(service: &RemoteConnectService) { return Err("Invalid username or password".to_string()); } *pairing_attempts.lock().await = (0, None); - Ok(session.user_id) + Ok(AccountPairingVerification::with_host_lease( + context.session.user_id, + account_lease, + )) } }) .await; @@ -572,12 +1211,35 @@ pub fn init_on_startup() { register_page_deploy_host(); register_page_publish_host(); tokio::spawn(async { + let startup_generation = account_context_generation(); // Restore persisted account session (if any) before anything else // so that auto-sync, device routing, and bot delegation work on restart. match session_store::load_session_detailed() { Ok(Some(loaded)) => { let user_id = loaded.user_id.clone(); - let relay_url = loaded.relay_url.clone(); + let relay_url = match normalize_relay_url(&loaded.relay_url) { + Ok(url) => url, + Err(error) => { + log::warn!("Ignoring invalid persisted relay URL: {error}"); + session_store::clear_session(); + sync_account_login_capability(false); + if let Err(error) = ensure_service().await { + log::warn!("Remote connect startup init failed: {error}"); + } + return; + } + }; + let Some(restore_guard) = + cancel_and_wait_if_account_current(startup_generation).await + else { + log::info!( + "Skipped persisted session restore after a newer account transition" + ); + if let Err(error) = ensure_service().await { + log::warn!("Remote connect startup init failed: {error}"); + } + return; + }; if let Some(device_id) = loaded.device_id.as_deref() { if let Err(e) = DeviceIdentity::adopt_account_device_id(device_id) { log::warn!("Failed to adopt restored session device_id: {e}"); @@ -588,13 +1250,16 @@ pub fn init_on_startup() { user_id: user_id.clone(), master_key: loaded.master_key, }; - *get_account_session().write().await = Some(session); - *get_account_relay_url().write().await = Some(relay_url.clone()); + *get_account_context().write().await = Some(AccountContextState { + session, + relay_url: relay_url.clone(), + }); sync_account_login_capability(true); // Keep the mirrored "Self-Hosted" server field in sync for // sessions restored from an older version without the mirror. set_self_hosted_form_url(Some(&relay_url)); log::info!("Restored account session for user {user_id}"); + drop(restore_guard); // Initialize the remote-connect service if not yet ready. if let Err(e) = ensure_service().await { @@ -1139,7 +1804,7 @@ pub async fn remote_connect_start( let service = guard.as_ref().ok_or("service not initialized")?; // Refresh account pairing context so a newly logged-in session is reflected // in the QR (`auth=account&user=...`) before the room is created. - if get_account_session().read().await.is_some() { + if read_account_context().await.is_ok() { register_account_pairing_context(service).await; } else { service.clear_account_pairing_context().await; @@ -1168,9 +1833,7 @@ pub async fn remote_connect_stop_bot() -> Result<(), String> { service.stop_bots().await; } // Remove persistence so the bot is not auto-restored - let mut data = bot::load_bot_persistence(); - data.connections.clear(); - bot::save_bot_persistence(&data); + bot::update_bot_persistence(|data| data.connections.clear()); Ok(()) } @@ -1208,9 +1871,7 @@ pub async fn remote_connect_get_form_state() -> Result Result<(), String> { - let mut data = bot::load_bot_persistence(); - data.form_state = request; - bot::save_bot_persistence(&data); + bot::update_bot_persistence(|data| data.form_state = request); Ok(()) } @@ -1329,9 +1990,7 @@ pub async fn remote_connect_set_bot_verbose_mode(verbose: bool) -> Result<(), St "remote_connect_set_bot_verbose_mode called with verbose={}", verbose ); - let mut data = bot::load_bot_persistence(); - data.verbose_mode = verbose; - bot::save_bot_persistence(&data); + bot::update_bot_persistence(|data| data.verbose_mode = verbose); log::info!("Saved bot verbose_mode={} to persistence", verbose); Ok(()) } @@ -1342,8 +2001,10 @@ pub async fn remote_connect_set_bot_verbose_mode(verbose: bool) -> Result<(), St /// The master key is deliberately NOT included — it stays in Rust memory. #[derive(Serialize, Deserialize, Clone)] pub struct AccountLoginResult { - pub token: String, pub user_id: String, + /// Opaque owner for the pending cloud/local decision. This is never an + /// account bearer token and is present only when a choice is required. + pub pending_login_id: Option, /// Whether the relay already has a cloud settings blob for this account. /// `true` = non-first login → the frontend should prompt the user before /// overwriting local settings. `false` = first login → auto-upload local. @@ -1365,6 +2026,11 @@ pub struct AccountAuthRequest { pub password: String, } +#[derive(Deserialize)] +pub struct PendingAccountLoginRequest { + pub pending_login_id: String, +} + fn current_device_identity() -> Result { DeviceIdentity::from_current_machine().map_err(|e| format!("detect device: {e}")) } @@ -1374,10 +2040,16 @@ fn current_device_identity() -> Result { /// must not restore a logged-in state. static PENDING_SYNC_CHOICE: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); +static PENDING_LOGIN_ID: OnceLock>> = OnceLock::new(); +static LAST_FINALIZED_PENDING_LOGIN: OnceLock< + std::sync::Mutex>, +> = OnceLock::new(); /// Persist the in-memory account session so restart restores login. async fn persist_account_session(device_id: Option<&str>) -> Result<(), String> { - let (session, relay_url) = read_account_context().await?; + // Login installs the context while its transition permit is still held; + // this internal persistence step intentionally reads that staged value. + let (session, relay_url) = read_account_context_raw().await?; session_store::save_session_with_device( &session.token, &session.user_id, @@ -1393,18 +2065,43 @@ async fn persist_account_session(device_id: Option<&str>) -> Result<(), String> /// /// Pair with `PENDING_SYNC_CHOICE` / `account_login`: never persist or emit /// logged-in before this runs when `has_cloud_settings` was true. Closing the -/// overwrite UI must `account_logout` instead of leaving a memory-only session. +/// overwrite UI must conditionally cancel its opaque pending owner instead of +/// leaving a memory-only session. #[tauri::command] -pub async fn account_finalize_login() -> Result<(), String> { +pub async fn account_finalize_login(request: PendingAccountLoginRequest) -> Result<(), String> { + if finalized_pending_login_is_current(&request.pending_login_id).await { + return Ok(()); + } + let _pending_guard = match lock_pending_login_for_finalize(&request.pending_login_id).await { + Ok(guard) => guard, + Err(error) => { + // A concurrent/retried call may arrive after the first invocation + // committed but before its transport response reached the UI. + if finalized_pending_login_is_current(&request.pending_login_id).await { + return Ok(()); + } + return Err(error); + } + }; + let account_generation = account_context_generation(); + let (session, _) = read_account_context_for_generation(account_generation).await?; + let finalized_owner = FinalizedPendingLoginOwner { + pending_login_id: request.pending_login_id, + account_generation, + account_token: session.token, + }; let device = current_device_identity()?; persist_account_session(Some(device.device_id.as_str())).await?; - PENDING_SYNC_CHOICE.store(false, std::sync::atomic::Ordering::Relaxed); + set_pending_login_id(None); TOKEN_EXPIRED.store(false, std::sync::atomic::Ordering::Relaxed); sync_account_login_capability(true); register_delegated_identity_providers().await; - let relay_url = get_account_relay_url().read().await.clone(); + let relay_url = read_account_context() + .await + .ok() + .map(|(_, relay_url)| relay_url); emit_account_event( "account://login-state", serde_json::json!({ @@ -1412,91 +2109,172 @@ pub async fn account_finalize_login() -> Result<(), String> { "relay_url": relay_url, }), ); + record_finalized_pending_login(finalized_owner); log::info!("Account login finalized (sync choice accepted)"); Ok(()) } +/// Abandon only the pending login identified by `pending_login_id`. A stale +/// component cleanup is a no-op and, importantly, does not begin an account +/// transition or increment the context generation. +#[tauri::command] +pub async fn account_cancel_pending_login( + request: PendingAccountLoginRequest, +) -> Result { + let transition_guard = ACCOUNT_CONTEXT_TRANSITION_LOCK.lock().await; + let generation = account_context_generation(); + if !account_context_is_current(generation) + || !pending_login_is_owned_by(&request.pending_login_id) + { + return Ok(false); + } + + let transition = AccountContextTransitionPermit::begin(); + let sync_guard = ACCOUNT_AUTO_SYNC_LOCK.lock().await; + bitfun_core::service::remote_connect::settings_sync::wait_for_sync_operations_idle().await; + let _transition_guard = AccountContextTransitionGuard { + sync_guard: Some(sync_guard), + transition: Some(transition), + transition_guard: Some(transition_guard), + }; + if !pending_login_is_owned_by(&request.pending_login_id) { + return Ok(false); + } + clear_account_login_state(true).await; + log::info!("Pending account login cancelled"); + Ok(true) +} + #[tauri::command] pub async fn account_login(request: AccountAuthRequest) -> Result { + // Keep the old account fully usable while credentials are verified. Only a + // successful candidate is allowed to begin the protected replacement + // transition and retire the old account's runtime state. + let _login_guard = ACCOUNT_LOGIN_LOCK.lock().await; + let expected_generation = account_context_generation(); + if !account_context_is_current(expected_generation) { + return Err("account context changed".to_string()); + } + let relay_url = normalize_relay_url(&request.relay_url)?; let device = current_device_identity()?; let client = AccountClient::new(); let session = client - .login( - &request.relay_url, - &request.username, - &request.password, - &device, - ) + .login(&relay_url, &request.username, &request.password, &device) .await .map_err(|e| format!("{e}"))?; // Check whether the relay already has a cloud settings blob for this // account. This tells the frontend whether to prompt before overwriting. - let has_cloud_settings = client - .fetch_settings(&request.relay_url, &session) - .await - .unwrap_or(None) - .is_some(); + let has_cloud_settings = + match cloud_settings_exist_from_probe(client.fetch_settings(&relay_url, &session).await) { + Ok(has_cloud_settings) => has_cloud_settings, + Err(error) => { + let message = error.to_string(); + revoke_login_candidate(&client, &relay_url, &session, "settings probe failure") + .await; + return Err(message); + } + }; + + let Some(mut transition_guard) = cancel_and_wait_if_account_current(expected_generation).await + else { + revoke_login_candidate(&client, &relay_url, &session, "account replacement race").await; + return Err("account context changed".to_string()); + }; + // The old account is now hidden, so account-backed tools must be hidden as + // well. A committed no-cloud login re-enables them after publication. + sync_account_login_capability(false); + let replaced_account = select_replaced_account_for_revocation( + get_account_context().read().await.clone(), + &relay_url, + &session.token, + ); + + // The candidate is authenticated and the old context is now hidden. Clear + // its device socket, presence, controllers, and account-pairing callbacks + // before publishing the replacement context. + stop_and_clear_device_routing("Account changed").await; + if let Some(service) = get_service_holder().read().await.as_ref() { + service.clear_account_pairing_context().await; + service.clear_trusted_mobile_identity().await; + service.clear_bot_delegated_identities().await; + } + // A replacement that still needs a sync choice must remain memory-only; + // never leave the prior account's persisted session restorable on crash. + session_store::clear_session(); + let pending_login_id = has_cloud_settings.then(|| uuid::Uuid::new_v4().to_string()); let result = AccountLoginResult { - token: session.token.clone(), user_id: session.user_id.clone(), + pending_login_id: pending_login_id.clone(), has_cloud_settings, }; - *get_account_session().write().await = Some(session); - *get_account_relay_url().write().await = Some(request.relay_url.clone()); + *get_account_context().write().await = Some(AccountContextState { + session, + relay_url: relay_url.clone(), + }); // Persist non-secret credentials for next startup pre-fill - save_credential_hint(&request.username, &request.relay_url); + save_credential_hint(&request.username, &relay_url); // Mirror the relay URL into the Remote Connect "Self-Hosted" server field // so phone pairing can ride the same relay the account is logged into. - set_self_hosted_form_url(Some(&request.relay_url)); + set_self_hosted_form_url(Some(&relay_url)); // Reset the token-expired flag on fresh login TOKEN_EXPIRED.store(false, std::sync::atomic::Ordering::Relaxed); if has_cloud_settings { // Hold the session in memory only until the user picks cloud vs local. // Persisting here would restore "logged in" after a kill with no choice. - PENDING_SYNC_CHOICE.store(true, std::sync::atomic::Ordering::Relaxed); - sync_account_login_capability(false); + set_pending_login_id(pending_login_id); + } else { + set_pending_login_id(None); + if let Err(e) = persist_account_session(Some(device.device_id.as_str())).await { + log::warn!("Failed to persist session: {e}"); + } + + register_delegated_identity_providers().await; + } + + // AccountClient revocation is transport-only and does not re-enter host + // lifecycle locks. Keep the transition lease until it finishes so no + // caller can replace B and then observe this command returning B's result. + revoke_replaced_account(&client, replaced_account).await; + // End context hiding before notifying listeners. Keep the transition mutex + // until after the event so a listener's immediate account-status probe sees + // this committed account instead of a transient logged-out state. + transition_guard.make_context_observable(); + if !has_cloud_settings { + sync_account_login_capability(true); + emit_account_event( + "account://login-state", + serde_json::json!({ + "logged_in": true, + "relay_url": relay_url, + }), + ); + } + if has_cloud_settings { log::info!( "Account authenticated pending sync choice: {} (has_cloud_settings=true)", result.user_id ); - return Ok(result); - } - - PENDING_SYNC_CHOICE.store(false, std::sync::atomic::Ordering::Relaxed); - if let Err(e) = persist_account_session(Some(device.device_id.as_str())).await { - log::warn!("Failed to persist session: {e}"); + } else { + log::info!( + "Account logged in: {} (has_cloud_settings=false)", + result.user_id + ); } - - register_delegated_identity_providers().await; - sync_account_login_capability(true); - - emit_account_event( - "account://login-state", - serde_json::json!({ - "logged_in": true, - "relay_url": request.relay_url, - }), - ); - - log::info!( - "Account logged in: {} (has_cloud_settings=false)", - result.user_id - ); Ok(result) } #[tauri::command] pub async fn account_status() -> Result { let pending = PENDING_SYNC_CHOICE.load(std::sync::atomic::Ordering::Relaxed); - let guard = get_account_session().read().await; - let logged_in = guard.is_some() && !pending; + let context = read_account_context().await.ok(); + let logged_in = context.is_some() && !pending; Ok(AccountStatus { logged_in, user_id: if logged_in { - guard.as_ref().map(|s| s.user_id.clone()) + context.map(|(session, _)| session.user_id) } else { None }, @@ -1508,29 +2286,60 @@ pub async fn account_status() -> Result { /// `revoke_relay_token` is false after deleting this device because the relay /// deletion already revoked the current token along with the device row. async fn clear_account_login(revoke_relay_token: bool) { + // Invalidate the active operation first, then wait for it to observe the + // cancellation and release its guard. This ensures no settings apply or + // progress event can happen after logout completes. + let _sync_guard = cancel_and_wait_for_account_auto_sync().await; + clear_account_login_state(revoke_relay_token).await; +} + +async fn clear_account_login_if_current( + expected_generation: u64, + expected_token: &str, + revoke_relay_token: bool, +) -> bool { + let Some(_transition_guard) = cancel_and_wait_if_account_current(expected_generation).await + else { + return false; + }; + let token_matches = get_account_context() + .read() + .await + .as_ref() + .is_some_and(|context| context.session.token == expected_token); + if !token_matches { + return false; + } + clear_account_login_state(revoke_relay_token).await; + true +} + +async fn clear_account_login_state(revoke_relay_token: bool) { + // Account transitions hide the context immediately; hide account-backed + // tools at the same boundary rather than after network cleanup completes. + sync_account_login_capability(false); // Disconnect device routing before clearing the session. + stop_and_clear_device_routing("Account logged out").await; if let Some(service) = get_service_holder().read().await.as_ref() { - service.stop_device_connection().await; service.clear_account_pairing_context().await; service.clear_trusted_mobile_identity().await; + service.clear_bot_delegated_identities().await; } if revoke_relay_token { // Best-effort relay revocation must not prevent local logout. - if let Ok((session, relay_url)) = read_account_context().await { + if let Ok((session, relay_url)) = read_account_context_raw().await { let _ = AccountClient::new() .revoke_token(&relay_url, &session) .await; } } - *get_account_session().write().await = None; - *get_account_relay_url().write().await = None; - PENDING_SYNC_CHOICE.store(false, std::sync::atomic::Ordering::Relaxed); + *get_account_context().write().await = None; + set_pending_login_id(None); clear_credential_hint(); session_store::clear_session(); // Clear the mirrored "Self-Hosted" server field on logout. set_self_hosted_form_url(None); TOKEN_EXPIRED.store(false, std::sync::atomic::Ordering::Relaxed); - sync_account_login_capability(false); emit_account_event( "account://login-state", serde_json::json!({ "logged_in": false }), @@ -1547,17 +2356,17 @@ pub async fn account_logout() -> Result<(), String> { /// Persist (or clear) the account relay URL in the Remote Connect /// "Self-Hosted" form field so the pairing UI follows account login state. fn set_self_hosted_form_url(url: Option<&str>) { - let mut data = bot::load_bot_persistence(); let value = url.unwrap_or_default(); - if data.form_state.custom_server_url != value { - data.form_state.custom_server_url = value.to_string(); - bot::save_bot_persistence(&data); - } + bot::update_bot_persistence(|data| { + if data.form_state.custom_server_url != value { + data.form_state.custom_server_url = value.to_string(); + } + }); } // ── P2: Device routing commands ────────────────────────────────────────── -#[derive(Serialize)] +#[derive(Clone, Debug, Serialize, PartialEq, Eq)] pub struct OnlineDeviceInfo { pub device_id: String, pub device_name: String, @@ -1568,7 +2377,9 @@ pub struct OnlineDeviceInfo { /// that logs presence updates; device messages are forwarded to the RemoteConnectService. #[tauri::command] pub async fn account_connect_devices() -> Result, String> { - let (session, relay_url) = read_account_context().await?; + let account_generation = account_context_generation(); + let sync_guard = lock_account_sync(account_generation).await?; + let (session, relay_url) = read_account_context_for_generation(account_generation).await?; let identity = current_device_identity()?; let device_name = identity.device_name.clone(); let holder = get_service_holder().read().await; @@ -1576,46 +2387,72 @@ pub async fn account_connect_devices() -> Result, String> .as_ref() .ok_or_else(|| "remote connect service not initialized".to_string())?; - // Skip reconnecting if the device WS is already active AND local identity - // already matches an online device. Otherwise reconnect so AuthOk can heal - // a drifted MAC-derived device_id (AuthOk is consumed inside start_device_connection). - if service.is_device_connected().await { - let devices = service.online_devices().await; - let local_id = DeviceIdentity::from_current_machine() - .ok() - .map(|d| d.device_id); - let local_known = local_id - .as_ref() - .is_some_and(|id| devices.iter().any(|d| d.device_id == *id)); - if local_known { - return Ok(devices - .into_iter() - .map(|d| OnlineDeviceInfo { - device_id: d.device_id, - device_name: d.device_name, - }) - .collect()); + let routing_lifecycle = DEVICE_ROUTING_LIFECYCLE_LOCK.write().await; + + // Reuse is allowed only when the active socket is explicitly owned by the + // current account generation and token. A service-level connected flag by + // itself may still describe the account that was just replaced. + if let Some(devices) = device_presence_for_account(account_generation, &session.token) { + let is_connected = service.is_device_connected().await; + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); } - log::info!( + if is_connected { + let local_id = DeviceIdentity::from_current_machine() + .ok() + .map(|d| d.device_id); + let local_known = local_id + .as_ref() + .is_some_and(|id| devices.iter().any(|d| d.device_id == *id)); + if local_known { + return Ok(devices); + } + log::info!( "Device WS connected but local device_id not in online set; reconnecting to heal identity" ); + } } - let (mut event_rx, auth_device_id) = match service + // Invalidate the prior loop before `start_device_connection` swaps the + // service client. Its compare-and-clear exit path must not touch this new + // connection's controllers or presence. + clear_device_routing_state(); + disconnect_peer_controllers("Device routing reconnecting").await; + + let (mut event_rx, auth_device_id, service_connection_id) = match service .start_device_connection(&relay_url, &session.token, &device_name) .await { Ok(result) => result, Err(e) => { let msg = format!("{e}"); + clear_device_routing_state(); + // Token invalidation re-enters the account transition path, so the + // current account-operation lease must be released first. + drop(routing_lifecycle); + drop(sync_guard); drop(holder); if error_indicates_expired_token(&msg) { - invalidate_local_account_session(&msg).await; + invalidate_local_account_session_if_current( + account_generation, + &session.token, + &msg, + ) + .await; } return Err(msg); } }; + if !account_context_matches(account_generation, &session.token).await { + service.stop_device_connection().await; + clear_device_routing_state(); + return Err("account context changed".to_string()); + } + let routing_owner = + new_device_routing_owner(account_generation, &session.token, service_connection_id); + install_device_routing_owner(routing_owner.clone()); + if let Err(e) = session_store::save_session_with_device( &session.token, &session.user_id, @@ -1628,10 +2465,14 @@ pub async fn account_connect_devices() -> Result, String> // Background task: consume events (presence / device messages / auth errors) // Note: AuthOk is consumed inside start_device_connection (adopt happens there). - let session_arc = get_account_session().clone(); + let event_session = session.clone(); + let event_owner = routing_owner.clone(); tokio::spawn(async move { use bitfun_core::service::remote_connect::relay_client::RelayEvent; - while let Some(event) = event_rx.recv().await { + 'routing_events: while let Some(event) = event_rx.recv().await { + if !device_routing_owner_is_current(&event_owner).await { + break; + } match event { RelayEvent::AuthOk { user_id, device_id } => { // Should not normally arrive — start_device_connection consumes AuthOk. @@ -1644,10 +2485,29 @@ pub async fn account_connect_devices() -> Result, String> // The socket can be rejected while the account panel is // closed. Fully invalidate local state here so desktop // capabilities and every UI surface agree immediately. - invalidate_local_account_session(&message).await; + invalidate_local_account_session_if_current( + account_generation, + &event_session.token, + &message, + ) + .await; break; } RelayEvent::DevicePresence { devices } => { + let Some(_routing_effect) = lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; + let presence = devices + .iter() + .map(|device| OnlineDeviceInfo { + device_id: device.device_id.clone(), + device_name: device.device_name.clone(), + }) + .collect(); + if !replace_device_presence_if_owner(&event_owner, presence) { + break 'routing_events; + } log::info!("Device presence updated: {} online", devices.len()); let online_device_ids = devices .iter() @@ -1664,6 +2524,9 @@ pub async fn account_connect_devices() -> Result, String> { log::warn!("Peer permission requests were not fully cancelled: {error}"); } + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } let pairs: Vec<(String, String)> = devices .iter() .map(|d| (d.device_id.clone(), d.device_name.clone())) @@ -1671,8 +2534,8 @@ pub async fn account_connect_devices() -> Result, String> emit_device_presence(&pairs); // Another device came online — pull cloud settings if needed. if devices.len() > 1 { - tokio::spawn(async { - pull_and_reconcile().await; + tokio::spawn(async move { + pull_and_reconcile(account_generation).await; }); } } @@ -1682,16 +2545,17 @@ pub async fn account_connect_devices() -> Result, String> encrypted_data, nonce, } => { - let session_guard = session_arc.read().await.clone(); - let Some(ref session) = session_guard else { - continue; - }; use bitfun_core::service::remote_connect::encryption::decrypt_from_base64; - match decrypt_from_base64(&session.master_key, &encrypted_data, &nonce) { + match decrypt_from_base64(&event_session.master_key, &encrypted_data, &nonce) { Ok(plaintext) => { use bitfun_core::service::remote_connect::remote_server::RemoteCommand; match serde_json::from_str::(&plaintext) { Ok(RemoteCommand::DeviceEvent { event, payload }) => { + let Some(_routing_effect) = + lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; // Controller receiving peer UI events — re-emit locally // under the same event name so PeerDeviceTransport listen works. log::debug!("DeviceEvent from {source_device_id}: {event}"); @@ -1703,6 +2567,11 @@ pub async fn account_connect_devices() -> Result, String> agent_type, workspace_path: _, }) => { + let Some(_routing_effect) = + lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; log::info!( "ExecuteOnDevice from {source_device_id}: \ session={:?} content_len={}", @@ -1745,6 +2614,9 @@ pub async fn account_connect_devices() -> Result, String> "DialogScheduler not available for ExecuteOnDevice" ); } + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } } Ok(RemoteCommand::SendSessionToDevice { session_data, @@ -1756,7 +2628,13 @@ pub async fn account_connect_devices() -> Result, String> session={session_id} bytes={}", session_data.len() ); - match import_session_bundle(&session_data).await { + let import_result = + import_session_bundle(&session_data, account_generation) + .await; + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } + match import_result { Ok(()) => { log::info!("Session {session_id} imported from device {source_device_id}"); } @@ -1768,27 +2646,45 @@ pub async fn account_connect_devices() -> Result, String> } } Ok(cmd) if source_device_id == "rpc" => { + let Some(_routing_effect) = + lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; // HTTP RPC request from another device via relay. // Execute the command locally and send back // the encrypted response (including errors). log::info!( "RPC request received from relay: corr={correlation_id}" ); - match execute_local_remote_command(&cmd).await { + let execution = execute_local_remote_command(&cmd).await; + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } + match execution { Ok(resp_value) => { - send_rpc_envelope(session, &correlation_id, resp_value) - .await; + send_rpc_envelope( + &event_owner, + &event_session, + &correlation_id, + resp_value, + ) + .await; } Err(e) => { log::warn!("RPC: execute command failed: {e}"); send_rpc_error( - session, + &event_owner, + &event_session, &correlation_id, format!("RPC execute failed: {e}"), ) .await; } } + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } } Ok(cmd) => { let _ = cmd; @@ -1797,12 +2693,21 @@ pub async fn account_connect_devices() -> Result, String> Err(e) => { log::warn!("Could not parse device command: {e}"); if source_device_id == "rpc" { + let Some(_routing_effect) = + lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; send_rpc_error( - session, + &event_owner, + &event_session, &correlation_id, format!("invalid RPC command: {e}"), ) .await; + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } } } } @@ -1810,17 +2715,33 @@ pub async fn account_connect_devices() -> Result, String> Err(e) => { log::warn!("Failed to decrypt device message: {e}"); if source_device_id == "rpc" { + let Some(_routing_effect) = + lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; send_rpc_error( - session, + &event_owner, + &event_session, &correlation_id, format!("failed to decrypt RPC request: {e}"), ) .await; + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } } } } } RelayEvent::Disconnected => { + let Some(_routing_effect) = lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; + if !replace_device_presence_if_owner(&event_owner, Vec::new()) { + break 'routing_events; + } log::info!("Device routing disconnected"); let request_ids = crate::api::peer_host_invoke::take_tracked_permission_requests(); @@ -1833,49 +2754,41 @@ pub async fn account_connect_devices() -> Result, String> { log::warn!("Peer permission requests were not fully cancelled: {error}"); } + if !device_routing_owner_is_current(&event_owner).await { + break 'routing_events; + } + emit_device_presence(&[]); } RelayEvent::Reconnected => { + let Some(_routing_effect) = lock_current_device_routing(&event_owner).await + else { + break 'routing_events; + }; log::info!("Device routing reconnected — AuthConnect re-sent by transport"); } _ => {} } } - let request_ids = crate::api::peer_host_invoke::disconnect_controllers(); - if let Err(error) = crate::api::peer_host_invoke::fail_closed_permission_requests( - request_ids, - "Peer device-routing stream closed", - ) - .await - { - log::warn!("Peer permission requests were not fully cancelled: {error}"); - } + finish_device_routing_event_loop(&event_owner).await; }); - let devices = service.online_devices().await; - Ok(devices - .into_iter() - .map(|d| OnlineDeviceInfo { - device_id: d.device_id, - device_name: d.device_name, - }) - .collect()) + drop(routing_lifecycle); + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); + } + Ok(device_presence_for_account(account_generation, &session.token).unwrap_or_default()) } /// Get the current online device list. #[tauri::command] pub async fn account_online_devices() -> Result, String> { - let holder = get_service_holder().read().await; - let service = holder - .as_ref() - .ok_or_else(|| "remote connect service not initialized".to_string())?; - let devices = service.online_devices().await; - Ok(devices - .into_iter() - .map(|d| OnlineDeviceInfo { - device_id: d.device_id, - device_name: d.device_name, - }) - .collect()) + let account_generation = account_context_generation(); + let (session, _) = read_account_context_for_generation(account_generation).await?; + let _lifecycle = DEVICE_ROUTING_LIFECYCLE_LOCK.read().await; + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); + } + Ok(device_presence_for_account(account_generation, &session.token).unwrap_or_default()) } /// Send an encrypted session to a peer device. The `session_json` is encrypted @@ -1886,11 +2799,8 @@ pub async fn account_send_session_to_device( session_id: String, session_json: String, ) -> Result<(), String> { - let (session, _) = read_account_context().await?; - let holder = get_service_holder().read().await; - let service = holder - .as_ref() - .ok_or_else(|| "remote connect service not initialized".to_string())?; + let account_generation = account_context_generation(); + let (session, _) = read_account_context_for_generation(account_generation).await?; // Wrap the raw session JSON in a SendSessionToDevice command envelope so the // receiving device knows what to do with the payload. @@ -1902,15 +2812,25 @@ pub async fn account_send_session_to_device( }) .map_err(|e| format!("serialize envelope: {e}"))?; + let _routing_effect = DEVICE_ROUTING_LIFECYCLE_LOCK.read().await; + let routing_owner = device_routing_owner_for_account(account_generation, &session.token) + .ok_or_else(|| "device routing not connected for current account".to_string())?; + if !device_routing_owner_is_current(&routing_owner).await { + return Err("device routing changed".to_string()); + } use bitfun_core::service::remote_connect::encryption::encrypt_to_base64; let (encrypted_data, nonce) = encrypt_to_base64(&session.master_key, &envelope).map_err(|e| format!("{e}"))?; let correlation_id = uuid::Uuid::new_v4().to_string(); - service - .send_device_message(&target_device_id, &correlation_id, &encrypted_data, &nonce) - .await - .map_err(|e| format!("{e}")) + send_device_message_with_routing_lease( + &routing_owner, + &target_device_id, + &correlation_id, + &encrypted_data, + &nonce, + ) + .await } // ── P4: Session / settings sync commands ───────────────────────────────── @@ -1952,6 +2872,8 @@ pub async fn account_fetch_synced_sessions() -> Result, Strin /// Delete a synced session blob from the relay. #[tauri::command] pub async fn account_delete_synced_session(session_id: String) -> Result<(), String> { + let generation = account_context_generation(); + let _sync_guard = lock_account_sync(generation).await?; let (session, relay_url) = read_account_context().await?; AccountClient::new() .delete_session(&relay_url, &session, &session_id) @@ -1966,6 +2888,8 @@ pub async fn account_delete_synced_session(session_id: String) -> Result<(), Str /// Upload settings blob (encrypted client-side with the master key). #[tauri::command] pub async fn account_sync_settings(settings_json: String) -> Result<(), String> { + let generation = account_context_generation(); + let _sync_guard = lock_account_sync(generation).await?; let (session, relay_url) = read_account_context().await?; bitfun_core::service::remote_connect::settings_sync::upload_settings_payload( &session, @@ -2049,6 +2973,8 @@ pub async fn account_export_local_session( app_state: State<'_, crate::api::app_state::AppState>, path_manager: State<'_, Arc>, ) -> Result<(), String> { + let generation = account_context_generation(); + let _sync_guard = lock_account_sync(generation).await?; let (acct_session, relay_url) = read_account_context().await?; let storage_path = @@ -2109,6 +3035,8 @@ pub async fn account_export_all_sessions( app_state: State<'_, crate::api::app_state::AppState>, path_manager: State<'_, Arc>, ) -> Result { + let generation = account_context_generation(); + let _sync_guard = lock_account_sync(generation).await?; let (acct_session, relay_url) = read_account_context().await?; let storage_path = @@ -2195,6 +3123,8 @@ pub async fn account_import_remote_sessions( app_state: State<'_, crate::api::app_state::AppState>, path_manager: State<'_, Arc>, ) -> Result, String> { + let generation = account_context_generation(); + let _sync_guard = lock_account_sync(generation).await?; let (acct_session, relay_url) = read_account_context().await?; let storage_path = @@ -2266,6 +3196,7 @@ pub async fn account_fetch_session_turns( app_state: State<'_, crate::api::app_state::AppState>, path_manager: State<'_, Arc>, ) -> Result { + let generation = account_context_generation(); // Soft-skip before any disk IO so accidental callers cannot fail-closed // Peer hydrate on metadata load errors. History comes from the peer host. if crate::api::peer_host_invoke::is_peer_controller_active() { @@ -2300,7 +3231,10 @@ pub async fn account_fetch_session_turns( return Ok(false); } - // Fetch the full bundle from the relay (which includes turns). + // Fetch the full bundle from the relay (which includes turns). Keep the + // account lease through the local commit so an account switch cannot write + // a stale account's history after it completes. + let _sync_guard = lock_account_sync(generation).await?; let (acct_session, relay_url) = read_account_context().await?; let fetched = AccountClient::new() .fetch_session(&relay_url, &acct_session, &session_id) @@ -2364,11 +3298,8 @@ pub async fn account_execute_on_device( agent_type: Option, workspace_path: Option, ) -> Result<(), String> { - let (session, _) = read_account_context().await?; - let holder = get_service_holder().read().await; - let service = holder - .as_ref() - .ok_or_else(|| "remote connect service not initialized".to_string())?; + let account_generation = account_context_generation(); + let (session, _) = read_account_context_for_generation(account_generation).await?; use bitfun_core::service::remote_connect::remote_server::RemoteCommand; let envelope = serde_json::to_string(&RemoteCommand::ExecuteOnDevice { @@ -2379,15 +3310,25 @@ pub async fn account_execute_on_device( }) .map_err(|e| format!("serialize envelope: {e}"))?; + let _routing_effect = DEVICE_ROUTING_LIFECYCLE_LOCK.read().await; + let routing_owner = device_routing_owner_for_account(account_generation, &session.token) + .ok_or_else(|| "device routing not connected for current account".to_string())?; + if !device_routing_owner_is_current(&routing_owner).await { + return Err("device routing changed".to_string()); + } use bitfun_core::service::remote_connect::encryption::encrypt_to_base64; let (encrypted_data, nonce) = encrypt_to_base64(&session.master_key, &envelope).map_err(|e| format!("{e}"))?; let correlation_id = uuid::Uuid::new_v4().to_string(); - service - .send_device_message(&target_device_id, &correlation_id, &encrypted_data, &nonce) - .await - .map_err(|e| format!("{e}")) + send_device_message_with_routing_lease( + &routing_owner, + &target_device_id, + &correlation_id, + &encrypted_data, + &nonce, + ) + .await } /// List all online devices in the account via the relay HTTP API. @@ -2402,18 +3343,22 @@ pub struct AccountDeviceInfo { #[tauri::command] pub async fn account_list_devices() -> Result, String> { - let (session, relay_url) = read_account_context().await?; + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; let client = AccountClient::new(); let devices = match client.list_devices(&relay_url, &session).await { Ok(devices) => devices, Err(e) => { let msg = format!("{e}"); if error_indicates_expired_token(&msg) { - invalidate_local_account_session(&msg).await; + invalidate_local_account_session_if_current(generation, &session.token, &msg).await; } return Err(msg); } }; + if !account_context_is_current(generation) { + return Err("account context changed".to_string()); + } Ok(devices .into_iter() .map(|d| AccountDeviceInfo { @@ -2428,7 +3373,8 @@ pub async fn account_list_devices() -> Result, String> { /// Remove a device from the account. #[tauri::command] pub async fn account_delete_device(targetDeviceId: String) -> Result<(), String> { - let (session, relay_url) = read_account_context().await?; + let generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(generation).await?; let is_current_device = current_device_identity()?.device_id == targetDeviceId; if let Err(error) = AccountClient::new() .delete_device(&relay_url, &session, &targetDeviceId) @@ -2436,14 +3382,20 @@ pub async fn account_delete_device(targetDeviceId: String) -> Result<(), String> { let message = error.to_string(); if error_indicates_expired_token(&message) { - invalidate_local_account_session(&message).await; + invalidate_local_account_session_if_current(generation, &session.token, &message).await; } return Err(message); } + if !account_context_is_current(generation) { + return Err("account context changed".to_string()); + } log::info!("Device {targetDeviceId} removed from account"); if is_current_device { - clear_account_login(false).await; - log::info!("Current device removed; local account session cleared"); + if clear_account_login_if_current(generation, &session.token, false).await { + log::info!("Current device removed; local account session cleared"); + } else { + return Err("account context changed".to_string()); + } } Ok(()) } @@ -2458,12 +3410,16 @@ pub async fn account_device_rpc( target_device_id: String, command_json: String, ) -> Result { - let (session, relay_url) = read_account_context().await?; + let account_generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(account_generation).await?; let client = AccountClient::new(); let response = client .device_rpc(&relay_url, &session, &target_device_id, &command_json) .await .map_err(|e| format!("{e}"))?; + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); + } Ok(response) } @@ -2471,14 +3427,46 @@ pub async fn account_device_rpc( /// Called by the frontend after pairing succeeds. #[tauri::command] pub async fn account_delegate_to_paired(correlation_id: String) -> Result { - let (session, relay_url) = read_account_context().await?; + let account_generation = account_context_generation(); + let (session, relay_url) = read_account_context_for_generation(account_generation).await?; let client = AccountClient::new(); + // Capture the room owner before requesting a token. A later secret check + // rejects a pairing that changed while the relay request was in flight. + let holder = get_service_holder().read().await; + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); + } + let service = holder + .as_ref() + .ok_or_else(|| "remote connect service not initialized".to_string())?; + let pairing_secret = service + .pairing_shared_secret() + .await + .ok_or_else(|| "no paired device".to_string())?; + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); + } + // 1. Get a delegated token from the relay let delegated = client .delegate_token(&relay_url, &session) .await .map_err(|e| format!("{e}"))?; + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); + } + if delegated.user_id != session.user_id { + return Err("delegated identity does not match the current account".to_string()); + } + + let current_pairing_secret = service.pairing_shared_secret().await; + if !account_context_matches(account_generation, &session.token).await { + return Err("account context changed".to_string()); + } + if current_pairing_secret.as_ref() != Some(&pairing_secret) { + return Err("paired device changed".to_string()); + } // 2. Build the delegated identity JSON (master_key as base64) use base64::{engine::general_purpose::STANDARD as B64, Engine}; @@ -2493,21 +3481,38 @@ pub async fn account_delegate_to_paired(correlation_id: String) -> Result, path_manager: State<'_, Arc>, ) -> Result { - account_auto_sync_inner( + if sync_operation_id == 0 { + return Err("sync operation id must be non-zero".to_string()); + } + // Capture the account generation before queueing. A logout or replacement + // login that wins the lock invalidates this call instead of letting a stale + // request start against the newly installed account. + let generation = account_context_generation(); + let _sync_guard = lock_account_sync(generation).await?; + ACTIVE_ACCOUNT_AUTO_SYNC_OPERATION_ID.store(sync_operation_id, Ordering::Release); + let result = account_auto_sync_inner( is_first_login, workspace_path, config_json, + sync_operation_id, app_state, path_manager, ) - .await + .await; + let _ = ACTIVE_ACCOUNT_AUTO_SYNC_OPERATION_ID.compare_exchange( + sync_operation_id, + 0, + Ordering::AcqRel, + Ordering::Acquire, + ); + result } async fn account_auto_sync_inner( is_first_login: bool, workspace_path: String, config_json: String, + sync_operation_id: u64, app_state: State<'_, crate::api::app_state::AppState>, path_manager: State<'_, Arc>, ) -> Result { @@ -2558,49 +3582,69 @@ async fn account_auto_sync_inner( sessions_imported: 0, }); } + ensure_account_auto_sync_current(sync_operation_id)?; let (acct_session, relay_url) = read_account_context().await?; let client = AccountClient::new(); use bitfun_core::service::remote_connect::settings_sync; // 1. Settings sync let settings_synced = if is_first_login { - emit_sync_progress("uploading_settings", 5, None, None, None); - settings_sync::upload_settings_payload(&acct_session, &relay_url, &config_json) - .await - .map_err(|e| format!("upload settings: {e}"))?; + emit_sync_progress(sync_operation_id, "uploading_settings", 5, None, None, None); + await_account_auto_sync( + sync_operation_id, + settings_sync::upload_settings_payload(&acct_session, &relay_url, &config_json), + ) + .await? + .map_err(|e| format!("upload settings: {e}"))?; + ensure_account_auto_sync_current(sync_operation_id)?; log::info!("First login: uploaded local settings to cloud"); - emit_sync_progress("settings_done", 15, None, None, None); + emit_sync_progress(sync_operation_id, "settings_done", 15, None, None, None); true } else { - emit_sync_progress("downloading_settings", 5, None, None, None); - let cloud = client - .fetch_settings_with_version(&relay_url, &acct_session) - .await - .map_err(|e| format!("fetch settings: {e}"))?; + emit_sync_progress( + sync_operation_id, + "downloading_settings", + 5, + None, + None, + None, + ); + let cloud = await_account_auto_sync( + sync_operation_id, + client.fetch_settings_with_version(&relay_url, &acct_session), + ) + .await? + .map_err(|e| format!("fetch settings: {e}"))?; + ensure_account_auto_sync_current(sync_operation_id)?; if let Some(blob) = cloud { - emit_sync_progress("applying_settings", 10, None, None, None); + emit_sync_progress(sync_operation_id, "applying_settings", 10, None, None, None); // Explicit user choice — always apply, even when the cursor says // this device already has this version. Applies into the global // config service, invalidates the AI client cache, reloads, and // emits `account://settings-applied`. - settings_sync::apply_settings_blob(&acct_session, &blob, true) - .await - .map_err(|e| format!("apply cloud config: {e}"))?; + await_account_auto_sync( + sync_operation_id, + settings_sync::apply_settings_blob(&acct_session, &blob, true), + ) + .await? + .map_err(|e| format!("apply cloud config: {e}"))?; + ensure_account_auto_sync_current(sync_operation_id)?; log::info!( "Applied cloud settings to local device (version={})", blob.version ); - emit_sync_progress("settings_done", 15, None, None, None); + emit_sync_progress(sync_operation_id, "settings_done", 15, None, None, None); true } else { - emit_sync_progress("settings_done", 15, None, None, None); + emit_sync_progress(sync_operation_id, "settings_done", 15, None, None, None); false } }; // 2. Session sync: upload local sessions only (backup). Do NOT import cloud // sessions into local disk — Remote peer mode reads the peer's live disk. - emit_sync_progress("listing_sessions", 18, None, None, None); + ensure_account_auto_sync_current(sync_operation_id)?; + emit_sync_progress(sync_operation_id, "listing_sessions", 18, None, None, None); let storage_path = desktop_effective_session_storage_path(&app_state, &workspace_path, None, None).await; let manager = PersistenceManager::new(path_manager.inner().clone()) @@ -2613,6 +3657,7 @@ async fn account_auto_sync_inner( let export_candidates = local_sessions.len(); emit_sync_progress( + sync_operation_id, "exporting_sessions", 20, Some(0), @@ -2623,6 +3668,7 @@ async fn account_auto_sync_inner( let mut sync_state_local = sync_state::load(&acct_session.user_id); let mut pending_uploads: Vec<(String, String, String)> = Vec::new(); for meta in local_sessions.iter() { + ensure_account_auto_sync_current(sync_operation_id)?; let turns = manager .load_session_turns(&storage_path, &meta.session_id) .await @@ -2650,7 +3696,14 @@ async fn account_auto_sync_inner( } let upload_total = pending_uploads.len(); - emit_sync_progress("exporting_sessions", 20, Some(0), Some(upload_total), None); + emit_sync_progress( + sync_operation_id, + "exporting_sessions", + 20, + Some(0), + Some(upload_total), + None, + ); let completed = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); let uploaded: Vec<(String, String, i64)> = stream::iter(pending_uploads) @@ -2660,16 +3713,29 @@ async fn account_auto_sync_inner( let acct_session = acct_session.clone(); let completed = completed.clone(); async move { - let result = client - .upload_session(&relay_url, &acct_session, &session_id, &bundle_json) - .await; + if ensure_account_auto_sync_current(sync_operation_id).is_err() { + return None; + } + let result = match await_account_auto_sync( + sync_operation_id, + client.upload_session(&relay_url, &acct_session, &session_id, &bundle_json), + ) + .await + { + Ok(result) => result, + Err(_) => return None, + }; let done = completed.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1; let percent = if upload_total == 0 { 95u8 } else { 20 + ((75 * done) / upload_total) as u8 }; + if ensure_account_auto_sync_current(sync_operation_id).is_err() { + return None; + } emit_sync_progress( + sync_operation_id, "exporting_sessions", percent.min(95), Some(done), @@ -2690,6 +3756,8 @@ async fn account_auto_sync_inner( .collect() .await; + ensure_account_auto_sync_current(sync_operation_id)?; + let exported = uploaded.len(); let mut max_uploaded_version = sync_state_local.last_session_since; for (session_id, hash, version) in uploaded { @@ -2703,8 +3771,17 @@ async fn account_auto_sync_inner( } let _ = sync_state::save(&acct_session.user_id, &sync_state_local); + ensure_session_backup_complete(upload_total, exported)?; + log::info!("Auto-sync: settings={settings_synced} exported={exported} imported=0"); - emit_sync_progress("done", 100, Some(exported), Some(0), None); + emit_sync_progress( + sync_operation_id, + "done", + 100, + Some(exported), + Some(0), + None, + ); Ok(AutoSyncResult { settings_synced, sessions_exported: exported, @@ -2712,6 +3789,15 @@ async fn account_auto_sync_inner( }) } +fn ensure_session_backup_complete(total: usize, uploaded: usize) -> Result<(), String> { + if uploaded == total { + return Ok(()); + } + Err(format!( + "session backup incomplete: uploaded {uploaded} of {total}; retry will resume remaining sessions" + )) +} + // ── Auto-sync: debounced upload on session changes ───────────────────────── // // Settings sync (debounced push + 30s pull) is owned by the shared engine in @@ -2750,8 +3836,30 @@ fn start_settings_sync_engine() { use bitfun_core::service::remote_connect::settings_sync; let hooks = settings_sync::SettingsSyncHooks { account_context: Some(std::sync::Arc::new(|| { - Box::pin(async { read_account_context().await.map_err(anyhow::Error::msg) }) + Box::pin(async { + if !background_account_sync_is_allowed() { + return Err(anyhow::anyhow!( + "account login is waiting for a settings choice" + )); + } + let generation = account_context_generation(); + if !account_context_is_current(generation) { + return Err(anyhow::anyhow!("account context is transitioning")); + } + let (account, relay_url) = + read_account_context().await.map_err(anyhow::Error::msg)?; + if !account_context_is_current(generation) { + return Err(anyhow::anyhow!("account context changed while reading")); + } + if !background_account_sync_is_allowed() { + return Err(anyhow::anyhow!( + "account login is waiting for a settings choice" + )); + } + Ok((account, relay_url, generation)) + }) })), + is_account_context_current: Some(std::sync::Arc::new(account_context_is_current)), should_pause: Some(std::sync::Arc::new(|| { crate::api::peer_host_invoke::is_peer_controller_active() })), @@ -2860,10 +3968,21 @@ async fn execute_debounced_sync( upserts: HashMap, deletes: std::collections::HashSet, ) { + if !background_account_sync_is_allowed() { + log::debug!("Debounced sync skipped while account login awaits a settings choice"); + return; + } if crate::api::peer_host_invoke::is_peer_controller_active() { log::debug!("Debounced sync skipped while peer controller mode is active"); return; } + let generation = account_context_generation(); + let Ok(_sync_guard) = lock_account_sync(generation).await else { + return; + }; + if !background_account_sync_is_allowed() { + return; + } // Need to be logged in let (acct_session, relay_url) = match read_account_context().await { Ok(ctx) => ctx, @@ -3100,7 +4219,14 @@ async fn execute_local_remote_command( /// Import a SessionBundle JSON into local storage. Tries all workspace session /// directories and writes to the first one found (or creates one if none exist). -async fn import_session_bundle(bundle_json: &str) -> anyhow::Result<()> { +async fn import_session_bundle(bundle_json: &str, account_generation: u64) -> anyhow::Result<()> { + let _sync_guard = lock_account_sync(account_generation) + .await + .map_err(anyhow::Error::msg)?; + // A queued event from a disconnected account must not write into a new + // account's local session view even if its encrypted payload was already + // received before the socket closed. + read_account_context().await.map_err(anyhow::Error::msg)?; let bundle: SessionBundle = serde_json::from_str(bundle_json)?; let path_manager = std::sync::Arc::new(bitfun_core::infrastructure::PathManager::new()?); @@ -3149,11 +4275,21 @@ async fn import_session_bundle(bundle_json: &str) -> anyhow::Result<()> { /// One-shot cloud settings pull, triggered when another same-account device /// comes online. The periodic pull lives in the shared settings sync engine. -async fn pull_and_reconcile() { +async fn pull_and_reconcile(account_generation: u64) { + if !background_account_sync_is_allowed() { + log::debug!("Pull: skip while account login awaits a settings choice"); + return; + } if crate::api::peer_host_invoke::is_peer_controller_active() { log::debug!("Pull: skip while peer controller mode is active"); return; } + let Ok(_sync_guard) = lock_account_sync(account_generation).await else { + return; + }; + if !background_account_sync_is_allowed() { + return; + } let Ok((acct_session, relay_url)) = read_account_context().await else { return; }; @@ -3170,6 +4306,287 @@ async fn pull_and_reconcile() { mod sync_state_tests { use super::*; + static ACCOUNT_CONTEXT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + #[test] + fn relay_url_normalization_removes_all_trailing_slashes() { + assert_eq!( + normalize_relay_url("https://relay.example.com///").unwrap(), + "https://relay.example.com" + ); + } + + #[test] + fn settings_probe_errors_are_not_treated_as_an_empty_cloud() { + assert!(!cloud_settings_exist_from_probe::(Ok(None)).unwrap()); + assert!(cloud_settings_exist_from_probe::(Ok(Some(1))).unwrap()); + assert_eq!( + cloud_settings_exist_from_probe::(Err("relay unavailable")), + Err("relay unavailable") + ); + } + + #[test] + fn login_result_exposes_only_an_opaque_pending_owner() { + let value = serde_json::to_value(AccountLoginResult { + user_id: "user-a".to_string(), + pending_login_id: Some("pending-a".to_string()), + has_cloud_settings: true, + }) + .unwrap(); + + assert_eq!(value["pending_login_id"], "pending-a"); + assert!(value.get("token").is_none()); + } + + #[test] + fn background_sync_is_fail_closed_while_login_choice_is_pending() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + set_pending_login_id(Some("pending-a".to_string())); + assert!(!background_account_sync_is_allowed()); + set_pending_login_id(None); + assert!(background_account_sync_is_allowed()); + } + + #[test] + fn replaced_token_revocation_never_selects_the_published_credential() { + let account = |token: &str, relay_url: &str| AccountContextState { + session: AccountSession { + token: token.to_string(), + user_id: "user-a".to_string(), + master_key: [7; 32], + }, + relay_url: relay_url.to_string(), + }; + + assert!(select_replaced_account_for_revocation( + Some(account("token-b", "https://relay.example.com")), + "https://relay.example.com", + "token-b", + ) + .is_none()); + let same_account_old_token = select_replaced_account_for_revocation( + Some(account("token-a", "https://relay.example.com")), + "https://relay.example.com", + "token-b", + ) + .expect("same-account relogin must revoke the old token"); + assert_eq!(same_account_old_token.session.token, "token-a"); + let other_relay = select_replaced_account_for_revocation( + Some(account("same-token", "https://relay-a.example.com")), + "https://relay-b.example.com", + "same-token", + ) + .expect("the same token text on a different relay is a distinct credential"); + assert_eq!(other_relay.relay_url, "https://relay-a.example.com"); + } + + #[tokio::test(flavor = "current_thread")] + async fn account_transition_cancels_an_in_flight_auto_sync_future() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let operation_id = u64::MAX - 41; + ACTIVE_ACCOUNT_AUTO_SYNC_OPERATION_ID.store(operation_id, Ordering::Release); + let waiter = tokio::spawn(async move { + await_account_auto_sync(operation_id, std::future::pending::<()>()).await + }); + tokio::task::yield_now().await; + + let permit = AccountContextTransitionPermit::begin(); + let result = tokio::time::timeout(std::time::Duration::from_secs(1), waiter) + .await + .expect("sync cancellation should not wait for the network timeout") + .expect("cancellation task should join"); + drop(permit); + assert_eq!(result.unwrap_err(), "account sync cancelled"); + } + + #[tokio::test(flavor = "current_thread")] + async fn external_account_reads_are_hidden_during_transition() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *get_account_context().write().await = Some(AccountContextState { + session: AccountSession { + token: "token-a".to_string(), + user_id: "user-a".to_string(), + master_key: [7; 32], + }, + relay_url: "https://relay.example.com".to_string(), + }); + assert!(read_account_context().await.is_ok()); + + let permit = AccountContextTransitionPermit::begin(); + assert!(read_account_context().await.is_err()); + assert!(read_account_context_raw().await.is_ok()); + drop(permit); + + *get_account_context().write().await = None; + } + + #[tokio::test(flavor = "current_thread")] + async fn login_event_probe_sees_context_before_transition_mutex_is_released() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let transition_guard = ACCOUNT_CONTEXT_TRANSITION_LOCK.lock().await; + let transition = AccountContextTransitionPermit::begin(); + let mut guard = AccountContextTransitionGuard { + sync_guard: None, + transition: Some(transition), + transition_guard: Some(transition_guard), + }; + + assert!(!account_context_is_current(account_context_generation())); + guard.make_context_observable(); + assert!(account_context_is_current(account_context_generation())); + assert!(ACCOUNT_CONTEXT_TRANSITION_LOCK.try_lock().is_err()); + + drop(guard); + assert!(ACCOUNT_CONTEXT_TRANSITION_LOCK.try_lock().is_ok()); + } + + #[tokio::test(flavor = "current_thread")] + async fn stale_pending_login_id_cannot_finalize_or_cancel_replacement() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *get_account_context().write().await = Some(AccountContextState { + session: AccountSession { + token: "token-b".to_string(), + user_id: "user-b".to_string(), + master_key: [9; 32], + }, + relay_url: "https://relay.example.com".to_string(), + }); + set_pending_login_id(Some("pending-b".to_string())); + let generation_before = account_context_generation(); + + assert!(lock_pending_login_for_finalize("pending-a").await.is_err()); + assert!(!account_cancel_pending_login(PendingAccountLoginRequest { + pending_login_id: "pending-a".to_string(), + }) + .await + .unwrap()); + assert_eq!(account_context_generation(), generation_before); + assert!(pending_login_is_owned_by("pending-b")); + + set_pending_login_id(None); + assert!(!PENDING_SYNC_CHOICE.load(std::sync::atomic::Ordering::Acquire)); + assert!(!pending_login_is_owned_by("pending-b")); + *get_account_context().write().await = None; + } + + #[tokio::test(flavor = "current_thread")] + async fn finalize_retry_after_commit_is_idempotent_only_for_the_same_account_owner() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *get_account_context().write().await = Some(AccountContextState { + session: AccountSession { + token: "token-a".to_string(), + user_id: "user-a".to_string(), + master_key: [7; 32], + }, + relay_url: "https://relay.example.com".to_string(), + }); + let generation = account_context_generation(); + record_finalized_pending_login(FinalizedPendingLoginOwner { + pending_login_id: "pending-a".to_string(), + account_generation: generation, + account_token: "token-a".to_string(), + }); + + assert!(finalized_pending_login_is_current("pending-a").await); + assert!(!finalized_pending_login_is_current("pending-b").await); + + let transition = AccountContextTransitionPermit::begin(); + assert!(!finalized_pending_login_is_current("pending-a").await); + drop(transition); + *get_account_context().write().await = None; + clear_last_finalized_pending_login(); + } + + #[test] + fn stale_routing_owner_cannot_clear_or_update_replacement() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + clear_device_routing_state(); + let owner_a = DeviceRoutingOwner { + account_generation: 10, + account_token: "token-a".to_string(), + connection_id: 1, + service_connection_id: 101, + }; + let owner_b = DeviceRoutingOwner { + account_generation: 12, + account_token: "token-b".to_string(), + connection_id: 2, + service_connection_id: 102, + }; + install_device_routing_owner(owner_a.clone()); + assert!(replace_device_presence_if_owner( + &owner_a, + vec![OnlineDeviceInfo { + device_id: "a-device".to_string(), + device_name: "A".to_string(), + }], + )); + + install_device_routing_owner(owner_b.clone()); + assert!(!replace_device_presence_if_owner( + &owner_a, + vec![OnlineDeviceInfo { + device_id: "late-a-device".to_string(), + device_name: "Late A".to_string(), + }], + )); + assert!(!clear_device_routing_if_owner(&owner_a)); + assert_eq!( + device_presence_for_account(owner_b.account_generation, &owner_b.account_token), + Some(Vec::new()) + ); + assert!(clear_device_routing_if_owner(&owner_b)); + } + + #[test] + fn routing_presence_is_bound_to_account_generation_and_token() { + let _test_guard = ACCOUNT_CONTEXT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + clear_device_routing_state(); + let owner = DeviceRoutingOwner { + account_generation: 20, + account_token: "token-current".to_string(), + connection_id: 3, + service_connection_id: 103, + }; + install_device_routing_owner(owner.clone()); + assert!(replace_device_presence_if_owner( + &owner, + vec![OnlineDeviceInfo { + device_id: "current-device".to_string(), + device_name: "Current".to_string(), + }], + )); + + assert!(device_presence_for_account(20, "token-current").is_some()); + assert!(device_presence_for_account(21, "token-current").is_none()); + assert!(device_presence_for_account(20, "token-replaced").is_none()); + clear_device_routing_state(); + } + + #[test] + fn partial_session_backup_is_not_reported_as_success() { + assert!(ensure_session_backup_complete(3, 3).is_ok()); + let error = ensure_session_backup_complete(3, 2).unwrap_err(); + assert!(error.contains("uploaded 2 of 3")); + } + #[test] fn content_hash_is_stable() { let a = sync_state::content_hash(r#"{"session_id":"x"}"#); diff --git a/src/apps/desktop/src/api/remote_workspace_policy.rs b/src/apps/desktop/src/api/remote_workspace_policy.rs index 77aeca4918..1c720a2195 100644 --- a/src/apps/desktop/src/api/remote_workspace_policy.rs +++ b/src/apps/desktop/src/api/remote_workspace_policy.rs @@ -121,6 +121,10 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] = "account_finalize_login", RemoteWorkspacePolicy::WorkspaceAgnostic, ), + ( + "account_cancel_pending_login", + RemoteWorkspacePolicy::WorkspaceAgnostic, + ), ("account_logout", RemoteWorkspacePolicy::WorkspaceAgnostic), ( "account_online_devices", @@ -1069,6 +1073,11 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] = ), ("open_remote_workspace", RemoteWorkspacePolicy::RemoteRouted), ("open_workspace", RemoteWorkspacePolicy::LegacyUnaudited), + ( + "page_create_open_link", + RemoteWorkspacePolicy::WorkspaceAgnostic, + ), + ("page_delete", RemoteWorkspacePolicy::WorkspaceAgnostic), ( "page_delete_version", RemoteWorkspacePolicy::WorkspaceAgnostic, @@ -1079,11 +1088,8 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] = "page_list_versions", RemoteWorkspacePolicy::WorkspaceAgnostic, ), - ("page_publish", RemoteWorkspacePolicy::WorkspaceAgnostic), - ( - "page_save_version", - RemoteWorkspacePolicy::WorkspaceAgnostic, - ), + ("page_publish", RemoteWorkspacePolicy::LocalOnly), + ("page_save_version", RemoteWorkspacePolicy::LocalOnly), ("page_unpublish", RemoteWorkspacePolicy::WorkspaceAgnostic), ("page_update", RemoteWorkspacePolicy::WorkspaceAgnostic), ("paste_files", RemoteWorkspacePolicy::LegacyUnaudited), diff --git a/src/apps/desktop/src/lib.rs b/src/apps/desktop/src/lib.rs index d91353fe36..6afddd3157 100644 --- a/src/apps/desktop/src/lib.rs +++ b/src/apps/desktop/src/lib.rs @@ -1326,6 +1326,7 @@ pub async fn run() { // Account API api::remote_connect_api::account_login, api::remote_connect_api::account_finalize_login, + api::remote_connect_api::account_cancel_pending_login, api::remote_connect_api::account_status, api::remote_connect_api::account_logout, api::remote_connect_api::account_connect_devices, @@ -1353,10 +1354,12 @@ pub async fn run() { api::pages_api::page_save_version, api::pages_api::page_list, api::pages_api::page_list_versions, + api::pages_api::page_create_open_link, api::pages_api::page_deploy, api::pages_api::page_delete_version, api::pages_api::page_update, api::pages_api::page_unpublish, + api::pages_api::page_delete, api::peer_host_invoke::peer_host_invoke_complete, api::peer_host_invoke::peer_control_attach, api::peer_host_invoke::peer_control_detach, diff --git a/src/apps/relay-server/src/bin/relay_admin.rs b/src/apps/relay-server/src/bin/relay_admin.rs index e3fe0ee615..18882a08ae 100644 --- a/src/apps/relay-server/src/bin/relay_admin.rs +++ b/src/apps/relay-server/src/bin/relay_admin.rs @@ -76,7 +76,7 @@ enum Command { async fn main() -> Result<()> { let cli = Cli::parse(); - let pool = bitfun_relay_service::db::connect(&cli.db).await?; + let pool = bitfun_relay_service::db::connect_for_admin(&cli.db).await?; match cli.command { Command::AddUser { username, password } => { diff --git a/src/apps/relay-server/src/lib.rs b/src/apps/relay-server/src/lib.rs index e56575d77c..f0f52bb352 100644 --- a/src/apps/relay-server/src/lib.rs +++ b/src/apps/relay-server/src/lib.rs @@ -4,8 +4,8 @@ //! on that crate directly; this facade preserves the existing import paths. pub use bitfun_relay_service::{ - admin, db, relay, routes, AppState, DiskAssetStore, MemoryAssetStore, ResponsePayload, - RoomManager, WebAssetStore, + admin, db, page_execution, relay, routes, AppState, DiskAssetStore, MemoryAssetStore, + ResponsePayload, RoomManager, WebAssetStore, }; /// Builds the shared relay router using this compatibility host's version. diff --git a/src/apps/relay-server/tests/library_compat.rs b/src/apps/relay-server/tests/library_compat.rs index 48aa0309de..8833354036 100644 --- a/src/apps/relay-server/tests/library_compat.rs +++ b/src/apps/relay-server/tests/library_compat.rs @@ -29,6 +29,11 @@ fn legacy_library_path_exposes_supported_relay_api() { asset_store: Arc::new(MemoryAssetStore::new()), db: None, page_data: None, + page_access_manager: Arc::new(routes::pages::PageAccessManager::new()), + page_upload_manager: Arc::new(routes::pages::PageUploadManager::new()), + page_execution_guard: Arc::new( + bitfun_relay_server::page_execution::PageExecutionGuard::new(), + ), login_rate_limiter: Arc::new(routes::auth::LoginRateLimiter::new()), device_manager: relay::DeviceManager::new(), cors_allow_origins: Arc::new(Vec::new()), diff --git a/src/crates/adapters/ai-adapters/Cargo.toml b/src/crates/adapters/ai-adapters/Cargo.toml index 8f2730765e..51e478d352 100644 --- a/src/crates/adapters/ai-adapters/Cargo.toml +++ b/src/crates/adapters/ai-adapters/Cargo.toml @@ -18,8 +18,11 @@ bitfun-core-types = { path = "../../contracts/core-types" } bitfun-services-core = { path = "../../services/services-core", default-features = false, optional = true } chrono = { workspace = true } dirs = { workspace = true, optional = true } +keyring-core = { workspace = true, optional = true } eventsource-stream = { workspace = true } futures = { workspace = true } +fs2 = { workspace = true, optional = true } +libc = { workspace = true, optional = true } log = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } @@ -31,13 +34,28 @@ tokio-util = { workspace = true } urlencoding = { workspace = true } uuid = { workspace = true, optional = true } +[target.'cfg(target_os = "macos")'.dependencies] +apple-native-keyring-store = { workspace = true, optional = true } + +[target.'cfg(target_os = "windows")'.dependencies] +windows-native-keyring-store = { workspace = true, optional = true } + +[target.'cfg(all(unix, not(any(target_os = "macos", target_os = "ios", target_os = "android"))))'.dependencies] +zbus-secret-service-keyring-store = { workspace = true, optional = true } + [features] subscription-auth = [ + "dep:apple-native-keyring-store", "dep:base64", "dep:bitfun-services-core", "dep:dirs", + "dep:fs2", + "dep:keyring-core", + "dep:libc", "dep:sha2", "dep:uuid", + "dep:windows-native-keyring-store", + "dep:zbus-secret-service-keyring-store", ] [dev-dependencies] diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/antigravity.rs b/src/crates/adapters/ai-adapters/src/subscription_auth/antigravity.rs index 30558815e0..0a05db26f7 100644 --- a/src/crates/adapters/ai-adapters/src/subscription_auth/antigravity.rs +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/antigravity.rs @@ -1,8 +1,9 @@ //! Antigravity (Google) subscription login and credential resolution. //! -//! Browser PKCE login against Google OAuth on the fixed loopback port `51121`, -//! then Bearer access to the Cloud Code Assist (`cloudcode-pa`) endpoint using -//! the `gemini-code-assist` request format. Constants mirror +//! Browser PKCE login against Google OAuth on a loopback listener (preferring +//! port `51121` and falling back to an ephemeral port), then Bearer access to +//! the Cloud Code Assist (`cloudcode-pa`) endpoint using the +//! `gemini-code-assist` request format. Constants mirror //! `opencode-antigravity-auth`. use super::store::{self, StoredCredential}; @@ -15,6 +16,7 @@ use tokio_util::sync::CancellationToken; const CLIENT_ID: &str = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"; const CALLBACK_PATH: &str = "/oauth-callback"; const CALLBACK_PORT: u16 = 51121; +const CALLBACK_PORTS: &[u16] = &[CALLBACK_PORT, 0]; const AUTHORIZE_URL: &str = "https://accounts.google.com/o/oauth2/v2/auth"; const TOKEN_URL: &str = "https://oauth2.googleapis.com/token"; const CODE_ASSIST_BASE_URL: &str = "https://cloudcode-pa.googleapis.com"; @@ -26,8 +28,8 @@ const DEFAULT_MODEL: &str = "gemini-3-pro-high"; const REFRESH_LEEWAY_MS: i64 = 5 * 60 * 1000; const STORE_KEY: &str = "antigravity"; -fn redirect_uri() -> String { - oauth_server::loopback_redirect_uri(CALLBACK_PORT, CALLBACK_PATH) +fn redirect_uri(port: u16) -> String { + oauth_server::loopback_redirect_uri(port, CALLBACK_PATH) } const SCOPES: &[&str] = &[ @@ -49,18 +51,27 @@ fn client_secret() -> String { /// Returns the platform-specific User-Agent and Client-Metadata platform token. fn platform_tokens() -> (String, &'static str) { - if cfg!(target_os = "windows") { - ("windows/amd64".to_string(), "WINDOWS") - } else if cfg!(target_os = "macos") { - let arch = if cfg!(target_arch = "aarch64") { - "darwin/arm64" - } else { - "darwin/amd64" - }; - (arch.to_string(), "MACOS") - } else { - ("linux/amd64".to_string(), "LINUX") - } + platform_tokens_for(std::env::consts::OS, std::env::consts::ARCH) +} + +fn platform_tokens_for(os: &str, arch: &str) -> (String, &'static str) { + let user_agent_os = match os { + "windows" => "windows", + "macos" => "darwin", + _ => "linux", + }; + let metadata_os = match os { + "windows" => "WINDOWS", + "macos" => "MACOS", + _ => "LINUX", + }; + let user_agent_arch = match arch { + "x86_64" => "amd64", + "aarch64" => "arm64", + "x86" => "386", + other => other, + }; + (format!("{user_agent_os}/{user_agent_arch}"), metadata_os) } #[derive(Debug, Deserialize)] @@ -180,10 +191,7 @@ fn metadata_from( } } -async fn persist_tokens(tokens: TokenResponse) -> Result<()> { - let _guard = super::store_lock(super::SubscriptionProvider::Antigravity) - .lock() - .await; +async fn persist_tokens(tokens: TokenResponse, expected_revision: u64) -> Result<()> { let access = tokens .access_token .clone() @@ -194,9 +202,9 @@ async fn persist_tokens(tokens: TokenResponse) -> Result<()> { .ok_or_else(|| anyhow!("antigravity token response missing refresh_token"))?; let expires = now_ms() + tokens.expires_in.unwrap_or(3600) * 1000; let metadata = metadata_from(&tokens, None); - let mut store = store::load().await.unwrap_or_default(); - store.insert( - STORE_KEY.to_string(), + let outcome = store::upsert_if_revision( + STORE_KEY, + expected_revision, StoredCredential::Oauth { refresh, access, @@ -204,35 +212,41 @@ async fn persist_tokens(tokens: TokenResponse) -> Result<()> { account_id: None, metadata, }, - ); - store::save(&store).await?; + ) + .await?; + super::require_current_store_revision(super::SubscriptionProvider::Antigravity, outcome)?; log::info!("antigravity subscription tokens saved"); Ok(()) } /// Starts the browser PKCE login flow, binding the loopback callback server. -pub(crate) async fn begin_login(cancel: CancellationToken) -> Result { +pub(crate) async fn begin_login( + cancel: CancellationToken, + expected_revision: u64, +) -> Result { let pkce = Pkce::generate(); let state = pkce::random_state(); - let redirect_uri = redirect_uri(); + let (listener, callback_port) = oauth_server::bind_loopback_ports(CALLBACK_PORTS).await?; + let redirect_uri = redirect_uri(callback_port); let authorization_url = build_authorize_url(&pkce, &state, &redirect_uri); - let listener = oauth_server::bind_loopback(CALLBACK_PORT).await?; let verifier = pkce.verifier.clone(); let runner = async move { - tokio::select! { - _ = cancel.cancelled() => Err(anyhow!("login cancelled")), - result = async { + super::authorize_then_persist( + super::SubscriptionProvider::Antigravity, + cancel, + async { let params = oauth_server::wait_for_callback(listener, CALLBACK_PATH, &state).await?; let code = params .get("code") .cloned() .ok_or_else(|| anyhow!("antigravity callback missing code"))?; - let tokens = exchange_code(&code, &verifier, &redirect_uri).await?; - persist_tokens(tokens).await - } => result, - } + exchange_code(&code, &verifier, &redirect_uri).await + }, + move |tokens| persist_tokens(tokens, expected_revision), + ) + .await }; Ok(StartedLogin { @@ -246,13 +260,9 @@ pub(crate) async fn begin_login(cancel: CancellationToken) -> Result Result<(String, i64)> { - let _guard = super::store_lock(super::SubscriptionProvider::Antigravity) - .lock() - .await; - let mut store = store::load().await.unwrap_or_default(); - let entry = store - .get(STORE_KEY) - .cloned() + let snapshot = store::load_entry_with_revision(STORE_KEY).await?; + let entry = snapshot + .credential .ok_or_else(|| anyhow!("Antigravity is not connected; sign in first"))?; let StoredCredential::Oauth { refresh: refresh_token, @@ -277,8 +287,9 @@ async fn ensure_fresh() -> Result<(String, i64)> { let new_refresh = refreshed.refresh_token.clone().unwrap_or(refresh_token); let new_expires = now_ms() + refreshed.expires_in.unwrap_or(3600) * 1000; let new_metadata = metadata_from(&refreshed, metadata); - store.insert( - STORE_KEY.to_string(), + let outcome = store::upsert_if_revision( + STORE_KEY, + snapshot.revision, StoredCredential::Oauth { refresh: new_refresh, access: new_access.clone(), @@ -286,8 +297,9 @@ async fn ensure_fresh() -> Result<(String, i64)> { account_id, metadata: new_metadata, }, - ); - store::save(&store).await?; + ) + .await?; + super::require_current_store_revision(super::SubscriptionProvider::Antigravity, outcome)?; log::info!("antigravity subscription tokens refreshed"); Ok((new_access, new_expires)) } @@ -326,12 +338,12 @@ pub(crate) fn suggested() -> (&'static str, &'static str, &'static str) { #[cfg(test)] mod tests { - use super::{build_authorize_url, redirect_uri}; + use super::{build_authorize_url, platform_tokens_for, redirect_uri}; use crate::subscription_auth::pkce::Pkce; #[test] fn uses_registered_localhost_redirect_uri() { - let redirect_uri = redirect_uri(); + let redirect_uri = redirect_uri(super::CALLBACK_PORT); assert_eq!(redirect_uri, "http://localhost:51121/oauth-callback"); let authorize_url = build_authorize_url(&Pkce::generate(), "state", &redirect_uri); @@ -339,4 +351,28 @@ mod tests { authorize_url.contains("redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback") ); } + + #[test] + fn reports_real_architecture_on_all_desktop_platforms() { + assert_eq!( + platform_tokens_for("windows", "x86_64"), + ("windows/amd64".to_string(), "WINDOWS") + ); + assert_eq!( + platform_tokens_for("windows", "aarch64"), + ("windows/arm64".to_string(), "WINDOWS") + ); + assert_eq!( + platform_tokens_for("macos", "aarch64"), + ("darwin/arm64".to_string(), "MACOS") + ); + assert_eq!( + platform_tokens_for("linux", "x86_64"), + ("linux/amd64".to_string(), "LINUX") + ); + assert_eq!( + platform_tokens_for("linux", "aarch64"), + ("linux/arm64".to_string(), "LINUX") + ); + } } diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/codex.rs b/src/crates/adapters/ai-adapters/src/subscription_auth/codex.rs index d92b95abfb..d7eb86c2cb 100644 --- a/src/crates/adapters/ai-adapters/src/subscription_auth/codex.rs +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/codex.rs @@ -138,10 +138,7 @@ fn metadata_from(tokens: &TokenResponse) -> Option { Some(serde_json::json!({ "email": email })) } -async fn persist_tokens(tokens: TokenResponse) -> Result<()> { - let _guard = super::store_lock(super::SubscriptionProvider::Codex) - .lock() - .await; +async fn persist_tokens(tokens: TokenResponse, expected_revision: u64) -> Result<()> { let access = tokens .access_token .clone() @@ -153,9 +150,9 @@ async fn persist_tokens(tokens: TokenResponse) -> Result<()> { let expires = now_ms() + tokens.expires_in.unwrap_or(3600) * 1000; let account_id = account_id_from(&tokens); let metadata = metadata_from(&tokens); - let mut store = store::load().await.unwrap_or_default(); - store.insert( - STORE_KEY.to_string(), + let outcome = store::upsert_if_revision( + STORE_KEY, + expected_revision, StoredCredential::Oauth { refresh, access, @@ -163,14 +160,18 @@ async fn persist_tokens(tokens: TokenResponse) -> Result<()> { account_id, metadata, }, - ); - store::save(&store).await?; + ) + .await?; + super::require_current_store_revision(super::SubscriptionProvider::Codex, outcome)?; log::info!("codex subscription tokens saved"); Ok(()) } /// Starts the browser PKCE login flow, binding the loopback callback server. -pub(crate) async fn begin_login(cancel: CancellationToken) -> Result { +pub(crate) async fn begin_login( + cancel: CancellationToken, + expected_revision: u64, +) -> Result { let pkce = Pkce::generate(); let state = super::pkce::random_state(); let (listener, callback_port) = oauth_server::bind_loopback_ports(CALLBACK_PORTS).await?; @@ -179,19 +180,21 @@ pub(crate) async fn begin_login(cancel: CancellationToken) -> Result Err(anyhow!("login cancelled")), - result = async { + super::authorize_then_persist( + super::SubscriptionProvider::Codex, + cancel, + async { let params = oauth_server::wait_for_callback(listener, CALLBACK_PATH, &state).await?; let code = params .get("code") .cloned() .ok_or_else(|| anyhow!("codex callback missing code"))?; - let tokens = exchange_code(&code, &verifier, &redirect_uri).await?; - persist_tokens(tokens).await - } => result, - } + exchange_code(&code, &verifier, &redirect_uri).await + }, + move |tokens| persist_tokens(tokens, expected_revision), + ) + .await }; Ok(StartedLogin { @@ -205,13 +208,9 @@ pub(crate) async fn begin_login(cancel: CancellationToken) -> Result Result<(String, Option, i64)> { - let _guard = super::store_lock(super::SubscriptionProvider::Codex) - .lock() - .await; - let mut store = store::load().await.unwrap_or_default(); - let entry = store - .get(STORE_KEY) - .cloned() + let snapshot = store::load_entry_with_revision(STORE_KEY).await?; + let entry = snapshot + .credential .ok_or_else(|| anyhow!("Codex is not connected; sign in first"))?; let StoredCredential::Oauth { refresh: refresh_token, @@ -237,8 +236,9 @@ async fn ensure_fresh() -> Result<(String, Option, i64)> { let new_expires = now_ms() + refreshed.expires_in.unwrap_or(3600) * 1000; let new_account_id = account_id_from(&refreshed).or(account_id); let new_metadata = metadata_from(&refreshed).or(metadata); - store.insert( - STORE_KEY.to_string(), + let outcome = store::upsert_if_revision( + STORE_KEY, + snapshot.revision, StoredCredential::Oauth { refresh: new_refresh, access: new_access.clone(), @@ -246,8 +246,9 @@ async fn ensure_fresh() -> Result<(String, Option, i64)> { account_id: new_account_id.clone(), metadata: new_metadata, }, - ); - store::save(&store).await?; + ) + .await?; + super::require_current_store_revision(super::SubscriptionProvider::Codex, outcome)?; log::info!("codex subscription tokens refreshed"); Ok((new_access, new_account_id, new_expires)) } diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/jwt.rs b/src/crates/adapters/ai-adapters/src/subscription_auth/jwt.rs index 99606ae0a9..b994b992a4 100644 --- a/src/crates/adapters/ai-adapters/src/subscription_auth/jwt.rs +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/jwt.rs @@ -50,7 +50,7 @@ pub(crate) fn chatgpt_account_id(token: &str) -> Option { #[cfg(test)] mod tests { use super::*; - use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; + use base64::engine::general_purpose::URL_SAFE_NO_PAD; fn make_token(payload: serde_json::Value) -> String { let header = URL_SAFE_NO_PAD.encode(b"{\"alg\":\"none\"}"); diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/mod.rs b/src/crates/adapters/ai-adapters/src/subscription_auth/mod.rs index fb58603f0b..50437fd972 100644 --- a/src/crates/adapters/ai-adapters/src/subscription_auth/mod.rs +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/mod.rs @@ -2,8 +2,9 @@ //! //! Lets BitFun sign in to another product's subscription (Codex/ChatGPT, //! Antigravity/Google, OpenCode Zen) with an OpenCode-style in-app OAuth flow, -//! and use the resulting tokens to authenticate AI requests. Tokens are stored -//! locally in `subscription_auth.json` (mode 0600) and refreshed on resolve. +//! and use the resulting tokens to authenticate AI requests. Secret material +//! is stored in the operating-system credential vault; the local JSON file +//! contains non-secret account metadata only. //! //! There is no upgrade path for the previous Codex/Gemini CLI disk-scan import. @@ -88,11 +89,27 @@ pub struct SubscriptionAccount { /// Unix seconds when the current credential expires (for UI display). pub expires_at: Option, pub connected: bool, + /// The account was known previously, but its secret is absent from the + /// system credential vault. The UI should ask the user to sign in again. + #[serde(default)] + pub reauthentication_required: bool, + /// The system credential vault is currently locked or unavailable. Unlike + /// a missing entry, this is retryable and should not request re-login. + #[serde(default)] + pub vault_unavailable: bool, pub suggested_format: String, pub suggested_base_url: String, pub suggested_model: String, } +/// Structured sign-out result. Metadata removal determines connection state; +/// native-vault deletion may be queued for a later retry. +#[derive(Debug, Clone, Serialize)] +pub struct SubscriptionLogoutResult { + pub cleanup_pending: bool, + pub warning: Option, +} + /// Runtime-resolved credential that overrides fields in the AI client config. #[derive(Debug, Clone)] pub struct ResolvedCredential { @@ -109,6 +126,7 @@ pub struct ResolvedCredential { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LoginStartResult { pub provider: SubscriptionProvider, + pub session_id: String, pub authorization_url: String, pub user_code: Option, pub instructions: String, @@ -128,6 +146,7 @@ pub enum LoginStatus { #[derive(Debug, Clone, Serialize)] pub struct LoginSessionSnapshot { pub provider: SubscriptionProvider, + pub session_id: String, pub status: LoginStatus, pub authorization_url: Option, pub user_code: Option, @@ -145,6 +164,8 @@ pub(crate) struct StartedLogin { } struct SessionState { + /// Client-generated UUID used to correlate start/status/cancel commands. + session_id: String, status: LoginStatus, authorization_url: Option, user_code: Option, @@ -160,6 +181,7 @@ impl SessionState { fn snapshot(&self, provider: SubscriptionProvider) -> LoginSessionSnapshot { LoginSessionSnapshot { provider, + session_id: self.session_id.clone(), status: self.status, authorization_url: self.authorization_url.clone(), user_code: self.user_code.clone(), @@ -180,9 +202,15 @@ fn next_generation() -> u64 { GENERATION.fetch_add(1, Ordering::Relaxed) } -/// Per-provider lock serializing credential-store read-modify-write cycles. -/// Token refresh persists a rotated refresh token, so a concurrent refresh or -/// logout must not interleave and overwrite the newer credentials. +fn validate_session_id(session_id: &str) -> Result<()> { + uuid::Uuid::parse_str(session_id) + .map(|_| ()) + .map_err(|_| anyhow!("subscription login session_id must be a valid UUID")) +} + +/// Per-provider commit barrier for login cancellation/replacement and logout. +/// Refresh deliberately does not hold this across an external request: its +/// durable revision CAS lets logout commit immediately and reject stale tokens. pub(crate) fn store_lock(provider: SubscriptionProvider) -> &'static tokio::sync::Mutex<()> { static CODEX: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); static ANTIGRAVITY: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); @@ -194,9 +222,57 @@ pub(crate) fn store_lock(provider: SubscriptionProvider) -> &'static tokio::sync } } +/// Runs the externally cancellable authorization/polling phase, then commits +/// the resulting credential without cancellation. Dropping a credential-vault +/// write can leave an orphan secret because blocking platform keyring calls +/// continue running after their Rust future is dropped. +pub(crate) async fn authorize_then_persist( + provider: SubscriptionProvider, + cancel: CancellationToken, + authorize: Authorize, + persist: Persist, +) -> Result<()> +where + Authorize: std::future::Future>, + Persist: FnOnce(T) -> PersistFuture, + PersistFuture: std::future::Future>, +{ + let credential = tokio::select! { + _ = cancel.cancelled() => return Err(anyhow!("login cancelled")), + result = tokio::time::timeout(LOGIN_TIMEOUT, authorize) => match result { + Ok(result) => result?, + Err(_) => return Err(anyhow!("Login timed out")), + }, + }; + // Logout/re-login cancels the generation before waiting on this same + // provider lock. Whichever side reaches the lock boundary first wins: + // an already-started commit finishes before logout deletes it, while a + // cancelled commit waiting on the lock is discarded before writing. + let _guard = store_lock(provider).lock().await; + if cancel.is_cancelled() { + return Err(anyhow!("login cancelled")); + } + persist(credential).await +} + +pub(crate) fn require_current_store_revision( + provider: SubscriptionProvider, + outcome: store::ConditionalCommitOutcome, +) -> Result { + match outcome { + store::ConditionalCommitOutcome::Committed { revision } => Ok(revision), + store::ConditionalCommitOutcome::Conflict { current_revision } => Err(anyhow!( + "{} credentials changed in another BitFun process (current revision {current_revision}); retry the operation", + provider.display_label() + )), + } +} + fn build_account( provider: SubscriptionProvider, entry: Option<&StoredCredential>, + reauthentication_required: bool, + vault_unavailable: bool, ) -> SubscriptionAccount { let (format, base_url, model) = provider.suggested(); let (connected, account, expires_at) = match entry { @@ -223,6 +299,8 @@ fn build_account( account, expires_at, connected, + reauthentication_required, + vault_unavailable, suggested_format: format.to_string(), suggested_base_url: base_url.to_string(), suggested_model: model.to_string(), @@ -230,78 +308,175 @@ fn build_account( } async fn account_snapshot(provider: SubscriptionProvider) -> SubscriptionAccount { - let store = store::load().await.unwrap_or_default(); - build_account(provider, store.get(provider.key())) + let state = store::load_with_state().await.unwrap_or_else(|error| { + log::warn!("load subscription credential state failed: {error:#}"); + store::LoadState { + credentials: store::Store::new(), + requires_reauthentication: std::collections::HashSet::new(), + vault_unavailable: std::collections::HashSet::new(), + provider_revisions: std::collections::HashMap::new(), + } + }); + build_account( + provider, + state.credentials.get(provider.key()), + state.requires_reauthentication.contains(provider.key()), + state.vault_unavailable.contains(provider.key()), + ) } /// Lists all providers with their current connection state. pub async fn list_accounts() -> Vec { - let store = store::load().await.unwrap_or_default(); + let state = store::load_with_state().await.unwrap_or_else(|error| { + log::warn!("load subscription credential state failed: {error:#}"); + store::LoadState { + credentials: store::Store::new(), + requires_reauthentication: std::collections::HashSet::new(), + vault_unavailable: std::collections::HashSet::new(), + provider_revisions: std::collections::HashMap::new(), + } + }); SubscriptionProvider::ALL .iter() - .map(|provider| build_account(*provider, store.get(provider.key()))) + .map(|provider| { + build_account( + *provider, + state.credentials.get(provider.key()), + state.requires_reauthentication.contains(provider.key()), + state.vault_unavailable.contains(provider.key()), + ) + }) .collect() } /// Starts a login session, cancelling any existing pending session for the /// same provider. Returns immediately with the authorization URL / user code. -pub async fn start_login(provider: SubscriptionProvider) -> Result { - if let Some(previous) = { +pub async fn start_login( + provider: SubscriptionProvider, + session_id: String, +) -> Result { + validate_session_id(&session_id)?; + let cancel = CancellationToken::new(); + let generation = next_generation(); + // Serialize the durable revision snapshot with any local refresh/commit and + // install the replacement session before releasing that boundary. A prior + // local login can then either finish before this snapshot or observe its + // cancellation; it cannot commit between the snapshot and replacement. + let provider_guard = store_lock(provider).lock().await; + let expected_revision = store::credential_revision(provider.key()).await?; + { let mut map = sessions() .lock() .map_err(|_| anyhow!("subscription login session lock poisoned"))?; - map.remove(&provider) - } { - previous.cancel.cancel(); + if let Some(previous) = map.insert( + provider, + SessionState { + session_id: session_id.clone(), + status: LoginStatus::Pending, + authorization_url: None, + user_code: None, + instructions: None, + error: None, + account: None, + cancel: cancel.clone(), + generation, + }, + ) { + previous.cancel.cancel(); + } } + drop(provider_guard); - let cancel = CancellationToken::new(); - let started = match provider { - SubscriptionProvider::Codex => codex::begin_login(cancel.clone()).await, - SubscriptionProvider::Antigravity => antigravity::begin_login(cancel.clone()).await, - SubscriptionProvider::Opencode => opencode::begin_login(cancel.clone()).await, - }?; + // The placeholder above makes cancellation visible even while a provider + // is still binding its callback listener or requesting a device code. + let begin = async { + match provider { + SubscriptionProvider::Codex => { + codex::begin_login(cancel.clone(), expected_revision).await + } + SubscriptionProvider::Antigravity => { + antigravity::begin_login(cancel.clone(), expected_revision).await + } + SubscriptionProvider::Opencode => { + opencode::begin_login(cancel.clone(), expected_revision).await + } + } + }; + let started_result = tokio::select! { + _ = cancel.cancelled() => Err(anyhow!("login cancelled")), + result = begin => result, + }; + let started = match started_result { + Ok(started) if !cancel.is_cancelled() => started, + Ok(_) => return Err(anyhow!("login cancelled")), + Err(error) => { + if let Ok(mut map) = sessions().lock() { + if let Some(state) = map.get_mut(&provider).filter(|state| { + state.generation == generation && state.session_id == session_id + }) { + state.status = if cancel.is_cancelled() { + LoginStatus::Cancelled + } else { + LoginStatus::Failed + }; + state.error = Some(format!("{error:#}")); + } + } + return Err(error); + } + }; let authorization_url = started.authorization_url.clone(); // Desktop opener rejects relative URLs ("Not allowed to open url /..."). // Every provider must return an absolute http(s) authorization URL. if !(authorization_url.starts_with("https://") || authorization_url.starts_with("http://")) { cancel.cancel(); + if let Ok(mut map) = sessions().lock() { + if let Some(state) = map + .get_mut(&provider) + .filter(|state| state.generation == generation && state.session_id == session_id) + { + state.status = LoginStatus::Failed; + state.error = Some( + "Subscription login returned a non-absolute authorization URL".to_string(), + ); + } + } return Err(anyhow!( "subscription login returned a non-absolute authorization URL: {authorization_url}" )); } let user_code = started.user_code.clone(); let instructions = started.instructions.clone(); - let generation = next_generation(); - { let mut map = sessions() .lock() .map_err(|_| anyhow!("subscription login session lock poisoned"))?; - map.insert( - provider, - SessionState { - status: LoginStatus::Pending, - authorization_url: Some(authorization_url.clone()), - user_code: user_code.clone(), - instructions: Some(instructions.clone()), - error: None, - account: None, - cancel: cancel.clone(), - generation, - }, - ); + let Some(state) = map.get_mut(&provider).filter(|state| { + state.generation == generation + && state.session_id == session_id + && !state.cancel.is_cancelled() + }) else { + cancel.cancel(); + return Err(anyhow!("login cancelled")); + }; + state.authorization_url = Some(authorization_url.clone()); + state.user_code = user_code.clone(); + state.instructions = Some(instructions.clone()); } let runner = started.runner; + let runner_session_id = session_id.clone(); tokio::spawn(async move { - let outcome = tokio::time::timeout(LOGIN_TIMEOUT, runner).await; - finalize_session(provider, generation, &cancel, outcome).await; + // Authorization timeout lives inside `authorize_then_persist`; once + // persistence begins it must not be dropped by a surrounding timeout. + let outcome: Result, tokio::time::error::Elapsed> = Ok(runner.await); + finalize_session(provider, &runner_session_id, generation, &cancel, outcome).await; }); Ok(LoginStartResult { provider, + session_id, authorization_url, user_code, instructions, @@ -310,6 +485,7 @@ pub async fn start_login(provider: SubscriptionProvider) -> Result, tokio::time::error::Elapsed>, @@ -319,8 +495,9 @@ async fn finalize_session( let is_current = sessions() .lock() .map(|map| { - map.get(&provider) - .is_some_and(|state| state.generation == generation) + map.get(&provider).is_some_and(|state| { + state.generation == generation && state.session_id == session_id + }) }) .unwrap_or(false); if !is_current { @@ -350,8 +527,22 @@ async fn finalize_session( } }; + update_session_if_current(provider, session_id, generation, status, error, account); +} + +fn update_session_if_current( + provider: SubscriptionProvider, + session_id: &str, + generation: u64, + status: LoginStatus, + error: Option, + account: Option, +) { if let Ok(mut map) = sessions().lock() { - if let Some(state) = map.get_mut(&provider) { + if let Some(state) = map + .get_mut(&provider) + .filter(|state| state.generation == generation && state.session_id == session_id) + { state.status = status; state.error = error; if account.is_some() { @@ -361,44 +552,68 @@ async fn finalize_session( } } -/// Returns the current login session snapshot for a provider. -pub async fn login_status(provider: SubscriptionProvider) -> LoginSessionSnapshot { - if let Ok(map) = sessions().lock() { - if let Some(state) = map.get(&provider) { - return state.snapshot(provider); - } - } - - let account = account_snapshot(provider).await; - let status = if account.connected { - LoginStatus::Authorized - } else { - LoginStatus::Failed - }; - LoginSessionSnapshot { - provider, - status, - authorization_url: None, - user_code: None, - instructions: None, - error: None, - account: account.connected.then_some(account), - } +/// Returns a login snapshot only when both provider and session id still refer +/// to the same current operation. +pub async fn login_status( + provider: SubscriptionProvider, + session_id: &str, +) -> Result { + validate_session_id(session_id)?; + let map = sessions() + .lock() + .map_err(|_| anyhow!("subscription login session lock poisoned"))?; + map.get(&provider) + .filter(|state| state.session_id == session_id) + .map(|state| state.snapshot(provider)) + .ok_or_else(|| anyhow!("subscription login session is no longer current")) } /// Cancels an in-flight login session for a provider. -pub async fn cancel_login(provider: SubscriptionProvider) { - if let Ok(mut map) = sessions().lock() { - if let Some(state) = map.get_mut(&provider) { - state.cancel.cancel(); - state.status = LoginStatus::Cancelled; - state.error = Some("Login cancelled".to_string()); +pub async fn cancel_login(provider: SubscriptionProvider, session_id: &str) -> Result<()> { + validate_session_id(session_id)?; + let wait_for_commit_barrier = if let Ok(mut map) = sessions().lock() { + if let Some(state) = map + .get_mut(&provider) + .filter(|state| state.session_id == session_id) + { + match state.status { + LoginStatus::Pending => { + state.cancel.cancel(); + state.status = LoginStatus::Cancelled; + state.error = Some("Login cancelled".to_string()); + true + } + // A duplicate cancel can observe the state update performed by + // the first cancel before that call reaches the credential + // commit barrier. It must join the same barrier instead of + // reporting completion while persistence may still succeed. + LoginStatus::Cancelled => true, + // Authorization may have already committed and finalized. + // Never rewrite an Authorized terminal state to Cancelled; + // that would disagree with the connected credential. + LoginStatus::Authorized | LoginStatus::Failed => false, + } + } else { + false } + } else { + return Err(anyhow!("subscription login session lock poisoned")); + }; + // A stale cancel from an older UI request is a no-op and must not wait on + // or interfere with the replacement session's commit. + if !wait_for_commit_barrier { + return Ok(()); } + // Act as a completion barrier for the commit phase. If cancellation wins + // the provider lock, the runner observes the cancelled token and skips its + // write. If persistence already owns the lock, let that atomic commit + // finish before reporting cancellation back to the UI. + let _guard = store_lock(provider).lock().await; + Ok(()) } /// Removes the stored credential for a provider. -pub async fn logout(provider: SubscriptionProvider) -> Result<()> { +pub async fn logout(provider: SubscriptionProvider) -> Result { // Cancel any in-flight login first so its runner cannot persist fresh // tokens after the logout completes. if let Ok(mut map) = sessions().lock() { @@ -407,12 +622,26 @@ pub async fn logout(provider: SubscriptionProvider) -> Result<()> { } } let _guard = store_lock(provider).lock().await; - let mut store = store::load().await.unwrap_or_default(); - store.remove(provider.key()); - store::save(&store).await?; + let outcome = store::remove(provider.key()).await?; drop(_guard); log::info!("subscription provider {} logged out", provider.key()); - Ok(()) + Ok(match outcome { + store::RemoveOutcome::Removed => SubscriptionLogoutResult { + cleanup_pending: false, + warning: None, + }, + store::RemoveOutcome::CleanupPending(warning) => { + log::warn!( + "subscription provider {} logged out with native credential cleanup pending: {}", + provider.key(), + warning + ); + SubscriptionLogoutResult { + cleanup_pending: true, + warning: Some(warning), + } + } + }) } /// Resolves a runtime credential for a provider, refreshing tokens if needed. @@ -435,6 +664,11 @@ mod tests { use super::store::{self, StoredCredential}; use super::*; + const STALE_LOGIN_CHILD_METADATA_ENV: &str = "BITFUN_SUBAUTH_CAS_CHILD_METADATA"; + const STALE_LOGIN_CHILD_LOADED_ENV: &str = "BITFUN_SUBAUTH_CAS_CHILD_LOADED"; + const STALE_LOGIN_CHILD_RESUME_ENV: &str = "BITFUN_SUBAUTH_CAS_CHILD_RESUME"; + const STALE_LOGIN_CHILD_OUTCOME_ENV: &str = "BITFUN_SUBAUTH_CAS_CHILD_OUTCOME"; + /// Serializes tests that rely on the process-global store path override. fn test_lock() -> &'static Mutex<()> { static LOCK: OnceLock> = OnceLock::new(); @@ -447,6 +681,10 @@ mod tests { dir.join("subscription_auth.json") } + fn test_session_id() -> String { + uuid::Uuid::new_v4().to_string() + } + #[test] fn subscription_provider_serde_roundtrip() { assert_eq!( @@ -486,6 +724,11 @@ mod tests { ); store::save(&store).await.unwrap(); + let metadata_file = std::fs::read_to_string(store_path_override_for_assertion()).unwrap(); + assert!(!metadata_file.contains("refresh-token")); + assert!(!metadata_file.contains("access-token")); + assert!(metadata_file.contains("user@example.com")); + let loaded = store::load().await.unwrap(); let entry = loaded.get("codex").expect("codex entry present"); match entry { @@ -506,6 +749,514 @@ mod tests { assert!(codex.connected); assert_eq!(codex.account.as_deref(), Some("user@example.com")); assert_eq!(codex.expires_at, Some(1_800_000_000)); + assert!(!codex.reauthentication_required); + } + + fn store_path_override_for_assertion() -> std::path::PathBuf { + super::store::store_path_for_test_assertion() + } + + #[tokio::test] + async fn legacy_plaintext_store_is_migrated_and_scrubbed() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + let legacy = serde_json::json!({ + "opencode": { + "type": "oauth", + "refresh": "legacy-refresh-secret", + "access": "legacy-access-secret", + "expires": 1_900_000_000_000_i64, + "metadata": { "email": "legacy@example.com" } + } + }); + std::fs::write(&path, serde_json::to_vec_pretty(&legacy).unwrap()).unwrap(); + + let loaded = store::load().await.unwrap(); + assert!(loaded.contains_key("opencode")); + + let migrated = std::fs::read_to_string(&path).unwrap(); + assert!(migrated.contains("\"version\": 2")); + assert!(migrated.contains("legacy@example.com")); + assert!(!migrated.contains("legacy-refresh-secret")); + assert!(!migrated.contains("legacy-access-secret")); + } + + #[tokio::test] + async fn legacy_migration_retries_after_temporary_vault_unavailability() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + let legacy = serde_json::json!({ + "opencode": { + "type": "oauth", + "refresh": "legacy-retry-refresh", + "access": "legacy-retry-access", + "expires": 1_900_000_000_000_i64 + } + }); + std::fs::write(&path, serde_json::to_vec_pretty(&legacy).unwrap()).unwrap(); + + store::set_test_vault_unavailable(true); + let deferred = store::load_with_state().await.unwrap(); + assert!(deferred.credentials.is_empty()); + assert!(deferred.vault_unavailable.contains("opencode")); + assert!(!deferred.requires_reauthentication.contains("opencode")); + let unchanged = std::fs::read_to_string(&path).unwrap(); + assert!(unchanged.contains("legacy-retry-refresh")); + assert!(unchanged.contains("legacy-retry-access")); + + store::set_test_vault_unavailable(false); + let migrated = store::load().await.unwrap(); + assert!(migrated.contains_key("opencode")); + let scrubbed = std::fs::read_to_string(&path).unwrap(); + assert!(scrubbed.contains("\"version\": 2")); + assert!(!scrubbed.contains("legacy-retry-refresh")); + assert!(!scrubbed.contains("legacy-retry-access")); + assert!(store::cleanup_journal_entries_for_assertion() + .await + .is_empty()); + } + + #[tokio::test] + async fn partial_chunk_write_is_durably_cleaned_after_retry() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + store::set_store_path_for_test(temp_store_path()); + store::set_test_vault_write_failure_after(Some(1)); + store::set_test_vault_delete_failure(true); + + let error = store::upsert( + "codex", + StoredCredential::Oauth { + refresh: "r".repeat(3_000), + access: "a".repeat(3_000), + expires: 1_900_000_000_000, + account_id: None, + metadata: None, + }, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("vault write failure")); + assert!(!store::test_vault_entries_for_assertion().is_empty()); + assert!(!store::cleanup_journal_entries_for_assertion() + .await + .is_empty()); + + store::set_test_vault_write_failure_after(None); + store::set_test_vault_delete_failure(false); + let loaded = store::load().await.unwrap(); + assert!(loaded.is_empty()); + assert!(store::test_vault_entries_for_assertion().is_empty()); + assert!(store::cleanup_journal_entries_for_assertion() + .await + .is_empty()); + } + + #[tokio::test] + async fn windows_post_commit_backup_cleanup_failure_does_not_fail_commit() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + let tmp = path.with_extension("tmp-one"); + let backup = path.with_extension("bak"); + std::fs::write(&path, b"legacy-plaintext-secret").unwrap(); + std::fs::write(&tmp, b"new-metadata").unwrap(); + + store::set_test_backup_cleanup_failure(&backup, true); + store::replace_metadata_file_windows(&tmp, &path) + .await + .unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), b"new-metadata"); + assert_eq!(std::fs::read(&backup).unwrap(), b"legacy-plaintext-secret"); + + store::set_test_backup_cleanup_failure(&backup, false); + let next_tmp = path.with_extension("tmp-two"); + std::fs::write(&next_tmp, b"newer-metadata").unwrap(); + store::replace_metadata_file_windows(&next_tmp, &path) + .await + .unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), b"newer-metadata"); + assert!(!backup.exists()); + } + + #[tokio::test] + async fn concurrent_provider_upserts_preserve_both_metadata_entries() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + + let codex = store::upsert( + "codex", + StoredCredential::Oauth { + refresh: "codex-refresh".to_string(), + access: "codex-access".to_string(), + expires: 1_900_000_000_000, + account_id: Some("codex-account".to_string()), + metadata: None, + }, + ); + let opencode = store::upsert( + "opencode", + StoredCredential::Oauth { + refresh: "opencode-refresh".to_string(), + access: "opencode-access".to_string(), + expires: 1_900_000_000_000, + account_id: None, + metadata: Some(serde_json::json!({ "email": "zen@example.com" })), + }, + ); + let (codex_result, opencode_result) = tokio::join!(codex, opencode); + codex_result.unwrap(); + opencode_result.unwrap(); + + let loaded = store::load().await.unwrap(); + assert!(loaded.contains_key("codex")); + assert!(loaded.contains_key("opencode")); + let metadata = std::fs::read_to_string(path).unwrap(); + assert!(metadata.contains("\"codex\"")); + assert!(metadata.contains("\"opencode\"")); + assert!(!metadata.contains("codex-access")); + assert!(!metadata.contains("opencode-access")); + } + + #[tokio::test] + async fn logout_tombstone_wins_over_a_refresh_paused_after_load() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + store::upsert( + "codex", + StoredCredential::Oauth { + refresh: "refresh-before-logout".to_string(), + access: "access-before-logout".to_string(), + expires: 1, + account_id: None, + metadata: None, + }, + ) + .await + .unwrap(); + + // Model a second process paused in the external refresh request after + // it has loaded the old credential and revision. + let (loaded_tx, loaded_rx) = tokio::sync::oneshot::channel(); + let (resume_tx, resume_rx) = tokio::sync::oneshot::channel(); + let stale_refresh = tokio::spawn(async move { + let snapshot = store::load_entry_with_revision("codex").await?; + loaded_tx + .send(snapshot.revision) + .map_err(|_| anyhow!("refresh load signal receiver dropped"))?; + resume_rx + .await + .map_err(|_| anyhow!("refresh resume signal sender dropped"))?; + store::upsert_if_revision( + "codex", + snapshot.revision, + StoredCredential::Oauth { + refresh: "stale-rotated-refresh".to_string(), + access: "stale-refreshed-access".to_string(), + expires: 1_900_000_000_000, + account_id: None, + metadata: None, + }, + ) + .await + }); + + let loaded_revision = loaded_rx.await.unwrap(); + let remove_outcome = store::remove("codex").await.unwrap(); + assert!(matches!(remove_outcome, store::RemoveOutcome::Removed)); + let logout_revision = store::credential_revision("codex").await.unwrap(); + assert!(logout_revision > loaded_revision); + resume_tx.send(()).unwrap(); + + let refresh_outcome = stale_refresh.await.unwrap().unwrap(); + assert_eq!( + refresh_outcome, + store::ConditionalCommitOutcome::Conflict { + current_revision: logout_revision, + } + ); + assert!(store::load_entry("codex").await.unwrap().is_none()); + + let metadata: serde_json::Value = + serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert!(metadata["accounts"].get("codex").is_none()); + assert_eq!( + metadata["provider_revisions"]["codex"].as_u64(), + Some(logout_revision) + ); + } + + #[tokio::test] + async fn cross_process_stale_login_cas_child() { + let Some(path) = + std::env::var_os(STALE_LOGIN_CHILD_METADATA_ENV).map(std::path::PathBuf::from) + else { + return; + }; + let loaded_path = std::path::PathBuf::from( + std::env::var_os(STALE_LOGIN_CHILD_LOADED_ENV).expect("child loaded marker path"), + ); + let resume_path = std::path::PathBuf::from( + std::env::var_os(STALE_LOGIN_CHILD_RESUME_ENV).expect("child resume marker path"), + ); + let outcome_path = std::path::PathBuf::from( + std::env::var_os(STALE_LOGIN_CHILD_OUTCOME_ENV).expect("child outcome marker path"), + ); + store::set_store_path_for_test(path); + + let login_revision = store::credential_revision("opencode").await.unwrap(); + assert_eq!(login_revision, 0); + std::fs::write(&loaded_path, b"loaded").unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while !resume_path.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("parent should release the stale login after logout"); + + let outcome = store::upsert_if_revision( + "opencode", + login_revision, + StoredCredential::Api { + key: "stale-login-key".to_string(), + metadata: None, + }, + ) + .await + .unwrap(); + let store::ConditionalCommitOutcome::Conflict { current_revision } = outcome else { + panic!("stale cross-process login unexpectedly committed: {outcome:?}"); + }; + std::fs::write(outcome_path, current_revision.to_string()).unwrap(); + } + + #[tokio::test] + async fn logout_of_an_absent_provider_invalidates_a_cross_process_login() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + let parent = path.parent().unwrap(); + let loaded_path = parent.join("child-loaded"); + let resume_path = parent.join("child-resume"); + let outcome_path = parent.join("child-outcome"); + let mut child = std::process::Command::new(std::env::current_exe().unwrap()) + .arg("--exact") + .arg("subscription_auth::tests::cross_process_stale_login_cas_child") + .arg("--nocapture") + .env(STALE_LOGIN_CHILD_METADATA_ENV, &path) + .env(STALE_LOGIN_CHILD_LOADED_ENV, &loaded_path) + .env(STALE_LOGIN_CHILD_RESUME_ENV, &resume_path) + .env(STALE_LOGIN_CHILD_OUTCOME_ENV, &outcome_path) + .spawn() + .expect("spawn stale-login child process"); + + tokio::time::timeout(Duration::from_secs(5), async { + while !loaded_path.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("child should capture the pre-logout revision"); + store::remove("opencode").await.unwrap(); + let logout_revision = store::credential_revision("opencode").await.unwrap(); + std::fs::write(&resume_path, b"resume").unwrap(); + + let status = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Some(status) = child.try_wait().expect("poll stale-login child process") { + break status; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("stale-login child should finish after resume"); + assert!(status.success(), "stale-login child failed: {status}"); + assert_eq!( + std::fs::read_to_string(outcome_path).unwrap(), + logout_revision.to_string() + ); + assert!(store::load_entry("opencode").await.unwrap().is_none()); + } + + #[tokio::test] + async fn v2_metadata_without_revision_map_remains_conditionally_writable() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + std::fs::write( + &path, + serde_json::to_vec_pretty(&serde_json::json!({ + "version": 2, + "accounts": {} + })) + .unwrap(), + ) + .unwrap(); + + assert_eq!(store::credential_revision("antigravity").await.unwrap(), 0); + let outcome = store::upsert_if_revision( + "antigravity", + 0, + StoredCredential::Api { + key: "compatible-key".to_string(), + metadata: None, + }, + ) + .await + .unwrap(); + assert_eq!( + outcome, + store::ConditionalCommitOutcome::Committed { revision: 1 } + ); + assert!(store::load_entry("antigravity").await.unwrap().is_some()); + } + + #[tokio::test] + async fn repeated_upsert_replaces_existing_metadata_file() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + + store::upsert( + "codex", + StoredCredential::Oauth { + refresh: "old-refresh".to_string(), + access: "old-access".to_string(), + expires: 1_800_000_000_000, + account_id: None, + metadata: None, + }, + ) + .await + .unwrap(); + store::upsert( + "codex", + StoredCredential::Oauth { + refresh: "new-refresh".to_string(), + access: "new-access".to_string(), + expires: 1_900_000_000_000, + account_id: Some("updated-account".to_string()), + metadata: None, + }, + ) + .await + .unwrap(); + + let loaded = store::load_entry("codex").await.unwrap().unwrap(); + match loaded { + StoredCredential::Oauth { + refresh, + access, + expires, + account_id, + .. + } => { + assert_eq!(refresh, "new-refresh"); + assert_eq!(access, "new-access"); + assert_eq!(expires, 1_900_000_000_000); + assert_eq!(account_id.as_deref(), Some("updated-account")); + } + _ => panic!("expected oauth credential"), + } + let metadata = std::fs::read_to_string(path).unwrap(); + assert!(!metadata.contains("old-access")); + assert!(!metadata.contains("new-access")); + assert!(metadata.contains("updated-account")); + } + + #[tokio::test] + async fn long_tokens_are_split_below_the_windows_vault_limit() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let path = temp_store_path(); + store::set_store_path_for_test(path.clone()); + let refresh = "r".repeat(5_000); + let access = "a".repeat(9_000); + + store::upsert( + "codex", + StoredCredential::Oauth { + refresh: refresh.clone(), + access: access.clone(), + expires: 1_900_000_000_000, + account_id: None, + metadata: None, + }, + ) + .await + .unwrap(); + + let entries = store::test_vault_entries_for_assertion(); + assert!(entries.len() > 2, "long tokens must use multiple entries"); + assert!(entries.keys().all(|name| name != "codex")); + assert!(entries.values().all(|part| part.len() <= 2_048)); + let loaded = store::load_entry("codex").await.unwrap().unwrap(); + match loaded { + StoredCredential::Oauth { + refresh: loaded_refresh, + access: loaded_access, + .. + } => { + assert_eq!(loaded_refresh, refresh); + assert_eq!(loaded_access, access); + } + _ => panic!("expected oauth credential"), + } + let metadata = std::fs::read_to_string(path).unwrap(); + assert!(!metadata.contains(&refresh)); + assert!(!metadata.contains(&access)); + } + + #[tokio::test] + async fn unavailable_vault_is_retryable_not_missing_credential() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + store::set_store_path_for_test(temp_store_path()); + store::upsert( + "opencode", + StoredCredential::Api { + key: "sk-present-but-locked".to_string(), + metadata: None, + }, + ) + .await + .unwrap(); + + store::set_test_vault_unavailable(true); + let state = store::load_with_state().await.unwrap(); + assert!(state.credentials.get("opencode").is_none()); + assert!(!state.requires_reauthentication.contains("opencode")); + assert!(state.vault_unavailable.contains("opencode")); + let error = store::load_entry("opencode").await.unwrap_err(); + assert!(error.to_string().contains("locked or unavailable")); + store::set_test_vault_unavailable(false); + + let restored = store::load_entry("opencode").await.unwrap(); + assert!(restored.is_some()); } #[tokio::test] @@ -529,15 +1280,88 @@ mod tests { assert!(loaded.get("opencode").is_none()); } + #[tokio::test] + async fn failed_logout_metadata_commit_preserves_usable_credential() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + store::set_store_path_for_test(temp_store_path()); + store::upsert( + "opencode", + StoredCredential::Api { + key: "sk-still-usable".to_string(), + metadata: None, + }, + ) + .await + .unwrap(); + let entries_before = store::test_vault_entries_for_assertion(); + + store::set_test_metadata_write_failure(true); + let error = store::remove("opencode").await.unwrap_err(); + assert!(error.to_string().contains("injected")); + store::set_test_metadata_write_failure(false); + + assert_eq!(store::test_vault_entries_for_assertion(), entries_before); + let loaded = store::load_entry("opencode").await.unwrap().unwrap(); + match loaded { + StoredCredential::Api { key, .. } => assert_eq!(key, "sk-still-usable"), + _ => panic!("expected api credential"), + } + } + + #[tokio::test] + async fn failed_logout_vault_delete_is_reported_and_retried() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + store::set_store_path_for_test(temp_store_path()); + store::upsert( + "opencode", + StoredCredential::Api { + key: "sk-pending-delete".to_string(), + metadata: None, + }, + ) + .await + .unwrap(); + + store::set_test_vault_delete_failure(true); + let outcome = logout(SubscriptionProvider::Opencode).await.unwrap(); + assert!(outcome.cleanup_pending); + assert!(outcome + .warning + .as_deref() + .is_some_and(|warning| warning.contains("cleanup is pending"))); + assert!(!store::test_vault_entries_for_assertion().is_empty()); + assert!(!store::cleanup_journal_entries_for_assertion() + .await + .is_empty()); + assert!(store::load_entry("opencode").await.unwrap().is_none()); + + store::set_test_vault_delete_failure(false); + assert!(store::load().await.unwrap().is_empty()); + assert!(store::test_vault_entries_for_assertion().is_empty()); + assert!(store::cleanup_journal_entries_for_assertion() + .await + .is_empty()); + } + #[tokio::test] async fn finalize_ignores_superseded_session() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); let provider = SubscriptionProvider::Codex; let stale_generation = next_generation(); + let stale_session_id = test_session_id(); + let current_session_id = test_session_id(); { let mut map = sessions().lock().unwrap(); map.insert( provider, SessionState { + session_id: current_session_id, status: LoginStatus::Pending, authorization_url: None, user_code: None, @@ -554,6 +1378,7 @@ mod tests { // pending session when it finishes. finalize_session( provider, + &stale_session_id, stale_generation, &CancellationToken::new(), Ok(Err(anyhow!("stale runner failed"))), @@ -568,4 +1393,293 @@ mod tests { }; assert_eq!(status, Some(LoginStatus::Pending)); } + + #[tokio::test] + async fn cancellation_does_not_drop_started_credential_persistence() { + let cancel = CancellationToken::new(); + let (persist_started_tx, persist_started_rx) = tokio::sync::oneshot::channel(); + let (allow_persist_tx, allow_persist_rx) = tokio::sync::oneshot::channel(); + + let task = tokio::spawn(authorize_then_persist( + SubscriptionProvider::Codex, + cancel.clone(), + async { Ok::<_, anyhow::Error>("authorized-token") }, + move |token| async move { + assert_eq!(token, "authorized-token"); + persist_started_tx.send(()).unwrap(); + allow_persist_rx.await.unwrap(); + Ok(()) + }, + )); + + persist_started_rx.await.unwrap(); + cancel.cancel(); + tokio::task::yield_now().await; + assert!(!task.is_finished()); + + allow_persist_tx.send(()).unwrap(); + assert!(task.await.unwrap().is_ok()); + } + + #[tokio::test] + async fn cancellation_before_authorization_skips_persistence() { + let cancel = CancellationToken::new(); + cancel.cancel(); + let persisted = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let persisted_for_task = persisted.clone(); + + let result = authorize_then_persist( + SubscriptionProvider::Opencode, + cancel, + std::future::pending::>(), + move |_| async move { + persisted_for_task.store(true, Ordering::SeqCst); + Ok(()) + }, + ) + .await; + + assert!(result.is_err()); + assert!(!persisted.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn cancelled_commit_waiting_for_provider_lock_does_not_persist() { + let provider = SubscriptionProvider::Antigravity; + let store_guard = store_lock(provider).lock().await; + let cancel = CancellationToken::new(); + let (authorized_tx, authorized_rx) = tokio::sync::oneshot::channel(); + let persisted = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let persisted_for_task = persisted.clone(); + + let task = tokio::spawn(authorize_then_persist( + provider, + cancel.clone(), + async move { + authorized_tx.send(()).unwrap(); + Ok::<_, anyhow::Error>("authorized-token") + }, + move |_| async move { + persisted_for_task.store(true, Ordering::SeqCst); + Ok(()) + }, + )); + + authorized_rx.await.unwrap(); + cancel.cancel(); + drop(store_guard); + + assert!(task.await.unwrap().is_err()); + assert!(!persisted.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn cancel_command_waits_for_the_commit_boundary() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let provider = SubscriptionProvider::Opencode; + let store_guard = store_lock(provider).lock().await; + let cancel = CancellationToken::new(); + let generation = next_generation(); + let session_id = test_session_id(); + { + let mut map = sessions().lock().unwrap(); + map.insert( + provider, + SessionState { + session_id: session_id.clone(), + status: LoginStatus::Pending, + authorization_url: None, + user_code: None, + instructions: None, + error: None, + account: None, + cancel: cancel.clone(), + generation, + }, + ); + } + + let task_session_id = session_id.clone(); + let task = + tokio::spawn(async move { cancel_login(provider, &task_session_id).await.unwrap() }); + tokio::task::yield_now().await; + assert!(cancel.is_cancelled()); + assert!(!task.is_finished()); + + drop(store_guard); + task.await.unwrap(); + let status = sessions() + .lock() + .unwrap() + .remove(&provider) + .map(|state| state.status); + assert_eq!(status, Some(LoginStatus::Cancelled)); + } + + #[tokio::test] + async fn duplicate_cancel_waits_for_the_same_commit_boundary() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let provider = SubscriptionProvider::Antigravity; + let store_guard = store_lock(provider).lock().await; + let cancel = CancellationToken::new(); + let session_id = test_session_id(); + { + let mut map = sessions().lock().unwrap(); + map.insert( + provider, + SessionState { + session_id: session_id.clone(), + status: LoginStatus::Pending, + authorization_url: None, + user_code: None, + instructions: None, + error: None, + account: None, + cancel: cancel.clone(), + generation: next_generation(), + }, + ); + } + + let first_session_id = session_id.clone(); + let first_cancel = + tokio::spawn(async move { cancel_login(provider, &first_session_id).await.unwrap() }); + tokio::task::yield_now().await; + assert!(cancel.is_cancelled()); + assert!(!first_cancel.is_finished()); + + let second_session_id = session_id.clone(); + let second_cancel = + tokio::spawn(async move { cancel_login(provider, &second_session_id).await.unwrap() }); + tokio::task::yield_now().await; + assert!(!second_cancel.is_finished()); + + drop(store_guard); + first_cancel.await.unwrap(); + second_cancel.await.unwrap(); + let status = sessions() + .lock() + .unwrap() + .remove(&provider) + .map(|state| state.status); + assert_eq!(status, Some(LoginStatus::Cancelled)); + } + + #[tokio::test] + async fn stale_cancel_does_not_cancel_replacement_session() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let provider = SubscriptionProvider::Opencode; + let stale_session_id = test_session_id(); + let current_session_id = test_session_id(); + let current_cancel = CancellationToken::new(); + { + let mut map = sessions().lock().unwrap(); + map.insert( + provider, + SessionState { + session_id: current_session_id.clone(), + status: LoginStatus::Pending, + authorization_url: None, + user_code: None, + instructions: None, + error: None, + account: None, + cancel: current_cancel.clone(), + generation: next_generation(), + }, + ); + } + + cancel_login(provider, &stale_session_id).await.unwrap(); + assert!(!current_cancel.is_cancelled()); + let snapshot = login_status(provider, ¤t_session_id).await.unwrap(); + assert_eq!(snapshot.session_id, current_session_id); + assert_eq!(snapshot.status, LoginStatus::Pending); + sessions().lock().unwrap().remove(&provider); + } + + #[tokio::test] + async fn cancel_does_not_rewrite_authorized_terminal_state() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let provider = SubscriptionProvider::Codex; + let session_id = test_session_id(); + let cancel = CancellationToken::new(); + { + let mut map = sessions().lock().unwrap(); + map.insert( + provider, + SessionState { + session_id: session_id.clone(), + status: LoginStatus::Authorized, + authorization_url: None, + user_code: None, + instructions: None, + error: None, + account: None, + cancel: cancel.clone(), + generation: next_generation(), + }, + ); + } + + cancel_login(provider, &session_id).await.unwrap(); + assert!(!cancel.is_cancelled()); + let snapshot = login_status(provider, &session_id).await.unwrap(); + assert_eq!(snapshot.status, LoginStatus::Authorized); + sessions().lock().unwrap().remove(&provider); + } + + #[test] + fn final_state_update_rechecks_generation_after_async_work() { + let _guard = test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let provider = SubscriptionProvider::Codex; + let old_generation = next_generation(); + let new_generation = next_generation(); + let old_session_id = test_session_id(); + let new_session_id = test_session_id(); + { + let mut map = sessions().lock().unwrap(); + map.insert( + provider, + SessionState { + session_id: new_session_id, + status: LoginStatus::Pending, + authorization_url: None, + user_code: None, + instructions: None, + error: None, + account: None, + cancel: CancellationToken::new(), + generation: new_generation, + }, + ); + } + + update_session_if_current( + provider, + &old_session_id, + old_generation, + LoginStatus::Authorized, + None, + None, + ); + + let status = { + let mut map = sessions().lock().unwrap(); + let status = map.get(&provider).map(|state| state.status); + map.remove(&provider); + status + }; + assert_eq!(status, Some(LoginStatus::Pending)); + } } diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/oauth_callback_locales.json b/src/crates/adapters/ai-adapters/src/subscription_auth/oauth_callback_locales.json new file mode 100644 index 0000000000..0432668850 --- /dev/null +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/oauth_callback_locales.json @@ -0,0 +1,26 @@ +{ + "en": { + "success_title": "Sign-in complete", + "success_message": "You are now signed in. You can close this window and return to BitFun.", + "error_title": "Sign-in failed", + "bad_request": "Bad request", + "missing_code": "Missing authorization code", + "invalid_state": "Invalid authorization state" + }, + "zh-CN": { + "success_title": "登录完成", + "success_message": "登录已完成。你可以关闭此窗口并返回 BitFun。", + "error_title": "登录失败", + "bad_request": "请求无效", + "missing_code": "缺少授权码", + "invalid_state": "授权状态无效" + }, + "zh-TW": { + "success_title": "登入完成", + "success_message": "登入已完成。你可以關閉此視窗並返回 BitFun。", + "error_title": "登入失敗", + "bad_request": "請求無效", + "missing_code": "缺少授權碼", + "invalid_state": "授權狀態無效" + } +} diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/oauth_server.rs b/src/crates/adapters/ai-adapters/src/subscription_auth/oauth_server.rs index 71b5ba2de2..bf33af4bee 100644 --- a/src/crates/adapters/ai-adapters/src/subscription_auth/oauth_server.rs +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/oauth_server.rs @@ -6,10 +6,16 @@ //! query parameters. use anyhow::{anyhow, Context, Result}; +use serde::Deserialize; use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; +const CALLBACK_READ_TIMEOUT: Duration = Duration::from_secs(2); +const MAX_CALLBACK_HEADER_BYTES: usize = 16 * 1024; + /// IPv4 loopback address used for the TCP listener. pub(crate) const LOOPBACK_BIND_HOST: &str = "127.0.0.1"; @@ -30,18 +36,8 @@ pub(crate) fn loopback_redirect_uri(port: u16, path: &str) -> String { format!("http://{LOOPBACK_REDIRECT_HOST}:{port}{path}") } -/// Binds the OAuth callback listener on [`LOOPBACK_BIND_HOST`]. -pub(crate) async fn bind_loopback(port: u16) -> Result { - TcpListener::bind((LOOPBACK_BIND_HOST, port)) - .await - .with_context(|| { - format!( - "bind OAuth callback on {LOOPBACK_BIND_HOST}:{port} (is another app using this port?)" - ) - }) -} - -/// Binds the first available registered callback port. +/// Binds the first available provider-supported callback port. A final `0` +/// entry requests an ephemeral port for desktop OAuth providers that permit it. /// /// Fallback is attempted only when a preferred port is already in use. The /// returned port must be used to construct both the authorize and token @@ -60,7 +56,7 @@ pub(crate) async fn bind_loopback_ports(ports: &[u16]) -> Result<(TcpListener, u .port(); if index > 0 { log::warn!( - "OAuth callback port {preferred_port} is unavailable; using registered fallback port {actual_port}" + "OAuth callback port {preferred_port} is unavailable; using fallback port {actual_port}" ); } return Ok((listener, actual_port)); @@ -90,18 +86,29 @@ pub(crate) async fn wait_for_callback( ) -> Result> { loop { let (mut stream, _) = listener.accept().await?; - let mut buf = vec![0u8; 8192]; - let n = match stream.read(&mut buf).await { - Ok(0) => continue, - Ok(n) => n, + let request_bytes = match read_http_request(&mut stream, CALLBACK_READ_TIMEOUT).await { + Ok(Some(request)) => request, + Ok(None) => continue, Err(err) => { log::debug!("subscription oauth callback read failed: {err}"); + write_response( + &mut stream, + 400, + &error_page(&callback_messages("en").bad_request, "en"), + ) + .await; continue; } }; - let request = String::from_utf8_lossy(&buf[..n]); + let request = String::from_utf8_lossy(&request_bytes); + let locale = preferred_locale(&request); let Some(request_line) = request.lines().next() else { - write_response(&mut stream, 400, &error_page("Bad request")).await; + write_response( + &mut stream, + 400, + &error_page(&callback_messages(locale).bad_request, locale), + ) + .await; continue; }; let target = request_line @@ -119,30 +126,79 @@ pub(crate) async fn wait_for_callback( } let params = parse_query(query); + // Ignore unsolicited loopback requests instead of letting a local + // process/browser probe terminate the real OAuth session. Validate + // state before accepting provider errors for the same reason. + match params.get("state") { + Some(state) if state == expected_state => {} + _ => { + write_response( + &mut stream, + 400, + &error_page(&callback_messages(locale).invalid_state, locale), + ) + .await; + continue; + } + } if let Some(error) = params.get("error") { let message = params .get("error_description") .cloned() .unwrap_or_else(|| error.clone()); - write_response(&mut stream, 200, &error_page(&message)).await; + write_response(&mut stream, 200, &error_page(&message, locale)).await; return Err(anyhow!("authorization failed: {message}")); } if params.get("code").map(String::is_empty).unwrap_or(true) { - write_response(&mut stream, 400, &error_page("Missing authorization code")).await; + write_response( + &mut stream, + 400, + &error_page(&callback_messages(locale).missing_code, locale), + ) + .await; return Err(anyhow!("authorization callback missing code")); } - match params.get("state") { - Some(state) if state == expected_state => {} - _ => { - write_response(&mut stream, 400, &error_page("Invalid state")).await; - return Err(anyhow!("authorization state mismatch")); - } - } - write_response(&mut stream, 200, &success_page()).await; + write_response(&mut stream, 200, &success_page(locale)).await; return Ok(params); } } +/// Reads one complete HTTP header block. Browser/TCP writes may split the +/// request line and headers across packets, while a local process can connect +/// and never send data; cap both total bytes and wall-clock read time so such a +/// connection cannot hold the callback listener indefinitely. +async fn read_http_request(stream: &mut R, timeout: Duration) -> Result>> +where + R: tokio::io::AsyncRead + Unpin, +{ + tokio::time::timeout(timeout, async { + let mut request = Vec::with_capacity(2048); + let mut chunk = [0u8; 2048]; + loop { + let read = stream.read(&mut chunk).await?; + if read == 0 { + if request.is_empty() { + return Ok(None); + } + return Err(anyhow!( + "OAuth callback connection closed before headers completed" + )); + } + if request.len() + read > MAX_CALLBACK_HEADER_BYTES { + return Err(anyhow!( + "OAuth callback headers exceed {MAX_CALLBACK_HEADER_BYTES} bytes" + )); + } + request.extend_from_slice(&chunk[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + return Ok(Some(request)); + } + } + }) + .await + .map_err(|_| anyhow!("OAuth callback request read timed out"))? +} + fn parse_query(query: &str) -> HashMap { let mut out = HashMap::new(); for pair in query.split('&') { @@ -181,25 +237,75 @@ async fn write_response(stream: &mut tokio::net::TcpStream, status: u16, body: & let _ = stream.flush().await; } -fn success_page() -> String { - result_page( - "Sign-in complete", - "You are now signed in. You can close this window and return to BitFun.", - ) +#[derive(Debug, Deserialize)] +struct CallbackMessages { + success_title: String, + success_message: String, + error_title: String, + bad_request: String, + missing_code: String, + invalid_state: String, } -fn error_page(message: &str) -> String { - result_page("Sign-in failed", message) +fn callback_locales() -> &'static HashMap { + static LOCALES: OnceLock> = OnceLock::new(); + LOCALES.get_or_init(|| { + serde_json::from_str(include_str!("oauth_callback_locales.json")) + .expect("embedded OAuth callback locales are valid JSON") + }) } -fn result_page(title: &str, message: &str) -> String { +fn callback_messages(locale: &str) -> &'static CallbackMessages { + callback_locales() + .get(locale) + .or_else(|| callback_locales().get("en")) + .expect("OAuth callback English locale is embedded") +} + +fn preferred_locale(request: &str) -> &'static str { + for line in request.lines() { + let Some((name, value)) = line.split_once(':') else { + continue; + }; + if !name.eq_ignore_ascii_case("accept-language") { + continue; + } + let language = value + .split(',') + .next() + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + if language.starts_with("zh-tw") || language.starts_with("zh-hk") { + return "zh-TW"; + } + if language.starts_with("zh") { + return "zh-CN"; + } + break; + } + "en" +} + +fn success_page(locale: &str) -> String { + let messages = callback_messages(locale); + result_page(locale, &messages.success_title, &messages.success_message) +} + +fn error_page(message: &str, locale: &str) -> String { + result_page(locale, &callback_messages(locale).error_title, message) +} + +fn result_page(language: &str, title: &str, message: &str) -> String { let message = escape_html(message); format!( - "{title}\ -

{title}

{message}

" ) } @@ -217,9 +323,12 @@ fn escape_html(text: &str) -> String { #[cfg(test)] mod tests { use super::{ - bind_loopback_ports, escape_html, loopback_redirect_uri, LOOPBACK_BIND_HOST, + bind_loopback_ports, callback_messages, escape_html, loopback_redirect_uri, + preferred_locale, read_http_request, success_page, wait_for_callback, LOOPBACK_BIND_HOST, LOOPBACK_REDIRECT_HOST, }; + use std::time::Duration; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; #[test] fn escapes_html_injection() { @@ -243,6 +352,105 @@ mod tests { assert_eq!(LOOPBACK_REDIRECT_HOST, "localhost"); } + #[tokio::test] + async fn invalid_state_does_not_terminate_the_real_callback_session() { + let listener = tokio::net::TcpListener::bind((LOOPBACK_BIND_HOST, 0)) + .await + .unwrap(); + let address = listener.local_addr().unwrap(); + let waiter = tokio::spawn(async move { + wait_for_callback(listener, "/auth/callback", "expected-state").await + }); + + let mut invalid = tokio::net::TcpStream::connect(address).await.unwrap(); + invalid + .write_all( + b"GET /auth/callback?error=denied&state=attacker-state HTTP/1.1\r\nHost: localhost\r\n\r\n", + ) + .await + .unwrap(); + let mut invalid_response = Vec::new(); + invalid.read_to_end(&mut invalid_response).await.unwrap(); + assert!(String::from_utf8_lossy(&invalid_response).contains("400 Bad Request")); + assert!(!waiter.is_finished()); + + let mut valid = tokio::net::TcpStream::connect(address).await.unwrap(); + valid + .write_all( + b"GET /auth/callback?code=real-code&state=expected-state HTTP/1.1\r\nHost: localhost\r\n\r\n", + ) + .await + .unwrap(); + let mut valid_response = Vec::new(); + valid.read_to_end(&mut valid_response).await.unwrap(); + assert!(String::from_utf8_lossy(&valid_response).contains("200 OK")); + + let params = waiter.await.unwrap().unwrap(); + assert_eq!(params.get("code").map(String::as_str), Some("real-code")); + } + + #[tokio::test] + async fn fragmented_callback_request_is_read_until_headers_complete() { + let (mut client, mut server) = tokio::io::duplex(4096); + let writer = tokio::spawn(async move { + client + .write_all(b"GET /auth/callback?code=fragmented") + .await + .unwrap(); + tokio::task::yield_now().await; + client + .write_all(b"-code&state=expected-state HTTP/1.1\r\nHost: local") + .await + .unwrap(); + tokio::task::yield_now().await; + client.write_all(b"host\r\n\r\n").await.unwrap(); + }); + + let request = read_http_request(&mut server, Duration::from_secs(1)) + .await + .unwrap() + .unwrap(); + writer.await.unwrap(); + assert_eq!( + String::from_utf8(request).unwrap(), + "GET /auth/callback?code=fragmented-code&state=expected-state HTTP/1.1\r\nHost: localhost\r\n\r\n" + ); + } + + #[tokio::test] + async fn stalled_callback_request_hits_read_timeout() { + let (_client, mut server) = tokio::io::duplex(64); + + let error = read_http_request(&mut server, Duration::from_millis(25)) + .await + .unwrap_err(); + assert!(error.to_string().contains("timed out")); + } + + #[test] + fn callback_page_uses_browser_language_and_color_scheme() { + assert_eq!( + preferred_locale("GET / HTTP/1.1\r\nAccept-Language: zh-CN,zh;q=0.9\r\n"), + "zh-CN" + ); + assert_eq!( + preferred_locale("GET / HTTP/1.1\r\nAccept-Language: zh-TW,zh;q=0.9\r\n"), + "zh-TW" + ); + assert_eq!( + preferred_locale("GET / HTTP/1.1\r\nAccept-Language: en-US,en;q=0.9\r\n"), + "en" + ); + assert_ne!( + callback_messages("zh-CN").success_title, + callback_messages("en").success_title + ); + let chinese = success_page("zh-CN"); + assert!(chinese.contains("lang=\"zh-CN\"")); + assert!(chinese.contains(&callback_messages("zh-CN").success_title)); + assert!(chinese.contains("prefers-color-scheme:dark")); + } + #[tokio::test] async fn falls_back_when_preferred_callback_port_is_occupied() { let occupied = tokio::net::TcpListener::bind((LOOPBACK_BIND_HOST, 0)) diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/opencode.rs b/src/crates/adapters/ai-adapters/src/subscription_auth/opencode.rs index 2370e792de..b690d13ca7 100644 --- a/src/crates/adapters/ai-adapters/src/subscription_auth/opencode.rs +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/opencode.rs @@ -177,15 +177,15 @@ async fn fetch_metadata(access: &str) -> serde_json::Value { serde_json::Value::Object(metadata) } -async fn persist_tokens(tokens: TokenResponse) -> Result<()> { - let _guard = super::store_lock(super::SubscriptionProvider::Opencode) - .lock() - .await; +async fn persist_tokens( + tokens: TokenResponse, + metadata: serde_json::Value, + expected_revision: u64, +) -> Result<()> { let expires = now_ms() + tokens.expires_in * 1000; - let metadata = fetch_metadata(&tokens.access_token).await; - let mut store = store::load().await.unwrap_or_default(); - store.insert( - STORE_KEY.to_string(), + let outcome = store::upsert_if_revision( + STORE_KEY, + expected_revision, StoredCredential::Oauth { refresh: tokens.refresh_token, access: tokens.access_token, @@ -193,8 +193,9 @@ async fn persist_tokens(tokens: TokenResponse) -> Result<()> { account_id: None, metadata: Some(metadata), }, - ); - store::save(&store).await?; + ) + .await?; + super::require_current_store_revision(super::SubscriptionProvider::Opencode, outcome)?; log::info!("opencode subscription tokens saved"); Ok(()) } @@ -238,7 +239,10 @@ fn absolute_verification_url(uri: &str) -> String { /// Starts the device-code login flow. The verification URL and user code are /// returned immediately; the runner polls in the background. -pub(crate) async fn begin_login(cancel: CancellationToken) -> Result { +pub(crate) async fn begin_login( + cancel: CancellationToken, + expected_revision: u64, +) -> Result { let device = request_device_code().await?; let interval = device.interval.unwrap_or(5).max(1); let device_code = device.device_code.clone(); @@ -246,14 +250,23 @@ pub(crate) async fn begin_login(cancel: CancellationToken) -> Result Err(anyhow!("login cancelled")), - result = async { + super::authorize_then_persist( + super::SubscriptionProvider::Opencode, + cancel, + async { let mut wait = interval; loop { tokio::time::sleep(Duration::from_secs(wait)).await; match poll_once(&device_code).await? { - DevicePoll::Authorized(tokens) => return persist_tokens(tokens).await, + DevicePoll::Authorized(tokens) => { + // Optional profile/org network calls belong to the + // cancellable authorization phase. The provider + // commit lock should cover only the credential + // store transaction, never up to 60 seconds of + // metadata fetching. + let metadata = fetch_metadata(&tokens.access_token).await; + return Ok((tokens, metadata)); + } DevicePoll::Pending => { wait = interval; } @@ -264,8 +277,10 @@ pub(crate) async fn begin_login(cancel: CancellationToken) -> Result result, - } + }, + move |(tokens, metadata)| persist_tokens(tokens, metadata, expected_revision), + ) + .await }; Ok(StartedLogin { @@ -278,13 +293,9 @@ pub(crate) async fn begin_login(cancel: CancellationToken) -> Result Result { - let _guard = super::store_lock(super::SubscriptionProvider::Opencode) - .lock() - .await; - let mut store = store::load().await.unwrap_or_default(); - let entry = store - .get(STORE_KEY) - .cloned() + let snapshot = store::load_entry_with_revision(STORE_KEY).await?; + let entry = snapshot + .credential .ok_or_else(|| anyhow!("OpenCode Zen is not connected; sign in first"))?; match entry { StoredCredential::Api { key, .. } => Ok(key), @@ -300,8 +311,9 @@ async fn ensure_fresh() -> Result { } let refreshed = refresh(&refresh_token).await?; let new_expires = now_ms() + refreshed.expires_in * 1000; - store.insert( - STORE_KEY.to_string(), + let outcome = store::upsert_if_revision( + STORE_KEY, + snapshot.revision, StoredCredential::Oauth { refresh: refreshed.refresh_token, access: refreshed.access_token.clone(), @@ -309,8 +321,9 @@ async fn ensure_fresh() -> Result { account_id, metadata, }, - ); - store::save(&store).await?; + ) + .await?; + super::require_current_store_revision(super::SubscriptionProvider::Opencode, outcome)?; log::info!("opencode subscription tokens refreshed"); Ok(refreshed.access_token) } diff --git a/src/crates/adapters/ai-adapters/src/subscription_auth/store.rs b/src/crates/adapters/ai-adapters/src/subscription_auth/store.rs index 089903086e..4a5774f40c 100644 --- a/src/crates/adapters/ai-adapters/src/subscription_auth/store.rs +++ b/src/crates/adapters/ai-adapters/src/subscription_auth/store.rs @@ -1,22 +1,33 @@ -//! On-disk persistence for subscription auth credentials. +//! Persistence for subscription-account credentials. //! -//! Tokens live in a single JSON document keyed by provider id -//! (`codex` / `antigravity` / `opencode`). The file is written with mode -//! `0600` on Unix so other local users cannot read the stored tokens. +//! OAuth tokens and API keys are stored in the operating-system credential +//! vault (macOS Keychain, Windows Credential Manager, or Linux Secret +//! Service). The JSON file contains only non-secret account metadata and +//! references used to discover the corresponding vault entries. //! //! Path: `{dirs::config_dir()}/bitfun/data/subscription_auth.json`. use anyhow::{anyhow, Context, Result}; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::path::PathBuf; -use std::sync::{OnceLock, RwLock}; - -/// A single stored credential for one provider. -/// -/// `oauth` credentials keep the refresh/access token pair plus the millisecond -/// epoch expiry; `api` credentials keep a static API key. Both may carry an -/// opaque `metadata` object (email, org info, etc.). +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::fmt; +use std::fs::{File, OpenOptions}; +use std::path::{Path, PathBuf}; +use std::sync::{Mutex, OnceLock, RwLock}; +use std::time::Duration; + +const STORE_VERSION: u8 = 2; +const CLEANUP_JOURNAL_VERSION: u8 = 1; +const KEYRING_SERVICE: &str = "openbitfun.bitfun.subscription-auth.v1"; +const STORE_FILE_LOCK_TIMEOUT: Duration = Duration::from_secs(10); +const STORE_FILE_LOCK_RETRY_INTERVAL: Duration = Duration::from_millis(25); +// Windows Credential Manager limits a generic credential blob to 2560 bytes. +// Leave headroom for platform-store implementations and split every logical +// secret so a long JWT or refresh token remains portable across all hosts. +const SECRET_CHUNK_BYTES: usize = 2_048; + +/// A single credential assembled in memory after its secret material has been +/// read from the platform credential vault. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub enum StoredCredential { @@ -37,27 +48,739 @@ pub enum StoredCredential { }, } -/// Provider id -> stored credential. +/// Provider id -> in-memory credential. pub type Store = HashMap; +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum CredentialMetadata { + Oauth { + expires: i64, + #[serde(default, skip_serializing_if = "Option::is_none")] + account_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + metadata: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + needs_reauthentication: bool, + /// Unique namespace for this committed set of vault chunks. `None` + /// denotes the legacy single-password entry keyed by provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + secret_set_id: Option, + #[serde(default, skip_serializing_if = "is_zero")] + refresh_parts: u32, + #[serde(default, skip_serializing_if = "is_zero")] + access_parts: u32, + }, + Api { + #[serde(default, skip_serializing_if = "Option::is_none")] + metadata: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + needs_reauthentication: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + secret_set_id: Option, + #[serde(default, skip_serializing_if = "is_zero")] + key_parts: u32, + }, +} + +fn is_zero(value: &u32) -> bool { + *value == 0 +} + +fn secret_part_count(secret: &str) -> u32 { + // Even an empty value has one explicit part, distinguishing a committed + // empty refresh token from a missing vault entry. + secret.len().max(1).div_ceil(SECRET_CHUNK_BYTES) as u32 +} + +impl CredentialMetadata { + fn from_credential(credential: &StoredCredential) -> Self { + match credential { + StoredCredential::Oauth { + refresh, + access, + expires, + account_id, + metadata, + .. + } => Self::Oauth { + expires: *expires, + account_id: account_id.clone(), + metadata: metadata.clone(), + needs_reauthentication: false, + secret_set_id: Some(uuid::Uuid::new_v4().simple().to_string()), + refresh_parts: secret_part_count(refresh), + access_parts: secret_part_count(access), + }, + StoredCredential::Api { key, metadata } => Self::Api { + metadata: metadata.clone(), + needs_reauthentication: false, + secret_set_id: Some(uuid::Uuid::new_v4().simple().to_string()), + key_parts: secret_part_count(key), + }, + } + } + + fn needs_reauthentication(&self) -> bool { + match self { + Self::Oauth { + needs_reauthentication, + .. + } + | Self::Api { + needs_reauthentication, + .. + } => *needs_reauthentication, + } + } + + fn combine(&self, secret: SecretMaterial) -> Option { + match (self, secret) { + ( + Self::Oauth { + expires, + account_id, + metadata, + .. + }, + SecretMaterial::Oauth { refresh, access }, + ) => Some(StoredCredential::Oauth { + refresh, + access, + expires: *expires, + account_id: account_id.clone(), + metadata: metadata.clone(), + }), + (Self::Api { metadata, .. }, SecretMaterial::Api { key }) => { + Some(StoredCredential::Api { + key, + metadata: metadata.clone(), + }) + } + _ => None, + } + } + + fn vault_entries(&self, provider: &str) -> Vec { + if self.needs_reauthentication() { + return Vec::new(); + } + match self { + Self::Oauth { + secret_set_id: Some(set_id), + refresh_parts, + access_parts, + .. + } => secret_entry_names(provider, set_id, "refresh", *refresh_parts) + .chain(secret_entry_names( + provider, + set_id, + "access", + *access_parts, + )) + .collect(), + Self::Api { + secret_set_id: Some(set_id), + key_parts, + .. + } => secret_entry_names(provider, set_id, "api-key", *key_parts).collect(), + // The old secure-vault representation used one password entry + // named exactly after the provider. + _ => vec![provider.to_string()], + } + } +} + +fn secret_entry_names<'a>( + provider: &'a str, + set_id: &'a str, + field: &'a str, + parts: u32, +) -> impl Iterator + 'a { + (0..parts).map(move |index| format!("{provider}/v2/{set_id}/{field}/{index}")) +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum SecretMaterial { + Oauth { refresh: String, access: String }, + Api { key: String }, +} + +impl From<&StoredCredential> for SecretMaterial { + fn from(value: &StoredCredential) -> Self { + match value { + StoredCredential::Oauth { + refresh, access, .. + } => Self::Oauth { + refresh: refresh.clone(), + access: access.clone(), + }, + StoredCredential::Api { key, .. } => Self::Api { key: key.clone() }, + } + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +struct SecureStoreFile { + version: u8, + #[serde(default)] + accounts: HashMap, + /// Monotonic provider epochs, retained even after an account is removed. + /// + /// The retained entry is a tombstone: an authorization or token refresh + /// that began in another BitFun process before logout must not be able to + /// publish its stale credential after logout has completed. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + provider_revisions: HashMap, +} + +impl SecureStoreFile { + fn active_vault_entries(&self) -> HashSet { + self.accounts + .iter() + .flat_map(|(provider, metadata)| metadata.vault_entries(provider)) + .collect() + } + + fn provider_revision(&self, provider: &str) -> u64 { + self.provider_revisions.get(provider).copied().unwrap_or(0) + } + + fn advance_provider_revision(&mut self, provider: &str) -> Result { + let next = self + .provider_revision(provider) + .checked_add(1) + .ok_or_else(|| anyhow!("subscription credential revision overflow for {provider}"))?; + self.provider_revisions.insert(provider.to_string(), next); + Ok(next) + } +} + +#[derive(Debug, Default, Serialize, Deserialize)] +struct CleanupJournal { + version: u8, + #[serde(default)] + entries: BTreeSet, +} + +#[derive(Debug)] +struct VaultUnavailableError(String); + +impl fmt::Display for VaultUnavailableError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.0) + } +} + +impl std::error::Error for VaultUnavailableError {} + +fn vault_unavailable(message: impl Into) -> anyhow::Error { + VaultUnavailableError(message.into()).into() +} + +fn is_vault_unavailable(error: &anyhow::Error) -> bool { + error + .chain() + .any(|cause| cause.downcast_ref::().is_some()) +} + +/// Result used by account discovery so a missing/locked vault entry is visible +/// to the UI instead of silently looking like a never-configured account. +pub(crate) struct LoadState { + pub credentials: Store, + pub requires_reauthentication: HashSet, + /// Metadata and secret entries exist, but the OS vault is currently + /// locked/unavailable. This is retryable and must not be shown as lost. + pub vault_unavailable: HashSet, + pub provider_revisions: HashMap, +} + +/// One credential and the durable provider epoch observed in the same store +/// transaction. Callers that perform network refresh work must conditionally +/// commit against this revision instead of blindly resurrecting stale state. +pub(crate) struct VersionedCredential { + pub credential: Option, + pub revision: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ConditionalCommitOutcome { + Committed { revision: u64 }, + Conflict { current_revision: u64 }, +} + +/// Result of removing an account from the committed metadata index. Native +/// vault cleanup can remain pending without making the account look connected. +#[derive(Debug)] +pub(crate) enum RemoveOutcome { + Removed, + CleanupPending(String), +} + fn store_path_override() -> &'static RwLock> { static OVERRIDE: OnceLock>> = OnceLock::new(); OVERRIDE.get_or_init(|| RwLock::new(None)) } -/// Overrides the store path for tests. Pass a temp-dir file path. +/// Test-only secret material, keyed by the overridden metadata path. Tests +/// must never read from or write to a developer's real system credential vault. +fn test_secrets() -> &'static Mutex>>> { + static SECRETS: OnceLock>>>> = OnceLock::new(); + SECRETS.get_or_init(|| Mutex::new(HashMap::new())) +} + +#[cfg(test)] +fn unavailable_test_vaults() -> &'static Mutex> { + static PATHS: OnceLock>> = OnceLock::new(); + PATHS.get_or_init(|| Mutex::new(HashSet::new())) +} + +#[cfg(test)] +fn failing_metadata_writes() -> &'static Mutex> { + static PATHS: OnceLock>> = OnceLock::new(); + PATHS.get_or_init(|| Mutex::new(HashSet::new())) +} + +#[cfg(test)] +fn failing_vault_writes_after() -> &'static Mutex> { + static WRITES: OnceLock>> = OnceLock::new(); + WRITES.get_or_init(|| Mutex::new(HashMap::new())) +} + +#[cfg(test)] +fn failing_vault_deletes() -> &'static Mutex> { + static PATHS: OnceLock>> = OnceLock::new(); + PATHS.get_or_init(|| Mutex::new(HashSet::new())) +} + +#[cfg(test)] +fn failing_backup_cleanup() -> &'static Mutex> { + static PATHS: OnceLock>> = OnceLock::new(); + PATHS.get_or_init(|| Mutex::new(HashSet::new())) +} + +fn native_keyring_lock() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) +} + +fn store_operation_lock() -> &'static tokio::sync::Mutex<()> { + static LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + &LOCK +} + +/// Cross-process owner for the metadata, cleanup journal, and credential-vault +/// transaction. The lock file is stable and is never removed, because replacing +/// it could let two processes lock different inodes for the same store path. +struct StoreFileLock { + file: File, + path: PathBuf, +} + +impl Drop for StoreFileLock { + fn drop(&mut self) { + if let Err(error) = fs2::FileExt::unlock(&self.file) { + log::warn!( + "release subscription credential transaction lock {} failed: {error}", + self.path.display() + ); + } + } +} + +/// Field order releases the OS lock before the in-process mutex, preserving the +/// global acquisition order for the next local transaction. +struct StoreTransactionGuard { + _file_lock: StoreFileLock, + _process_lock: tokio::sync::MutexGuard<'static, ()>, +} + +fn store_lock_path(metadata_path: &Path) -> PathBuf { + metadata_path.with_extension("lock") +} + +#[cfg(unix)] +fn configure_store_lock_open(options: &mut OpenOptions) { + use std::os::unix::fs::OpenOptionsExt; + + options.mode(0o600).custom_flags(libc::O_NOFOLLOW); +} + +#[cfg(windows)] +fn configure_store_lock_open(options: &mut OpenOptions) { + use std::os::windows::fs::OpenOptionsExt; + + // Open the reparse point itself so the regular-file check below rejects it + // instead of following it to a different transaction-lock file. + const FILE_FLAG_OPEN_REPARSE_POINT: u32 = 0x0020_0000; + options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT); +} + +#[cfg(not(any(unix, windows)))] +fn configure_store_lock_open(_options: &mut OpenOptions) {} + +#[cfg(windows)] +fn store_lock_is_reparse_point(metadata: &std::fs::Metadata) -> bool { + use std::os::windows::fs::MetadataExt; + + const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x400; + metadata.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0 +} + +#[cfg(not(windows))] +fn store_lock_is_reparse_point(_metadata: &std::fs::Metadata) -> bool { + false +} + +fn open_store_lock_file(path: &Path) -> Result { + match std::fs::symlink_metadata(path) { + Ok(metadata) + if !metadata.is_file() + || metadata.file_type().is_symlink() + || store_lock_is_reparse_point(&metadata) => + { + return Err(anyhow!( + "subscription credential transaction lock {} must be a regular file, not a directory, symlink, or reparse point", + path.display() + )); + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(error).with_context(|| { + format!( + "inspect subscription credential transaction lock path {}", + path.display() + ) + }); + } + } + + let mut options = OpenOptions::new(); + options.create(true).truncate(false).read(true).write(true); + configure_store_lock_open(&mut options); + let file = options.open(path).with_context(|| { + format!( + "open subscription credential transaction lock {}", + path.display() + ) + })?; + let metadata = file.metadata().with_context(|| { + format!( + "inspect subscription credential transaction lock {}", + path.display() + ) + })?; + if !metadata.is_file() || store_lock_is_reparse_point(&metadata) { + return Err(anyhow!( + "subscription credential transaction lock {} must be a regular file, not a directory, symlink, or reparse point", + path.display() + )); + } + + let path_metadata = std::fs::symlink_metadata(path).with_context(|| { + format!( + "inspect subscription credential transaction lock path {}", + path.display() + ) + })?; + if !path_metadata.is_file() + || path_metadata.file_type().is_symlink() + || store_lock_is_reparse_point(&path_metadata) + { + return Err(anyhow!( + "subscription credential transaction lock {} must be an ordinary file", + path.display() + )); + } + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + + file.set_permissions(std::fs::Permissions::from_mode(0o600)) + .with_context(|| { + format!( + "restrict subscription credential transaction lock permissions {}", + path.display() + ) + })?; + } + Ok(file) +} + +async fn acquire_store_file_lock_with_timeout( + metadata_path: &Path, + timeout: Duration, +) -> Result { + let lock_path = store_lock_path(metadata_path); + let parent = lock_path.parent().ok_or_else(|| { + anyhow!( + "subscription credential transaction lock has no parent directory: {}", + lock_path.display() + ) + })?; + let deadline = tokio::time::Instant::now() + timeout; + tokio::time::timeout(timeout, tokio::fs::create_dir_all(parent)) + .await + .map_err(|_| { + anyhow!( + "timed out creating subscription credential transaction lock directory {}", + parent.display() + ) + })? + .with_context(|| { + format!( + "create subscription credential transaction lock directory {}", + parent.display() + ) + })?; + + let open_path = lock_path.clone(); + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err(anyhow!( + "timed out waiting for subscription credential transaction lock {}", + lock_path.display() + )); + } + let file = tokio::time::timeout( + remaining, + tokio::task::spawn_blocking(move || open_store_lock_file(&open_path)), + ) + .await + .map_err(|_| { + anyhow!( + "timed out opening subscription credential transaction lock {}", + lock_path.display() + ) + })? + .context("join subscription credential transaction lock open task")??; + + loop { + match fs2::FileExt::try_lock_exclusive(&file) { + Ok(()) => { + if tokio::time::Instant::now() >= deadline { + let _ = fs2::FileExt::unlock(&file); + return Err(anyhow!( + "timed out waiting for subscription credential transaction lock {}; another BitFun process may be updating credentials", + lock_path.display() + )); + } + return Ok(StoreFileLock { + file, + path: lock_path, + }); + } + Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err(anyhow!( + "timed out waiting for subscription credential transaction lock {}; another BitFun process may be updating credentials", + lock_path.display() + )); + } + tokio::time::sleep(remaining.min(STORE_FILE_LOCK_RETRY_INTERVAL)).await; + } + Err(error) => { + return Err(error).with_context(|| { + format!( + "acquire subscription credential transaction lock {}", + lock_path.display() + ) + }); + } + } + } +} + +async fn acquire_store_transaction() -> Result<(PathBuf, StoreTransactionGuard)> { + // Lock order is process mutex first, OS file lock second everywhere. The + // unlocked helpers below prevent recursive acquisition during migration. + let process_lock = store_operation_lock().lock().await; + let path = store_path()?; + let file_lock = acquire_store_file_lock_with_timeout(&path, STORE_FILE_LOCK_TIMEOUT).await?; + Ok(( + path, + StoreTransactionGuard { + _file_lock: file_lock, + _process_lock: process_lock, + }, + )) +} + +/// Overrides the metadata path for tests. The override also switches secret +/// persistence to the process-local test vault above. pub fn set_store_path_for_test(path: PathBuf) { if let Ok(mut guard) = store_path_override().write() { *guard = Some(path); } } -fn store_path() -> Result { - if let Ok(guard) = store_path_override().read() { - if let Some(path) = guard.as_ref() { - return Ok(path.clone()); +#[cfg(test)] +pub(crate) fn store_path_for_test_assertion() -> PathBuf { + overridden_store_path().expect("subscription test store path is configured") +} + +#[cfg(test)] +pub(crate) fn test_vault_entries_for_assertion() -> HashMap> { + let path = store_path_for_test_assertion(); + test_secrets() + .lock() + .ok() + .and_then(|vault| vault.get(&path).cloned()) + .unwrap_or_default() +} + +#[cfg(test)] +pub(crate) fn set_test_vault_unavailable(unavailable: bool) { + let path = store_path_for_test_assertion(); + if let Ok(mut paths) = unavailable_test_vaults().lock() { + if unavailable { + paths.insert(path); + } else { + paths.remove(&path); } } +} + +#[cfg(test)] +pub(crate) fn set_test_metadata_write_failure(fail: bool) { + let path = store_path_for_test_assertion(); + if let Ok(mut paths) = failing_metadata_writes().lock() { + if fail { + paths.insert(path); + } else { + paths.remove(&path); + } + } +} + +#[cfg(test)] +pub(crate) fn set_test_vault_write_failure_after(successful_writes: Option) { + let path = store_path_for_test_assertion(); + if let Ok(mut failures) = failing_vault_writes_after().lock() { + match successful_writes { + Some(count) => { + failures.insert(path, count); + } + None => { + failures.remove(&path); + } + } + } +} + +#[cfg(test)] +pub(crate) fn set_test_vault_delete_failure(fail: bool) { + let path = store_path_for_test_assertion(); + if let Ok(mut paths) = failing_vault_deletes().lock() { + if fail { + paths.insert(path); + } else { + paths.remove(&path); + } + } +} + +#[cfg(test)] +pub(crate) fn set_test_backup_cleanup_failure(path: &Path, fail: bool) { + if let Ok(mut paths) = failing_backup_cleanup().lock() { + if fail { + paths.insert(path.to_path_buf()); + } else { + paths.remove(path); + } + } +} + +#[cfg(test)] +fn test_vault_is_unavailable(path: &Path) -> bool { + unavailable_test_vaults() + .lock() + .map(|paths| paths.contains(path)) + .unwrap_or(true) +} + +#[cfg(not(test))] +fn test_vault_is_unavailable(_path: &Path) -> bool { + false +} + +#[cfg(test)] +fn metadata_write_should_fail(path: &Path) -> bool { + failing_metadata_writes() + .lock() + .map(|paths| paths.contains(path)) + .unwrap_or(true) +} + +#[cfg(test)] +fn vault_write_should_fail(path: &Path) -> bool { + failing_vault_writes_after() + .lock() + .map(|mut failures| { + let Some(remaining) = failures.get_mut(path) else { + return false; + }; + if *remaining == 0 { + true + } else { + *remaining -= 1; + false + } + }) + .unwrap_or(true) +} + +#[cfg(not(test))] +fn vault_write_should_fail(_path: &Path) -> bool { + false +} + +#[cfg(test)] +fn vault_delete_should_fail(path: &Path) -> bool { + failing_vault_deletes() + .lock() + .map(|paths| paths.contains(path)) + .unwrap_or(true) +} + +#[cfg(not(test))] +fn vault_delete_should_fail(_path: &Path) -> bool { + false +} + +#[cfg(test)] +fn backup_cleanup_should_fail(path: &Path) -> bool { + failing_backup_cleanup() + .lock() + .map(|paths| paths.contains(path)) + .unwrap_or(true) +} + +#[cfg(all(not(test), windows))] +fn backup_cleanup_should_fail(_path: &Path) -> bool { + false +} + +#[cfg(not(test))] +fn metadata_write_should_fail(_path: &Path) -> bool { + false +} + +fn overridden_store_path() -> Option { + store_path_override() + .read() + .ok() + .and_then(|guard| guard.clone()) +} + +fn store_path() -> Result { + if let Some(path) = overridden_store_path() { + return Ok(path); + } let base = dirs::config_dir().ok_or_else(|| anyhow!("system config directory unavailable"))?; Ok(base .join("bitfun") @@ -65,39 +788,823 @@ fn store_path() -> Result { .join("subscription_auth.json")) } -/// Loads the full credential store. Returns an empty store when the file does -/// not exist yet. -pub async fn load() -> Result { - let path = store_path()?; +async fn read_bytes(path: &Path) -> Result>> { + #[cfg(windows)] + restore_windows_backup_if_needed(path).await?; if !path.exists() { - return Ok(Store::new()); + return Ok(None); } - let bytes = tokio::fs::read(&path) + let bytes = tokio::fs::read(path) .await - .with_context(|| format!("read subscription auth store at {}", path.display()))?; - if bytes.is_empty() { - return Ok(Store::new()); + .with_context(|| format!("read subscription auth metadata at {}", path.display()))?; + Ok((!bytes.is_empty()).then_some(bytes)) +} + +/// Recover the old metadata index if the process stopped after rotating the +/// destination but before moving the new temp file into place. If the new file +/// is already present, scrub the stale backup instead; it can contain legacy +/// plaintext credentials from a migration. +#[cfg(windows)] +async fn restore_windows_backup_if_needed(path: &Path) -> Result<()> { + if path.exists() { + cleanup_stale_backup(path).await; + return Ok(()); } - serde_json::from_slice(&bytes) - .with_context(|| format!("parse subscription auth store at {}", path.display())) + let backup = path.with_extension("bak"); + if !backup.exists() { + return Ok(()); + } + tokio::fs::rename(&backup, path).await.with_context(|| { + format!( + "restore interrupted subscription auth metadata {} -> {}", + backup.display(), + path.display() + ) + }) } -/// Persists the full credential store with restrictive permissions. -pub async fn save(store: &Store) -> Result<()> { - let path = store_path()?; +#[cfg(windows)] +async fn cleanup_stale_backup(path: &Path) { + let backup = path.with_extension("bak"); + if !backup.exists() { + return; + } + if let Err(error) = scrub_and_remove_file(&backup).await { + log::warn!( + "remove stale subscription auth metadata backup failed; cleanup will retry later: {error:#}" + ); + } +} + +#[cfg(any(windows, test))] +async fn scrub_and_remove_file(path: &Path) -> Result<()> { + if backup_cleanup_should_fail(path) { + return Err(anyhow!( + "injected subscription metadata backup cleanup failure" + )); + } + match tokio::fs::OpenOptions::new() + .write(true) + .truncate(true) + .open(path) + .await + { + Ok(file) => file + .sync_all() + .await + .with_context(|| format!("sync scrubbed file {}", path.display()))?, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error).with_context(|| format!("scrub file {}", path.display())), + } + tokio::fs::remove_file(path) + .await + .with_context(|| format!("remove scrubbed file {}", path.display())) +} + +fn parse_secure_file(bytes: &[u8], path: &Path) -> Result { + let file: SecureStoreFile = serde_json::from_slice(bytes) + .with_context(|| format!("parse subscription auth metadata at {}", path.display()))?; + if file.version != STORE_VERSION { + return Err(anyhow!( + "unsupported subscription auth metadata version {} at {}", + file.version, + path.display() + )); + } + Ok(file) +} + +async fn read_secure_file(path: &Path) -> Result { + let Some(bytes) = read_bytes(path).await? else { + return Ok(SecureStoreFile { + version: STORE_VERSION, + accounts: HashMap::new(), + provider_revisions: HashMap::new(), + }); + }; + parse_secure_file(&bytes, path) +} + +fn open_native_keyring_entry(entry_name: &str) -> std::result::Result { + if keyring_core::get_default_store().is_none() { + #[cfg(target_os = "macos")] + let store = apple_native_keyring_store::keychain::Store::new(); + #[cfg(target_os = "windows")] + let store = windows_native_keyring_store::Store::new(); + #[cfg(all( + unix, + not(any(target_os = "macos", target_os = "ios", target_os = "android")) + ))] + let store = zbus_secret_service_keyring_store::Store::new(); + #[cfg(not(any( + target_os = "macos", + target_os = "windows", + all( + unix, + not(any(target_os = "macos", target_os = "ios", target_os = "android")) + ) + )))] + let store: keyring_core::Result> = + Err(keyring_core::Error::NoDefaultStore); + + // Unlike the keyring v1 facade, failed initialization leaves no sticky + // once flag. A later UI retry can reconnect to Linux Secret Service. + let store = + store.map_err(|error| format!("initialize system credential store: {error}"))?; + keyring_core::set_default_store(store); + } + keyring_core::Entry::new(KEYRING_SERVICE, entry_name) + .map_err(|error| format!("open system credential entry: {error}")) +} + +async fn get_secret_bytes(entry_name: &str) -> Result>> { + if let Some(path) = overridden_store_path() { + if test_vault_is_unavailable(&path) { + return Err(vault_unavailable("subscription test vault unavailable")); + } + return test_secrets() + .lock() + .map_err(|_| anyhow!("subscription test vault lock poisoned")) + .map(|vault| { + vault + .get(&path) + .and_then(|items| items.get(entry_name)) + .cloned() + }); + } + + let entry_name = entry_name.to_string(); + tokio::task::spawn_blocking(move || { + let _guard = native_keyring_lock() + .lock() + .map_err(|_| "subscription keyring lock poisoned".to_string())?; + let entry = open_native_keyring_entry(&entry_name)?; + match entry.get_secret() { + Ok(secret) => Ok(Some(secret)), + Err(keyring_core::Error::NoEntry) => Ok(None), + Err(err) => Err(format!("read system credential entry: {err}")), + } + }) + .await + .context("join system credential read task")? + .map_err(vault_unavailable) +} + +/// Reads the v1 combined JSON entry. It was written through the password API, +/// which uses a platform-specific text encoding on Windows, so it cannot be +/// safely read through `get_secret` there. +async fn get_legacy_password(provider: &str) -> Result> { + if let Some(path) = overridden_store_path() { + if test_vault_is_unavailable(&path) { + return Err(vault_unavailable("subscription test vault unavailable")); + } + return test_secrets() + .lock() + .map_err(|_| anyhow!("subscription test vault lock poisoned")) + .map(|vault| { + vault + .get(&path) + .and_then(|items| items.get(provider)) + .and_then(|bytes| String::from_utf8(bytes.clone()).ok()) + }); + } + + let provider = provider.to_string(); + tokio::task::spawn_blocking(move || { + let _guard = native_keyring_lock() + .lock() + .map_err(|_| "subscription keyring lock poisoned".to_string())?; + let entry = open_native_keyring_entry(&provider)?; + match entry.get_password() { + Ok(secret) => Ok(Some(secret)), + Err(keyring_core::Error::NoEntry) => Ok(None), + Err(err) => Err(format!("read legacy system credential entry: {err}")), + } + }) + .await + .context("join legacy system credential read task")? + .map_err(vault_unavailable) +} + +async fn set_secret_bytes(entry_name: &str, secret: Vec) -> Result<()> { + if secret.len() > SECRET_CHUNK_BYTES { + return Err(anyhow!( + "subscription credential chunk exceeds portable size limit: {} bytes", + secret.len() + )); + } + if let Some(path) = overridden_store_path() { + if test_vault_is_unavailable(&path) { + return Err(vault_unavailable("subscription test vault unavailable")); + } + if vault_write_should_fail(&path) { + return Err(vault_unavailable( + "injected subscription test vault write failure", + )); + } + let mut vault = test_secrets() + .lock() + .map_err(|_| anyhow!("subscription test vault lock poisoned"))?; + vault + .entry(path) + .or_default() + .insert(entry_name.to_string(), secret); + return Ok(()); + } + + let entry_name = entry_name.to_string(); + tokio::task::spawn_blocking(move || { + let _guard = native_keyring_lock() + .lock() + .map_err(|_| "subscription keyring lock poisoned".to_string())?; + let entry = open_native_keyring_entry(&entry_name)?; + entry + .set_secret(&secret) + .map_err(|err| format!("write system credential entry: {err}")) + }) + .await + .context("join system credential write task")? + .map_err(vault_unavailable) +} + +async fn delete_secret_entry(entry_name: &str) -> Result<()> { + if let Some(path) = overridden_store_path() { + if test_vault_is_unavailable(&path) { + return Err(vault_unavailable("subscription test vault unavailable")); + } + if vault_delete_should_fail(&path) { + return Err(vault_unavailable( + "injected subscription test vault delete failure", + )); + } + if let Ok(mut vault) = test_secrets().lock() { + if let Some(items) = vault.get_mut(&path) { + items.remove(entry_name); + } + } + return Ok(()); + } + + let entry_name = entry_name.to_string(); + tokio::task::spawn_blocking(move || { + let _guard = native_keyring_lock() + .lock() + .map_err(|_| "subscription keyring lock poisoned".to_string())?; + let entry = open_native_keyring_entry(&entry_name)?; + match entry.delete_credential() { + Ok(()) | Err(keyring_core::Error::NoEntry) => Ok(()), + Err(err) => Err(format!("delete system credential entry: {err}")), + } + }) + .await + .context("join system credential delete task")? + .map_err(vault_unavailable) +} + +fn secret_chunks(secret: &str) -> Vec> { + if secret.is_empty() { + return vec![Vec::new()]; + } + secret + .as_bytes() + .chunks(SECRET_CHUNK_BYTES) + .map(<[u8]>::to_vec) + .collect() +} + +async fn read_chunked_field( + provider: &str, + set_id: &str, + field: &str, + parts: u32, +) -> Result> { + if parts == 0 { + return Ok(None); + } + let mut bytes = Vec::new(); + for entry_name in secret_entry_names(provider, set_id, field, parts) { + let Some(mut part) = get_secret_bytes(&entry_name).await? else { + return Ok(None); + }; + bytes.append(&mut part); + } + match String::from_utf8(bytes) { + Ok(secret) => Ok(Some(secret)), + Err(error) => { + log::warn!( + "subscription credential vault chunks are invalid for provider {provider} field {field}: {error}" + ); + Ok(None) + } + } +} + +async fn read_secret_material( + provider: &str, + metadata: &CredentialMetadata, +) -> Result> { + match metadata { + CredentialMetadata::Oauth { + secret_set_id: Some(set_id), + refresh_parts, + access_parts, + .. + } => { + let Some(refresh) = + read_chunked_field(provider, set_id, "refresh", *refresh_parts).await? + else { + return Ok(None); + }; + let Some(access) = + read_chunked_field(provider, set_id, "access", *access_parts).await? + else { + return Ok(None); + }; + Ok(Some(SecretMaterial::Oauth { refresh, access })) + } + CredentialMetadata::Api { + secret_set_id: Some(set_id), + key_parts, + .. + } => Ok(read_chunked_field(provider, set_id, "api-key", *key_parts) + .await? + .map(|key| SecretMaterial::Api { key })), + // Backward-compatible read of the original combined JSON password. + _ => { + let Some(secret) = get_legacy_password(provider).await? else { + return Ok(None); + }; + match serde_json::from_str(&secret) { + Ok(material) => Ok(Some(material)), + Err(error) => { + log::warn!( + "legacy subscription credential vault entry is invalid for provider {provider}: {error}" + ); + Ok(None) + } + } + } + } +} + +async fn write_secret_material( + provider: &str, + metadata: &CredentialMetadata, + credential: &StoredCredential, +) -> Result<()> { + let (set_id, fields): (&str, Vec<(&str, &str)>) = match (metadata, credential) { + ( + CredentialMetadata::Oauth { + secret_set_id: Some(set_id), + .. + }, + StoredCredential::Oauth { + refresh, access, .. + }, + ) => (set_id, vec![("refresh", refresh), ("access", access)]), + ( + CredentialMetadata::Api { + secret_set_id: Some(set_id), + .. + }, + StoredCredential::Api { key, .. }, + ) => (set_id, vec![("api-key", key)]), + _ => return Err(anyhow!("subscription credential metadata type mismatch")), + }; + + for (field, value) in fields { + for (index, chunk) in secret_chunks(value).into_iter().enumerate() { + let entry_name = format!("{provider}/v2/{set_id}/{field}/{index}"); + set_secret_bytes(&entry_name, chunk).await?; + } + } + Ok(()) +} + +fn cleanup_journal_path(metadata_path: &Path) -> PathBuf { + metadata_path.with_extension("cleanup.json") +} + +async fn read_cleanup_journal(metadata_path: &Path) -> Result { + let path = cleanup_journal_path(metadata_path); + let Some(bytes) = read_bytes(&path).await? else { + return Ok(CleanupJournal { + version: CLEANUP_JOURNAL_VERSION, + entries: BTreeSet::new(), + }); + }; + let journal: CleanupJournal = serde_json::from_slice(&bytes).with_context(|| { + format!( + "parse subscription credential cleanup journal {}", + path.display() + ) + })?; + if journal.version != CLEANUP_JOURNAL_VERSION { + return Err(anyhow!( + "unsupported subscription credential cleanup journal version {} at {}", + journal.version, + path.display() + )); + } + Ok(journal) +} + +async fn write_cleanup_journal(metadata_path: &Path, journal: &CleanupJournal) -> Result<()> { + let path = cleanup_journal_path(metadata_path); + if let Some(parent) = path.parent() { + tokio::fs::create_dir_all(parent).await.with_context(|| { + format!( + "create subscription credential cleanup journal directory {}", + parent.display() + ) + })?; + } + let bytes = serde_json::to_vec_pretty(journal)?; + write_atomic(&path, &bytes).await +} + +/// Durably records vault entries before an operation can make them unreachable. +/// Reconciliation always compares the journal with the committed metadata, so +/// a crash before metadata commit cannot delete the still-active credential. +async fn schedule_cleanup(metadata_path: &Path, entries: I) -> Result<()> +where + I: IntoIterator, +{ + let mut journal = read_cleanup_journal(metadata_path).await?; + let before = journal.entries.len(); + journal.entries.extend(entries); + if journal.entries.len() != before { + write_cleanup_journal(metadata_path, &journal).await?; + } + Ok(()) +} + +/// Deletes only entries that are not referenced by the committed metadata. +/// Failed deletions remain durable in the journal for a later startup/retry. +async fn reconcile_cleanup_journal( + metadata_path: &Path, + active_entries: &HashSet, +) -> Result<()> { + let mut journal = read_cleanup_journal(metadata_path).await?; + if journal.entries.is_empty() { + return Ok(()); + } + + let mut first_error = None; + let pending: Vec = journal.entries.iter().cloned().collect(); + for entry_name in pending { + if active_entries.contains(&entry_name) { + journal.entries.remove(&entry_name); + continue; + } + match delete_secret_entry(&entry_name).await { + Ok(()) => { + journal.entries.remove(&entry_name); + } + Err(error) => { + if first_error.is_none() { + first_error = Some(error); + } + } + } + } + write_cleanup_journal(metadata_path, &journal).await?; + if let Some(error) = first_error { + Err(error.context("subscription credential cleanup remains pending")) + } else { + Ok(()) + } +} + +#[cfg(test)] +pub(crate) async fn cleanup_journal_entries_for_assertion() -> BTreeSet { + let path = store_path_for_test_assertion(); + read_cleanup_journal(&path) + .await + .map(|journal| journal.entries) + .unwrap_or_default() +} + +async fn write_secure_file(path: &Path, file: &SecureStoreFile) -> Result<()> { + if metadata_write_should_fail(path) { + return Err(anyhow!("injected subscription metadata write failure")); + } if let Some(parent) = path.parent() { tokio::fs::create_dir_all(parent) .await - .with_context(|| format!("create subscription auth store dir {}", parent.display()))?; + .with_context(|| format!("create subscription auth directory {}", parent.display()))?; } - let bytes = serde_json::to_vec_pretty(store)?; - write_atomic(&path, &bytes).await + let bytes = serde_json::to_vec_pretty(file)?; + write_atomic(path, &bytes).await +} + +async fn migrate_legacy_store(path: &Path, legacy: Store) -> Result { + let mut secure = SecureStoreFile { + version: STORE_VERSION, + accounts: HashMap::new(), + provider_revisions: HashMap::new(), + }; + + for (provider, credential) in &legacy { + secure.accounts.insert( + provider.clone(), + CredentialMetadata::from_credential(credential), + ); + } + let new_entries = secure.active_vault_entries(); + // Commit cleanup intent before the first vault write so a process stop can + // never strand partially-written migration chunks. + schedule_cleanup(path, new_entries.iter().cloned()).await?; + + for (provider, credential) in &legacy { + let metadata = secure + .accounts + .get(provider) + .expect("legacy migration metadata exists"); + if let Err(error) = write_secret_material(provider, metadata, credential).await { + if let Err(cleanup_error) = reconcile_cleanup_journal(path, &HashSet::new()).await { + log::warn!( + "cleanup after interrupted subscription credential migration remains pending: {cleanup_error:#}" + ); + } + + // A locked or temporarily unavailable native vault must not turn + // a retryable migration into permanent credential loss. Leave the + // legacy file untouched and retry after the vault is available. + if is_vault_unavailable(&error) { + log::warn!( + "subscription credential vault migration deferred because the vault is unavailable: {error:#}" + ); + return Ok(LoadState { + credentials: Store::new(), + requires_reauthentication: HashSet::new(), + vault_unavailable: legacy.keys().cloned().collect(), + provider_revisions: HashMap::new(), + }); + } + return Err(error.context("migrate subscription credential to the system vault")); + } + } + + if let Err(error) = write_secure_file(path, &secure).await { + if let Err(cleanup_error) = reconcile_cleanup_journal(path, &HashSet::new()).await { + log::warn!( + "cleanup after failed subscription credential metadata migration remains pending: {cleanup_error:#}" + ); + } + return Err(error); + } + if let Err(error) = reconcile_cleanup_journal(path, &new_entries).await { + // The entries are active after this commit. A later load can safely + // remove their journal markers without deleting the credential. + log::warn!( + "finalize subscription credential migration cleanup journal failed; retrying later: {error:#}" + ); + } + log::info!("subscription credentials migrated to the system credential vault"); + Ok(LoadState { + credentials: legacy, + requires_reauthentication: HashSet::new(), + vault_unavailable: HashSet::new(), + provider_revisions: HashMap::new(), + }) +} + +/// Loads credentials plus vault availability state. Legacy plaintext files are +/// migrated in place and immediately rewritten without secret fields. +async fn load_with_state_unlocked(path: &Path) -> Result { + let Some(bytes) = read_bytes(path).await? else { + if let Err(error) = reconcile_cleanup_journal(path, &HashSet::new()).await { + log::warn!( + "subscription credential cleanup remains pending while no accounts are configured: {error:#}" + ); + } + return Ok(LoadState { + credentials: Store::new(), + requires_reauthentication: HashSet::new(), + vault_unavailable: HashSet::new(), + provider_revisions: HashMap::new(), + }); + }; + + let secure = match parse_secure_file(&bytes, path) { + Ok(file) => file, + Err(secure_error) => match serde_json::from_slice::(&bytes) { + Ok(legacy) => return migrate_legacy_store(path, legacy).await, + Err(_) => return Err(secure_error), + }, + }; + + let active_entries = secure.active_vault_entries(); + let provider_revisions = secure.provider_revisions.clone(); + if let Err(error) = reconcile_cleanup_journal(path, &active_entries).await { + log::warn!("subscription credential cleanup remains pending: {error:#}"); + } + + let mut credentials = Store::new(); + let mut requires_reauthentication = HashSet::new(); + let mut vault_unavailable = HashSet::new(); + for (provider, metadata) in secure.accounts { + if metadata.needs_reauthentication() { + requires_reauthentication.insert(provider); + continue; + } + let material = match read_secret_material(&provider, &metadata).await { + Ok(Some(material)) => material, + Ok(None) => { + requires_reauthentication.insert(provider); + continue; + } + Err(error) => { + log::warn!( + "subscription credential vault is unavailable for provider {provider}: {error:#}" + ); + vault_unavailable.insert(provider); + continue; + } + }; + if let Some(credential) = metadata.combine(material) { + credentials.insert(provider, credential); + } else { + requires_reauthentication.insert(provider); + } + } + Ok(LoadState { + credentials, + requires_reauthentication, + vault_unavailable, + provider_revisions, + }) +} + +/// Serializes discovery with migrations and metadata mutations so callers +/// never observe or race a partially rewritten credential index. +pub(crate) async fn load_with_state() -> Result { + let (path, _transaction) = acquire_store_transaction().await?; + load_with_state_unlocked(&path).await +} + +/// Loads all credentials that are currently available from the system vault. +pub async fn load() -> Result { + Ok(load_with_state().await?.credentials) +} + +/// Loads one provider credential without exposing its secret in the metadata +/// file. `None` means the provider needs a new sign-in. +pub async fn load_entry(provider: &str) -> Result> { + Ok(load_entry_with_revision(provider).await?.credential) +} + +/// Loads a credential and its provider tombstone revision from one locked +/// metadata/vault snapshot. +pub(crate) async fn load_entry_with_revision(provider: &str) -> Result { + let mut state = load_with_state().await?; + if state.vault_unavailable.contains(provider) { + return Err(anyhow!( + "system credential vault is locked or unavailable; unlock it and retry" + )); + } + Ok(VersionedCredential { + credential: state.credentials.remove(provider), + revision: state.provider_revisions.get(provider).copied().unwrap_or(0), + }) +} + +/// Captures the provider epoch before a long-running authorization begins. +/// Logout retains and advances this value even when no credential is present. +pub(crate) async fn credential_revision(provider: &str) -> Result { + let state = load_with_state().await?; + Ok(state.provider_revisions.get(provider).copied().unwrap_or(0)) +} + +/// Inserts or replaces a provider credential. Secret material is committed to +/// the platform vault before the non-secret metadata advertises the entry. +pub async fn upsert(provider: &str, credential: StoredCredential) -> Result<()> { + match upsert_internal(provider, credential, None).await? { + ConditionalCommitOutcome::Committed { .. } => Ok(()), + ConditionalCommitOutcome::Conflict { .. } => { + unreachable!("unconditional subscription credential commit cannot conflict") + } + } +} + +/// Commits only if no other process changed or logged out this provider since +/// `expected_revision` was captured. The comparison, vault write, metadata +/// revision advance, and metadata commit all run under the store OS lock. +pub(crate) async fn upsert_if_revision( + provider: &str, + expected_revision: u64, + credential: StoredCredential, +) -> Result { + upsert_internal(provider, credential, Some(expected_revision)).await +} + +async fn upsert_internal( + provider: &str, + credential: StoredCredential, + expected_revision: Option, +) -> Result { + let (path, _transaction) = acquire_store_transaction().await?; + // Trigger one-time migration before modifying an older file. + let _ = load_with_state_unlocked(&path).await?; + let mut file = read_secure_file(&path).await?; + let current_revision = file.provider_revision(provider); + if expected_revision.is_some_and(|expected| expected != current_revision) { + return Ok(ConditionalCommitOutcome::Conflict { current_revision }); + } + let next_revision = file.advance_provider_revision(provider)?; + let active_before = file.active_vault_entries(); + let previous = file.accounts.get(provider).cloned(); + let metadata = CredentialMetadata::from_credential(&credential); + let mut cleanup_entries = metadata.vault_entries(provider); + let mut previous_entries = previous + .as_ref() + .map(|value| value.vault_entries(provider)) + .unwrap_or_default(); + if previous_entries.is_empty() { + previous_entries.push(provider.to_string()); + } + cleanup_entries.extend(previous_entries); + schedule_cleanup(&path, cleanup_entries).await?; + + if let Err(error) = write_secret_material(provider, &metadata, &credential).await { + if let Err(cleanup_error) = reconcile_cleanup_journal(&path, &active_before).await { + log::warn!( + "cleanup after failed subscription credential write remains pending for provider {provider}: {cleanup_error:#}" + ); + } + return Err(error); + } + file.accounts.insert(provider.to_string(), metadata.clone()); + if let Err(error) = write_secure_file(&path, &file).await { + if let Err(cleanup_error) = reconcile_cleanup_journal(&path, &active_before).await { + log::warn!( + "cleanup after failed subscription metadata commit remains pending for provider {provider}: {cleanup_error:#}" + ); + } + return Err(error); + } + if let Err(error) = reconcile_cleanup_journal(&path, &file.active_vault_entries()).await { + log::warn!( + "remove superseded subscription credential chunks remains pending for provider {provider}: {error:#}" + ); + } + Ok(ConditionalCommitOutcome::Committed { + revision: next_revision, + }) +} + +/// Removes one provider from both the native vault and metadata index. +pub(crate) async fn remove(provider: &str) -> Result { + let (path, _transaction) = acquire_store_transaction().await?; + let _ = load_with_state_unlocked(&path).await?; + let mut file = read_secure_file(&path).await?; + // Advance even when the account is already absent. The persisted entry is + // the logout tombstone that rejects authorizations captured beforehand in + // another process. + file.advance_provider_revision(provider)?; + let active_before = file.active_vault_entries(); + let previous = file.accounts.get(provider).cloned(); + let mut cleanup_entries = previous + .as_ref() + .map(|metadata| metadata.vault_entries(provider)) + .unwrap_or_default(); + if cleanup_entries.is_empty() { + cleanup_entries.push(provider.to_string()); + } + // Persist cleanup intent before metadata can stop referencing the vault + // entries. Reconciliation protects them if the metadata commit fails. + schedule_cleanup(&path, cleanup_entries).await?; + file.accounts.remove(provider); + if let Err(error) = write_secure_file(&path, &file).await { + if let Err(cleanup_error) = reconcile_cleanup_journal(&path, &active_before).await { + log::warn!( + "cleanup journal reconciliation after failed sign-out remains pending for provider {provider}: {cleanup_error:#}" + ); + } + return Err(error); + } + if let Err(error) = reconcile_cleanup_journal(&path, &file.active_vault_entries()).await { + return Ok(RemoveOutcome::CleanupPending(format!( + "native credential cleanup is pending; unlock the credential vault and retry: {error:#}" + ))); + } + Ok(RemoveOutcome::Removed) } -/// Writes `bytes` to `path` atomically (temp file + rename) so a crash -/// mid-write cannot corrupt the token store. On Unix the temp file is created -/// with mode `0600` up front, so tokens are never briefly world-readable. -async fn write_atomic(path: &std::path::Path, bytes: &[u8]) -> Result<()> { +/// Persists all supplied credentials. Kept for compatibility with focused +/// tests; production refresh/login paths should call [`upsert`] for one +/// provider so concurrent providers cannot overwrite each other's tokens. +pub async fn save(store: &Store) -> Result<()> { + for (provider, credential) in store { + upsert(provider, credential.clone()).await?; + } + Ok(()) +} + +/// Writes `bytes` atomically (temp file + rename). Although the v2 file is +/// non-secret, restrictive Unix permissions protect account metadata too. +async fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> { use tokio::io::AsyncWriteExt; let tmp = path.with_extension("tmp"); @@ -126,12 +1633,314 @@ async fn write_atomic(path: &std::path::Path, bytes: &[u8]) -> Result<()> { .await .with_context(|| format!("sync subscription auth temp file {}", tmp.display()))?; } - tokio::fs::rename(&tmp, path).await.with_context(|| { + replace_metadata_file(&tmp, path).await +} + +#[cfg(not(windows))] +async fn replace_metadata_file(tmp: &Path, path: &Path) -> Result<()> { + tokio::fs::rename(tmp, path).await.with_context(|| { format!( - "rename subscription auth store {} -> {}", + "rename subscription auth metadata {} -> {}", tmp.display(), path.display() ) - })?; + }) +} + +/// Windows does not reliably replace an existing destination with `rename`. +/// Rotate the prior metadata file to a backup first and restore it if moving +/// the newly-synced temp file into place fails. +#[cfg(windows)] +async fn replace_metadata_file(tmp: &Path, path: &Path) -> Result<()> { + replace_metadata_file_windows(tmp, path).await +} + +#[cfg(any(windows, test))] +pub(crate) async fn replace_metadata_file_windows(tmp: &Path, path: &Path) -> Result<()> { + let backup = path.with_extension("bak"); + if backup.exists() { + // A stale backup can be a pre-v2 plaintext credential file. Scrub it + // before reusing the deterministic recovery path. + scrub_and_remove_file(&backup).await.with_context(|| { + format!( + "remove stale subscription auth metadata backup {}", + backup.display() + ) + })?; + } + + let had_existing = match tokio::fs::rename(path, &backup).await { + Ok(()) => true, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => false, + Err(error) => { + return Err(error).with_context(|| { + format!( + "rotate subscription auth metadata {} -> {}", + path.display(), + backup.display() + ) + }); + } + }; + + if let Err(error) = tokio::fs::rename(tmp, path).await { + if had_existing { + if let Err(restore_error) = tokio::fs::rename(&backup, path).await { + return Err(anyhow!( + "replace subscription auth metadata failed: {error}; restoring {} also failed: {restore_error}", + backup.display() + )); + } + } + return Err(error).with_context(|| { + format!( + "rename subscription auth metadata {} -> {}", + tmp.display(), + path.display() + ) + }); + } + + if had_existing { + // The new metadata is already committed. Backup cleanup must not be + // reported as a failed commit: callers would otherwise delete the new + // vault chunks. Keep it for a scrub-and-delete retry on the next read. + if let Err(error) = scrub_and_remove_file(&backup).await { + log::warn!( + "remove committed subscription auth metadata backup failed; cleanup will retry later: {error:#}" + ); + } + } Ok(()) } + +#[cfg(test)] +mod file_lock_tests { + use super::*; + + const LOCK_CHILD_METADATA_ENV: &str = "BITFUN_SUBAUTH_LOCK_CHILD_METADATA"; + const LOCK_CHILD_STARTED_ENV: &str = "BITFUN_SUBAUTH_LOCK_CHILD_STARTED"; + const LOCK_CHILD_OBSERVED_ENV: &str = "BITFUN_SUBAUTH_LOCK_CHILD_OBSERVED"; + + fn temporary_metadata_path(label: &str) -> PathBuf { + std::env::temp_dir() + .join(format!( + "bitfun-subauth-lock-{label}-{}", + uuid::Uuid::new_v4() + )) + .join("subscription_auth.json") + } + + #[tokio::test] + async fn independent_file_descriptors_contend_for_the_store_lock() { + let metadata_path = temporary_metadata_path("descriptor-contention"); + let first = acquire_store_file_lock_with_timeout(&metadata_path, Duration::from_secs(1)) + .await + .expect("first descriptor should acquire the transaction lock"); + + let error = + match acquire_store_file_lock_with_timeout(&metadata_path, Duration::from_millis(100)) + .await + { + Ok(_) => panic!("an independent descriptor must not bypass the held OS lock"), + Err(error) => error, + }; + assert!(error.to_string().contains("timed out waiting")); + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + + let mode = std::fs::metadata(store_lock_path(&metadata_path)) + .expect("lock metadata") + .permissions() + .mode() + & 0o777; + assert_eq!(mode, 0o600); + } + + drop(first); + let _second = acquire_store_file_lock_with_timeout(&metadata_path, Duration::from_secs(1)) + .await + .expect("lock should be released when its guard drops"); + } + + #[tokio::test] + async fn second_file_transaction_cannot_reconcile_before_metadata_commit() { + let metadata_path = temporary_metadata_path("commit-order"); + let first = acquire_store_file_lock_with_timeout(&metadata_path, Duration::from_secs(1)) + .await + .expect("first transaction lock"); + let (attempting_tx, attempting_rx) = tokio::sync::oneshot::channel(); + let second_path = metadata_path.clone(); + let second = tokio::spawn(async move { + let _ = attempting_tx.send(()); + let _lock = + acquire_store_file_lock_with_timeout(&second_path, Duration::from_secs(5)).await?; + reconcile_cleanup_journal(&second_path, &HashSet::new()).await?; + let committed = read_secure_file(&second_path).await?; + Ok::(committed.version) + }); + + attempting_rx.await.expect("second transaction started"); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + !second.is_finished(), + "reconciliation must remain behind the first transaction lock" + ); + assert!(!metadata_path.exists()); + + write_secure_file( + &metadata_path, + &SecureStoreFile { + version: STORE_VERSION, + accounts: HashMap::new(), + provider_revisions: HashMap::new(), + }, + ) + .await + .expect("commit metadata while the first transaction owns the lock"); + drop(first); + + let observed_version = tokio::time::timeout(Duration::from_secs(2), second) + .await + .expect("second transaction should proceed after commit") + .expect("second transaction task") + .expect("second transaction result"); + assert_eq!(observed_version, STORE_VERSION); + } + + #[tokio::test] + async fn cross_process_lock_child() { + let Some(metadata_path) = std::env::var_os(LOCK_CHILD_METADATA_ENV).map(PathBuf::from) + else { + return; + }; + let started_path = PathBuf::from( + std::env::var_os(LOCK_CHILD_STARTED_ENV).expect("child started marker path"), + ); + let observed_path = PathBuf::from( + std::env::var_os(LOCK_CHILD_OBSERVED_ENV).expect("child observed marker path"), + ); + std::fs::write(&started_path, b"started").expect("write child started marker"); + + let _lock = acquire_store_file_lock_with_timeout(&metadata_path, Duration::from_secs(5)) + .await + .expect("child transaction lock"); + // Keep a lock-regression failure inside the process-local test vault; + // it must never touch the developer's native credential store. + set_store_path_for_test(metadata_path.clone()); + let committed = read_secure_file(&metadata_path) + .await + .expect("child reads committed metadata"); + reconcile_cleanup_journal(&metadata_path, &committed.active_vault_entries()) + .await + .expect("child reconciliation"); + std::fs::write(observed_path, committed.accounts.len().to_string()) + .expect("write child observation"); + } + + #[tokio::test] + async fn separate_process_waits_for_metadata_commit_before_reconciliation() { + let metadata_path = temporary_metadata_path("cross-process-commit-order"); + let first = acquire_store_file_lock_with_timeout(&metadata_path, Duration::from_secs(1)) + .await + .expect("parent transaction lock"); + let parent = metadata_path.parent().expect("metadata parent"); + let started_path = parent.join("child-started"); + let observed_path = parent.join("child-observed"); + let committed_metadata = CredentialMetadata::Api { + metadata: None, + needs_reauthentication: false, + secret_set_id: Some("committed-set".to_string()), + key_parts: 1, + }; + write_cleanup_journal( + &metadata_path, + &CleanupJournal { + version: CLEANUP_JOURNAL_VERSION, + entries: committed_metadata + .vault_entries("codex") + .into_iter() + .collect(), + }, + ) + .await + .expect("stage cleanup intent before vault write and metadata commit"); + let mut child = std::process::Command::new(std::env::current_exe().expect("test binary")) + .arg("--exact") + .arg("subscription_auth::store::file_lock_tests::cross_process_lock_child") + .arg("--nocapture") + .env(LOCK_CHILD_METADATA_ENV, &metadata_path) + .env(LOCK_CHILD_STARTED_ENV, &started_path) + .env(LOCK_CHILD_OBSERVED_ENV, &observed_path) + .spawn() + .expect("spawn lock-contending child process"); + + tokio::time::timeout(Duration::from_secs(2), async { + while !started_path.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("child should begin its transaction attempt"); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + !observed_path.exists(), + "child process must not reconcile while the parent transaction is uncommitted" + ); + + write_secure_file( + &metadata_path, + &SecureStoreFile { + version: STORE_VERSION, + accounts: HashMap::from([("codex".to_string(), committed_metadata)]), + provider_revisions: HashMap::new(), + }, + ) + .await + .expect("parent metadata commit"); + drop(first); + + let status = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Some(status) = child.try_wait().expect("poll child process") { + break status; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("child should finish after the parent releases the lock"); + assert!(status.success(), "child transaction failed: {status}"); + assert_eq!( + std::fs::read_to_string(observed_path).expect("child observation"), + "1" + ); + assert!(read_cleanup_journal(&metadata_path) + .await + .expect("read reconciled journal") + .entries + .is_empty()); + } + + #[tokio::test] + async fn transaction_lock_rejects_non_regular_paths() { + let metadata_path = temporary_metadata_path("regular-file-check"); + let lock_path = store_lock_path(&metadata_path); + std::fs::create_dir_all(&lock_path).expect("create directory at lock path"); + + let error = match acquire_store_file_lock_with_timeout( + &metadata_path, + Duration::from_secs(1), + ) + .await + { + Ok(_) => panic!("a directory must not be accepted as a lock file"), + Err(error) => error, + }; + let message = format!("{error:#}"); + assert!(message.contains("transaction lock")); + assert!(message.contains("directory") || message.contains("regular file")); + } +} diff --git a/src/crates/assembly/core/src/agentic/tools/account_login_capability.rs b/src/crates/assembly/core/src/agentic/tools/account_login_capability.rs index ac9058deab..e5931506ff 100644 --- a/src/crates/assembly/core/src/agentic/tools/account_login_capability.rs +++ b/src/crates/assembly/core/src/agentic/tools/account_login_capability.rs @@ -2,6 +2,9 @@ use std::sync::atomic::{AtomicBool, Ordering}; +#[cfg(test)] +use std::sync::{Mutex, MutexGuard}; + static ACCOUNT_LOGIN_AVAILABLE: AtomicBool = AtomicBool::new(false); /// Mark whether the current process has a fully logged-in BitFun account session. @@ -13,12 +16,37 @@ pub fn account_login_available() -> bool { ACCOUNT_LOGIN_AVAILABLE.load(Ordering::SeqCst) } +#[cfg(test)] +static ACCOUNT_LOGIN_TEST_LOCK: Mutex<()> = Mutex::new(()); + +#[cfg(test)] +pub(crate) struct AccountLoginTestGuard { + _guard: MutexGuard<'static, ()>, +} + +#[cfg(test)] +impl Drop for AccountLoginTestGuard { + fn drop(&mut self) { + set_account_login_available(false); + } +} + +#[cfg(test)] +pub(crate) fn lock_account_login_for_test() -> AccountLoginTestGuard { + let guard = ACCOUNT_LOGIN_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + set_account_login_available(false); + AccountLoginTestGuard { _guard: guard } +} + #[cfg(test)] mod tests { use super::*; #[test] fn toggles_account_login_availability() { + let _guard = lock_account_login_for_test(); set_account_login_available(false); assert!(!account_login_available()); set_account_login_available(true); diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/page_deploy_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/page_deploy_tool.rs index d0312db44d..d326201ec7 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/page_deploy_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/page_deploy_tool.rs @@ -31,9 +31,9 @@ impl Tool for PageDeployTool { Ok( r#"Switch the production pointer of an existing BitFun Page to a previously saved version_id (rollback or promote a prior version). -Requires a logged-in BitFun account. This tool is only available after account login. To create or update page content and publish, use PagePublish instead — do not ask the user for a version_id they do not have, and do not mention a Page management scene. +Requires a logged-in BitFun account. This tool is only available after account login. To create or update page content and publish, use PagePublish instead. Existing versions can also be reviewed from the Pages scene. -Input: slug (page path id), version_id (immutable saved version from a prior PagePublish). Returns absolute `url` plus url_path / deployed_version_id. +Input: slug (page path id), version_id (immutable saved version from a prior PagePublish). Returns absolute `url` plus url_path / deployed_version_id. Public links can be shared directly. Private and relay links must be opened or copied through the Pages scene/tool card so the browser receives a scoped one-time access handoff. When telling the user the link: paste the full absolute URL and put a trailing space after it (before any punctuation or newline). @@ -70,10 +70,43 @@ Preview a version at /p/{username}/{slug}/@v/{version_id}."# fn permission_intents( &self, - _input: &Value, + input: &Value, _context: &ToolUseContext, ) -> BitFunResult> { - Ok(Vec::new()) + let slug = input + .get("slug") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(""); + let version_id = input + .get("version_id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(""); + let resource = format!("page:{slug}; production-version={version_id}"); + let mut intent = PermissionIntent::new("page_deploy", vec![resource]); + intent.save_resources.clear(); + intent.display_metadata.insert( + "permissionScope".to_string(), + Value::String("account".to_string()), + ); + intent + .display_metadata + .insert("requiresFreshApproval".to_string(), Value::Bool(true)); + intent.display_metadata.insert( + "pageOperation".to_string(), + Value::String("deploy".to_string()), + ); + intent + .display_metadata + .insert("pageSlug".to_string(), Value::String(slug.to_string())); + intent.display_metadata.insert( + "pageVersion".to_string(), + Value::String(version_id.to_string()), + ); + Ok(vec![intent]) } async fn is_available_in_context(&self, _context: Option<&ToolUseContext>) -> bool { @@ -116,15 +149,11 @@ Preview a version at /p/{username}/{slug}/@v/{version_id}."# .filter(|s| !s.is_empty()) .or_else(|| result.get("url_path").and_then(|v| v.as_str())) .unwrap_or(""); - // Trailing space after URL keeps chat linkifiers from eating the next char. - let assistant = if url.is_empty() { - format!("Deployed BitFun Page '{slug}' version '{version_id}' to production.") - } else { - format!( - "Deployed BitFun Page '{slug}' version '{version_id}'. Production URL: {url} \n\ - Share this full absolute URL with the user, and keep a trailing space after the URL." - ) - }; + let visibility = result + .get("visibility") + .and_then(Value::as_str) + .unwrap_or("private"); + let assistant = deploy_result_for_assistant(&slug, &version_id, visibility, url); Ok(vec![ToolResult::Result { data: result, @@ -134,14 +163,35 @@ Preview a version at /p/{username}/{slug}/@v/{version_id}."# } } +fn deploy_result_for_assistant( + slug: &str, + version_id: &str, + visibility: &str, + url: &str, +) -> String { + if visibility != "public" { + return format!( + "Deployed BitFun Page '{slug}' version '{version_id}' with {visibility} visibility. Open or copy it from the Pages scene/tool card so BitFun can create a scoped browser-access link; do not share the raw Page URL." + ); + } + // Trailing space after URL keeps chat linkifiers from eating the next char. + if url.is_empty() { + format!("Deployed public BitFun Page '{slug}' version '{version_id}' to production.") + } else { + format!( + "Deployed public BitFun Page '{slug}' version '{version_id}'. Production URL: {url} \n\ + Share this full absolute URL with the user, and keep a trailing space after the URL." + ) + } +} + #[cfg(test)] mod tests { use super::*; - use crate::agentic::tools::account_login_capability::set_account_login_available; + use crate::agentic::tools::account_login_capability::{ + lock_account_login_for_test, set_account_login_available, + }; use std::collections::HashMap; - use std::sync::Mutex; - - static LOGIN_GATE: Mutex<()> = Mutex::new(()); fn empty_context() -> ToolUseContext { ToolUseContext { @@ -161,7 +211,7 @@ mod tests { #[tokio::test] async fn login_gate_controls_availability_and_execution() { - let _guard = LOGIN_GATE.lock().unwrap(); + let _guard = lock_account_login_for_test(); let tool = PageDeployTool::new(); set_account_login_available(false); @@ -180,4 +230,47 @@ mod tests { .expect_err("should reject without login"); assert!(err.to_string().contains("logged-in")); } + + #[test] + fn permission_intent_names_page_and_target_version() { + let tool = PageDeployTool::new(); + let intents = tool + .permission_intents( + &json!({ "slug": "demo", "version_id": "v2026" }), + &empty_context(), + ) + .expect("permission intent"); + + assert_eq!(intents.len(), 1); + assert_eq!(intents[0].action, "page_deploy"); + assert_eq!( + intents[0].resources, + vec!["page:demo; production-version=v2026"] + ); + assert!(intents[0].save_resources.is_empty()); + assert_eq!( + intents[0].display_metadata.get("permissionScope"), + Some(&Value::String("account".to_string())) + ); + assert_eq!( + intents[0].display_metadata.get("requiresFreshApproval"), + Some(&Value::Bool(true)) + ); + assert_eq!( + intents[0].display_metadata.get("pageOperation"), + Some(&Value::String("deploy".to_string())) + ); + } + + #[test] + fn private_deploy_result_never_advertises_the_raw_url() { + let message = deploy_result_for_assistant( + "demo", + "v1", + "private", + "https://relay.example/p/alice/demo", + ); + assert!(!message.contains("https://")); + assert!(message.contains("scoped browser-access link")); + } } diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/page_publish_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/page_publish_tool.rs index ea5293ce2e..348dfea8ce 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/page_publish_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/page_publish_tool.rs @@ -1,10 +1,10 @@ //! PagePublish tool — create/update a BitFun Page from inline files or a directory, -//! then optionally deploy to production (default: deploy). +//! then optionally deploy to production. use std::collections::HashMap; use crate::agentic::tools::account_login_capability::account_login_available; -use crate::agentic::tools::framework::{Tool, ToolResult, ToolUseContext}; +use crate::agentic::tools::framework::{PermissionIntent, Tool, ToolResult, ToolUseContext}; use crate::agentic::tools::page_publish_host::{invoke_page_publish, PagePublishHostRequest}; use crate::util::errors::{BitFunError, BitFunResult}; use async_trait::async_trait; @@ -32,26 +32,26 @@ impl Tool for PagePublishTool { async fn description(&self) -> BitFunResult { Ok( - r#"Publish a BitFun Page to the account relay: upload content, freeze an immutable version, and deploy to production by default. + r#"Publish a BitFun Page to the account relay: upload content, freeze an immutable version, and optionally deploy it to production. -Requires a logged-in BitFun account. This tool is only available after account login. There is no separate Page management scene — create and ship pages from the conversation with this tool. +Requires a logged-in BitFun account. This tool is only available after account login. Published Pages can be reviewed and managed later from the Pages scene. When you produce self-contained publishable web content (landing page, docs site, or a Page with server/worker.js) and the user is logged in, proactively ask whether they want it published to BitFun Page (suggest a slug and visibility). If they already said publish/deploy/上线, proceed with permission confirmation. IMPORTANT — content source: - Prefer `files` (inline path→UTF-8 content). For agent-authored pages, pass HTML/JS directly in `files` and call PagePublish. Do NOT Write/Edit page files into the user workspace just to publish, and do NOT create folders like bitfun-page/ unless the user explicitly asked to keep a local copy. -- Use `directory` only when the user already has (or explicitly wants) page sources on disk in the workspace. +- Use `directory` only when the user already has (or explicitly wants) page sources on disk in a local workspace. Remote workspaces must use `files` so BitFun never mistakes a remote path for a local path. Input: - slug (required): page path id (lowercase letters, digits, hyphens) -- visibility: private | relay | public (default public) +- visibility: private | relay | public (default private) - title?, note? -- deploy: boolean (default true). false = save version only; still returns absolute preview_url +- deploy: boolean (default false). false = save version only; still returns absolute preview_url - Exactly one of: - files: object map of relative path → UTF-8 file content (default/preferred). Must include index.html and/or server/worker.js - directory: existing local workspace path (only when user wants on-disk sources) -Returns version_id, absolute `url` / `preview_url` (plus relative paths), deployed_version_id when deployed. +Returns version_id, absolute `url` / `preview_url` (plus relative paths), deployed_version_id when deployed. Public links can be shared directly. Private and relay links must be opened or copied through the Pages scene/tool card so the browser receives a scoped one-time access handoff; never share their raw URL as if it were independently accessible. When telling the user the link: paste the full absolute URL and put a trailing space after it (before any punctuation or newline), so chat linkifiers do not swallow the next character. @@ -77,7 +77,7 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a "visibility": { "type": "string", "enum": ["private", "relay", "public"], - "description": "Page visibility. Defaults to public." + "description": "Page visibility. Defaults to private. Use public only when the user explicitly intends to share the page publicly." }, "title": { "type": "string", @@ -89,7 +89,7 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a }, "deploy": { "type": "boolean", - "description": "Deploy to production after saving. Defaults to true." + "description": "Deploy to production after saving. Defaults to false; set true only when the user explicitly asked to publish or deploy." }, "files": { "type": "object", @@ -98,7 +98,7 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a }, "directory": { "type": "string", - "description": "Only when page sources already exist (or user asked to keep them) on disk. Mutually exclusive with files." + "description": "Only when page sources already exist (or user asked to keep them) on disk in a local workspace. Remote workspaces must use files. Mutually exclusive with files." } } }) @@ -108,6 +108,60 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a false } + fn permission_intents( + &self, + input: &Value, + _context: &ToolUseContext, + ) -> BitFunResult> { + let slug = input + .get("slug") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(""); + let visibility = input + .get("visibility") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("private"); + let deploy = input + .get("deploy") + .and_then(Value::as_bool) + .unwrap_or(false); + let resource = format!( + "page:{slug}; visibility={visibility}; deploy={}", + if deploy { + "production" + } else { + "saved-version-only" + } + ); + let mut intent = PermissionIntent::new("page_publish", vec![resource]); + // Publishing decisions should stay per-call: a remembered wildcard grant could + // otherwise hide a later change from private preview to public production. + intent.save_resources.clear(); + intent.display_metadata.insert( + "permissionScope".to_string(), + Value::String("account".to_string()), + ); + intent + .display_metadata + .insert("requiresFreshApproval".to_string(), Value::Bool(true)); + intent.display_metadata.insert( + "pageOperation".to_string(), + Value::String(if deploy { "publish" } else { "save" }.to_string()), + ); + intent + .display_metadata + .insert("pageSlug".to_string(), Value::String(slug.to_string())); + intent.display_metadata.insert( + "pageVisibility".to_string(), + Value::String(visibility.to_string()), + ); + Ok(vec![intent]) + } + async fn is_available_in_context(&self, _context: Option<&ToolUseContext>) -> bool { account_login_available() } @@ -115,7 +169,7 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a async fn call_impl( &self, input: &Value, - _context: &ToolUseContext, + context: &ToolUseContext, ) -> BitFunResult> { if !account_login_available() { return Err(BitFunError::tool( @@ -136,7 +190,7 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a .and_then(|v| v.as_str()) .map(str::trim) .filter(|s| !s.is_empty()) - .unwrap_or("public") + .unwrap_or("private") .to_string(); let title = input @@ -155,7 +209,7 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a let deploy = input .get("deploy") .and_then(|v| v.as_bool()) - .unwrap_or(true); + .unwrap_or(false); let directory = input .get("directory") @@ -164,6 +218,19 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a .filter(|s| !s.is_empty()) .map(str::to_string); + if directory.is_some() && context.workspace.is_none() { + return Err(BitFunError::tool( + "PagePublish directory requires a local workspace; use inline files when no workspace is open" + .to_string(), + )); + } + if directory.is_some() && context.is_remote() { + return Err(BitFunError::tool( + "PagePublish cannot read a remote workspace directory through the local desktop host; use inline files" + .to_string(), + )); + } + let files = parse_files_map(input.get("files"))?; match (&directory, &files) { @@ -182,7 +249,7 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a let result = invoke_page_publish(PagePublishHostRequest { slug: slug.clone(), - visibility, + visibility: visibility.clone(), title, note, deploy, @@ -213,24 +280,8 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a .and_then(|v| v.as_bool()) .unwrap_or(false); - // Trailing space after URL keeps chat linkifiers from eating the next char. - let assistant = if deployed { - if url.is_empty() { - format!("Published BitFun Page '{slug}' version '{version_id}' to production.") - } else { - format!( - "Published BitFun Page '{slug}' version '{version_id}'. Production URL: {url} \n\ - Share this full absolute URL with the user, and keep a trailing space after the URL." - ) - } - } else if preview.is_empty() { - format!("Saved BitFun Page '{slug}' version '{version_id}' (not deployed).") - } else { - format!( - "Saved BitFun Page '{slug}' version '{version_id}' (not deployed). Preview URL: {preview} \n\ - Share this full absolute URL with the user, and keep a trailing space after the URL." - ) - }; + let assistant = + publish_result_for_assistant(&slug, version_id, &visibility, deployed, url, preview); Ok(vec![ToolResult::Result { data: result, @@ -240,6 +291,45 @@ Use PageDeploy only to switch an already-saved version_id (rollback / promote a } } +fn publish_result_for_assistant( + slug: &str, + version_id: &str, + visibility: &str, + deployed: bool, + url: &str, + preview: &str, +) -> String { + if visibility != "public" { + let state = if deployed { + "published to production" + } else { + "saved without changing production" + }; + return format!( + "BitFun Page '{slug}' version '{version_id}' was {state} with {visibility} visibility. Open or copy it from the Pages scene/tool card so BitFun can create a scoped browser-access link; do not share the raw Page URL." + ); + } + + // Trailing space after URL keeps chat linkifiers from eating the next char. + if deployed { + if url.is_empty() { + format!("Published public BitFun Page '{slug}' version '{version_id}' to production.") + } else { + format!( + "Published public BitFun Page '{slug}' version '{version_id}'. Production URL: {url} \n\ + Share this full absolute URL with the user, and keep a trailing space after the URL." + ) + } + } else if preview.is_empty() { + format!("Saved public BitFun Page '{slug}' version '{version_id}' (not deployed).") + } else { + format!( + "Saved public BitFun Page '{slug}' version '{version_id}' (not deployed). Preview URL: {preview} \n\ + Share this full absolute URL with the user, and keep a trailing space after the URL." + ) + } +} + fn parse_files_map(value: Option<&Value>) -> BitFunResult>> { let Some(value) = value else { return Ok(None); @@ -268,10 +358,9 @@ fn parse_files_map(value: Option<&Value>) -> BitFunResult = Mutex::new(()); + use crate::agentic::tools::account_login_capability::{ + lock_account_login_for_test, set_account_login_available, + }; fn empty_context() -> ToolUseContext { ToolUseContext { @@ -291,7 +380,7 @@ mod tests { #[tokio::test] async fn login_gate_controls_availability_and_execution() { - let _guard = LOGIN_GATE.lock().unwrap(); + let _guard = lock_account_login_for_test(); let tool = PagePublishTool::new(); set_account_login_available(false); @@ -316,7 +405,7 @@ mod tests { #[tokio::test] async fn rejects_missing_source_when_logged_in() { - let _guard = LOGIN_GATE.lock().unwrap(); + let _guard = lock_account_login_for_test(); let tool = PagePublishTool::new(); set_account_login_available(true); let err = tool @@ -326,4 +415,110 @@ mod tests { assert!(err.to_string().contains("directory or files")); set_account_login_available(false); } + + #[test] + fn permission_intent_exposes_safe_defaults_and_publish_scope() { + let tool = PagePublishTool::new(); + let intents = tool + .permission_intents( + &json!({ + "slug": "release-notes", + "files": { "index.html": "" } + }), + &empty_context(), + ) + .expect("permission intent"); + + assert_eq!(intents.len(), 1); + assert_eq!(intents[0].action, "page_publish"); + assert_eq!( + intents[0].resources, + vec!["page:release-notes; visibility=private; deploy=saved-version-only"] + ); + assert!(intents[0].save_resources.is_empty()); + assert_eq!( + intents[0].display_metadata.get("permissionScope"), + Some(&Value::String("account".to_string())) + ); + assert_eq!( + intents[0].display_metadata.get("requiresFreshApproval"), + Some(&Value::Bool(true)) + ); + assert_eq!( + intents[0].display_metadata.get("pageOperation"), + Some(&Value::String("save".to_string())) + ); + + let public_deploy = tool + .permission_intents( + &json!({ + "slug": "launch", + "visibility": "public", + "deploy": true, + "files": { "index.html": "" } + }), + &empty_context(), + ) + .expect("public deploy permission intent"); + assert_eq!( + public_deploy[0].resources, + vec!["page:launch; visibility=public; deploy=production"] + ); + assert_eq!( + public_deploy[0].display_metadata.get("pageOperation"), + Some(&Value::String("publish".to_string())) + ); + } + + #[tokio::test] + async fn directory_source_requires_a_local_workspace() { + let _guard = lock_account_login_for_test(); + let tool = PagePublishTool::new(); + set_account_login_available(true); + + let no_workspace_error = tool + .call_impl( + &json!({ "slug": "demo", "directory": "page" }), + &empty_context(), + ) + .await + .expect_err("directory without workspace should be rejected"); + assert!(no_workspace_error.to_string().contains("local workspace")); + + let mut remote_context = empty_context(); + remote_context.workspace = Some(crate::agentic::WorkspaceBinding::new_remote( + None, + std::path::PathBuf::from("/srv/page"), + "connection-1".to_string(), + "Remote".to_string(), + crate::service::remote_ssh::workspace_state::WorkspaceSessionIdentity { + hostname: "remote.example".to_string(), + logical_workspace_path: "/srv/page".to_string(), + remote_connection_id: Some("connection-1".to_string()), + }, + )); + let remote_error = tool + .call_impl( + &json!({ "slug": "demo", "directory": "." }), + &remote_context, + ) + .await + .expect_err("remote directory should be rejected"); + assert!(remote_error.to_string().contains("remote workspace")); + set_account_login_available(false); + } + + #[test] + fn private_result_never_advertises_the_raw_url() { + let message = publish_result_for_assistant( + "demo", + "v1", + "private", + true, + "https://relay.example/p/alice/demo", + "https://relay.example/p/alice/demo/@v/v1", + ); + assert!(!message.contains("https://")); + assert!(message.contains("scoped browser-access link")); + } } diff --git a/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs b/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs index dadee63a78..ea8eb24dab 100644 --- a/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs +++ b/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs @@ -500,6 +500,40 @@ fn permission_project_path(context: &ToolUseContext) -> BitFunResult { .to_string()) } +const ACCOUNT_PERMISSION_SCOPE: &str = "account"; +const ACCOUNT_PERMISSION_PROJECT_ID: &str = "__bitfun_account_actions__"; +const ACCOUNT_PERMISSION_PROJECT_PATH: &str = "BitFun account"; + +fn permission_scope( + context: &ToolUseContext, + intents: &[PermissionIntent], +) -> BitFunResult<(String, String)> { + if context.workspace.is_some() { + return Ok(( + permission_project_id(context)?, + permission_project_path(context)?, + )); + } + + let account_scoped = intents.iter().all(|intent| { + intent + .display_metadata + .get("permissionScope") + .and_then(serde_json::Value::as_str) + == Some(ACCOUNT_PERMISSION_SCOPE) + }); + if account_scoped { + return Ok(( + ACCOUNT_PERMISSION_PROJECT_ID.to_string(), + ACCOUNT_PERMISSION_PROJECT_PATH.to_string(), + )); + } + + Err(BitFunError::validation( + "A workspace is required for file permissions".to_string(), + )) +} + fn permission_resource_case_sensitivity( context: &ToolUseContext, ) -> PermissionResourceCaseSensitivity { @@ -564,10 +598,21 @@ fn permission_intent_effect( } } - if intent.resources.is_empty() { + let effect = if intent.resources.is_empty() { PermissionEffect::Ask } else { aggregate + }; + if effect != PermissionEffect::Deny + && intent + .display_metadata + .get("requiresFreshApproval") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + { + PermissionEffect::Ask + } else { + effect } } @@ -623,8 +668,7 @@ impl ToolPipeline { return Ok(PermissionPlanDraft::Allowed); } - let project_id = permission_project_id(&context)?; - let project_path = permission_project_path(&context)?; + let (project_id, project_path) = permission_scope(&context, &intents)?; let permission_rules = task.options.permission_rules.clone(); let case_sensitivity = permission_resource_case_sensitivity(&context); let round_id = task.context.round_id.clone(); @@ -2209,6 +2253,65 @@ mod tests { ); } + #[test] + fn account_scoped_fresh_approval_works_without_a_workspace_and_ignores_allow_rules() { + let mut intent = PermissionIntent::new( + "page_publish", + vec!["page:demo; visibility=private; deploy=saved-version-only".to_string()], + ); + intent.display_metadata.insert( + "permissionScope".to_string(), + json!(ACCOUNT_PERMISSION_SCOPE), + ); + intent + .display_metadata + .insert("requiresFreshApproval".to_string(), json!(true)); + let context = ToolUseContext::for_tool_listing(None, None); + assert_eq!( + permission_scope(&context, &[intent.clone()]).expect("account scope"), + ( + ACCOUNT_PERMISSION_PROJECT_ID.to_string(), + ACCOUNT_PERMISSION_PROJECT_PATH.to_string(), + ) + ); + + let allow = vec![PermissionRule::new( + "page_publish", + "*", + PermissionEffect::Allow, + )]; + assert_eq!( + permission_intent_effect( + &intent, + &allow, + &[], + PermissionResourceCaseSensitivity::Sensitive, + ), + PermissionEffect::Ask + ); + let deny = vec![PermissionRule::new( + "page_publish", + "*", + PermissionEffect::Deny, + )]; + assert_eq!( + permission_intent_effect( + &intent, + &deny, + &[], + PermissionResourceCaseSensitivity::Sensitive, + ), + PermissionEffect::Deny + ); + } + + #[test] + fn ordinary_permission_intents_still_require_a_workspace() { + let context = ToolUseContext::for_tool_listing(None, None); + let intent = PermissionIntent::new("edit", vec!["src/main.rs".to_string()]); + assert!(permission_scope(&context, &[intent]).is_err()); + } + struct StaticTestTool { name: String, response: serde_json::Value, diff --git a/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs b/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs index 1ea27e33b6..8b152a4d61 100644 --- a/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs +++ b/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs @@ -840,9 +840,7 @@ async fn confirm_then_run( } async fn set_verbose(state: &mut BotChatState, on: bool, s: &'static BotStrings) -> HandleResult { - let mut data = super::load_bot_persistence(); - data.verbose_mode = on; - super::save_bot_persistence(&data); + super::update_bot_persistence(|data| data.verbose_mode = on); let body = if on { s.verbose_enabled @@ -2294,20 +2292,19 @@ async fn route_pending( let (device_id, device_name) = options[n - 1].clone(); if device_id == "local" { // Switch back to local - state.active_remote_device = None; - state.current_session_id = None; + state.select_local_device(); let body = s.devices_switched_local.to_string(); let mut view = main_menu_view(state, s); view = view.with_body(body); result_from_menu(state, view) } else { // Switch to remote device - state.active_remote_device = - Some(crate::service::remote_connect::bot::RemoteDeviceTarget { + state.select_remote_device( + crate::service::remote_connect::bot::RemoteDeviceTarget { device_id: device_id.clone(), device_name: device_name.clone(), - }); - state.current_session_id = None; + }, + ); let body = format!("{}: {}", s.devices_switched_to, device_name); let mut view = main_menu_view(state, s); view = view.with_body(body); diff --git a/src/crates/assembly/core/src/service/remote_connect/bot/feishu.rs b/src/crates/assembly/core/src/service/remote_connect/bot/feishu.rs index 92fc002f71..f6b2bc8874 100644 --- a/src/crates/assembly/core/src/service/remote_connect/bot/feishu.rs +++ b/src/crates/assembly/core/src/service/remote_connect/bot/feishu.rs @@ -21,7 +21,9 @@ use super::command_router::{ parse_command, welcome_message, BotAction, BotChatState, BotInteractionHandler, BotInteractiveRequest, BotLanguage, BotMessageSender, HandleResult, }; -use super::{load_bot_persistence, save_bot_persistence, BotConfig, SavedBotConnection}; +use super::{ + load_bot_persistence, update_bot_persistence, BotConfig, BotRuntimeFence, SavedBotConnection, +}; use crate::service::remote_connect::remote_server::ImageAttachment; #[derive(Debug, Clone)] @@ -33,6 +35,7 @@ pub struct FeishuBot { api: FeishuBotApi, pending_pairings: Arc>>, chat_states: Arc>>, + runtime_fence: BotRuntimeFence, } impl FeishuBot { @@ -85,18 +88,53 @@ impl FeishuBot { } pub fn new(config: FeishuConfig) -> Self { + Self::new_fenced(config, BotRuntimeFence::standalone()) + } + + pub(crate) fn new_fenced(config: FeishuConfig, runtime_fence: BotRuntimeFence) -> Self { Self { api: FeishuBotApi::new(config), pending_pairings: Arc::new(RwLock::new(HashMap::new())), chat_states: Arc::new(RwLock::new(HashMap::new())), + runtime_fence, } } - pub async fn restore_chat_state(&self, chat_id: &str, state: BotChatState) { - self.chat_states - .write() - .await - .insert(chat_id.to_string(), state); + pub async fn restore_chat_state(&self, chat_id: &str, mut state: BotChatState) { + state.prepare_for_restore(); + let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); + states.insert(chat_id.to_string(), state); + let restored = states + .get(chat_id) + .cloned() + .expect("restored Feishu state should exist"); + drop(states); + self.persist_chat_state(chat_id, &restored).await; + } + + pub async fn clear_delegated_identities(&self) { + match tokio::time::timeout( + std::time::Duration::from_millis(100), + self.chat_states.write(), + ) + .await + { + Ok(mut states) => { + self.runtime_fence.clear_states(&mut states); + let snapshots: Vec<_> = states + .iter() + .map(|(chat_id, state)| (chat_id.clone(), state.clone())) + .collect(); + drop(states); + for (chat_id, state) in snapshots { + self.persist_chat_state(&chat_id, &state).await; + } + } + Err(_) => { + warn!("Feishu account identity clear deferred behind an in-flight command"); + } + } } pub async fn send_message(&self, chat_id: &str, content: &str) -> Result<()> { @@ -277,13 +315,19 @@ impl FeishuBot { if self.verify_pairing_code(trimmed).await { info!("Feishu pairing successful, chat_id={chat_id}"); let mut state = BotChatState::new(chat_id.clone()); + let identity_epoch = self.runtime_fence.identity_epoch(); let result = complete_im_bot_pairing(&mut state).await; - self.send_handle_result(&chat_id, &result).await.ok(); - self.chat_states - .write() - .await - .insert(chat_id.clone(), state.clone()); + if !self.runtime_fence.is_lifecycle_current() { + return None; + } + let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); + self.runtime_fence + .sanitize_after_epoch(identity_epoch, &mut state); + states.insert(chat_id.clone(), state.clone()); + drop(states); self.persist_chat_state(&chat_id, &state).await; + self.send_handle_result(&chat_id, &result).await.ok(); return Some(chat_id); } else { @@ -488,7 +532,11 @@ impl FeishuBot { text: &str, images: Vec, ) { + if !self.runtime_fence.is_lifecycle_current() { + return; + } let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); let state = states.entry(chat_id.to_string()).or_insert_with(|| { let mut s = BotChatState::new(chat_id.to_string()); s.paired = true; @@ -506,7 +554,13 @@ impl FeishuBot { } if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) { if self.verify_pairing_code(trimmed).await { + let identity_epoch = self.runtime_fence.identity_epoch(); let result = complete_im_bot_pairing(state).await; + self.runtime_fence + .sanitize_after_epoch(identity_epoch, state); + if !self.runtime_fence.is_lifecycle_current() { + return; + } self.send_handle_result(chat_id, &result).await.ok(); self.persist_chat_state(chat_id, state).await; return; @@ -526,9 +580,16 @@ impl FeishuBot { let cmd = parse_command(text); let result = handle_command(state, cmd, images).await; - self.persist_chat_state(chat_id, state).await; + self.runtime_fence.reconcile_states(&mut states); + if let Some(state) = states.get(chat_id) { + self.persist_chat_state(chat_id, state).await; + } drop(states); + if !self.runtime_fence.is_lifecycle_current() { + return; + } + self.send_handle_result(chat_id, &result).await.ok(); if let Some(forward) = result.forward_to_session { @@ -573,7 +634,11 @@ impl FeishuBot { } async fn deliver_interaction(&self, chat_id: &str, interaction: BotInteractiveRequest) { + if !self.runtime_fence.is_lifecycle_current() { + return; + } let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); let state = states.entry(chat_id.to_string()).or_insert_with(|| { let mut s = BotChatState::new(chat_id.to_string()); s.paired = true; @@ -593,18 +658,20 @@ impl FeishuBot { } async fn persist_chat_state(&self, chat_id: &str, state: &BotChatState) { - let mut data = load_bot_persistence(); - data.upsert(SavedBotConnection { + let snapshot = self.runtime_fence.persistence_snapshot(state); + let connection = SavedBotConnection { bot_type: "feishu".to_string(), chat_id: chat_id.to_string(), config: BotConfig::Feishu { app_id: self.api.config().app_id.clone(), app_secret: self.api.config().app_secret.clone(), }, - chat_state: state.clone(), + chat_state: snapshot, connected_at: chrono::Utc::now().timestamp(), + }; + self.runtime_fence.commit_if_current(|| { + update_bot_persistence(|data| data.upsert(connection)); }); - save_bot_persistence(&data); } } diff --git a/src/crates/assembly/core/src/service/remote_connect/bot/mod.rs b/src/crates/assembly/core/src/service/remote_connect/bot/mod.rs index 2ee4ca9831..fc2fa6bc4b 100644 --- a/src/crates/assembly/core/src/service/remote_connect/bot/mod.rs +++ b/src/crates/assembly/core/src/service/remote_connect/bot/mod.rs @@ -16,11 +16,213 @@ pub use bitfun_services_integrations::remote_connect::bot::{ auto_push_failed_message, auto_push_intro, auto_push_skip_too_large_message, collect_auto_push_files, detect_mime_type, extract_computer_file_paths, extract_downloadable_file_paths, format_file_size, get_file_metadata, load_bot_persistence, - read_workspace_file, resolve_workspace_path, save_bot_persistence, AutoPushFile, BotConfig, - BotLanguage, BotPairingInfo, BotPersistenceData, MenuItem, MenuItemStyle, MenuView, - RemoteConnectFormState, RemoteDeviceTarget, SavedBotConnection, WorkspaceFileContent, + read_workspace_file, resolve_workspace_path, save_bot_persistence, update_bot_persistence, + AutoPushFile, BotConfig, BotLanguage, BotPairingInfo, BotPersistenceData, MenuItem, + MenuItemStyle, MenuView, RemoteConnectFormState, RemoteDeviceTarget, SavedBotConnection, + WorkspaceFileContent, }; pub use command_router::{ set_delegated_identity_provider, BotChatState, ForwardRequest, ForwardedTurnResult, HandleResult, }; + +use std::collections::HashMap; +use std::hash::Hash; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex as StdMutex}; + +/// Serializes persistence commits with lifecycle generation changes for one +/// bot platform. A retired polling task can finish later, but it can no longer +/// overwrite the replacement instance's persisted state. +#[derive(Default)] +pub(crate) struct BotSlotFence { + generation: StdMutex, +} + +impl BotSlotFence { + pub(crate) fn advance(&self) -> u64 { + let mut generation = self + .generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *generation = generation.wrapping_add(1); + if *generation == 0 { + *generation = 1; + } + *generation + } + + pub(crate) fn is_current(&self, expected_generation: u64) -> bool { + *self + .generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + == expected_generation + } + + fn commit_if_current( + &self, + expected_generation: u64, + commit: impl FnOnce() -> R, + ) -> Option { + let generation = self + .generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if *generation != expected_generation { + return None; + } + Some(commit()) + } +} + +/// Shared account-identity epoch plus a per-platform lifecycle owner. Network +/// work is not awaited during account replacement; instead every state commit +/// is fenced and sanitized if it started under an older account epoch. +pub(crate) struct BotRuntimeFence { + account_identity_epoch: Arc, + observed_identity_epoch: AtomicU64, + slot: Arc, + lifecycle_generation: u64, +} + +impl BotRuntimeFence { + pub(crate) fn new( + account_identity_epoch: Arc, + slot: Arc, + lifecycle_generation: u64, + ) -> Self { + let observed_identity_epoch = account_identity_epoch.load(Ordering::Acquire); + Self { + account_identity_epoch, + observed_identity_epoch: AtomicU64::new(observed_identity_epoch), + slot, + lifecycle_generation, + } + } + + pub(crate) fn standalone() -> Self { + let account_identity_epoch = Arc::new(AtomicU64::new(0)); + let slot = Arc::new(BotSlotFence::default()); + let lifecycle_generation = slot.advance(); + Self::new(account_identity_epoch, slot, lifecycle_generation) + } + + pub(crate) fn identity_epoch(&self) -> u64 { + self.account_identity_epoch.load(Ordering::Acquire) + } + + pub(crate) fn is_lifecycle_current(&self) -> bool { + self.slot.is_current(self.lifecycle_generation) + } + + pub(crate) fn reconcile_states(&self, states: &mut HashMap) -> bool + where + K: Eq + Hash, + { + let current = self.identity_epoch(); + if self.observed_identity_epoch.load(Ordering::Acquire) == current { + return false; + } + for state in states.values_mut() { + state.clear_delegated_identity(); + } + self.observed_identity_epoch + .store(current, Ordering::Release); + true + } + + pub(crate) fn clear_states(&self, states: &mut HashMap) + where + K: Eq + Hash, + { + let current = self.identity_epoch(); + for state in states.values_mut() { + state.clear_delegated_identity(); + } + self.observed_identity_epoch + .store(current, Ordering::Release); + } + + pub(crate) fn sanitize_after_epoch(&self, started_epoch: u64, state: &mut BotChatState) { + if started_epoch != self.identity_epoch() { + state.clear_delegated_identity(); + } + } + + pub(crate) fn persistence_snapshot(&self, state: &BotChatState) -> BotChatState { + let mut snapshot = state.clone(); + if self.observed_identity_epoch.load(Ordering::Acquire) != self.identity_epoch() { + snapshot.clear_delegated_identity(); + } + snapshot + } + + pub(crate) fn commit_if_current(&self, commit: impl FnOnce() -> R) -> Option { + self.slot + .commit_if_current(self.lifecycle_generation, commit) + } +} + +#[cfg(test)] +mod lifecycle_tests { + use super::*; + + #[test] + fn account_epoch_sanitizes_a_late_state_commit() { + let account_epoch = Arc::new(AtomicU64::new(3)); + let slot = Arc::new(BotSlotFence::default()); + let generation = slot.advance(); + let fence = BotRuntimeFence::new(account_epoch.clone(), slot, generation); + let mut states = HashMap::from([("chat", { + let mut state = BotChatState::new("chat".into()); + state.relay_url = Some("https://relay-a.example".into()); + state.set_delegated_identity("token-a".into(), vec![1; 32]); + state + })]); + + account_epoch.fetch_add(1, Ordering::AcqRel); + + assert!(fence.reconcile_states(&mut states)); + let state = states.get("chat").expect("chat state should remain"); + assert!(state.relay_url.is_none()); + assert!(!state.has_delegated_identity()); + } + + #[test] + fn account_epoch_rejects_identity_returned_by_late_pairing() { + let account_epoch = Arc::new(AtomicU64::new(11)); + let slot = Arc::new(BotSlotFence::default()); + let generation = slot.advance(); + let fence = BotRuntimeFence::new(account_epoch.clone(), slot, generation); + let started_epoch = fence.identity_epoch(); + let mut late_state = BotChatState::new("late-chat".into()); + late_state.relay_url = Some("https://relay-a.example".into()); + late_state.set_delegated_identity("token-a".into(), vec![2; 32]); + + account_epoch.fetch_add(1, Ordering::AcqRel); + fence.sanitize_after_epoch(started_epoch, &mut late_state); + + assert!(late_state.relay_url.is_none()); + assert!(!late_state.has_delegated_identity()); + } + + #[test] + fn retired_slot_cannot_commit_after_replacement() { + let account_epoch = Arc::new(AtomicU64::new(0)); + let slot = Arc::new(BotSlotFence::default()); + let old_generation = slot.advance(); + let old = BotRuntimeFence::new(account_epoch.clone(), slot.clone(), old_generation); + let new_generation = slot.advance(); + let replacement = BotRuntimeFence::new(account_epoch, slot, new_generation); + let committed = AtomicU64::new(0); + + assert!(old + .commit_if_current(|| committed.store(1, Ordering::Release)) + .is_none()); + assert!(replacement + .commit_if_current(|| committed.store(2, Ordering::Release)) + .is_some()); + assert_eq!(committed.load(Ordering::Acquire), 2); + } +} diff --git a/src/crates/assembly/core/src/service/remote_connect/bot/telegram.rs b/src/crates/assembly/core/src/service/remote_connect/bot/telegram.rs index 7083f6230c..26a2b1d155 100644 --- a/src/crates/assembly/core/src/service/remote_connect/bot/telegram.rs +++ b/src/crates/assembly/core/src/service/remote_connect/bot/telegram.rs @@ -19,7 +19,9 @@ use super::command_router::{ parse_command, welcome_message, BotAction, BotChatState, BotInteractionHandler, BotInteractiveRequest, BotLanguage, BotMessageSender, HandleResult, }; -use super::{load_bot_persistence, save_bot_persistence, BotConfig, SavedBotConnection}; +use super::{ + load_bot_persistence, update_bot_persistence, BotConfig, BotRuntimeFence, SavedBotConnection, +}; use crate::service::remote_connect::remote_server::ImageAttachment; pub struct TelegramBot { @@ -27,6 +29,7 @@ pub struct TelegramBot { pending_pairings: Arc>>, last_update_id: Arc>, chat_states: Arc>>, + runtime_fence: BotRuntimeFence, } #[derive(Debug, Clone)] @@ -52,17 +55,55 @@ impl TelegramBot { } pub fn new(config: TelegramConfig) -> Self { + Self::new_fenced(config, BotRuntimeFence::standalone()) + } + + pub(crate) fn new_fenced(config: TelegramConfig, runtime_fence: BotRuntimeFence) -> Self { Self { api: TelegramBotApi::new(config), pending_pairings: Arc::new(RwLock::new(HashMap::new())), last_update_id: Arc::new(RwLock::new(0)), chat_states: Arc::new(RwLock::new(HashMap::new())), + runtime_fence, } } /// Restore a previously paired chat so the bot skips the pairing step. - pub async fn restore_chat_state(&self, chat_id: i64, state: BotChatState) { - self.chat_states.write().await.insert(chat_id, state); + pub async fn restore_chat_state(&self, chat_id: i64, mut state: BotChatState) { + state.prepare_for_restore(); + let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); + states.insert(chat_id, state); + let restored = states + .get(&chat_id) + .cloned() + .expect("restored Telegram state should exist"); + drop(states); + self.persist_chat_state(chat_id, &restored).await; + } + + pub async fn clear_delegated_identities(&self) { + match tokio::time::timeout( + std::time::Duration::from_millis(100), + self.chat_states.write(), + ) + .await + { + Ok(mut states) => { + self.runtime_fence.clear_states(&mut states); + let snapshots: Vec<_> = states + .iter() + .map(|(chat_id, state)| (*chat_id, state.clone())) + .collect(); + drop(states); + for (chat_id, state) in snapshots { + self.persist_chat_state(chat_id, &state).await; + } + } + Err(_) => { + warn!("Telegram account identity clear deferred behind an in-flight command"); + } + } } pub async fn send_message(&self, chat_id: i64, text: &str) -> Result<()> { @@ -248,11 +289,17 @@ impl TelegramBot { if self.verify_pairing_code(trimmed).await { info!("Telegram pairing successful, chat_id={chat_id}"); let mut state = BotChatState::new(chat_id.to_string()); + let identity_epoch = self.runtime_fence.identity_epoch(); let result = complete_im_bot_pairing(&mut state).await; - self.chat_states - .write() - .await - .insert(chat_id, state.clone()); + if *stop_rx.borrow() || !self.runtime_fence.is_lifecycle_current() { + return Err(anyhow!("bot lifecycle replaced during pairing")); + } + let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); + self.runtime_fence + .sanitize_after_epoch(identity_epoch, &mut state); + states.insert(chat_id, state.clone()); + drop(states); self.persist_chat_state(chat_id, &state).await; self.send_handle_result(chat_id, &result).await; self.set_bot_commands().await.ok(); @@ -324,7 +371,11 @@ impl TelegramBot { text: &str, images: Vec, ) { + if !self.runtime_fence.is_lifecycle_current() { + return; + } let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); let state = states.entry(chat_id).or_insert_with(|| { let mut s = BotChatState::new(chat_id.to_string()); s.paired = true; @@ -342,8 +393,14 @@ impl TelegramBot { } if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) { if self.verify_pairing_code(trimmed).await { + let identity_epoch = self.runtime_fence.identity_epoch(); let result = complete_im_bot_pairing(state).await; + self.runtime_fence + .sanitize_after_epoch(identity_epoch, state); self.persist_chat_state(chat_id, state).await; + if !self.runtime_fence.is_lifecycle_current() { + return; + } self.send_handle_result(chat_id, &result).await; self.set_bot_commands().await.ok(); return; @@ -363,9 +420,16 @@ impl TelegramBot { let cmd = parse_command(text); let result = handle_command(state, cmd, images).await; - self.persist_chat_state(chat_id, state).await; + self.runtime_fence.reconcile_states(&mut states); + if let Some(state) = states.get(&chat_id) { + self.persist_chat_state(chat_id, state).await; + } drop(states); + if !self.runtime_fence.is_lifecycle_current() { + return; + } + self.send_handle_result(chat_id, &result).await; if let Some(forward) = result.forward_to_session { @@ -401,7 +465,11 @@ impl TelegramBot { } async fn deliver_interaction(&self, chat_id: i64, interaction: BotInteractiveRequest) { + if !self.runtime_fence.is_lifecycle_current() { + return; + } let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); let state = states.entry(chat_id).or_insert_with(|| { let mut s = BotChatState::new(chat_id.to_string()); s.paired = true; @@ -421,16 +489,38 @@ impl TelegramBot { } async fn persist_chat_state(&self, chat_id: i64, state: &BotChatState) { - let mut data = load_bot_persistence(); - data.upsert(SavedBotConnection { + let snapshot = self.runtime_fence.persistence_snapshot(state); + let connection = SavedBotConnection { bot_type: "telegram".to_string(), chat_id: chat_id.to_string(), config: BotConfig::Telegram { bot_token: self.api.config().bot_token.clone(), }, - chat_state: state.clone(), + chat_state: snapshot, connected_at: chrono::Utc::now().timestamp(), + }; + self.runtime_fence.commit_if_current(|| { + update_bot_persistence(|data| data.upsert(connection)); + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn account_clear_is_bounded_by_epoch_when_a_command_holds_state_lock() { + let bot = TelegramBot::new(TelegramConfig { + bot_token: "test-token".to_string(), }); - save_bot_persistence(&data); + let _in_flight_command = bot.chat_states.write().await; + + tokio::time::timeout( + std::time::Duration::from_millis(500), + bot.clear_delegated_identities(), + ) + .await + .expect("account replacement must not wait indefinitely for bot network work"); } } diff --git a/src/crates/assembly/core/src/service/remote_connect/bot/weixin.rs b/src/crates/assembly/core/src/service/remote_connect/bot/weixin.rs index d6ae33781e..10d1546fe0 100644 --- a/src/crates/assembly/core/src/service/remote_connect/bot/weixin.rs +++ b/src/crates/assembly/core/src/service/remote_connect/bot/weixin.rs @@ -24,7 +24,9 @@ use super::command_router::{ parse_command, welcome_message, BotChatState, BotInteractionHandler, BotInteractiveRequest, BotMessageSender, HandleResult, }; -use super::{load_bot_persistence, save_bot_persistence, BotConfig, SavedBotConnection}; +use super::{ + load_bot_persistence, update_bot_persistence, BotConfig, BotRuntimeFence, SavedBotConnection, +}; use crate::service::remote_connect::remote_server::ImageAttachment; const LONG_POLL_TIMEOUT_SECS: u64 = 36; @@ -39,6 +41,7 @@ pub struct WeixinBot { pending_pairings: Arc>>, chat_states: Arc>>, context_tokens: Arc>>, + runtime_fence: BotRuntimeFence, } pub async fn weixin_qr_start(base_url_override: Option) -> Result { @@ -54,19 +57,54 @@ pub async fn weixin_qr_poll( impl WeixinBot { pub fn new(config: WeixinConfig) -> Self { + Self::new_fenced(config, BotRuntimeFence::standalone()) + } + + pub(crate) fn new_fenced(config: WeixinConfig, runtime_fence: BotRuntimeFence) -> Self { Self { api: Arc::new(WeixinProviderClient::new(config)), pending_pairings: Arc::new(RwLock::new(HashMap::new())), chat_states: Arc::new(RwLock::new(HashMap::new())), context_tokens: Arc::new(RwLock::new(HashMap::new())), + runtime_fence, } } - pub async fn restore_chat_state(&self, peer_id: &str, state: BotChatState) { - self.chat_states - .write() - .await - .insert(peer_id.to_string(), state); + pub async fn restore_chat_state(&self, peer_id: &str, mut state: BotChatState) { + state.prepare_for_restore(); + let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); + states.insert(peer_id.to_string(), state); + let restored = states + .get(peer_id) + .cloned() + .expect("restored Weixin state should exist"); + drop(states); + self.persist_chat_state(peer_id, &restored).await; + } + + pub async fn clear_delegated_identities(&self) { + match tokio::time::timeout( + std::time::Duration::from_millis(100), + self.chat_states.write(), + ) + .await + { + Ok(mut states) => { + self.runtime_fence.clear_states(&mut states); + let snapshots: Vec<_> = states + .iter() + .map(|(peer_id, state)| (peer_id.clone(), state.clone())) + .collect(); + drop(states); + for (peer_id, state) in snapshots { + self.persist_chat_state(&peer_id, &state).await; + } + } + Err(_) => { + warn!("Weixin account identity clear deferred behind an in-flight command"); + } + } } pub async fn register_pairing(&self, pairing_code: &str) -> Result<()> { @@ -244,9 +282,9 @@ impl WeixinBot { } async fn persist_chat_state(&self, peer_id: &str, state: &BotChatState) { - let config = self.api.config(); - let mut data = load_bot_persistence(); - data.upsert(SavedBotConnection { + let config = self.api.config().clone(); + let snapshot = self.runtime_fence.persistence_snapshot(state); + let connection = SavedBotConnection { bot_type: "weixin".to_string(), chat_id: peer_id.to_string(), config: BotConfig::Weixin { @@ -254,10 +292,12 @@ impl WeixinBot { base_url: config.base_url.clone(), bot_account_id: config.bot_account_id.clone(), }, - chat_state: state.clone(), + chat_state: snapshot, connected_at: chrono::Utc::now().timestamp(), + }; + self.runtime_fence.commit_if_current(|| { + update_bot_persistence(|data| data.upsert(connection)); }); - save_bot_persistence(&data); } pub async fn wait_for_pairing( @@ -339,11 +379,17 @@ impl WeixinBot { if self.verify_pairing_code(&text).await { info!("Weixin pairing successful peer={peer}"); let mut state = BotChatState::new(peer.clone()); + let identity_epoch = self.runtime_fence.identity_epoch(); let result = complete_im_bot_pairing(&mut state).await; - self.chat_states - .write() - .await - .insert(peer.clone(), state.clone()); + if *stop_rx.borrow() || !self.runtime_fence.is_lifecycle_current() { + return Err(anyhow!("bot lifecycle replaced during pairing")); + } + let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); + self.runtime_fence + .sanitize_after_epoch(identity_epoch, &mut state); + states.insert(peer.clone(), state.clone()); + drop(states); self.persist_chat_state(&peer, &state).await; self.send_handle_result(&peer, &result).await; @@ -483,7 +529,11 @@ impl WeixinBot { text: &str, images: Vec, ) { + if !self.runtime_fence.is_lifecycle_current() { + return; + } let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); let state = states.entry(peer_id.clone()).or_insert_with(|| { let mut state = BotChatState::new(peer_id.clone()); state.paired = true; @@ -501,9 +551,15 @@ impl WeixinBot { } if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) { if self.verify_pairing_code(trimmed).await { + let identity_epoch = self.runtime_fence.identity_epoch(); let result = complete_im_bot_pairing(state).await; + self.runtime_fence + .sanitize_after_epoch(identity_epoch, state); self.persist_chat_state(&peer_id, state).await; drop(states); + if !self.runtime_fence.is_lifecycle_current() { + return; + } self.send_handle_result(&peer_id, &result).await; return; } @@ -528,9 +584,16 @@ impl WeixinBot { let command = parse_command(text); let result = handle_command(state, command, images).await; - self.persist_chat_state(&peer_id, state).await; + self.runtime_fence.reconcile_states(&mut states); + if let Some(state) = states.get(&peer_id) { + self.persist_chat_state(&peer_id, state).await; + } drop(states); + if !self.runtime_fence.is_lifecycle_current() { + return; + } + self.send_handle_result(&peer_id, &result).await; if let Some(forward) = result.forward_to_session { @@ -580,7 +643,11 @@ impl WeixinBot { } async fn deliver_interaction(&self, peer_id: String, interaction: BotInteractiveRequest) { + if !self.runtime_fence.is_lifecycle_current() { + return; + } let mut states = self.chat_states.write().await; + self.runtime_fence.reconcile_states(&mut states); let state = states.entry(peer_id.clone()).or_insert_with(|| { let mut state = BotChatState::new(peer_id.clone()); state.paired = true; diff --git a/src/crates/assembly/core/src/service/remote_connect/mod.rs b/src/crates/assembly/core/src/service/remote_connect/mod.rs index 29eb417338..ed653aaa78 100644 --- a/src/crates/assembly/core/src/service/remote_connect/mod.rs +++ b/src/crates/assembly/core/src/service/remote_connect/mod.rs @@ -64,6 +64,7 @@ use bitfun_services_integrations::remote_connect::upload_mobile_web_to_relay; use embedded_relay_host::EmbeddedRelayHost; use log::{debug, error, info}; use serde::{Deserialize, Serialize}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use tokio::sync::{Mutex, RwLock}; @@ -137,6 +138,145 @@ struct TrustedMobileIdentity { user_id: String, } +#[derive(Clone, Debug, Eq, PartialEq)] +struct RoomConnectionOwner { + generation: u64, + room_id: String, +} + +fn room_owner_is_current( + active: &Option, + expected: &RoomConnectionOwner, +) -> bool { + active.as_ref() == Some(expected) +} + +fn clear_room_owner_if_current( + active: &mut Option, + expected: &RoomConnectionOwner, +) -> bool { + if !room_owner_is_current(active, expected) { + return false; + } + *active = None; + true +} + +/// Successful account pairing verification together with an optional host +/// lease. The service keeps the lease alive until the trusted identity and +/// paired server are committed, so an account transition cannot clear state +/// and then be overwritten by a retiring verifier. +pub struct AccountPairingVerification { + user_id: String, + _host_lease: Option>, +} + +impl std::fmt::Debug for AccountPairingVerification { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("AccountPairingVerification") + .field("user_id", &self.user_id) + .finish_non_exhaustive() + } +} + +impl AccountPairingVerification { + pub fn new(user_id: String) -> Self { + Self { + user_id, + _host_lease: None, + } + } + + pub fn with_host_lease(user_id: String, lease: L) -> Self + where + L: Send + 'static, + { + Self { + user_id, + _host_lease: Some(Box::new(lease)), + } + } + + pub fn user_id(&self) -> &str { + &self.user_id + } +} + +/// Delegated account credentials together with the host account lease that +/// authorized them. The lease is retained after the provider returns and is +/// released only after the encrypted room response has been sent. +pub struct DelegatedIdentityAuthorization { + token: String, + user_id: String, + master_key: [u8; 32], + host_lease: Option>, +} + +impl DelegatedIdentityAuthorization { + pub fn new(token: String, user_id: String, master_key: [u8; 32]) -> Self { + Self { + token, + user_id, + master_key, + host_lease: None, + } + } + + pub fn with_host_lease( + token: String, + user_id: String, + master_key: [u8; 32], + lease: L, + ) -> Self + where + L: Send + 'static, + { + Self { + token, + user_id, + master_key, + host_lease: Some(Box::new(lease)), + } + } + + fn into_response(self, local_device_id: &str) -> DelegatedIdentityResolution { + use base64::{engine::general_purpose::STANDARD as B64, Engine}; + + let Self { + token, + user_id, + master_key, + host_lease, + } = self; + DelegatedIdentityResolution { + response: remote_server::RemoteResponse::DelegateIdentity { + token, + user_id, + master_key: B64.encode(master_key), + device_id: local_device_id.to_string(), + }, + _host_lease: host_lease, + } + } +} + +struct DelegatedIdentityResolution { + response: remote_server::RemoteResponse, + _host_lease: Option>, +} + +impl DelegatedIdentityResolution { + fn error(message: impl Into) -> Self { + Self { + response: remote_server::RemoteResponse::Error { + message: message.into(), + }, + _host_lease: None, + } + } +} + /// Unified Remote Connect service that orchestrates all connection methods. pub struct RemoteConnectService { config: RemoteConnectConfig, @@ -147,8 +287,15 @@ pub struct RemoteConnectService { active_method: Arc>>, ngrok_tunnel: Arc>>, embedded_relay_host: Arc, - relay_lifecycle: Mutex<()>, + relay_lifecycle: Arc>, + room_connection_generation: AtomicU64, + active_room_owner: Arc>>, // Bot handles live independently of relay connections + bot_lifecycle: Arc>, + bot_account_identity_epoch: Arc, + bot_telegram_slot: Arc, + bot_feishu_slot: Arc, + bot_weixin_slot: Arc, bot_telegram_handle: Arc>>, bot_feishu_handle: Arc>>, bot_weixin_handle: Arc>>, @@ -164,9 +311,12 @@ pub struct RemoteConnectService { /// Account-authenticated device-routing relay client (P2). Independent from /// the room-pairing relay_client above; connects after account login. device_relay_client: Arc>>, + device_relay_lifecycle: Arc>, + device_connection_generation: AtomicU64, + active_device_connection_id: Arc>>, /// Latest online-device presence for the account (P2). online_devices: Arc>>, - /// Callback that provides a delegated identity (token + master_key) + /// Callback that provides a delegated identity and its host account lease /// for paired mobile/IM clients. Set by the desktop layer after account /// login. Resolved on demand when a paired client sends /// `get_delegated_identity` over the room channel. @@ -178,10 +328,14 @@ pub struct RemoteConnectService { account_pairing_verifier: Arc>>, } -/// Provider returning `(token, master_key, relay_url)` for the paired client. +/// Provider returning authorized delegated credentials for the paired client. type DelegatedIdentityFn = Arc< dyn Fn() -> std::pin::Pin< - Box> + Send + Sync>, + Box< + dyn std::future::Future> + + Send + + Sync, + >, > + Send + Sync, >; @@ -193,7 +347,11 @@ type AccountPairingVerifierFn = Arc< String, String, ) -> std::pin::Pin< - Box> + Send + Sync>, + Box< + dyn std::future::Future> + + Send + + Sync, + >, > + Send + Sync, >; @@ -215,7 +373,14 @@ impl RemoteConnectService { active_method: Arc::new(RwLock::new(None)), ngrok_tunnel: Arc::new(RwLock::new(None)), embedded_relay_host, - relay_lifecycle: Mutex::new(()), + relay_lifecycle: Arc::new(Mutex::new(())), + room_connection_generation: AtomicU64::new(0), + active_room_owner: Arc::new(RwLock::new(None)), + bot_lifecycle: Arc::new(Mutex::new(())), + bot_account_identity_epoch: Arc::new(AtomicU64::new(0)), + bot_telegram_slot: Arc::new(bot::BotSlotFence::default()), + bot_feishu_slot: Arc::new(bot::BotSlotFence::default()), + bot_weixin_slot: Arc::new(bot::BotSlotFence::default()), bot_telegram_handle: Arc::new(RwLock::new(None)), bot_feishu_handle: Arc::new(RwLock::new(None)), bot_weixin_handle: Arc::new(RwLock::new(None)), @@ -225,6 +390,9 @@ impl RemoteConnectService { bot_connected_info: Arc::new(RwLock::new(None)), trusted_mobile_identity: Arc::new(RwLock::new(None)), device_relay_client: Arc::new(RwLock::new(None)), + device_relay_lifecycle: Arc::new(Mutex::new(())), + device_connection_generation: AtomicU64::new(0), + active_device_connection_id: Arc::new(RwLock::new(None)), online_devices: Arc::new(RwLock::new(Vec::new())), delegated_identity_fn: Arc::new(RwLock::new(None)), account_pairing_username: Arc::new(RwLock::new(None)), @@ -233,11 +401,11 @@ impl RemoteConnectService { } /// Set the delegated identity provider (called by desktop after login). - /// Returns `(token, master_key, relay_url)` for the paired client. + /// Returns delegated credentials with a host account lease for the paired client. pub async fn set_delegated_identity_provider(&self, f: F) where F: Fn() -> Fut + Send + Sync + 'static, - Fut: std::future::Future> + Fut: std::future::Future> + Send + Sync + 'static, @@ -256,7 +424,10 @@ impl RemoteConnectService { pub async fn set_account_pairing_verifier(&self, f: F) where F: Fn(String, String) -> Fut + Send + Sync + 'static, - Fut: std::future::Future> + Send + Sync + 'static, + Fut: std::future::Future> + + Send + + Sync + + 'static, { *self.account_pairing_verifier.write().await = Some(Arc::new(move |username, password| { Box::pin(f(username, password)) @@ -275,6 +446,40 @@ impl RemoteConnectService { *self.trusted_mobile_identity.write().await = None; } + /// Clear credentials and remote-device selections cached by long-lived IM + /// bot chats. Bot connections intentionally survive account logout, but + /// their delegated authority must not. + pub async fn clear_bot_delegated_identities(&self) { + // Increment before waiting for lifecycle ownership. In-flight bot work + // can observe the epoch immediately and is forbidden from committing + // account-bound state even if a provider request does not return. + self.bot_account_identity_epoch + .fetch_add(1, Ordering::AcqRel); + let _lifecycle = self.bot_lifecycle.lock().await; + let telegram = self.telegram_bot.read().await.clone(); + let feishu = self.feishu_bot.read().await.clone(); + let weixin = self.weixin_bot.read().await.clone(); + + tokio::join!( + async move { + if let Some(bot) = telegram { + bot.clear_delegated_identities().await; + } + }, + async move { + if let Some(bot) = feishu { + bot.clear_delegated_identities().await; + } + }, + async move { + if let Some(bot) = weixin { + bot.clear_delegated_identities().await; + } + }, + ); + bitfun_services_integrations::remote_connect::bot::clear_persisted_bot_account_contexts(); + } + pub fn device_identity(&self) -> &DeviceIdentity { &self.device_identity } @@ -319,7 +524,7 @@ impl RemoteConnectService { async fn resolve_pairing_user_id( account_pairing_verifier: &Arc>>, response: &pairing::PairingResponse, - ) -> std::result::Result { + ) -> std::result::Result { let verifier = account_pairing_verifier.read().await.clone(); let Some(verify) = verifier else { // The mobile submitted account credentials (QR advertised @@ -328,7 +533,6 @@ impl RemoteConnectService { if response .password .as_deref() - .map(str::trim) .is_some_and(|value| !value.is_empty()) { return Err( @@ -342,7 +546,7 @@ impl RemoteConnectService { .map(str::trim) .filter(|value| !value.is_empty()) .ok_or_else(|| "Missing user ID".to_string())?; - return Ok(user_id.to_string()); + return Ok(AccountPairingVerification::new(user_id.to_string())); }; let username = response @@ -354,7 +558,6 @@ impl RemoteConnectService { let password = response .password .as_deref() - .map(str::trim) .filter(|value| !value.is_empty()) .ok_or_else(|| "Missing password".to_string())?; @@ -374,32 +577,31 @@ impl RemoteConnectService { delegated_identity_fn: &Arc>>, trusted_mobile_identity: &Arc>>, local_device_id: &str, - ) -> remote_server::RemoteResponse { + ) -> DelegatedIdentityResolution { let trusted_identity = trusted_mobile_identity.read().await.clone(); let Some(trusted_identity) = trusted_identity else { - return remote_server::RemoteResponse::Error { - message: "Pairing authorization expired; scan a new QR code".to_string(), - }; + return DelegatedIdentityResolution::error( + "Pairing authorization expired; scan a new QR code", + ); }; let provider = delegated_identity_fn.read().await.clone(); let Some(get_identity) = provider else { - return remote_server::RemoteResponse::Error { - message: "Desktop is not logged into a BitFun account".to_string(), - }; + return DelegatedIdentityResolution::error( + "Desktop is not logged into a BitFun account", + ); }; - let Some((token, master_key, _relay_url)) = get_identity().await else { - return remote_server::RemoteResponse::Error { - message: "Desktop is not logged into a BitFun account".to_string(), - }; + let Some(authorization) = get_identity().await else { + return DelegatedIdentityResolution::error( + "Desktop is not logged into a BitFun account", + ); }; - use base64::{engine::general_purpose::STANDARD as B64, Engine}; - info!("Delegated identity resolved for paired client"); - remote_server::RemoteResponse::DelegateIdentity { - token, - user_id: trusted_identity.user_id, - master_key: B64.encode(master_key), - device_id: local_device_id.to_string(), + if authorization.user_id != trusted_identity.user_id { + return DelegatedIdentityResolution::error( + "Paired mobile identity no longer matches the desktop account", + ); } + info!("Delegated identity resolved for paired client"); + authorization.into_response(local_device_id) } async fn send_pairing_error_response( @@ -650,16 +852,35 @@ impl RemoteConnectService { *self.active_method.write().await = Some(method.clone()); *self.relay_client.write().await = Some(client); + let room_owner = RoomConnectionOwner { + generation: self + .room_connection_generation + .fetch_add(1, Ordering::AcqRel) + + 1, + room_id: qr_payload.room_id.clone(), + }; + *self.active_room_owner.write().await = Some(room_owner.clone()); let pairing_arc = self.pairing.clone(); let relay_arc = self.relay_client.clone(); let server_arc = self.remote_server.clone(); + let active_method_arc = self.active_method.clone(); + let room_lifecycle = self.relay_lifecycle.clone(); + let active_room_owner = self.active_room_owner.clone(); let trusted_mobile_identity_arc = self.trusted_mobile_identity.clone(); let delegated_identity_fn_arc = self.delegated_identity_fn.clone(); let account_pairing_verifier_arc = self.account_pairing_verifier.clone(); let local_device_id = self.device_identity.device_id.clone(); tokio::spawn(async move { while let Some(event) = event_rx.recv().await { + // Lease only this event's effects. `start`/`stop_relay` can + // acquire the lifecycle mutex between events, and a retiring + // loop observes the owner mismatch before touching shared + // pairing/client/server state. + let _room_effect = room_lifecycle.lock().await; + if !room_owner_is_current(&*active_room_owner.read().await, &room_owner) { + break; + } match event { relay_client::RelayEvent::PairRequest { correlation_id, @@ -703,7 +924,7 @@ impl RemoteConnectService { Ok((cmd, request_id)) => { handled_as_active_command = true; debug!("Remote command decrypted"); - let response = if matches!( + let response_resolution = if matches!( cmd, remote_server::RemoteCommand::GetDelegatedIdentity ) { @@ -714,10 +935,16 @@ impl RemoteConnectService { ) .await } else { - server.dispatch(&cmd).await + DelegatedIdentityResolution { + response: server.dispatch(&cmd).await, + _host_lease: None, + } }; match server - .encrypt_response(&response, request_id.as_deref()) + .encrypt_response( + &response_resolution.response, + request_id.as_deref(), + ) { Ok((enc, resp_nonce)) => { if let Some(ref client) = *relay_arc.read().await { @@ -734,6 +961,10 @@ impl RemoteConnectService { error!("Failed to encrypt response: {e}"); } } + // `response_resolution` owns the account lease for + // delegated credentials. Keep it alive through both + // encryption and the awaited room send above. + drop(response_resolution); } Err(e) => { debug!( @@ -769,7 +1000,7 @@ impl RemoteConnectService { .await; continue; } - let canonical_user_id = match RemoteConnectService::resolve_pairing_user_id( + let account_verification = match RemoteConnectService::resolve_pairing_user_id( &account_pairing_verifier_arc, &response, ) @@ -793,6 +1024,8 @@ impl RemoteConnectService { continue; } }; + let canonical_user_id = + account_verification.user_id().to_string(); let mobile_install_id = response .mobile_install_id .clone() @@ -869,6 +1102,9 @@ impl RemoteConnectService { // frame would be dropped. Paired clients // pull it via `get_delegated_identity`. } + // Keep the host's account lease through every + // successful pairing commit and response-side effect. + drop(account_verification); } Ok(false) => { error!("Pairing verification failed"); @@ -914,6 +1150,19 @@ impl RemoteConnectService { _ => {} } } + + // Stream exit is itself generation-fenced. In particular, an old + // loop closing after reconnect must not disconnect the new room. + let _room_effect = room_lifecycle.lock().await; + let mut owner = active_room_owner.write().await; + if clear_room_owner_if_current(&mut *owner, &room_owner) { + drop(owner); + *relay_arc.write().await = None; + pairing_arc.write().await.disconnect().await; + *server_arc.write().await = None; + *active_method_arc.write().await = None; + *trusted_mobile_identity_arc.write().await = None; + } }); let state = pairing.state().await; @@ -936,6 +1185,7 @@ impl RemoteConnectService { } async fn start_bot_connection(&self, method: &ConnectionMethod) -> Result { + let _lifecycle = self.bot_lifecycle.lock().await; let pairing_code = PairingProtocol::generate_bot_pairing_code(); let bot_link = match method { @@ -943,14 +1193,20 @@ impl RemoteConnectService { match &self.config.bot_telegram { Some(bot::BotConfig::Telegram { bot_token }) if !bot_token.is_empty() => { // Stop any existing Telegram bot + let generation = self.bot_telegram_slot.advance(); if let Some(handle) = self.bot_telegram_handle.write().await.take() { handle.stop(); } - let tg_bot = Arc::new(bot::telegram::TelegramBot::new( + let tg_bot = Arc::new(bot::telegram::TelegramBot::new_fenced( bot::telegram::TelegramConfig { bot_token: bot_token.clone(), }, + bot::BotRuntimeFence::new( + self.bot_account_identity_epoch.clone(), + self.bot_telegram_slot.clone(), + generation, + ), )); tg_bot.register_pairing(&pairing_code).await?; @@ -960,6 +1216,8 @@ impl RemoteConnectService { let bot_for_pair = tg_bot.clone(); let bot_for_loop = tg_bot.clone(); let tg_bot_ref = self.telegram_bot.clone(); + let bot_lifecycle = self.bot_lifecycle.clone(); + let bot_slot = self.bot_telegram_slot.clone(); *tg_bot_ref.write().await = Some(tg_bot.clone()); @@ -967,12 +1225,14 @@ impl RemoteConnectService { let mut stop_rx = stop_rx; match bot_for_pair.wait_for_pairing(&mut stop_rx).await { Ok(chat_id) => { + let lifecycle = bot_lifecycle.lock().await; // Guard against the race where stop_bots() cleared // bot_connected_info between pairing completing and // this task running. - if !*stop_rx.borrow() { + if !*stop_rx.borrow() && bot_slot.is_current(generation) { *bot_connected_info.write().await = Some(format!("Telegram({chat_id})")); + drop(lifecycle); info!("Telegram bot paired, starting message loop"); bot_for_loop.run_message_loop(stop_rx).await; } else { @@ -1001,15 +1261,22 @@ impl RemoteConnectService { Some(bot::BotConfig::Feishu { app_id, app_secret }) if !app_id.is_empty() && !app_secret.is_empty() => { + let generation = self.bot_feishu_slot.advance(); if let Some(handle) = self.bot_feishu_handle.write().await.take() { handle.stop(); } - let fs_bot = - Arc::new(bot::feishu::FeishuBot::new(bot::feishu::FeishuConfig { + let fs_bot = Arc::new(bot::feishu::FeishuBot::new_fenced( + bot::feishu::FeishuConfig { app_id: app_id.clone(), app_secret: app_secret.clone(), - })); + }, + bot::BotRuntimeFence::new( + self.bot_account_identity_epoch.clone(), + self.bot_feishu_slot.clone(), + generation, + ), + )); fs_bot.register_pairing(&pairing_code).await?; let (stop_tx, stop_rx) = tokio::sync::watch::channel(false); @@ -1018,6 +1285,8 @@ impl RemoteConnectService { let bot_for_pair = fs_bot.clone(); let bot_for_loop = fs_bot.clone(); let fs_bot_ref = self.feishu_bot.clone(); + let bot_lifecycle = self.bot_lifecycle.clone(); + let bot_slot = self.bot_feishu_slot.clone(); *fs_bot_ref.write().await = Some(fs_bot.clone()); @@ -1025,12 +1294,14 @@ impl RemoteConnectService { let mut stop_rx = stop_rx; match bot_for_pair.wait_for_pairing(&mut stop_rx).await { Ok(chat_id) => { + let lifecycle = bot_lifecycle.lock().await; // Guard against the race where stop_bots() cleared // bot_connected_info between pairing completing and // this task running. - if !*stop_rx.borrow() { + if !*stop_rx.borrow() && bot_slot.is_current(generation) { *bot_connected_info.write().await = Some(format!("Feishu({chat_id})")); + drop(lifecycle); info!("Feishu bot paired, starting message loop"); bot_for_loop.run_message_loop(stop_rx).await; } else { @@ -1062,6 +1333,7 @@ impl RemoteConnectService { base_url, bot_account_id, }) if !ilink_token.is_empty() && !bot_account_id.is_empty() => { + let generation = self.bot_weixin_slot.advance(); if let Some(handle) = self.bot_weixin_handle.write().await.take() { handle.stop(); } @@ -1076,7 +1348,14 @@ impl RemoteConnectService { bot_account_id: bot_account_id.clone(), }; - let wx_bot = Arc::new(bot::weixin::WeixinBot::new(wx_cfg)); + let wx_bot = Arc::new(bot::weixin::WeixinBot::new_fenced( + wx_cfg, + bot::BotRuntimeFence::new( + self.bot_account_identity_epoch.clone(), + self.bot_weixin_slot.clone(), + generation, + ), + )); wx_bot.register_pairing(&pairing_code).await?; let (stop_tx, stop_rx) = tokio::sync::watch::channel(false); @@ -1085,6 +1364,8 @@ impl RemoteConnectService { let bot_for_pair = wx_bot.clone(); let bot_for_loop = wx_bot.clone(); let wx_bot_ref = self.weixin_bot.clone(); + let bot_lifecycle = self.bot_lifecycle.clone(); + let bot_slot = self.bot_weixin_slot.clone(); *wx_bot_ref.write().await = Some(wx_bot.clone()); @@ -1092,9 +1373,11 @@ impl RemoteConnectService { let mut stop_rx = stop_rx; match bot_for_pair.wait_for_pairing(&mut stop_rx).await { Ok(peer_id) => { - if !*stop_rx.borrow() { + let lifecycle = bot_lifecycle.lock().await; + if !*stop_rx.borrow() && bot_slot.is_current(generation) { *bot_connected_info.write().await = Some(format!("Weixin({peer_id})")); + drop(lifecycle); info!("Weixin bot paired, starting message loop"); bot_for_loop.run_message_loop(stop_rx).await; } else { @@ -1139,16 +1422,23 @@ impl RemoteConnectService { /// Restore a previously paired bot from persistence. /// Skips the pairing step and directly starts the message loop. pub async fn restore_bot(&self, saved: &bot::SavedBotConnection) -> Result<()> { + let _lifecycle = self.bot_lifecycle.lock().await; match saved.config { bot::BotConfig::Telegram { ref bot_token } => { + let generation = self.bot_telegram_slot.advance(); if let Some(handle) = self.bot_telegram_handle.write().await.take() { handle.stop(); } - let tg_bot = Arc::new(bot::telegram::TelegramBot::new( + let tg_bot = Arc::new(bot::telegram::TelegramBot::new_fenced( bot::telegram::TelegramConfig { bot_token: bot_token.clone(), }, + bot::BotRuntimeFence::new( + self.bot_account_identity_epoch.clone(), + self.bot_telegram_slot.clone(), + generation, + ), )); let chat_id: i64 = saved.chat_id.parse().map_err(|_| { @@ -1175,14 +1465,22 @@ impl RemoteConnectService { ref app_id, ref app_secret, } => { + let generation = self.bot_feishu_slot.advance(); if let Some(handle) = self.bot_feishu_handle.write().await.take() { handle.stop(); } - let fs_bot = Arc::new(bot::feishu::FeishuBot::new(bot::feishu::FeishuConfig { - app_id: app_id.clone(), - app_secret: app_secret.clone(), - })); + let fs_bot = Arc::new(bot::feishu::FeishuBot::new_fenced( + bot::feishu::FeishuConfig { + app_id: app_id.clone(), + app_secret: app_secret.clone(), + }, + bot::BotRuntimeFence::new( + self.bot_account_identity_epoch.clone(), + self.bot_feishu_slot.clone(), + generation, + ), + )); fs_bot .restore_chat_state(&saved.chat_id, saved.chat_state.clone()) @@ -1208,6 +1506,7 @@ impl RemoteConnectService { ref base_url, ref bot_account_id, } => { + let generation = self.bot_weixin_slot.advance(); if let Some(handle) = self.bot_weixin_handle.write().await.take() { handle.stop(); } @@ -1222,7 +1521,14 @@ impl RemoteConnectService { bot_account_id: bot_account_id.clone(), }; - let wx_bot = Arc::new(bot::weixin::WeixinBot::new(wx_cfg)); + let wx_bot = Arc::new(bot::weixin::WeixinBot::new_fenced( + wx_cfg, + bot::BotRuntimeFence::new( + self.bot_account_identity_epoch.clone(), + self.bot_weixin_slot.clone(), + generation, + ), + )); wx_bot .restore_chat_state(&saved.chat_id, saved.chat_state.clone()) .await; @@ -1258,6 +1564,9 @@ impl RemoteConnectService { } async fn stop_relay_inner(&self) { + // Fence the retiring event loop before disconnecting its client. A + // late event or stream-exit cleanup must not mutate the next room. + *self.active_room_owner.write().await = None; if let Some(ref client) = *self.relay_client.read().await { client.disconnect().await; } @@ -1279,16 +1588,20 @@ impl RemoteConnectService { /// Stop all bot connections. pub async fn stop_bots(&self) { + let _lifecycle = self.bot_lifecycle.lock().await; + self.bot_telegram_slot.advance(); if let Some(handle) = self.bot_telegram_handle.write().await.take() { handle.stop(); } *self.telegram_bot.write().await = None; + self.bot_feishu_slot.advance(); if let Some(handle) = self.bot_feishu_handle.write().await.take() { handle.stop(); } *self.feishu_bot.write().await = None; + self.bot_weixin_slot.advance(); if let Some(handle) = self.bot_weixin_handle.write().await.take() { handle.stop(); } @@ -1371,7 +1684,8 @@ impl RemoteConnectService { ) } - /// Start account device routing. Returns `(event_rx, authenticated_device_id)`. + /// Start account device routing. Returns + /// `(event_rx, authenticated_device_id, connection_id)`. /// /// `AuthOk` is consumed here (not forwarded) so callers must use the returned /// `authenticated_device_id` — and this method adopts it into the persisted @@ -1384,9 +1698,11 @@ impl RemoteConnectService { ) -> Result<( tokio::sync::mpsc::UnboundedReceiver, String, + u64, )> { + let _lifecycle = self.device_relay_lifecycle.lock().await; // Disconnect previous device connection if any. - self.stop_device_connection().await; + self.stop_device_connection_inner().await; let ws_url = format!( "{}/ws", @@ -1445,11 +1761,23 @@ impl RemoteConnectService { let online_arc = self.online_devices.clone(); let device_client_arc = self.device_relay_client.clone(); + let device_lifecycle = self.device_relay_lifecycle.clone(); + let active_connection_id = self.active_device_connection_id.clone(); + let connection_id = self + .device_connection_generation + .fetch_add(1, Ordering::AcqRel) + + 1; + *device_client_arc.write().await = Some(client); + *active_connection_id.write().await = Some(connection_id); // Spawn event forwarder that updates presence state; the raw event stream // is also forwarded to a new channel for the caller to consume. let (forward_tx, forward_rx) = tokio::sync::mpsc::unbounded_channel(); tokio::spawn(async move { while let Some(event) = event_rx.recv().await { + let _effect = device_lifecycle.lock().await; + if *active_connection_id.read().await != Some(connection_id) { + break; + } match &event { relay_client::RelayEvent::DevicePresence { devices } => { *online_arc.write().await = devices.clone(); @@ -1461,14 +1789,27 @@ impl RemoteConnectService { } let _ = forward_tx.send(event); } + let _effect = device_lifecycle.lock().await; + let mut active = active_connection_id.write().await; + if *active == Some(connection_id) { + *active = None; + drop(active); + *device_client_arc.write().await = None; + online_arc.write().await.clear(); + } }); - *device_client_arc.write().await = Some(client); - Ok((forward_rx, authenticated_device_id)) + Ok((forward_rx, authenticated_device_id, connection_id)) } /// Disconnect the account-authenticated device-routing connection. pub async fn stop_device_connection(&self) { + let _lifecycle = self.device_relay_lifecycle.lock().await; + self.stop_device_connection_inner().await; + } + + async fn stop_device_connection_inner(&self) { + *self.active_device_connection_id.write().await = None; if let Some(client) = self.device_relay_client.write().await.take() { client.disconnect().await; } @@ -1492,6 +1833,31 @@ impl RemoteConnectService { .await } + /// Send through the exact device-routing client captured by the caller. + /// The lifecycle lease prevents connection replacement between the owner + /// check and the transport write. + pub async fn send_device_message_if_connection( + &self, + expected_connection_id: u64, + target_device_id: &str, + correlation_id: &str, + encrypted_data: &str, + nonce: &str, + ) -> Result { + let _lifecycle = self.device_relay_lifecycle.lock().await; + if *self.active_device_connection_id.read().await != Some(expected_connection_id) { + return Ok(false); + } + let guard = self.device_relay_client.read().await; + let client = guard + .as_ref() + .ok_or_else(|| anyhow::anyhow!("device routing not connected"))?; + client + .send_device_message(target_device_id, correlation_id, encrypted_data, nonce) + .await?; + Ok(true) + } + /// Current online devices in the account (presence list). pub async fn online_devices(&self) -> Vec { self.online_devices.read().await.clone() @@ -1514,6 +1880,68 @@ impl RemoteConnectService { .await } + /// Send only if the relay connection still belongs to the pairing secret + /// captured by the caller. Holding the relay lifecycle and pairing read + /// leases across the send prevents a concurrently replaced room from + /// receiving a response prepared for the previous room. + pub async fn send_room_response_if_pairing_secret( + &self, + expected_secret: &[u8; 32], + correlation_id: &str, + encrypted_data: &str, + nonce: &str, + ) -> Result { + let _lifecycle = self.relay_lifecycle.lock().await; + let pairing = self.pairing.read().await; + if pairing.shared_secret() != Some(expected_secret) { + return Ok(false); + } + let guard = self.relay_client.read().await; + let client = guard + .as_ref() + .ok_or_else(|| anyhow::anyhow!("relay client not connected"))?; + client + .send_relay_response(correlation_id, encrypted_data, nonce) + .await?; + Ok(true) + } + + /// Room-first variant for host-authorized secret-bearing responses. The + /// returned authorization lease stays alive through the transport write; + /// hosts can use it to keep account replacement from completing without + /// introducing an account-lock -> room-lock inversion. + pub async fn send_room_response_if_pairing_secret_authorized( + &self, + expected_secret: &[u8; 32], + correlation_id: &str, + encrypted_data: &str, + nonce: &str, + authorize: F, + ) -> Result + where + F: FnOnce() -> Fut, + Fut: std::future::Future>, + L: Send, + { + let _lifecycle = self.relay_lifecycle.lock().await; + let pairing = self.pairing.read().await; + if pairing.shared_secret() != Some(expected_secret) { + return Ok(false); + } + let _authorization = authorize().await.map_err(anyhow::Error::msg)?; + if pairing.shared_secret() != Some(expected_secret) { + return Ok(false); + } + let guard = self.relay_client.read().await; + let client = guard + .as_ref() + .ok_or_else(|| anyhow::anyhow!("relay client not connected"))?; + client + .send_relay_response(correlation_id, encrypted_data, nonce) + .await?; + Ok(true) + } + /// Get the pairing shared secret (for encrypting delegate identity). pub async fn pairing_shared_secret(&self) -> Option<[u8; 32]> { self.pairing.read().await.shared_secret().copied() @@ -1540,14 +1968,41 @@ mod tests { Arc::new(RwLock::new(None)) } + #[test] + fn retiring_room_owner_cannot_clear_replacement() { + let owner_a = RoomConnectionOwner { + generation: 1, + room_id: "room-a".to_string(), + }; + let owner_b = RoomConnectionOwner { + generation: 2, + room_id: "room-b".to_string(), + }; + let mut active = Some(owner_b.clone()); + + assert!(!clear_room_owner_if_current(&mut active, &owner_a)); + assert_eq!(active, Some(owner_b.clone())); + assert!(clear_room_owner_if_current(&mut active, &owner_b)); + assert_eq!(active, None); + } + fn verifier_returning( canonical_user_id: &'static str, ) -> Arc>> { Arc::new(RwLock::new(Some( Arc::new(move |_username: String, _password: String| { - Box::pin(async move { Ok(canonical_user_id.to_string()) }) + Box::pin(async move { + Ok(AccountPairingVerification::new( + canonical_user_id.to_string(), + )) + }) as std::pin::Pin< - Box> + Send + Sync>, + Box< + dyn std::future::Future< + Output = Result, + > + Send + + Sync, + >, > }) as AccountPairingVerifierFn, ))) @@ -1592,7 +2047,7 @@ mod tests { &pairing_response(Some("local-user"), None), ) .await; - assert_eq!(result.unwrap(), "local-user"); + assert_eq!(result.unwrap().user_id(), "local-user"); } #[tokio::test] @@ -1603,22 +2058,28 @@ mod tests { ) .await .expect("verification should succeed"); - assert_eq!(canonical, "canonical-user-123"); + assert_eq!(canonical.user_id(), "canonical-user-123"); + let canonical_user_id = canonical.user_id().to_string(); let trusted = Arc::new(RwLock::new(None)); - let identity = - RemoteConnectService::validate_mobile_identity(&trusted, "install-1", &canonical) - .await - .expect("first pairing binds the identity"); + let identity = RemoteConnectService::validate_mobile_identity( + &trusted, + "install-1", + &canonical_user_id, + ) + .await + .expect("first pairing binds the identity"); assert_eq!(identity.user_id, "canonical-user-123"); RemoteConnectService::persist_mobile_identity(&trusted, identity).await; // Reconnect with the same canonical id still matches. - assert!( - RemoteConnectService::validate_mobile_identity(&trusted, "install-1", &canonical) - .await - .is_ok() - ); + assert!(RemoteConnectService::validate_mobile_identity( + &trusted, + "install-1", + &canonical_user_id, + ) + .await + .is_ok()); // A different account user id is rejected against the bound identity. assert!(RemoteConnectService::validate_mobile_identity( &trusted, @@ -1629,6 +2090,35 @@ mod tests { .is_err()); } + #[tokio::test] + async fn account_pairing_preserves_password_whitespace_for_verification() { + let verifier = Arc::new(RwLock::new(Some( + Arc::new(|_username: String, password: String| { + Box::pin(async move { + assert_eq!(password, " secret with spaces "); + Ok(AccountPairingVerification::new( + "canonical-user-123".to_string(), + )) + }) + as std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, + > + Send + + Sync, + >, + > + }) as AccountPairingVerifierFn, + ))); + let verified = RemoteConnectService::resolve_pairing_user_id( + &verifier, + &pairing_response(Some("alice"), Some(" secret with spaces ")), + ) + .await + .expect("pairing should pass the exact password to the verifier"); + assert_eq!(verified.user_id(), "canonical-user-123"); + } + #[tokio::test] async fn delegated_identity_requires_a_trusted_pairing_before_minting_credentials() { let provider_called = Arc::new(AtomicBool::new(false)); @@ -1637,7 +2127,11 @@ mod tests { let called = called.clone(); Box::pin(async move { called.store(true, Ordering::SeqCst); - Some(("account-token".to_string(), [7_u8; 32], "relay".to_string())) + Some(DelegatedIdentityAuthorization::new( + "account-token".to_string(), + "account-user".to_string(), + [7_u8; 32], + )) }) }); let provider = Arc::new(RwLock::new(Some(provider))); @@ -1651,7 +2145,7 @@ mod tests { .await; assert!(matches!( - response, + response.response, remote_server::RemoteResponse::Error { .. } )); assert!(!provider_called.load(Ordering::SeqCst)); @@ -1660,7 +2154,13 @@ mod tests { #[tokio::test] async fn delegated_identity_uses_the_account_bound_during_pairing() { let provider: DelegatedIdentityFn = Arc::new(|| { - Box::pin(async { Some(("account-token".to_string(), [7_u8; 32], "relay".to_string())) }) + Box::pin(async { + Some(DelegatedIdentityAuthorization::new( + "account-token".to_string(), + "paired-user".to_string(), + [7_u8; 32], + )) + }) }); let provider = Arc::new(RwLock::new(Some(provider))); let trusted = Arc::new(RwLock::new(Some(TrustedMobileIdentity { @@ -1676,7 +2176,7 @@ mod tests { .await; assert!(matches!( - response, + response.response, remote_server::RemoteResponse::DelegateIdentity { user_id, device_id, @@ -1684,4 +2184,85 @@ mod tests { } if user_id == "paired-user" && device_id == "desktop-1" )); } + + #[tokio::test] + async fn delegated_identity_rejects_a_provider_for_another_account() { + let provider: DelegatedIdentityFn = Arc::new(|| { + Box::pin(async { + Some(DelegatedIdentityAuthorization::new( + "account-token".to_string(), + "replacement-user".to_string(), + [7_u8; 32], + )) + }) + }); + let provider = Arc::new(RwLock::new(Some(provider))); + let trusted = Arc::new(RwLock::new(Some(TrustedMobileIdentity { + mobile_install_id: "install-1".to_string(), + user_id: "paired-user".to_string(), + }))); + + let response = RemoteConnectService::resolve_delegated_identity_response( + &provider, + &trusted, + "desktop-1", + ) + .await; + assert!(matches!( + response.response, + remote_server::RemoteResponse::Error { message } + if message.contains("no longer matches") + )); + } + + #[tokio::test] + async fn delegated_identity_keeps_account_lease_until_response_is_released() { + let account_lifecycle = Arc::new(Mutex::new(())); + let provider_lifecycle = account_lifecycle.clone(); + let provider: DelegatedIdentityFn = Arc::new(move || { + let provider_lifecycle = provider_lifecycle.clone(); + Box::pin(async move { + let lease = provider_lifecycle.lock_owned().await; + Some(DelegatedIdentityAuthorization::with_host_lease( + "account-token".to_string(), + "paired-user".to_string(), + [7_u8; 32], + lease, + )) + }) + }); + let provider = Arc::new(RwLock::new(Some(provider))); + let trusted = Arc::new(RwLock::new(Some(TrustedMobileIdentity { + mobile_install_id: "install-1".to_string(), + user_id: "paired-user".to_string(), + }))); + + let response = RemoteConnectService::resolve_delegated_identity_response( + &provider, + &trusted, + "desktop-1", + ) + .await; + assert!(matches!( + &response.response, + remote_server::RemoteResponse::DelegateIdentity { .. } + )); + assert!( + tokio::time::timeout( + std::time::Duration::from_millis(20), + account_lifecycle.clone().lock_owned(), + ) + .await + .is_err(), + "account replacement must remain blocked while the response is in flight" + ); + + drop(response); + tokio::time::timeout( + std::time::Duration::from_secs(1), + account_lifecycle.lock_owned(), + ) + .await + .expect("account replacement should proceed after the response is released"); + } } diff --git a/src/crates/assembly/core/src/service/remote_connect/settings_sync.rs b/src/crates/assembly/core/src/service/remote_connect/settings_sync.rs index c111b9c72a..22450c9669 100644 --- a/src/crates/assembly/core/src/service/remote_connect/settings_sync.rs +++ b/src/crates/assembly/core/src/service/remote_connect/settings_sync.rs @@ -26,7 +26,7 @@ use std::time::Duration; use anyhow::{anyhow, Result}; use log::{debug, warn}; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, Notify}; use bitfun_services_integrations::remote_connect::account::{ error_indicates_expired_token, AccountClient, AccountSession, SettingsBlob, @@ -40,7 +40,7 @@ pub const SETTINGS_PUSH_DEBOUNCE: Duration = Duration::from_secs(5); /// Account context needed for every relay call: the session (token + /// master_key) and the relay base URL. -pub type AccountContext = (AccountSession, String); +pub type AccountContext = (AccountSession, String, u64); type AccountContextFn = dyn Fn() -> std::pin::Pin> + Send>> + Send @@ -54,6 +54,10 @@ pub struct SettingsSyncHooks { /// logged out. Required for the background loop; one-shot helpers take /// the context explicitly. pub account_context: Option>, + /// Confirms that a context generation captured before an async relay call + /// still belongs to the active account. Hosts bump the generation before + /// logout or replacement login. + pub is_account_context_current: Option bool + Send + Sync>>, /// When true, push and pull are paused (Desktop: Peer controller mode). pub should_pause: Option bool + Send + Sync>>, /// Fired after cloud settings were applied to the local config. @@ -73,6 +77,7 @@ static PUSH_TX: OnceLock> = OnceLock::new(); /// user-chosen direction. A counter (not a flag) so concurrent ops do not /// clear each other's in-flight state on completion. static SYNC_OPS_IN_FLIGHT: AtomicUsize = AtomicUsize::new(0); +static SYNC_OPS_IDLE: Notify = Notify::const_new(); /// RAII guard that marks a sync upload/apply as in flight. struct SyncOpGuard; @@ -84,13 +89,29 @@ impl SyncOpGuard { } impl Drop for SyncOpGuard { fn drop(&mut self) { - SYNC_OPS_IN_FLIGHT.fetch_sub(1, Ordering::SeqCst); + if SYNC_OPS_IN_FLIGHT.fetch_sub(1, Ordering::SeqCst) == 1 { + SYNC_OPS_IDLE.notify_waiters(); + } + } +} + +/// Wait until every settings upload/apply critical section has completed. +/// Hosts call this after invalidating their account generation and before +/// completing logout or installing a replacement account. +pub async fn wait_for_sync_operations_idle() { + loop { + let notified = SYNC_OPS_IDLE.notified(); + if SYNC_OPS_IN_FLIGHT.load(Ordering::SeqCst) == 0 { + return; + } + notified.await; } } fn hooks() -> &'static SettingsSyncHooks { static DEFAULT: SettingsSyncHooks = SettingsSyncHooks { account_context: None, + is_account_context_current: None, should_pause: None, on_settings_applied: None, on_settings_pushed: None, @@ -99,6 +120,14 @@ fn hooks() -> &'static SettingsSyncHooks { HOOKS.get().unwrap_or(&DEFAULT) } +fn is_account_context_current(generation: u64) -> bool { + hooks() + .is_account_context_current + .as_ref() + .map(|check| check(generation)) + .unwrap_or(true) +} + fn should_pause() -> bool { hooks().should_pause.as_ref().map(|f| f()).unwrap_or(false) } @@ -183,11 +212,29 @@ pub async fn upload_settings_payload( relay_url: &str, payload: &str, ) -> Result { + upload_settings_payload_for_generation(account, relay_url, payload, None).await +} + +async fn upload_settings_payload_for_generation( + account: &AccountSession, + relay_url: &str, + payload: &str, + generation: Option, +) -> Result { + if generation.is_some_and(|value| !is_account_context_current(value)) { + return Err(anyhow!("account context changed before settings upload")); + } let _op = SyncOpGuard::begin(); + if generation.is_some_and(|value| !is_account_context_current(value)) { + return Err(anyhow!("account context changed before settings upload")); + } let client = AccountClient::new(); // The version returned here is the exact one stored on the relay — // recording it keeps the next pull from re-applying our own upload. let version = client.upload_settings(relay_url, account, payload).await?; + if generation.is_some_and(|value| !is_account_context_current(value)) { + return Err(anyhow!("account context changed during settings upload")); + } let hash = settings_content_hash(payload).unwrap_or_default(); record_settings_cursor(&account.user_id, version, hash); debug!("Settings sync: uploaded settings (version={version})"); @@ -198,6 +245,14 @@ pub async fn upload_settings_payload( /// Export the current config and upload it when the content differs from the /// last uploaded/applied blob. Returns `true` when an upload happened. pub async fn push_settings_now(account: &AccountSession, relay_url: &str) -> Result { + push_settings_now_for_generation(account, relay_url, None).await +} + +async fn push_settings_now_for_generation( + account: &AccountSession, + relay_url: &str, + generation: Option, +) -> Result { let config_service = crate::service::config::get_global_config_service() .await .map_err(|e| anyhow!("config service: {e}"))?; @@ -207,6 +262,10 @@ pub async fn push_settings_now(account: &AccountSession, relay_url: &str) -> Res .map_err(|e| anyhow!("export config: {e}"))?; let payload = serde_json::to_string(&exported).map_err(|e| anyhow!("serialize config: {e}"))?; + if generation.is_some_and(|value| !is_account_context_current(value)) { + return Err(anyhow!("account context changed while exporting settings")); + } + let hash = settings_content_hash(&payload)?; let known = sync_state::load_settings_cursor(&account.user_id); if known.hash == hash && known.version != 0 { @@ -214,7 +273,7 @@ pub async fn push_settings_now(account: &AccountSession, relay_url: &str) -> Res return Ok(false); } - upload_settings_payload(account, relay_url, &payload).await?; + upload_settings_payload_for_generation(account, relay_url, &payload, generation).await?; Ok(true) } @@ -226,6 +285,15 @@ pub async fn apply_settings_blob( account: &AccountSession, blob: &SettingsBlob, force: bool, +) -> Result { + apply_settings_blob_for_generation(account, blob, force, None).await +} + +async fn apply_settings_blob_for_generation( + account: &AccountSession, + blob: &SettingsBlob, + force: bool, + generation: Option, ) -> Result { if !force { let known = sync_state::load_settings_cursor(&account.user_id); @@ -233,7 +301,13 @@ pub async fn apply_settings_blob( return Ok(false); } } + if generation.is_some_and(|value| !is_account_context_current(value)) { + return Err(anyhow!("account context changed before settings apply")); + } let _op = SyncOpGuard::begin(); + if generation.is_some_and(|value| !is_account_context_current(value)) { + return Err(anyhow!("account context changed before settings apply")); + } let inner_config = inner_config_value(&blob.plaintext)?; let config_service = crate::service::config::get_global_config_service() @@ -273,6 +347,14 @@ pub async fn apply_settings_blob( /// when new settings were applied; `Ok(false)` also when no cloud settings /// exist yet. pub async fn pull_and_apply_settings(account: &AccountSession, relay_url: &str) -> Result { + pull_and_apply_settings_for_generation(account, relay_url, None).await +} + +async fn pull_and_apply_settings_for_generation( + account: &AccountSession, + relay_url: &str, + generation: Option, +) -> Result { let client = AccountClient::new(); let Some(blob) = client .fetch_settings_with_version(relay_url, account) @@ -280,7 +362,10 @@ pub async fn pull_and_apply_settings(account: &AccountSession, relay_url: &str) else { return Ok(false); }; - apply_settings_blob(account, &blob, false).await + if generation.is_some_and(|value| !is_account_context_current(value)) { + return Err(anyhow!("account context changed during settings pull")); + } + apply_settings_blob_for_generation(account, &blob, false, generation).await } async fn account_context() -> Result { @@ -296,11 +381,17 @@ async fn push_from_loop() { debug!("Settings sync: push paused by host app"); return; } - let (account, relay_url) = match account_context().await { + let (account, relay_url, generation) = match account_context().await { Ok(ctx) => ctx, Err(_) => return, // logged out — silently skip }; - if let Err(e) = push_settings_now(&account, &relay_url).await { + if !is_account_context_current(generation) { + return; + } + if let Err(e) = push_settings_now_for_generation(&account, &relay_url, Some(generation)).await { + if !is_account_context_current(generation) { + return; + } note_relay_error(&e, "push"); } } @@ -314,11 +405,19 @@ async fn pull_from_loop() { debug!("Settings sync: pull skipped while an upload/apply is in flight"); return; } - let (account, relay_url) = match account_context().await { + let (account, relay_url, generation) = match account_context().await { Ok(ctx) => ctx, Err(_) => return, // logged out — silently skip }; - if let Err(e) = pull_and_apply_settings(&account, &relay_url).await { + if !is_account_context_current(generation) { + return; + } + if let Err(e) = + pull_and_apply_settings_for_generation(&account, &relay_url, Some(generation)).await + { + if !is_account_context_current(generation) { + return; + } note_relay_error(&e, "pull"); } } diff --git a/src/crates/services/page-function-runtime/src/lib.rs b/src/crates/services/page-function-runtime/src/lib.rs index 4101f019bb..2eab200978 100644 --- a/src/crates/services/page-function-runtime/src/lib.rs +++ b/src/crates/services/page-function-runtime/src/lib.rs @@ -7,6 +7,7 @@ //! host bindings are synchronous. use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -97,19 +98,34 @@ pub fn run_fetch( timeout: Duration, ) -> Result { let started = Instant::now(); + let interrupted = Arc::new(AtomicBool::new(false)); let runtime = Runtime::new().map_err(|e| PageFunctionError::Init(e.to_string()))?; runtime.set_memory_limit(16 * 1024 * 1024); runtime.set_max_stack_size(256 * 1024); + let interrupt_flag = Arc::clone(&interrupted); + runtime.set_interrupt_handler(Some(Box::new(move || { + let timed_out = started.elapsed() >= timeout; + if timed_out { + interrupt_flag.store(true, Ordering::Relaxed); + } + timed_out + }))); let context = Context::full(&runtime).map_err(|e| PageFunctionError::Init(e.to_string()))?; let json_out: String = context.with(|ctx| { - if started.elapsed() > timeout { + if started.elapsed() >= timeout { return Err(PageFunctionError::Timeout(timeout)); } - ctx.eval::<(), _>(worker_source) - .map_err(|e| PageFunctionError::Eval(format!("{e}")))?; + ctx.eval::<(), _>(worker_source).map_err(|e| { + execution_error( + &interrupted, + started, + timeout, + PageFunctionError::Eval(format!("{e}")), + ) + })?; // Ensure fetch exists before wrapping. let globals = ctx.globals(); @@ -150,7 +166,14 @@ pub fn run_fetch( }; "#, ) - .map_err(|e| PageFunctionError::Eval(format!("wrap fetch: {e}")))?; + .map_err(|e| { + execution_error( + &interrupted, + started, + timeout, + PageFunctionError::Eval(format!("wrap fetch: {e}")), + ) + })?; let invoke: Function = globals .get("__bitfun_invoke") @@ -160,17 +183,29 @@ pub fn run_fetch( let req_obj = request_to_object(&ctx, request)?; let maybe: MaybePromise = invoke .call((req_obj, env)) - .map_err(|e| PageFunctionError::Handler(format!("{e}")))?; + .map_err(|e| { + execution_error( + &interrupted, + started, + timeout, + PageFunctionError::Handler(format!("{e}")), + ) + })?; // Drive QuickJS microtasks until the (maybe) promise settles, respecting timeout. let out = loop { - if started.elapsed() > timeout { + if started.elapsed() >= timeout { return Err(PageFunctionError::Timeout(timeout)); } match maybe.result::() { Some(Ok(s)) => break s, Some(Err(e)) => { - return Err(PageFunctionError::Handler(format!("async fetch failed: {e}"))); + return Err(execution_error( + &interrupted, + started, + timeout, + PageFunctionError::Handler(format!("async fetch failed: {e}")), + )); } None => { if !ctx.execute_pending_job() { @@ -190,6 +225,19 @@ pub fn run_fetch( .map_err(|e| PageFunctionError::Handler(format!("invalid fetch response JSON: {e}"))) } +fn execution_error( + interrupted: &AtomicBool, + started: Instant, + timeout: Duration, + fallback: PageFunctionError, +) -> PageFunctionError { + if interrupted.load(Ordering::Relaxed) || started.elapsed() >= timeout { + PageFunctionError::Timeout(timeout) + } else { + fallback + } +} + fn build_env_object<'js>( ctx: &Ctx<'js>, host: Arc, @@ -209,16 +257,22 @@ fn build_env_object<'js>( KV: { get: function(k) { var r = JSON.parse(hostCall("kv_get", String(k), "", "")); + if (r.error) throw new Error(r.error); return r.v; }, put: function(k, v) { - JSON.parse(hostCall("kv_put", String(k), String(v), "")); + var r = JSON.parse(hostCall("kv_put", String(k), String(v), "")); + if (!r.ok) throw new Error(r.error || "KV put failed"); }, delete: function(k) { - return JSON.parse(hostCall("kv_delete", String(k), "", "")).ok; + var r = JSON.parse(hostCall("kv_delete", String(k), "", "")); + if (r.error) throw new Error(r.error); + return r.ok; }, list: function() { - return JSON.parse(hostCall("kv_list", "", "", "")).keys; + var r = JSON.parse(hostCall("kv_list", "", "", "")); + if (r.error) throw new Error(r.error); + return r.keys; } }, DB: { @@ -231,14 +285,19 @@ fn build_env_object<'js>( }, BLOBS: { put: function(id, contentType, dataB64) { - return JSON.parse(hostCall("blob_put", String(id), String(contentType||"application/octet-stream"), String(dataB64))).ok; + var r = JSON.parse(hostCall("blob_put", String(id), String(contentType||"application/octet-stream"), String(dataB64))); + if (!r.ok) throw new Error(r.error || "Blob put failed"); + return true; }, get: function(id) { var r = JSON.parse(hostCall("blob_get", String(id), "", "")); + if (r.error) throw new Error(r.error); return r.found ? { contentType: r.contentType, data: r.data } : null; }, delete: function(id) { - return JSON.parse(hostCall("blob_delete", String(id), "", "")).ok; + var r = JSON.parse(hostCall("blob_delete", String(id), "", "")); + if (r.error) throw new Error(r.error); + return r.ok; } }, ASSETS: { @@ -273,14 +332,14 @@ fn host_call(host: &dyn PageHost, op: &str, a: &str, b: &str, c: &str) -> String }, "kv_delete" => match host.kv_delete(a) { Ok(ok) => format!(r#"{{"ok":{ok}}}"#), - Err(_) => r#"{"ok":false}"#.into(), + Err(e) => format!(r#"{{"ok":false,"error":{}}}"#, json_str(&e)), }, "kv_list" => match host.kv_list() { Ok(keys) => format!( r#"{{"keys":{}}}"#, serde_json::to_string(&keys).unwrap_or_else(|_| "[]".into()) ), - Err(_) => r#"{"keys":[]}"#.into(), + Err(e) => format!(r#"{{"keys":[],"error":{}}}"#, json_str(&e)), }, "db_execute" => host .db_execute(a, b) @@ -303,7 +362,7 @@ fn host_call(host: &dyn PageHost, op: &str, a: &str, b: &str, c: &str) -> String }, "blob_delete" => match host.blob_delete(a) { Ok(ok) => format!(r#"{{"ok":{ok}}}"#), - Err(_) => r#"{"ok":false}"#.into(), + Err(e) => format!(r#"{{"ok":false,"error":{}}}"#, json_str(&e)), }, "assets_fetch" => match host.assets_get(a) { Ok(Some((ct, bytes))) => { @@ -497,4 +556,46 @@ mod tests { Some("application/json") ); } + + #[test] + fn top_level_infinite_loop_is_interrupted() { + let timeout = Duration::from_millis(25); + let started = Instant::now(); + let err = run_fetch( + "while (true) {}", + &test_request(), + Arc::new(MemoryPageHost::default()), + timeout, + ) + .unwrap_err(); + + assert!(matches!(err, PageFunctionError::Timeout(value) if value == timeout)); + assert!(started.elapsed() < Duration::from_secs(1)); + } + + #[test] + fn fetch_handler_infinite_loop_is_interrupted() { + let timeout = Duration::from_millis(25); + let started = Instant::now(); + let err = run_fetch( + "function fetch() { while (true) {} }", + &test_request(), + Arc::new(MemoryPageHost::default()), + timeout, + ) + .unwrap_err(); + + assert!(matches!(err, PageFunctionError::Timeout(value) if value == timeout)); + assert!(started.elapsed() < Duration::from_secs(1)); + } + + fn test_request() -> FetchRequest { + FetchRequest { + method: "GET".into(), + url: "https://example/p/u/s/".into(), + path: "/".into(), + headers: HashMap::new(), + body: None, + } + } } diff --git a/src/crates/services/relay-service/Cargo.toml b/src/crates/services/relay-service/Cargo.toml index 5134e90459..5b959a203f 100644 --- a/src/crates/services/relay-service/Cargo.toml +++ b/src/crates/services/relay-service/Cargo.toml @@ -34,7 +34,7 @@ libsqlite3-sys = { version = "0.30", features = ["bundled"] } argon2 = "0.5" aes-gcm = "0.10" bitfun-page-function-runtime = { path = "../page-function-runtime" } -rusqlite = { version = "0.32", features = ["bundled"] } +rusqlite = { version = "0.32", features = ["bundled", "hooks", "limits"] } [dev-dependencies] tower = { version = "0.5", features = ["util"] } diff --git a/src/crates/services/relay-service/src/db.rs b/src/crates/services/relay-service/src/db.rs index ef1b44ac9e..3dc922b1ff 100644 --- a/src/crates/services/relay-service/src/db.rs +++ b/src/crates/services/relay-service/src/db.rs @@ -7,12 +7,19 @@ use anyhow::{anyhow, Result}; use chrono::Utc; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous}; -use sqlx::{Pool, Sqlite}; +use sqlx::{Pool, QueryBuilder, Sqlite}; use std::str::FromStr; use std::time::Duration; pub type DbPool = Pool; +pub const MAX_PAGE_KV_KEY_BYTES: usize = 256; +pub const MAX_PAGE_KV_VALUE_BYTES: usize = 64 * 1024; +pub const MAX_PAGE_KV_ENTRIES: i64 = 1_024; +pub const MAX_USER_KV_ENTRIES: i64 = 10_000; +pub const MAX_PAGE_KV_BYTES: i64 = 5 * 1024 * 1024; +pub const MAX_USER_KV_BYTES: i64 = 25 * 1024 * 1024; + const SCHEMA: &str = r#" CREATE TABLE IF NOT EXISTS users ( user_id TEXT PRIMARY KEY, @@ -68,6 +75,7 @@ CREATE TABLE IF NOT EXISTS sync_settings ( CREATE TABLE IF NOT EXISTS pages ( user_id TEXT NOT NULL REFERENCES users(user_id), slug TEXT NOT NULL, + generation TEXT NOT NULL DEFAULT '', visibility TEXT NOT NULL DEFAULT 'private', title TEXT NOT NULL DEFAULT '', file_count INTEGER NOT NULL DEFAULT 0, @@ -115,12 +123,28 @@ const MIGRATE_PAGES_DEPLOYED_VERSION: &str = r#" ALTER TABLE pages ADD COLUMN deployed_version_id TEXT; "#; +const MIGRATE_PAGES_GENERATION: &str = r#" +ALTER TABLE pages ADD COLUMN generation TEXT NOT NULL DEFAULT ''; +"#; + const MIGRATE_AUTH_TOKEN_KIND: &str = r#" ALTER TABLE auth_tokens ADD COLUMN token_kind TEXT NOT NULL DEFAULT 'device'; "#; /// Open (or create) the SQLite database and ensure the schema exists. pub async fn connect(db_path: &str) -> Result { + connect_with_presence_reset(db_path, true).await +} + +/// Open the shared database for the out-of-process administration CLI. +/// Unlike a server-process startup, an admin connection must preserve the +/// live relay's durable presence projection: even `list-users` may run via +/// `docker exec` while authenticated WebSockets remain active. +pub async fn connect_for_admin(db_path: &str) -> Result { + connect_with_presence_reset(db_path, false).await +} + +async fn connect_with_presence_reset(db_path: &str, reset_presence: bool) -> Result { let options = SqliteConnectOptions::from_str(&format!("sqlite://{db_path}"))? .create_if_missing(true) .foreign_keys(true) @@ -136,6 +160,21 @@ pub async fn connect(db_path: &str) -> Result { let _ = sqlx::query(MIGRATE_PAGES_DEPLOYED_VERSION) .execute(&pool) .await; + if let Err(error) = sqlx::query(MIGRATE_PAGES_GENERATION).execute(&pool).await { + if !error.to_string().contains("duplicate column name") { + return Err(anyhow!("migrate page generations: {error}")); + } + } + // A Page generation is an authorization boundary. It must change when a + // `(user_id, slug)` row is deleted and later recreated, so legacy rows get + // a random value once instead of sharing a constant migration default. + sqlx::query( + "UPDATE pages SET generation = lower(hex(randomblob(16))) \ + WHERE generation = ''", + ) + .execute(&pool) + .await + .map_err(|e| anyhow!("initialize page generations: {e}"))?; // Existing tokens predate delegation scopes and are full device tokens. if let Err(error) = sqlx::query(MIGRATE_AUTH_TOKEN_KIND).execute(&pool).await { if !error.to_string().contains("duplicate column name") { @@ -149,13 +188,16 @@ pub async fn connect(db_path: &str) -> Result { .execute(&pool) .await .map_err(|e| anyhow!("clean expired auth tokens: {e}"))?; - // Online presence is owned by the in-memory connection registry. A fresh - // process has no live sockets, so never carry stale online flags across a - // restart or crash. - sqlx::query("UPDATE devices SET online = 0") - .execute(&pool) - .await - .map_err(|e| anyhow!("reset stale device presence: {e}"))?; + if reset_presence { + // Online presence is owned by the in-memory connection registry. A + // fresh *server* process has no live sockets, so never carry stale + // online flags across a restart or crash. Administration processes + // intentionally skip this reset because the server may still run. + sqlx::query("UPDATE devices SET online = 0") + .execute(&pool) + .await + .map_err(|e| anyhow!("reset stale device presence: {e}"))?; + } tracing::info!("Account database initialized at {db_path}"); Ok(pool) } @@ -552,12 +594,11 @@ impl DeviceRow { let now = Utc::now().timestamp(); sqlx::query( "INSERT INTO devices (device_id, user_id, device_name, public_key, last_seen_at, online) \ - VALUES (?, ?, ?, ?, ?, 1) \ + VALUES (?, ?, ?, ?, ?, 0) \ ON CONFLICT(user_id, device_id) DO UPDATE SET \ device_name = excluded.device_name, \ public_key = excluded.public_key, \ - last_seen_at = excluded.last_seen_at, \ - online = 1", + last_seen_at = excluded.last_seen_at", ) .bind(device_id) .bind(user_id) @@ -754,6 +795,39 @@ impl AuthToken { Ok(Some(auth_token)) } + /// Batch-load active device tokens for the server-side revocation reaper. + /// Chunking keeps the query below SQLite's bind-variable limit even when a + /// relay has many connected devices. + pub async fn find_valid_device_tokens( + pool: &DbPool, + tokens: &[String], + ) -> Result> { + const TOKEN_QUERY_CHUNK_SIZE: usize = 200; + let now = Utc::now().timestamp(); + let mut valid = Vec::new(); + for chunk in tokens.chunks(TOKEN_QUERY_CHUNK_SIZE) { + let mut query = QueryBuilder::::new( + "SELECT token, user_id, device_id, token_kind, created_at, expires_at \ + FROM auth_tokens WHERE token_kind = 'device' AND expires_at > ", + ); + query.push_bind(now); + query.push(" AND token IN ("); + let mut separated = query.separated(", "); + for token in chunk { + separated.push_bind(token); + } + separated.push_unseparated(")"); + valid.extend( + query + .build_query_as::() + .fetch_all(pool) + .await + .map_err(|e| anyhow!("batch find active device tokens: {e}"))?, + ); + } + Ok(valid) + } + /// Revoke (delete) all tokens belonging to a specific device. pub async fn revoke_by_device(pool: &DbPool, user_id: &str, device_id: &str) -> Result<()> { sqlx::query("DELETE FROM auth_tokens WHERE user_id = ? AND device_id = ?") @@ -983,6 +1057,7 @@ impl PageVisibility { pub struct PageRow { pub user_id: String, pub slug: String, + pub generation: String, pub visibility: String, pub title: String, pub file_count: i64, @@ -998,6 +1073,7 @@ pub struct PageWithUsername { pub user_id: String, pub username: String, pub slug: String, + pub generation: String, pub visibility: String, pub title: String, pub file_count: i64, @@ -1020,12 +1096,36 @@ pub struct PageVersionRow { pub created_at: i64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SavePageVersionOutcome { + Saved, + PageLimitReached, + VersionLimitReached, + GenerationMismatch, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeletePageVersionOutcome { + Deleted, + Deployed, + NotFound, + GenerationMismatch, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PageMutationOutcome { + Applied(T), + NotFound, + GenerationMismatch, +} + impl PageRow { pub fn visibility_enum(&self) -> Option { PageVisibility::parse(&self.visibility) } - const SELECT_COLS: &'static str = "user_id, slug, visibility, title, file_count, total_bytes, \ + const SELECT_COLS: &'static str = + "user_id, slug, generation, visibility, title, file_count, total_bytes, \ deployed_version_id, created_at, updated_at"; /// Ensure a page metadata row exists (draft upload / first save). @@ -1037,11 +1137,12 @@ impl PageRow { title: &str, ) -> Result<()> { let now = Utc::now().timestamp(); + let generation = new_page_generation(); sqlx::query( "INSERT INTO pages \ - (user_id, slug, visibility, title, file_count, total_bytes, deployed_version_id, \ + (user_id, slug, generation, visibility, title, file_count, total_bytes, deployed_version_id, \ created_at, updated_at) \ - VALUES (?, ?, ?, ?, 0, 0, NULL, ?, ?) \ + VALUES (?, ?, ?, ?, ?, 0, 0, NULL, ?, ?) \ ON CONFLICT(user_id, slug) DO UPDATE SET \ visibility = excluded.visibility, \ title = CASE WHEN excluded.title = '' THEN pages.title ELSE excluded.title END, \ @@ -1049,6 +1150,7 @@ impl PageRow { ) .bind(user_id) .bind(slug) + .bind(generation) .bind(visibility.as_str()) .bind(title) .bind(now) @@ -1093,17 +1195,155 @@ impl PageRow { Ok(row.0) } + /// Atomically create/update Page metadata and insert its immutable version. + /// A failed version insert or quota check must never leak title/visibility + /// changes onto the previously deployed Page. + #[allow(clippy::too_many_arguments)] + pub async fn save_version_with_meta( + pool: &DbPool, + user_id: &str, + slug: &str, + visibility: PageVisibility, + title: &str, + version_id: &str, + file_count: i64, + total_bytes: i64, + has_worker: bool, + note: &str, + max_pages_per_user: i64, + max_versions_per_page: i64, + expected_generation: Option<&str>, + create: bool, + ) -> Result { + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("begin save page version transaction: {e}"))?; + + let page_exists: Option<(String,)> = + sqlx::query_as("SELECT generation FROM pages WHERE user_id = ? AND slug = ?") + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("read page before saving version: {e}"))?; + let version_count: (i64,) = + sqlx::query_as("SELECT COUNT(*) FROM page_versions WHERE user_id = ? AND slug = ?") + .bind(user_id) + .bind(slug) + .fetch_one(&mut *tx) + .await + .map_err(|e| anyhow!("count page versions before save: {e}"))?; + if version_count.0 >= max_versions_per_page { + return Ok(SavePageVersionOutcome::VersionLimitReached); + } + + let now = Utc::now().timestamp(); + if let Some((generation,)) = page_exists { + if create || expected_generation != Some(generation.as_str()) { + return Ok(SavePageVersionOutcome::GenerationMismatch); + } + let result = sqlx::query( + "UPDATE pages SET visibility = ?, title = ?, updated_at = ? \ + WHERE user_id = ? AND slug = ? AND generation = ?", + ) + .bind(visibility.as_str()) + .bind(title) + .bind(now) + .bind(user_id) + .bind(slug) + .bind(&generation) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("update page metadata while saving version: {e}"))?; + if result.rows_affected() == 0 { + return Err(anyhow!("page disappeared while saving version")); + } + } else { + if !create || expected_generation.is_some() { + return Ok(SavePageVersionOutcome::GenerationMismatch); + } + let page_count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pages WHERE user_id = ?") + .bind(user_id) + .fetch_one(&mut *tx) + .await + .map_err(|e| anyhow!("count pages before save: {e}"))?; + if page_count.0 >= max_pages_per_user { + return Ok(SavePageVersionOutcome::PageLimitReached); + } + let inserted = sqlx::query( + "INSERT INTO pages \ + (user_id, slug, generation, visibility, title, file_count, total_bytes, \ + deployed_version_id, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, 0, 0, NULL, ?, ?) \ + ON CONFLICT(user_id, slug) DO NOTHING", + ) + .bind(user_id) + .bind(slug) + .bind(new_page_generation()) + .bind(visibility.as_str()) + .bind(title) + .bind(now) + .bind(now) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("create page while saving version: {e}"))?; + if inserted.rows_affected() == 0 { + return Ok(SavePageVersionOutcome::GenerationMismatch); + } + } + + sqlx::query( + "INSERT INTO page_versions \ + (user_id, slug, version_id, title, file_count, total_bytes, has_worker, note, created_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(user_id) + .bind(slug) + .bind(version_id) + .bind(title) + .bind(file_count) + .bind(total_bytes) + .bind(has_worker as i64) + .bind(note) + .bind(now) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("insert page version with metadata: {e}"))?; + + tx.commit() + .await + .map_err(|e| anyhow!("commit saved page version: {e}"))?; + Ok(SavePageVersionOutcome::Saved) + } + pub async fn update_meta( pool: &DbPool, user_id: &str, slug: &str, + expected_generation: &str, visibility: Option, title: Option<&str>, - ) -> Result { - let existing = Self::get(pool, user_id, slug).await?; + ) -> Result> { + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("begin update page transaction: {e}"))?; + let existing = sqlx::query_as::<_, PageRow>(&format!( + "SELECT {} FROM pages WHERE user_id = ? AND slug = ?", + Self::SELECT_COLS + )) + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("read page before metadata update: {e}"))?; let Some(page) = existing else { - return Ok(false); + return Ok(PageMutationOutcome::NotFound); }; + if page.generation != expected_generation { + return Ok(PageMutationOutcome::GenerationMismatch); + } let now = Utc::now().timestamp(); let new_vis = visibility .map(|v| v.as_str().to_string()) @@ -1111,17 +1351,34 @@ impl PageRow { let new_title = title.map(|t| t.to_string()).unwrap_or(page.title); let result = sqlx::query( "UPDATE pages SET visibility = ?, title = ?, updated_at = ? \ - WHERE user_id = ? AND slug = ?", + WHERE user_id = ? AND slug = ? AND generation = ?", ) .bind(&new_vis) .bind(&new_title) .bind(now) .bind(user_id) .bind(slug) - .execute(pool) + .bind(expected_generation) + .execute(&mut *tx) .await .map_err(|e| anyhow!("update page meta: {e}"))?; - Ok(result.rows_affected() > 0) + if result.rows_affected() == 0 { + return Ok(PageMutationOutcome::GenerationMismatch); + } + let updated = sqlx::query_as::<_, PageRow>(&format!( + "SELECT {} FROM pages WHERE user_id = ? AND slug = ? AND generation = ?", + Self::SELECT_COLS + )) + .bind(user_id) + .bind(slug) + .bind(expected_generation) + .fetch_one(&mut *tx) + .await + .map_err(|e| anyhow!("read updated page metadata: {e}"))?; + tx.commit() + .await + .map_err(|e| anyhow!("commit update page transaction: {e}"))?; + Ok(PageMutationOutcome::Applied(updated)) } pub async fn set_deployed_version( @@ -1153,32 +1410,268 @@ impl PageRow { Ok(()) } - pub async fn delete(pool: &DbPool, user_id: &str, slug: &str) -> Result { + /// Atomically attach the synthetic `v1` record used by the pre-versioned + /// asset-layout migration. The generation fence prevents a stale migration + /// from inserting an orphan version or updating a delete/recreated Page. + #[allow(clippy::too_many_arguments)] + pub async fn migrate_legacy_version( + pool: &DbPool, + user_id: &str, + slug: &str, + expected_generation: &str, + version_id: &str, + title: &str, + file_count: i64, + total_bytes: i64, + has_worker: bool, + ) -> Result> { + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("begin legacy page migration transaction: {e}"))?; + let page: Option<(String, Option)> = sqlx::query_as( + "SELECT generation, deployed_version_id FROM pages WHERE user_id = ? AND slug = ?", + ) + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("read page before legacy migration: {e}"))?; + let Some((generation, deployed_version_id)) = page else { + return Ok(PageMutationOutcome::NotFound); + }; + if generation != expected_generation { + return Ok(PageMutationOutcome::GenerationMismatch); + } + let version_count: (i64,) = + sqlx::query_as("SELECT COUNT(*) FROM page_versions WHERE user_id = ? AND slug = ?") + .bind(user_id) + .bind(slug) + .fetch_one(&mut *tx) + .await + .map_err(|e| anyhow!("count versions before legacy migration: {e}"))?; + if deployed_version_id.is_some() || version_count.0 > 0 { + return Ok(PageMutationOutcome::Applied(false)); + } + + let now = Utc::now().timestamp(); + let inserted = sqlx::query( + "INSERT OR IGNORE INTO page_versions \ + (user_id, slug, version_id, title, file_count, total_bytes, has_worker, note, created_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, 'migrated', ?)", + ) + .bind(user_id) + .bind(slug) + .bind(version_id) + .bind(title) + .bind(file_count) + .bind(total_bytes) + .bind(has_worker as i64) + .bind(now) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("insert migrated legacy page version: {e}"))?; + if inserted.rows_affected() == 0 { + return Ok(PageMutationOutcome::Applied(false)); + } + let updated = sqlx::query( + "UPDATE pages SET deployed_version_id = ?, updated_at = ? \ + WHERE user_id = ? AND slug = ? AND generation = ? AND deployed_version_id IS NULL", + ) + .bind(version_id) + .bind(now) + .bind(user_id) + .bind(slug) + .bind(expected_generation) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("deploy migrated legacy page version: {e}"))?; + if updated.rows_affected() == 0 { + return Ok(PageMutationOutcome::GenerationMismatch); + } + tx.commit() + .await + .map_err(|e| anyhow!("commit legacy page migration: {e}"))?; + Ok(PageMutationOutcome::Applied(true)) + } + + /// Deploy a version and return the resulting Page from one transaction. + /// The version lookup and production-pointer update cannot race a version + /// deletion performed through the paired transactional API below. + pub async fn deploy_version( + pool: &DbPool, + user_id: &str, + slug: &str, + version_id: &str, + expected_generation: &str, + ) -> Result> { + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("begin deploy page version transaction: {e}"))?; + let page_generation: Option<(String,)> = + sqlx::query_as("SELECT generation FROM pages WHERE user_id = ? AND slug = ?") + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("read page before deploy: {e}"))?; + let Some((page_generation,)) = page_generation else { + return Ok(PageMutationOutcome::NotFound); + }; + if page_generation != expected_generation { + return Ok(PageMutationOutcome::GenerationMismatch); + } + + let version = sqlx::query_as::<_, PageVersionRow>( + "SELECT user_id, slug, version_id, title, file_count, total_bytes, has_worker, note, created_at \ + FROM page_versions WHERE user_id = ? AND slug = ? AND version_id = ?", + ) + .bind(user_id) + .bind(slug) + .bind(version_id) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("get page version for deploy: {e}"))?; + let Some(version) = version else { + return Ok(PageMutationOutcome::NotFound); + }; + + let now = Utc::now().timestamp(); + let updated = sqlx::query( + "UPDATE pages SET deployed_version_id = ?, file_count = ?, total_bytes = ?, \ + title = CASE WHEN ? = '' THEN title ELSE ? END, updated_at = ? \ + WHERE user_id = ? AND slug = ? AND generation = ?", + ) + .bind(&version.version_id) + .bind(version.file_count) + .bind(version.total_bytes) + .bind(&version.title) + .bind(&version.title) + .bind(now) + .bind(user_id) + .bind(slug) + .bind(expected_generation) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("deploy page version: {e}"))?; + if updated.rows_affected() == 0 { + return Ok(PageMutationOutcome::GenerationMismatch); + } + let page = sqlx::query_as::<_, PageRow>(&format!( + "SELECT {} FROM pages WHERE user_id = ? AND slug = ? AND generation = ?", + Self::SELECT_COLS + )) + .bind(user_id) + .bind(slug) + .bind(expected_generation) + .fetch_one(&mut *tx) + .await + .map_err(|e| anyhow!("read deployed page: {e}"))?; + tx.commit() + .await + .map_err(|e| anyhow!("commit deployed page version: {e}"))?; + Ok(PageMutationOutcome::Applied(page)) + } + + pub async fn clear_deployed_version( + pool: &DbPool, + user_id: &str, + slug: &str, + expected_generation: &str, + ) -> Result> { + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("begin unpublish page transaction: {e}"))?; + let now = Utc::now().timestamp(); + let result = sqlx::query( + "UPDATE pages SET deployed_version_id = NULL, updated_at = ? \ + WHERE user_id = ? AND slug = ? AND generation = ?", + ) + .bind(now) + .bind(user_id) + .bind(slug) + .bind(expected_generation) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("clear deployed version: {e}"))?; + if result.rows_affected() == 0 { + let exists: Option<(String,)> = + sqlx::query_as("SELECT generation FROM pages WHERE user_id = ? AND slug = ?") + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("read page after failed unpublish: {e}"))?; + return Ok(if exists.is_some() { + PageMutationOutcome::GenerationMismatch + } else { + PageMutationOutcome::NotFound + }); + } + tx.commit() + .await + .map_err(|e| anyhow!("commit unpublish page transaction: {e}"))?; + Ok(PageMutationOutcome::Applied(())) + } + + pub async fn delete( + pool: &DbPool, + user_id: &str, + slug: &str, + expected_generation: &str, + ) -> Result> { + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("begin page delete transaction: {e}"))?; + let existing: Option<(String,)> = + sqlx::query_as("SELECT generation FROM pages WHERE user_id = ? AND slug = ?") + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("read page before delete: {e}"))?; + let Some((generation,)) = existing else { + return Ok(PageMutationOutcome::NotFound); + }; + if generation != expected_generation { + return Ok(PageMutationOutcome::GenerationMismatch); + } sqlx::query("DELETE FROM page_kv WHERE user_id = ? AND slug = ?") .bind(user_id) .bind(slug) - .execute(pool) + .execute(&mut *tx) .await .map_err(|e| anyhow!("delete page_kv: {e}"))?; sqlx::query("DELETE FROM page_blobs WHERE user_id = ? AND slug = ?") .bind(user_id) .bind(slug) - .execute(pool) + .execute(&mut *tx) .await .map_err(|e| anyhow!("delete page_blobs: {e}"))?; sqlx::query("DELETE FROM page_versions WHERE user_id = ? AND slug = ?") .bind(user_id) .bind(slug) - .execute(pool) + .execute(&mut *tx) .await .map_err(|e| anyhow!("delete page_versions: {e}"))?; - let result = sqlx::query("DELETE FROM pages WHERE user_id = ? AND slug = ?") - .bind(user_id) - .bind(slug) - .execute(pool) + let result = + sqlx::query("DELETE FROM pages WHERE user_id = ? AND slug = ? AND generation = ?") + .bind(user_id) + .bind(slug) + .bind(expected_generation) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("delete page: {e}"))?; + if result.rows_affected() == 0 { + return Ok(PageMutationOutcome::GenerationMismatch); + } + tx.commit() .await - .map_err(|e| anyhow!("delete page: {e}"))?; - Ok(result.rows_affected() > 0) + .map_err(|e| anyhow!("commit page delete transaction: {e}"))?; + Ok(PageMutationOutcome::Applied(())) } /// Resolve a page by public URL components `(username, slug)`. @@ -1188,7 +1681,7 @@ impl PageRow { slug: &str, ) -> Result> { let row = sqlx::query_as::<_, PageWithUsername>( - "SELECT p.user_id, u.username, p.slug, p.visibility, p.title, \ + "SELECT p.user_id, u.username, p.slug, p.generation, p.visibility, p.title, \ p.file_count, p.total_bytes, p.deployed_version_id, p.created_at, p.updated_at \ FROM pages p JOIN users u ON u.user_id = p.user_id \ WHERE u.username = ? AND p.slug = ?", @@ -1305,18 +1798,80 @@ impl PageVersionRow { .map_err(|e| anyhow!("delete page version: {e}"))?; Ok(result.rows_affected() > 0) } + + pub async fn delete_if_not_deployed( + pool: &DbPool, + user_id: &str, + slug: &str, + version_id: &str, + expected_generation: &str, + ) -> Result { + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("begin delete page version transaction: {e}"))?; + let page: Option<(String, Option)> = sqlx::query_as( + "SELECT generation, deployed_version_id FROM pages WHERE user_id = ? AND slug = ?", + ) + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("read page before version delete: {e}"))?; + let Some((generation, deployed_version_id)) = page else { + return Ok(DeletePageVersionOutcome::NotFound); + }; + if generation != expected_generation { + return Ok(DeletePageVersionOutcome::GenerationMismatch); + } + if deployed_version_id.as_deref() == Some(version_id) { + return Ok(DeletePageVersionOutcome::Deployed); + } + let result = sqlx::query( + "DELETE FROM page_versions WHERE user_id = ? AND slug = ? AND version_id = ? \ + AND EXISTS (SELECT 1 FROM pages WHERE user_id = ? AND slug = ? AND generation = ?)", + ) + .bind(user_id) + .bind(slug) + .bind(version_id) + .bind(user_id) + .bind(slug) + .bind(expected_generation) + .execute(&mut *tx) + .await + .map_err(|e| anyhow!("delete undeployed page version: {e}"))?; + if result.rows_affected() == 0 { + return Ok(DeletePageVersionOutcome::NotFound); + } + tx.commit() + .await + .map_err(|e| anyhow!("commit page version delete: {e}"))?; + Ok(DeletePageVersionOutcome::Deleted) + } } /// Page KV helpers (mutable runtime data, keyed by page not version). pub mod page_kv { use super::*; + fn validate_key(key: &str) -> Result<()> { + if key.is_empty() || key.len() > MAX_PAGE_KV_KEY_BYTES || key.chars().any(char::is_control) + { + return Err(anyhow!( + "page KV key must be non-empty, control-free, and at most {} bytes", + MAX_PAGE_KV_KEY_BYTES + )); + } + Ok(()) + } + pub async fn get( pool: &DbPool, user_id: &str, slug: &str, key: &str, ) -> Result> { + validate_key(key)?; let row: Option<(String,)> = sqlx::query_as("SELECT value FROM page_kv WHERE user_id = ? AND slug = ? AND key = ?") .bind(user_id) @@ -1335,6 +1890,63 @@ pub mod page_kv { key: &str, value: &str, ) -> Result<()> { + validate_key(key)?; + if value.len() > MAX_PAGE_KV_VALUE_BYTES { + return Err(anyhow!( + "page KV value exceeds the {} byte operation limit", + MAX_PAGE_KV_VALUE_BYTES + )); + } + + let mut tx = pool + .begin() + .await + .map_err(|e| anyhow!("page_kv begin quota transaction: {e}"))?; + let page_exists: Option<(i64,)> = + sqlx::query_as("SELECT 1 FROM pages WHERE user_id = ? AND slug = ?") + .bind(user_id) + .bind(slug) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("page_kv verify page exists: {e}"))?; + if page_exists.is_none() { + return Err(anyhow!("page no longer exists")); + } + let old_bytes: Option<(i64,)> = sqlx::query_as( + "SELECT length(CAST(key AS BLOB)) + length(CAST(value AS BLOB)) \ + FROM page_kv WHERE user_id = ? AND slug = ? AND key = ?", + ) + .bind(user_id) + .bind(slug) + .bind(key) + .fetch_optional(&mut *tx) + .await + .map_err(|e| anyhow!("page_kv read existing size: {e}"))?; + let page_usage: (i64, i64) = sqlx::query_as( + "SELECT COUNT(*), COALESCE(SUM(length(CAST(key AS BLOB)) + \ + length(CAST(value AS BLOB))), 0) FROM page_kv WHERE user_id = ? AND slug = ?", + ) + .bind(user_id) + .bind(slug) + .fetch_one(&mut *tx) + .await + .map_err(|e| anyhow!("page_kv read page quota: {e}"))?; + let user_usage: (i64, i64) = sqlx::query_as( + "SELECT COUNT(*), COALESCE(SUM(length(CAST(key AS BLOB)) + \ + length(CAST(value AS BLOB))), 0) FROM page_kv WHERE user_id = ?", + ) + .bind(user_id) + .fetch_one(&mut *tx) + .await + .map_err(|e| anyhow!("page_kv read account quota: {e}"))?; + enforce_quota( + page_usage, + user_usage, + old_bytes.map_or(0, |row| row.0), + key.len().saturating_add(value.len()) as i64, + old_bytes.is_none(), + )?; + let now = Utc::now().timestamp(); sqlx::query( "INSERT INTO page_kv (user_id, slug, key, value, updated_at) VALUES (?, ?, ?, ?, ?) \ @@ -1345,13 +1957,17 @@ pub mod page_kv { .bind(key) .bind(value) .bind(now) - .execute(pool) + .execute(&mut *tx) .await .map_err(|e| anyhow!("page_kv put: {e}"))?; + tx.commit() + .await + .map_err(|e| anyhow!("page_kv commit: {e}"))?; Ok(()) } pub async fn delete(pool: &DbPool, user_id: &str, slug: &str, key: &str) -> Result { + validate_key(key)?; let result = sqlx::query("DELETE FROM page_kv WHERE user_id = ? AND slug = ? AND key = ?") .bind(user_id) .bind(slug) @@ -1372,6 +1988,78 @@ pub mod page_kv { .map_err(|e| anyhow!("page_kv list: {e}"))?; Ok(rows.into_iter().map(|r| r.0).collect()) } + + fn enforce_quota( + page_usage: (i64, i64), + user_usage: (i64, i64), + replaced_bytes: i64, + added_bytes: i64, + is_new: bool, + ) -> Result<()> { + let added_entries = i64::from(is_new); + if page_usage.0.saturating_add(added_entries) > MAX_PAGE_KV_ENTRIES { + return Err(anyhow!("page KV entry quota exceeded")); + } + if user_usage.0.saturating_add(added_entries) > MAX_USER_KV_ENTRIES { + return Err(anyhow!("account KV entry quota exceeded")); + } + if page_usage + .1 + .saturating_sub(replaced_bytes) + .saturating_add(added_bytes) + > MAX_PAGE_KV_BYTES + { + return Err(anyhow!("page KV byte quota exceeded")); + } + if user_usage + .1 + .saturating_sub(replaced_bytes) + .saturating_add(added_bytes) + > MAX_USER_KV_BYTES + { + return Err(anyhow!("account KV byte quota exceeded")); + } + Ok(()) + } + + #[cfg(test)] + mod quota_tests { + use super::*; + + #[test] + fn quota_projection_handles_insert_and_overwrite() { + assert!(enforce_quota( + (MAX_PAGE_KV_ENTRIES, 100), + (MAX_PAGE_KV_ENTRIES, 100), + 10, + 10, + false, + ) + .is_ok()); + assert!(enforce_quota( + (MAX_PAGE_KV_ENTRIES, 100), + (MAX_PAGE_KV_ENTRIES, 100), + 0, + 1, + true, + ) + .is_err()); + assert!( + enforce_quota((1, MAX_PAGE_KV_BYTES), (1, MAX_PAGE_KV_BYTES), 1, 2, false,) + .is_err() + ); + assert!(enforce_quota((1, 1), (1, MAX_USER_KV_BYTES), 0, 1, false,).is_err()); + } + + #[tokio::test] + async fn put_rejects_oversized_single_values_before_writing() { + let pool = connect(":memory:").await.unwrap(); + let value = "x".repeat(MAX_PAGE_KV_VALUE_BYTES + 1); + let err = put(&pool, "u1", "site", "key", &value).await.unwrap_err(); + assert!(err.to_string().contains("operation limit")); + assert!(get(&pool, "u1", "site", "key").await.unwrap().is_none()); + } + } } /// Legacy single-room asset key (pre-versioning). Used only for one-time migration. @@ -1389,6 +2077,11 @@ pub fn page_version_asset_key(user_id: &str, slug: &str, version_id: &str) -> St format!("pages/{user_id}/{slug}/v/{version_id}") } +fn new_page_generation() -> String { + let bytes: [u8; 16] = rand::random(); + bytes.iter().map(|byte| format!("{byte:02x}")).collect() +} + /// Generate a new version id (`v` + 12 hex chars). pub fn new_page_version_id() -> String { let bytes: [u8; 6] = rand::random(); @@ -1408,6 +2101,28 @@ mod tests { pool } + #[tokio::test] + async fn admin_connection_preserves_live_server_presence_projection() { + let temp = tempfile::tempdir().unwrap(); + let db_path = temp.path().join("relay.db"); + let db_path = db_path.to_str().unwrap(); + let runtime_pool = connect(db_path).await.unwrap(); + UserRow::create(&runtime_pool, "u1", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + DeviceRow::upsert(&runtime_pool, "d1", "u1", "Laptop", None) + .await + .unwrap(); + DeviceRow::set_online(&runtime_pool, "u1", "d1", true) + .await + .unwrap(); + + let admin_pool = connect_for_admin(db_path).await.unwrap(); + let rows = DeviceRow::list_by_user(&admin_pool, "u1").await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].online, 1); + } + #[test] fn lockout_schedule() { let now = 1000; @@ -1665,6 +2380,77 @@ mod tests { let _ = std::fs::remove_file(db_path); } + #[tokio::test] + async fn legacy_pages_receive_nonempty_authorization_generations() { + let db_path = std::env::temp_dir().join(format!( + "bitfun-relay-page-generation-migration-{}-{}.db", + std::process::id(), + rand::random::() + )); + let db_path_text = db_path.to_string_lossy().to_string(); + let legacy_options = SqliteConnectOptions::from_str(&format!("sqlite://{db_path_text}")) + .unwrap() + .create_if_missing(true) + .foreign_keys(false); + let legacy = SqlitePoolOptions::new() + .max_connections(1) + .connect_with(legacy_options) + .await + .unwrap(); + sqlx::query( + "CREATE TABLE users (\ + user_id TEXT PRIMARY KEY, username TEXT UNIQUE NOT NULL,\ + salt TEXT NOT NULL, kdf_salt TEXT NOT NULL, argon2_params TEXT NOT NULL,\ + password_hash TEXT NOT NULL, wrapped_master_key TEXT NOT NULL,\ + failed_attempts INTEGER NOT NULL DEFAULT 0, locked_until INTEGER NOT NULL DEFAULT 0,\ + created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL\ + )", + ) + .execute(&legacy) + .await + .unwrap(); + sqlx::query( + "CREATE TABLE pages (\ + user_id TEXT NOT NULL REFERENCES users(user_id), slug TEXT NOT NULL,\ + visibility TEXT NOT NULL DEFAULT 'private', title TEXT NOT NULL DEFAULT '',\ + file_count INTEGER NOT NULL DEFAULT 0, total_bytes INTEGER NOT NULL DEFAULT 0,\ + created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL,\ + PRIMARY KEY (user_id, slug)\ + )", + ) + .execute(&legacy) + .await + .unwrap(); + sqlx::query( + "INSERT INTO users \ + (user_id, username, salt, kdf_salt, argon2_params, password_hash, \ + wrapped_master_key, created_at, updated_at) \ + VALUES ('u1', 'alice', 's', 'ks', '{}', 'hash', 'wmk', 1, 1)", + ) + .execute(&legacy) + .await + .unwrap(); + sqlx::query( + "INSERT INTO pages \ + (user_id, slug, visibility, title, file_count, total_bytes, created_at, updated_at) \ + VALUES ('u1', 'legacy', 'private', 'Legacy', 0, 0, 1, 1)", + ) + .execute(&legacy) + .await + .unwrap(); + legacy.close().await; + + let migrated = connect(&db_path_text).await.unwrap(); + let page = PageRow::get(&migrated, "u1", "legacy") + .await + .unwrap() + .unwrap(); + assert_eq!(page.generation.len(), 32); + assert!(page.generation.bytes().all(|byte| byte.is_ascii_hexdigit())); + migrated.close().await; + let _ = std::fs::remove_file(db_path); + } + #[tokio::test] async fn page_ensure_version_deploy_and_resolve() { let pool = setup().await; @@ -1710,10 +2496,210 @@ mod tests { Some("v") ); - assert!(PageRow::delete(&pool, "u1", "my-site").await.unwrap()); + assert_eq!( + PageRow::delete(&pool, "u1", "my-site", &listed[0].generation) + .await + .unwrap(), + PageMutationOutcome::Applied(()) + ); assert!(PageRow::get(&pool, "u1", "my-site") .await .unwrap() .is_none()); } + + #[tokio::test] + async fn page_version_metadata_and_lifecycle_mutations_are_atomic() { + let pool = setup().await; + UserRow::create(&pool, "u1", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + PageRow::ensure( + &pool, + "u1", + "atomic-site", + PageVisibility::Private, + "Old title", + ) + .await + .unwrap(); + let original = PageRow::get(&pool, "u1", "atomic-site") + .await + .unwrap() + .unwrap(); + PageVersionRow::insert( + &pool, + "u1", + "atomic-site", + "vold", + "Old title", + 1, + 10, + false, + "", + ) + .await + .unwrap(); + + let limited = PageRow::save_version_with_meta( + &pool, + "u1", + "atomic-site", + PageVisibility::Public, + "New title", + "vnew", + 2, + 20, + false, + "new", + 50, + 1, + Some(&original.generation), + false, + ) + .await + .unwrap(); + assert_eq!(limited, SavePageVersionOutcome::VersionLimitReached); + let unchanged = PageRow::get(&pool, "u1", "atomic-site") + .await + .unwrap() + .unwrap(); + assert_eq!(unchanged.visibility, "private"); + assert_eq!(unchanged.title, "Old title"); + assert_eq!(unchanged.generation, original.generation); + assert!(PageVersionRow::get(&pool, "u1", "atomic-site", "vnew") + .await + .unwrap() + .is_none()); + + assert_eq!( + PageRow::save_version_with_meta( + &pool, + "u1", + "atomic-site", + PageVisibility::Public, + "New title", + "vnew", + 2, + 20, + false, + "new", + 50, + 2, + Some(&original.generation), + false, + ) + .await + .unwrap(), + SavePageVersionOutcome::Saved + ); + let updated = PageRow::get(&pool, "u1", "atomic-site") + .await + .unwrap() + .unwrap(); + assert_eq!(updated.visibility, "public"); + assert_eq!(updated.title, "New title"); + assert_eq!(updated.generation, original.generation); + + let deployed = + match PageRow::deploy_version(&pool, "u1", "atomic-site", "vnew", &original.generation) + .await + .unwrap() + { + PageMutationOutcome::Applied(page) => page, + outcome => panic!("unexpected deploy outcome: {outcome:?}"), + }; + assert_eq!(deployed.deployed_version_id.as_deref(), Some("vnew")); + assert_eq!( + PageVersionRow::delete_if_not_deployed( + &pool, + "u1", + "atomic-site", + "vnew", + &original.generation, + ) + .await + .unwrap(), + DeletePageVersionOutcome::Deployed + ); + assert_eq!( + PageVersionRow::delete_if_not_deployed( + &pool, + "u1", + "atomic-site", + "vold", + &original.generation, + ) + .await + .unwrap(), + DeletePageVersionOutcome::Deleted + ); + + assert_eq!( + PageRow::delete(&pool, "u1", "atomic-site", &original.generation) + .await + .unwrap(), + PageMutationOutcome::Applied(()) + ); + PageRow::ensure( + &pool, + "u1", + "atomic-site", + PageVisibility::Private, + "Recreated", + ) + .await + .unwrap(); + let recreated = PageRow::get(&pool, "u1", "atomic-site") + .await + .unwrap() + .unwrap(); + assert_ne!(recreated.generation, original.generation); + + assert_eq!( + PageRow::migrate_legacy_version( + &pool, + "u1", + "atomic-site", + &original.generation, + "v1", + "Old title", + 1, + 10, + false, + ) + .await + .unwrap(), + PageMutationOutcome::GenerationMismatch + ); + assert!(PageVersionRow::list_for_page(&pool, "u1", "atomic-site") + .await + .unwrap() + .is_empty()); + let after_stale_migration = PageRow::get(&pool, "u1", "atomic-site") + .await + .unwrap() + .unwrap(); + assert_eq!(after_stale_migration.title, "Recreated"); + assert!(after_stale_migration.deployed_version_id.is_none()); + assert!(matches!( + PageRow::update_meta( + &pool, + "u1", + "atomic-site", + &original.generation, + Some(PageVisibility::Public), + Some("stale"), + ) + .await + .unwrap(), + PageMutationOutcome::GenerationMismatch + )); + assert_eq!( + PageRow::delete(&pool, "u1", "atomic-site", &original.generation) + .await + .unwrap(), + PageMutationOutcome::GenerationMismatch + ); + } } diff --git a/src/crates/services/relay-service/src/lib.rs b/src/crates/services/relay-service/src/lib.rs index eb3aaedd2d..a3e25dea7a 100644 --- a/src/crates/services/relay-service/src/lib.rs +++ b/src/crates/services/relay-service/src/lib.rs @@ -12,6 +12,7 @@ pub mod admin; pub mod db; pub mod page_data; +pub mod page_execution; pub mod relay; pub mod routes; @@ -1050,6 +1051,9 @@ pub fn build_relay_router_with_page_data_and_origins( asset_store, db, page_data, + page_access_manager: Arc::new(routes::pages::PageAccessManager::new()), + page_upload_manager: Arc::new(routes::pages::PageUploadManager::new()), + page_execution_guard: Arc::new(crate::page_execution::PageExecutionGuard::new()), login_rate_limiter: std::sync::Arc::new(crate::routes::auth::LoginRateLimiter::new()), device_manager: crate::relay::DeviceManager::new(), cors_allow_origins: Arc::new(cors_allow_origins.clone()), diff --git a/src/crates/services/relay-service/src/page_data.rs b/src/crates/services/relay-service/src/page_data.rs index 4db0c52452..06f84ffd70 100644 --- a/src/crates/services/relay-service/src/page_data.rs +++ b/src/crates/services/relay-service/src/page_data.rs @@ -1,78 +1,198 @@ -//! Per-page mutable runtime data (KV / SQLite / Blobs), keyed by (user_id, slug). -//! Survives version deploy/rollback; separate from immutable version assets. +//! Per-page mutable runtime data (KV / SQLite / Blobs), keyed by +//! (user_id, slug, generation). Survives version deploy/rollback, while a +//! delete/recreate cycle receives a fresh, isolated generation. use std::path::{Path, PathBuf}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; use anyhow::{anyhow, Result}; use base64::{engine::general_purpose::STANDARD as B64, Engine}; use bitfun_page_function_runtime::{PageHost, PageMeta}; -use chrono::Utc; +use dashmap::DashMap; +use rusqlite::hooks::{AuthAction, AuthContext, Authorization}; +use rusqlite::limits::Limit; +use sha2::{Digest, Sha256}; use tokio::runtime::Handle; -use crate::db::{page_kv, DbPool}; +use crate::db::{page_kv, DbPool, PageRow}; -/// Root directory for page-data (`{base}/{user_id}/{slug}/...`). +pub const MAX_BLOB_ID_BYTES: usize = 128; +pub const MAX_BLOB_BYTES: usize = 4 * 1024 * 1024; +pub const MAX_BLOB_FILES_PER_PAGE: u64 = 2_048; +pub const MAX_BLOB_FILES_PER_USER: u64 = 10_000; +pub const MAX_MUTABLE_BYTES_PER_PAGE: u64 = 64 * 1024 * 1024; +pub const MAX_MUTABLE_BYTES_PER_USER: u64 = 256 * 1024 * 1024; +pub const MAX_PAGE_DB_BYTES: u64 = 20 * 1024 * 1024; +pub const MAX_DB_SQL_BYTES: usize = 64 * 1024; +pub const MAX_DB_PARAMS_BYTES: usize = 256 * 1024; +pub const MAX_DB_PARAMS: usize = 256; +pub const MAX_DB_QUERY_ROWS: usize = 1_000; +pub const MAX_DB_QUERY_BYTES: usize = 2 * 1024 * 1024; +const MAX_DB_VALUE_BYTES: i32 = 2 * 1024 * 1024; +const LEGACY_BLOB_LAYOUT_MARKER: &str = ".legacy-blob-layout-v1"; +#[cfg(not(test))] +const MAX_DB_OPERATION_TIME: Duration = Duration::from_secs(1); +#[cfg(test)] +const MAX_DB_OPERATION_TIME: Duration = Duration::from_millis(50); + +/// Root directory for page-data +/// (`{base}/{user_id}/{slug}/{generation}/...`). #[derive(Clone)] pub struct PageDataStore { base_dir: PathBuf, + user_mutation_locks: Arc>>>, } impl PageDataStore { pub fn new(base_dir: impl Into) -> Self { let base_dir = base_dir.into(); let _ = std::fs::create_dir_all(&base_dir); - Self { base_dir } + Self { + base_dir, + user_mutation_locks: Arc::new(DashMap::new()), + } } pub fn base_dir(&self) -> &Path { &self.base_dir } - pub fn page_dir(&self, user_id: &str, slug: &str) -> PathBuf { + fn page_root(&self, user_id: &str, slug: &str) -> PathBuf { self.base_dir.join(user_id).join(slug) } - pub fn db_path(&self, user_id: &str, slug: &str) -> PathBuf { - self.page_dir(user_id, slug).join("db.sqlite") + fn page_dir(&self, user_id: &str, slug: &str, generation: &str) -> PathBuf { + self.page_root(user_id, slug).join(generation) + } + + fn db_path(&self, user_id: &str, slug: &str, generation: &str) -> PathBuf { + self.page_dir(user_id, slug, generation).join("db.sqlite") } - pub fn blobs_dir(&self, user_id: &str, slug: &str) -> PathBuf { - self.page_dir(user_id, slug).join("blobs") + fn blobs_dir(&self, user_id: &str, slug: &str, generation: &str) -> PathBuf { + self.page_dir(user_id, slug, generation).join("blobs-v2") } pub fn cleanup_page(&self, user_id: &str, slug: &str) { - let dir = self.page_dir(user_id, slug); - if dir.exists() { + let lock = self.user_mutation_lock(user_id); + let Ok(_guard) = lock.lock() else { + return; + }; + if let Ok(Some(dir)) = self.existing_page_root(user_id, slug) { let _ = std::fs::remove_dir_all(&dir); } } - fn ensure_page_dir(&self, user_id: &str, slug: &str) -> Result { - let dir = self.page_dir(user_id, slug); - std::fs::create_dir_all(&dir).map_err(|e| anyhow!("create page-data dir: {e}"))?; - Ok(dir) + /// Assign any pre-generation PageData layout to `generation` and make the + /// generation directory ready. Calling this before deleting the Page row + /// ensures that a crash cannot let a later Page with the same slug adopt + /// the old mutable data. + pub fn prepare_generation( + &self, + user_id: &str, + slug: &str, + generation: &str, + ) -> Result { + let lock = self.user_mutation_lock(user_id); + let _guard = lock + .lock() + .map_err(|_| anyhow!("page-data mutation lock poisoned"))?; + self.ensure_generation_dir_locked(user_id, slug, generation) + } + + fn ensure_generation_dir_locked( + &self, + user_id: &str, + slug: &str, + generation: &str, + ) -> Result { + validate_generation(generation)?; + ensure_directory(&self.base_dir)?; + ensure_directory(&self.base_dir.join(user_id))?; + let root = self.page_root(user_id, slug); + ensure_directory(&root)?; + + let existing_generations = generation_directories(&root)?; + let may_adopt_legacy = existing_generations.is_empty() + || existing_generations + .iter() + .all(|existing| existing == generation); + let target = self.page_dir(user_id, slug, generation); + ensure_directory(&target)?; + if may_adopt_legacy { + migrate_legacy_page_data(&root, &target)?; + } + Ok(target) + } + + fn existing_page_root(&self, user_id: &str, slug: &str) -> Result> { + if existing_directory(&self.base_dir)?.is_none() + || existing_directory(&self.base_dir.join(user_id))?.is_none() + { + return Ok(None); + } + existing_directory(&self.page_root(user_id, slug)) + } + + fn user_mutation_lock(&self, user_id: &str) -> Arc> { + Arc::clone( + self.user_mutation_locks + .entry(user_id.to_string()) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .value(), + ) } pub fn blob_put( &self, user_id: &str, slug: &str, + generation: &str, blob_id: &str, content_type: &str, data: &[u8], ) -> Result<()> { - if blob_id.contains("..") || blob_id.contains('/') || blob_id.contains('\\') { - return Err(anyhow!("invalid blob id")); + validate_blob_id(blob_id)?; + validate_content_type(content_type)?; + if data.len() > MAX_BLOB_BYTES { + return Err(anyhow!( + "blob exceeds the {} byte operation limit", + MAX_BLOB_BYTES + )); } - let blobs = self.blobs_dir(user_id, slug); - std::fs::create_dir_all(&blobs).map_err(|e| anyhow!("create blobs dir: {e}"))?; - let path = blobs.join(blob_id); + let lock = self.user_mutation_lock(user_id); + let _guard = lock + .lock() + .map_err(|_| anyhow!("page-data mutation lock poisoned"))?; + let page_dir = self.ensure_generation_dir_locked(user_id, slug, generation)?; + let blobs = self.blobs_dir(user_id, slug, generation); + ensure_directory(&blobs)?; + ensure_directory(&blobs.join(".metadata"))?; + migrate_legacy_blob_locked(&page_dir, &blobs, blob_id)?; + let storage_name = blob_storage_name(blob_id); + let path = blobs.join(&storage_name); + let meta = blob_meta_path(&blobs, &storage_name); + + if !path.exists() { + let page_blob_files = page_blob_file_count(&page_dir)?; + let user_blob_files = user_blob_file_count(&self.base_dir.join(user_id))?; + if page_blob_files >= MAX_BLOB_FILES_PER_PAGE { + return Err(anyhow!("page blob file quota exceeded")); + } + if user_blob_files >= MAX_BLOB_FILES_PER_USER { + return Err(anyhow!("account blob file quota exceeded")); + } + } + + let replaced_bytes = file_len(&path).saturating_add(file_len(&meta)); + let added_bytes = (data.len() as u64).saturating_add(content_type.len() as u64); + let page_bytes = directory_size(&page_dir)?; + let user_bytes = directory_size(&self.base_dir.join(user_id))?; + enforce_storage_quota(page_bytes, user_bytes, replaced_bytes, added_bytes)?; + std::fs::write(&path, data).map_err(|e| anyhow!("write blob: {e}"))?; - let meta = blobs.join(format!("{blob_id}.meta")); std::fs::write(&meta, content_type).map_err(|e| anyhow!("write blob meta: {e}"))?; - let _ = content_type; - let _ = Utc::now(); Ok(()) } @@ -80,32 +200,64 @@ impl PageDataStore { &self, user_id: &str, slug: &str, + generation: &str, blob_id: &str, ) -> Result)>> { - if blob_id.contains("..") || blob_id.contains('/') || blob_id.contains('\\') { - return Err(anyhow!("invalid blob id")); + validate_blob_id(blob_id)?; + let lock = self.user_mutation_lock(user_id); + let _guard = lock + .lock() + .map_err(|_| anyhow!("page-data mutation lock poisoned"))?; + let page_dir = self.ensure_generation_dir_locked(user_id, slug, generation)?; + let blobs = self.blobs_dir(user_id, slug, generation); + ensure_directory(&blobs)?; + ensure_directory(&blobs.join(".metadata"))?; + migrate_legacy_blob_locked(&page_dir, &blobs, blob_id)?; + let storage_name = blob_storage_name(blob_id); + let path = blobs.join(&storage_name); + let metadata = match std::fs::symlink_metadata(&path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(anyhow!("read blob metadata: {error}")), + }; + if metadata.file_type().is_symlink() { + return Err(anyhow!("blob path must not be a symlink")); } - let path = self.blobs_dir(user_id, slug).join(blob_id); - if !path.is_file() { + if !metadata.is_file() { return Ok(None); } + let size = metadata.len(); + if size > MAX_BLOB_BYTES as u64 { + return Err(anyhow!("stored blob exceeds the read limit")); + } let data = std::fs::read(&path).map_err(|e| anyhow!("read blob: {e}"))?; - let meta = self - .blobs_dir(user_id, slug) - .join(format!("{blob_id}.meta")); - let content_type = std::fs::read_to_string(&meta) - .unwrap_or_else(|_| "application/octet-stream".to_string()); + let meta = blob_meta_path(&blobs, &storage_name); + let content_type = read_small_metadata(&meta) + .or_else(|| legacy_blob_content_type(&page_dir, blob_id)) + .unwrap_or_else(|| "application/octet-stream".to_string()); Ok(Some((content_type, data))) } - pub fn blob_delete(&self, user_id: &str, slug: &str, blob_id: &str) -> Result { - if blob_id.contains("..") || blob_id.contains('/') || blob_id.contains('\\') { - return Err(anyhow!("invalid blob id")); - } - let path = self.blobs_dir(user_id, slug).join(blob_id); - let meta = self - .blobs_dir(user_id, slug) - .join(format!("{blob_id}.meta")); + pub fn blob_delete( + &self, + user_id: &str, + slug: &str, + generation: &str, + blob_id: &str, + ) -> Result { + validate_blob_id(blob_id)?; + let lock = self.user_mutation_lock(user_id); + let _guard = lock + .lock() + .map_err(|_| anyhow!("page-data mutation lock poisoned"))?; + let page_dir = self.ensure_generation_dir_locked(user_id, slug, generation)?; + let blobs = self.blobs_dir(user_id, slug, generation); + ensure_directory(&blobs)?; + ensure_directory(&blobs.join(".metadata"))?; + migrate_legacy_blob_locked(&page_dir, &blobs, blob_id)?; + let storage_name = blob_storage_name(blob_id); + let path = blobs.join(&storage_name); + let meta = blob_meta_path(&blobs, &storage_name); let existed = path.exists(); let _ = std::fs::remove_file(&path); let _ = std::fs::remove_file(&meta); @@ -116,18 +268,40 @@ impl PageDataStore { &self, user_id: &str, slug: &str, + generation: &str, sql: &str, params_json: &str, ) -> Result { - self.ensure_page_dir(user_id, slug)?; - let path = self.db_path(user_id, slug); - let conn = rusqlite::Connection::open(&path).map_err(|e| anyhow!("open page db: {e}"))?; - let params: Vec = - serde_json::from_str(params_json).unwrap_or_else(|_| Vec::new()); - let mut stmt = conn.prepare(sql).map_err(|e| anyhow!("prepare: {e}"))?; - let changes = stmt - .execute(rusqlite::params_from_iter(params.iter().map(json_to_sql))) - .map_err(|e| anyhow!("execute: {e}"))?; + validate_db_input(sql, params_json)?; + let params = parse_db_params(params_json)?; + let lock = self.user_mutation_lock(user_id); + let _guard = lock + .lock() + .map_err(|_| anyhow!("page-data mutation lock poisoned"))?; + self.ensure_generation_dir_locked(user_id, slug, generation)?; + let path = self.db_path(user_id, slug, generation); + ensure_regular_file_or_missing(&path)?; + let mut conn = + rusqlite::Connection::open(&path).map_err(|e| anyhow!("open page db: {e}"))?; + configure_connection(&conn); + configure_database_quota(self, &conn, user_id, slug, generation, &path)?; + + let tx = conn + .transaction() + .map_err(|e| anyhow!("begin transaction: {e}"))?; + tx.authorizer(Some(authorize_execute)); + let result = tx.execute( + sql, + rusqlite::params_from_iter(params.iter().map(json_to_sql)), + ); + tx.authorizer(None::) -> Authorization>); + let changes = result.map_err(|e| anyhow!("execute: {e}"))?; + let logical_bytes = database_logical_bytes(&tx)?; + if logical_bytes > MAX_PAGE_DB_BYTES { + return Err(anyhow!("page database quota exceeded")); + } + tx.commit().map_err(|e| anyhow!("commit: {e}"))?; + enforce_current_storage_quota(self, user_id, slug, generation)?; Ok(serde_json::json!({ "ok": true, "changes": changes }).to_string()) } @@ -135,24 +309,42 @@ impl PageDataStore { &self, user_id: &str, slug: &str, + generation: &str, sql: &str, params_json: &str, ) -> Result { - self.ensure_page_dir(user_id, slug)?; - let path = self.db_path(user_id, slug); + validate_db_input(sql, params_json)?; + let params = parse_db_params(params_json)?; + let lock = self.user_mutation_lock(user_id); + let _guard = lock + .lock() + .map_err(|_| anyhow!("page-data mutation lock poisoned"))?; + self.ensure_generation_dir_locked(user_id, slug, generation)?; + let path = self.db_path(user_id, slug, generation); + ensure_regular_file_or_missing(&path)?; let conn = rusqlite::Connection::open(&path).map_err(|e| anyhow!("open page db: {e}"))?; - let params: Vec = - serde_json::from_str(params_json).unwrap_or_else(|_| Vec::new()); + configure_connection(&conn); + conn.authorizer(Some(authorize_query)); let mut stmt = conn.prepare(sql).map_err(|e| anyhow!("prepare: {e}"))?; + if !stmt.readonly() { + return Err(anyhow!("DB.query only accepts read-only statements")); + } let col_count = stmt.column_count(); let col_names: Vec = (0..col_count) .map(|i| stmt.column_name(i).unwrap_or("?").to_string()) .collect(); let mut rows_out = Vec::new(); + let mut response_bytes = 0usize; let mut rows = stmt .query(rusqlite::params_from_iter(params.iter().map(json_to_sql))) .map_err(|e| anyhow!("query: {e}"))?; while let Some(row) = rows.next().map_err(|e| anyhow!("row: {e}"))? { + if rows_out.len() >= MAX_DB_QUERY_ROWS { + return Err(anyhow!( + "query exceeds the {} row result limit", + MAX_DB_QUERY_ROWS + )); + } let mut obj = serde_json::Map::new(); for (i, name) in col_names.iter().enumerate() { let val = match row.get_ref(i).map_err(|e| anyhow!("get: {e}"))? { @@ -166,12 +358,639 @@ impl PageDataStore { }; obj.insert(name.clone(), val); } - rows_out.push(serde_json::Value::Object(obj)); + let value = serde_json::Value::Object(obj); + let row_bytes = serde_json::to_vec(&value) + .map_err(|e| anyhow!("serialize query row: {e}"))? + .len(); + if response_bytes.saturating_add(row_bytes) > MAX_DB_QUERY_BYTES { + return Err(anyhow!( + "query exceeds the {} byte result limit", + MAX_DB_QUERY_BYTES + )); + } + response_bytes = response_bytes.saturating_add(row_bytes); + rows_out.push(value); } Ok(serde_json::json!({ "ok": true, "rows": rows_out }).to_string()) } } +fn validate_generation(generation: &str) -> Result<()> { + if generation.len() != 32 || !generation.bytes().all(|byte| byte.is_ascii_hexdigit()) { + return Err(anyhow!("invalid page-data generation")); + } + Ok(()) +} + +fn is_generation_name(name: &std::ffi::OsStr) -> bool { + name.to_str().is_some_and(|value| { + value.len() == 32 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) + }) +} + +fn generation_directories(page_root: &Path) -> Result> { + let mut generations = Vec::new(); + for entry in std::fs::read_dir(page_root) + .map_err(|e| anyhow!("read PageData root {}: {e}", page_root.display()))? + { + let entry = entry.map_err(|e| anyhow!("read PageData entry: {e}"))?; + if !is_generation_name(&entry.file_name()) { + continue; + } + let metadata = std::fs::symlink_metadata(entry.path()) + .map_err(|e| anyhow!("inspect PageData generation: {e}"))?; + if metadata.file_type().is_symlink() { + return Err(anyhow!("PageData generation must not be a symlink")); + } + if !metadata.is_dir() { + return Err(anyhow!("PageData generation path must be a directory")); + } + generations.push(entry.file_name().to_string_lossy().into_owned()); + } + Ok(generations) +} + +fn migrate_legacy_page_data(page_root: &Path, generation_dir: &Path) -> Result<()> { + for entry_name in [ + "db.sqlite", + "db.sqlite-wal", + "db.sqlite-shm", + "db.sqlite-journal", + "blobs", + ] { + let source = page_root.join(entry_name); + let metadata = match std::fs::symlink_metadata(&source) { + Ok(metadata) => metadata, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue, + Err(error) => { + return Err(anyhow!( + "inspect legacy PageData path {}: {error}", + source.display() + )) + } + }; + if metadata.file_type().is_symlink() { + return Err(anyhow!( + "legacy PageData path must not be a symlink: {}", + source.display() + )); + } + if entry_name == "blobs" { + create_legacy_blob_layout_marker(generation_dir)?; + } + let target = generation_dir.join(entry_name); + if target.exists() { + return Err(anyhow!( + "legacy PageData migration target already exists: {}", + target.display() + )); + } + std::fs::rename(&source, &target).map_err(|e| { + anyhow!( + "move legacy PageData {} to generation: {e}", + source.display() + ) + })?; + } + Ok(()) +} + +fn create_legacy_blob_layout_marker(generation_dir: &Path) -> Result<()> { + let marker = generation_dir.join(LEGACY_BLOB_LAYOUT_MARKER); + ensure_regular_file_or_missing(&marker)?; + match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&marker) + { + Ok(_) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(()), + Err(error) => Err(anyhow!( + "create legacy PageData blob-layout marker {}: {error}", + marker.display() + )), + } +} + +fn has_legacy_blob_layout(generation_dir: &Path) -> Result { + let marker = generation_dir.join(LEGACY_BLOB_LAYOUT_MARKER); + match std::fs::symlink_metadata(&marker) { + Ok(metadata) if metadata.file_type().is_symlink() => Err(anyhow!( + "legacy PageData blob-layout marker must not be a symlink" + )), + Ok(metadata) if metadata.is_file() => Ok(true), + Ok(_) => Err(anyhow!( + "legacy PageData blob-layout marker must be a regular file" + )), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(anyhow!( + "inspect legacy PageData blob-layout marker: {error}" + )), + } +} + +/// Move one pre-hash blob into the portable hashed layout while the caller +/// holds the account mutation lock. A full-directory rewrite cannot safely +/// distinguish `foo.meta` (legacy metadata for `foo`) from a real blob whose +/// id is `foo.meta`, so migration is intentionally lazy and id-scoped. +fn migrate_legacy_blob_locked( + generation_dir: &Path, + blobs_dir: &Path, + blob_id: &str, +) -> Result<()> { + if !has_legacy_blob_layout(generation_dir)? || !is_legacy_blob_id(blob_id) { + return Ok(()); + } + + let storage_name = blob_storage_name(blob_id); + let legacy_blobs_dir = generation_dir.join("blobs"); + + migrate_legacy_blob_file( + &legacy_blobs_dir.join(blob_id), + &blobs_dir.join(&storage_name), + "data", + )?; + migrate_legacy_blob_file( + &legacy_blobs_dir.join(".metadata").join(blob_id), + &blob_meta_path(blobs_dir, &storage_name), + "metadata", + )?; + Ok(()) +} + +fn migrate_legacy_blob_file(source: &Path, target: &Path, kind: &str) -> Result<()> { + let source_metadata = match std::fs::symlink_metadata(source) { + Ok(metadata) => metadata, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(anyhow!("inspect legacy blob {kind}: {error}")), + }; + if source_metadata.file_type().is_symlink() || !source_metadata.is_file() { + return Err(anyhow!("legacy blob {kind} must be a regular file")); + } + + match std::fs::symlink_metadata(target) { + Ok(target_metadata) => { + if target_metadata.file_type().is_symlink() || !target_metadata.is_file() { + return Err(anyhow!("hashed blob {kind} must be a regular file")); + } + if !regular_files_equal(source, target, source_metadata.len(), target_metadata.len())? { + return Err(anyhow!( + "legacy blob {kind} conflicts with existing hashed storage" + )); + } + std::fs::remove_file(source) + .map_err(|error| anyhow!("remove duplicate legacy blob {kind}: {error}"))?; + Ok(()) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + std::fs::rename(source, target) + .map_err(|error| anyhow!("move legacy blob {kind} to hashed storage: {error}")) + } + Err(error) => Err(anyhow!("inspect hashed blob {kind}: {error}")), + } +} + +fn regular_files_equal(left: &Path, right: &Path, left_len: u64, right_len: u64) -> Result { + if left_len != right_len { + return Ok(false); + } + let left = std::fs::read(left).map_err(|error| anyhow!("read legacy blob: {error}"))?; + let right = std::fs::read(right).map_err(|error| anyhow!("read hashed blob: {error}"))?; + Ok(left == right) +} + +fn legacy_blob_content_type(generation_dir: &Path, blob_id: &str) -> Option { + if !is_legacy_blob_id(blob_id) || !has_legacy_blob_layout(generation_dir).ok()? { + return None; + } + let legacy_blobs_dir = generation_dir.join("blobs"); + read_small_metadata(&legacy_blobs_dir.join(".metadata").join(blob_id)) + // The sibling form is read-only during migration: deleting or moving + // it could destroy a distinct blob whose id happens to end in `.meta`. + .or_else(|| read_small_metadata(&legacy_blobs_dir.join(format!("{blob_id}.meta")))) +} + +fn ensure_directory(path: &Path) -> Result<()> { + match std::fs::symlink_metadata(path) { + Ok(metadata) if metadata.file_type().is_symlink() => Err(anyhow!( + "page-data directory must not be a symlink: {}", + path.display() + )), + Ok(metadata) if metadata.is_dir() => Ok(()), + Ok(_) => Err(anyhow!( + "page-data directory path is not a directory: {}", + path.display() + )), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => std::fs::create_dir(path) + .map_err(|e| anyhow!("create page-data directory {}: {e}", path.display())), + Err(error) => Err(anyhow!( + "inspect page-data directory {}: {error}", + path.display() + )), + } +} + +fn existing_directory(path: &Path) -> Result> { + match std::fs::symlink_metadata(path) { + Ok(metadata) if metadata.file_type().is_symlink() => Err(anyhow!( + "page-data directory must not be a symlink: {}", + path.display() + )), + Ok(metadata) if metadata.is_dir() => Ok(Some(path.to_path_buf())), + Ok(_) => Err(anyhow!( + "page-data directory path is not a directory: {}", + path.display() + )), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(anyhow!( + "inspect page-data directory {}: {error}", + path.display() + )), + } +} + +fn ensure_regular_file_or_missing(path: &Path) -> Result<()> { + match std::fs::symlink_metadata(path) { + Ok(metadata) if metadata.file_type().is_symlink() => Err(anyhow!( + "page-data file must not be a symlink: {}", + path.display() + )), + Ok(metadata) if metadata.is_file() => Ok(()), + Ok(_) => Err(anyhow!( + "page-data file path is not a regular file: {}", + path.display() + )), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(anyhow!( + "inspect page-data file {}: {error}", + path.display() + )), + } +} + +fn validate_blob_id(blob_id: &str) -> Result<()> { + if blob_id.is_empty() + || blob_id.len() > MAX_BLOB_ID_BYTES + || blob_id.chars().any(char::is_control) + { + return Err(anyhow!("invalid blob id")); + } + Ok(()) +} + +fn is_legacy_blob_id(blob_id: &str) -> bool { + if blob_id.is_empty() + || blob_id.len() > MAX_BLOB_ID_BYTES + || matches!(blob_id, "." | ".." | ".metadata") + || blob_id.contains('/') + || blob_id.contains('\\') + || blob_id.chars().any(char::is_control) + { + return false; + } + + #[cfg(windows)] + { + if blob_id + .chars() + .any(|character| matches!(character, '<' | '>' | ':' | '"' | '|' | '?' | '*')) + || blob_id.ends_with(' ') + || blob_id.ends_with('.') + { + return false; + } + let stem = blob_id + .split('.') + .next() + .unwrap_or_default() + .to_ascii_uppercase(); + if matches!(stem.as_str(), "CON" | "PRN" | "AUX" | "NUL") + || (stem.len() == 4 + && (stem.starts_with("COM") || stem.starts_with("LPT")) + && stem.as_bytes()[3].is_ascii_digit() + && stem.as_bytes()[3] != b'0') + { + return false; + } + } + + true +} + +fn blob_storage_name(blob_id: &str) -> String { + Sha256::digest(blob_id.as_bytes()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn validate_content_type(content_type: &str) -> Result<()> { + if content_type.is_empty() + || content_type.len() > 256 + || content_type.chars().any(char::is_control) + { + return Err(anyhow!("invalid blob content type")); + } + Ok(()) +} + +fn blob_meta_path(blobs_dir: &Path, storage_name: &str) -> PathBuf { + blobs_dir.join(".metadata").join(storage_name) +} + +fn file_len(path: &Path) -> u64 { + path.metadata().map_or(0, |metadata| metadata.len()) +} + +fn read_small_metadata(path: &Path) -> Option { + if file_len(path) > 256 { + return None; + } + std::fs::read_to_string(path).ok() +} + +fn directory_size(path: &Path) -> Result { + if !path.exists() { + return Ok(0); + } + let metadata = std::fs::symlink_metadata(path) + .map_err(|e| anyhow!("inspect page-data path {}: {e}", path.display()))?; + if metadata.file_type().is_symlink() { + return Err(anyhow!( + "page-data path must not contain symlinks: {}", + path.display() + )); + } + if metadata.is_file() { + return Ok(metadata.len()); + } + if !metadata.is_dir() { + return Ok(0); + } + + let mut total = 0u64; + for entry in std::fs::read_dir(path) + .map_err(|e| anyhow!("read page-data directory {}: {e}", path.display()))? + { + let entry = entry.map_err(|e| anyhow!("read page-data entry: {e}"))?; + total = total + .checked_add(directory_size(&entry.path())?) + .ok_or_else(|| anyhow!("page-data size overflow"))?; + } + Ok(total) +} + +fn direct_file_count(path: &Path) -> Result { + if !path.exists() { + return Ok(0); + } + let mut count = 0u64; + for entry in std::fs::read_dir(path) + .map_err(|e| anyhow!("read blob directory {}: {e}", path.display()))? + { + let entry = entry.map_err(|e| anyhow!("read blob entry: {e}"))?; + let metadata = std::fs::symlink_metadata(entry.path()) + .map_err(|e| anyhow!("inspect blob entry: {e}"))?; + if metadata.file_type().is_symlink() { + return Err(anyhow!("blob directory must not contain symlinks")); + } + if metadata.is_file() { + count = count + .checked_add(1) + .ok_or_else(|| anyhow!("blob file count overflow"))?; + } + } + Ok(count) +} + +fn page_blob_file_count(page_dir: &Path) -> Result { + direct_file_count(&page_dir.join("blobs"))? + .checked_add(direct_file_count(&page_dir.join("blobs-v2"))?) + .ok_or_else(|| anyhow!("page blob file count overflow")) +} + +fn user_blob_file_count(user_dir: &Path) -> Result { + if !user_dir.exists() { + return Ok(0); + } + let mut count = 0u64; + for page in + std::fs::read_dir(user_dir).map_err(|e| anyhow!("read account page-data directory: {e}"))? + { + let page = page.map_err(|e| anyhow!("read account page-data entry: {e}"))?; + let metadata = std::fs::symlink_metadata(page.path()) + .map_err(|e| anyhow!("inspect account page-data entry: {e}"))?; + if metadata.file_type().is_symlink() { + return Err(anyhow!("page-data directory must not contain symlinks")); + } + if metadata.is_dir() { + // Count the pre-generation layout until it has been migrated. + count = count + .checked_add(page_blob_file_count(&page.path())?) + .ok_or_else(|| anyhow!("account blob file count overflow"))?; + for generation in std::fs::read_dir(page.path()) + .map_err(|e| anyhow!("read PageData generation directory: {e}"))? + { + let generation = + generation.map_err(|e| anyhow!("read PageData generation entry: {e}"))?; + if !is_generation_name(&generation.file_name()) { + continue; + } + let generation_metadata = std::fs::symlink_metadata(generation.path()) + .map_err(|e| anyhow!("inspect PageData generation entry: {e}"))?; + if generation_metadata.file_type().is_symlink() { + return Err(anyhow!("PageData generation must not be a symlink")); + } + if !generation_metadata.is_dir() { + return Err(anyhow!("PageData generation path must be a directory")); + } + count = count + .checked_add(page_blob_file_count(&generation.path())?) + .ok_or_else(|| anyhow!("account blob file count overflow"))?; + } + } + } + Ok(count) +} + +fn enforce_storage_quota( + page_bytes: u64, + user_bytes: u64, + replaced_bytes: u64, + added_bytes: u64, +) -> Result<()> { + let projected_page = page_bytes + .saturating_sub(replaced_bytes) + .checked_add(added_bytes) + .ok_or_else(|| anyhow!("page-data size overflow"))?; + if projected_page > MAX_MUTABLE_BYTES_PER_PAGE { + return Err(anyhow!( + "page mutable storage exceeds the {} byte quota", + MAX_MUTABLE_BYTES_PER_PAGE + )); + } + let projected_user = user_bytes + .saturating_sub(replaced_bytes) + .checked_add(added_bytes) + .ok_or_else(|| anyhow!("account page-data size overflow"))?; + if projected_user > MAX_MUTABLE_BYTES_PER_USER { + return Err(anyhow!( + "account mutable storage exceeds the {} byte quota", + MAX_MUTABLE_BYTES_PER_USER + )); + } + Ok(()) +} + +fn enforce_current_storage_quota( + store: &PageDataStore, + user_id: &str, + slug: &str, + generation: &str, +) -> Result<()> { + enforce_storage_quota( + directory_size(&store.page_dir(user_id, slug, generation))?, + directory_size(&store.base_dir.join(user_id))?, + 0, + 0, + ) +} + +fn validate_db_input(sql: &str, params_json: &str) -> Result<()> { + if sql.trim().is_empty() || sql.len() > MAX_DB_SQL_BYTES { + return Err(anyhow!( + "SQL must be non-empty and at most {} bytes", + MAX_DB_SQL_BYTES + )); + } + if params_json.len() > MAX_DB_PARAMS_BYTES { + return Err(anyhow!( + "SQL parameters exceed the {} byte limit", + MAX_DB_PARAMS_BYTES + )); + } + Ok(()) +} + +fn parse_db_params(params_json: &str) -> Result> { + let params: Vec = + serde_json::from_str(params_json).map_err(|e| anyhow!("invalid SQL parameters: {e}"))?; + if params.len() > MAX_DB_PARAMS { + return Err(anyhow!( + "SQL parameters exceed the {} item limit", + MAX_DB_PARAMS + )); + } + Ok(params) +} + +fn configure_connection(conn: &rusqlite::Connection) { + conn.set_limit(Limit::SQLITE_LIMIT_LENGTH, MAX_DB_VALUE_BYTES); + conn.set_limit(Limit::SQLITE_LIMIT_SQL_LENGTH, MAX_DB_SQL_BYTES as i32); + conn.set_limit(Limit::SQLITE_LIMIT_COLUMN, 128); + conn.set_limit(Limit::SQLITE_LIMIT_EXPR_DEPTH, 64); + conn.set_limit(Limit::SQLITE_LIMIT_COMPOUND_SELECT, 16); + conn.set_limit(Limit::SQLITE_LIMIT_FUNCTION_ARG, 32); + conn.set_limit(Limit::SQLITE_LIMIT_ATTACHED, 0); + conn.set_limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER, MAX_DB_PARAMS as i32); + conn.set_limit(Limit::SQLITE_LIMIT_TRIGGER_DEPTH, 8); + conn.set_limit(Limit::SQLITE_LIMIT_WORKER_THREADS, 0); + let started = Instant::now(); + conn.progress_handler( + 1_000, + Some(move || started.elapsed() >= MAX_DB_OPERATION_TIME), + ); + let _ = conn.busy_timeout(Duration::from_millis(250)); +} + +fn configure_database_quota( + store: &PageDataStore, + conn: &rusqlite::Connection, + user_id: &str, + slug: &str, + generation: &str, + db_path: &Path, +) -> Result<()> { + let page_bytes = directory_size(&store.page_dir(user_id, slug, generation))?; + let user_bytes = directory_size(&store.base_dir.join(user_id))?; + enforce_storage_quota(page_bytes, user_bytes, 0, 0)?; + + let current_db_bytes = file_len(db_path); + if current_db_bytes > MAX_PAGE_DB_BYTES { + return Err(anyhow!("page database already exceeds its quota")); + } + let remaining_page = MAX_MUTABLE_BYTES_PER_PAGE.saturating_sub(page_bytes); + let remaining_user = MAX_MUTABLE_BYTES_PER_USER.saturating_sub(user_bytes); + let allowed_db_bytes = + MAX_PAGE_DB_BYTES.min(current_db_bytes.saturating_add(remaining_page.min(remaining_user))); + let page_size: u64 = conn + .query_row("PRAGMA page_size", [], |row| row.get(0)) + .map_err(|e| anyhow!("read database page size: {e}"))?; + let max_pages = (allowed_db_bytes / page_size).max(1); + let _: u64 = conn + .query_row(&format!("PRAGMA max_page_count = {max_pages}"), [], |row| { + row.get(0) + }) + .map_err(|e| anyhow!("set database quota: {e}"))?; + Ok(()) +} + +fn database_logical_bytes(conn: &rusqlite::Connection) -> Result { + let page_count: u64 = conn + .query_row("PRAGMA page_count", [], |row| row.get(0)) + .map_err(|e| anyhow!("read database page count: {e}"))?; + let page_size: u64 = conn + .query_row("PRAGMA page_size", [], |row| row.get(0)) + .map_err(|e| anyhow!("read database page size: {e}"))?; + Ok(page_count.saturating_mul(page_size)) +} + +fn authorize_execute(context: AuthContext<'_>) -> Authorization { + match context.action { + AuthAction::Attach { .. } + | AuthAction::Detach { .. } + | AuthAction::Pragma { .. } + | AuthAction::Transaction { .. } + | AuthAction::Savepoint { .. } + | AuthAction::CreateTempIndex { .. } + | AuthAction::CreateTempTable { .. } + | AuthAction::CreateTempTrigger { .. } + | AuthAction::CreateTempView { .. } + | AuthAction::DropTempIndex { .. } + | AuthAction::DropTempTable { .. } + | AuthAction::DropTempTrigger { .. } + | AuthAction::DropTempView { .. } + | AuthAction::CreateVtable { .. } + | AuthAction::DropVtable { .. } + | AuthAction::Analyze { .. } + | AuthAction::Reindex { .. } + | AuthAction::Unknown { .. } => Authorization::Deny, + AuthAction::Function { function_name } if forbidden_db_function(function_name) => { + Authorization::Deny + } + _ => Authorization::Allow, + } +} + +fn authorize_query(context: AuthContext<'_>) -> Authorization { + match context.action { + AuthAction::Read { .. } | AuthAction::Select | AuthAction::Recursive => { + Authorization::Allow + } + AuthAction::Function { function_name } if !forbidden_db_function(function_name) => { + Authorization::Allow + } + _ => Authorization::Deny, + } +} + +fn forbidden_db_function(function_name: &str) -> bool { + matches!( + function_name.to_ascii_lowercase().as_str(), + "load_extension" | "readfile" | "writefile" | "edit" | "shell" + ) +} + fn json_to_sql(v: &serde_json::Value) -> rusqlite::types::Value { match v { serde_json::Value::Null => rusqlite::types::Value::Null, @@ -196,6 +1015,7 @@ pub struct RelayPageHost { pub page_data: PageDataStore, pub user_id: String, pub slug: String, + pub generation: String, pub meta: PageMeta, pub asset_store: Arc, pub asset_key: String, @@ -220,6 +1040,21 @@ impl RelayPageHost { } } } + + fn ensure_current_page(&self) -> Result<(), String> { + let db = Arc::clone(&self.db); + let user_id = self.user_id.clone(); + let slug = self.slug.clone(); + let generation = self.generation.clone(); + self.block_on_kv(async move { + let page = PageRow::get(&db, &user_id, &slug).await?; + if page.is_some_and(|page| page.generation == generation) { + Ok(()) + } else { + Err(anyhow!("page no longer exists")) + } + }) + } } impl PageHost for RelayPageHost { @@ -232,6 +1067,11 @@ impl PageHost for RelayPageHost { } fn kv_put(&self, key: &str, value: &str) -> Result<(), String> { + self.ensure_current_page()?; + let lock = self.page_data.user_mutation_lock(&self.user_id); + let _guard = lock + .lock() + .map_err(|_| "page-data mutation lock poisoned".to_string())?; let db = Arc::clone(&self.db); let user_id = self.user_id.clone(); let slug = self.slug.clone(); @@ -241,6 +1081,7 @@ impl PageHost for RelayPageHost { } fn kv_delete(&self, key: &str) -> Result { + self.ensure_current_page()?; let db = Arc::clone(&self.db); let user_id = self.user_id.clone(); let slug = self.slug.clone(); @@ -256,28 +1097,55 @@ impl PageHost for RelayPageHost { } fn db_execute(&self, sql: &str, params_json: &str) -> Result { + self.ensure_current_page()?; self.page_data - .db_execute(&self.user_id, &self.slug, sql, params_json) + .db_execute( + &self.user_id, + &self.slug, + &self.generation, + sql, + params_json, + ) .map_err(|e| e.to_string()) } fn db_query(&self, sql: &str, params_json: &str) -> Result { self.page_data - .db_query(&self.user_id, &self.slug, sql, params_json) + .db_query( + &self.user_id, + &self.slug, + &self.generation, + sql, + params_json, + ) .map_err(|e| e.to_string()) } fn blob_put(&self, blob_id: &str, content_type: &str, data_b64: &str) -> Result<(), String> { + self.ensure_current_page()?; + const MAX_ENCODED_BLOB_BYTES: usize = MAX_BLOB_BYTES.div_ceil(3) * 4; + if data_b64.len() > MAX_ENCODED_BLOB_BYTES { + return Err(format!( + "encoded blob exceeds the {MAX_ENCODED_BLOB_BYTES} byte operation limit" + )); + } let data = B64.decode(data_b64).map_err(|e| e.to_string())?; self.page_data - .blob_put(&self.user_id, &self.slug, blob_id, content_type, &data) + .blob_put( + &self.user_id, + &self.slug, + &self.generation, + blob_id, + content_type, + &data, + ) .map_err(|e| e.to_string()) } fn blob_get(&self, blob_id: &str) -> Result, String> { match self .page_data - .blob_get(&self.user_id, &self.slug, blob_id) + .blob_get(&self.user_id, &self.slug, &self.generation, blob_id) .map_err(|e| e.to_string())? { Some((ct, data)) => Ok(Some((ct, B64.encode(data)))), @@ -286,8 +1154,9 @@ impl PageHost for RelayPageHost { } fn blob_delete(&self, blob_id: &str) -> Result { + self.ensure_current_page()?; self.page_data - .blob_delete(&self.user_id, &self.slug, blob_id) + .blob_delete(&self.user_id, &self.slug, &self.generation, blob_id) .map_err(|e| e.to_string()) } @@ -323,3 +1192,431 @@ fn mime_from_path(p: &str) -> &'static str { pub fn default_page_data_dir(room_web_dir: &Path) -> PathBuf { room_web_dir.join("page-data") } + +#[cfg(test)] +mod tests { + use super::*; + + const GENERATION_A: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + const GENERATION_B: &str = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; + + #[test] + fn storage_projection_enforces_page_and_account_quotas() { + assert!(enforce_storage_quota( + MAX_MUTABLE_BYTES_PER_PAGE, + MAX_MUTABLE_BYTES_PER_PAGE, + 10, + 10, + ) + .is_ok()); + assert!(enforce_storage_quota( + MAX_MUTABLE_BYTES_PER_PAGE, + MAX_MUTABLE_BYTES_PER_PAGE, + 0, + 1, + ) + .is_err()); + assert!(enforce_storage_quota(0, MAX_MUTABLE_BYTES_PER_USER, 0, 1,).is_err()); + } + + #[test] + fn blob_limits_are_enforced_and_metadata_names_do_not_collide() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + let oversized = vec![0u8; MAX_BLOB_BYTES + 1]; + assert!(store + .blob_put( + "u1", + "site", + GENERATION_A, + "too-big", + "application/octet-stream", + &oversized + ) + .unwrap_err() + .to_string() + .contains("operation limit")); + + store + .blob_put("u1", "site", GENERATION_A, "item", "text/plain", b"first") + .unwrap(); + store + .blob_put( + "u1", + "site", + GENERATION_A, + "item.meta", + "application/octet-stream", + b"second", + ) + .unwrap(); + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap() + .unwrap(), + ("text/plain".to_string(), b"first".to_vec()) + ); + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, "item.meta") + .unwrap() + .unwrap(), + ("application/octet-stream".to_string(), b"second".to_vec()) + ); + } + + #[test] + fn blob_ids_use_portable_collision_free_internal_names() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + for (id, value) in [ + ("A", b"upper".as_slice()), + ("a", b"lower".as_slice()), + ("é", b"composed".as_slice()), + ("e\u{301}", b"decomposed".as_slice()), + ] { + store + .blob_put("u1", "site", GENERATION_A, id, "text/plain", value) + .unwrap(); + } + for (id, expected) in [ + ("A", b"upper".as_slice()), + ("a", b"lower".as_slice()), + ("é", b"composed".as_slice()), + ("e\u{301}", b"decomposed".as_slice()), + ] { + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, id) + .unwrap() + .unwrap() + .1, + expected + ); + } + assert_ne!(blob_storage_name("A"), blob_storage_name("a")); + assert_ne!(blob_storage_name("é"), blob_storage_name("e\u{301}")); + } + + #[test] + fn legacy_blob_layout_is_lazily_migrated_and_remains_mutable() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + let legacy_blobs = temp.path().join("u1/site/blobs"); + std::fs::create_dir_all(legacy_blobs.join(".metadata")).unwrap(); + std::fs::write(legacy_blobs.join("item"), b"legacy").unwrap(); + std::fs::write(legacy_blobs.join(".metadata/item"), "text/plain").unwrap(); + + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap() + .unwrap(), + ("text/plain".to_string(), b"legacy".to_vec()) + ); + let migrated_blobs = store.blobs_dir("u1", "site", GENERATION_A); + assert!(!store + .page_dir("u1", "site", GENERATION_A) + .join("blobs/item") + .exists()); + assert!(migrated_blobs.join(blob_storage_name("item")).is_file()); + assert!(store + .page_dir("u1", "site", GENERATION_A) + .join(LEGACY_BLOB_LAYOUT_MARKER) + .is_file()); + + // Repeating migration is a no-op; writes and deletes continue through + // the portable storage name after the upgrade. + assert!(store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap() + .is_some()); + store + .blob_put( + "u1", + "site", + GENERATION_A, + "item", + "application/json", + b"{}", + ) + .unwrap(); + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap() + .unwrap(), + ("application/json".to_string(), b"{}".to_vec()) + ); + assert!(store + .blob_delete("u1", "site", GENERATION_A, "item") + .unwrap()); + assert!(store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap() + .is_none()); + } + + #[test] + fn legacy_sibling_metadata_is_read_without_destroying_another_possible_blob() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + let legacy_blobs = temp.path().join("u1/site/blobs"); + std::fs::create_dir_all(&legacy_blobs).unwrap(); + std::fs::write(legacy_blobs.join("item"), b"legacy").unwrap(); + std::fs::write(legacy_blobs.join("item.meta"), "text/custom").unwrap(); + + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap() + .unwrap(), + ("text/custom".to_string(), b"legacy".to_vec()) + ); + let legacy_sibling = store + .page_dir("u1", "site", GENERATION_A) + .join("blobs/item.meta"); + assert!(legacy_sibling.is_file()); + + store + .blob_put("u1", "site", GENERATION_A, "item", "text/new", b"new") + .unwrap(); + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap() + .unwrap() + .0, + "text/new" + ); + assert!(store + .blob_delete("u1", "site", GENERATION_A, "item") + .unwrap()); + assert!(legacy_sibling.is_file()); + } + + #[test] + fn legacy_blob_conflict_never_overwrites_different_hashed_data() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + let legacy_blobs = temp.path().join("u1/site/blobs"); + std::fs::create_dir_all(&legacy_blobs).unwrap(); + let hashed_name = blob_storage_name("item"); + std::fs::write(legacy_blobs.join("item"), b"legacy").unwrap(); + let current_blobs = temp + .path() + .join("u1/site") + .join(GENERATION_A) + .join("blobs-v2"); + std::fs::create_dir_all(¤t_blobs).unwrap(); + std::fs::write(current_blobs.join(&hashed_name), b"current").unwrap(); + + let error = store + .blob_get("u1", "site", GENERATION_A, "item") + .unwrap_err(); + assert!(error.to_string().contains("conflicts")); + let migrated_blobs = store.blobs_dir("u1", "site", GENERATION_A); + assert_eq!( + std::fs::read( + store + .page_dir("u1", "site", GENERATION_A) + .join("blobs/item") + ) + .unwrap(), + b"legacy" + ); + assert_eq!( + std::fs::read(migrated_blobs.join(hashed_name)).unwrap(), + b"current" + ); + } + + #[test] + fn legacy_id_equal_to_another_ids_hash_stays_isolated_in_v2_storage() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + let id_b = "B"; + let id_a = blob_storage_name(id_b); + let legacy_blobs = temp.path().join("u1/site/blobs"); + std::fs::create_dir_all(legacy_blobs.join(".metadata")).unwrap(); + std::fs::write(legacy_blobs.join(&id_a), b"legacy-a").unwrap(); + std::fs::write(legacy_blobs.join(".metadata").join(&id_a), "text/a").unwrap(); + + store + .blob_put("u1", "site", GENERATION_A, id_b, "text/b", b"current-b") + .unwrap(); + let page_dir = store.page_dir("u1", "site", GENERATION_A); + assert_eq!(page_blob_file_count(&page_dir).unwrap(), 2); + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, &id_a) + .unwrap() + .unwrap(), + ("text/a".to_string(), b"legacy-a".to_vec()) + ); + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, id_b) + .unwrap() + .unwrap(), + ("text/b".to_string(), b"current-b".to_vec()) + ); + assert_eq!(page_blob_file_count(&page_dir).unwrap(), 2); + + assert!(store.blob_delete("u1", "site", GENERATION_A, id_b).unwrap()); + assert!(store + .blob_get("u1", "site", GENERATION_A, id_b) + .unwrap() + .is_none()); + assert_eq!( + store + .blob_get("u1", "site", GENERATION_A, &id_a) + .unwrap() + .unwrap() + .1, + b"legacy-a" + ); + assert!(store + .blob_delete("u1", "site", GENERATION_A, &id_a) + .unwrap()); + } + + #[test] + fn ids_that_were_not_legacy_filenames_use_only_hashed_storage() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + store + .blob_put( + "u1", + "site", + GENERATION_A, + "folder/item:portable", + "text/plain", + b"value", + ) + .unwrap(); + let blobs = store.blobs_dir("u1", "site", GENERATION_A); + assert!(blobs + .join(blob_storage_name("folder/item:portable")) + .is_file()); + assert!(!is_legacy_blob_id("folder/item:portable")); + } + + #[test] + fn database_denies_file_attachment_and_caps_query_rows() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path().join("page-data")); + let attached = temp.path().join("outside.sqlite"); + let error = store + .db_execute( + "u1", + "site", + GENERATION_A, + &format!("ATTACH DATABASE '{}' AS outside", attached.display()), + "[]", + ) + .unwrap_err(); + assert!( + error.to_string().contains("not authorized") + || error.to_string().contains("too many attached") + ); + assert!(!attached.exists()); + + let error = store + .db_query( + "u1", + "site", + GENERATION_A, + "WITH RECURSIVE items(n) AS (VALUES(1) UNION ALL SELECT n + 1 FROM items WHERE n <= 1000) SELECT n FROM items", + "[]", + ) + .unwrap_err(); + assert!(error.to_string().contains("row result limit")); + } + + #[test] + fn database_execute_accepts_normal_schema_and_data_changes() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + store + .db_execute( + "u1", + "site", + GENERATION_A, + "CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT)", + "[]", + ) + .unwrap(); + store + .db_execute( + "u1", + "site", + GENERATION_A, + "INSERT INTO notes (body) VALUES (?)", + r#"["hello"]"#, + ) + .unwrap(); + let output = store + .db_query("u1", "site", GENERATION_A, "SELECT body FROM notes", "[]") + .unwrap(); + let json: serde_json::Value = serde_json::from_str(&output).unwrap(); + assert_eq!(json["rows"][0]["body"], "hello"); + } + + #[test] + fn delete_recreate_generation_cannot_observe_old_page_data() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + store + .blob_put("u1", "site", GENERATION_A, "secret", "text/plain", b"old") + .unwrap(); + store + .db_execute( + "u1", + "site", + GENERATION_A, + "CREATE TABLE old_data (value TEXT)", + "[]", + ) + .unwrap(); + + // Simulate a process crash after relational deletion but before the + // best-effort filesystem cleanup. The recreated Page gets generation B. + assert!(store + .blob_get("u1", "site", GENERATION_B, "secret") + .unwrap() + .is_none()); + let output = store + .db_query( + "u1", + "site", + GENERATION_B, + "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'old_data'", + "[]", + ) + .unwrap(); + let json: serde_json::Value = serde_json::from_str(&output).unwrap(); + assert_eq!(json["rows"], serde_json::json!([])); + assert!(store.page_dir("u1", "site", GENERATION_A).exists()); + assert!(store.page_dir("u1", "site", GENERATION_B).exists()); + } + + #[test] + fn database_long_running_query_is_interrupted() { + let temp = tempfile::tempdir().unwrap(); + let store = PageDataStore::new(temp.path()); + let started = Instant::now(); + let error = store + .db_query( + "u1", + "site", + GENERATION_A, + "WITH RECURSIVE items(n) AS (VALUES(1) UNION ALL SELECT n + 1 FROM items WHERE n < 1000000000) SELECT sum(n) FROM items", + "[]", + ) + .unwrap_err(); + assert!(error.to_string().contains("interrupted")); + assert!(started.elapsed() < Duration::from_secs(1)); + } +} diff --git a/src/crates/services/relay-service/src/page_execution.rs b/src/crates/services/relay-service/src/page_execution.rs new file mode 100644 index 0000000000..7d4f4d4772 --- /dev/null +++ b/src/crates/services/relay-service/src/page_execution.rs @@ -0,0 +1,310 @@ +//! Admission control for public Page Function execution. + +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use dashmap::DashMap; +use tokio::sync::{ + OwnedRwLockReadGuard, OwnedRwLockWriteGuard, OwnedSemaphorePermit, RwLock, Semaphore, +}; + +pub const MAX_PAGE_FUNCTION_REQUEST_BODY_BYTES: usize = 1024 * 1024; +const GLOBAL_CONCURRENCY: usize = 64; +const USER_CONCURRENCY: usize = 16; +const PAGE_CONCURRENCY: usize = 8; +const USER_REQUESTS_PER_WINDOW: usize = 3_000; +const PAGE_REQUESTS_PER_WINDOW: usize = 600; +const RATE_WINDOW: Duration = Duration::from_secs(60); +const MAX_TRACKED_IDENTITIES: usize = 10_000; + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum PageExecutionRejection { + Busy, + RateLimited, +} + +#[derive(Clone, Copy)] +struct Limits { + global_concurrency: usize, + user_concurrency: usize, + page_concurrency: usize, + user_requests_per_window: usize, + page_requests_per_window: usize, + rate_window: Duration, +} + +impl Default for Limits { + fn default() -> Self { + Self { + global_concurrency: GLOBAL_CONCURRENCY, + user_concurrency: USER_CONCURRENCY, + page_concurrency: PAGE_CONCURRENCY, + user_requests_per_window: USER_REQUESTS_PER_WINDOW, + page_requests_per_window: PAGE_REQUESTS_PER_WINDOW, + rate_window: RATE_WINDOW, + } + } +} + +/// Process-local guard against one account or page exhausting blocking workers. +pub struct PageExecutionGuard { + limits: Limits, + global: Arc, + users: DashMap>, + pages: DashMap>, + page_lifecycles: DashMap>>, + rate_windows: Mutex>>, +} + +impl Default for PageExecutionGuard { + fn default() -> Self { + Self::new() + } +} + +impl PageExecutionGuard { + pub fn new() -> Self { + Self::with_limits(Limits::default()) + } + + fn with_limits(limits: Limits) -> Self { + Self { + global: Arc::new(Semaphore::new(limits.global_concurrency)), + users: DashMap::new(), + pages: DashMap::new(), + page_lifecycles: DashMap::new(), + rate_windows: Mutex::new(HashMap::new()), + limits, + } + } + + pub fn try_acquire( + &self, + user_id: &str, + slug: &str, + ) -> Result { + self.check_rate(user_id, slug)?; + + let global = Arc::clone(&self.global) + .try_acquire_owned() + .map_err(|_| PageExecutionRejection::Busy)?; + let user = self.user_semaphore(user_id); + let user = user + .try_acquire_owned() + .map_err(|_| PageExecutionRejection::Busy)?; + let page = self.page_semaphore(user_id, slug); + let page = page + .try_acquire_owned() + .map_err(|_| PageExecutionRejection::Busy)?; + + Ok(PageExecutionPermit { + _global: global, + _user: user, + _page: page, + }) + } + + fn user_semaphore(&self, user_id: &str) -> Arc { + self.prune_idle_semaphores(); + Arc::clone( + self.users + .entry(user_id.to_string()) + .or_insert_with(|| Arc::new(Semaphore::new(self.limits.user_concurrency))) + .value(), + ) + } + + fn page_semaphore(&self, user_id: &str, slug: &str) -> Arc { + self.prune_idle_semaphores(); + let key = format!("{user_id}\0{slug}"); + Arc::clone( + self.pages + .entry(key) + .or_insert_with(|| Arc::new(Semaphore::new(self.limits.page_concurrency))) + .value(), + ) + } + + fn page_lifecycle(&self, user_id: &str, slug: &str) -> Arc> { + self.prune_idle_semaphores(); + let key = format!("{user_id}\0{slug}"); + Arc::clone( + self.page_lifecycles + .entry(key) + .or_insert_with(|| Arc::new(RwLock::new(()))) + .value(), + ) + } + + /// Hold while serving one Page request. Delete and other lifecycle writes + /// wait for all running workers, while unrelated Pages remain independent. + pub async fn acquire_page_read(&self, user_id: &str, slug: &str) -> OwnedRwLockReadGuard<()> { + self.page_lifecycle(user_id, slug).read_owned().await + } + + /// Hold while deleting or atomically changing Page/version lifecycle + /// state. New workers wait until the mutation completes and then re-resolve + /// the Page row before execution. + pub async fn acquire_page_write(&self, user_id: &str, slug: &str) -> OwnedRwLockWriteGuard<()> { + self.page_lifecycle(user_id, slug).write_owned().await + } + + fn prune_idle_semaphores(&self) { + if self.users.len() > MAX_TRACKED_IDENTITIES { + self.users + .retain(|_, semaphore| Arc::strong_count(semaphore) > 1); + } + if self.pages.len() > MAX_TRACKED_IDENTITIES { + self.pages + .retain(|_, semaphore| Arc::strong_count(semaphore) > 1); + } + if self.page_lifecycles.len() > MAX_TRACKED_IDENTITIES { + self.page_lifecycles + .retain(|_, lifecycle| Arc::strong_count(lifecycle) > 1); + } + } + + fn check_rate(&self, user_id: &str, slug: &str) -> Result<(), PageExecutionRejection> { + let now = Instant::now(); + let cutoff = now.checked_sub(self.limits.rate_window).unwrap_or(now); + let mut windows = self + .rate_windows + .lock() + .map_err(|_| PageExecutionRejection::Busy)?; + if windows.len() > MAX_TRACKED_IDENTITIES * 2 { + windows.retain(|_, entries| { + entries.retain(|timestamp| *timestamp > cutoff); + !entries.is_empty() + }); + } + + record_request( + &mut windows, + format!("user\0{user_id}"), + cutoff, + now, + self.limits.user_requests_per_window, + )?; + if let Err(error) = record_request( + &mut windows, + format!("page\0{user_id}\0{slug}"), + cutoff, + now, + self.limits.page_requests_per_window, + ) { + if let Some(user_window) = windows.get_mut(&format!("user\0{user_id}")) { + if user_window.back() == Some(&now) { + user_window.pop_back(); + } + } + return Err(error); + } + Ok(()) + } +} + +fn record_request( + windows: &mut HashMap>, + key: String, + cutoff: Instant, + now: Instant, + limit: usize, +) -> Result<(), PageExecutionRejection> { + let entries = windows.entry(key).or_default(); + while entries + .front() + .is_some_and(|timestamp| *timestamp <= cutoff) + { + entries.pop_front(); + } + if entries.len() >= limit { + return Err(PageExecutionRejection::RateLimited); + } + entries.push_back(now); + Ok(()) +} + +#[derive(Debug)] +pub struct PageExecutionPermit { + _global: OwnedSemaphorePermit, + _user: OwnedSemaphorePermit, + _page: OwnedSemaphorePermit, +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_guard() -> PageExecutionGuard { + PageExecutionGuard::with_limits(Limits { + global_concurrency: 3, + user_concurrency: 2, + page_concurrency: 1, + user_requests_per_window: 100, + page_requests_per_window: 100, + rate_window: Duration::from_secs(60), + }) + } + + #[test] + fn page_concurrency_is_isolated_and_permits_recover() { + let guard = test_guard(); + let first = guard.try_acquire("user", "one").unwrap(); + assert_eq!( + guard.try_acquire("user", "one").unwrap_err(), + PageExecutionRejection::Busy + ); + let second_page = guard.try_acquire("user", "two").unwrap(); + assert_eq!( + guard.try_acquire("user", "three").unwrap_err(), + PageExecutionRejection::Busy + ); + drop((first, second_page)); + assert!(guard.try_acquire("user", "one").is_ok()); + } + + #[test] + fn page_and_account_rate_windows_are_enforced() { + let guard = PageExecutionGuard::with_limits(Limits { + global_concurrency: 3, + user_concurrency: 2, + page_concurrency: 1, + user_requests_per_window: 4, + page_requests_per_window: 2, + rate_window: Duration::from_secs(60), + }); + drop(guard.try_acquire("user", "one").unwrap()); + drop(guard.try_acquire("user", "one").unwrap()); + assert_eq!( + guard.try_acquire("user", "one").unwrap_err(), + PageExecutionRejection::RateLimited + ); + + drop(guard.try_acquire("user", "two").unwrap()); + drop(guard.try_acquire("user", "two").unwrap()); + assert_eq!( + guard.try_acquire("user", "three").unwrap_err(), + PageExecutionRejection::RateLimited + ); + assert!(guard.try_acquire("other", "one").is_ok()); + } + + #[tokio::test] + async fn page_lifecycle_write_waits_for_running_readers() { + let guard = Arc::new(test_guard()); + let read = guard.acquire_page_read("user", "site").await; + let waiting_guard = Arc::clone(&guard); + let mut writer = + tokio::spawn(async move { waiting_guard.acquire_page_write("user", "site").await }); + assert!(tokio::time::timeout(Duration::from_millis(20), &mut writer) + .await + .is_err()); + drop(read); + let write = tokio::time::timeout(Duration::from_secs(1), writer) + .await + .unwrap() + .unwrap(); + drop(write); + } +} diff --git a/src/crates/services/relay-service/src/relay/device_manager.rs b/src/crates/services/relay-service/src/relay/device_manager.rs index 983cccae29..44eb927d95 100644 --- a/src/crates/services/relay-service/src/relay/device_manager.rs +++ b/src/crates/services/relay-service/src/relay/device_manager.rs @@ -11,7 +11,10 @@ //! arrives via WS from a device, the pending future is resolved. use dashmap::DashMap; -use std::sync::Arc; +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, MutexGuard, +}; use tokio::sync::{mpsc, oneshot, watch, OwnedSemaphorePermit, Semaphore}; use tracing::{debug, info}; @@ -23,6 +26,23 @@ pub const MAX_PENDING_DEVICE_RPCS: usize = 1024; struct DeviceConn { #[allow(dead_code)] conn_id: ConnId, + /// Bearer token that authenticated this exact socket. A device may have + /// more than one still-valid token (for example while a replacement login + /// is being verified), so logout must not disconnect the socket merely + /// because its `(user_id, device_id)` matches another token. + auth_token: String, + device_name: String, + tx: mpsc::Sender, + force_close_tx: watch::Sender, +} + +/// A socket that passed the initial token lookup but has not yet completed the +/// post-registration revocation check. Pending sockets are deliberately absent +/// from presence and routing, and do not replace an active socket. +struct PendingDeviceConn { + user_id: String, + device_id: String, + auth_token: String, device_name: String, tx: mpsc::Sender, force_close_tx: watch::Sender, @@ -46,22 +66,44 @@ pub struct RpcResponse { /// Tracks online devices grouped by `user_id` so that `device_to_device` /// messages can be routed within an account without exposing other accounts. pub struct DeviceManager { + /// Serializes presence mutations with authoritative snapshot broadcasts. + /// + /// DashMap keeps individual registry operations safe, but a presence + /// message is a compound operation: capture the current membership and + /// enqueue that snapshot to every current member. Without this gate, an + /// older logout/disconnect task can resume after a reconnect and publish + /// a stale snapshot after the reconnect's newer one. + presence_gate: Mutex<()>, + /// Serializes registry ownership changes with their SQLite `online` + /// projection across async route handlers. This is intentionally global: + /// connects/logouts are rare, and one lock avoids an unbounded per-device + /// lock registry while making the durable last-writer deterministic. + presence_projection_gate: tokio::sync::Mutex<()>, /// user_id → (device_id → DeviceConn) users: DashMap>, /// conn_id → (user_id, device_id) for cleanup on disconnect. conn_to_device: DashMap, + /// conn_id → provisional device connection awaiting its final token check. + pending_connections: DashMap, /// correlation_id → pending RPC response sender (for HTTP→WS→HTTP bridge). pending_rpcs: DashMap, pending_rpc_permits: Arc, + /// Starts the database-backed token revalidator exactly once, lazily from + /// the first WebSocket handled inside a Tokio runtime. + token_revalidator_started: AtomicBool, } impl DeviceManager { pub fn new() -> Arc { Arc::new(Self { + presence_gate: Mutex::new(()), + presence_projection_gate: tokio::sync::Mutex::new(()), users: DashMap::new(), conn_to_device: DashMap::new(), + pending_connections: DashMap::new(), pending_rpcs: DashMap::new(), pending_rpc_permits: Arc::new(Semaphore::new(MAX_PENDING_DEVICE_RPCS)), + token_revalidator_started: AtomicBool::new(false), }) } @@ -73,11 +115,13 @@ impl DeviceManager { &self, user_id: &str, device_id: &str, + auth_token: &str, device_name: &str, conn_id: ConnId, tx: mpsc::Sender, force_close_tx: watch::Sender, ) -> Vec<(String, String)> { + let _presence_guard = self.lock_presence(); // Remove any stale conn mapping for this conn first. if let Some((_, (old_user, old_device))) = self.conn_to_device.remove(&conn_id) { if let Some(user_devices) = self.users.get(&old_user) { @@ -114,6 +158,7 @@ impl DeviceManager { device_id.to_string(), DeviceConn { conn_id, + auth_token: auth_token.to_string(), device_name: device_name.to_string(), tx, force_close_tx, @@ -129,9 +174,106 @@ impl DeviceManager { others } + /// Stage a connection while an async post-registration token check runs. + /// It cannot receive routed messages or presence and cannot evict an + /// already-authorized connection for the same physical device. + pub fn register_pending( + &self, + user_id: &str, + device_id: &str, + auth_token: &str, + device_name: &str, + conn_id: ConnId, + tx: mpsc::Sender, + force_close_tx: watch::Sender, + ) { + let _presence_guard = self.lock_presence(); + if let Some((_, previous)) = self.pending_connections.remove(&conn_id) { + let _ = previous.force_close_tx.send(true); + } + self.pending_connections.insert( + conn_id, + PendingDeviceConn { + user_id: user_id.to_string(), + device_id: device_id.to_string(), + auth_token: auth_token.to_string(), + device_name: device_name.to_string(), + tx, + force_close_tx, + }, + ); + } + + /// Promote the exact pending connection to the single active owner for + /// `(user_id, device_id)`. Returns false if logout/delete already removed + /// it while the final database lookup was in flight. + pub fn activate_pending_with_initial_message( + &self, + user_id: &str, + device_id: &str, + auth_token: &str, + conn_id: ConnId, + initial_text: &str, + ) -> bool { + let _presence_guard = self.lock_presence(); + let pending = self + .pending_connections + .remove_if(&conn_id, |_, pending| { + pending.user_id == user_id + && pending.device_id == device_id + && pending.auth_token == auth_token + }) + .map(|(_, pending)| pending); + let Some(pending) = pending else { + return false; + }; + + // Queue AuthOk while the candidate is still invisible and while the + // same presence gate excludes snapshot broadcasts. Once membership is + // published below, every later DevicePresence is necessarily behind + // AuthOk in this socket's FIFO queue. + if pending + .tx + .try_send(OutboundMessage::text(initial_text)) + .is_err() + { + let _ = pending.force_close_tx.send(true); + return false; + } + + let entry = self.users.entry(user_id.to_string()).or_default(); + if let Some((prior_conn_id, prior_close_tx)) = entry + .get(device_id) + .map(|prior| (prior.conn_id, prior.force_close_tx.clone())) + { + if prior_conn_id != conn_id { + self.conn_to_device.remove(&prior_conn_id); + let _ = prior_close_tx.send(true); + } + } + entry.insert( + device_id.to_string(), + DeviceConn { + conn_id, + auth_token: pending.auth_token, + device_name: pending.device_name, + tx: pending.tx, + force_close_tx: pending.force_close_tx, + }, + ); + self.conn_to_device + .insert(conn_id, (user_id.to_string(), device_id.to_string())); + info!( + "Device {device_id} activated for user {user_id} ({} online)", + entry.len() + ); + true + } + /// Remove a device on disconnect. Returns the `(user_id, device_id)` that /// was removed, if any (for presence/DB cleanup by the caller). pub fn unregister(&self, conn_id: ConnId) -> Option<(String, String)> { + let _presence_guard = self.lock_presence(); let removed = self.conn_to_device.remove(&conn_id); if let Some((_, (user_id, device_id))) = &removed { if let Some(user_devices) = self.users.get(user_id) { @@ -152,25 +294,100 @@ impl DeviceManager { return None; } } + if let Some((_, pending)) = self.pending_connections.remove(&conn_id) { + return Some((pending.user_id, pending.device_id)); + } removed.map(|(_, v)| v) } + /// Reject one provisional AuthConnect without affecting an already-active + /// socket that may legitimately use the same still-valid token. + pub fn disconnect_pending(&self, conn_id: ConnId) -> bool { + let _presence_guard = self.lock_presence(); + let Some((_, pending)) = self.pending_connections.remove(&conn_id) else { + return false; + }; + let _ = pending.force_close_tx.send(true); + true + } + /// Immediately revoke a device's in-memory authorization and close its /// WebSocket through a dedicated control channel. The close signal cannot /// be starved by a saturated outbound data queue. pub fn disconnect_device(&self, user_id: &str, device_id: &str) -> bool { - let removed = self + let _presence_guard = self.lock_presence(); + let active = self .users .get(user_id) .and_then(|devices| devices.remove(device_id).map(|(_, device)| device)); - let Some(device) = removed else { - return false; - }; + if let Some(device) = active.as_ref() { + self.conn_to_device.remove(&device.conn_id); + let _ = device.force_close_tx.send(true); + } + let pending_ids: Vec<_> = self + .pending_connections + .iter() + .filter(|entry| entry.user_id == user_id && entry.device_id == device_id) + .map(|entry| *entry.key()) + .collect(); + let mut removed_pending = false; + for pending_id in pending_ids { + if let Some((_, pending)) = self.pending_connections.remove(&pending_id) { + removed_pending = true; + let _ = pending.force_close_tx.send(true); + } + } + let removed = active.is_some() || removed_pending; + if removed { + debug!("Force-disconnected device {device_id}"); + } + removed + } - self.conn_to_device.remove(&device.conn_id); - let _ = device.force_close_tx.send(true); - debug!("Force-disconnected device {device_id}"); - true + /// Disconnect only when the current socket was authenticated by + /// `auth_token`. Revoking a second token for the same machine must leave a + /// still-valid replacement/previous socket untouched. + pub fn disconnect_device_if_token( + &self, + user_id: &str, + device_id: &str, + auth_token: &str, + ) -> bool { + let _presence_guard = self.lock_presence(); + // `remove_if` keeps the token comparison and removal inside one shard + // lock. A reconnect cannot replace socket A with socket B between a + // successful comparison against A and the removal of the entry. + let active = self.users.get(user_id).and_then(|devices| { + devices + .remove_if(device_id, |_, device| device.auth_token == auth_token) + .map(|(_, device)| device) + }); + if let Some(device) = active.as_ref() { + self.conn_to_device.remove(&device.conn_id); + let _ = device.force_close_tx.send(true); + } + let pending_ids: Vec<_> = self + .pending_connections + .iter() + .filter(|entry| { + entry.user_id == user_id + && entry.device_id == device_id + && entry.auth_token == auth_token + }) + .map(|entry| *entry.key()) + .collect(); + let mut removed_pending = false; + for pending_id in pending_ids { + if let Some((_, pending)) = self.pending_connections.remove(&pending_id) { + removed_pending = true; + let _ = pending.force_close_tx.send(true); + } + } + let removed = active.is_some() || removed_pending; + if removed { + debug!("Force-disconnected device {device_id} for its revoked token"); + } + removed } /// Route a raw JSON text message to `target_device_id` within `user_id`. @@ -195,6 +412,11 @@ impl DeviceManager { /// List currently online `(device_id, device_name)` for a user (for /// presence broadcasts). pub fn online_devices(&self, user_id: &str) -> Vec<(String, String)> { + let _presence_guard = self.lock_presence(); + self.online_devices_unlocked(user_id) + } + + fn online_devices_unlocked(&self, user_id: &str) -> Vec<(String, String)> { self.users .get(user_id) .map(|d| { @@ -205,8 +427,50 @@ impl DeviceManager { .unwrap_or_default() } + /// Whether the in-memory routing registry currently has an owner for this + /// device. The registry is the live-connection authority; durable `online` + /// flags are projected from it. + pub fn is_device_online(&self, user_id: &str, device_id: &str) -> bool { + self.users + .get(user_id) + .is_some_and(|devices| devices.contains_key(device_id)) + } + + /// Snapshot all active socket credentials for the periodic revocation + /// revalidator. Pending candidates are excluded because their activation + /// path already performs two serialized database checks. + pub fn active_device_credentials(&self) -> Vec<(ConnId, String, String, String)> { + let _presence_guard = self.lock_presence(); + self.users + .iter() + .flat_map(|user| { + let user_id = user.key().clone(); + user.value() + .iter() + .map(move |device| { + ( + device.conn_id, + user_id.clone(), + device.key().clone(), + device.auth_token.clone(), + ) + }) + .collect::>() + }) + .collect() + } + + /// Return true only for the first caller. The background worker is + /// intentionally launched by the WebSocket entrypoint rather than the + /// synchronous router builder, which may run before a Tokio runtime exists. + pub fn claim_token_revalidator_start(&self) -> bool { + self.token_revalidator_started + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + } + pub fn connection_count(&self) -> usize { - self.conn_to_device.len() + self.conn_to_device.len() + self.pending_connections.len() } pub fn pending_rpc_count(&self) -> usize { @@ -220,21 +484,45 @@ impl DeviceManager { self.conn_to_device.get(&conn_id).map(|e| e.value().clone()) } - /// Broadcast a message to *all* online devices of a user except the sender. - pub fn broadcast_except(&self, user_id: &str, exclude_device_id: &str, text: &str) { + pub fn has_connection(&self, conn_id: ConnId) -> bool { + self.conn_to_device.contains_key(&conn_id) + || self.pending_connections.contains_key(&conn_id) + } + + pub async fn lock_presence_projection(&self) -> tokio::sync::MutexGuard<'_, ()> { + self.presence_projection_gate.lock().await + } + + /// Build and broadcast one authoritative presence snapshot to every + /// currently registered device. Membership cannot change between the + /// snapshot and the queue writes, and delayed lifecycle tasks always + /// rebuild from current state instead of publishing a captured old list. + pub fn broadcast_current_presence(&self, user_id: &str, build_message: F) + where + F: FnOnce(&[(String, String)]) -> Option, + { + let _presence_guard = self.lock_presence(); + let devices = self.online_devices_unlocked(user_id); + let Some(text) = build_message(&devices) else { + return; + }; let Some(user_devices) = self.users.get(user_id) else { return; }; for entry in user_devices.iter() { - if entry.key() != exclude_device_id { - let tx = entry.tx.clone(); - let msg = OutboundMessage::text(text); - // best-effort; don't block the caller on a slow peer - let _ = tx.try_send(msg); - } + let tx = entry.tx.clone(); + let msg = OutboundMessage::text(&text); + // best-effort; don't block the caller on a slow peer + let _ = tx.try_send(msg); } } + fn lock_presence(&self) -> MutexGuard<'_, ()> { + self.presence_gate + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + // ── HTTP RPC bridge ──────────────────────────────────────────────── /// Register a pending RPC response keyed by `correlation_id`. @@ -311,12 +599,12 @@ mod tests { let tx = dummy_tx(); let (close_tx_1, mut close_rx_1) = watch::channel(false); - mgr.register("user-1", "dev-1", "A", 1, tx.clone(), close_tx_1); + mgr.register("user-1", "dev-1", "token-1", "A", 1, tx.clone(), close_tx_1); assert_eq!(mgr.online_devices("user-1").len(), 1); // Reconnect with a new conn id. let (close_tx_2, _close_rx_2) = watch::channel(false); - mgr.register("user-1", "dev-1", "A", 2, tx, close_tx_2); + mgr.register("user-1", "dev-1", "token-2", "A", 2, tx, close_tx_2); assert_eq!(mgr.online_devices("user-1").len(), 1); assert_eq!(mgr.conn_mapping(2), Some(("user-1".into(), "dev-1".into()))); assert_eq!(mgr.conn_mapping(1), None); @@ -337,7 +625,7 @@ mod tests { let mgr = DeviceManager::new(); let (tx, _rx) = mpsc::channel(8); let (close_tx, mut close_rx) = watch::channel(false); - mgr.register("user-1", "dev-1", "A", 1, tx, close_tx); + mgr.register("user-1", "dev-1", "token-1", "A", 1, tx, close_tx); assert!(mgr.disconnect_device("user-1", "dev-1")); assert!(mgr.conn_mapping(1).is_none()); @@ -346,6 +634,113 @@ mod tests { assert!(*close_rx.borrow()); } + #[tokio::test] + async fn token_scoped_disconnect_does_not_close_another_login_socket() { + let mgr = DeviceManager::new(); + let (tx, _rx) = mpsc::channel(8); + let (close_tx, mut close_rx) = watch::channel(false); + mgr.register("user-1", "dev-1", "active-token", "A", 1, tx, close_tx); + + assert!(!mgr.disconnect_device_if_token("user-1", "dev-1", "candidate-token")); + assert_eq!(mgr.conn_mapping(1), Some(("user-1".into(), "dev-1".into()))); + assert!(!*close_rx.borrow_and_update()); + + assert!(mgr.disconnect_device_if_token("user-1", "dev-1", "active-token")); + assert!(mgr.conn_mapping(1).is_none()); + close_rx.changed().await.expect("close signal"); + assert!(*close_rx.borrow()); + } + + #[tokio::test] + async fn revoked_pending_candidate_is_hidden_and_preserves_active_owner() { + let mgr = DeviceManager::new(); + let (active_tx, _active_rx) = mpsc::channel(8); + let (active_close_tx, mut active_close_rx) = watch::channel(false); + mgr.register( + "user-1", + "dev-1", + "active-token", + "Active", + 1, + active_tx, + active_close_tx, + ); + + let (candidate_tx, _candidate_rx) = mpsc::channel(8); + let (candidate_close_tx, mut candidate_close_rx) = watch::channel(false); + mgr.register_pending( + "user-1", + "dev-1", + "revoked-token", + "Candidate", + 2, + candidate_tx, + candidate_close_tx, + ); + + assert_eq!(mgr.connection_count(), 2); + assert_eq!( + mgr.online_devices("user-1"), + vec![("dev-1".to_string(), "Active".to_string())] + ); + assert_eq!(mgr.conn_mapping(1), Some(("user-1".into(), "dev-1".into()))); + assert!(mgr.conn_mapping(2).is_none()); + assert!(!*active_close_rx.borrow_and_update()); + + assert!(mgr.disconnect_device_if_token("user-1", "dev-1", "revoked-token")); + candidate_close_rx.changed().await.expect("candidate close"); + assert!(*candidate_close_rx.borrow()); + assert!(!*active_close_rx.borrow_and_update()); + assert_eq!(mgr.connection_count(), 1); + assert_eq!(mgr.conn_mapping(1), Some(("user-1".into(), "dev-1".into()))); + assert!(!mgr.activate_pending_with_initial_message( + "user-1", + "dev-1", + "revoked-token", + 2, + "auth-ok", + )); + } + + #[tokio::test] + async fn delayed_presence_broadcast_uses_current_membership() { + let mgr = DeviceManager::new(); + let (old_tx, _old_rx) = mpsc::channel(8); + let (old_close_tx, _old_close_rx) = watch::channel(false); + mgr.register( + "user-1", + "old-device", + "old-token", + "Old", + 1, + old_tx, + old_close_tx, + ); + + let (current_tx, mut current_rx) = mpsc::channel(8); + let (current_close_tx, _current_close_rx) = watch::channel(false); + mgr.register( + "user-1", + "current-device", + "current-token", + "Current", + 2, + current_tx, + current_close_tx, + ); + assert!(mgr.disconnect_device("user-1", "old-device")); + + // Model an older lifecycle task resuming after the replacement. It + // must serialize the live registry, not a list captured before await. + mgr.broadcast_current_presence("user-1", |devices| serde_json::to_string(devices).ok()); + let message = current_rx.recv().await.expect("presence snapshot"); + let devices: Vec<(String, String)> = serde_json::from_str(&message.text).unwrap(); + assert_eq!( + devices, + vec![("current-device".to_string(), "Current".to_string())] + ); + } + #[tokio::test] async fn rpc_response_must_come_from_the_expected_account_and_device() { let mgr = DeviceManager::new(); diff --git a/src/crates/services/relay-service/src/routes/api.rs b/src/crates/services/relay-service/src/routes/api.rs index 6388ee4055..0c1d8d2e35 100644 --- a/src/crates/services/relay-service/src/routes/api.rs +++ b/src/crates/services/relay-service/src/routes/api.rs @@ -65,6 +65,13 @@ pub struct AppState { pub db: Option>, /// Optional per-page mutable data root (KV/SQLite/blobs). Required for Page Functions data plane. pub page_data: Option, + /// Short-lived browser handoff tickets and page-scoped grants. These stay + /// process-local so no account credential is persisted or exposed in URLs. + pub page_access_manager: Arc, + /// Manifest-bound Page draft upload sessions and per-Page serialization. + pub page_upload_manager: Arc, + /// Global/account/page admission control for public Page Function execution. + pub page_execution_guard: Arc, /// Per-IP rate limiter for auth endpoints (brute-force protection). pub login_rate_limiter: Arc, /// Per-user online device registry for account-based device routing. @@ -622,6 +629,9 @@ mod tests { asset_store: Arc::new(MemoryAssetStore::new()), db: None, page_data: None, + page_access_manager: Arc::new(crate::routes::pages::PageAccessManager::new()), + page_upload_manager: Arc::new(crate::routes::pages::PageUploadManager::new()), + page_execution_guard: Arc::new(crate::page_execution::PageExecutionGuard::new()), login_rate_limiter: Arc::new(crate::routes::auth::LoginRateLimiter::new()), device_manager: crate::relay::DeviceManager::new(), cors_allow_origins: Arc::new(Vec::new()), diff --git a/src/crates/services/relay-service/src/routes/auth.rs b/src/crates/services/relay-service/src/routes/auth.rs index f30defe931..dff61d3ef0 100644 --- a/src/crates/services/relay-service/src/routes/auth.rs +++ b/src/crates/services/relay-service/src/routes/auth.rs @@ -398,6 +398,21 @@ pub async fn logout(State(state): State, headers: HeaderMap) -> Status }; match AuthToken::find(db, &token).await { Ok(Some(auth)) => { + let _presence_projection_guard = state.device_manager.lock_presence_projection().await; + // The first lookup determines which lifecycle to serialize. Check + // the exact token again under that lifecycle boundary so a + // concurrent device deletion cannot leave this handler operating + // on stale authorization state. + let current = match AuthToken::find(db, &token).await { + Ok(Some(current)) + if current.user_id == auth.user_id + && current.device_id == auth.device_id + && current.token_kind == auth.token_kind => + { + current + } + _ => return StatusCode::UNAUTHORIZED, + }; // Delete the token row if let Err(error) = sqlx::query("DELETE FROM auth_tokens WHERE token = ?") .bind(&token) @@ -410,30 +425,44 @@ pub async fn logout(State(state): State, headers: HeaderMap) -> Status // A delegated control token borrows the desktop's device id but // does not own its live connection. Logging out that client must // never disconnect or mark the desktop offline. - if auth.is_device_token() { - let _ = crate::db::DeviceRow::set_online(db, &auth.user_id, &auth.device_id, false) - .await; - state - .device_manager - .disconnect_device(&auth.user_id, &auth.device_id); - let devices = state + if current.is_device_token() { + state.device_manager.disconnect_device_if_token( + &auth.user_id, + &auth.device_id, + &token, + ); + // A failed token match means another login currently owns the + // same machine's socket. Keep its durable presence online; + // otherwise a rejected candidate login could make the still- + // connected prior account appear offline. + if !state .device_manager - .online_devices(&auth.user_id) - .into_iter() - .map( - |(device_id, device_name)| crate::routes::websocket::DevicePresenceEntry { - device_id, - device_name, - }, - ) - .collect(); - let presence = - crate::routes::websocket::OutboundProtocol::DevicePresence { devices }; - if let Ok(message) = serde_json::to_string(&presence) { - state - .device_manager - .broadcast_except(&auth.user_id, &auth.device_id, &message); + .is_device_online(&auth.user_id, &auth.device_id) + { + let _ = + crate::db::DeviceRow::set_online(db, &auth.user_id, &auth.device_id, false) + .await; } + drop(_presence_projection_guard); + state + .device_manager + .broadcast_current_presence(&auth.user_id, |devices| { + let devices = devices + .iter() + .map(|(device_id, device_name)| { + crate::routes::websocket::DevicePresenceEntry { + device_id: device_id.clone(), + device_name: device_name.clone(), + } + }) + .collect(); + serde_json::to_string( + &crate::routes::websocket::OutboundProtocol::DevicePresence { devices }, + ) + .ok() + }); + } else { + drop(_presence_projection_guard); } tracing::info!("Account token revoked for device_id={}", auth.device_id); StatusCode::NO_CONTENT diff --git a/src/crates/services/relay-service/src/routes/devices.rs b/src/crates/services/relay-service/src/routes/devices.rs index b1e1a3b876..c92e427a4d 100644 --- a/src/crates/services/relay-service/src/routes/devices.rs +++ b/src/crates/services/relay-service/src/routes/devices.rs @@ -87,7 +87,7 @@ async fn list_devices( headers: HeaderMap, ) -> Result>, StatusCode> { let auth = validate_user(&state, &headers).await?; - let user_id = auth.user_id; + let user_id = auth.user_id.clone(); // Get online devices from DeviceManager (in-memory) let online = state.device_manager.online_devices(&user_id); @@ -225,13 +225,24 @@ async fn delete_device( if !auth.is_device_token() { return Err(StatusCode::FORBIDDEN); } - let user_id = auth.user_id; + let user_id = auth.user_id.clone(); if !is_valid_device_id(&target_device_id) { return Err(StatusCode::BAD_REQUEST); } let db = state.db.as_ref().ok_or(StatusCode::NOT_IMPLEMENTED)?; + let _presence_projection_guard = state.device_manager.lock_presence_projection().await; + let current_auth = crate::db::AuthToken::find(db, &auth.token) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::UNAUTHORIZED)?; + if !current_auth.is_device_token() + || current_auth.user_id != user_id + || current_auth.device_id != auth.device_id + { + return Err(StatusCode::UNAUTHORIZED); + } // Revoke the target's auth tokens before removing its device row. The DB // helper also scopes the deletion to this account and performs both writes @@ -255,23 +266,22 @@ async fn delete_device( state .device_manager .disconnect_device(&user_id, &target_device_id); + drop(_presence_projection_guard); - let online = state.device_manager.online_devices(&user_id); - let presence = online - .iter() - .map( - |(device_id, device_name)| crate::routes::websocket::DevicePresenceEntry { - device_id: device_id.clone(), - device_name: device_name.clone(), - }, - ) - .collect(); - let presence_json = - serde_json::to_string(&OutboundProtocol::DevicePresence { devices: presence }) - .unwrap_or_default(); state .device_manager - .broadcast_except(&user_id, &target_device_id, &presence_json); + .broadcast_current_presence(&user_id, |devices| { + let devices = devices + .iter() + .map( + |(device_id, device_name)| crate::routes::websocket::DevicePresenceEntry { + device_id: device_id.clone(), + device_name: device_name.clone(), + }, + ) + .collect(); + serde_json::to_string(&OutboundProtocol::DevicePresence { devices }).ok() + }); tracing::info!("Device {target_device_id} removed from account {user_id}"); Ok(StatusCode::NO_CONTENT) diff --git a/src/crates/services/relay-service/src/routes/pages.rs b/src/crates/services/relay-service/src/routes/pages.rs index 5209905feb..7590bb0ba2 100644 --- a/src/crates/services/relay-service/src/routes/pages.rs +++ b/src/crates/services/relay-service/src/routes/pages.rs @@ -4,23 +4,27 @@ //! Preview: `/p/{user}/{slug}/@v/{version}/...` //! Production: `/p/{user}/{slug}/...` (deployed version only). -use axum::extract::{DefaultBodyLimit, Path, State}; -use axum::http::{header, HeaderMap, Method, StatusCode}; -use axum::response::IntoResponse; +use axum::extract::{DefaultBodyLimit, Path, Query, State}; +use axum::http::{header, HeaderMap, HeaderValue, Method, StatusCode}; +use axum::response::{IntoResponse, Redirect, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; use base64::{engine::general_purpose::STANDARD as B64, Engine}; use bitfun_page_function_runtime::{ - run_fetch, FetchRequest, PageMeta, DEFAULT_TIMEOUT, WORKER_ENTRY_PATH, + run_fetch, FetchRequest, PageFunctionError, PageMeta, DEFAULT_TIMEOUT, WORKER_ENTRY_PATH, }; +use dashmap::DashMap; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use std::collections::HashMap; -use std::sync::Arc; +use std::hash::{DefaultHasher, Hash, Hasher}; +use std::sync::{Arc, Mutex as StdMutex}; +use std::time::{Duration, Instant}; use crate::db::{ new_page_version_id, page_draft_asset_key, page_legacy_asset_key, page_version_asset_key, - PageRow, PageVersionRow, PageVisibility, PageWithUsername, UserRow, + DeletePageVersionOutcome, PageMutationOutcome, PageRow, PageVersionRow, PageVisibility, + PageWithUsername, SavePageVersionOutcome, UserRow, }; use crate::page_data::RelayPageHost; use crate::routes::api::AppState; @@ -31,7 +35,303 @@ pub const MAX_PAGES_PER_USER: i64 = 50; pub const MAX_VERSIONS_PER_PAGE: i64 = 30; pub const MAX_PAGE_BYTES: u64 = 100 * 1024 * 1024; pub const MAX_FILE_BYTES: u64 = 10 * 1024 * 1024; -pub const PAGE_UPLOAD_BODY_LIMIT: usize = 12 * 1024 * 1024; +pub const MAX_PAGE_FILES: usize = 4096; +const MAX_PAGE_TITLE_CHARS: usize = 120; +const MAX_PAGE_TITLE_BYTES: usize = 512; +const MAX_PAGE_NOTE_CHARS: usize = 1000; +const MAX_PAGE_NOTE_BYTES: usize = 4096; +// A 10 MiB file expands to roughly 13.34 MiB in base64 before JSON framing. +pub const PAGE_UPLOAD_BODY_LIMIT: usize = 15 * 1024 * 1024; +const PAGE_OPEN_TICKET_TTL: Duration = Duration::from_secs(60); +const PAGE_BROWSER_GRANT_TTL: Duration = Duration::from_secs(10 * 60); +const MAX_PAGE_OPEN_TICKETS: usize = 4096; +const MAX_PAGE_OPEN_TICKETS_PER_USER: usize = 256; +const MAX_PAGE_BROWSER_GRANTS: usize = 8192; +const MAX_PAGE_BROWSER_GRANTS_PER_USER: usize = 512; +const PAGE_ACCESS_COOKIE: &str = "bitfun_page_access"; +const PAGE_UPLOAD_SESSION_TTL: Duration = Duration::from_secs(15 * 60); +const MAX_PAGE_UPLOAD_SESSIONS: usize = 4096; +const MAX_PAGE_UPLOAD_SESSIONS_PER_USER: usize = 64; +const PAGE_UPLOAD_LOCK_SHARDS: usize = 64; + +#[derive(Clone)] +struct PageAccessScope { + user_id: String, + username: String, + slug: String, + generation: String, + version_id: Option, + expires_at: Instant, +} + +/// Process-local, time-bounded grants used to hand an authenticated Page URL +/// from the desktop client to the user's external browser without putting the +/// account bearer token in a URL. +pub struct PageAccessManager { + open_tickets: DashMap, + browser_grants: DashMap, + capacity_lock: StdMutex<()>, +} + +#[derive(Clone)] +struct PageUploadSession { + user_id: String, + upload_id: Option, + draft_key: String, + manifest: HashMap, + finalized: bool, + title: Option, + visibility: Option, + expected_generation: Option, + create: bool, + expires_at: Instant, +} + +/// Tracks the one active, manifest-bound upload session for each Page. Drafts +/// use upload-specific asset namespaces, so superseded/concurrent uploads can +/// never mix files before an immutable version is frozen. +pub struct PageUploadManager { + sessions: DashMap, + locks: Vec>, + capacity_lock: StdMutex<()>, +} + +impl Default for PageUploadManager { + fn default() -> Self { + Self::new() + } +} + +impl PageUploadManager { + pub fn new() -> Self { + Self { + sessions: DashMap::new(), + locks: (0..PAGE_UPLOAD_LOCK_SHARDS) + .map(|_| tokio::sync::Mutex::new(())) + .collect(), + capacity_lock: StdMutex::new(()), + } + } + + fn lock_index(&self, page_key: &str) -> usize { + let mut hasher = DefaultHasher::new(); + page_key.hash(&mut hasher); + hasher.finish() as usize % self.locks.len() + } + + fn lock_for(&self, page_key: &str) -> &tokio::sync::Mutex<()> { + &self.locks[self.lock_index(page_key)] + } + + async fn prune_expired(&self, now: Instant) -> Vec { + let mut expired_drafts = Vec::new(); + // Acquire the exact shard used by every handler before removing its + // sessions. An upload that passed its lease check and is still doing + // storage/DB work therefore cannot be pruned from another Page. + for shard_index in 0..self.locks.len() { + let _guard = self.locks[shard_index].lock().await; + let expired_keys = self + .sessions + .iter() + .filter(|entry| { + self.lock_index(entry.key()) == shard_index && entry.expires_at <= now + }) + .map(|entry| entry.key().clone()) + .collect::>(); + for key in expired_keys { + if let Some((_, session)) = self.sessions.remove(&key) { + expired_drafts.push(session.draft_key); + } + } + } + expired_drafts + } +} + +impl Default for PageAccessManager { + fn default() -> Self { + Self { + open_tickets: DashMap::new(), + browser_grants: DashMap::new(), + capacity_lock: StdMutex::new(()), + } + } +} + +impl PageAccessManager { + pub fn new() -> Self { + Self::default() + } + + fn prune_expired(&self, now: Instant) { + self.open_tickets.retain(|_, scope| scope.expires_at > now); + self.browser_grants + .retain(|_, scope| scope.expires_at > now); + } + + fn issue_open_ticket( + &self, + user_id: String, + username: String, + slug: String, + generation: String, + version_id: Option, + ) -> Result { + let _capacity_guard = self + .capacity_lock + .lock() + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let now = Instant::now(); + self.prune_expired(now); + let user_tickets = self + .open_tickets + .iter() + .filter(|entry| entry.user_id == user_id) + .count(); + if self.open_tickets.len() >= MAX_PAGE_OPEN_TICKETS + || user_tickets >= MAX_PAGE_OPEN_TICKETS_PER_USER + { + return Err(StatusCode::TOO_MANY_REQUESTS); + } + let ticket = random_page_access_token(); + self.open_tickets.insert( + ticket.clone(), + PageAccessScope { + user_id, + username, + slug, + generation, + version_id, + expires_at: now + PAGE_OPEN_TICKET_TTL, + }, + ); + Ok(ticket) + } + + fn exchange_ticket( + &self, + ticket: &str, + ) -> Result, StatusCode> { + let _capacity_guard = self + .capacity_lock + .lock() + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let now = Instant::now(); + self.prune_expired(now); + let Some(scope) = self + .open_tickets + .get(ticket) + .map(|entry| entry.value().clone()) + else { + return Ok(None); + }; + if scope.expires_at <= now { + self.open_tickets.remove(ticket); + return Ok(None); + } + let user_grants = self + .browser_grants + .iter() + .filter(|entry| entry.user_id == scope.user_id) + .count(); + if self.browser_grants.len() >= MAX_PAGE_BROWSER_GRANTS + || user_grants >= MAX_PAGE_BROWSER_GRANTS_PER_USER + { + return Err(StatusCode::TOO_MANY_REQUESTS); + } + let Some((_, mut scope)) = self.open_tickets.remove(ticket) else { + return Ok(None); + }; + let grant = random_page_access_token(); + scope.expires_at = now + PAGE_BROWSER_GRANT_TTL; + self.browser_grants.insert(grant.clone(), scope.clone()); + Ok(Some((grant, scope))) + } + + fn authorizes_page( + &self, + headers: &HeaderMap, + user_id: &str, + slug: &str, + generation: &str, + version_id: Option<&str>, + ) -> bool { + let now = Instant::now(); + let mut authorized = false; + for cookie_header in headers.get_all(header::COOKIE) { + let Ok(value) = cookie_header.to_str() else { + continue; + }; + for cookie in value.split(';') { + let Some((name, token)) = cookie.trim().split_once('=') else { + continue; + }; + if name != PAGE_ACCESS_COOKIE { + continue; + } + let Some(scope) = self.browser_grants.get(token) else { + continue; + }; + if scope.expires_at > now + && scope.user_id == user_id + && scope.slug == slug + && scope.generation == generation + && scope.version_id.as_deref() == version_id + { + authorized = true; + break; + } + } + if authorized { + break; + } + } + if !authorized { + self.prune_expired(now); + } + authorized + } + + fn revoke_page(&self, user_id: &str, slug: &str) { + self.open_tickets + .retain(|_, scope| scope.user_id != user_id || scope.slug != slug); + self.browser_grants + .retain(|_, scope| scope.user_id != user_id || scope.slug != slug); + } +} + +fn random_page_access_token() -> String { + format!( + "{}{}", + uuid::Uuid::new_v4().simple(), + uuid::Uuid::new_v4().simple() + ) +} + +fn is_valid_page_upload_id(upload_id: &str) -> bool { + upload_id.len() == 32 && upload_id.bytes().all(|byte| byte.is_ascii_hexdigit()) +} + +fn is_valid_page_title(title: &str) -> bool { + !title.trim().is_empty() + && title.len() <= MAX_PAGE_TITLE_BYTES + && title.chars().count() <= MAX_PAGE_TITLE_CHARS + && !title.chars().any(char::is_control) +} + +fn is_valid_page_note(note: &str) -> bool { + note.len() <= MAX_PAGE_NOTE_BYTES + && note.chars().count() <= MAX_PAGE_NOTE_CHARS + && !note.chars().any(char::is_control) +} + +fn page_upload_session_key(user_id: &str, slug: &str) -> String { + format!("{user_id}\0{slug}") +} + +fn page_upload_draft_key(user_id: &str, slug: &str, upload_id: &str) -> String { + format!("pages/{user_id}/{slug}/draft/{upload_id}") +} fn is_valid_slug(slug: &str) -> bool { let bytes = slug.as_bytes(); @@ -49,6 +349,89 @@ fn is_valid_slug(slug: &str) -> bool { .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || *b == b'-') } +fn is_valid_page_generation(generation: &str) -> bool { + generation.len() == 32 && generation.bytes().all(|byte| byte.is_ascii_hexdigit()) +} + +fn validate_optional_expected_generation(value: Option<&str>) -> Result, StatusCode> { + match value { + None | Some("") => Ok(None), + Some(generation) if is_valid_page_generation(generation) => Ok(value), + Some(_) => Err(StatusCode::BAD_REQUEST), + } +} + +/// Bind a management request to the Page generation visible while the caller +/// holds the Page lifecycle lock. Pre-generation clients omit the fence and +/// intentionally operate on that current generation; generation-aware clients +/// retain the explicit mismatch fence across delete/recreate ABA transitions. +async fn bind_page_generation_locked( + db: &crate::db::DbPool, + user_id: &str, + slug: &str, + expected_generation: Option<&str>, +) -> Result { + let page = PageRow::get(db, user_id, slug) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::NOT_FOUND)?; + if expected_generation.is_some_and(|expected| expected != page.generation) { + return Err(StatusCode::CONFLICT); + } + Ok(page) +} + +fn generation_intent_matches( + page: Option<&PageRow>, + expected_generation: Option<&str>, + create: bool, +) -> bool { + match page { + Some(page) => { + !create + && expected_generation.is_some_and(|expected| { + is_valid_page_generation(expected) && expected == page.generation + }) + } + None => create && expected_generation.is_none(), + } +} + +fn is_legacy_upload_generation_intent(expected_generation: Option<&str>, create: bool) -> bool { + expected_generation.is_none() && !create +} + +/// Resolve the upload intent at the manifest-check boundary. Clients from +/// before the generation protocol omit both fields; bind those requests to +/// the exact Page generation observed by the relay (or to an explicit create +/// intent when the slug is absent). The server-owned session carries that +/// binding through upload and freeze, so this compatibility path cannot turn +/// into an unfenced mutation after a delete/recreate ABA transition. +fn resolve_upload_generation_intent( + page: Option<&PageRow>, + expected_generation: Option<&str>, + create: bool, +) -> Option<(Option, bool)> { + if is_legacy_upload_generation_intent(expected_generation, create) { + let resolved_generation = page.map(|page| page.generation.clone()); + let resolved_create = page.is_none(); + return generation_intent_matches(page, resolved_generation.as_deref(), resolved_create) + .then_some((resolved_generation, resolved_create)); + } + generation_intent_matches(page, expected_generation, create) + .then(|| (expected_generation.map(str::to_string), create)) +} + +fn upload_request_matches_session_intent( + expected_generation: Option<&str>, + create: bool, + session: &PageUploadSession, +) -> bool { + is_legacy_upload_generation_intent(expected_generation, create) + || (session.expected_generation.as_deref() == expected_generation + && session.create == create) +} + fn require_db(state: &AppState) -> Result<&crate::db::DbPool, StatusCode> { state .db @@ -73,18 +456,26 @@ pub fn pages_router() -> Router { get(list_versions).post(freeze_version), ) .route("/api/pages/{slug}/deploy", post(deploy_version)) + .route("/api/pages/{slug}/unpublish", post(unpublish_page)) .route( "/api/pages/{slug}/versions/{version_id}", axum::routing::delete(delete_version), ) .route( "/api/pages/{slug}", - axum::routing::patch(update_page).delete(delete_page), + post(create_open_ticket) + .patch(update_page) + .delete(delete_page), ) + .route("/api/page-open/{ticket}", get(exchange_open_ticket)) // Preview routes (more specific first). .route( "/p/{username}/{slug}/@v/{version_id}", - get(serve_preview_root).post(serve_preview_root), + get(serve_preview_root) + .post(serve_preview_root) + .layer(DefaultBodyLimit::max( + crate::page_execution::MAX_PAGE_FUNCTION_REQUEST_BODY_BYTES, + )), ) .route( "/p/{username}/{slug}/@v/{version_id}/{*path}", @@ -92,11 +483,18 @@ pub fn pages_router() -> Router { .post(serve_preview_path) .put(serve_preview_path) .delete(serve_preview_path) - .patch(serve_preview_path), + .patch(serve_preview_path) + .layer(DefaultBodyLimit::max( + crate::page_execution::MAX_PAGE_FUNCTION_REQUEST_BODY_BYTES, + )), ) .route( "/p/{username}/{slug}", - get(serve_prod_root).post(serve_prod_root), + get(serve_prod_root) + .post(serve_prod_root) + .layer(DefaultBodyLimit::max( + crate::page_execution::MAX_PAGE_FUNCTION_REQUEST_BODY_BYTES, + )), ) .route( "/p/{username}/{slug}/{*path}", @@ -104,7 +502,10 @@ pub fn pages_router() -> Router { .post(serve_prod_path) .put(serve_prod_path) .delete(serve_prod_path) - .patch(serve_prod_path), + .patch(serve_prod_path) + .layer(DefaultBodyLimit::max( + crate::page_execution::MAX_PAGE_FUNCTION_REQUEST_BODY_BYTES, + )), ) } @@ -120,11 +521,21 @@ pub struct FileManifestEntry { #[derive(Deserialize)] pub struct CheckPageFilesRequest { pub slug: String, + #[serde(default)] + pub upload_id: Option, + #[serde(default)] + pub expected_generation: Option, + #[serde(default)] + pub create: bool, pub files: Vec, } #[derive(Serialize)] pub struct CheckPageFilesResponse { + #[serde(skip_serializing_if = "Option::is_none")] + pub upload_id: Option, + pub expected_generation: Option, + pub create: bool, pub needed: Vec, pub existing_count: usize, pub total_count: usize, @@ -140,6 +551,12 @@ pub struct UploadFileEntry { pub struct UploadPageFilesRequest { pub slug: String, #[serde(default)] + pub upload_id: Option, + #[serde(default)] + pub expected_generation: Option, + #[serde(default)] + pub create: bool, + #[serde(default)] pub title: String, pub visibility: String, pub files: HashMap, @@ -149,6 +566,12 @@ pub struct UploadPageFilesRequest { #[derive(Deserialize)] pub struct FreezeVersionRequest { + #[serde(default)] + pub upload_id: Option, + #[serde(default)] + pub expected_generation: Option, + #[serde(default)] + pub create: bool, #[serde(default)] pub title: String, #[serde(default)] @@ -158,11 +581,14 @@ pub struct FreezeVersionRequest { #[derive(Deserialize)] pub struct DeployRequest { pub version_id: String, + #[serde(default)] + pub expected_generation: Option, } #[derive(Serialize)] pub struct PageInfo { pub slug: String, + pub generation: String, pub visibility: String, pub title: String, pub file_count: i64, @@ -176,6 +602,7 @@ pub struct PageInfo { #[derive(Serialize)] pub struct VersionInfo { + pub generation: String, pub version_id: String, pub title: String, pub file_count: i64, @@ -189,12 +616,148 @@ pub struct VersionInfo { #[derive(Deserialize)] pub struct UpdatePageRequest { + #[serde(default)] + pub expected_generation: Option, pub visibility: Option, pub title: Option, } +#[derive(Deserialize)] +pub struct PageOpenRequest { + #[serde(default)] + pub expected_generation: Option, + #[serde(default)] + pub version_id: Option, +} + +#[derive(Deserialize)] +pub struct ExpectedGenerationQuery { + #[serde(default)] + pub expected_generation: Option, +} + +#[derive(Serialize)] +pub struct PageOpenResponse { + pub open_url_path: String, + pub expires_in_seconds: u64, +} + // ── Management ────────────────────────────────────────────────────────── +async fn create_open_ticket( + State(state): State, + headers: HeaderMap, + Path(slug): Path, + Json(body): Json, +) -> Result, StatusCode> { + let auth = validate_auth(&state, &headers).await?; + let db = require_db(&state)?; + if !is_valid_slug(&slug) { + return Err(StatusCode::BAD_REQUEST); + } + let expected_generation = + validate_optional_expected_generation(body.expected_generation.as_deref())?; + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_read(&auth.user_id, &slug) + .await; + let page = bind_page_generation_locked(db, &auth.user_id, &slug, expected_generation).await?; + let username = UserRow::find_by_username_for_user_id(db, &auth.user_id) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::NOT_FOUND)?; + if !is_safe_page_username(&username) { + return Err(StatusCode::BAD_REQUEST); + } + + let version_id = body + .version_id + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + if let Some(version_id) = version_id.as_deref() { + PageVersionRow::get(db, &auth.user_id, &slug, version_id) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::NOT_FOUND)?; + } else if page.deployed_version_id.is_none() { + return Err(StatusCode::NOT_FOUND); + } + + let ticket = state.page_access_manager.issue_open_ticket( + auth.user_id, + username, + slug, + page.generation, + version_id, + )?; + Ok(Json(PageOpenResponse { + open_url_path: format!("/api/page-open/{ticket}"), + expires_in_seconds: PAGE_OPEN_TICKET_TTL.as_secs(), + })) +} + +async fn exchange_open_ticket( + State(state): State, + headers: HeaderMap, + Path(ticket): Path, +) -> Result { + if ticket.len() != 64 || !ticket.bytes().all(|byte| byte.is_ascii_hexdigit()) { + return Err(StatusCode::NOT_FOUND); + } + let (grant, scope) = state + .page_access_manager + .exchange_ticket(&ticket)? + .ok_or(StatusCode::NOT_FOUND)?; + let page_path = format!("/p/{}/{}", scope.username, scope.slug); + let target = match scope.version_id { + Some(version_id) => format!("{page_path}/@v/{version_id}"), + None => page_path.clone(), + }; + let secure = request_used_https(&headers); + let cookie = format!( + "{PAGE_ACCESS_COOKIE}={grant}; Path={target}; Max-Age={}; HttpOnly; SameSite=Lax{}", + PAGE_BROWSER_GRANT_TTL.as_secs(), + if secure { "; Secure" } else { "" } + ); + let mut response = Redirect::temporary(&target).into_response(); + response.headers_mut().insert( + header::SET_COOKIE, + HeaderValue::from_str(&cookie).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?, + ); + response + .headers_mut() + .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store")); + Ok(response) +} + +fn request_used_https(headers: &HeaderMap) -> bool { + headers + .get("forwarded") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| { + value + .split(';') + .any(|part| part.trim().eq_ignore_ascii_case("proto=https")) + }) + || headers + .get("x-forwarded-proto") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| { + value + .split(',') + .next() + .is_some_and(|proto| proto.trim().eq_ignore_ascii_case("https")) + }) +} + +fn is_safe_page_username(username: &str) -> bool { + !username.is_empty() + && username.len() <= 128 + && username + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.')) +} + async fn list_pages( State(state): State, headers: HeaderMap, @@ -214,11 +777,18 @@ async fn list_pages( for p in pages { // One-time legacy migration: old room without versions. maybe_migrate_legacy_page(&state, &p).await; + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_read(&p.user_id, &p.slug) + .await; let page = PageRow::get(db, &p.user_id, &p.slug) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .unwrap_or(p); - infos.push(page_to_info(&page, &username)); + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + // A row deleted after the initial list query must not be resurrected + // in the response from the stale in-memory `p` snapshot. + if let Some(page) = page { + infos.push(page_to_info(&page, &username)); + } } Ok(Json(infos)) } @@ -230,11 +800,24 @@ async fn check_page_files( ) -> Result, StatusCode> { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; - if !is_valid_slug(&body.slug) { + if !is_valid_slug(&body.slug) + || body + .upload_id + .as_deref() + .is_some_and(|upload_id| !is_valid_page_upload_id(upload_id)) + || body + .expected_generation + .as_deref() + .is_some_and(|generation| !is_valid_page_generation(generation)) + { return Err(StatusCode::BAD_REQUEST); } let mut total_bytes: u64 = 0; + let mut manifest = HashMap::with_capacity(body.files.len()); + if body.files.is_empty() || body.files.len() > MAX_PAGE_FILES { + return Err(StatusCode::BAD_REQUEST); + } for entry in &body.files { if crate::validated_asset_relative_path(&entry.path).is_err() || !crate::is_valid_content_hash(&entry.hash) @@ -245,6 +828,12 @@ async fn check_page_files( return Err(StatusCode::PAYLOAD_TOO_LARGE); } total_bytes = total_bytes.saturating_add(entry.size); + if manifest + .insert(entry.path.clone(), entry.hash.clone()) + .is_some() + { + return Err(StatusCode::BAD_REQUEST); + } } if total_bytes > MAX_PAGE_BYTES { return Err(StatusCode::PAYLOAD_TOO_LARGE); @@ -253,6 +842,12 @@ async fn check_page_files( let existing = PageRow::get(db, &auth.user_id, &body.slug) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let (expected_generation, create) = resolve_upload_generation_intent( + existing.as_ref(), + body.expected_generation.as_deref(), + body.create, + ) + .ok_or(StatusCode::CONFLICT)?; if existing.is_none() { let count = PageRow::count_for_user(db, &auth.user_id) .await @@ -262,13 +857,92 @@ async fn check_page_files( } } - let draft_key = page_draft_asset_key(&auth.user_id, &body.slug); + for stale_draft in state + .page_upload_manager + .prune_expired(Instant::now()) + .await + { + state.asset_store.cleanup_room(&stale_draft); + } + + let page_key = page_upload_session_key(&auth.user_id, &body.slug); + let _page_guard = state.page_upload_manager.lock_for(&page_key).lock().await; + let current_page = PageRow::get(db, &auth.user_id, &body.slug) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + if !generation_intent_matches( + current_page.as_ref(), + expected_generation.as_deref(), + create, + ) { + return Err(StatusCode::CONFLICT); + } + let draft_key = body.upload_id.as_deref().map_or_else( + || page_draft_asset_key(&auth.user_id, &body.slug), + |upload_id| page_upload_draft_key(&auth.user_id, &body.slug, upload_id), + ); + if let Some((_, previous)) = state.page_upload_manager.sessions.remove(&page_key) { + state.asset_store.cleanup_room(&previous.draft_key); + } + // Reclaim upload-specific drafts left by a prior process crash. The + // in-memory store uses exact namespaces (the active one was removed + // above); the disk store recursively clears this Page's draft subtree. + state + .asset_store + .cleanup_room(&page_draft_asset_key(&auth.user_id, &body.slug)); + // A retry using the same upload id must also start from the submitted + // manifest, never from partial mappings left by the failed attempt. + state.asset_store.cleanup_room(&draft_key); let asset_store = Arc::clone(&state.asset_store); + let processing_draft_key = draft_key.clone(); + let response_upload_id = body.upload_id.clone(); + let response_generation = expected_generation.clone(); + let response_create = create; let response = tokio::task::spawn_blocking(move || { - process_check_page_files(asset_store, &draft_key, body.files) + process_check_page_files( + asset_store, + &processing_draft_key, + response_upload_id, + response_generation, + response_create, + body.files, + ) }) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)??; + + let _capacity_guard = state + .page_upload_manager + .capacity_lock + .lock() + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let user_sessions = state + .page_upload_manager + .sessions + .iter() + .filter(|entry| entry.user_id == auth.user_id) + .count(); + if state.page_upload_manager.sessions.len() >= MAX_PAGE_UPLOAD_SESSIONS + || user_sessions >= MAX_PAGE_UPLOAD_SESSIONS_PER_USER + { + state.asset_store.cleanup_room(&draft_key); + return Err(StatusCode::TOO_MANY_REQUESTS); + } + state.page_upload_manager.sessions.insert( + page_key, + PageUploadSession { + user_id: auth.user_id, + upload_id: body.upload_id.clone(), + draft_key, + manifest, + finalized: false, + title: None, + visibility: None, + expected_generation, + create, + expires_at: Instant::now() + PAGE_UPLOAD_SESSION_TTL, + }, + ); Ok(Json(response)) } @@ -279,7 +953,19 @@ async fn upload_page_files( ) -> Result, StatusCode> { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; - if !is_valid_slug(&body.slug) { + if !is_valid_slug(&body.slug) + || body + .upload_id + .as_deref() + .is_some_and(|upload_id| !is_valid_page_upload_id(upload_id)) + || body + .expected_generation + .as_deref() + .is_some_and(|generation| !is_valid_page_generation(generation)) + { + return Err(StatusCode::BAD_REQUEST); + } + if !body.title.is_empty() && !is_valid_page_title(&body.title) { return Err(StatusCode::BAD_REQUEST); } let visibility = PageVisibility::parse(&body.visibility).ok_or(StatusCode::BAD_REQUEST)?; @@ -290,16 +976,54 @@ async fn upload_page_files( { return Err(StatusCode::BAD_REQUEST); } - let approx = (entry.content.len() as u64).saturating_mul(3) / 4; - if approx > MAX_FILE_BYTES { + let max_encoded_file_bytes = MAX_FILE_BYTES.div_ceil(3).saturating_mul(4); + if entry.content.len() as u64 > max_encoded_file_bytes { return Err(StatusCode::PAYLOAD_TOO_LARGE); } } - let existing = PageRow::get(db, &auth.user_id, &body.slug) + let slug = body.slug.clone(); + let upload_id = body.upload_id.clone(); + let finalize = body.finalize; + let page_key = page_upload_session_key(&auth.user_id, &slug); + let _page_guard = state.page_upload_manager.lock_for(&page_key).lock().await; + let session = state + .page_upload_manager + .sessions + .get(&page_key) + .map(|entry| entry.value().clone()) + .ok_or(StatusCode::CONFLICT)?; + if session.upload_id != upload_id + || !upload_request_matches_session_intent( + body.expected_generation.as_deref(), + body.create, + &session, + ) + { + return Err(StatusCode::CONFLICT); + } + if session.expires_at <= Instant::now() { + state.page_upload_manager.sessions.remove(&page_key); + state.asset_store.cleanup_room(&session.draft_key); + return Err(StatusCode::GONE); + } + for (path, entry) in &body.files { + if session.manifest.get(path) != Some(&entry.hash) { + return Err(StatusCode::CONFLICT); + } + } + + let existing = PageRow::get(db, &auth.user_id, &slug) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - if existing.is_none() { + if !generation_intent_matches( + existing.as_ref(), + session.expected_generation.as_deref(), + session.create, + ) { + return Err(StatusCode::CONFLICT); + } + if finalize && existing.is_none() { let count = PageRow::count_for_user(db, &auth.user_id) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; @@ -312,15 +1036,11 @@ async fn upload_page_files( existing .as_ref() .map(|p| p.title.clone()) - .unwrap_or_else(|| body.slug.clone()) + .unwrap_or_else(|| slug.clone()) } else { body.title.clone() }; - PageRow::ensure(db, &auth.user_id, &body.slug, visibility, &title) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - let draft_key = page_draft_asset_key(&auth.user_id, &body.slug); + let draft_key = session.draft_key.clone(); let asset_store = Arc::clone(&state.asset_store); let files = body.files; let stored = tokio::task::spawn_blocking(move || { @@ -329,12 +1049,43 @@ async fn upload_page_files( .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)??; + let mut active = state + .page_upload_manager + .sessions + .get_mut(&page_key) + .ok_or(StatusCode::CONFLICT)?; + if active.upload_id != upload_id + || active.expected_generation != session.expected_generation + || active.create != session.create + { + return Err(StatusCode::CONFLICT); + } + // Every accepted batch renews the upload lease. Long but progressing + // uploads must not expire relative to the initial manifest check. + active.expires_at = Instant::now() + PAGE_UPLOAD_SESSION_TTL; + + if finalize { + let actual_manifest = state + .asset_store + .list_room_entries(&session.draft_key) + .into_iter() + .collect::>(); + if actual_manifest != session.manifest { + return Err(StatusCode::CONFLICT); + } + active.finalized = true; + active.title = Some(title); + active.visibility = Some(visibility); + } + drop(active); + Ok(Json(serde_json::json!({ "status": "ok", "files_stored": stored, - "slug": body.slug, + "slug": slug, + "upload_id": upload_id, "draft": true, - "finalize": body.finalize, + "finalize": finalize, }))) } @@ -346,14 +1097,63 @@ async fn freeze_version( ) -> Result, StatusCode> { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; - if !is_valid_slug(&slug) { + if !is_valid_slug(&slug) + || body + .upload_id + .as_deref() + .is_some_and(|upload_id| !is_valid_page_upload_id(upload_id)) + || body + .expected_generation + .as_deref() + .is_some_and(|generation| !is_valid_page_generation(generation)) + { return Err(StatusCode::BAD_REQUEST); } + if (!body.title.is_empty() && !is_valid_page_title(&body.title)) + || !is_valid_page_note(&body.note) + { + return Err(StatusCode::BAD_REQUEST); + } + + let page_key = page_upload_session_key(&auth.user_id, &slug); + let _page_guard = state.page_upload_manager.lock_for(&page_key).lock().await; + let session = state + .page_upload_manager + .sessions + .get(&page_key) + .map(|entry| entry.value().clone()) + .ok_or(StatusCode::CONFLICT)?; + if session.upload_id != body.upload_id + || !upload_request_matches_session_intent( + body.expected_generation.as_deref(), + body.create, + &session, + ) + || !session.finalized + { + return Err(StatusCode::CONFLICT); + } + if session.expires_at <= Instant::now() { + state.page_upload_manager.sessions.remove(&page_key); + state.asset_store.cleanup_room(&session.draft_key); + return Err(StatusCode::GONE); + } - let page = PageRow::get(db, &auth.user_id, &slug) + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(&auth.user_id, &slug) + .await; + + let current_page = PageRow::get(db, &auth.user_id, &slug) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::NOT_FOUND)?; + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + if !generation_intent_matches( + current_page.as_ref(), + session.expected_generation.as_deref(), + session.create, + ) { + return Err(StatusCode::CONFLICT); + } let count = PageVersionRow::count_for_page(db, &auth.user_id, &slug) .await @@ -362,7 +1162,7 @@ async fn freeze_version( return Err(StatusCode::TOO_MANY_REQUESTS); } - let draft_key = page_draft_asset_key(&auth.user_id, &slug); + let draft_key = session.draft_key.clone(); if !state.asset_store.has_room_files(&draft_key) { return Err(StatusCode::BAD_REQUEST); } @@ -372,10 +1172,14 @@ async fn freeze_version( let asset_store = Arc::clone(&state.asset_store); let draft_key_c = draft_key.clone(); let version_key_c = version_key.clone(); - tokio::task::spawn_blocking(move || asset_store.copy_room(&draft_key_c, &version_key_c)) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let copied = + tokio::task::spawn_blocking(move || asset_store.copy_room(&draft_key_c, &version_key_c)) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + if copied.is_err() { + state.asset_store.cleanup_room(&version_key); + return Err(StatusCode::INTERNAL_SERVER_ERROR); + } let entries = state.asset_store.list_room_entries(&version_key); let has_worker = entries.iter().any(|(p, _)| p == WORKER_ENTRY_PATH); @@ -394,27 +1198,51 @@ async fn freeze_version( return Err(StatusCode::PAYLOAD_TOO_LARGE); } let total_bytes = total_bytes as i64; + let session_title = session.title.clone().ok_or(StatusCode::CONFLICT)?; + let visibility = session.visibility.ok_or(StatusCode::CONFLICT)?; let title = if body.title.is_empty() { - page.title.clone() + session_title } else { body.title.clone() }; - PageVersionRow::insert( + let save_outcome = PageRow::save_version_with_meta( db, &auth.user_id, &slug, - &version_id, + visibility, &title, + &version_id, file_count, total_bytes, has_worker, &body.note, + MAX_PAGES_PER_USER, + MAX_VERSIONS_PER_PAGE, + session.expected_generation.as_deref(), + session.create, ) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + .await; + match save_outcome { + Ok(SavePageVersionOutcome::Saved) => {} + Ok( + SavePageVersionOutcome::PageLimitReached | SavePageVersionOutcome::VersionLimitReached, + ) => { + state.asset_store.cleanup_room(&version_key); + return Err(StatusCode::TOO_MANY_REQUESTS); + } + Ok(SavePageVersionOutcome::GenerationMismatch) => { + state.asset_store.cleanup_room(&version_key); + return Err(StatusCode::CONFLICT); + } + Err(_) => { + state.asset_store.cleanup_room(&version_key); + return Err(StatusCode::INTERNAL_SERVER_ERROR); + } + } - // Clear draft after successful freeze. + // Consume the exact finalized upload session after a successful freeze. + state.page_upload_manager.sessions.remove(&page_key); state.asset_store.cleanup_room(&draft_key); let username = UserRow::find_by_username_for_user_id(db, &auth.user_id) @@ -422,8 +1250,14 @@ async fn freeze_version( .ok() .flatten() .unwrap_or_default(); + let generation = PageRow::get(db, &auth.user_id, &slug) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::CONFLICT)? + .generation; Ok(Json(VersionInfo { + generation, version_id: version_id.clone(), title, file_count, @@ -440,18 +1274,30 @@ async fn list_versions( State(state): State, headers: HeaderMap, Path(slug): Path, + Query(query): Query, ) -> Result>, StatusCode> { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; if !is_valid_slug(&slug) { return Err(StatusCode::BAD_REQUEST); } - let page = PageRow::get(db, &auth.user_id, &slug) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::NOT_FOUND)?; - maybe_migrate_legacy_page(&state, &page).await; - + let expected_generation = + validate_optional_expected_generation(query.expected_generation.as_deref())?; + // Legacy migration itself needs the write lease. Keep it through the + // Page/version snapshot so an omitted generation is bound exactly once + // and cannot drift across a local delete/recreate transition. + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(&auth.user_id, &slug) + .await; + let page = bind_page_generation_locked(db, &auth.user_id, &slug, expected_generation).await?; + let bound_generation = page.generation.clone(); + maybe_migrate_legacy_page_locked(&state, &auth.user_id, &slug, &bound_generation).await; + // The migration is generation-fenced in SQLite. Re-read before assembling + // the response as a second guard against another relay process replacing + // the Page while the migration performed storage work. + let page = + bind_page_generation_locked(db, &auth.user_id, &slug, Some(&bound_generation)).await?; let versions = PageVersionRow::list_for_page(db, &auth.user_id, &slug) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; @@ -466,6 +1312,7 @@ async fn list_versions( versions .into_iter() .map(|v| VersionInfo { + generation: page.generation.clone(), preview_url_path: format!("/p/{username}/{slug}/@v/{}", v.version_id), deployed: deployed.as_deref() == Some(v.version_id.as_str()), version_id: v.version_id, @@ -488,31 +1335,26 @@ async fn deploy_version( ) -> Result, StatusCode> { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; + let expected_generation = + validate_optional_expected_generation(body.expected_generation.as_deref())?; if !is_valid_slug(&slug) { return Err(StatusCode::BAD_REQUEST); } - let version = PageVersionRow::get(db, &auth.user_id, &slug, &body.version_id) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::NOT_FOUND)?; - - PageRow::set_deployed_version( - db, - &auth.user_id, - &slug, - &version.version_id, - version.file_count, - version.total_bytes, - &version.title, - ) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - let page = PageRow::get(db, &auth.user_id, &slug) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::NOT_FOUND)?; - let username = UserRow::find_by_username_for_user_id(db, &auth.user_id) + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(&auth.user_id, &slug) + .await; + let page = bind_page_generation_locked(db, &auth.user_id, &slug, expected_generation).await?; + let page = + match PageRow::deploy_version(db, &auth.user_id, &slug, &body.version_id, &page.generation) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + { + PageMutationOutcome::Applied(page) => page, + PageMutationOutcome::NotFound => return Err(StatusCode::NOT_FOUND), + PageMutationOutcome::GenerationMismatch => return Err(StatusCode::CONFLICT), + }; + let username = UserRow::find_by_username_for_user_id(db, &auth.user_id) .await .ok() .flatten() @@ -520,28 +1362,67 @@ async fn deploy_version( Ok(Json(page_to_info(&page, &username))) } -async fn delete_version( +async fn unpublish_page( State(state): State, headers: HeaderMap, - Path((slug, version_id)): Path<(String, String)>, -) -> Result, StatusCode> { + Path(slug): Path, + Query(query): Query, +) -> Result { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; + let expected_generation = + validate_optional_expected_generation(query.expected_generation.as_deref())?; if !is_valid_slug(&slug) { return Err(StatusCode::BAD_REQUEST); } - let page = PageRow::get(db, &auth.user_id, &slug) + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(&auth.user_id, &slug) + .await; + let page = bind_page_generation_locked(db, &auth.user_id, &slug, expected_generation).await?; + let outcome = PageRow::clear_deployed_version(db, &auth.user_id, &slug, &page.generation) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::NOT_FOUND)?; - if page.deployed_version_id.as_deref() == Some(version_id.as_str()) { + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + match outcome { + PageMutationOutcome::Applied(()) => {} + PageMutationOutcome::NotFound => return Err(StatusCode::NOT_FOUND), + PageMutationOutcome::GenerationMismatch => return Err(StatusCode::CONFLICT), + } + Ok(StatusCode::NO_CONTENT) +} + +async fn delete_version( + State(state): State, + headers: HeaderMap, + Path((slug, version_id)): Path<(String, String)>, + Query(query): Query, +) -> Result, StatusCode> { + let auth = validate_auth(&state, &headers).await?; + let db = require_db(&state)?; + let expected_generation = + validate_optional_expected_generation(query.expected_generation.as_deref())?; + if !is_valid_slug(&slug) { return Err(StatusCode::BAD_REQUEST); } - let deleted = PageVersionRow::delete(db, &auth.user_id, &slug, &version_id) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - if !deleted { - return Err(StatusCode::NOT_FOUND); + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(&auth.user_id, &slug) + .await; + let page = bind_page_generation_locked(db, &auth.user_id, &slug, expected_generation).await?; + match PageVersionRow::delete_if_not_deployed( + db, + &auth.user_id, + &slug, + &version_id, + &page.generation, + ) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + { + DeletePageVersionOutcome::Deleted => {} + DeletePageVersionOutcome::Deployed => return Err(StatusCode::BAD_REQUEST), + DeletePageVersionOutcome::NotFound => return Err(StatusCode::NOT_FOUND), + DeletePageVersionOutcome::GenerationMismatch => return Err(StatusCode::CONFLICT), } let key = page_version_asset_key(&auth.user_id, &slug, &version_id); state.asset_store.cleanup_room(&key); @@ -558,23 +1439,42 @@ async fn update_page( ) -> Result, StatusCode> { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; + let expected_generation = + validate_optional_expected_generation(body.expected_generation.as_deref())?; if !is_valid_slug(&slug) { return Err(StatusCode::BAD_REQUEST); } + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(&auth.user_id, &slug) + .await; + let page = bind_page_generation_locked(db, &auth.user_id, &slug, expected_generation).await?; let visibility = match body.visibility.as_deref() { Some(v) => Some(PageVisibility::parse(v).ok_or(StatusCode::BAD_REQUEST)?), None => None, }; - let updated = PageRow::update_meta(db, &auth.user_id, &slug, visibility, body.title.as_deref()) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - if !updated { - return Err(StatusCode::NOT_FOUND); + if body + .title + .as_deref() + .is_some_and(|title| !is_valid_page_title(title)) + { + return Err(StatusCode::BAD_REQUEST); } - let page = PageRow::get(db, &auth.user_id, &slug) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::NOT_FOUND)?; + let page = match PageRow::update_meta( + db, + &auth.user_id, + &slug, + &page.generation, + visibility, + body.title.as_deref(), + ) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + { + PageMutationOutcome::Applied(page) => page, + PageMutationOutcome::NotFound => return Err(StatusCode::NOT_FOUND), + PageMutationOutcome::GenerationMismatch => return Err(StatusCode::CONFLICT), + }; let username = UserRow::find_by_username_for_user_id(db, &auth.user_id) .await .ok() @@ -587,15 +1487,47 @@ async fn delete_page( State(state): State, headers: HeaderMap, Path(slug): Path, + Query(query): Query, ) -> Result, StatusCode> { let auth = validate_auth(&state, &headers).await?; let db = require_db(&state)?; + let expected_generation = + validate_optional_expected_generation(query.expected_generation.as_deref())?; if !is_valid_slug(&slug) { return Err(StatusCode::BAD_REQUEST); } + let page_key = page_upload_session_key(&auth.user_id, &slug); + let _page_guard = state.page_upload_manager.lock_for(&page_key).lock().await; + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(&auth.user_id, &slug) + .await; + let page = bind_page_generation_locked(db, &auth.user_id, &slug, expected_generation).await?; let versions = PageVersionRow::list_for_page(db, &auth.user_id, &slug) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + if let Some(store) = &state.page_data { + store + .prepare_generation(&auth.user_id, &slug, &page.generation) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + } + // Commit the relational deletion first. Filesystem/object cleanup is + // idempotent, whereas deleting assets before a failed multi-table DB + // mutation could leave a partially present Page with missing content. + // Mutable PageData has already been assigned to this Page generation, so a + // crash after this commit cannot expose it to a recreated Page generation. + let deleted = PageRow::delete(db, &auth.user_id, &slug, &page.generation) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + match deleted { + PageMutationOutcome::Applied(()) => {} + PageMutationOutcome::NotFound => return Err(StatusCode::NOT_FOUND), + PageMutationOutcome::GenerationMismatch => return Err(StatusCode::CONFLICT), + } + state.page_access_manager.revoke_page(&auth.user_id, &slug); + if let Some((_, active_upload)) = state.page_upload_manager.sessions.remove(&page_key) { + state.asset_store.cleanup_room(&active_upload.draft_key); + } for v in &versions { state.asset_store.cleanup_room(&page_version_asset_key( &auth.user_id, @@ -612,12 +1544,6 @@ async fn delete_page( if let Some(store) = &state.page_data { store.cleanup_page(&auth.user_id, &slug); } - let deleted = PageRow::delete(db, &auth.user_id, &slug) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - if !deleted { - return Err(StatusCode::NOT_FOUND); - } Ok(Json(serde_json::json!({ "status": "ok", "slug": slug }))) } @@ -706,22 +1632,54 @@ async fn serve_page( return Err(StatusCode::NOT_FOUND); } + let initial_page = PageRow::get_by_username(db, username, slug) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::NOT_FOUND)?; + // Legacy migration needs the lifecycle write lock. Modern Pages must stay + // on the shared serving path: taking that write lock for every production + // request would serialize traffic behind any running Page Function. Only + // a row with no deployment, no version, and legacy assets is a migration + // candidate; the locked helper revalidates the generation and all of these + // facts before committing. + let should_migrate_legacy = version_override.is_none() + && initial_page.deployed_version_id.is_none() + && PageVersionRow::count_for_page(db, &initial_page.user_id, slug) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + == 0 + && state + .asset_store + .has_room_files(&page_legacy_asset_key(&initial_page.user_id, slug)); + if should_migrate_legacy { + maybe_migrate_legacy_page_by_ids( + &state, + &initial_page.user_id, + slug, + &initial_page.generation, + ) + .await; + } + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_read(&initial_page.user_id, slug) + .await; + // A request may have resolved the Page just before a delete obtained the + // write lock. Re-resolve under the read lock so it cannot execute stale + // worker code or write mutable data after deletion/recreation. let page = PageRow::get_by_username(db, username, slug) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? .ok_or(StatusCode::NOT_FOUND)?; - enforce_visibility(&state, &headers, &page).await?; + enforce_visibility(&state, &headers, &page, version_override).await?; // Resolve version. let version_id = if let Some(v) = version_override { v.to_string() } else { - maybe_migrate_legacy_page_by_ids(&state, &page.user_id, slug).await; - let refreshed = PageRow::get(db, &page.user_id, slug) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::NOT_FOUND)?; - refreshed.deployed_version_id.ok_or(StatusCode::NOT_FOUND)? + page.deployed_version_id + .clone() + .ok_or(StatusCode::NOT_FOUND)? }; let version = PageVersionRow::get(db, &page.user_id, slug, &version_id) @@ -794,6 +1752,13 @@ async fn serve_with_worker( raw_path: &str, body: axum::body::Bytes, ) -> Result { + if body.len() > crate::page_execution::MAX_PAGE_FUNCTION_REQUEST_BODY_BYTES { + return Err(StatusCode::PAYLOAD_TOO_LARGE); + } + let _execution_permit = state + .page_execution_guard + .try_acquire(&page.user_id, &page.slug) + .map_err(|_| StatusCode::TOO_MANY_REQUESTS)?; let page_data = state.page_data.clone().ok_or(StatusCode::NOT_IMPLEMENTED)?; let db = state.db.clone().ok_or(StatusCode::NOT_IMPLEMENTED)?; let asset_store = Arc::clone(&state.asset_store); @@ -846,6 +1811,7 @@ async fn serve_with_worker( page_data, user_id: page.user_id.clone(), slug: page.slug.clone(), + generation: page.generation.clone(), meta: PageMeta { username: page.username.clone(), slug: page.slug.clone(), @@ -875,7 +1841,11 @@ async fn serve_with_worker( .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? .map_err(|e| { tracing::warn!("Page function error: {e}"); - StatusCode::BAD_GATEWAY + if matches!(e, PageFunctionError::Timeout(_)) { + StatusCode::GATEWAY_TIMEOUT + } else { + StatusCode::BAD_GATEWAY + } })?; // If worker returns 404 for document GET, fall back to static assets. @@ -908,7 +1878,9 @@ async fn serve_with_worker( axum::http::HeaderName::try_from(k), axum::http::HeaderValue::try_from(v), ) { - builder = builder.header(name, val); + if should_forward_page_worker_response_header(&name) { + builder = builder.header(name, val); + } } } builder @@ -916,6 +1888,38 @@ async fn serve_with_worker( .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR) } +/// Filter worker-controlled response headers that can mutate origin-wide +/// browser policy/state or connection framing. This is defense in depth: Page +/// documents still share an origin with the relay until hosting moves to a +/// dedicated Page origin. Ordinary representation and CORS headers remain +/// available to Page authors. +fn should_forward_page_worker_response_header(name: &axum::http::HeaderName) -> bool { + !matches!( + name.as_str(), + "accept-ch" + | "alt-svc" + | "clear-site-data" + | "connection" + | "content-length" + | "critical-ch" + | "keep-alive" + | "nel" + | "proxy-authenticate" + | "proxy-authorization" + | "report-to" + | "reporting-endpoints" + | "service-worker-allowed" + | "set-cookie" + | "set-cookie2" + | "strict-transport-security" + | "te" + | "trailer" + | "transfer-encoding" + | "upgrade" + | "x-content-type-options" + ) +} + // ── Helpers ───────────────────────────────────────────────────────────── fn page_to_info(page: &PageRow, username: &str) -> PageInfo { @@ -933,6 +1937,7 @@ fn page_to_info(page: &PageRow, username: &str) -> PageInfo { }); PageInfo { slug: page.slug.clone(), + generation: page.generation.clone(), visibility: page.visibility.clone(), title: page.title.clone(), file_count: page.file_count, @@ -946,16 +1951,37 @@ fn page_to_info(page: &PageRow, username: &str) -> PageInfo { } async fn maybe_migrate_legacy_page(state: &AppState, page: &PageRow) { - maybe_migrate_legacy_page_by_ids(state, &page.user_id, &page.slug).await; + maybe_migrate_legacy_page_by_ids(state, &page.user_id, &page.slug, &page.generation).await; +} + +async fn maybe_migrate_legacy_page_by_ids( + state: &AppState, + user_id: &str, + slug: &str, + expected_generation: &str, +) { + let _lifecycle_guard = state + .page_execution_guard + .acquire_page_write(user_id, slug) + .await; + maybe_migrate_legacy_page_locked(state, user_id, slug, expected_generation).await; } -async fn maybe_migrate_legacy_page_by_ids(state: &AppState, user_id: &str, slug: &str) { +async fn maybe_migrate_legacy_page_locked( + state: &AppState, + user_id: &str, + slug: &str, + expected_generation: &str, +) { let Some(db) = state.db.as_ref() else { return; }; let Ok(Some(page)) = PageRow::get(db, user_id, slug).await else { return; }; + if page.generation != expected_generation { + return; + } if page.deployed_version_id.is_some() { return; } @@ -971,40 +1997,72 @@ async fn maybe_migrate_legacy_page_by_ids(state: &AppState, user_id: &str, slug: } let version_id = "v1".to_string(); let version_key = page_version_asset_key(user_id, slug, &version_id); - if state.asset_store.copy_room(&legacy, &version_key).is_err() { + let legacy_entries = state + .asset_store + .list_room_entries(&legacy) + .into_iter() + .collect::>(); + let target_preexisted = state.asset_store.has_room_files(&version_key); + if target_preexisted { + let target_entries = state + .asset_store + .list_room_entries(&version_key) + .into_iter() + .collect::>(); + if target_entries != legacy_entries { + tracing::warn!( + "Skipped legacy Page migration because synthetic v1 assets already differ: {user_id}/{slug}" + ); + return; + } + } else if state.asset_store.copy_room(&legacy, &version_key).is_err() { return; } let entries = state.asset_store.list_room_entries(&version_key); let has_worker = entries.iter().any(|(p, _)| p == WORKER_ENTRY_PATH); - let _ = PageVersionRow::insert( + let migration = PageRow::migrate_legacy_version( db, user_id, slug, + expected_generation, &version_id, &page.title, page.file_count, page.total_bytes, has_worker, - "migrated", - ) - .await; - let _ = PageRow::set_deployed_version( - db, - user_id, - slug, - &version_id, - page.file_count, - page.total_bytes, - &page.title, ) .await; - tracing::info!("Migrated legacy page {user_id}/{slug} to version v1"); + if matches!(migration, Ok(PageMutationOutcome::Applied(true))) { + tracing::info!("Migrated legacy page {user_id}/{slug} to version v1"); + return; + } + + if !target_preexisted { + // A concurrent process may have committed the same v1 between the + // asset copy and our fenced transaction. Preserve the room only when + // the database now references it; otherwise this attempt owns and + // cleans the uncommitted synthetic target. + let target_is_committed = match PageRow::get(db, user_id, slug).await { + Ok(Some(current)) if current.deployed_version_id.as_deref() == Some(&version_id) => { + PageVersionRow::get(db, user_id, slug, &version_id) + .await + .ok() + .flatten() + .is_some() + } + _ => false, + }; + if !target_is_committed { + state.asset_store.cleanup_room(&version_key); + } + } } async fn enforce_visibility( state: &AppState, headers: &HeaderMap, page: &PageWithUsername, + version_id: Option<&str>, ) -> Result<(), StatusCode> { let visibility = page .visibility_enum() @@ -1012,19 +2070,39 @@ async fn enforce_visibility( match visibility { PageVisibility::Public => Ok(()), PageVisibility::Relay => { - resolve_viewer(state, headers).await?; - Ok(()) + if resolve_viewer(state, headers).await.is_ok() + || state.page_access_manager.authorizes_page( + headers, + &page.user_id, + &page.slug, + &page.generation, + version_id, + ) + { + Ok(()) + } else { + Err(StatusCode::UNAUTHORIZED) + } } PageVisibility::Private => { // Return NOT_FOUND for any auth failure so the existence of a // private page is not revealed to anonymous or foreign viewers. - let viewer = resolve_viewer(state, headers) - .await - .map_err(|_| StatusCode::NOT_FOUND)?; - if viewer.user_id != page.user_id { - return Err(StatusCode::NOT_FOUND); + if let Ok(viewer) = resolve_viewer(state, headers).await { + if viewer.user_id == page.user_id { + return Ok(()); + } + } + if state.page_access_manager.authorizes_page( + headers, + &page.user_id, + &page.slug, + &page.generation, + version_id, + ) { + Ok(()) + } else { + Err(StatusCode::NOT_FOUND) } - Ok(()) } } } @@ -1040,8 +2118,11 @@ async fn resolve_viewer(state: &AppState, headers: &HeaderMap) -> Result, asset_key: &str, + upload_id: Option, + expected_generation: Option, + create: bool, files: Vec, -) -> CheckPageFilesResponse { +) -> Result { let mut needed = Vec::new(); let mut existing_count = 0usize; let total_count = files.len(); @@ -1054,16 +2135,21 @@ fn process_check_page_files( } if asset_store.has_content(&entry.hash) { existing_count += 1; - let _ = asset_store.map_to_room(asset_key, &entry.path, &entry.hash); + asset_store + .map_to_room(asset_key, &entry.path, &entry.hash) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; } else { needed.push(entry.path); } } - CheckPageFilesResponse { + Ok(CheckPageFilesResponse { + upload_id, + expected_generation, + create, needed, existing_count, total_count, - } + }) } fn process_upload_page_files( @@ -1139,14 +2225,60 @@ fn mime_from_path(p: &str) -> &'static str { #[cfg(test)] mod tests { use super::*; - use crate::db::{connect, AuthToken, DeviceRow, UserRow}; + use crate::db::{connect, page_kv, AuthToken, DeviceRow, PageRow, UserRow}; use crate::relay::RoomManager; use crate::MemoryAssetStore; use axum::body::to_bytes; use axum::http::{Request, StatusCode}; use tower::ServiceExt; - async fn setup_app() -> (axum::Router, String, String) { + #[test] + fn management_generation_validation_only_normalizes_the_legacy_empty_value() { + let valid = "0123456789abcdef0123456789abcdef"; + assert_eq!(validate_optional_expected_generation(None).unwrap(), None); + assert_eq!( + validate_optional_expected_generation(Some("")).unwrap(), + None + ); + assert_eq!( + validate_optional_expected_generation(Some(valid)).unwrap(), + Some(valid) + ); + assert_eq!( + validate_optional_expected_generation(Some(" ")).unwrap_err(), + StatusCode::BAD_REQUEST + ); + assert_eq!( + validate_optional_expected_generation(Some("not-a-generation")).unwrap_err(), + StatusCode::BAD_REQUEST + ); + } + + #[test] + fn worker_response_headers_cannot_mutate_origin_wide_browser_state() { + for name in [ + "service-worker-allowed", + "set-cookie", + "clear-site-data", + "strict-transport-security", + "reporting-endpoints", + "transfer-encoding", + ] { + let name = axum::http::HeaderName::from_bytes(name.as_bytes()).unwrap(); + assert!(!should_forward_page_worker_response_header(&name), "{name}"); + } + + for name in [ + "content-type", + "cache-control", + "access-control-allow-origin", + ] { + let name = axum::http::HeaderName::from_bytes(name.as_bytes()).unwrap(); + assert!(should_forward_page_worker_response_header(&name), "{name}"); + } + } + + async fn setup_app_with_pool() -> (axum::Router, String, String, Arc) { let pool = connect(":memory:").await.unwrap(); let pool = Arc::new(pool); UserRow::create(&pool, "u1", "alice", "s", "ks", "{}", "hash", "wmk") @@ -1171,34 +2303,108 @@ mod tests { RoomManager::new(), Arc::new(MemoryAssetStore::new()), std::time::Instant::now(), - Some(pool), + Some(Arc::clone(&pool)), "test", Some(page_data_dir), ); - (app, tok_alice.token, tok_bob.token) + (app, tok_alice.token, tok_bob.token, pool) } - async fn save_and_deploy(app: &axum::Router, token: &str, slug: &str, html: &str) -> String { - save_and_deploy_with_visibility(app, token, slug, html, "public").await + async fn setup_app() -> (axum::Router, String, String) { + let (app, alice, bob, _) = setup_app_with_pool().await; + (app, alice, bob) } - async fn save_and_deploy_with_visibility( + async fn test_page_generation(app: &axum::Router, token: &str, slug: &str) -> Option { + let response = app + .clone() + .oneshot( + Request::builder() + .uri("/api/pages") + .header("Authorization", format!("Bearer {token}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let pages: Vec = serde_json::from_slice(&body).unwrap(); + pages.into_iter().find_map(|page| { + (page["slug"].as_str() == Some(slug)) + .then(|| page["generation"].as_str().map(str::to_string)) + .flatten() + }) + } + + async fn begin_test_upload( + app: &axum::Router, + token: &str, + slug: &str, + files: &[(&str, &[u8])], + ) -> (String, serde_json::Map) { + let upload_id = uuid::Uuid::new_v4().simple().to_string(); + let expected_generation = test_page_generation(app, token, slug).await; + let create = expected_generation.is_none(); + let mut manifest = Vec::new(); + let mut upload_files = serde_json::Map::new(); + for (path, content) in files { + let hash = hex_sha256(content); + manifest.push(serde_json::json!({ + "path": path, + "hash": hash, + "size": content.len(), + })); + upload_files.insert( + (*path).to_string(), + serde_json::json!({ "content": B64.encode(content), "hash": hash }), + ); + } + let check = serde_json::json!({ + "slug": slug, + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, + "files": manifest, + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(check.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + (upload_id, upload_files) + } + + async fn finish_test_upload( app: &axum::Router, token: &str, slug: &str, - html: &str, visibility: &str, + upload_id: &str, + files: serde_json::Map, ) -> String { - let hash = hex_sha256(html.as_bytes()); - let b64 = B64.encode(html.as_bytes()); + let expected_generation = test_page_generation(app, token, slug).await; + let create = expected_generation.is_none(); let upload = serde_json::json!({ "slug": slug, + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, "title": slug, "visibility": visibility, + "files": files, "finalize": true, - "files": { "index.html": { "content": b64, "hash": hash } } }); - let resp = app + let response = app .clone() .oneshot( Request::builder() @@ -1211,10 +2417,15 @@ mod tests { ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(response.status(), StatusCode::OK); - let freeze = serde_json::json!({ "title": slug, "note": "test" }); - let resp = app + let freeze = serde_json::json!({ + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, + "title": slug, + }); + let response = app .clone() .oneshot( Request::builder() @@ -1227,68 +2438,283 @@ mod tests { ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - let body = to_bytes(resp.into_body(), usize::MAX).await.unwrap(); - let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); - let version_id = v["version_id"].as_str().unwrap().to_string(); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let version: serde_json::Value = serde_json::from_slice(&body).unwrap(); + version["version_id"].as_str().unwrap().to_string() + } - let deploy = serde_json::json!({ "version_id": version_id }); - let resp = app + async fn save_version_with_legacy_generation_intent( + app: &axum::Router, + token: &str, + slug: &str, + html: &[u8], + ) -> (String, String) { + let upload_id = uuid::Uuid::new_v4().simple().to_string(); + let generation_before_check = test_page_generation(app, token, slug).await; + let hash = hex_sha256(html); + let check = serde_json::json!({ + "slug": slug, + "upload_id": upload_id, + "files": [{ "path": "index.html", "hash": hash, "size": html.len() }], + }); + let response = app .clone() .oneshot( Request::builder() .method("POST") - .uri(format!("/api/pages/{slug}/deploy")) + .uri("/api/pages/check-files") .header("Authorization", format!("Bearer {token}")) .header("content-type", "application/json") - .body(axum::body::Body::from(deploy.to_string())) + .body(axum::body::Body::from(check.to_string())) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - version_id - } + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let checked: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(checked["upload_id"].as_str(), Some(upload_id.as_str())); + assert_eq!( + checked["expected_generation"].as_str(), + generation_before_check.as_deref(), + ); + assert_eq!( + checked["create"].as_bool(), + Some(generation_before_check.is_none()) + ); - #[tokio::test] - async fn save_does_not_publish_until_deploy() { - let (app, alice, _) = setup_app().await; - let hash = hex_sha256(b"draft"); - let b64 = B64.encode(b"draft"); + // This is the pre-generation client wire shape: upload and freeze do + // not echo the new fields returned by check-files. let upload = serde_json::json!({ - "slug": "staged", - "title": "staged", - "visibility": "public", - "files": { "index.html": { "content": b64, "hash": hash } } + "slug": slug, + "upload_id": upload_id, + "title": slug, + "visibility": "private", + "files": { + "index.html": { "content": B64.encode(html), "hash": hash } + }, + "finalize": true, }); - let resp = app + let response = app .clone() .oneshot( Request::builder() .method("POST") .uri("/api/pages/upload-files") - .header("Authorization", format!("Bearer {alice}")) + .header("Authorization", format!("Bearer {token}")) .header("content-type", "application/json") .body(axum::body::Body::from(upload.to_string())) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(response.status(), StatusCode::OK); - let resp = app + let freeze = serde_json::json!({ + "upload_id": upload_id, + "title": slug, + "note": "legacy client", + }); + let response = app .clone() .oneshot( Request::builder() - .uri("/p/alice/staged") - .body(axum::body::Body::empty()) + .method("POST") + .uri(format!("/api/pages/{slug}/versions")) + .header("Authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(freeze.to_string())) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::NOT_FOUND); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let frozen: serde_json::Value = serde_json::from_slice(&body).unwrap(); + ( + frozen["generation"].as_str().unwrap().to_string(), + frozen["version_id"].as_str().unwrap().to_string(), + ) + } - let version_id = save_and_deploy(&app, &alice, "staged2", "live").await; + async fn save_and_deploy(app: &axum::Router, token: &str, slug: &str, html: &str) -> String { + save_and_deploy_with_visibility(app, token, slug, html, "public").await + } + + async fn save_and_deploy_with_visibility( + app: &axum::Router, + token: &str, + slug: &str, + html: &str, + visibility: &str, + ) -> String { + let hash = hex_sha256(html.as_bytes()); + let b64 = B64.encode(html.as_bytes()); + let upload_id = uuid::Uuid::new_v4().simple().to_string(); + let expected_generation = test_page_generation(app, token, slug).await; + let create = expected_generation.is_none(); + let check = serde_json::json!({ + "slug": slug, + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, + "files": [{ "path": "index.html", "hash": hash, "size": html.len() }] + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(check.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let upload = serde_json::json!({ + "slug": slug, + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, + "title": slug, + "visibility": visibility, + "finalize": true, + "files": { "index.html": { "content": b64, "hash": hash } } + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let freeze = serde_json::json!({ + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, + "title": slug, + "note": "test" + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(format!("/api/pages/{slug}/versions")) + .header("Authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(freeze.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let body = to_bytes(resp.into_body(), usize::MAX).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let version_id = v["version_id"].as_str().unwrap().to_string(); + let generation = v["generation"].as_str().unwrap(); + + let deploy = serde_json::json!({ + "version_id": version_id, + "expected_generation": generation, + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(format!("/api/pages/{slug}/deploy")) + .header("Authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(deploy.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + version_id + } + + #[tokio::test] + async fn save_does_not_publish_until_deploy() { + let (app, alice, _) = setup_app().await; + let hash = hex_sha256(b"draft"); + let b64 = B64.encode(b"draft"); + let upload_id = uuid::Uuid::new_v4().simple().to_string(); + let check = serde_json::json!({ + "slug": "staged", + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "files": [{ + "path": "index.html", + "hash": hash, + "size": b"draft".len() + }] + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(check.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let upload = serde_json::json!({ + "slug": "staged", + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "title": "staged", + "visibility": "public", + "files": { "index.html": { "content": b64, "hash": hash } } + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let resp = app + .clone() + .oneshot( + Request::builder() + .uri("/p/alice/staged") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + let version_id = save_and_deploy(&app, &alice, "staged2", "live").await; let resp = app .clone() .oneshot( @@ -1299,51 +2725,1435 @@ mod tests { ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(resp.status(), StatusCode::OK); + + let resp = app + .oneshot( + Request::builder() + .uri(format!("/p/alice/staged2/@v/{version_id}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + + #[tokio::test] + async fn republish_manifest_removes_files_not_present_in_the_new_version() { + let (app, alice, _) = setup_app().await; + let first_files = [ + ("index.html", b"first".as_slice()), + ("removed.txt", b"must not survive".as_slice()), + ]; + let (first_upload, first_map) = + begin_test_upload(&app, &alice, "clean-republish", &first_files).await; + let first_version = finish_test_upload( + &app, + &alice, + "clean-republish", + "public", + &first_upload, + first_map, + ) + .await; + assert_eq!( + get_page( + &app, + &format!("/p/alice/clean-republish/@v/{first_version}/removed.txt"), + None, + ) + .await, + StatusCode::OK + ); + + let second_files = [("index.html", b"second".as_slice())]; + let (second_upload, second_map) = + begin_test_upload(&app, &alice, "clean-republish", &second_files).await; + let second_version = finish_test_upload( + &app, + &alice, + "clean-republish", + "public", + &second_upload, + second_map, + ) + .await; + let response = app + .clone() + .oneshot( + Request::builder() + .uri(format!( + "/p/alice/clean-republish/@v/{second_version}/removed.txt" + )) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + // Static Page routing falls back to the new index.html for unknown + // paths. The old file content must not survive in the new manifest. + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + assert_eq!(String::from_utf8_lossy(&body), "second"); + } + + #[tokio::test] + async fn superseded_upload_session_cannot_mix_or_freeze_files() { + let (app, alice, _) = setup_app().await; + let first_files = [("index.html", b"first".as_slice())]; + let second_files = [("index.html", b"second".as_slice())]; + let (first_upload, first_map) = + begin_test_upload(&app, &alice, "concurrent", &first_files).await; + let (second_upload, second_map) = + begin_test_upload(&app, &alice, "concurrent", &second_files).await; + + let stale_upload = serde_json::json!({ + "slug": "concurrent", + "upload_id": first_upload, + "visibility": "private", + "files": first_map, + "finalize": true, + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(stale_upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::CONFLICT); + + let _ = finish_test_upload( + &app, + &alice, + "concurrent", + "private", + &second_upload, + second_map, + ) + .await; + let stale_freeze = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/concurrent/versions") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ "upload_id": first_upload }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(stale_freeze.status(), StatusCode::CONFLICT); + } + + #[tokio::test] + async fn legacy_upload_requests_without_upload_id_remain_compatible() { + let (app, alice, _) = setup_app().await; + let html = b"legacy"; + let hash = hex_sha256(html); + let check = serde_json::json!({ + "slug": "legacy-client", + "expected_generation": null, + "create": true, + "files": [{ "path": "index.html", "hash": hash, "size": html.len() }] + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(check.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let check_response: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(check_response.get("upload_id").is_none()); + + let upload = serde_json::json!({ + "slug": "legacy-client", + "expected_generation": null, + "create": true, + "title": "Legacy client", + "visibility": "private", + "files": { + "index.html": { "content": B64.encode(html), "hash": hash } + }, + "finalize": true + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/legacy-client/versions") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "expected_generation": null, + "create": true, + "title": "Legacy client", + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn pre_generation_upload_client_is_bound_to_new_and_existing_page_generations() { + let (app, alice, _) = setup_app().await; + let (created_generation, first_version) = save_version_with_legacy_generation_intent( + &app, + &alice, + "rolling-upgrade", + b"first", + ) + .await; + let (updated_generation, second_version) = save_version_with_legacy_generation_intent( + &app, + &alice, + "rolling-upgrade", + b"second", + ) + .await; + + assert_eq!(updated_generation, created_generation); + assert_ne!(second_version, first_version); + let response = app + .oneshot( + Request::builder() + .uri(format!( + "/api/pages/rolling-upgrade/versions?expected_generation={created_generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let versions: Vec = serde_json::from_slice(&body).unwrap(); + assert_eq!(versions.len(), 2); + assert!(versions + .iter() + .all(|version| version["generation"] == created_generation)); + } + + #[tokio::test] + async fn pre_generation_management_client_binds_each_request_to_the_current_page() { + let (app, alice, _) = setup_app().await; + let (generation, first_version) = save_version_with_legacy_generation_intent( + &app, + &alice, + "legacy-management", + b"first", + ) + .await; + let (same_generation, second_version) = save_version_with_legacy_generation_intent( + &app, + &alice, + "legacy-management", + b"second", + ) + .await; + assert_eq!(same_generation, generation); + + let deployed = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/legacy-management/deploy") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "version_id": first_version, + "expected_generation": "", + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(deployed.status(), StatusCode::OK); + + let open_ticket = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/legacy-management") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(open_ticket.status(), StatusCode::OK); + + let listed = app + .clone() + .oneshot( + Request::builder() + .uri("/api/pages/legacy-management/versions?expected_generation=") + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(listed.status(), StatusCode::OK); + let listed_body = to_bytes(listed.into_body(), usize::MAX).await.unwrap(); + let listed_versions: Vec = serde_json::from_slice(&listed_body).unwrap(); + assert_eq!(listed_versions.len(), 2); + assert!(listed_versions + .iter() + .all(|version| version["generation"] == generation)); + + let updated = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/pages/legacy-management") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ "title": "Updated by legacy client" }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(updated.status(), StatusCode::OK); + let updated_body = to_bytes(updated.into_body(), usize::MAX).await.unwrap(); + let updated_page: serde_json::Value = serde_json::from_slice(&updated_body).unwrap(); + assert_eq!(updated_page["generation"], generation); + assert_eq!(updated_page["title"], "Updated by legacy client"); + + let unpublished = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/legacy-management/unpublish") + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(unpublished.status(), StatusCode::NO_CONTENT); + + let redeployed = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/legacy-management/deploy") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ "version_id": second_version }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(redeployed.status(), StatusCode::OK); + + let deleted_version = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri(format!( + "/api/pages/legacy-management/versions/{first_version}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(deleted_version.status(), StatusCode::OK); + + let deleted_page = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri("/api/pages/legacy-management") + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(deleted_page.status(), StatusCode::OK); + assert!(test_page_generation(&app, &alice, "legacy-management") + .await + .is_none()); + } + + #[tokio::test] + async fn expired_session_pruning_waits_for_the_target_page_lock() { + let manager = Arc::new(PageUploadManager::new()); + let key = page_upload_session_key("u1", "locked"); + manager.sessions.insert( + key.clone(), + PageUploadSession { + user_id: "u1".to_string(), + upload_id: Some("a".repeat(32)), + draft_key: "pages/u1/locked/draft/id".to_string(), + manifest: HashMap::new(), + finalized: false, + title: None, + visibility: None, + expected_generation: None, + create: true, + expires_at: Instant::now() - Duration::from_secs(1), + }, + ); + let page_guard = manager.lock_for(&key).lock().await; + let pruning_manager = Arc::clone(&manager); + let mut pruning = + tokio::spawn(async move { pruning_manager.prune_expired(Instant::now()).await }); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut pruning) + .await + .is_err() + ); + manager.sessions.get_mut(&key).unwrap().expires_at = + Instant::now() + PAGE_UPLOAD_SESSION_TTL; + drop(page_guard); + assert!(pruning.await.unwrap().is_empty()); + assert!(manager.sessions.contains_key(&key)); + } + + #[tokio::test] + async fn upload_sessions_are_bounded_per_account_without_starving_others() { + let (app, alice, bob) = setup_app().await; + let hash = hex_sha256(b"x"); + for index in 0..MAX_PAGE_UPLOAD_SESSIONS_PER_USER { + let check = serde_json::json!({ + "slug": format!("site-{index}"), + "upload_id": format!("{index:032x}"), + "expected_generation": null, + "create": true, + "files": [{ "path": "index.html", "hash": hash, "size": 1 }] + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(check.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + let overflow = serde_json::json!({ + "slug": "site-overflow", + "upload_id": "f".repeat(32), + "expected_generation": null, + "create": true, + "files": [{ "path": "index.html", "hash": hash, "size": 1 }] + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(overflow.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + + let other_user = serde_json::json!({ + "slug": "other-user-site", + "upload_id": "e".repeat(32), + "expected_generation": null, + "create": true, + "files": [{ "path": "index.html", "hash": hash, "size": 1 }] + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {bob}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(other_user.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn page_metadata_and_max_file_upload_bounds_match_the_http_contract() { + assert!(is_valid_page_title("A valid title")); + assert!(!is_valid_page_title("bad\ntitle")); + assert!(!is_valid_page_title(&"x".repeat(MAX_PAGE_TITLE_CHARS + 1))); + assert!(is_valid_page_note("")); + assert!(!is_valid_page_note("bad\tnote")); + assert!(!is_valid_page_note(&"x".repeat(MAX_PAGE_NOTE_BYTES + 1))); + + let bytes = vec![0_u8; MAX_FILE_BYTES as usize]; + let encoded = B64.encode(&bytes); + let payload = serde_json::json!({ + "slug": "max-file", + "upload_id": "a".repeat(32), + "title": "Max file", + "visibility": "private", + "files": { + "index.html": { "content": encoded, "hash": hex_sha256(&bytes) } + }, + "finalize": true + }) + .to_string(); + assert!(payload.len() <= PAGE_UPLOAD_BODY_LIMIT); + + let (app, alice, _) = setup_app().await; + let files = [("index.html", bytes.as_slice())]; + let (upload_id, files) = begin_test_upload(&app, &alice, "max-file", &files).await; + let upload = serde_json::json!({ + "slug": "max-file", + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "title": "Max file", + "visibility": "private", + "files": files, + "finalize": true + }); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn page_metadata_bounds_are_enforced_by_upload_freeze_and_update_routes() { + let (app, alice, _) = setup_app().await; + let files = [("index.html", b"bounded".as_slice())]; + let (upload_id, upload_files) = + begin_test_upload(&app, &alice, "bounded-meta", &files).await; + let invalid_upload = serde_json::json!({ + "slug": "bounded-meta", + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "title": "bad\ntitle", + "visibility": "private", + "files": upload_files, + "finalize": true + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(invalid_upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + + let (upload_id, upload_files) = + begin_test_upload(&app, &alice, "bounded-meta", &files).await; + let valid_upload = serde_json::json!({ + "slug": "bounded-meta", + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "title": "Bounded", + "visibility": "private", + "files": upload_files, + "finalize": true + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(valid_upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/bounded-meta/versions") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "note": "bad\tnote", + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/bounded-meta/versions") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "note": "valid", + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let frozen: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let generation = frozen["generation"].as_str().unwrap(); + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/pages/bounded-meta") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "expected_generation": generation, + "title": "x".repeat(MAX_PAGE_TITLE_CHARS + 1), + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + + #[tokio::test] + async fn worker_can_use_kv() { + let (app, alice, _) = setup_app().await; + let worker = r#" + function fetch(request, env) { + if (request.path === "/api/hello") { + env.KV.put("msg", "hi"); + return { status: 200, headers: { "content-type": "text/plain" }, body: env.KV.get("msg") }; + } + return env.ASSETS.fetch("index.html"); + } + "#; + let files = [ + ("index.html", b"static".as_slice()), + ("server/worker.js", worker.as_bytes()), + ]; + let (upload_id, map) = begin_test_upload(&app, &alice, "fn", &files).await; + let upload = serde_json::json!({ + "slug": "fn", + "upload_id": upload_id, + "expected_generation": null, + "create": true, + "title": "fn", + "visibility": "public", + "files": map, + "finalize": true + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/upload-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(upload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let freeze = serde_json::json!({ + "upload_id": upload_id, + "expected_generation": null, + "create": true, + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/fn/versions") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(freeze.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let body = to_bytes(resp.into_body(), usize::MAX).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let version_id = v["version_id"].as_str().unwrap(); + let generation = v["generation"].as_str().unwrap(); + assert_eq!(v["has_worker"], true); + + let deploy = serde_json::json!({ + "version_id": version_id, + "expected_generation": generation, + }); + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/fn/deploy") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(deploy.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let resp = app + .oneshot( + Request::builder() + .uri(format!("/p/alice/fn/api/hello")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let body = to_bytes(resp.into_body(), usize::MAX).await.unwrap(); + assert_eq!(String::from_utf8_lossy(&body), "hi"); + } + + #[tokio::test] + async fn deployed_page_request_does_not_wait_for_running_worker_read_lease() { + let (app, alice, _, pool) = setup_app_with_pool().await; + let worker = r#" + function fetch(request, env) { + if (request.path === "/hold") { + env.KV.put("hold-started", "yes"); + const until = Date.now() + 1200; + while (Date.now() < until) {} + return { status: 200, headers: { "content-type": "text/plain" }, body: "held" }; + } + return { status: 200, headers: { "content-type": "text/plain" }, body: "fast" }; + } + "#; + let files = [("server/worker.js", worker.as_bytes())]; + let (upload_id, map) = begin_test_upload(&app, &alice, "shared-read", &files).await; + let version = + finish_test_upload(&app, &alice, "shared-read", "public", &upload_id, map).await; + let generation = test_page_generation(&app, &alice, "shared-read") + .await + .unwrap(); + let deploy = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/shared-read/deploy") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "version_id": version, + "expected_generation": generation, + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(deploy.status(), StatusCode::OK); + + let holding_request = tokio::spawn({ + let app = app.clone(); + async move { + app.oneshot( + Request::builder() + .uri("/p/alice/shared-read/hold") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap() + } + }); + let mut observed_start = false; + for _ in 0..200 { + if page_kv::get(&pool, "u1", "shared-read", "hold-started") + .await + .unwrap() + .as_deref() + == Some("yes") + { + observed_start = true; + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!(observed_start, "holding worker did not start"); + + let fast_response = tokio::time::timeout( + Duration::from_millis(400), + app.clone().oneshot( + Request::builder() + .uri("/p/alice/shared-read/fast") + .body(axum::body::Body::empty()) + .unwrap(), + ), + ) + .await + .expect("deployed Page request waited for an existing read lease") + .unwrap(); + assert_eq!(fast_response.status(), StatusCode::OK); + let body = to_bytes(fast_response.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!(body.as_ref(), b"fast"); + assert_eq!(holding_request.await.unwrap().status(), StatusCode::OK); + } + + #[tokio::test] + async fn page_delete_waits_for_running_worker_and_leaves_no_orphan_kv() { + let (app, alice, _, pool) = setup_app_with_pool().await; + let worker = r#" + function fetch(request, env) { + env.KV.put("started", "yes"); + const until = Date.now() + 150; + while (Date.now() < until) {} + env.KV.put("late", "yes"); + return { status: 200, headers: { "content-type": "text/plain" }, body: "done" }; + } + "#; + let files = [("server/worker.js", worker.as_bytes())]; + let (upload_id, map) = begin_test_upload(&app, &alice, "delete-running", &files).await; + let version = + finish_test_upload(&app, &alice, "delete-running", "public", &upload_id, map).await; + let generation = test_page_generation(&app, &alice, "delete-running") + .await + .unwrap(); + let deploy = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/delete-running/deploy") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "version_id": version, + "expected_generation": generation, + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(deploy.status(), StatusCode::OK); + + let worker_app = app.clone(); + let worker_request = tokio::spawn(async move { + worker_app + .oneshot( + Request::builder() + .uri("/p/alice/delete-running/run") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap() + }); + let mut observed_start = false; + for _ in 0..200 { + if page_kv::get(&pool, "u1", "delete-running", "started") + .await + .unwrap() + .as_deref() + == Some("yes") + { + observed_start = true; + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!(observed_start, "worker did not start before delete"); + + let mut deletion = tokio::spawn({ + let app = app.clone(); + let alice = alice.clone(); + let generation = generation.clone(); + async move { + app.oneshot( + Request::builder() + .method("DELETE") + .uri(format!( + "/api/pages/delete-running?expected_generation={generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap() + } + }); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut deletion) + .await + .is_err() + ); + assert_eq!(worker_request.await.unwrap().status(), StatusCode::OK); + assert_eq!(deletion.await.unwrap().status(), StatusCode::OK); + assert!(PageRow::get(&pool, "u1", "delete-running") + .await + .unwrap() + .is_none()); + assert!(page_kv::get(&pool, "u1", "delete-running", "late") + .await + .unwrap() + .is_none()); + } + + async fn get_page(app: &axum::Router, uri: &str, token: Option<&str>) -> StatusCode { + let mut builder = Request::builder().uri(uri); + if let Some(token) = token { + builder = builder.header("Authorization", format!("Bearer {token}")); + } + let resp = app + .clone() + .oneshot(builder.body(axum::body::Body::empty()).unwrap()) + .await + .unwrap(); + resp.status() + } + + async fn exchange_page_cookie( + app: &axum::Router, + token: &str, + slug: &str, + version_id: Option<&str>, + ) -> (String, String) { + let generation = test_page_generation(app, token, slug) + .await + .expect("Page must exist before creating an open ticket"); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(format!("/api/pages/{slug}")) + .header("Authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "version_id": version_id, + "expected_generation": generation, + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let response = app + .clone() + .oneshot( + Request::builder() + .uri(json["open_url_path"].as_str().unwrap()) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::TEMPORARY_REDIRECT); + let set_cookie = response + .headers() + .get(header::SET_COOKIE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + let cookie = set_cookie.split(';').next().unwrap().to_string(); + (cookie, set_cookie) + } + + #[tokio::test] + async fn visibility_matrix_private_relay_public() { + let (app, alice, bob) = setup_app().await; + save_and_deploy_with_visibility(&app, &alice, "priv", "p", "private").await; + save_and_deploy_with_visibility(&app, &alice, "rel", "r", "relay").await; + save_and_deploy(&app, &alice, "pub", "u").await; + + // Private: hidden from anonymous and foreign users, visible to owner. + assert_eq!( + get_page(&app, "/p/alice/priv", None).await, + StatusCode::NOT_FOUND + ); + assert_eq!( + get_page(&app, "/p/alice/priv", Some(&bob)).await, + StatusCode::NOT_FOUND + ); + assert_eq!( + get_page(&app, "/p/alice/priv", Some(&alice)).await, + StatusCode::OK + ); + // Query-string tokens must be ignored (same-origin JS could steal them). + assert_eq!( + get_page(&app, &format!("/p/alice/priv?access_token={alice}"), None).await, + StatusCode::NOT_FOUND + ); + + // Relay: any authenticated relay user may view, anonymous may not. + assert_eq!( + get_page(&app, "/p/alice/rel", None).await, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + get_page(&app, "/p/alice/rel", Some(&bob)).await, + StatusCode::OK + ); + assert_eq!( + get_page(&app, "/p/alice/rel", Some(&alice)).await, + StatusCode::OK + ); + + // Public: no credentials needed. + assert_eq!(get_page(&app, "/p/alice/pub", None).await, StatusCode::OK); + } + + #[tokio::test] + async fn one_time_open_ticket_exchanges_for_scoped_http_only_cookie() { + let (app, alice, bob) = setup_app().await; + let private_version = save_and_deploy_with_visibility( + &app, + &alice, + "private-open", + "p", + "private", + ) + .await; + save_and_deploy_with_visibility(&app, &alice, "other-private", "o", "private") + .await; + + let unauthorized = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/private-open") + .header("content-type", "application/json") + .body(axum::body::Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(unauthorized.status(), StatusCode::UNAUTHORIZED); + + let foreign = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/private-open") + .header("Authorization", format!("Bearer {bob}")) + .header("content-type", "application/json") + .body(axum::body::Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(foreign.status(), StatusCode::NOT_FOUND); + + let ticket_response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/private-open") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "expected_generation": test_page_generation( + &app, + &alice, + "private-open", + ) + .await + .unwrap(), + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(ticket_response.status(), StatusCode::OK); + let body = to_bytes(ticket_response.into_body(), usize::MAX) + .await + .unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let open_path = json["open_url_path"].as_str().unwrap(); + assert!(!open_path.contains(&alice)); + + let exchange = app + .clone() + .oneshot( + Request::builder() + .uri(open_path) + .header("x-forwarded-proto", "https") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(exchange.status(), StatusCode::TEMPORARY_REDIRECT); + assert_eq!( + exchange.headers().get(header::LOCATION).unwrap(), + "/p/alice/private-open" + ); + let cookie = exchange + .headers() + .get(header::SET_COOKIE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + assert!(cookie.contains("HttpOnly")); + assert!(cookie.contains("SameSite=Lax")); + assert!(cookie.contains("Secure")); + assert!(cookie.contains("Path=/p/alice/private-open")); + assert!(!cookie.contains(&alice)); + let browser_cookie = cookie.split(';').next().unwrap(); + + let replay = app + .clone() + .oneshot( + Request::builder() + .uri(open_path) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(replay.status(), StatusCode::NOT_FOUND); + + let authorized = app + .clone() + .oneshot( + Request::builder() + .uri("/p/alice/private-open") + .header(header::COOKIE, browser_cookie) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(authorized.status(), StatusCode::OK); - let resp = app + // A production ticket is not a blanket grant for immutable preview + // routes, even when the preview belongs to the same Page. + let wrong_version_scope = app + .clone() .oneshot( Request::builder() - .uri(format!("/p/alice/staged2/@v/{version_id}")) + .uri(format!("/p/alice/private-open/@v/{private_version}")) + .header(header::COOKIE, browser_cookie) .body(axum::body::Body::empty()) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(wrong_version_scope.status(), StatusCode::NOT_FOUND); + + let wrong_page = app + .oneshot( + Request::builder() + .uri("/p/alice/other-private") + .header(header::COOKIE, browser_cookie) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(wrong_page.status(), StatusCode::NOT_FOUND); } #[tokio::test] - async fn worker_can_use_kv() { + async fn production_and_preview_browser_grants_can_coexist() { let (app, alice, _) = setup_app().await; - let worker = r#" - function fetch(request, env) { - if (request.path === "/api/hello") { - env.KV.put("msg", "hi"); - return { status: 200, headers: { "content-type": "text/plain" }, body: env.KV.get("msg") }; - } - return env.ASSETS.fetch("index.html"); - } - "#; - let files = [ - ("index.html", b"static".as_slice()), - ("server/worker.js", worker.as_bytes()), - ]; - let mut map = serde_json::Map::new(); - for (path, content) in files { - let hash = hex_sha256(content); - map.insert( - path.to_string(), - serde_json::json!({ "content": B64.encode(content), "hash": hash }), - ); + let version = save_and_deploy_with_visibility( + &app, + &alice, + "cookie-scope", + "private", + "private", + ) + .await; + let (production_cookie, production_set_cookie) = + exchange_page_cookie(&app, &alice, "cookie-scope", None).await; + let (preview_cookie, preview_set_cookie) = + exchange_page_cookie(&app, &alice, "cookie-scope", Some(&version)).await; + assert!(production_set_cookie.contains("Path=/p/alice/cookie-scope;")); + assert!(preview_set_cookie.contains(&format!("Path=/p/alice/cookie-scope/@v/{version};"))); + + let combined = format!("{preview_cookie}; {production_cookie}"); + let production = app + .clone() + .oneshot( + Request::builder() + .uri("/p/alice/cookie-scope") + .header(header::COOKIE, &combined) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(production.status(), StatusCode::OK); + let preview = app + .oneshot( + Request::builder() + .uri(format!("/p/alice/cookie-scope/@v/{version}")) + .header(header::COOKIE, combined) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(preview.status(), StatusCode::OK); + } + + #[test] + fn access_grants_are_generation_bound_and_capacity_is_per_user() { + let manager = PageAccessManager::new(); + let ticket = manager + .issue_open_ticket( + "u1".into(), + "alice".into(), + "site".into(), + "generation-one".into(), + None, + ) + .unwrap(); + let (grant, _) = manager.exchange_ticket(&ticket).unwrap().unwrap(); + let mut headers = HeaderMap::new(); + headers.insert( + header::COOKIE, + HeaderValue::from_str(&format!("{PAGE_ACCESS_COOKIE}={grant}")).unwrap(), + ); + assert!(manager.authorizes_page(&headers, "u1", "site", "generation-one", None,)); + assert!(!manager.authorizes_page(&headers, "u1", "site", "generation-two", None,)); + + let bounded = PageAccessManager::new(); + for index in 0..MAX_PAGE_OPEN_TICKETS_PER_USER { + bounded + .issue_open_ticket( + "u1".into(), + "alice".into(), + format!("site-{index}"), + "generation".into(), + None, + ) + .unwrap(); } - let upload = serde_json::json!({ - "slug": "fn", - "title": "fn", - "visibility": "public", - "files": map - }); - let resp = app + assert_eq!( + bounded + .issue_open_ticket( + "u1".into(), + "alice".into(), + "overflow".into(), + "generation".into(), + None, + ) + .unwrap_err(), + StatusCode::TOO_MANY_REQUESTS + ); + assert!(bounded + .issue_open_ticket( + "u2".into(), + "bob".into(), + "other".into(), + "generation".into(), + None, + ) + .is_ok()); + } + + #[tokio::test] + async fn deleting_and_recreating_a_page_rejects_the_old_browser_grant() { + let (app, alice, _) = setup_app().await; + save_and_deploy_with_visibility(&app, &alice, "recreated", "old", "private") + .await; + let (old_cookie, _) = exchange_page_cookie(&app, &alice, "recreated", None).await; + let old_generation = test_page_generation(&app, &alice, "recreated") + .await + .unwrap(); + let deleted = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri(format!( + "/api/pages/recreated?expected_generation={old_generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(deleted.status(), StatusCode::OK); + save_and_deploy_with_visibility(&app, &alice, "recreated", "new", "private") + .await; + let response = app + .oneshot( + Request::builder() + .uri("/p/alice/recreated") + .header(header::COOKIE, old_cookie) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + + #[tokio::test] + async fn recreated_page_rejects_all_stale_generation_management_requests() { + let (app, alice, _) = setup_app().await; + save_and_deploy_with_visibility(&app, &alice, "aba", "old", "private").await; + let old_generation = test_page_generation(&app, &alice, "aba").await.unwrap(); + + // Leave an upload request owned by A in flight. Deleting A must make + // this request unusable even after a new Page B takes the same slug. + let stale_files = [("index.html", b"stale upload".as_slice())]; + let (stale_upload_id, stale_upload_files) = + begin_test_upload(&app, &alice, "aba", &stale_files).await; + + let deleted = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri(format!( + "/api/pages/aba?expected_generation={old_generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(deleted.status(), StatusCode::OK); + + let new_version = + save_and_deploy_with_visibility(&app, &alice, "aba", "new", "private") + .await; + let new_generation = test_page_generation(&app, &alice, "aba").await.unwrap(); + assert_ne!(new_generation, old_generation); + + let stale_upload = app .clone() .oneshot( Request::builder() @@ -1351,117 +4161,227 @@ mod tests { .uri("/api/pages/upload-files") .header("Authorization", format!("Bearer {alice}")) .header("content-type", "application/json") - .body(axum::body::Body::from(upload.to_string())) + .body(axum::body::Body::from( + serde_json::json!({ + "slug": "aba", + "upload_id": stale_upload_id, + "title": "stale", + "visibility": "private", + "files": stale_upload_files, + "finalize": true, + }) + .to_string(), + )) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - let freeze = serde_json::json!({}); - let resp = app + assert_eq!(stale_upload.status(), StatusCode::CONFLICT); + + let stale_update = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/pages/aba") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "expected_generation": old_generation, + "title": "stale overwrite", + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(stale_update.status(), StatusCode::CONFLICT); + + let stale_versions = app + .clone() + .oneshot( + Request::builder() + .uri(format!( + "/api/pages/aba/versions?expected_generation={old_generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(stale_versions.status(), StatusCode::CONFLICT); + + let stale_open = app .clone() .oneshot( Request::builder() .method("POST") - .uri("/api/pages/fn/versions") + .uri("/api/pages/aba") .header("Authorization", format!("Bearer {alice}")) .header("content-type", "application/json") - .body(axum::body::Body::from(freeze.to_string())) + .body(axum::body::Body::from( + serde_json::json!({ "expected_generation": old_generation }).to_string(), + )) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - let body = to_bytes(resp.into_body(), usize::MAX).await.unwrap(); - let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); - let version_id = v["version_id"].as_str().unwrap(); - assert_eq!(v["has_worker"], true); + assert_eq!(stale_open.status(), StatusCode::CONFLICT); - let deploy = serde_json::json!({ "version_id": version_id }); - let resp = app + let stale_deploy = app .clone() .oneshot( Request::builder() .method("POST") - .uri("/api/pages/fn/deploy") + .uri("/api/pages/aba/deploy") .header("Authorization", format!("Bearer {alice}")) .header("content-type", "application/json") - .body(axum::body::Body::from(deploy.to_string())) + .body(axum::body::Body::from( + serde_json::json!({ + "version_id": new_version, + "expected_generation": old_generation, + }) + .to_string(), + )) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(stale_deploy.status(), StatusCode::CONFLICT); - let resp = app + let stale_unpublish = app + .clone() .oneshot( Request::builder() - .uri(format!("/p/alice/fn/api/hello")) + .method("POST") + .uri(format!( + "/api/pages/aba/unpublish?expected_generation={old_generation}" + )) + .header("Authorization", format!("Bearer {alice}")) .body(axum::body::Body::empty()) .unwrap(), ) .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - let body = to_bytes(resp.into_body(), usize::MAX).await.unwrap(); - assert_eq!(String::from_utf8_lossy(&body), "hi"); - } + assert_eq!(stale_unpublish.status(), StatusCode::CONFLICT); - async fn get_page(app: &axum::Router, uri: &str, token: Option<&str>) -> StatusCode { - let mut builder = Request::builder().uri(uri); - if let Some(token) = token { - builder = builder.header("Authorization", format!("Bearer {token}")); - } - let resp = app + let stale_version_delete = app .clone() - .oneshot(builder.body(axum::body::Body::empty()).unwrap()) + .oneshot( + Request::builder() + .method("DELETE") + .uri(format!( + "/api/pages/aba/versions/{new_version}?expected_generation={old_generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) .await .unwrap(); - resp.status() - } + assert_eq!(stale_version_delete.status(), StatusCode::CONFLICT); - #[tokio::test] - async fn visibility_matrix_private_relay_public() { - let (app, alice, bob) = setup_app().await; - save_and_deploy_with_visibility(&app, &alice, "priv", "p", "private").await; - save_and_deploy_with_visibility(&app, &alice, "rel", "r", "relay").await; - save_and_deploy(&app, &alice, "pub", "u").await; + let stale_check = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/pages/check-files") + .header("Authorization", format!("Bearer {alice}")) + .header("content-type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "slug": "aba", + "upload_id": uuid::Uuid::new_v4().simple().to_string(), + "expected_generation": old_generation, + "create": false, + "files": [{ + "path": "index.html", + "hash": hex_sha256(b"stale"), + "size": 5, + }], + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(stale_check.status(), StatusCode::CONFLICT); - // Private: hidden from anonymous and foreign users, visible to owner. - assert_eq!( - get_page(&app, "/p/alice/priv", None).await, - StatusCode::NOT_FOUND - ); - assert_eq!( - get_page(&app, "/p/alice/priv", Some(&bob)).await, - StatusCode::NOT_FOUND - ); - assert_eq!( - get_page(&app, "/p/alice/priv", Some(&alice)).await, - StatusCode::OK - ); - // Query-string tokens must be ignored (same-origin JS could steal them). + let stale_delete = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri(format!( + "/api/pages/aba?expected_generation={old_generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(stale_delete.status(), StatusCode::CONFLICT); assert_eq!( - get_page(&app, &format!("/p/alice/priv?access_token={alice}"), None).await, - StatusCode::NOT_FOUND + test_page_generation(&app, &alice, "aba").await.as_deref(), + Some(new_generation.as_str()), ); + } - // Relay: any authenticated relay user may view, anonymous may not. - assert_eq!( - get_page(&app, "/p/alice/rel", None).await, - StatusCode::UNAUTHORIZED - ); + #[tokio::test] + async fn unpublish_stops_production_without_deleting_versions() { + let (app, alice, _) = setup_app().await; + let version_id = save_and_deploy(&app, &alice, "pause-me", "live").await; + let generation = test_page_generation(&app, &alice, "pause-me") + .await + .unwrap(); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(format!( + "/api/pages/pause-me/unpublish?expected_generation={generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NO_CONTENT); assert_eq!( - get_page(&app, "/p/alice/rel", Some(&bob)).await, - StatusCode::OK + get_page(&app, "/p/alice/pause-me", None).await, + StatusCode::NOT_FOUND ); assert_eq!( - get_page(&app, "/p/alice/rel", Some(&alice)).await, + get_page(&app, &format!("/p/alice/pause-me/@v/{version_id}"), None,).await, StatusCode::OK ); - // Public: no credentials needed. - assert_eq!(get_page(&app, "/p/alice/pub", None).await, StatusCode::OK); + let versions = app + .oneshot( + Request::builder() + .uri(format!( + "/api/pages/pause-me/versions?expected_generation={generation}" + )) + .header("Authorization", format!("Bearer {alice}")) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(versions.status(), StatusCode::OK); + let body = to_bytes(versions.into_body(), usize::MAX).await.unwrap(); + let list: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(list.as_array().unwrap().len(), 1); + assert_eq!(list[0]["deployed"], false); } #[tokio::test] @@ -1474,19 +4394,16 @@ mod tests { } "#; let files = [("server/worker.js", worker.as_bytes())]; - let mut map = serde_json::Map::new(); - for (path, content) in files { - let hash = hex_sha256(content); - map.insert( - path.to_string(), - serde_json::json!({ "content": B64.encode(content), "hash": hash }), - ); - } + let (upload_id, map) = begin_test_upload(&app, &alice, "hdr", &files).await; let upload = serde_json::json!({ "slug": "hdr", + "upload_id": upload_id, + "expected_generation": null, + "create": true, "title": "hdr", "visibility": "relay", - "files": map + "files": map, + "finalize": true }); let resp = app .clone() @@ -1510,7 +4427,14 @@ mod tests { .uri("/api/pages/hdr/versions") .header("Authorization", format!("Bearer {alice}")) .header("content-type", "application/json") - .body(axum::body::Body::from("{}")) + .body(axum::body::Body::from( + serde_json::json!({ + "upload_id": upload_id, + "expected_generation": null, + "create": true, + }) + .to_string(), + )) .unwrap(), ) .await @@ -1518,7 +4442,11 @@ mod tests { let body = to_bytes(resp.into_body(), usize::MAX).await.unwrap(); let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); let version_id = v["version_id"].as_str().unwrap(); - let deploy = serde_json::json!({ "version_id": version_id }); + let generation = v["generation"].as_str().unwrap(); + let deploy = serde_json::json!({ + "version_id": version_id, + "expected_generation": generation, + }); let resp = app .clone() .oneshot( @@ -1574,4 +4502,21 @@ mod tests { assert_eq!(page["file_count"].as_i64().unwrap(), 1); assert_eq!(page["total_bytes"].as_i64().unwrap(), html.len() as i64); } + + #[tokio::test] + async fn page_request_body_limit_is_applied_before_worker_execution() { + let (app, _, _) = setup_app().await; + let body = vec![0u8; crate::page_execution::MAX_PAGE_FUNCTION_REQUEST_BODY_BYTES + 1]; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/p/alice/missing/api") + .body(axum::body::Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE); + } } diff --git a/src/crates/services/relay-service/src/routes/websocket.rs b/src/crates/services/relay-service/src/routes/websocket.rs index 3478a891a4..885fa5a074 100644 --- a/src/crates/services/relay-service/src/routes/websocket.rs +++ b/src/crates/services/relay-service/src/routes/websocket.rs @@ -14,7 +14,11 @@ use axum::{ }; use futures_util::{SinkExt, StreamExt}; use serde::{Deserialize, Serialize}; -use std::time::{Duration, Instant}; +use std::{ + collections::HashSet, + sync::Arc, + time::{Duration, Instant}, +}; use tokio::sync::{ mpsc::{self, error::TrySendError}, watch, @@ -33,6 +37,7 @@ const MAX_PUBLIC_KEY_BYTES: usize = 512; const MAX_NONCE_BYTES: usize = 256; const RATE_LIMIT_WINDOW: Duration = Duration::from_secs(60); const MAX_MESSAGES_PER_WINDOW: u32 = 600; +const DEVICE_TOKEN_REVALIDATION_INTERVAL: Duration = Duration::from_secs(5); struct ConnectionRateLimiter { window_started: Instant, @@ -197,6 +202,7 @@ fn is_websocket_origin_allowed(headers: &HeaderMap, allowed_origins: &[String]) } async fn handle_socket(socket: WebSocket, state: AppState) { + ensure_device_token_revalidator(&state); let (mut ws_sender, mut ws_receiver) = socket.split(); let (out_tx, mut out_rx) = mpsc::channel::(OUTBOUND_QUEUE_CAPACITY); let (force_close_tx, mut force_close_rx) = watch::channel(false); @@ -207,6 +213,7 @@ async fn handle_socket(socket: WebSocket, state: AppState) { info!("WebSocket connected: conn_id={conn_id}"); let mut writer_force_close_rx = force_close_rx.clone(); + let writer_failure_close_tx = force_close_tx.clone(); let write_task = tokio::spawn(async move { loop { tokio::select! { @@ -225,6 +232,11 @@ async fn handle_socket(socket: WebSocket, state: AppState) { .await .is_err() { + // The read half can remain open after a write-half + // failure. Wake the owner loop so it promptly removes + // routing/presence instead of leaving a half-open + // device online until token expiry. + let _ = writer_failure_close_tx.send(true); break; } } @@ -234,6 +246,7 @@ async fn handle_socket(socket: WebSocket, state: AppState) { loop { let msg_result = tokio::select! { + biased; changed = force_close_rx.changed() => { if changed.is_ok() && *force_close_rx.borrow() { info!("Revoked WebSocket connection: conn_id={conn_id}"); @@ -298,25 +311,116 @@ async fn handle_socket(socket: WebSocket, state: AppState) { } state.room_manager.on_disconnect(conn_id); + let _presence_projection_guard = state.device_manager.lock_presence_projection().await; if let Some((user_id, device_id)) = state.device_manager.unregister(conn_id) { // Best-effort: mark the device offline in the DB and notify peers. - if let Some(db) = state.db.as_ref() { - let _ = crate::db::DeviceRow::set_online(db, &user_id, &device_id, false).await; + if !state.device_manager.is_device_online(&user_id, &device_id) { + if let Some(db) = state.db.as_ref() { + let _ = crate::db::DeviceRow::set_online(db, &user_id, &device_id, false).await; + } } - let remaining = state.device_manager.online_devices(&user_id); - let presence = build_presence(&remaining); - state.device_manager.broadcast_except( - &user_id, - &device_id, - &serde_json::to_string(&OutboundProtocol::DevicePresence { devices: presence }) - .unwrap_or_default(), - ); + state + .device_manager + .broadcast_current_presence(&user_id, |devices| { + serde_json::to_string(&OutboundProtocol::DevicePresence { + devices: build_presence(devices), + }) + .ok() + }); } + drop(_presence_projection_guard); drop(out_tx); let _ = write_task.await; info!("WebSocket disconnected: conn_id={conn_id}"); } +fn ensure_device_token_revalidator(state: &AppState) { + let Some(db) = state.db.clone() else { + return; + }; + if !state.device_manager.claim_token_revalidator_start() { + return; + } + let device_manager = Arc::clone(&state.device_manager); + tokio::spawn(async move { + let mut interval = tokio::time::interval(DEVICE_TOKEN_REVALIDATION_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // Consume the immediate first tick. AuthConnect performs its own two + // checks; this worker is for later out-of-process revocation. + interval.tick().await; + loop { + interval.tick().await; + if let Err(error) = revalidate_active_device_tokens_once(&db, &device_manager).await { + warn!(%error, "Failed to revalidate active device tokens"); + } + } + }); +} + +async fn revalidate_active_device_tokens_once( + db: &crate::db::DbPool, + device_manager: &crate::relay::DeviceManager, +) -> anyhow::Result { + let active = device_manager.active_device_credentials(); + if active.is_empty() { + return Ok(0); + } + let tokens = active + .iter() + .map(|(_, _, _, token)| token.clone()) + .collect::>(); + let valid = crate::db::AuthToken::find_valid_device_tokens(db, &tokens) + .await? + .into_iter() + .map(|token| (token.user_id, token.device_id, token.token)) + .collect::>(); + let revoked = active + .into_iter() + .filter(|(_, user_id, device_id, token)| { + !valid.contains(&(user_id.clone(), device_id.clone(), token.clone())) + }) + .collect::>(); + if revoked.is_empty() { + return Ok(0); + } + + // A batch snapshot can race a same-process login/logout. Re-check each + // missing exact token under the lifecycle gate before removing its socket. + let _presence_projection_guard = device_manager.lock_presence_projection().await; + let mut affected_users = HashSet::new(); + let mut disconnected = 0; + for (_conn_id, user_id, device_id, token) in revoked { + match registered_device_token_is_current(db, &token, &user_id, &device_id).await { + Ok(true) => continue, + Err(error) => { + warn!(%error, %user_id, %device_id, "Failed exact token revalidation"); + continue; + } + Ok(false) => {} + } + if device_manager.disconnect_device_if_token(&user_id, &device_id, &token) { + disconnected += 1; + affected_users.insert(user_id.clone()); + if !device_manager.is_device_online(&user_id, &device_id) { + if let Err(error) = + crate::db::DeviceRow::set_online(db, &user_id, &device_id, false).await + { + warn!(%error, %user_id, %device_id, "Failed to project revoked device offline"); + } + } + } + } + for user_id in affected_users { + device_manager.broadcast_current_presence(&user_id, |devices| { + serde_json::to_string(&OutboundProtocol::DevicePresence { + devices: build_presence(devices), + }) + .ok() + }); + } + Ok(disconnected) +} + async fn handle_text_message( text: &str, conn_id: ConnId, @@ -438,7 +542,7 @@ async fn handle_text_message( InboundMessage::AuthConnect { token, device_name } => { if state.room_manager.has_connection(conn_id) - || state.device_manager.conn_mapping(conn_id).is_some() + || state.device_manager.has_connection(conn_id) { return reject_protocol(out_tx, "connection is already authenticated"); } @@ -474,25 +578,39 @@ async fn handle_text_message( }, ); } - // Mark the device online in the DB and the in-memory registry. - let _ = crate::db::DeviceRow::upsert( - db, - &auth.device_id, - &auth.user_id, - &device_name, - None, - ) - .await; - let _ = - crate::db::DeviceRow::set_online(db, &auth.user_id, &auth.device_id, true).await; - let _others = state.device_manager.register( + // Make the candidate visible to token-scoped logout without + // exposing it to presence/routing or replacing the current active + // socket until the mandatory final token check succeeds. + state.device_manager.register_pending( &auth.user_id, &auth.device_id, + &token, &device_name, conn_id, out_tx.clone(), force_close_tx.clone(), ); + + let activated = match activate_pending_device_if_authorized( + db, + &state.device_manager, + &token, + &auth.user_id, + &auth.device_id, + &device_name, + conn_id, + ) + .await + { + Ok(activated) => activated, + Err(error) => { + warn!(%error, "Failed to activate authenticated device"); + false + } + }; + if !activated { + return false; + } let expires_in = auth .expires_at .saturating_sub(chrono::Utc::now().timestamp()) @@ -502,30 +620,17 @@ async fn handle_text_message( tokio::time::sleep(Duration::from_secs(expires_in)).await; let _ = expiry_close_tx.send(true); })); - send_json( - out_tx, - &OutboundProtocol::AuthOk { - user_id: auth.user_id.clone(), - device_id: auth.device_id.clone(), - }, - ) - .await; // Full presence (including self) so clients can treat the snapshot - // as authoritative rather than an incremental patch. - let all_online = state.device_manager.online_devices(&auth.user_id); - let presence = build_presence(&all_online); - send_json_best_effort( - out_tx, - &OutboundProtocol::DevicePresence { - devices: presence.clone(), - }, - ); - state.device_manager.broadcast_except( - &auth.user_id, - &auth.device_id, - &serde_json::to_string(&OutboundProtocol::DevicePresence { devices: presence }) - .unwrap_or_default(), - ); + // as authoritative rather than an incremental patch. Snapshot and + // enqueue are serialized with membership mutations. + state + .device_manager + .broadcast_current_presence(&auth.user_id, |devices| { + serde_json::to_string(&OutboundProtocol::DevicePresence { + devices: build_presence(devices), + }) + .ok() + }); true } @@ -632,6 +737,165 @@ fn build_presence(devices: &[(String, String)]) -> Vec { .collect() } +async fn registered_device_token_is_current( + db: &crate::db::DbPool, + token: &str, + expected_user_id: &str, + expected_device_id: &str, +) -> anyhow::Result { + Ok(crate::db::AuthToken::find(db, token) + .await? + .is_some_and(|current| { + current.is_device_token() + && current.user_id == expected_user_id + && current.device_id == expected_device_id + })) +} + +async fn project_device_offline_if_unowned( + db: &crate::db::DbPool, + device_manager: &crate::relay::DeviceManager, + user_id: &str, + device_id: &str, +) -> anyhow::Result<()> { + if !device_manager.is_device_online(user_id, device_id) { + crate::db::DeviceRow::set_online(db, user_id, device_id, false).await?; + } + Ok(()) +} + +async fn reconcile_device_after_token_disconnect( + db: &crate::db::DbPool, + device_manager: &crate::relay::DeviceManager, + user_id: &str, + device_id: &str, +) -> anyhow::Result<()> { + let projection_result = + project_device_offline_if_unowned(db, device_manager, user_id, device_id).await; + device_manager.broadcast_current_presence(user_id, |devices| { + serde_json::to_string(&OutboundProtocol::DevicePresence { + devices: build_presence(devices), + }) + .ok() + }); + projection_result +} + +/// Complete the activation after the durable device projection has been +/// written. The caller must hold `presence_projection_gate`, keeping the +/// second token check, in-memory promotion, and any rollback atomic with +/// logout, device deletion, and socket cleanup. +async fn complete_pending_device_activation_if_authorized( + db: &crate::db::DbPool, + device_manager: &crate::relay::DeviceManager, + token: &str, + expected_user_id: &str, + expected_device_id: &str, + conn_id: ConnId, +) -> anyhow::Result { + // The durable writes can wait on SQLite. Re-check wall-clock expiry + // immediately before the synchronous promotion so a token that expired + // during that wait never becomes routable or receives AuthOk. + let still_current = + match registered_device_token_is_current(db, token, expected_user_id, expected_device_id) + .await + { + Ok(current) => current, + Err(error) => { + device_manager.disconnect_pending(conn_id); + let _ = project_device_offline_if_unowned( + db, + device_manager, + expected_user_id, + expected_device_id, + ) + .await; + return Err(error); + } + }; + if !still_current { + device_manager.disconnect_device_if_token(expected_user_id, expected_device_id, token); + reconcile_device_after_token_disconnect( + db, + device_manager, + expected_user_id, + expected_device_id, + ) + .await?; + return Ok(false); + } + + // Pending remains invisible until every durable write succeeds. Promotion + // is synchronous under the same lifecycle gate, so clients never route to + // a socket whose AuthConnect may still fail on database projection. + let auth_ok = serde_json::to_string(&OutboundProtocol::AuthOk { + user_id: expected_user_id.to_string(), + device_id: expected_device_id.to_string(), + })?; + if !device_manager.activate_pending_with_initial_message( + expected_user_id, + expected_device_id, + token, + conn_id, + &auth_ok, + ) { + project_device_offline_if_unowned(db, device_manager, expected_user_id, expected_device_id) + .await?; + return Ok(false); + } + Ok(true) +} + +async fn activate_pending_device_if_authorized( + db: &crate::db::DbPool, + device_manager: &crate::relay::DeviceManager, + token: &str, + expected_user_id: &str, + expected_device_id: &str, + device_name: &str, + conn_id: ConnId, +) -> anyhow::Result { + // Serialize the final token lookup, durable device update, activation, and + // online projection with logout/delete/socket cleanup. No persistent side + // effect occurs before this lookup, so a deleted device cannot be revived + // by an AuthConnect request that passed only the earlier lookup. + let _presence_projection_guard = device_manager.lock_presence_projection().await; + if !registered_device_token_is_current(db, token, expected_user_id, expected_device_id).await? { + device_manager.disconnect_device_if_token(expected_user_id, expected_device_id, token); + reconcile_device_after_token_disconnect( + db, + device_manager, + expected_user_id, + expected_device_id, + ) + .await?; + return Ok(false); + } + + if let Err(error) = + crate::db::DeviceRow::upsert(db, expected_device_id, expected_user_id, device_name, None) + .await + { + device_manager.disconnect_pending(conn_id); + return Err(error); + } + if let Err(error) = + crate::db::DeviceRow::set_online(db, expected_user_id, expected_device_id, true).await + { + device_manager.disconnect_pending(conn_id); + return Err(error); + } + complete_pending_device_activation_if_authorized( + db, + device_manager, + token, + expected_user_id, + expected_device_id, + conn_id, + ) + .await +} + async fn send_json(tx: &mpsc::Sender, msg: &T) -> bool { match serde_json::to_string(msg) { Ok(json) => send_outbound_message(tx, OutboundMessage::text(json)).await, @@ -667,12 +931,16 @@ fn generate_room_id() -> String { #[cfg(test)] mod tests { use super::{ + activate_pending_device_if_authorized, complete_pending_device_activation_if_authorized, is_valid_display_text, is_valid_encrypted_payload, is_valid_identifier, - is_websocket_origin_allowed, send_json_best_effort, ConnectionRateLimiter, + is_websocket_origin_allowed, registered_device_token_is_current, + revalidate_active_device_tokens_once, send_json_best_effort, ConnectionRateLimiter, OutboundProtocol, MAX_MESSAGES_PER_WINDOW, }; + use crate::db::{connect, AuthToken, DeviceRow, UserRow}; + use crate::relay::DeviceManager; use axum::http::{header, HeaderMap}; - use tokio::sync::mpsc; + use tokio::sync::{mpsc, watch}; #[test] fn best_effort_control_response_does_not_block_on_full_queue() { @@ -731,4 +999,364 @@ mod tests { assert!(!is_websocket_origin_allowed(&headers, &[])); } } + + #[tokio::test] + async fn token_revoked_between_initial_validation_and_activation_is_rejected() { + let db = connect(":memory:").await.unwrap(); + UserRow::create(&db, "owner", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + DeviceRow::upsert(&db, "device-a", "owner", "Device A", None) + .await + .unwrap(); + let token = AuthToken::create(&db, "owner", "device-a") + .await + .unwrap() + .token; + + // This is the first AuthConnect lookup. Pause the conceptual request + // after it, then let logout delete the row before registration's + // mandatory second lookup. + assert!(AuthToken::find(&db, &token).await.unwrap().is_some()); + sqlx::query("DELETE FROM auth_tokens WHERE token = ?") + .bind(&token) + .execute(&db) + .await + .unwrap(); + + let manager = DeviceManager::new(); + let (tx, _rx) = mpsc::channel(4); + let (close_tx, mut close_rx) = watch::channel(false); + manager.register_pending("owner", "device-a", &token, "Device A", 1, tx, close_tx); + assert!(manager.online_devices("owner").is_empty()); + assert!(manager.conn_mapping(1).is_none()); + + assert!( + !registered_device_token_is_current(&db, &token, "owner", "device-a") + .await + .unwrap() + ); + assert!(!activate_pending_device_if_authorized( + &db, &manager, &token, "owner", "device-a", "Device A", 1, + ) + .await + .unwrap()); + close_rx.changed().await.unwrap(); + assert!(*close_rx.borrow()); + assert!(manager.conn_mapping(1).is_none()); + } + + #[tokio::test] + async fn expired_token_disconnects_active_and_pending_without_ghost_online_projection() { + let db = connect(":memory:").await.unwrap(); + UserRow::create(&db, "owner", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + DeviceRow::upsert(&db, "device-a", "owner", "Device A", None) + .await + .unwrap(); + DeviceRow::set_online(&db, "owner", "device-a", true) + .await + .unwrap(); + let token = AuthToken::create(&db, "owner", "device-a") + .await + .unwrap() + .token; + + let manager = DeviceManager::new(); + let (active_tx, _active_rx) = mpsc::channel(4); + let (active_close_tx, mut active_close_rx) = watch::channel(false); + manager.register( + "owner", + "device-a", + &token, + "Device A", + 1, + active_tx, + active_close_tx, + ); + let (pending_tx, _pending_rx) = mpsc::channel(4); + let (pending_close_tx, mut pending_close_rx) = watch::channel(false); + manager.register_pending( + "owner", + "device-a", + &token, + "Device A", + 2, + pending_tx, + pending_close_tx, + ); + let (peer_tx, mut peer_rx) = mpsc::channel(4); + let (peer_close_tx, _peer_close_rx) = watch::channel(false); + manager.register( + "owner", + "device-c", + "independent-token", + "Device C", + 3, + peer_tx, + peer_close_tx, + ); + + sqlx::query("UPDATE auth_tokens SET expires_at = ? WHERE token = ?") + .bind(chrono::Utc::now().timestamp()) + .bind(&token) + .execute(&db) + .await + .unwrap(); + assert!(!activate_pending_device_if_authorized( + &db, &manager, &token, "owner", "device-a", "Device A", 2, + ) + .await + .unwrap()); + + active_close_rx.changed().await.unwrap(); + pending_close_rx.changed().await.unwrap(); + assert!(*active_close_rx.borrow()); + assert!(*pending_close_rx.borrow()); + assert!(manager.conn_mapping(1).is_none()); + assert!(manager.conn_mapping(2).is_none()); + assert_eq!( + manager.conn_mapping(3), + Some(("owner".into(), "device-c".into())) + ); + assert!(!manager.route_message("owner", "device-a", "opaque")); + let presence: serde_json::Value = + serde_json::from_str(&peer_rx.recv().await.unwrap().text).unwrap(); + assert_eq!(presence["type"], "device_presence"); + assert_eq!(presence["devices"].as_array().unwrap().len(), 1); + assert_eq!(presence["devices"][0]["device_id"], "device-c"); + let rows = DeviceRow::list_by_user(&db, "owner").await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].online, 0); + } + + #[tokio::test] + async fn external_token_revocation_reaper_disconnects_idle_device_and_updates_presence() { + let db = connect(":memory:").await.unwrap(); + UserRow::create(&db, "owner", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + DeviceRow::upsert(&db, "device-a", "owner", "Device A", None) + .await + .unwrap(); + DeviceRow::upsert(&db, "device-c", "owner", "Device C", None) + .await + .unwrap(); + DeviceRow::set_online(&db, "owner", "device-a", true) + .await + .unwrap(); + DeviceRow::set_online(&db, "owner", "device-c", true) + .await + .unwrap(); + let revoked_token = AuthToken::create(&db, "owner", "device-a") + .await + .unwrap() + .token; + let peer_token = AuthToken::create(&db, "owner", "device-c") + .await + .unwrap() + .token; + + let manager = DeviceManager::new(); + let (revoked_tx, _revoked_rx) = mpsc::channel(4); + let (revoked_close_tx, mut revoked_close_rx) = watch::channel(false); + manager.register( + "owner", + "device-a", + &revoked_token, + "Device A", + 1, + revoked_tx, + revoked_close_tx, + ); + let (peer_tx, mut peer_rx) = mpsc::channel(4); + let (peer_close_tx, _peer_close_rx) = watch::channel(false); + manager.register( + "owner", + "device-c", + &peer_token, + "Device C", + 2, + peer_tx, + peer_close_tx, + ); + + // Simulate reset-password/delete-user tooling mutating the shared DB + // from outside the running relay process while both sockets are idle. + sqlx::query("DELETE FROM auth_tokens WHERE token = ?") + .bind(&revoked_token) + .execute(&db) + .await + .unwrap(); + assert_eq!( + revalidate_active_device_tokens_once(&db, &manager) + .await + .unwrap(), + 1 + ); + + revoked_close_rx.changed().await.unwrap(); + assert!(*revoked_close_rx.borrow()); + assert!(manager.conn_mapping(1).is_none()); + assert_eq!( + manager.conn_mapping(2), + Some(("owner".into(), "device-c".into())) + ); + let presence: serde_json::Value = + serde_json::from_str(&peer_rx.recv().await.unwrap().text).unwrap(); + assert_eq!(presence["devices"].as_array().unwrap().len(), 1); + assert_eq!(presence["devices"][0]["device_id"], "device-c"); + let rows = DeviceRow::list_by_user(&db, "owner").await.unwrap(); + assert_eq!( + rows.iter() + .find(|row| row.device_id == "device-a") + .unwrap() + .online, + 0 + ); + assert_eq!( + rows.iter() + .find(|row| row.device_id == "device-c") + .unwrap() + .online, + 1 + ); + } + + #[tokio::test] + async fn token_expiring_during_durable_projection_never_receives_auth_ok() { + let db = connect(":memory:").await.unwrap(); + UserRow::create(&db, "owner", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + DeviceRow::upsert(&db, "device-a", "owner", "Device A", None) + .await + .unwrap(); + let token = AuthToken::create(&db, "owner", "device-a") + .await + .unwrap() + .token; + let manager = DeviceManager::new(); + let (tx, mut rx) = mpsc::channel(4); + let (close_tx, mut close_rx) = watch::channel(false); + manager.register_pending("owner", "device-a", &token, "Device A", 1, tx, close_tx); + + let _projection_guard = manager.lock_presence_projection().await; + assert!( + registered_device_token_is_current(&db, &token, "owner", "device-a") + .await + .unwrap() + ); + DeviceRow::upsert(&db, "device-a", "owner", "Device A", None) + .await + .unwrap(); + DeviceRow::set_online(&db, "owner", "device-a", true) + .await + .unwrap(); + sqlx::query("UPDATE auth_tokens SET expires_at = ? WHERE token = ?") + .bind(chrono::Utc::now().timestamp()) + .bind(&token) + .execute(&db) + .await + .unwrap(); + + assert!(!complete_pending_device_activation_if_authorized( + &db, &manager, &token, "owner", "device-a", 1, + ) + .await + .unwrap()); + close_rx.changed().await.unwrap(); + assert!(*close_rx.borrow()); + assert!(rx.try_recv().is_err(), "AuthOk must not be queued"); + assert!(manager.conn_mapping(1).is_none()); + let rows = DeviceRow::list_by_user(&db, "owner").await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].online, 0); + } + + #[tokio::test] + async fn stale_auth_connect_cannot_recreate_a_deleted_device() { + let db = connect(":memory:").await.unwrap(); + UserRow::create(&db, "owner", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + DeviceRow::upsert(&db, "device-a", "owner", "Device A", None) + .await + .unwrap(); + let token = AuthToken::create(&db, "owner", "device-a") + .await + .unwrap() + .token; + assert!(AuthToken::find(&db, &token).await.unwrap().is_some()); + + assert!(DeviceRow::delete_for_user(&db, "owner", "device-a") + .await + .unwrap()); + let manager = DeviceManager::new(); + let (tx, _rx) = mpsc::channel(4); + let (close_tx, mut close_rx) = watch::channel(false); + manager.register_pending("owner", "device-a", &token, "Device A", 1, tx, close_tx); + + assert!(!activate_pending_device_if_authorized( + &db, &manager, &token, "owner", "device-a", "Device A", 1, + ) + .await + .unwrap()); + close_rx.changed().await.unwrap(); + assert!(*close_rx.borrow()); + assert!(DeviceRow::list_by_user(&db, "owner") + .await + .unwrap() + .is_empty()); + } + + #[tokio::test] + async fn pending_device_becomes_routable_only_after_durable_projection() { + let db = connect(":memory:").await.unwrap(); + UserRow::create(&db, "owner", "alice", "s", "ks", "{}", "hash", "wmk") + .await + .unwrap(); + DeviceRow::upsert(&db, "device-a", "owner", "Device A", None) + .await + .unwrap(); + let token = AuthToken::create(&db, "owner", "device-a") + .await + .unwrap() + .token; + let manager = DeviceManager::new(); + let (tx, mut rx) = mpsc::channel(4); + let (close_tx, _close_rx) = watch::channel(false); + manager.register_pending("owner", "device-a", &token, "Device A", 1, tx, close_tx); + + assert!(manager.online_devices("owner").is_empty()); + assert!(!manager.route_message("owner", "device-a", "opaque")); + assert!(activate_pending_device_if_authorized( + &db, &manager, &token, "owner", "device-a", "Device A", 1, + ) + .await + .unwrap()); + + assert_eq!( + manager.conn_mapping(1), + Some(("owner".into(), "device-a".into())) + ); + assert_eq!(manager.online_devices("owner").len(), 1); + let rows = DeviceRow::list_by_user(&db, "owner").await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].online, 1); + + manager.broadcast_current_presence("owner", |devices| { + serde_json::to_string(&OutboundProtocol::DevicePresence { + devices: super::build_presence(devices), + }) + .ok() + }); + let auth_ok: serde_json::Value = + serde_json::from_str(&rx.recv().await.unwrap().text).unwrap(); + let presence: serde_json::Value = + serde_json::from_str(&rx.recv().await.unwrap().text).unwrap(); + assert_eq!(auth_ok["type"], "auth_ok"); + assert_eq!(presence["type"], "device_presence"); + } } diff --git a/src/crates/services/services-integrations/Cargo.toml b/src/crates/services/services-integrations/Cargo.toml index 8328a947fb..782027a281 100644 --- a/src/crates/services/services-integrations/Cargo.toml +++ b/src/crates/services/services-integrations/Cargo.toml @@ -238,6 +238,7 @@ ssh_config = ["dep:ssh_config"] [dev-dependencies] tempfile = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } [[test]] name = "debug_log_owner_contracts" diff --git a/src/crates/services/services-integrations/src/remote_connect.rs b/src/crates/services/services-integrations/src/remote_connect.rs index 79d467905a..4238749cad 100644 --- a/src/crates/services/services-integrations/src/remote_connect.rs +++ b/src/crates/services/services-integrations/src/remote_connect.rs @@ -50,11 +50,12 @@ pub use ngrok::{ cleanup_all_ngrok, detect_running_ngrok, is_ngrok_available, start_ngrok_tunnel, NgrokTunnel, }; pub use page_upload::{ - delete_page_version_on_relay, deploy_page_version_on_relay, join_relay_url, - list_page_versions_from_relay, list_pages_from_relay, publish_page_content_on_relay, - publish_page_to_relay, save_page_version_from_inline_files, save_page_version_to_relay, - unpublish_page_from_relay, update_page_on_relay, PageContentPublishResult, PageInfo, - PagePublishResult, PageSaveVersionResult, PageVersionInfo, + create_page_open_link_on_relay, delete_page_from_relay, delete_page_version_on_relay, + deploy_page_version_on_relay, join_relay_url, list_page_versions_from_relay, + list_pages_from_relay, publish_page_content_on_relay, publish_page_to_relay, + save_page_version_from_inline_files, save_page_version_to_relay, unpublish_page_from_relay, + update_page_on_relay, PageContentPublishResult, PageInfo, PageOpenLink, PagePublishResult, + PageSaveVersionResult, PageVersionInfo, }; pub use pairing::{PairingChallenge, PairingProtocol, PairingResponse, PairingState, QrPayload}; pub use qr_generator::QrGenerator; diff --git a/src/crates/services/services-integrations/src/remote_connect/bot/mod.rs b/src/crates/services/services-integrations/src/remote_connect/bot/mod.rs index a1af531d10..15a9d0a7da 100644 --- a/src/crates/services/services-integrations/src/remote_connect/bot/mod.rs +++ b/src/crates/services/services-integrations/src/remote_connect/bot/mod.rs @@ -12,6 +12,7 @@ pub mod telegram; pub mod weixin; use serde::{Deserialize, Serialize}; +use std::sync::{Mutex as StdMutex, OnceLock}; pub use command::{parse_command, BotCommand}; pub use feishu::{FeishuBotApi, FeishuConfig}; @@ -486,6 +487,14 @@ pub fn auto_push_failed_message(language: BotLanguage, file_name: &str, err: &st const REMOTE_CONNECT_PERSISTENCE_FILENAME: &str = "remote_connect_persistence.json"; const LEGACY_BOT_PERSISTENCE_FILENAME: &str = "bot_connections.json"; +static BOT_PERSISTENCE_LOCK: OnceLock> = OnceLock::new(); + +fn bot_persistence_lock() -> std::sync::MutexGuard<'static, ()> { + BOT_PERSISTENCE_LOCK + .get_or_init(|| StdMutex::new(())) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} fn bot_persistence_path() -> Option { dirs::home_dir().map(|home| { @@ -494,17 +503,31 @@ fn bot_persistence_path() -> Option { }) } +fn bot_persistence_backup_path(path: &std::path::Path) -> std::path::PathBuf { + let file_name = path + .file_name() + .map(|name| name.to_string_lossy().to_string()) + .unwrap_or_else(|| REMOTE_CONNECT_PERSISTENCE_FILENAME.to_string()); + path.with_file_name(format!("{file_name}.bak")) +} + fn legacy_bot_persistence_path() -> Option { dirs::home_dir().map(|home| home.join(".bitfun").join(LEGACY_BOT_PERSISTENCE_FILENAME)) } -pub fn load_bot_persistence() -> BotPersistenceData { +fn load_bot_persistence_unlocked() -> BotPersistenceData { let Some(path) = bot_persistence_path() else { return BotPersistenceData::default(); }; match std::fs::read_to_string(&path) { Ok(data) => serde_json::from_str(&data).unwrap_or_default(), Err(_) => { + // A backup without the canonical file means the process stopped + // during the Windows replace dance. Fail closed instead of + // restoring the legacy file or a pre-clear account context. + if bot_persistence_backup_path(&path).exists() { + return BotPersistenceData::default(); + } let Some(legacy_path) = legacy_bot_persistence_path() else { return BotPersistenceData::default(); }; @@ -516,23 +539,128 @@ pub fn load_bot_persistence() -> BotPersistenceData { } } -pub fn save_bot_persistence(data: &BotPersistenceData) { +fn write_bot_persistence_atomic(path: &std::path::Path, json: &[u8]) -> std::io::Result<()> { + use std::io::Write; + + let parent = path.parent().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "bot persistence path has no parent directory", + ) + })?; + std::fs::create_dir_all(parent)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700))?; + } + + let file_name = path + .file_name() + .map(|name| name.to_string_lossy().to_string()) + .unwrap_or_else(|| REMOTE_CONNECT_PERSISTENCE_FILENAME.to_string()); + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let temp_path = parent.join(format!(".{file_name}.{}.{}.tmp", std::process::id(), nonce)); + let backup_path = bot_persistence_backup_path(path); + + let write_result = (|| -> std::io::Result<()> { + let mut options = std::fs::OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + let mut temp = options.open(&temp_path)?; + temp.write_all(json)?; + temp.sync_all()?; + drop(temp); + + // POSIX rename replaces atomically. On Windows std::fs::rename does + // not replace an existing destination, so retain the old file as a + // rollback backup until the new file is installed. + if std::fs::rename(&temp_path, path).is_ok() { + let _ = std::fs::remove_file(&backup_path); + return Ok(()); + } + + if !path.exists() { + return std::fs::rename(&temp_path, path); + } + if backup_path.exists() { + std::fs::remove_file(&backup_path)?; + } + std::fs::rename(path, &backup_path)?; + match std::fs::rename(&temp_path, path) { + Ok(()) => { + let _ = std::fs::remove_file(&backup_path); + Ok(()) + } + Err(install_error) => { + let _ = std::fs::rename(&backup_path, path); + Err(install_error) + } + } + })(); + + if write_result.is_err() { + let _ = std::fs::remove_file(&temp_path); + } + write_result +} + +fn save_bot_persistence_unlocked(data: &BotPersistenceData) { let Some(path) = bot_persistence_path() else { return; }; - if let Some(parent) = path.parent() { - let _ = std::fs::create_dir_all(parent); - } if let Ok(json) = serde_json::to_string_pretty(data) { - if let Err(e) = std::fs::write(&path, json) { + if let Err(e) = write_bot_persistence_atomic(&path, json.as_bytes()) { log::error!("Failed to save bot persistence: {e}"); } } } +pub fn load_bot_persistence() -> BotPersistenceData { + let _persistence = bot_persistence_lock(); + load_bot_persistence_unlocked() +} + +pub fn save_bot_persistence(data: &BotPersistenceData) { + let _persistence = bot_persistence_lock(); + save_bot_persistence_unlocked(data); +} + +/// Apply one read-modify-write transaction to the shared bot persistence +/// document. Individual platform loops run concurrently, so separate +/// `load`/`save` calls can otherwise overwrite another bot's update. +pub fn update_bot_persistence(update: impl FnOnce(&mut BotPersistenceData) -> R) -> R { + let _persistence = bot_persistence_lock(); + let mut data = load_bot_persistence_unlocked(); + let result = update(&mut data); + save_bot_persistence_unlocked(&data); + result +} + +/// Remove persisted account-device context for every bot connection. This is +/// called as part of the desktop account epoch transition even when a bot is +/// between stop/restore and therefore has no live chat-state map to clear. +pub fn clear_persisted_bot_account_contexts() { + update_bot_persistence(|data| { + for connection in &mut data.connections { + connection.chat_state.clear_delegated_identity(); + } + }); +} + #[cfg(test)] mod tests { - use super::{collect_auto_push_files, extract_downloadable_file_paths, resolve_workspace_path}; + use super::{ + bot_persistence_backup_path, collect_auto_push_files, extract_downloadable_file_paths, + resolve_workspace_path, write_bot_persistence_atomic, + }; fn make_temp_workspace() -> (std::path::PathBuf, std::path::PathBuf, std::path::PathBuf) { let base = std::env::temp_dir().join(format!( @@ -547,6 +675,34 @@ mod tests { (base, workspace, report) } + #[test] + fn bot_persistence_replaces_existing_file_via_same_directory_temp() { + let root = tempfile::tempdir().expect("temporary persistence directory"); + let path = root.path().join("remote_connect_persistence.json"); + write_bot_persistence_atomic(&path, br#"{"version":1}"#) + .expect("initial persistence write"); + write_bot_persistence_atomic(&path, br#"{"version":2}"#) + .expect("replacement persistence write"); + + assert_eq!( + std::fs::read_to_string(&path).expect("read installed persistence"), + r#"{"version":2}"# + ); + assert!(!bot_persistence_backup_path(&path).exists()); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + assert_eq!( + std::fs::metadata(&path) + .expect("persistence metadata") + .permissions() + .mode() + & 0o777, + 0o600 + ); + } + } + #[test] fn resolves_relative_paths_within_workspace_root() { let (base, workspace, report) = make_temp_workspace(); diff --git a/src/crates/services/services-integrations/src/remote_connect/bot/state.rs b/src/crates/services/services-integrations/src/remote_connect/bot/state.rs index 0ac61463ea..8e98df3351 100644 --- a/src/crates/services/services-integrations/src/remote_connect/bot/state.rs +++ b/src/crates/services/services-integrations/src/remote_connect/bot/state.rs @@ -168,6 +168,13 @@ pub struct BotChatState { /// Cleared by `/devices` → pick "local" or selecting an offline device. #[serde(skip, default)] pub active_remote_device: Option, + /// Records that the serializable workspace/session fields currently belong + /// to an account-routed device. The device target itself and delegated + /// credentials intentionally stay in memory only, but this marker must be + /// persisted so a restart cannot reinterpret a remote path/session as a + /// local desktop context. + #[serde(default, skip_serializing_if = "is_false")] + pub account_remote_context: bool, } /// A remote device the bot has switched to. All subsequent bot commands @@ -197,6 +204,7 @@ impl BotChatState { delegated_token: None, delegated_master_key: None, active_remote_device: None, + account_remote_context: false, } } @@ -252,6 +260,71 @@ impl BotChatState { pub fn has_delegated_identity(&self) -> bool { self.delegated_token.is_some() && self.delegated_master_key.is_some() } + + /// Switch command routing to an account device and fence every workspace + /// or session selection made for the previous target. + pub fn select_remote_device(&mut self, target: RemoteDeviceTarget) { + self.clear_device_scoped_context(); + self.active_remote_device = Some(target); + self.account_remote_context = true; + } + + /// Return command routing to this desktop without leaking a remote + /// workspace/session into local execution. + pub fn select_local_device(&mut self) { + let was_remote = self.active_remote_device.take().is_some() || self.account_remote_context; + self.account_remote_context = false; + if was_remote { + self.clear_device_scoped_context(); + } + } + + /// Sanitize state loaded from disk. Delegated authority is never restored, + /// and a persisted remote-context marker makes its workspace/session + /// selections invalid on a fresh process. + pub fn prepare_for_restore(&mut self) { + if self.account_remote_context { + self.clear_delegated_identity(); + } + } + + /// Remove every account-bound value when the desktop account changes. + /// A bot chat can outlive many desktop login sessions, so retaining these + /// fields would either keep controlling the previous account or replay an + /// old device selection with the replacement account's credentials. + pub fn clear_delegated_identity(&mut self) { + self.relay_url = None; + self.delegated_token = None; + self.delegated_master_key = None; + + let was_remote = self.active_remote_device.take().is_some() || self.account_remote_context; + self.account_remote_context = false; + let was_selecting_device = matches!( + self.pending_action, + Some(PendingAction::SelectDevice { .. }) + ); + if was_remote { + // Workspace/session selections made while targeting a remote + // device are owned by that account and must not fall through to + // local execution after the remote target is cleared. + self.clear_device_scoped_context(); + } + if was_remote || was_selecting_device { + self.clear_pending(); + self.last_menu_commands.clear(); + } + } + + fn clear_device_scoped_context(&mut self) { + self.current_workspace = None; + self.current_assistant = None; + self.current_assistant_name = None; + self.current_session_id = None; + } +} + +fn is_false(value: &bool) -> bool { + !*value } fn now_secs() -> i64 { @@ -406,6 +479,95 @@ mod tests { assert_eq!(state.pending_invalid_count, 0); } + #[test] + fn clearing_delegated_identity_drops_remote_account_state() { + let mut state = BotChatState::new("chat".into()); + state.relay_url = Some("https://relay-a.example".into()); + state.set_delegated_identity("token-a".into(), vec![7; 32]); + state.select_remote_device(RemoteDeviceTarget { + device_id: "device-a".into(), + device_name: "Device A".into(), + }); + state.current_workspace = Some(BotWorkspaceRef::local("/remote/a")); + state.current_assistant = Some("/remote/a".into()); + state.current_assistant_name = Some("Remote A".into()); + state.current_session_id = Some("session-a".into()); + state.set_pending(PendingAction::SelectSession { + options: vec![], + page: 0, + has_more: false, + }); + state.last_menu_commands = vec!["/sessions".into()]; + + state.clear_delegated_identity(); + + assert!(state.relay_url.is_none()); + assert!(!state.has_delegated_identity()); + assert!(state.active_remote_device.is_none()); + assert!(state.current_workspace.is_none()); + assert!(state.current_assistant.is_none()); + assert!(state.current_assistant_name.is_none()); + assert!(state.current_session_id.is_none()); + assert!(state.pending_action.is_none()); + assert!(state.last_menu_commands.is_empty()); + } + + #[test] + fn persisted_remote_context_is_sanitized_during_restore() { + let mut state = BotChatState::new("chat".into()); + state.paired = true; + state.select_remote_device(RemoteDeviceTarget { + device_id: "device-a".into(), + device_name: "Device A".into(), + }); + state.current_workspace = Some(BotWorkspaceRef::local("/remote/a")); + state.current_assistant = Some("/remote/a".into()); + state.current_assistant_name = Some("Remote A".into()); + state.current_session_id = Some("session-a".into()); + + let encoded = serde_json::to_string(&state).expect("state should serialize"); + assert!(encoded.contains("account_remote_context")); + assert!(!encoded.contains("device-a")); + + let mut restored: BotChatState = + serde_json::from_str(&encoded).expect("state should deserialize"); + assert!(restored.active_remote_device.is_none()); + assert!(restored.account_remote_context); + + restored.prepare_for_restore(); + + assert!(!restored.account_remote_context); + assert!(restored.current_workspace.is_none()); + assert!(restored.current_assistant.is_none()); + assert!(restored.current_assistant_name.is_none()); + assert!(restored.current_session_id.is_none()); + + let reencoded = serde_json::to_string(&restored).expect("restored state should serialize"); + let recovered_again: BotChatState = + serde_json::from_str(&reencoded).expect("clean state should deserialize"); + assert!(!recovered_again.account_remote_context); + assert!(recovered_again.current_workspace.is_none()); + assert!(recovered_again.current_session_id.is_none()); + } + + #[test] + fn switching_back_to_local_drops_remote_context() { + let mut state = BotChatState::new("chat".into()); + state.select_remote_device(RemoteDeviceTarget { + device_id: "device-a".into(), + device_name: "Device A".into(), + }); + state.current_workspace = Some(BotWorkspaceRef::local("/remote/a")); + state.current_session_id = Some("session-a".into()); + + state.select_local_device(); + + assert!(state.active_remote_device.is_none()); + assert!(!state.account_remote_context); + assert!(state.current_workspace.is_none()); + assert!(state.current_session_id.is_none()); + } + #[test] fn current_workspace_deserializes_legacy_path_string() { let state: BotChatState = serde_json::from_value(serde_json::json!({ diff --git a/src/crates/services/services-integrations/src/remote_connect/page_upload.rs b/src/crates/services/services-integrations/src/remote_connect/page_upload.rs index cdf4f0fa04..a7ae5001fa 100644 --- a/src/crates/services/services-integrations/src/remote_connect/page_upload.rs +++ b/src/crates/services/services-integrations/src/remote_connect/page_upload.rs @@ -1,7 +1,7 @@ //! BitFun Page incremental upload client (Save Version → Deploy). use anyhow::{anyhow, Result}; -use log::info; +use log::{info, warn}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; use std::path::Path; @@ -9,6 +9,7 @@ use std::path::Path; const MAX_UPLOAD_BATCH_BASE64_BYTES: usize = 256 * 1024; const MAX_PAGE_BYTES: u64 = 100 * 1024 * 1024; const MAX_FILE_BYTES: u64 = 10 * 1024 * 1024; +const MAX_PAGE_FILES: usize = 4096; #[derive(Debug, Clone, PartialEq, Eq, Serialize)] struct PageUploadManifestEntry { @@ -24,9 +25,67 @@ struct CollectedPageFile { hash: String, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum UploadSessionMode { + ManifestBound, + LegacyRelay, +} + +fn validate_upload_session_echo( + response: &serde_json::Value, + expected_upload_id: &str, +) -> Result { + match response + .get("upload_id") + .and_then(serde_json::Value::as_str) + { + Some(actual) if actual == expected_upload_id => Ok(UploadSessionMode::ManifestBound), + Some(_) => Err(anyhow!("check-files returned a mismatched upload session")), + // Relays predating manifest-bound upload sessions ignore the request's + // unknown `upload_id` field and omit it from the response. Their upload + // and freeze DTOs likewise ignore the extra field, so retaining the + // legacy flow preserves rolling-upgrade compatibility. + None => Ok(UploadSessionMode::LegacyRelay), + } +} + +fn validate_generation_intent_echo( + response: &serde_json::Value, + expected_generation: Option<&str>, + create: bool, +) -> Result { + match (response.get("expected_generation"), response.get("create")) { + // Relays with upload-id sessions but predating the Page generation + // protocol omit both fields. Their DTOs ignore the extra request + // fields throughout upload/freeze, so this remains a valid rolling + // upgrade path (without claiming generation protection). + (None, None) => Ok(false), + (Some(actual_generation), Some(actual_create)) => { + let expected_generation = expected_generation + .map_or(serde_json::Value::Null, |generation| { + serde_json::Value::String(generation.to_string()) + }); + if actual_generation == &expected_generation && actual_create.as_bool() == Some(create) + { + Ok(true) + } else { + Err(anyhow!( + "check-files returned a mismatched Page generation intent" + )) + } + } + // A partial echo is neither an old contract nor a trustworthy new one. + _ => Err(anyhow!( + "check-files returned an incomplete Page generation intent" + )), + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PageInfo { pub slug: String, + #[serde(default)] + pub generation: String, pub visibility: String, pub title: String, pub file_count: i64, @@ -42,6 +101,8 @@ pub struct PageInfo { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PageVersionInfo { + #[serde(default)] + pub generation: String, pub version_id: String, pub title: String, pub file_count: i64, @@ -53,9 +114,22 @@ pub struct PageVersionInfo { pub preview_url_path: String, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PageOpenLink { + pub open_url: String, + pub expires_in_seconds: u64, +} + +#[derive(Debug, Deserialize)] +struct PageOpenLinkRelayResponse { + open_url_path: String, + expires_in_seconds: u64, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PageSaveVersionResult { pub slug: String, + pub generation: String, pub visibility: String, pub title: String, pub version_id: String, @@ -70,6 +144,7 @@ pub struct PageSaveVersionResult { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PageContentPublishResult { pub slug: String, + pub generation: String, pub visibility: String, pub title: String, pub version_id: String, @@ -163,6 +238,7 @@ pub async fn publish_page_content_on_relay( let preview_url = join_relay_url(relay_url, &saved.preview_url_path); return Ok(PageContentPublishResult { slug: saved.slug, + generation: saved.generation, visibility: saved.visibility, title: saved.title, version_id: saved.version_id, @@ -178,20 +254,27 @@ pub async fn publish_page_content_on_relay( }); } - let info = deploy_page_version_on_relay(relay_url, token, &saved.slug, &saved.version_id) - .await - .map_err(|e| { - anyhow!( - "deploy failed: {e}. Version {} was saved on the relay but is not live; \ + let info = deploy_page_version_on_relay( + relay_url, + token, + &saved.slug, + &saved.version_id, + &saved.generation, + ) + .await + .map_err(|e| { + anyhow!( + "deploy failed: {e}. Version {} was saved on the relay but is not live; \ retry deploying version {} instead of re-publishing from scratch", - saved.version_id, - saved.version_id - ) - })?; + saved.version_id, + saved.version_id + ) + })?; let preview_url = join_relay_url(relay_url, &saved.preview_url_path); let url = join_relay_url(relay_url, &info.url_path); Ok(PageContentPublishResult { slug: saved.slug, + generation: info.generation, visibility: info.visibility, title: info.title, version_id: saved.version_id.clone(), @@ -270,12 +353,29 @@ async fn save_page_version_from_collected_files( size: f.content.len() as u64, }) .collect(); + let upload_id = uuid::Uuid::new_v4().simple().to_string(); + let existing_page = list_pages_from_relay(relay_url, token) + .await? + .into_iter() + .find(|page| page.slug == slug); + let create = existing_page.is_none(); + let expected_generation = existing_page + .as_ref() + .map(|page| page.generation.as_str()) + .filter(|generation| !generation.is_empty()) + .map(str::to_string); let check_url = format!("{relay_base}/api/pages/check-files"); let check_resp = client .post(&check_url) .header("Authorization", &auth) - .json(&serde_json::json!({ "slug": slug, "files": manifest })) + .json(&serde_json::json!({ + "slug": slug, + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, + "files": manifest, + })) .timeout(std::time::Duration::from_secs(30)) .send() .await @@ -289,6 +389,17 @@ async fn save_page_version_from_collected_files( .json() .await .map_err(|e| anyhow!("parse check-files response: {e}"))?; + let upload_mode = validate_upload_session_echo(&check_body, &upload_id)?; + if upload_mode == UploadSessionMode::LegacyRelay { + warn!( + "Page relay does not support manifest-bound upload sessions; using the legacy upload flow" + ); + } else if !validate_generation_intent_echo(&check_body, expected_generation.as_deref(), create)? + { + warn!( + "Page relay does not support generation-bound uploads; using its legacy generation flow" + ); + } let needed: Vec = check_body["needed"] .as_array() .map(|items| { @@ -301,7 +412,17 @@ async fn save_page_version_from_collected_files( if !needed.is_empty() { upload_needed_page_files( - &client, relay_base, &auth, slug, &title, visibility, &all_files, &needed, + &client, + relay_base, + &auth, + slug, + &upload_id, + expected_generation.as_deref(), + create, + &title, + visibility, + &all_files, + &needed, ) .await?; } else { @@ -311,6 +432,9 @@ async fn save_page_version_from_collected_files( &format!("{relay_base}/api/pages/upload-files"), &auth, slug, + &upload_id, + expected_generation.as_deref(), + create, &title, visibility, &HashMap::new(), @@ -325,6 +449,9 @@ async fn save_page_version_from_collected_files( .post(&freeze_url) .header("Authorization", &auth) .json(&serde_json::json!({ + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, "title": title, "note": note.unwrap_or(""), })) @@ -344,6 +471,7 @@ async fn save_page_version_from_collected_files( Ok(PageSaveVersionResult { slug: slug.to_string(), + generation: version.generation, visibility: visibility.to_string(), title, version_id: version.version_id, @@ -392,13 +520,15 @@ pub async fn list_page_versions_from_relay( relay_url: &str, token: &str, slug: &str, + expected_generation: &str, ) -> Result> { validate_slug(slug)?; let client = reqwest::Client::new(); let url = format!( - "{}/api/pages/{}/versions", + "{}/api/pages/{}/versions?expected_generation={}", relay_url.trim_end_matches('/'), - slug + slug, + expected_generation ); let resp = client .get(&url) @@ -418,11 +548,53 @@ pub async fn list_page_versions_from_relay( .map_err(|e| anyhow!("parse list versions: {e}"))?) } +/// Create a short-lived, one-time browser handoff URL for a production Page or +/// a specific immutable preview. The account bearer token is sent only in the +/// authenticated request header and is never embedded in the returned URL. +pub async fn create_page_open_link_on_relay( + relay_url: &str, + token: &str, + slug: &str, + version_id: Option<&str>, + expected_generation: &str, +) -> Result { + validate_slug(slug)?; + let client = reqwest::Client::new(); + let url = format!("{}/api/pages/{}", relay_url.trim_end_matches('/'), slug); + let resp = client + .post(&url) + .header("Authorization", format!("Bearer {token}")) + .json(&serde_json::json!({ + "version_id": version_id, + "expected_generation": expected_generation, + })) + .timeout(std::time::Duration::from_secs(15)) + .send() + .await + .map_err(|e| anyhow!("create page open link failed: {e}"))?; + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(anyhow!( + "create page open link failed: HTTP {status} — {body}" + )); + } + let result: PageOpenLinkRelayResponse = resp + .json() + .await + .map_err(|e| anyhow!("parse page open link response: {e}"))?; + Ok(PageOpenLink { + open_url: join_relay_url(relay_url, &result.open_url_path), + expires_in_seconds: result.expires_in_seconds, + }) +} + pub async fn deploy_page_version_on_relay( relay_url: &str, token: &str, slug: &str, version_id: &str, + expected_generation: &str, ) -> Result { validate_slug(slug)?; let client = reqwest::Client::new(); @@ -434,7 +606,10 @@ pub async fn deploy_page_version_on_relay( let resp = client .post(&url) .header("Authorization", format!("Bearer {token}")) - .json(&serde_json::json!({ "version_id": version_id })) + .json(&serde_json::json!({ + "version_id": version_id, + "expected_generation": expected_generation, + })) .timeout(std::time::Duration::from_secs(15)) .send() .await @@ -455,14 +630,16 @@ pub async fn delete_page_version_on_relay( token: &str, slug: &str, version_id: &str, + expected_generation: &str, ) -> Result<()> { validate_slug(slug)?; let client = reqwest::Client::new(); let url = format!( - "{}/api/pages/{}/versions/{}", + "{}/api/pages/{}/versions/{}?expected_generation={}", relay_url.trim_end_matches('/'), slug, - version_id + version_id, + expected_generation ); let resp = client .delete(&url) @@ -483,6 +660,7 @@ pub async fn update_page_on_relay( relay_url: &str, token: &str, slug: &str, + expected_generation: &str, visibility: Option<&str>, title: Option<&str>, ) -> Result { @@ -493,6 +671,10 @@ pub async fn update_page_on_relay( let client = reqwest::Client::new(); let url = format!("{}/api/pages/{}", relay_url.trim_end_matches('/'), slug); let mut body = serde_json::Map::new(); + body.insert( + "expected_generation".into(), + serde_json::json!(expected_generation), + ); if let Some(v) = visibility { body.insert("visibility".into(), serde_json::json!(v)); } @@ -518,12 +700,22 @@ pub async fn update_page_on_relay( .map_err(|e| anyhow!("parse update page: {e}"))?) } -pub async fn unpublish_page_from_relay(relay_url: &str, token: &str, slug: &str) -> Result<()> { +pub async fn unpublish_page_from_relay( + relay_url: &str, + token: &str, + slug: &str, + expected_generation: &str, +) -> Result<()> { validate_slug(slug)?; let client = reqwest::Client::new(); - let url = format!("{}/api/pages/{}", relay_url.trim_end_matches('/'), slug); + let url = format!( + "{}/api/pages/{}/unpublish?expected_generation={}", + relay_url.trim_end_matches('/'), + slug, + expected_generation + ); let resp = client - .delete(&url) + .post(&url) .header("Authorization", format!("Bearer {token}")) .timeout(std::time::Duration::from_secs(15)) .send() @@ -537,6 +729,35 @@ pub async fn unpublish_page_from_relay(relay_url: &str, token: &str, slug: &str) Ok(()) } +pub async fn delete_page_from_relay( + relay_url: &str, + token: &str, + slug: &str, + expected_generation: &str, +) -> Result<()> { + validate_slug(slug)?; + let client = reqwest::Client::new(); + let url = format!( + "{}/api/pages/{}?expected_generation={}", + relay_url.trim_end_matches('/'), + slug, + expected_generation + ); + let resp = client + .delete(&url) + .header("Authorization", format!("Bearer {token}")) + .timeout(std::time::Duration::from_secs(15)) + .send() + .await + .map_err(|e| anyhow!("delete page failed: {e}"))?; + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(anyhow!("delete page failed: HTTP {status} — {body}")); + } + Ok(()) +} + fn validate_slug(slug: &str) -> Result<()> { let bytes = slug.as_bytes(); if bytes.is_empty() || bytes.len() > 64 { @@ -571,19 +792,22 @@ fn validate_visibility(v: &str) -> Result<()> { } fn collect_page_files(base: &Path) -> Result> { - if !base.is_dir() { + let base_metadata = std::fs::symlink_metadata(base) + .map_err(|_| anyhow!("page directory does not exist: {}", base.display()))?; + if base_metadata.file_type().is_symlink() || !base_metadata.is_dir() { return Err(anyhow!("page directory does not exist: {}", base.display())); } - let has_index = base.join("index.html").exists(); - let has_worker = base.join("server").join("worker.js").exists(); + let has_index = is_regular_page_source_file(&base.join("index.html"))?; + let has_worker = is_regular_page_source_file(&base.join("server").join("worker.js"))?; if !has_index && !has_worker { return Err(anyhow!( "page directory must contain index.html and/or server/worker.js: {}", base.display() )); } + let canonical_base = std::fs::canonicalize(base)?; let mut all_files = Vec::new(); - collect_files_with_hash(base, base, &mut all_files)?; + collect_files_with_hash(&canonical_base, &canonical_base, &mut all_files)?; Ok(all_files) } @@ -593,6 +817,12 @@ fn collect_inline_page_files(files: &HashMap) -> Result MAX_PAGE_FILES { + return Err(anyhow!( + "page exceeds file count limit ({} > {MAX_PAGE_FILES})", + files.len() + )); + } let has_index = files.contains_key("index.html"); let has_worker = files.contains_key("server/worker.js"); @@ -640,18 +870,35 @@ fn collect_files_with_hash( for entry in std::fs::read_dir(dir)? { let entry = entry?; let path = entry.path(); - if path.is_dir() { - collect_files_with_hash(base, &path, out)?; - } else if path.is_file() { - let rel = path + let file_type = entry.file_type()?; + if file_type.is_symlink() { + return Err(anyhow!( + "symbolic links are not allowed in Page sources: {}", + path.display() + )); + } + let canonical_path = std::fs::canonicalize(&path)?; + if !canonical_path.starts_with(base) { + return Err(anyhow!( + "Page source entry resolves outside its directory: {}", + path.display() + )); + } + if file_type.is_dir() { + collect_files_with_hash(base, &canonical_path, out)?; + } else if file_type.is_file() { + if out.len() >= MAX_PAGE_FILES { + return Err(anyhow!("page exceeds file count limit ({MAX_PAGE_FILES})")); + } + let rel = canonical_path .strip_prefix(base) - .unwrap_or(&path) + .unwrap_or(&canonical_path) .to_string_lossy() .replace('\\', "/"); if rel.contains("..") { continue; } - let content = std::fs::read(&path)?; + let content = std::fs::read(&canonical_path)?; if content.len() as u64 > MAX_FILE_BYTES { return Err(anyhow!( "file exceeds size limit: {rel} ({} > {MAX_FILE_BYTES} bytes)", @@ -666,17 +913,37 @@ fn collect_files_with_hash( content, hash, }); + } else { + return Err(anyhow!( + "unsupported Page source entry type: {}", + path.display() + )); } } Ok(()) } +fn is_regular_page_source_file(path: &Path) -> Result { + match std::fs::symlink_metadata(path) { + Ok(metadata) if metadata.file_type().is_symlink() => Err(anyhow!( + "symbolic links are not allowed in Page sources: {}", + path.display() + )), + Ok(metadata) => Ok(metadata.is_file()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(error.into()), + } +} + #[allow(clippy::too_many_arguments)] async fn upload_needed_page_files( client: &reqwest::Client, relay_base: &str, auth: &str, slug: &str, + upload_id: &str, + expected_generation: Option<&str>, + create: bool, title: &str, visibility: &str, all_files: &[CollectedPageFile], @@ -710,6 +977,9 @@ async fn upload_needed_page_files( &url, auth, slug, + upload_id, + expected_generation, + create, title, visibility, ¤t_batch, @@ -730,6 +1000,9 @@ async fn upload_needed_page_files( &url, auth, slug, + upload_id, + expected_generation, + create, title, visibility, ¤t_batch, @@ -747,6 +1020,9 @@ async fn post_upload_batch( url: &str, auth: &str, slug: &str, + upload_id: &str, + expected_generation: Option<&str>, + create: bool, title: &str, visibility: &str, files: &HashMap, @@ -758,6 +1034,9 @@ async fn post_upload_batch( .header("Authorization", auth) .json(&serde_json::json!({ "slug": slug, + "upload_id": upload_id, + "expected_generation": expected_generation, + "create": create, "title": title, "visibility": visibility, "files": files, @@ -803,6 +1082,61 @@ mod tests { ); } + #[test] + fn upload_session_echo_accepts_legacy_relay_but_rejects_mismatch() { + let legacy = serde_json::json!({ "needed": ["index.html"] }); + assert_eq!( + validate_upload_session_echo(&legacy, "abc").unwrap(), + UploadSessionMode::LegacyRelay + ); + + let current = serde_json::json!({ "upload_id": "abc", "needed": [] }); + assert_eq!( + validate_upload_session_echo(¤t, "abc").unwrap(), + UploadSessionMode::ManifestBound + ); + + let mismatch = serde_json::json!({ "upload_id": "other", "needed": [] }); + assert!(validate_upload_session_echo(&mismatch, "abc").is_err()); + } + + #[test] + fn generation_intent_echo_accepts_old_relay_only_when_both_fields_are_absent() { + let old_relay = serde_json::json!({ "upload_id": "abc", "needed": [] }); + assert!(!validate_generation_intent_echo(&old_relay, Some("generation-a"), false).unwrap()); + + let current = serde_json::json!({ + "upload_id": "abc", + "expected_generation": "generation-a", + "create": false, + "needed": [], + }); + assert!(validate_generation_intent_echo(¤t, Some("generation-a"), false).unwrap()); + + let create = serde_json::json!({ + "upload_id": "abc", + "expected_generation": null, + "create": true, + "needed": [], + }); + assert!(validate_generation_intent_echo(&create, None, true).unwrap()); + + let partial = serde_json::json!({ + "upload_id": "abc", + "expected_generation": "generation-a", + "needed": [], + }); + assert!(validate_generation_intent_echo(&partial, Some("generation-a"), false).is_err()); + + let mismatch = serde_json::json!({ + "upload_id": "abc", + "expected_generation": "generation-b", + "create": false, + "needed": [], + }); + assert!(validate_generation_intent_echo(&mismatch, Some("generation-a"), false).is_err()); + } + #[test] fn collect_allows_worker_without_index() { let base = @@ -835,6 +1169,42 @@ mod tests { assert!(collect_inline_page_files(&empty_entry).is_err()); } + #[cfg(unix)] + #[test] + fn directory_source_rejects_symlinks_to_external_files() { + use std::os::unix::fs::symlink; + + let root = std::env::temp_dir().join(format!( + "bitfun-page-symlink-external-{}", + uuid::Uuid::new_v4() + )); + let page = root.join("page"); + std::fs::create_dir_all(&page).unwrap(); + std::fs::write(page.join("index.html"), b"").unwrap(); + std::fs::write(root.join("secret.txt"), b"secret").unwrap(); + symlink(root.join("secret.txt"), page.join("secret.txt")).unwrap(); + + let error = collect_page_files(&page).expect_err("external symlink must be rejected"); + assert!(error.to_string().contains("symbolic links are not allowed")); + let _ = std::fs::remove_dir_all(root); + } + + #[cfg(unix)] + #[test] + fn directory_source_rejects_symlink_loops() { + use std::os::unix::fs::symlink; + + let root = + std::env::temp_dir().join(format!("bitfun-page-symlink-loop-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + std::fs::write(root.join("index.html"), b"").unwrap(); + symlink(&root, root.join("loop")).unwrap(); + + let error = collect_page_files(&root).expect_err("symlink loop must be rejected"); + assert!(error.to_string().contains("symbolic links are not allowed")); + let _ = std::fs::remove_dir_all(root); + } + #[tokio::test] async fn publish_source_xor_is_enforced() { let err = publish_page_content_on_relay( diff --git a/src/crates/services/services-integrations/src/remote_connect/relay_client.rs b/src/crates/services/services-integrations/src/remote_connect/relay_client.rs index b269a2b73d..f5c1f2d2fb 100644 --- a/src/crates/services/services-integrations/src/remote_connect/relay_client.rs +++ b/src/crates/services/services-integrations/src/remote_connect/relay_client.rs @@ -35,6 +35,8 @@ pub fn ensure_rustls_crypto_provider() { type WsStream = tokio_tungstenite::WebSocketStream>; +const RELAY_DIAL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15); + /// Messages in the relay protocol (both directions). #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] @@ -566,38 +568,93 @@ async fn dial(ws_url: &str) -> Result { #[cfg(windows)] { - let request = ws_url - .into_client_request() - .map_err(|e| anyhow!("dial {ws_url}: build request failed: {e}"))?; - - // Wrap TLS connector construction in catch_unwind so that a panic - // (e.g. duplicate CryptoProvider install) is converted to an error - // instead of unwinding the tokio task and potentially crashing the - // process. - let connector = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - build_windows_rustls_connector() - })) - .map_err(|_| anyhow!("dial {ws_url}: TLS connector construction panicked"))??; - - let (stream, _) = tokio_tungstenite::connect_async_tls_with_config( - request, - Some(config), - false, - Some(connector), - ) + await_dial(ws_url, async move { + let request = ws_url + .into_client_request() + .map_err(|e| anyhow!("dial {ws_url}: build request failed: {e}"))?; + + // Wrap TLS connector construction in catch_unwind so that a panic + // (e.g. duplicate CryptoProvider install) is converted to an error + // instead of unwinding the tokio task and potentially crashing the + // process. + let connector = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + build_windows_rustls_connector() + })) + .map_err(|_| anyhow!("dial {ws_url}: TLS connector construction panicked"))??; + + let (stream, _) = tokio_tungstenite::connect_async_tls_with_config( + request, + Some(config), + false, + Some(connector), + ) + .await + .map_err(|e| anyhow!("dial {ws_url}: {e}"))?; + Ok(stream) + }) .await - .map_err(|e| anyhow!("dial {ws_url}: {e}"))?; - Ok(stream) } #[cfg(not(windows))] { // Non-Windows uses tokio-tungstenite's built-in rustls connector. // CryptoProvider must already be installed (see ensure_rustls_crypto_provider). - let (stream, _) = tokio_tungstenite::connect_async_with_config(ws_url, Some(config), false) - .await - .map_err(|e| anyhow!("dial {ws_url}: {e}"))?; - Ok(stream) + await_dial(ws_url, async move { + let (stream, _) = + tokio_tungstenite::connect_async_with_config(ws_url, Some(config), false) + .await + .map_err(|e| anyhow!("dial {ws_url}: {e}"))?; + Ok(stream) + }) + .await + } +} + +async fn await_dial(ws_url: &str, dial_future: F) -> Result +where + F: std::future::Future>, +{ + tokio::time::timeout(RELAY_DIAL_TIMEOUT, dial_future) + .await + .map_err(|_| { + anyhow!( + "dial {ws_url}: connection timed out after {} seconds", + RELAY_DIAL_TIMEOUT.as_secs() + ) + })? +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn dial_timeout_bounds_a_pending_connection_attempt() { + let result = await_dial( + "wss://relay.example.invalid/ws", + std::future::pending::>(), + ) + .await; + + let error = result.expect_err("pending dial must be bounded by the connection timeout"); + assert_eq!( + error.to_string(), + "dial wss://relay.example.invalid/ws: connection timed out after 15 seconds" + ); + } + + #[tokio::test] + async fn dial_timeout_preserves_connection_errors() { + let result = await_dial::<(), _>( + "wss://relay.example.invalid/ws", + std::future::ready(Err(anyhow!("dial failed before timeout"))), + ) + .await; + + assert_eq!( + result.expect_err("dial error must be returned").to_string(), + "dial failed before timeout" + ); } } diff --git a/src/crates/services/services-integrations/src/remote_connect/session_store.rs b/src/crates/services/services-integrations/src/remote_connect/session_store.rs index 46adb54ac5..dddcc72eee 100644 --- a/src/crates/services/services-integrations/src/remote_connect/session_store.rs +++ b/src/crates/services/services-integrations/src/remote_connect/session_store.rs @@ -12,6 +12,7 @@ //! payload `{ token, user_id, master_key_b64, relay_url }`. use std::path::PathBuf; +use std::sync::{OnceLock, RwLock}; use std::{fs::OpenOptions, io::Write}; use aes_gcm::aead::{Aead, KeyInit, OsRng}; @@ -24,6 +25,33 @@ use sha2::{Digest, Sha256}; const NONCE_SIZE: usize = 12; +fn session_store_directory_override() -> &'static RwLock> { + static OVERRIDE: OnceLock>> = OnceLock::new(); + OVERRIDE.get_or_init(|| RwLock::new(None)) +} + +/// Redirects session, local-key, and credential-hint files for integration +/// tests. Tests must use this instead of changing HOME or touching real login +/// state shared by Desktop and CLI. +pub fn set_session_store_directory_for_test(path: PathBuf) { + let mut override_path = session_store_directory_override() + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *override_path = Some(path); +} + +fn session_store_directory() -> Result { + let override_path = session_store_directory_override() + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(path) = override_path.as_ref() { + return Ok(path.clone()); + } + drop(override_path); + let home = dirs::home_dir().ok_or_else(|| anyhow!("cannot determine home directory"))?; + Ok(home.join(".bitfun")) +} + /// The on-disk JSON payload (plaintext before encryption). #[derive(Serialize, Deserialize)] struct SessionPayload { @@ -40,13 +68,11 @@ struct SessionPayload { /// Resolve the persistent session file path. fn session_file_path() -> Result { - let home = dirs::home_dir().ok_or_else(|| anyhow!("cannot determine home directory"))?; - Ok(home.join(".bitfun").join("account_session.enc")) + Ok(session_store_directory()?.join("account_session.enc")) } fn session_key_file_path() -> Result { - let home = dirs::home_dir().ok_or_else(|| anyhow!("cannot determine home directory"))?; - Ok(home.join(".bitfun").join("account_session.key")) + Ok(session_store_directory()?.join("account_session.key")) } /// Atomically replace a secret-bearing file and restrict it to the current @@ -408,8 +434,7 @@ pub fn clear_session() { // Shared by Desktop and CLI so login forms pre-fill the same values. fn credential_hint_path() -> Result { - let home = dirs::home_dir().ok_or_else(|| anyhow!("cannot determine home directory"))?; - Ok(home.join(".bitfun").join("account_hint.json")) + Ok(session_store_directory()?.join("account_hint.json")) } /// Non-secret login pre-fill (never stores password or master key). diff --git a/src/mobile-web/src/App.tsx b/src/mobile-web/src/App.tsx index fa7c90f370..1af17e1e81 100644 --- a/src/mobile-web/src/App.tsx +++ b/src/mobile-web/src/App.tsx @@ -7,6 +7,7 @@ import { ErrorBoundary } from './components/ErrorBoundary'; import { I18nProvider, useI18n } from './i18n'; import { RelayHttpClient } from './services/RelayHttpClient'; import { RemoteSessionManager } from './services/RemoteSessionManager'; +import { reconcileDelegatedAccountOwner } from './services/delegatedAccountOwner'; import { ThemeProvider } from './theme'; import { useConnectionHealth } from './hooks/useConnectionHealth'; import { useMobileStore } from './services/store'; @@ -40,6 +41,7 @@ const AppContent: React.FC = () => { const [chatAutoFocus, setChatAutoFocus] = useState(false); const connectionHealth = useMobileStore((state) => state.connectionHealth); const clientRef = useRef(null); + const delegatedOwnerUnlistenRef = useRef<(() => void) | null>(null); const sessionMgrRef = useRef(null); const [sessionMgr, setSessionMgr] = useState(null); @@ -109,7 +111,26 @@ const AppContent: React.FC = () => { const handlePaired = useCallback( (client: RelayHttpClient, sessionMgr: RemoteSessionManager) => { + delegatedOwnerUnlistenRef.current?.(); clientRef.current = client; + delegatedOwnerUnlistenRef.current = client.onDelegatedAccountOwnerChange((change) => { + if (clientRef.current !== client) return; + const ownerScopedStateWasReset = reconcileDelegatedAccountOwner(change); + if (!ownerScopedStateWasReset) return; + + // A detail page can retain local IDs in addition to Zustand state. + // Return to the session root before any stale completion can render + // data from the replacement account. + clearTimeout(timerRef.current); + setActiveSessionId(null); + setActiveSessionName('Session'); + setChatAutoFocus(false); + setPrevPage(null); + setNavDir(null); + pageStackRef.current = ['pairing', 'sessions']; + history.replaceState({ page: 'sessions' }, ''); + setPage('sessions'); + }, { emitCurrent: true }); sessionMgrRef.current = sessionMgr; setSessionMgr(sessionMgr); pageStackRef.current = ['pairing', 'sessions']; @@ -184,6 +205,9 @@ const AppContent: React.FC = () => { }, [navigateTo]); const handleDisconnect = useCallback(() => { + delegatedOwnerUnlistenRef.current?.(); + delegatedOwnerUnlistenRef.current = null; + clientRef.current?.resetConnectionIdentity(); clientRef.current = null; sessionMgrRef.current = null; setSessionMgr(null); @@ -199,6 +223,11 @@ const AppContent: React.FC = () => { setPage('pairing'); }, []); + useEffect(() => () => { + delegatedOwnerUnlistenRef.current?.(); + delegatedOwnerUnlistenRef.current = null; + }, []); + const isAnimating = navDir !== null; const currentPage: Page = page; diff --git a/src/mobile-web/src/hooks/useConnectionHealth.ts b/src/mobile-web/src/hooks/useConnectionHealth.ts index 1edac87844..154c870a54 100644 --- a/src/mobile-web/src/hooks/useConnectionHealth.ts +++ b/src/mobile-web/src/hooks/useConnectionHealth.ts @@ -1,9 +1,14 @@ import { useEffect, useRef } from 'react'; -import { RemoteSessionManager } from '../services/RemoteSessionManager'; +import { isDelegatedIdentityChangedError } from '../services/RelayHttpClient'; +import { + isRemoteControlTargetChangedError, + RemoteSessionManager, +} from '../services/RemoteSessionManager'; import { useMobileStore } from '../services/store'; const PING_INTERVAL = 15000; const PING_TIMEOUT = 10000; +const OWNERSHIP_RETRY_DELAY = 250; function pingWithTimeout(mgr: RemoteSessionManager, ms: number): Promise { let timeoutId: ReturnType | undefined; @@ -23,31 +28,57 @@ export function useConnectionHealth(sessionMgr: RemoteSessionManager | null) { useEffect(() => { let cancelled = false; + let loopGeneration = 0; if (!sessionMgr) { setConnectionHealth('unpaired'); return; } - setConnectionHealth('checking'); + const schedule = (generation: number, delay: number) => { + if (cancelled || generation !== loopGeneration) return; + if (timerRef.current) clearTimeout(timerRef.current); + timerRef.current = setTimeout(() => { + void loop(generation); + }, delay); + }; - const loop = async () => { + const loop = async (generation: number) => { + if (cancelled || generation !== loopGeneration) return; try { await pingWithTimeout(sessionMgr, PING_TIMEOUT); - if (!cancelled) setConnectionHealth('connected'); - } catch { - if (!cancelled) setConnectionHealth('unreachable'); + if (cancelled || generation !== loopGeneration) return; + setConnectionHealth('connected'); + schedule(generation, PING_INTERVAL); + } catch (error: unknown) { + if (cancelled || generation !== loopGeneration) return; + if ( + isRemoteControlTargetChangedError(error) + || isDelegatedIdentityChangedError(error) + ) { + setConnectionHealth('checking'); + schedule(generation, OWNERSHIP_RETRY_DELAY); + return; + } + setConnectionHealth('unreachable'); + schedule(generation, PING_INTERVAL); } + }; - if (!cancelled) { - timerRef.current = setTimeout(loop, PING_INTERVAL); - } + const restart = () => { + loopGeneration += 1; + if (timerRef.current) clearTimeout(timerRef.current); + setConnectionHealth('checking'); + void loop(loopGeneration); }; - loop(); + const unlisten = sessionMgr.onControlTargetChange(restart); + restart(); return () => { cancelled = true; + loopGeneration += 1; + unlisten(); if (timerRef.current) clearTimeout(timerRef.current); }; }, [sessionMgr, setConnectionHealth]); diff --git a/src/mobile-web/src/hooks/useControlTargetEpoch.ts b/src/mobile-web/src/hooks/useControlTargetEpoch.ts new file mode 100644 index 0000000000..b93368ae9f --- /dev/null +++ b/src/mobile-web/src/hooks/useControlTargetEpoch.ts @@ -0,0 +1,20 @@ +import { useCallback, useSyncExternalStore } from 'react'; +import type { RemoteSessionManager } from '../services/RemoteSessionManager'; + +/** + * React subscription for the transport-owned control-target generation. + * useSyncExternalStore also closes the render-to-subscribe missed-event window + * and gives StrictMode setup/cleanup deterministic semantics. + */ +export function useControlTargetEpoch(sessionMgr: RemoteSessionManager): number { + const subscribe = useCallback( + (onStoreChange: () => void) => sessionMgr.onControlTargetChange(onStoreChange), + [sessionMgr], + ); + const getSnapshot = useCallback( + () => sessionMgr.controlTargetEpoch, + [sessionMgr], + ); + + return useSyncExternalStore(subscribe, getSnapshot, getSnapshot); +} diff --git a/src/mobile-web/src/i18n/I18nProvider.tsx b/src/mobile-web/src/i18n/I18nProvider.tsx index 6bbb16a78c..25dd58c26b 100644 --- a/src/mobile-web/src/i18n/I18nProvider.tsx +++ b/src/mobile-web/src/i18n/I18nProvider.tsx @@ -17,6 +17,7 @@ interface I18nContextValue { toggleLanguage: () => void; t: (key: string, params?: TranslateParams) => string; formatDate: (date: Date | number, options?: Intl.DateTimeFormatOptions) => string; + formatRelativeTime: (date: Date | number) => string; } const STORAGE_KEY = 'bitfun-mobile-language'; @@ -99,8 +100,22 @@ export const I18nContext = createContext({ toggleLanguage: () => {}, t: (key) => key, formatDate: (date, options) => new Intl.DateTimeFormat(DEFAULT_LANGUAGE, options).format(date), + formatRelativeTime: (date) => formatRelativeTime(DEFAULT_LANGUAGE, date), }); +function formatRelativeTime(language: MobileLanguage, date: Date | number): string { + const timestamp = date instanceof Date ? date.getTime() : date; + const diffSeconds = (timestamp - Date.now()) / 1000; + const absoluteSeconds = Math.abs(diffSeconds); + const formatter = new Intl.RelativeTimeFormat(language, { numeric: 'auto' }); + + if (absoluteSeconds < 60) return formatter.format(Math.round(diffSeconds), 'second'); + if (absoluteSeconds < 3600) return formatter.format(Math.round(diffSeconds / 60), 'minute'); + if (absoluteSeconds < 86_400) return formatter.format(Math.round(diffSeconds / 3600), 'hour'); + if (absoluteSeconds < 2_592_000) return formatter.format(Math.round(diffSeconds / 86_400), 'day'); + return new Intl.DateTimeFormat(language, { dateStyle: 'medium', timeStyle: 'short' }).format(timestamp); +} + export const I18nProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => { const [language, setLanguageState] = useState(detectInitialLanguage); @@ -127,6 +142,7 @@ export const I18nProvider: React.FC<{ children: React.ReactNode }> = ({ children toggleLanguage, t: (key, params) => translate(language, key, params), formatDate: (date, options) => new Intl.DateTimeFormat(language, options).format(date), + formatRelativeTime: (date) => formatRelativeTime(language, date), }), [language, setLanguage, toggleLanguage]); return ( @@ -137,4 +153,3 @@ export const I18nProvider: React.FC<{ children: React.ReactNode }> = ({ children }; export type { MobileLanguage, TranslateParams }; - diff --git a/src/mobile-web/src/i18n/messages.ts b/src/mobile-web/src/i18n/messages.ts index d0a17dbf84..4488fd1bfb 100644 --- a/src/mobile-web/src/i18n/messages.ts +++ b/src/mobile-web/src/i18n/messages.ts @@ -176,6 +176,7 @@ export const messages: Record = { noDevices: 'No devices found', online: 'Online', offline: 'Offline', + lastSeen: 'Last seen {time}', current: 'Current', pairedDesktop: 'Paired', switchFailed: 'Failed to switch device', @@ -364,6 +365,7 @@ export const messages: Record = { noDevices: '未找到设备', online: '在线', offline: '离线', + lastSeen: '上次在线：{time}', current: '当前', pairedDesktop: '已配对', switchFailed: '切换设备失败', @@ -552,6 +554,7 @@ export const messages: Record = { noDevices: '未找到設備', online: '在線', offline: '離線', + lastSeen: '上次在線：{time}', current: '目前', pairedDesktop: '已配對', switchFailed: '切換設備失敗', diff --git a/src/mobile-web/src/pages/ChatPage.tsx b/src/mobile-web/src/pages/ChatPage.tsx index c6325b370e..2b1ab4c1d5 100644 --- a/src/mobile-web/src/pages/ChatPage.tsx +++ b/src/mobile-web/src/pages/ChatPage.tsx @@ -28,7 +28,10 @@ import typescript from 'react-syntax-highlighter/dist/esm/languages/prism/typesc import yaml from 'react-syntax-highlighter/dist/esm/languages/prism/yaml'; import { useI18n } from '../i18n'; import { messages } from '../i18n/messages'; +import { useControlTargetEpoch } from '../hooks/useControlTargetEpoch'; import { + isRemoteControlTargetChangedError, + RemoteControlTargetChangedError, RemoteSessionManager, SessionPoller, type PollResponse, @@ -42,6 +45,14 @@ import { import { useMobileStore } from '../services/store'; import { useTheme } from '../theme'; +function reportRemoteSessionError( + error: unknown, + setError: (message: string) => void, +): void { + if (isRemoteControlTargetChangedError(error)) return; + setError(error instanceof Error ? error.message : String(error)); +} + const SYNTAX_LANGUAGES = { bash, c, @@ -1653,6 +1664,7 @@ function renderActiveTurnItems( now: number, sessionMgr: RemoteSessionManager, setError: (e: string) => void, + isTargetCurrent: () => boolean, onAnswer: (toolId: string, answers: any) => Promise, onFileDownload?: (path: string, onProgress?: (downloaded: number, total: number) => void) => Promise, onGetFileInfo?: (path: string) => Promise<{ name: string; size: number; mimeType: string }>, @@ -1660,7 +1672,10 @@ function renderActiveTurnItems( const items = filterSubagentItems(rawItems); const askEntries = items.filter(item => isPendingAskUserQuestion(item.tool)); const onCancel = (toolId: string) => { - sessionMgr.cancelTool(toolId, 'User cancelled').catch(err => { setError(String(err)); }); + if (!isTargetCurrent()) return; + sessionMgr.cancelTool(toolId, 'User cancelled').catch((error) => { + reportRemoteSessionError(error, setError); + }); }; if (askEntries.length === 0) { @@ -2086,12 +2101,56 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, const fileInputRef = useRef(null); const inputBarRef = useRef(null); const pollerRef = useRef(null); - + const messagesRequestSeqRef = useRef(0); const [isLoadingMore, setIsLoadingMore] = useState(false); const [hasMore, setHasMore] = useState(true); const isLoadingMoreRef = useRef(false); const hasMoreRef = useRef(true); + const controlTargetEpoch = useControlTargetEpoch(sessionMgr); + const chatTargetOwnerRef = useRef({ + sessionMgr, + sessionId, + epoch: controlTargetEpoch, + active: true, + }); + if ( + chatTargetOwnerRef.current.sessionMgr !== sessionMgr + || chatTargetOwnerRef.current.sessionId !== sessionId + || chatTargetOwnerRef.current.epoch !== controlTargetEpoch + ) { + chatTargetOwnerRef.current = { + sessionMgr, + sessionId, + epoch: controlTargetEpoch, + active: true, + }; + } + + const captureChatTargetEpoch = useCallback((): number | null => { + const owner = chatTargetOwnerRef.current; + if ( + !owner.active + || owner.sessionMgr !== sessionMgr + || owner.sessionId !== sessionId + || owner.epoch !== sessionMgr.controlTargetEpoch + ) { + return null; + } + return owner.epoch; + }, [controlTargetEpoch, sessionId, sessionMgr]); + + const isChatTargetCurrent = useCallback((epoch: number | null): boolean => { + const owner = chatTargetOwnerRef.current; + return epoch !== null + && owner.active + && owner.sessionMgr === sessionMgr + && owner.sessionId === sessionId + && owner.epoch === epoch + && sessionMgr.controlTargetEpoch === epoch; + }, [controlTargetEpoch, sessionId, sessionMgr]); + const modelSelectionInitializedRef = useRef(false); + const modelCatalogRequestSeqRef = useRef(0); const messagesEndRef = useRef(null); const messagesContainerRef = useRef(null); const [expandedMsgIds, setExpandedMsgIds] = useState>(new Set()); @@ -2103,23 +2162,84 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, const msgLongPressTimerRef = useRef>(); const msgLongPressPosRef = useRef({ x: 0, y: 0 }); const msgToastTimerRef = useRef>(); + const committedChatTargetRef = useRef({ sessionMgr, sessionId, epoch: controlTargetEpoch }); + + useLayoutEffect(() => { + const previous = committedChatTargetRef.current; + const targetChanged = previous.sessionMgr !== sessionMgr + || previous.sessionId !== sessionId + || previous.epoch !== controlTargetEpoch; + const owner = chatTargetOwnerRef.current; + owner.active = owner.sessionMgr === sessionMgr + && owner.sessionId === sessionId + && owner.epoch === controlTargetEpoch + && sessionMgr.controlTargetEpoch === controlTargetEpoch; + if (targetChanged) { + messagesRequestSeqRef.current += 1; + modelCatalogRequestSeqRef.current += 1; + isLoadingMoreRef.current = false; + hasMoreRef.current = true; + setIsLoadingMore(false); + setHasMore(true); + setModelUpdating(false); + setImageAnalyzing(false); + setOptimisticMsg(null); + modelSelectionInitializedRef.current = false; + setModelCatalog(null); + setSelectedModelId('auto'); + setMessages(sessionId, []); + setMenuMessage(null); + setDeletingMsg(false); + setActionToast(null); + setInfoToast(null); + setExpandedMsgIds(new Set()); + setShowScrollToBottom(false); + setActiveTurn(null); + if (msgLongPressTimerRef.current) { + clearTimeout(msgLongPressTimerRef.current); + msgLongPressTimerRef.current = undefined; + } + if (msgToastTimerRef.current) { + clearTimeout(msgToastTimerRef.current); + msgToastTimerRef.current = undefined; + } + pollerRef.current?.stop(); + pollerRef.current = null; + } + committedChatTargetRef.current = { sessionMgr, sessionId, epoch: controlTargetEpoch }; + return () => { + owner.active = false; + messagesRequestSeqRef.current += 1; + modelCatalogRequestSeqRef.current += 1; + pollerRef.current?.stop(); + }; + }, [controlTargetEpoch, sessionId, sessionMgr, setActiveTurn, setMessages]); const isStreaming = activeTurn != null && activeTurn.status === 'active'; const [now, setNow] = useState(() => Date.now()); const handleAnswerQuestion = useCallback(async (toolId: string, answers: any) => { + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) throw new RemoteControlTargetChangedError(); try { await sessionMgr.answerQuestion(toolId, answers); + if (!isChatTargetCurrent(targetEpoch)) throw new RemoteControlTargetChangedError(); } catch (err) { - setError(String(err)); + reportRemoteSessionError(err, setError); throw err; } - }, [sessionMgr, setError]); + }, [captureChatTargetEpoch, isChatTargetCurrent, sessionMgr, setError]); /** Fetch metadata for a workspace file before the user confirms the download. */ const handleGetFileInfo = useCallback( - (filePath: string) => sessionMgr.getFileInfo(filePath, sessionId), - [sessionId, sessionMgr], + async (filePath: string) => { + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) throw new RemoteControlTargetChangedError(); + const info = await sessionMgr.getFileInfo(filePath, sessionId); + if (!isChatTargetCurrent(targetEpoch)) throw new RemoteControlTargetChangedError(); + return info; + }, + [captureChatTargetEpoch, isChatTargetCurrent, sessionId, sessionMgr], ); /** Download a workspace file referenced by a `computer://` link. */ @@ -2127,12 +2247,17 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, filePath: string, onProgress?: (downloaded: number, total: number) => void, ) => { + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) return; try { const { name, contentBase64, mimeType } = await sessionMgr.readFile( filePath, sessionId, - onProgress, + (downloaded, total) => { + if (isChatTargetCurrent(targetEpoch)) onProgress?.(downloaded, total); + }, ); + if (!isChatTargetCurrent(targetEpoch)) return; const byteCharacters = atob(contentBase64); const byteNumbers = new Uint8Array(byteCharacters.length); for (let i = 0; i < byteCharacters.length; i++) { @@ -2149,14 +2274,21 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, URL.revokeObjectURL(url); } catch (err) { // Use the backend's message directly; it's already user-readable. - const msg = err instanceof Error ? err.message : String(err); - setError(msg); + reportRemoteSessionError(err, setError); + throw err; } - }, [sessionId, sessionMgr, setError]); + }, [captureChatTargetEpoch, isChatTargetCurrent, sessionId, sessionMgr, setError]); const loadModelCatalog = useCallback(async () => { + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) return null; + const requestSeq = ++modelCatalogRequestSeqRef.current; try { const catalog = await sessionMgr.getModelCatalog(sessionId); + if ( + requestSeq !== modelCatalogRequestSeqRef.current + || !isChatTargetCurrent(targetEpoch) + ) return null; setModelCatalog(catalog); if (!modelSelectionInitializedRef.current) { const preferredSelection = resolvePreferredModelSelection(loadLastSelectedModelId(), catalog); @@ -2165,6 +2297,10 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, if (preferredSelection.modelId && preferredSelection.modelId !== sessionModelId) { const normalizedModelId = await sessionMgr.setSessionModel(sessionId, preferredSelection.modelId); + if ( + requestSeq !== modelCatalogRequestSeqRef.current + || !isChatTargetCurrent(targetEpoch) + ) return null; setSelectedModelId(normalizedModelId || 'auto'); if (preferredSelection.fellBackToAuto && (!normalizedModelId || normalizedModelId === 'auto')) { persistLastSelectedModelId('auto'); @@ -2179,24 +2315,30 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, } return catalog; } catch (err) { - setError(err instanceof Error ? err.message : String(err)); + if ( + requestSeq === modelCatalogRequestSeqRef.current + && isChatTargetCurrent(targetEpoch) + ) reportRemoteSessionError(err, setError); return null; } - }, [sessionId, sessionMgr, setError]); + }, [captureChatTargetEpoch, isChatTargetCurrent, sessionId, sessionMgr, setError]); const handleSelectModel = useCallback(async (modelId: string) => { if (modelUpdating || isStreaming || imageAnalyzing) return; + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) return; setModelUpdating(true); try { const normalizedModelId = await sessionMgr.setSessionModel(sessionId, modelId); + if (!isChatTargetCurrent(targetEpoch)) return; setSelectedModelId(normalizedModelId || 'auto'); persistLastSelectedModelId(normalizedModelId || 'auto'); } catch (err) { - setError(err instanceof Error ? err.message : String(err)); + reportRemoteSessionError(err, setError); } finally { - setModelUpdating(false); + if (isChatTargetCurrent(targetEpoch)) setModelUpdating(false); } - }, [imageAnalyzing, isStreaming, modelUpdating, sessionId, sessionMgr, setError]); + }, [captureChatTargetEpoch, imageAnalyzing, isChatTargetCurrent, isStreaming, modelUpdating, sessionId, sessionMgr, setError]); useEffect(() => { if (!isStreaming) return; @@ -2217,11 +2359,18 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, }, [infoToast]); const loadMessages = useCallback(async (beforeId?: string) => { - if (isLoadingMoreRef.current || (!hasMoreRef.current && beforeId)) return; + if (beforeId && (isLoadingMoreRef.current || !hasMoreRef.current)) return; + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) return; + const requestSeq = ++messagesRequestSeqRef.current; try { isLoadingMoreRef.current = true; setIsLoadingMore(true); const resp = await sessionMgr.getSessionMessages(sessionId, 50, beforeId); + if ( + requestSeq !== messagesRequestSeqRef.current + || !isChatTargetCurrent(targetEpoch) + ) return; if (beforeId) { const currentMsgs = getMessages(sessionId); setMessages(sessionId, [...resp.messages, ...currentMsgs]); @@ -2231,12 +2380,20 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, setHasMore(resp.has_more); hasMoreRef.current = resp.has_more; } catch (e: any) { - setError(e.message); + if ( + requestSeq === messagesRequestSeqRef.current + && isChatTargetCurrent(targetEpoch) + ) reportRemoteSessionError(e, setError); } finally { - isLoadingMoreRef.current = false; - setIsLoadingMore(false); + if ( + requestSeq === messagesRequestSeqRef.current + && isChatTargetCurrent(targetEpoch) + ) { + isLoadingMoreRef.current = false; + setIsLoadingMore(false); + } } - }, [sessionMgr, sessionId, setMessages, setError, getMessages]); + }, [captureChatTargetEpoch, getMessages, isChatTargetCurrent, sessionId, sessionMgr, setError, setMessages]); // ── Message long-press context menu ────────────────────────────── const clearMsgLongPressTimer = () => { @@ -2286,6 +2443,8 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, const handleResendMessage = useCallback(async () => { if (!menuMessage || menuMessage.role !== 'user') return; + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) return; const text = sanitizeMessageText(menuMessage.content); if (!text) return; setMenuMessage(null); @@ -2302,11 +2461,12 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, : undefined; try { await sessionMgr.sendMessage(sessionId, text, agentMode, imageContexts); + if (!isChatTargetCurrent(targetEpoch)) return; pollerRef.current?.nudge(); } catch (e: any) { - setError(e.message); + reportRemoteSessionError(e, setError); } - }, [menuMessage, sessionMgr, sessionId, agentMode, setError]); + }, [agentMode, captureChatTargetEpoch, isChatTargetCurrent, menuMessage, sessionId, sessionMgr, setError]); const handleDeleteMessage = useCallback(async () => { if (!menuMessage) return; @@ -2368,6 +2528,7 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, // Initial load + start poller const initialScrollDone = useRef(false); const pendingInitialScroll = useRef(false); + const chatInitSeqRef = useRef(0); useEffect(() => { modelSelectionInitializedRef.current = false; hasMoreRef.current = true; @@ -2381,11 +2542,22 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, useEffect(() => { initialScrollDone.current = false; pendingInitialScroll.current = false; + const initSeq = ++chatInitSeqRef.current; + let cancelled = false; + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) return; + const isInitCurrent = () => ( + !cancelled + && chatInitSeqRef.current === initSeq + && isChatTargetCurrent(targetEpoch) + ); Promise.all([loadMessages(), loadModelCatalog()]).then(([_, initialCatalog]) => { + if (!isInitCurrent()) return; const initialMsgCount = useMobileStore.getState().getMessages(sessionId).length; pendingInitialScroll.current = true; const poller = new SessionPoller(sessionMgr, sessionId, (resp: PollResponse) => { + if (!isInitCurrent()) return; if (resp.new_messages && resp.new_messages.length > 0) { appendNewMessages(sessionId, resp.new_messages); } @@ -2397,6 +2569,7 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, const localCount = useMobileStore.getState().getMessages(sessionId).length; if (localCount !== resp.total_msg_count) { sessionMgr.getSessionMessages(sessionId, 200).then(fresh => { + if (!isInitCurrent()) return; useMobileStore.getState().setMessages(sessionId, fresh.messages); }).catch(() => {}); } @@ -2417,11 +2590,23 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, }); return () => { + cancelled = true; + if (chatInitSeqRef.current === initSeq) chatInitSeqRef.current += 1; pollerRef.current?.stop(); pollerRef.current = null; setActiveTurn(null); }; - }, [sessionId, sessionMgr, loadMessages, loadModelCatalog, appendNewMessages, setActiveTurn, updateSessionName]); + }, [ + appendNewMessages, + captureChatTargetEpoch, + isChatTargetCurrent, + loadMessages, + loadModelCatalog, + sessionId, + sessionMgr, + setActiveTurn, + updateSessionName, + ]); const prevMsgCountRef = useRef(0); @@ -2484,6 +2669,8 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, const text = input.trim(); const imgs = pendingImages; if ((!text && imgs.length === 0) || imageAnalyzing) return; + const targetEpoch = captureChatTargetEpoch(); + if (targetEpoch === null) return; const wasStreaming = isStreaming; setInput(''); setPendingImages([]); @@ -2520,17 +2707,20 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, agentMode, imageContexts, ); + if (!isChatTargetCurrent(targetEpoch)) return; pollerRef.current?.nudge(); if (wasStreaming) { setInfoToast(t('chat.messageQueued')); } } catch (e: any) { - setError(e.message); + reportRemoteSessionError(e, setError); } finally { - setImageAnalyzing(false); - setOptimisticMsg(null); + if (isChatTargetCurrent(targetEpoch)) { + setImageAnalyzing(false); + setOptimisticMsg(null); + } } - }, [agentMode, imageAnalyzing, input, isStreaming, pendingImages, sessionId, sessionMgr, setError, t]); + }, [agentMode, captureChatTargetEpoch, imageAnalyzing, input, isChatTargetCurrent, isStreaming, pendingImages, sessionId, sessionMgr, setError, t]); const handleImageSelect = useCallback(() => { fileInputRef.current?.click(); @@ -2623,6 +2813,7 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, }; const handleCancel = async () => { + if (captureChatTargetEpoch() === null) return; try { await sessionMgr.cancelTask(sessionId, activeTurn?.turn_id); } catch { @@ -2825,7 +3016,16 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, return (
{turnIsActive - ? renderActiveTurnItems(turn.items, now, sessionMgr, setError, handleAnswerQuestion, handleFileDownload, handleGetFileInfo) + ? renderActiveTurnItems( + turn.items, + now, + sessionMgr, + setError, + () => captureChatTargetEpoch() !== null, + handleAnswerQuestion, + handleFileDownload, + handleGetFileInfo, + ) : renderOrderedItems(turn.items, now, undefined, undefined, handleFileDownload, handleGetFileInfo)} {turnIsActive && !turn.thinking && !turn.text && turn.tools.length === 0 && (
@@ -2848,7 +3048,10 @@ const ChatPage: React.FC = ({ sessionMgr, sessionId, sessionName, ] : []; const onCancel = (toolId: string) => { - sessionMgr.cancelTool(toolId, t('common.cancel')).catch(err => { setError(String(err)); }); + if (captureChatTargetEpoch() === null) return; + sessionMgr.cancelTool(toolId, t('common.cancel')).catch((error) => { + reportRemoteSessionError(error, setError); + }); }; return ( diff --git a/src/mobile-web/src/pages/DevicesPage.tsx b/src/mobile-web/src/pages/DevicesPage.tsx index b5ac28f310..f8198ca620 100644 --- a/src/mobile-web/src/pages/DevicesPage.tsx +++ b/src/mobile-web/src/pages/DevicesPage.tsx @@ -8,7 +8,10 @@ */ import React, { useState, useEffect, useCallback, useMemo, useRef } from 'react'; -import { RelayHttpClient } from '../services/RelayHttpClient'; +import { + RelayHttpClient, + isDelegatedIdentityChangedError, +} from '../services/RelayHttpClient'; import { useI18n } from '../i18n'; import { useMobileStore } from '../services/store'; @@ -57,7 +60,7 @@ const NoIdentityIcon = () => ( ); const DevicesPage: React.FC = ({ client, onBack }) => { - const { t } = useI18n(); + const { t, formatRelativeTime } = useI18n(); const { setControlTarget, resetForDeviceSwitch } = useMobileStore(); const [devices, setDevices] = useState([]); const [identityReady, setIdentityReady] = useState(client.hasDelegatedIdentity); @@ -66,6 +69,9 @@ const DevicesPage: React.FC = ({ client, onBack }) => { const [switchingId, setSwitchingId] = useState(null); const [error, setError] = useState(null); const mountedRef = useRef(true); + const identityRequestRef = useRef(0); + const devicesRequestRef = useRef(0); + const switchRequestRef = useRef(0); const sortedDevices = useMemo(() => [...devices].sort((left, right) => { const leftCurrent = left.device_id === client.pairedDeviceId; const rightCurrent = right.device_id === client.pairedDeviceId; @@ -90,19 +96,31 @@ const DevicesPage: React.FC = ({ client, onBack }) => { mountedRef.current = true; return () => { mountedRef.current = false; + identityRequestRef.current += 1; + devicesRequestRef.current += 1; + switchRequestRef.current += 1; }; }, []); const refreshDevices = useCallback(async () => { if (!client.hasDelegatedIdentity) return; + const requestId = ++devicesRequestRef.current; + const isCurrent = () => ( + mountedRef.current + && devicesRequestRef.current === requestId + ); try { const list = await client.listDevices(); - if (!mountedRef.current) return; + if (!isCurrent()) return; setDevices(list); setError(null); setIdentityReady(true); } catch (e: unknown) { - if (!mountedRef.current) return; + if (!isCurrent()) return; + // RelayHttpClient fences every response against its committed identity. + // A concurrent account refresh therefore makes this request stale rather + // than user-visible, while a successful 401 refresh + retry remains valid. + if (isDelegatedIdentityChangedError(e)) return; const message = String((e as { message?: string })?.message || e); if (message.includes('No delegated identity')) { setIdentityReady(false); @@ -117,6 +135,7 @@ const DevicesPage: React.FC = ({ client, onBack }) => { // its account after this mobile session was paired. Force-refresh so a // desktop account switch is reflected without re-scanning. const ensureIdentity = useCallback(async (force = false) => { + const requestId = ++identityRequestRef.current; setIdentityChecking(true); setError(null); let granted = false; @@ -124,13 +143,16 @@ const DevicesPage: React.FC = ({ client, onBack }) => { granted = await client.requestDelegatedIdentity({ force: force || !client.hasDelegatedIdentity }); } catch (e: unknown) { granted = false; - if (mountedRef.current) setError(friendlyError(e, 'devices.identityFailed')); + if (mountedRef.current && identityRequestRef.current === requestId) { + setError(friendlyError(e, 'devices.identityFailed')); + } } - if (mountedRef.current) { + if (mountedRef.current && identityRequestRef.current === requestId) { setIdentityReady(granted); setIdentityChecking(false); + return granted; } - return granted; + return false; }, [client, friendlyError]); useEffect(() => { @@ -162,6 +184,15 @@ const DevicesPage: React.FC = ({ client, onBack }) => { const selectDevice = useCallback(async (d: DeviceInfo) => { if (!d.online || switchingId) return; if (client.pairedDeviceId === d.device_id) return; + const requestId = ++switchRequestRef.current; + const accountEpoch = client.delegatedAccountEpoch; + let expectedTargetEpoch = client.controlTargetEpoch; + const isCurrent = () => ( + mountedRef.current + && switchRequestRef.current === requestId + && client.delegatedAccountEpoch === accountEpoch + && client.controlTargetEpoch === expectedTargetEpoch + ); setSwitchingId(d.device_id); setError(null); try { @@ -171,10 +202,12 @@ const DevicesPage: React.FC = ({ client, onBack }) => { command: 'peer_mode_ping', args: {}, }); + if (!isCurrent()) return; if (ping.resp === 'host_invoke_result' && ping.ok === false) { throw new Error(ping.error || t('devices.switchFailed')); } - client.pairedDeviceId = d.device_id; + client.setPairedDeviceId(d.device_id); + expectedTargetEpoch = client.controlTargetEpoch; resetForDeviceSwitch(); setControlTarget({ deviceId: d.device_id, @@ -183,7 +216,8 @@ const DevicesPage: React.FC = ({ client, onBack }) => { }); onBack(); } catch (e: unknown) { - if (!mountedRef.current) return; + if (!isCurrent()) return; + if (isDelegatedIdentityChangedError(e)) return; const message = String((e as { message?: string })?.message || e); if (message.includes('No delegated identity')) { setIdentityReady(false); @@ -192,7 +226,9 @@ const DevicesPage: React.FC = ({ client, onBack }) => { setError(friendlyError(e, 'devices.switchFailed')); } } finally { - if (mountedRef.current) setSwitchingId(null); + if (mountedRef.current && switchRequestRef.current === requestId) { + setSwitchingId(null); + } } }, [client, friendlyError, onBack, resetForDeviceSwitch, setControlTarget, switchingId, t]); @@ -270,7 +306,11 @@ const DevicesPage: React.FC = ({ client, onBack }) => { - {d.online ? t('devices.online') : t('devices.offline')} + {d.online + ? t('devices.online') + : d.last_seen_at + ? t('devices.lastSeen', { time: formatRelativeTime(d.last_seen_at * 1000) }) + : t('devices.offline')} {d.device_id.slice(0, 8)} @@ -301,7 +341,7 @@ const DevicesPage: React.FC = ({ client, onBack }) => { type="button" className={`devices-page__refresh-btn ${loading || identityChecking ? 'is-loading' : ''}`} onClick={handleManualRefresh} - disabled={loading || !!switchingId} + disabled={loading || identityChecking || !!switchingId} aria-label={t('devices.refresh')} title={t('devices.refresh')} > diff --git a/src/mobile-web/src/pages/PairingPage.tsx b/src/mobile-web/src/pages/PairingPage.tsx index ee1cf8f585..1f29580ce2 100644 --- a/src/mobile-web/src/pages/PairingPage.tsx +++ b/src/mobile-web/src/pages/PairingPage.tsx @@ -262,9 +262,16 @@ const PairingPage: React.FC = ({ onPaired }) => { const homeDeviceId = client.homeDeviceId; if (delegated && homeDeviceId) { store.setControlTarget({ deviceId: homeDeviceId, deviceName: null, isHome: true }); + const accountEpoch = client.delegatedAccountEpoch; + const target = client.getControlTargetSnapshot(); void client .listDevices() .then((devices) => { + if ( + client.delegatedAccountEpoch !== accountEpoch + || !client.isControlTargetCurrent(target) + || client.pairedDeviceId !== homeDeviceId + ) return; const home = devices.find((d) => d.device_id === homeDeviceId); if (home) { useMobileStore.getState().setControlTarget({ diff --git a/src/mobile-web/src/pages/SessionListPage.tsx b/src/mobile-web/src/pages/SessionListPage.tsx index 611ecc76a3..4e959488ca 100644 --- a/src/mobile-web/src/pages/SessionListPage.tsx +++ b/src/mobile-web/src/pages/SessionListPage.tsx @@ -1,7 +1,13 @@ -import React, { useEffect, useRef, useCallback, useState } from 'react'; +import React, { useEffect, useLayoutEffect, useRef, useCallback, useState } from 'react'; import LanguageToggleButton from '../components/LanguageToggleButton'; +import { useControlTargetEpoch } from '../hooks/useControlTargetEpoch'; import { useI18n } from '../i18n'; -import { RemoteSessionManager, type RecentWorkspaceEntry, type SessionInfo } from '../services/RemoteSessionManager'; +import { + isRemoteControlTargetChangedError, + RemoteSessionManager, + type RecentWorkspaceEntry, + type SessionInfo, +} from '../services/RemoteSessionManager'; import { useMobileStore } from '../services/store'; import { useTheme } from '../theme'; import logoIcon from '../assets/Logo-ICON.png'; @@ -18,6 +24,31 @@ interface SessionListPageProps { onOpenDevices?: () => void; } +type SessionListTargetOwner = { + sessionMgr: RemoteSessionManager; + epoch: number; + active: boolean; +}; + +/** + * Resolve the epoch owned by one render. The explicit renderedEpoch check is + * what prevents an old timer/poll closure from borrowing a newer mutable ref + * owner during the render-to-passive-cleanup window. + */ +export function captureSessionListOwnerEpoch( + owner: SessionListTargetOwner, + sessionMgr: RemoteSessionManager, + renderedEpoch: number, +): number | null { + if ( + !owner.active + || owner.sessionMgr !== sessionMgr + || owner.epoch !== renderedEpoch + || sessionMgr.controlTargetEpoch !== renderedEpoch + ) return null; + return renderedEpoch; +} + function formatTime( unixStr: string, formatDate: (date: Date | number, options?: Intl.DateTimeFormatOptions) => string, @@ -163,6 +194,8 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS const [creating, setCreating] = useState(false); const [loading, setLoading] = useState(false); const [loadingMore, setLoadingMore] = useState(false); + const [targetInitializing, setTargetInitializing] = useState(true); + const targetInitializingRef = useRef(true); const [hasMore, setHasMore] = useState(false); const [displayMode, setDisplayMode] = useState(() => { const hint = useMobileStore.getState().pairedDisplayMode; @@ -198,6 +231,39 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS const longPressPosRef = useRef({ x: 0, y: 0 }); const longPressTriggeredRef = useRef(false); const toastTimerRef = useRef>(); + const controlTargetEpoch = useControlTargetEpoch(sessionMgr); + const sessionListOwnerRef = useRef({ + sessionMgr, + epoch: controlTargetEpoch, + active: true, + }); + if ( + sessionListOwnerRef.current.sessionMgr !== sessionMgr + || sessionListOwnerRef.current.epoch !== controlTargetEpoch + ) { + sessionListOwnerRef.current = { + sessionMgr, + epoch: controlTargetEpoch, + active: true, + }; + } + + const captureSessionListEpoch = useCallback((): number | null => { + return captureSessionListOwnerEpoch( + sessionListOwnerRef.current, + sessionMgr, + controlTargetEpoch, + ); + }, [controlTargetEpoch, sessionMgr]); + + const isSessionListCurrent = useCallback((epoch: number | null): boolean => { + const owner = sessionListOwnerRef.current; + return epoch !== null + && owner.active + && owner.sessionMgr === sessionMgr + && owner.epoch === epoch + && sessionMgr.controlTargetEpoch === epoch; + }, [controlTargetEpoch, sessionMgr]); const hasSearchQuery = searchQuery.trim().length > 0; // Show the resume card as soon as session data is available — don't gate it @@ -264,34 +330,44 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS const handleRename = useCallback(async () => { if (!renameTarget || !renameValue.trim()) return; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; setRenaming(true); try { await sessionMgr.renameSession(renameTarget.session_id, renameValue.trim()); + if (!isSessionListCurrent(targetEpoch)) return; useMobileStore.getState().updateSessionName(renameTarget.session_id, renameValue.trim()); setRenameTarget(null); setMenuSession(null); } catch (e: any) { - showToast(e.message || t('sessions.renameFailed')); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + showToast(e.message || t('sessions.renameFailed')); + } } finally { - setRenaming(false); + if (isSessionListCurrent(targetEpoch)) setRenaming(false); } - }, [renameTarget, renameValue, sessionMgr, showToast, t]); + }, [captureSessionListEpoch, isSessionListCurrent, renameTarget, renameValue, sessionMgr, showToast, t]); const handleDelete = useCallback(async () => { if (!deleteConfirmTarget) return; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; setDeleting(true); try { await sessionMgr.deleteSession(deleteConfirmTarget.session_id); + if (!isSessionListCurrent(targetEpoch)) return; useMobileStore.getState().removeSession(deleteConfirmTarget.session_id); setDeleteConfirmTarget(null); setMenuSession(null); showToast(t('sessions.deleted')); } catch (e: any) { - showToast(e.message || t('sessions.deleteFailed')); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + showToast(e.message || t('sessions.deleteFailed')); + } } finally { - setDeleting(false); + if (isSessionListCurrent(targetEpoch)) setDeleting(false); } - }, [deleteConfirmTarget, sessionMgr, showToast, t]); + }, [captureSessionListEpoch, deleteConfirmTarget, isSessionListCurrent, sessionMgr, showToast, t]); const [pullDistance, setPullDistance] = useState(0); const [refreshing, setRefreshing] = useState(false); @@ -302,10 +378,76 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS const touchStartY = useRef(0); const isPulling = useRef(false); + const committedSessionListTargetRef = useRef({ sessionMgr, epoch: controlTargetEpoch }); + useLayoutEffect(() => { + const previous = committedSessionListTargetRef.current; + const targetChanged = previous.sessionMgr !== sessionMgr + || previous.epoch !== controlTargetEpoch; + const owner = sessionListOwnerRef.current; + owner.active = owner.sessionMgr === sessionMgr + && owner.epoch === controlTargetEpoch + && sessionMgr.controlTargetEpoch === controlTargetEpoch; + if (targetChanged) { + targetInitializingRef.current = true; + setTargetInitializing(true); + listRequestSeqRef.current += 1; + setLoading(false); + setLoadingMore(false); + setRefreshing(false); + clearLongPressTimer(); + longPressTriggeredRef.current = false; + isPulling.current = false; + setPullDistance(0); + if (toastTimerRef.current) { + clearTimeout(toastTimerRef.current); + toastTimerRef.current = undefined; + } + setCreating(false); + setRenaming(false); + setDeleting(false); + setAssistantList([]); + setWorkspaceList([]); + setShowAssistantPicker(false); + setShowWorkspacePicker(false); + setMenuSession(null); + setRenameTarget(null); + setRenameValue(''); + setDeleteConfirmTarget(null); + setActionToast(null); + setShowDisconnectConfirm(false); + setSearchQuery(''); + setDisplayMode('pro'); + setHasMore(false); + setSessions([]); + setCurrentWorkspace(null); + setCurrentAssistant(null); + setPairedDisplayMode(null); + setError(null); + offsetRef.current = 0; + initLoadedPathRef.current = undefined; + } + committedSessionListTargetRef.current = { sessionMgr, epoch: controlTargetEpoch }; + return () => { + owner.active = false; + listRequestSeqRef.current += 1; + }; + }, [ + controlTargetEpoch, + sessionMgr, + setCurrentAssistant, + setCurrentWorkspace, + setError, + setPairedDisplayMode, + setSessions, + ]); + // Load assistant list when entering assistant mode const loadAssistantList = useCallback(async () => { + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return undefined; try { const assistants = await sessionMgr.listAssistants(); + if (!isSessionListCurrent(targetEpoch)) return undefined; setAssistantList(assistants); // Set default assistant if none selected if (!currentAssistant && assistants.length > 0) { @@ -315,17 +457,23 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS } return currentAssistant?.path; } catch (e: any) { - setError(e.message); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + setError(e.message); + } return undefined; } - }, [sessionMgr, currentAssistant, setCurrentAssistant, setError]); + }, [captureSessionListEpoch, currentAssistant, isSessionListCurrent, sessionMgr, setCurrentAssistant, setError]); const loadFirstPage = useCallback(async ( workspacePath: string | undefined, query = '', identity?: { remoteConnectionId?: string; remoteSshHost?: string }, ) => { + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; const requestSeq = ++listRequestSeqRef.current; + // A new first page owns the complete list and supersedes pagination. + setLoadingMore(false); setLoading(true); offsetRef.current = 0; try { @@ -336,29 +484,43 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS query, identity, ); - if (requestSeq !== listRequestSeqRef.current) return; + if ( + requestSeq !== listRequestSeqRef.current + || !isSessionListCurrent(targetEpoch) + ) return; setSessions(resp.sessions); setHasMore(resp.has_more); offsetRef.current = resp.sessions.length; } catch (e: any) { - if (requestSeq !== listRequestSeqRef.current) return; - setError(e.message); + if ( + requestSeq !== listRequestSeqRef.current + || !isSessionListCurrent(targetEpoch) + ) return; + if (!isRemoteControlTargetChangedError(e)) setError(e.message); } finally { - if (requestSeq === listRequestSeqRef.current) { + if ( + requestSeq === listRequestSeqRef.current + && isSessionListCurrent(targetEpoch) + ) { setLoading(false); } } - }, [sessionMgr, setSessions, setError]); + }, [captureSessionListEpoch, isSessionListCurrent, sessionMgr, setError, setSessions]); // Load workspace list for Pro mode picker const loadWorkspaceList = useCallback(async () => { + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; try { const workspaces = await sessionMgr.listRecentWorkspaces(); + if (!isSessionListCurrent(targetEpoch)) return; setWorkspaceList(workspaces); } catch (e: any) { - setError(e.message); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + setError(e.message); + } } - }, [sessionMgr, setError]); + }, [captureSessionListEpoch, isSessionListCurrent, sessionMgr, setError]); const handleSelectWorkspace = useCallback(async (workspace: { path: string; @@ -366,11 +528,15 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS remote_connection_id?: string; remote_ssh_host?: string; }) => { + if (targetInitializingRef.current) return; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; try { const result = await sessionMgr.setWorkspace(workspace.path, { remoteConnectionId: workspace.remote_connection_id, remoteSshHost: workspace.remote_ssh_host, }); + if (!isSessionListCurrent(targetEpoch)) return; if (result.success) { const path = result.path || workspace.path; const remoteConnectionId = @@ -393,19 +559,25 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS setError(result.error || 'Failed to set workspace'); } } catch (e: any) { - setError(e.message); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + setError(e.message); + } } - }, [sessionMgr, setCurrentWorkspace, setError, loadFirstPage, searchQuery]); + }, [captureSessionListEpoch, isSessionListCurrent, loadFirstPage, searchQuery, sessionMgr, setCurrentWorkspace, setError]); const trySelectFirstProWorkspace = useCallback(async (): Promise => { + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return false; try { const list = await sessionMgr.listRecentWorkspaces(); + if (!isSessionListCurrent(targetEpoch)) return false; const candidate = pickFirstProWorkspace(list); if (!candidate) return false; const result = await sessionMgr.setWorkspace(candidate.path, { remoteConnectionId: candidate.remote_connection_id, remoteSshHost: candidate.remote_ssh_host, }); + if (!isSessionListCurrent(targetEpoch)) return false; if (result.success) { const path = result.path || candidate.path; const remoteConnectionId = @@ -423,22 +595,26 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS remote_ssh_host: remoteSshHost, }); await loadFirstPage(path, searchQuery, identity); - return true; + return isSessionListCurrent(targetEpoch); } setError(result.error || t('workspace.failedToSetWorkspace')); return false; } catch (e: any) { - setError(e.message); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + setError(e.message); + } return false; } - }, [sessionMgr, setCurrentWorkspace, setError, loadFirstPage, searchQuery, t]); + }, [captureSessionListEpoch, isSessionListCurrent, loadFirstPage, searchQuery, sessionMgr, setCurrentWorkspace, setError, t]); const loadNextPage = useCallback(async ( workspacePath: string | undefined, query = '', identity?: { remoteConnectionId?: string; remoteSshHost?: string }, ) => { - if (loadingMore || !hasMore) return; + if (loading || loadingMore || !hasMore) return; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; const requestSeq = listRequestSeqRef.current; setLoadingMore(true); try { @@ -449,24 +625,38 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS query, identity, ); - if (requestSeq !== listRequestSeqRef.current) return; + if ( + requestSeq !== listRequestSeqRef.current + || !isSessionListCurrent(targetEpoch) + ) return; appendSessions(resp.sessions); setHasMore(resp.has_more); offsetRef.current += resp.sessions.length; } catch (e: any) { - if (requestSeq !== listRequestSeqRef.current) return; - setError(e.message); + if ( + requestSeq !== listRequestSeqRef.current + || !isSessionListCurrent(targetEpoch) + ) return; + if (!isRemoteControlTargetChangedError(e)) setError(e.message); } finally { - setLoadingMore(false); + if ( + requestSeq === listRequestSeqRef.current + && isSessionListCurrent(targetEpoch) + ) setLoadingMore(false); } - }, [sessionMgr, appendSessions, setError, loadingMore, hasMore]); + }, [appendSessions, captureSessionListEpoch, hasMore, isSessionListCurrent, loading, loadingMore, sessionMgr, setError]); useEffect(() => { let cancelled = false; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; + const isInitCurrent = () => ( + !cancelled && isSessionListCurrent(targetEpoch) + ); const init = async () => { try { const info = await sessionMgr.getWorkspaceInfo(); - if (cancelled) return; + if (!isInitCurrent()) return; if (info.workspace_kind === 'assistant' && info.path) { setCurrentAssistant({ path: info.path, @@ -478,6 +668,7 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS initLoadedPathRef.current = info.path; await loadFirstPage(info.path); } else { + setDisplayMode('pro'); const ws = info.has_workspace ? info : null; setCurrentWorkspace(ws); if (ws?.path) { @@ -491,21 +682,36 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS } } } catch (e: any) { - if (!cancelled) setError(e.message); + if (isInitCurrent() && !isRemoteControlTargetChangedError(e)) setError(e.message); } finally { - if (!cancelled) setPairedDisplayMode(null); + if (isInitCurrent()) { + setPairedDisplayMode(null); + setLoading(false); + targetInitializingRef.current = false; + setTargetInitializing(false); + } } }; init(); return () => { cancelled = true; }; // eslint-disable-next-line react-hooks/exhaustive-deps - }, []); + }, [controlTargetEpoch]); const refreshData = useCallback(async () => { + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; const requestSeq = ++listRequestSeqRef.current; + // Refresh replaces both a first-page request and pagination. Their stale + // finally blocks intentionally cannot publish, so this owner must also + // settle the flags it superseded. + setLoadingMore(false); try { if (displayMode === 'pro') { const info = await sessionMgr.getWorkspaceInfo(); + if ( + requestSeq !== listRequestSeqRef.current + || !isSessionListCurrent(targetEpoch) + ) return; if (info.workspace_kind === 'assistant') { setCurrentWorkspace(null); setSessions([]); @@ -519,20 +725,35 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS remoteConnectionId: ws?.remote_connection_id, remoteSshHost: ws?.remote_ssh_host, }); - if (requestSeq !== listRequestSeqRef.current) return; + if ( + requestSeq !== listRequestSeqRef.current + || !isSessionListCurrent(targetEpoch) + ) return; setSessions(resp.sessions); setHasMore(resp.has_more); offsetRef.current = resp.sessions.length; } else { // Assistant mode: use currentAssistant path const resp = await sessionMgr.listSessions(currentAssistant?.path, PAGE_SIZE, 0, searchQuery); - if (requestSeq !== listRequestSeqRef.current) return; + if ( + requestSeq !== listRequestSeqRef.current + || !isSessionListCurrent(targetEpoch) + ) return; setSessions(resp.sessions); setHasMore(resp.has_more); offsetRef.current = resp.sessions.length; } } catch { /* ignore */ } - }, [sessionMgr, setSessions, setCurrentWorkspace, currentAssistant?.path, displayMode, searchQuery]); + finally { + if ( + requestSeq === listRequestSeqRef.current + && isSessionListCurrent(targetEpoch) + ) { + setLoading(false); + setLoadingMore(false); + } + } + }, [captureSessionListEpoch, currentAssistant?.path, displayMode, isSessionListCurrent, searchQuery, sessionMgr, setCurrentWorkspace, setSessions]); useEffect(() => { const poll = setInterval(refreshData, 10000); @@ -592,14 +813,16 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS const handleTouchEnd = useCallback(async () => { if (!isPulling.current) return; isPulling.current = false; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; if (pullDistance >= PULL_THRESHOLD) { setRefreshing(true); setPullDistance(PULL_THRESHOLD); await refreshData(); - setRefreshing(false); + if (isSessionListCurrent(targetEpoch)) setRefreshing(false); } - setPullDistance(0); - }, [pullDistance, refreshData]); + if (isSessionListCurrent(targetEpoch)) setPullDistance(0); + }, [captureSessionListEpoch, isSessionListCurrent, pullDistance, refreshData]); const handleScroll = useCallback((e: React.UIEvent) => { const el = e.currentTarget; @@ -624,7 +847,9 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS ]); const handleCreate = useCallback(async (agentType: string) => { - if (creating) return; + if (creating || targetInitializingRef.current) return; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; setCreating(true); try { // For assistant mode (Claw), use currentAssistant.path @@ -635,9 +860,11 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS : { remoteConnectionId: currentWorkspace?.remote_connection_id, remoteSshHost: currentWorkspace?.remote_ssh_host, - }; + }; const id = await sessionMgr.createSession(agentType, undefined, workspacePath, identity); + if (!isSessionListCurrent(targetEpoch)) return; await loadFirstPage(workspacePath, searchQuery, identity); + if (!isSessionListCurrent(targetEpoch)) return; const label = isClawAgent(agentType) ? t('sessions.remoteClawSession') : isCoworkAgent(agentType) @@ -645,17 +872,21 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS : t('sessions.remoteCodeSession'); onSelectSession(id, label, true); } catch (e: any) { - setError(e.message); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + setError(e.message); + } } finally { - setCreating(false); + if (isSessionListCurrent(targetEpoch)) setCreating(false); } }, [ creating, + captureSessionListEpoch, currentWorkspace?.path, currentWorkspace?.remote_connection_id, currentWorkspace?.remote_ssh_host, currentAssistant?.path, displayMode, + isSessionListCurrent, loadFirstPage, onSelectSession, searchQuery, @@ -665,10 +896,14 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS ]); const handleSelectMode = useCallback(async (mode: DisplayMode) => { + if (targetInitializingRef.current) return; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; setDisplayMode(mode); setShowAssistantPicker(false); if (mode === 'assistant') { const assistantPath = await loadAssistantList(); + if (!isSessionListCurrent(targetEpoch)) return; loadFirstPage(assistantPath, searchQuery); } else { if (currentWorkspace?.path) { @@ -680,18 +915,24 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS await trySelectFirstProWorkspace(); } } - }, [currentWorkspace?.path, loadFirstPage, loadAssistantList, searchQuery, trySelectFirstProWorkspace]); + }, [captureSessionListEpoch, currentWorkspace?.path, isSessionListCurrent, loadAssistantList, loadFirstPage, searchQuery, trySelectFirstProWorkspace]); const handleSelectAssistant = useCallback(async (assistant: { path: string; name: string; assistant_id?: string }) => { + if (targetInitializingRef.current) return; + const targetEpoch = captureSessionListEpoch(); + if (targetEpoch === null) return; try { await sessionMgr.setAssistant(assistant.path); + if (!isSessionListCurrent(targetEpoch)) return; setCurrentAssistant(assistant); setShowAssistantPicker(false); loadFirstPage(assistant.path, searchQuery); } catch (e: any) { - setError(e.message); + if (isSessionListCurrent(targetEpoch) && !isRemoteControlTargetChangedError(e)) { + setError(e.message); + } } - }, [sessionMgr, setCurrentAssistant, setError, loadFirstPage, searchQuery]); + }, [captureSessionListEpoch, isSessionListCurrent, loadFirstPage, searchQuery, sessionMgr, setCurrentAssistant, setError]); const workspaceDisplayName = currentWorkspace?.project_name || t('sessions.noWorkspaceSelected'); const assistantDisplayName = currentAssistant?.name || t('shared.agents.default'); @@ -809,6 +1050,7 @@ const SessionListPage: React.FC = ({ sessionMgr, onSelectS )}
- {loading && sessions.length === 0 && ( + {(loading || targetInitializing) && sessions.length === 0 && (
{t('sessions.loadingSessions')}
)} - {!loading && sessions.length === 0 && !hasSearchQuery && ( + {!loading && !targetInitializing && sessions.length === 0 && !hasSearchQuery && (
{t('sessions.noSessions')}
)} - {!loading && sessions.length === 0 && hasSearchQuery && ( + {!loading && !targetInitializing && sessions.length === 0 && hasSearchQuery && (
{t('sessions.emptySearch')}
)} diff --git a/src/mobile-web/src/services/RelayHttpClient.ts b/src/mobile-web/src/services/RelayHttpClient.ts index 8635deb164..9b43580b5f 100644 --- a/src/mobile-web/src/services/RelayHttpClient.ts +++ b/src/mobile-web/src/services/RelayHttpClient.ts @@ -16,21 +16,100 @@ import { type MobileKeyPair, } from './E2EEncryption'; +interface DelegatedIdentitySnapshot { + token: string; + masterKey: Uint8Array; + userId: string | null; + homeDeviceId: string | null; + generation: number; +} + +interface DelegatedAccountIdentity { + userId: string | null; + masterKey: Uint8Array; + homeDeviceId: string | null; +} + +export type DelegatedAccountOwnerChange = { + kind: 'initial' | 'replacement' | 'unavailable'; + epoch: number; + userId: string | null; + homeDeviceId: string | null; +}; + +export type ControlTargetSnapshot = Readonly<{ + deviceId: string | null; + homeDeviceId: string | null; + epoch: number; +}>; + +export class DelegatedIdentityChangedError extends Error { + constructor(message = 'Delegated identity changed') { + super(message); + this.name = 'DelegatedIdentityChangedError'; + } +} + +export class DelegatedAccountChangedError extends DelegatedIdentityChangedError { + constructor() { + super('Delegated account changed'); + this.name = 'DelegatedAccountChangedError'; + } +} + +export function isDelegatedIdentityChangedError( + value: unknown, +): value is DelegatedIdentityChangedError { + return value instanceof DelegatedIdentityChangedError; +} + +function equalBytesConstantTime(left: Uint8Array, right: Uint8Array): boolean { + if (left.length !== right.length) return false; + let difference = 0; + for (let index = 0; index < left.length; index += 1) { + difference |= left[index] ^ right[index]; + } + return difference === 0; +} + +function delegatedAccountChanged( + previous: DelegatedAccountIdentity | null, + next: DelegatedAccountIdentity, +): boolean { + if (!previous) return true; + if (!equalBytesConstantTime(previous.masterKey, next.masterKey)) return true; + if (previous.userId !== null && next.userId !== null) { + return previous.userId !== next.userId; + } + // Older Desktop builds omit user_id. The account master key is the stable + // identifier in that case; homeDeviceId separates distinct paired homes. + return previous.homeDeviceId !== next.homeDeviceId; +} + export class RelayHttpClient { private relayUrl: string; private roomId: string; private sharedKey: Uint8Array | null = null; private keyPair: MobileKeyPair | null = null; - /** Delegated account identity (token + master_key) from the paired desktop. */ - public delegatedToken: string | null = null; - public delegatedMasterKey: Uint8Array | null = null; + /** Delegated credentials are committed as one immutable generation. */ + private delegatedIdentity: DelegatedIdentitySnapshot | null = null; + private delegatedIdentityRequestEpoch = 0; + private delegatedIdentityGenerationValue = 0; + private delegatedIdentityRefreshOwner: number | null = null; + private delegatedAccountIdentity: DelegatedAccountIdentity | null = null; + private delegatedAccountEpochValue = 0; + private delegatedAccountOwnerListeners = new Set<( + change: DelegatedAccountOwnerChange, + ) => void>(); /** The current control-target device_id (for sendDeviceRpc). */ - public pairedDeviceId: string | null = null; + private pairedDeviceIdValue: string | null = null; + private controlTargetEpochValue = 0; + private controlTargetListeners = new Set<(snapshot: ControlTargetSnapshot) => void>(); /** The QR-paired desktop's device_id (the "home" device of this session). */ public homeDeviceId: string | null = null; constructor(relayUrl: string, roomId: string) { - this.relayUrl = relayUrl.replace(/\/$/, ''); + this.relayUrl = relayUrl.replace(/\/+$/, ''); this.roomId = roomId; } @@ -164,34 +243,153 @@ export class RelayHttpClient { */ async requestDelegatedIdentity(options?: { force?: boolean }): Promise { if (!options?.force && this.hasDelegatedIdentity) return true; + const requestGeneration = ++this.delegatedIdentityRequestEpoch; if (options?.force) { - this.clearDelegatedIdentity(); + // Keep the last committed credential present for target routing, but + // suspend its use until this refresh settles. On transport failure the + // last confirmed identity becomes usable again. + this.delegatedIdentityRefreshOwner = requestGeneration; } - const resp = await this.sendCommand<{ - resp: string; - token?: string; - master_key?: string; - device_id?: string; - message?: string; - }>({ cmd: 'get_delegated_identity' }); - if (resp?.resp === 'delegate_identity' && resp.token && resp.master_key) { - this.delegatedToken = resp.token; - this.delegatedMasterKey = fromB64(resp.master_key); - if (resp.device_id) { - this.homeDeviceId = resp.device_id; - if (!this.pairedDeviceId) { - this.pairedDeviceId = resp.device_id; + try { + const resp = await this.sendCommand<{ + resp: string; + token?: string; + master_key?: string; + user_id?: string; + device_id?: string; + message?: string; + }>({ cmd: 'get_delegated_identity' }); + if (this.delegatedIdentityRequestEpoch !== requestGeneration) return false; + if (resp?.resp === 'delegate_identity' && resp.token && resp.master_key) { + const homeDeviceId = resp.device_id ?? null; + const nextIdentity: DelegatedIdentitySnapshot = { + token: resp.token, + masterKey: fromB64(resp.master_key), + userId: resp.user_id ?? null, + homeDeviceId, + generation: ++this.delegatedIdentityGenerationValue, + }; + this.delegatedIdentity = nextIdentity; + + const nextAccountIdentity: DelegatedAccountIdentity = { + userId: nextIdentity.userId, + masterKey: nextIdentity.masterKey.slice(), + homeDeviceId: nextIdentity.homeDeviceId, + }; + const accountChanged = delegatedAccountChanged( + this.delegatedAccountIdentity, + nextAccountIdentity, + ); + const hadAccountIdentity = this.delegatedAccountIdentity !== null; + if (accountChanged) this.delegatedAccountEpochValue += 1; + this.delegatedAccountIdentity = nextAccountIdentity; + + const previousHomeDeviceId = this.homeDeviceId; + const wasUsingDefaultRoom = this.pairedDeviceIdValue === null + || this.pairedDeviceIdValue === previousHomeDeviceId; + this.homeDeviceId = homeDeviceId; + if (accountChanged) { + // Every semantic account-owner commit starts a new UI/data + // generation, including a late initial delegation. Explicit React + // epoch subscribers re-initialize same-owner screens safely. + this.setPairedDeviceId(homeDeviceId); + } else if (wasUsingDefaultRoom) { + // Token/home metadata refresh for the same committed account does + // not change the effective QR-room route. + this.pairedDeviceIdValue = homeDeviceId; + } + if (accountChanged) { + this.emitDelegatedAccountOwnerChange({ + kind: hadAccountIdentity ? 'replacement' : 'initial', + epoch: this.delegatedAccountEpochValue, + userId: nextIdentity.userId, + homeDeviceId, + }); } + return true; + } + this.commitDelegatedAccountUnavailable(); + return false; + } finally { + if (this.delegatedIdentityRefreshOwner === requestGeneration) { + this.delegatedIdentityRefreshOwner = null; } - return true; } - return false; } /** Drop cached delegated credentials so the next request can refresh them. */ clearDelegatedIdentity(): void { - this.delegatedToken = null; - this.delegatedMasterKey = null; + this.delegatedIdentityRequestEpoch += 1; + this.delegatedIdentityRefreshOwner = null; + if (this.delegatedIdentity) this.delegatedIdentityGenerationValue += 1; + this.delegatedIdentity = null; + } + + /** Fully discard account/control-target state when the mobile disconnects. */ + resetConnectionIdentity(): void { + const targetEpoch = this.controlTargetEpochValue; + this.clearDelegatedIdentity(); + this.commitDelegatedAccountUnavailable(); + if (this.controlTargetEpochValue === targetEpoch) { + // Disconnect is an explicit ownership boundary even when this session + // never received delegated credentials and was using only the QR room. + this.setPairedDeviceId(null); + } + } + + /** + * Observe semantic delegated-account owner changes. Token-only refreshes do + * not emit. The listener is synchronous with the credential commit so UI + * state is cleared before an operation can publish data for the new owner. + */ + onDelegatedAccountOwnerChange( + listener: (change: DelegatedAccountOwnerChange) => void, + options?: { emitCurrent?: boolean }, + ): () => void { + this.delegatedAccountOwnerListeners.add(listener); + if (options?.emitCurrent && this.delegatedAccountIdentity) { + listener({ + kind: 'initial', + epoch: this.delegatedAccountEpochValue, + userId: this.delegatedAccountIdentity.userId, + homeDeviceId: this.delegatedAccountIdentity.homeDeviceId, + }); + } + return () => this.delegatedAccountOwnerListeners.delete(listener); + } + + private emitDelegatedAccountOwnerChange(change: DelegatedAccountOwnerChange): void { + for (const listener of this.delegatedAccountOwnerListeners) { + listener(change); + } + } + + private commitDelegatedAccountUnavailable(): void { + if (this.delegatedIdentity) this.delegatedIdentityGenerationValue += 1; + this.delegatedIdentity = null; + if (!this.delegatedAccountIdentity) { + const wasUsingDefaultRoom = this.pairedDeviceIdValue === null + || this.pairedDeviceIdValue === this.homeDeviceId; + this.homeDeviceId = null; + if (wasUsingDefaultRoom) { + // A late "not logged in" result with no committed owner changes no + // route. Do not manufacture a target event that can freeze consumers. + this.pairedDeviceIdValue = null; + } else { + this.setPairedDeviceId(null); + } + return; + } + this.delegatedAccountEpochValue += 1; + this.delegatedAccountIdentity = null; + this.homeDeviceId = null; + this.setPairedDeviceId(null); + this.emitDelegatedAccountOwnerChange({ + kind: 'unavailable', + epoch: this.delegatedAccountEpochValue, + userId: null, + homeDeviceId: null, + }); } /** @@ -236,15 +434,106 @@ export class RelayHttpClient { } get hasDelegatedIdentity(): boolean { - return this.delegatedToken !== null && this.delegatedMasterKey !== null; + return this.delegatedIdentity !== null; + } + + get pairedDeviceId(): string | null { + return this.pairedDeviceIdValue; + } + + /** + * Commit a control target and advance its ownership epoch. Advancing even + * when the device id repeats is intentional: A -> B -> A must invalidate + * requests that were issued during the first A ownership interval. + */ + setPairedDeviceId(deviceId: string | null): void { + this.pairedDeviceIdValue = deviceId; + this.controlTargetEpochValue += 1; + const snapshot = this.getControlTargetSnapshot(); + for (const listener of this.controlTargetListeners) { + listener(snapshot); + } + } + + get controlTargetEpoch(): number { + return this.controlTargetEpochValue; + } + + getControlTargetSnapshot(): ControlTargetSnapshot { + return { + deviceId: this.pairedDeviceIdValue, + homeDeviceId: this.homeDeviceId, + epoch: this.controlTargetEpochValue, + }; + } + + isControlTargetCurrent(snapshot: ControlTargetSnapshot): boolean { + // The epoch represents route ownership. Device/home ids are immutable + // routing inputs captured by the request, but metadata-only home binding + // intentionally leaves an in-flight QR-room request current. + return snapshot.epoch === this.controlTargetEpochValue; + } + + onControlTargetChange( + listener: (snapshot: ControlTargetSnapshot) => void, + ): () => void { + this.controlTargetListeners.add(listener); + return () => this.controlTargetListeners.delete(listener); + } + + get delegatedIdentityGeneration(): number { + return this.delegatedIdentityGenerationValue; + } + + /** Changes only when the delegated account/home identity changes. */ + get delegatedAccountEpoch(): number { + return this.delegatedAccountEpochValue; + } + + /** Canonical delegated account user, or null for legacy Desktop responses. */ + get delegatedUserId(): string | null { + return this.delegatedIdentity?.userId ?? null; + } + + private requireDelegatedIdentity(): DelegatedIdentitySnapshot { + const identity = this.delegatedIdentity; + if (!identity) throw new Error('No delegated identity'); + return identity; + } + + private isDelegatedIdentityCurrent(identity: DelegatedIdentitySnapshot): boolean { + return this.delegatedIdentity?.generation === identity.generation + && this.delegatedIdentityGenerationValue === identity.generation + && this.delegatedIdentityRefreshOwner === null; + } + + private ensureDelegatedIdentityCurrent( + identity: DelegatedIdentitySnapshot, + accountEpoch: number, + ): void { + if (!this.isDelegatedIdentityCurrent(identity)) { + throw this.delegatedIdentityChangeError(accountEpoch); + } + } + + private delegatedIdentityChangeError(accountEpoch: number): DelegatedIdentityChangedError { + return this.delegatedAccountEpochValue === accountEpoch + ? new DelegatedIdentityChangedError() + : new DelegatedAccountChangedError(); } /** * Refresh delegated identity from the paired desktop after a 401, then * retry the caller once. */ - private async refreshDelegatedIdentityAfterUnauthorized(): Promise { - this.clearDelegatedIdentity(); + private async refreshDelegatedIdentityAfterUnauthorized( + failedIdentity: DelegatedIdentitySnapshot, + ): Promise { + // A 401 from an old generation must not clear credentials that were + // delegated by a newer desktop account in the meantime. + if (!this.isDelegatedIdentityCurrent(failedIdentity)) { + return this.hasDelegatedIdentity; + } try { return await this.requestDelegatedIdentity({ force: true }); } catch { @@ -258,10 +547,9 @@ export class RelayHttpClient { * On HTTP 401, refreshes identity from the paired desktop and retries once. */ async listDevices(): Promise> { - return this.withDelegatedAuthRetry(async () => { - if (!this.delegatedToken) throw new Error('No delegated identity'); + return this.withDelegatedAuthRetry(async (identity) => { const resp = await this.fetchWithTimeout(`${this.relayUrl}/api/devices`, { - headers: { 'Authorization': `Bearer ${this.delegatedToken}` }, + headers: { 'Authorization': `Bearer ${identity.token}` }, }, 20_000); if (!resp.ok) { const err = new Error(`List devices failed: HTTP ${resp.status}`) as Error & { @@ -271,7 +559,7 @@ export class RelayHttpClient { throw err; } return resp.json(); - }); + }, { allowAccountReplacementRetry: true }); } /** @@ -281,14 +569,10 @@ export class RelayHttpClient { * On HTTP 401, refreshes identity from the paired desktop and retries once. */ async sendDeviceRpc(targetDeviceId: string, command: object): Promise { - return this.withDelegatedAuthRetry(async () => { - if (!this.delegatedToken || !this.delegatedMasterKey) { - throw new Error('No delegated identity'); - } - + return this.withDelegatedAuthRetry(async (identity) => { const plaintext = JSON.stringify(command); const { data: encData, nonce: encNonce } = await encrypt( - this.delegatedMasterKey, + identity.masterKey, plaintext, ); @@ -298,7 +582,7 @@ export class RelayHttpClient { method: 'POST', headers: { 'Content-Type': 'application/json', - 'Authorization': `Bearer ${this.delegatedToken}`, + 'Authorization': `Bearer ${identity.token}`, }, body: JSON.stringify({ encrypted_data: encData, nonce: encNonce }), }, @@ -314,7 +598,7 @@ export class RelayHttpClient { } const data = await resp.json(); const decrypted = await decrypt( - this.delegatedMasterKey, + identity.masterKey, data.encrypted_data, data.nonce, ); @@ -323,13 +607,28 @@ export class RelayHttpClient { throw new Error(parsed.message || 'Remote error'); } return parsed as T; - }); + }, { allowAccountReplacementRetry: false }); } - private async withDelegatedAuthRetry(operation: () => Promise): Promise { + private async withDelegatedAuthRetry( + operation: (identity: DelegatedIdentitySnapshot) => Promise, + options: { allowAccountReplacementRetry: boolean }, + ): Promise { + let identity = this.requireDelegatedIdentity(); + const accountEpoch = this.delegatedAccountEpochValue; + // A forced refresh keeps the last committed bytes only so routing remains + // explicit; it suspends their authority. Fence before invoking the caller + // because checking only after a device RPC returns is too late for commands + // that may already have produced side effects on the old account. + this.ensureDelegatedIdentityCurrent(identity, accountEpoch); try { - return await operation(); + const result = await operation(identity); + this.ensureDelegatedIdentityCurrent(identity, accountEpoch); + return result; } catch (e: unknown) { + if (!this.isDelegatedIdentityCurrent(identity)) { + throw this.delegatedIdentityChangeError(accountEpoch); + } const status = (e as { status?: number })?.status; const message = String((e as { message?: string })?.message || e); const unauthorized = @@ -338,11 +637,22 @@ export class RelayHttpClient { || message.includes('Unauthorized'); if (!unauthorized) throw e; - const refreshed = await this.refreshDelegatedIdentityAfterUnauthorized(); + const refreshed = await this.refreshDelegatedIdentityAfterUnauthorized(identity); if (!refreshed) { throw new Error('No delegated identity'); } - return operation(); + if ( + !options.allowAccountReplacementRetry + && this.delegatedAccountEpochValue !== accountEpoch + ) { + // Never replay an A-owned RPC target/command with B's credentials. + // The account-owner listener has already cleared A's UI state. + throw new DelegatedAccountChangedError(); + } + identity = this.requireDelegatedIdentity(); + const result = await operation(identity); + this.ensureDelegatedIdentityCurrent(identity, this.delegatedAccountEpochValue); + return result; } } diff --git a/src/mobile-web/src/services/RemoteSessionManager.ts b/src/mobile-web/src/services/RemoteSessionManager.ts index a2c571f693..e95cc61030 100644 --- a/src/mobile-web/src/services/RemoteSessionManager.ts +++ b/src/mobile-web/src/services/RemoteSessionManager.ts @@ -8,7 +8,23 @@ * - On tab activation: immediate poll to catch up on missed changes */ -import { RelayHttpClient } from './RelayHttpClient'; +import { + RelayHttpClient, + type ControlTargetSnapshot, +} from './RelayHttpClient'; + +export class RemoteControlTargetChangedError extends Error { + constructor() { + super('Remote control target changed'); + this.name = 'RemoteControlTargetChangedError'; + } +} + +export function isRemoteControlTargetChangedError( + value: unknown, +): value is RemoteControlTargetChangedError { + return value instanceof RemoteControlTargetChangedError; +} export interface WorkspaceInfo { has_workspace: boolean; @@ -158,28 +174,57 @@ export class RemoteSessionManager { this.client = client; } - private async request(cmd: object): Promise { + get controlTargetEpoch(): number { + return this.client.controlTargetEpoch; + } + + onControlTargetChange(listener: () => void): () => void { + return this.client.onControlTargetChange(listener); + } + + private ensureControlTargetCurrent(snapshot: ControlTargetSnapshot): void { + if (!this.client.isControlTargetCurrent(snapshot)) { + throw new RemoteControlTargetChangedError(); + } + } + + private async request( + cmd: object, + target: ControlTargetSnapshot = this.client.getControlTargetSnapshot(), + ): Promise { + // A caller may bind several transport requests into one logical operation + // (for example, a chunked file download). Fence before any transport call + // so a stale operation cannot send its next step to the replacement target. + this.ensureControlTargetCurrent(target); const requestId = `req_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`; const cmdWithId = { ...cmd, _request_id: requestId }; // The QR-paired desktop keeps the proven room channel. Only a switched // control target (another same-account device) is reached through the // relay device RPC API using the delegated identity. - const targetDeviceId = this.client.pairedDeviceId; + const targetDeviceId = target.deviceId; const isRemoteTarget = - this.client.hasDelegatedIdentity - && !!targetDeviceId - && targetDeviceId !== this.client.homeDeviceId; - let resp: T; - if (isRemoteTarget && targetDeviceId) { - resp = await this.client.sendDeviceRpc(targetDeviceId, cmdWithId); - } else { - resp = await this.client.sendCommand(cmdWithId); - } - const respAny = resp as any; - if (respAny.resp === 'error') { - throw new Error(respAny.message || 'Unknown error'); + !!targetDeviceId + && targetDeviceId !== target.homeDeviceId; + try { + let resp: T; + if (isRemoteTarget && targetDeviceId) { + resp = await this.client.sendDeviceRpc(targetDeviceId, cmdWithId); + } else { + resp = await this.client.sendCommand(cmdWithId); + } + this.ensureControlTargetCurrent(target); + const respAny = resp as any; + if (respAny.resp === 'error') { + throw new Error(respAny.message || 'Unknown error'); + } + return resp; + } catch (error: unknown) { + // Suppress both successful and failed completions after a target switch. + // The epoch check (rather than device id alone) also closes A -> B -> A + // ABA races. + this.ensureControlTargetCurrent(target); + throw error; } - return resp; } async getWorkspaceInfo(): Promise { @@ -456,6 +501,7 @@ export class RemoteSessionManager { let fileName = ''; let mimeType = ''; let totalSize = 0; + const target = this.client.getControlTargetSnapshot(); // eslint-disable-next-line no-constant-condition while (true) { @@ -473,7 +519,8 @@ export class RemoteSessionManager { session_id: sessionId ?? undefined, offset, limit: CHUNK_SIZE, - }); + }, target); + this.ensureControlTargetCurrent(target); chunks.push(resp.chunk_base64); fileName = resp.name; @@ -486,6 +533,8 @@ export class RemoteSessionManager { if (offset >= totalSize || resp.chunk_size === 0) break; } + this.ensureControlTargetCurrent(target); + return { name: fileName, contentBase64: chunks.join(''), diff --git a/src/mobile-web/src/services/delegatedAccountOwner.ts b/src/mobile-web/src/services/delegatedAccountOwner.ts new file mode 100644 index 0000000000..6b8aadf262 --- /dev/null +++ b/src/mobile-web/src/services/delegatedAccountOwner.ts @@ -0,0 +1,46 @@ +import type { DelegatedAccountOwnerChange } from './RelayHttpClient'; +import { useMobileStore } from './store'; + +/** + * Reconcile a transport-level delegated-account commit with the mobile UI. + * Returns true when cached workspace/session/chat state belonged to a previous + * owner and callers should leave any detail page that may still reference it. + */ +export function reconcileDelegatedAccountOwner( + change: DelegatedAccountOwnerChange, +): boolean { + const store = useMobileStore.getState(); + const initialOwnerConflicts = change.kind === 'initial' + && change.userId !== null + && store.authenticatedUserId !== null + && store.authenticatedUserId !== change.userId; + const ownerWasReplaced = change.kind === 'replacement' + || change.kind === 'unavailable' + || initialOwnerConflicts; + + if (ownerWasReplaced) { + store.resetForDeviceSwitch(); + } + + if (change.kind === 'unavailable') { + store.setAuthenticatedUserId(null); + store.setControlTarget(null); + return true; + } + + if (change.userId !== null || change.kind === 'replacement') { + // A current Desktop reports userId. Legacy responses may omit it; only a + // confirmed replacement is allowed to clear an already known old owner. + store.setAuthenticatedUserId(change.userId); + } + + if (ownerWasReplaced || store.controlTarget === null) { + store.setControlTarget(change.homeDeviceId ? { + deviceId: change.homeDeviceId, + deviceName: null, + isHome: true, + } : null); + } + + return ownerWasReplaced; +} diff --git a/src/mobile-web/src/styles/components/language-toggle.scss b/src/mobile-web/src/styles/components/language-toggle.scss index 7b322b0388..c0fae2ec9e 100644 --- a/src/mobile-web/src/styles/components/language-toggle.scss +++ b/src/mobile-web/src/styles/components/language-toggle.scss @@ -2,8 +2,8 @@ display: inline-flex; align-items: center; justify-content: center; - min-width: 32px; - height: 32px; + min-width: 44px; + height: 44px; padding: 0 10px; border: 1px solid var(--border-subtle); border-radius: 999px; @@ -21,4 +21,3 @@ background: var(--element-bg-base); } } - diff --git a/src/mobile-web/src/styles/components/pairing.scss b/src/mobile-web/src/styles/components/pairing.scss index d70ad7143a..327b63399c 100644 --- a/src/mobile-web/src/styles/components/pairing.scss +++ b/src/mobile-web/src/styles/components/pairing.scss @@ -6,16 +6,21 @@ flex-direction: column; align-items: center; justify-content: center; - height: 100%; + min-height: 100%; + min-height: 100dvh; gap: var(--size-gap-6); - padding: var(--size-gap-8); + padding: + max(var(--size-gap-8), env(safe-area-inset-top, 0px)) + max(var(--size-gap-8), env(safe-area-inset-right, 0px)) + max(var(--size-gap-8), env(safe-area-inset-bottom, 0px)) + max(var(--size-gap-8), env(safe-area-inset-left, 0px)); animation: fadeIn var(--motion-slow) motion.$easing-decelerate; } .pairing-page__actions { position: absolute; - top: var(--size-gap-4); - right: var(--size-gap-4); + top: calc(var(--size-gap-4) + env(safe-area-inset-top, 0px)); + right: calc(var(--size-gap-4) + env(safe-area-inset-right, 0px)); display: flex; align-items: center; gap: var(--size-gap-2); @@ -25,8 +30,8 @@ display: inline-flex; align-items: center; justify-content: center; - width: 32px; - height: 32px; + width: 44px; + height: 44px; padding: 0; border: 1px solid var(--border-subtle); border-radius: 50%; diff --git a/src/mobile-web/src/styles/global.scss b/src/mobile-web/src/styles/global.scss index c46f836c7c..be2bac3361 100644 --- a/src/mobile-web/src/styles/global.scss +++ b/src/mobile-web/src/styles/global.scss @@ -9,6 +9,7 @@ html, body, #root { height: 100%; + height: 100dvh; width: 100%; overflow: hidden; } @@ -25,6 +26,7 @@ body { .mobile-app { height: 100%; + height: 100dvh; width: 100%; background: var(--color-bg-primary); color: var(--color-text-primary); diff --git a/src/web-ui/src/app/components/NavPanel/MainNav.tsx b/src/web-ui/src/app/components/NavPanel/MainNav.tsx index b25e71aaf7..82b4822a2d 100644 --- a/src/web-ui/src/app/components/NavPanel/MainNav.tsx +++ b/src/web-ui/src/app/components/NavPanel/MainNav.tsx @@ -13,7 +13,7 @@ import React, { useCallback, useState, useMemo, useEffect, useRef } from 'react'; import { createPortal } from 'react-dom'; -import { Plus, FolderOpen, FolderPlus, History, Check, User, Users, Puzzle, Blocks, ChevronDown, Search } from 'lucide-react'; +import { Plus, FolderOpen, FolderPlus, History, Check, User, Users, Puzzle, Blocks, ChevronDown, Search, PanelsTopLeft } from 'lucide-react'; import { Tooltip } from '@/component-library'; import { useApp } from '../../hooks/useApp'; import { useSceneManager } from '../../hooks/useSceneManager'; @@ -75,6 +75,7 @@ const MainNav: React.FC = ({ const activeTabId = useSceneStore(s => s.activeTabId); const setSelectedAssistantWorkspaceId = useMyAgentStore((s) => s.setSelectedAssistantWorkspaceId); const { t } = useI18n('common'); + const { t: tPages } = useI18n('scenes/pages'); const { currentWorkspace, loading: workspaceLoading, @@ -680,6 +681,16 @@ const MainNav: React.FC = ({ {/* ── Bottom: MiniApp ───────────────────────── */}
+
{ + try { + return await remoteConnectAPI.accountCancelPendingLogin(pendingLoginId); + } catch (firstError) { + log.warn('pending login cancel response was ambiguous; retrying', firstError); + return await remoteConnectAPI.accountCancelPendingLogin(pendingLoginId); + } +} + function syncPhaseLabel( t: (key: string, options?: Record) => string, phase: AccountSyncPhase, @@ -106,9 +121,11 @@ export const AccountPanel: React.FC = ({ const { enterPeerMode } = usePeerDeviceMode(); const syncStatus = useAccountSyncStore((s) => s.status); const syncProgress = useAccountSyncStore((s) => s.progress); + const lastSyncIsFirstLogin = useAccountSyncStore((s) => s.lastSyncIsFirstLogin); const setSyncing = useAccountSyncStore((s) => s.setSyncing); const setSyncDone = useAccountSyncStore((s) => s.setDone); const setSyncFailed = useAccountSyncStore((s) => s.setFailed); + const clearSync = useAccountSyncStore((s) => s.clear); const [username, setUsername] = useState(''); const [password, setPassword] = useState(''); @@ -127,13 +144,32 @@ export const AccountPanel: React.FC = ({ /** Relay URL of the current account session, shown in the devices view. */ const [accountRelayUrl, setAccountRelayUrl] = useState(''); const [copiedServerUrl, setCopiedServerUrl] = useState(false); + /** Account epoch whose presence events may update the device list. */ + const [activeAccountEpoch, setActiveAccountEpoch] = useState(null); const refreshTimer = useRef | null>(null); + /** Reject late responses after unmount or an account login/logout transition. */ + const mountedRef = useRef(false); + const accountEpochRef = useRef(0); + const refreshRequestRef = useRef(0); /** Prevent overlapping background syncs from rapid clicks. */ const syncInFlightRef = useRef(false); - /** Track the overwrite view for the unmount logout invariant. */ + /** Opaque backend owner ID for the memory-only overwrite decision. */ + const pendingLoginIdRef = useRef(null); + /** Track the overwrite view for conditional unmount cleanup. */ const viewRef = useRef(view); viewRef.current = view; + const invalidateAccountRequests = useCallback(() => { + accountEpochRef.current += 1; + refreshRequestRef.current += 1; + setActiveAccountEpoch(null); + return accountEpochRef.current; + }, []); + + const isAccountEpochCurrent = useCallback((epoch: number) => ( + mountedRef.current && accountEpochRef.current === epoch + ), []); + const sortedDevices = useMemo(() => [...devices].sort((left, right) => { const leftLocal = left.device_id === localDeviceId; const rightLocal = right.device_id === localDeviceId; @@ -143,6 +179,7 @@ export const AccountPanel: React.FC = ({ }), [devices, localDeviceId]); const resetState = useCallback(() => { + setActiveAccountEpoch(null); setDevices([]); setLocalDeviceId(null); setDevicesReady(false); @@ -154,25 +191,29 @@ export const AccountPanel: React.FC = ({ const handleCopyRelayUrl = useCallback(async () => { if (!accountRelayUrl) return; - try { - await navigator.clipboard.writeText(accountRelayUrl); + const copied = await copyTextToClipboard(accountRelayUrl); + if (copied) { setCopiedServerUrl(true); window.setTimeout(() => setCopiedServerUrl(false), 1500); - } catch (e) { - log.warn('copy relay url failed', e); + } else { + warning(t('accountLogin.copyServerFailed')); } - }, [accountRelayUrl]); + }, [accountRelayUrl, t, warning]); - const handleSessionExpired = useCallback(async (_error: unknown) => { - try { - await remoteConnectAPI.accountLogout(); - } catch (e) { - log.warn('logout after session expiry failed', e); - } + const handleSessionExpired = useCallback(async (_error: unknown, expectedEpoch: number) => { + if (!isAccountEpochCurrent(expectedEpoch)) return; + invalidateAccountRequests(); + // Invalidate detached retries before the logout request yields control. + syncInFlightRef.current = false; + pendingLoginIdRef.current = null; + clearSync(); + // Authenticated backend commands invalidate only the generation/token that + // produced their 401. Do not issue a second unconditional logout here: a + // late frontend response must never clear a newer login. resetState(); setView('login'); setError(t('accountLogin.sessionExpired')); - }, [resetState, t]); + }, [clearSync, invalidateAccountRequests, isAccountEpochCurrent, resetState, t]); const markRelayUnreachable = useCallback(() => { setDevicesReady(false); @@ -180,43 +221,55 @@ export const AccountPanel: React.FC = ({ }, [t]); const refreshDevices = useCallback(async () => { + const epoch = accountEpochRef.current; + const requestId = ++refreshRequestRef.current; + const isCurrent = () => ( + isAccountEpochCurrent(epoch) && refreshRequestRef.current === requestId + ); try { let list = await remoteConnectAPI.accountListDevices(); + if (!isCurrent()) return; const localOffline = list.some(d => d.device_id === localDeviceId && !d.online); if (localOffline && localDeviceId) { await new Promise(r => setTimeout(r, 1500)); + if (!isCurrent()) return; list = await remoteConnectAPI.accountListDevices(); + if (!isCurrent()) return; } setDevices(list); setDevicesReady(true); setRelayError(null); } catch (e) { + if (!isCurrent()) return; log.warn('refreshDevices failed', e); if (isAccountAuthFailure(e)) { - await handleSessionExpired(e); + await handleSessionExpired(e, epoch); } else { markRelayUnreachable(); } } - }, [localDeviceId, handleSessionExpired, markRelayUnreachable]); + }, [localDeviceId, handleSessionExpired, isAccountEpochCurrent, markRelayUnreachable]); const handleRetryConnect = useCallback(async () => { + const epoch = accountEpochRef.current; setLoading(true); setRelayError(null); try { await remoteConnectAPI.accountConnectDevices(); + if (!isAccountEpochCurrent(epoch)) return; await refreshDevices(); } catch (err) { log.warn('retry connect failed', err); + if (!isAccountEpochCurrent(epoch)) return; if (isAccountAuthFailure(err)) { - await handleSessionExpired(err); + await handleSessionExpired(err, epoch); return; } markRelayUnreachable(); } finally { - setLoading(false); + if (isAccountEpochCurrent(epoch)) setLoading(false); } - }, [handleSessionExpired, markRelayUnreachable, refreshDevices]); + }, [handleSessionExpired, isAccountEpochCurrent, markRelayUnreachable, refreshDevices]); const applyPresenceOnline = useCallback((onlineDevices: Array<{ device_id: string; device_name: string }>) => { const onlineIds = new Set(onlineDevices.map(d => d.device_id)); @@ -260,29 +313,40 @@ export const AccountPanel: React.FC = ({ /** Connect presence + load the device list for an active account session. */ const initializeDevices = useCallback(async () => { + const epoch = accountEpochRef.current; try { await remoteConnectAPI.accountConnectDevices(); + if (!isAccountEpochCurrent(epoch)) return; // Re-read after AuthOk may have adopted the account-bound device_id. try { const info = await remoteConnectAPI.getDeviceInfo(); + if (!isAccountEpochCurrent(epoch)) return; setLocalDeviceId(info.device_id); } catch (e) { log.warn('getDeviceInfo after connect failed', e); } } catch (err) { + if (!isAccountEpochCurrent(epoch)) return; log.warn('accountConnectDevices failed', err); if (isAccountAuthFailure(err)) { - await handleSessionExpired(err); + await handleSessionExpired(err, epoch); return; } markRelayUnreachable(); } + if (!isAccountEpochCurrent(epoch)) return; void refreshDevices(); startDevicePolling(); - }, [handleSessionExpired, markRelayUnreachable, refreshDevices, startDevicePolling]); + }, [handleSessionExpired, isAccountEpochCurrent, markRelayUnreachable, refreshDevices, startDevicePolling]); useEffect(() => { + mountedRef.current = true; ensureAccountSyncProgressListener(); + return () => { + mountedRef.current = false; + accountEpochRef.current += 1; + refreshRequestRef.current += 1; + }; }, []); // Unmounting (dialog close or group switch) during the sync-choice step @@ -292,45 +356,74 @@ export const AccountPanel: React.FC = ({ useEffect(() => { return () => { if (viewRef.current === 'overwrite') { - void remoteConnectAPI.accountLogout().catch((e) => { - log.warn('logout on overwrite abandon failed', e); - }); + syncInFlightRef.current = false; + clearSync(); + const pendingLoginId = pendingLoginIdRef.current; + if (pendingLoginId) { + void cancelPendingLoginWithRetry(pendingLoginId) + .then(() => { + if (pendingLoginIdRef.current === pendingLoginId) { + pendingLoginIdRef.current = null; + } + }) + .catch((e) => { + log.warn('pending login cancel on overwrite abandon failed', e); + }); + } } }; - }, []); + }, [clearSync]); useEffect(() => { + const epoch = accountEpochRef.current; remoteConnectAPI.getDeviceInfo().then((info) => { - setLocalDeviceId(info.device_id); + if (isAccountEpochCurrent(epoch)) setLocalDeviceId(info.device_id); }).catch((e) => { log.warn('getDeviceInfo failed', e); }); remoteConnectAPI.accountGetCredentialHint().then((hint: AccountHint | null) => { - if (hint) { setUsername(hint.username); setAuthServer(hint.relay_url); setAccountRelayUrl(hint.relay_url); } + if (hint && isAccountEpochCurrent(epoch)) { + setUsername(hint.username); + setAuthServer(hint.relay_url); + setAccountRelayUrl(hint.relay_url); + } }); remoteConnectAPI.accountStatus().then(async (status) => { - if (status.logged_in && status.user_id) { + if (isAccountEpochCurrent(epoch) && status.logged_in && status.user_id) { + setActiveAccountEpoch(epoch); setView('devices'); await initializeDevices(); } + }).catch((e) => { + // A failed status probe must not synthesize a logged-out transition. + log.warn('account status initialization failed', e); }); - const unlistenPresence = api.listen<{ devices: Array<{ device_id: string; device_name: string }> }>( - 'account://device-presence', - (payload) => { - if (payload?.devices) { - applyPresenceOnline(payload.devices); - } - }, - ); - return () => { if (refreshTimer.current) { clearInterval(refreshTimer.current); refreshTimer.current = null; } - unlistenPresence(); }; }, [ - applyPresenceOnline, initializeDevices, + isAccountEpochCurrent, ]); + // Subscribe only while a specific account epoch is active. The callback + // captures that epoch; invalidation flips the ref synchronously, so an old + // listener cannot update the next account before React runs its cleanup. + useEffect(() => { + if (activeAccountEpoch === null) return undefined; + const subscribedEpoch = activeAccountEpoch; + const unlistenPresence = api.listen<{ + devices: Array<{ device_id: string; device_name: string }>; + }>( + 'account://device-presence', + (payload) => { + if (isAccountEpochCurrent(subscribedEpoch) && payload?.devices) { + applyPresenceOnline(payload.devices); + } + }, + ); + return unlistenPresence; + }, [activeAccountEpoch, applyPresenceOnline, isAccountEpochCurrent]); + const validate = useCallback(() => { if (!username.trim() || !password || !authServer.trim()) { setError(t('accountLogin.emptyFields')); @@ -359,7 +452,11 @@ export const AccountPanel: React.FC = ({ } syncInFlightRef.current = true; ensureAccountSyncProgressListener(); - setSyncing(); + setSyncing(isFirstLogin); + const operationId = useAccountSyncStore.getState().operationId; + const isCurrentOperation = () => ( + useAccountSyncStore.getState().operationId === operationId + ); info(t('accountLogin.syncStarted')); // Connect device presence immediately so the device list can populate @@ -373,6 +470,7 @@ export const AccountPanel: React.FC = ({ let configJson = '{}'; if (isFirstLogin) { useAccountSyncStore.getState().applyProgress({ + operation_id: operationId, phase: 'uploading_settings', percent: 2, }); @@ -382,22 +480,32 @@ export const AccountPanel: React.FC = ({ } catch (e) { log.warn('export config failed', e); } + if (!isCurrentOperation()) return; } const wp = workspacePath || '/'; const maxAttempts = 3; let result: Awaited> | null = null; let lastError: unknown = null; for (let attempt = 1; attempt <= maxAttempts; attempt += 1) { + if (!isCurrentOperation()) return; try { - result = await remoteConnectAPI.accountAutoSync(isFirstLogin, wp, configJson); + result = await remoteConnectAPI.accountAutoSync( + isFirstLogin, + wp, + configJson, + operationId, + ); + if (!isCurrentOperation()) return; lastError = null; break; } catch (e) { + if (!isCurrentOperation()) return; lastError = e; log.warn(`Auto-sync attempt ${attempt}/${maxAttempts} failed`, e); if (attempt < maxAttempts) { info(t('accountLogin.syncRetrying', { attempt, max: maxAttempts })); await new Promise((resolve) => setTimeout(resolve, 2000 * attempt)); + if (!isCurrentOperation()) return; } } } @@ -406,28 +514,35 @@ export const AccountPanel: React.FC = ({ ? lastError : new Error(String(lastError ?? 'auto-sync failed')); } + if (!isCurrentOperation()) return; log.info( `Auto-sync done: settings=${result.settings_synced} exported=${result.sessions_exported}`, ); if (result.settings_synced && !isFirstLogin) { + if (!isCurrentOperation()) return; try { await configAPI.reloadConfig(); + if (!isCurrentOperation()) return; configManager.clearCache(); success(t('accountLogin.settingsApplied')); } catch (e) { log.warn('reloadConfig after sync failed', e); } } + if (!isCurrentOperation()) return; setSyncDone(result); success(t('accountLogin.syncDone', { exported: result.sessions_exported, })); } catch (e) { + if (!isCurrentOperation()) return; log.error('Auto-sync failed', e); setSyncFailed(e instanceof Error ? e.message : String(e)); warning(t('accountLogin.syncFailed')); } finally { - syncInFlightRef.current = false; + if (isCurrentOperation()) { + syncInFlightRef.current = false; + } } })(); }, [ @@ -441,39 +556,87 @@ export const AccountPanel: React.FC = ({ workspacePath, ]); + const handleRetrySync = useCallback(() => { + if (syncStatus !== 'failed' || syncInFlightRef.current) return; + startBackgroundSync(lastSyncIsFirstLogin ?? false); + }, [lastSyncIsFirstLogin, startBackgroundSync, syncStatus]); + /** Landing path after a completed login: devices view + background sync. */ - const completeLogin = useCallback((relayUrl: string, isFirstLogin: boolean) => { + const completeLogin = useCallback(( + relayUrl: string, + isFirstLogin: boolean, + accountEpoch: number, + ) => { + if (!isAccountEpochCurrent(accountEpoch)) return; + setActiveAccountEpoch(accountEpoch); setAccountRelayUrl(relayUrl); setView('devices'); void initializeDevices(); startBackgroundSync(isFirstLogin); - }, [initializeDevices, startBackgroundSync]); + }, [initializeDevices, isAccountEpochCurrent, startBackgroundSync]); const performLogin = useCallback(async (server: string, user: string, pass: string) => { + const epoch = invalidateAccountRequests(); + // Invalidate detached sync retries before the backend begins replacing the + // account. The store operation id fences any completion from the old run. + syncInFlightRef.current = false; + clearSync(); setLoading(true); setError(null); try { + const stalePendingLoginId = pendingLoginIdRef.current; + if (stalePendingLoginId) { + await cancelPendingLoginWithRetry(stalePendingLoginId); + if (pendingLoginIdRef.current === stalePendingLoginId) { + pendingLoginIdRef.current = null; + } + if (!isAccountEpochCurrent(epoch)) return; + } const result = await remoteConnectAPI.accountLogin(server, user, pass); + if (!isAccountEpochCurrent(epoch)) { + if (result.pending_login_id) { + await cancelPendingLoginWithRetry(result.pending_login_id); + } + return; + } if (result.has_cloud_settings) { + if (!result.pending_login_id) { + throw new Error(t('accountLogin.sessionExpired')); + } + pendingLoginIdRef.current = result.pending_login_id; setView('overwrite'); setLoading(false); return; } success(t('accountLogin.loginSuccess', { user_id: result.user_id })); - completeLogin(server, true); + completeLogin(server, true, epoch); } catch (e: unknown) { + if (!isAccountEpochCurrent(epoch)) return; setError(e instanceof Error ? e.message : String(e)); - } finally { setLoading(false); } - }, [completeLogin, success, t]); + } finally { + // The account session has its own token after this call; retaining the + // password in React state while the device list is open is unnecessary. + if (isAccountEpochCurrent(epoch)) { + setPassword(''); + setLoading(false); + } + } + }, [clearSync, completeLogin, invalidateAccountRequests, isAccountEpochCurrent, success, t]); const handleLogin = useCallback(async () => { if (!validate()) return; const relayUrl = parseRelayServer(authServer); if (!relayUrl) return; const isLoopback = ['localhost', '127.0.0.1', '[::1]', '::1'].includes(relayUrl.hostname); - if (relayUrl.protocol === 'http:' - && !isLoopback - && !window.confirm(t('accountLogin.insecureServerConfirm'))) { - return; + if (relayUrl.protocol === 'http:' && !isLoopback) { + const confirmed = await confirmWarning( + t('accountLogin.insecureServerTitle'), + t('accountLogin.insecureServerConfirm'), + { + confirmText: t('accountLogin.continueInsecure'), + cancelText: t('accountLogin.cancel'), + }, + ); + if (!confirmed) return; } await performLogin(authServer.trim(), username.trim(), password); }, [validate, authServer, username, password, performLogin, t]); @@ -491,27 +654,60 @@ export const AccountPanel: React.FC = ({ }, [performLogin]); const finalizeAndSync = useCallback(async (isFirstLogin: boolean) => { + const epoch = accountEpochRef.current; + const pendingLoginId = pendingLoginIdRef.current; + if (!pendingLoginId) { + setError(t('accountLogin.sessionExpired')); + return; + } setLoading(true); setError(null); try { - await remoteConnectAPI.accountFinalizeLogin(); + try { + await remoteConnectAPI.accountFinalizeLogin(pendingLoginId); + } catch (firstError) { + if (!isAccountEpochCurrent(epoch)) return; + // The backend commit may have succeeded even when its transport + // response was lost. Retrying the same opaque owner is idempotent and + // cannot authorize a replacement account generation. + log.warn('pending login finalize response was ambiguous; retrying', firstError); + await remoteConnectAPI.accountFinalizeLogin(pendingLoginId); + } + if (!isAccountEpochCurrent(epoch)) return; + if (pendingLoginIdRef.current === pendingLoginId) { + pendingLoginIdRef.current = null; + } success(t('accountLogin.loginSuccess', { user_id: username })); - completeLogin(authServer.trim(), isFirstLogin); + completeLogin(authServer.trim(), isFirstLogin, epoch); } catch (e: unknown) { + if (!isAccountEpochCurrent(epoch)) return; if (isAccountAuthFailure(e)) { - await handleSessionExpired(e); - } else { - setError(e instanceof Error ? e.message : String(e)); + await handleSessionExpired(e, epoch); + return; } - try { await remoteConnectAPI.accountLogout(); } catch (logoutErr) { - log.warn('logout after finalize failure failed', logoutErr); + setError(e instanceof Error ? e.message : String(e)); + // Stop any detached work before accountLogout can yield. + syncInFlightRef.current = false; + clearSync(); + const cleanupEpoch = invalidateAccountRequests(); + try { + await cancelPendingLoginWithRetry(pendingLoginId); + if (pendingLoginIdRef.current === pendingLoginId) { + pendingLoginIdRef.current = null; + } + } catch (cancelErr) { + log.warn('pending login cancel after finalize failure failed', cancelErr); + if (isAccountEpochCurrent(cleanupEpoch)) setLoading(false); + return; } + if (!isAccountEpochCurrent(cleanupEpoch)) return; resetState(); setView('login'); - } finally { setLoading(false); + } finally { + if (isAccountEpochCurrent(epoch)) setLoading(false); } - }, [authServer, completeLogin, handleSessionExpired, resetState, success, t, username]); + }, [authServer, clearSync, completeLogin, handleSessionExpired, invalidateAccountRequests, isAccountEpochCurrent, resetState, success, t, username]); const handleConfirmOverwrite = useCallback(() => { void finalizeAndSync(false); @@ -522,32 +718,83 @@ export const AccountPanel: React.FC = ({ }, [finalizeAndSync]); const handleCancelOverwrite = useCallback(async () => { - try { await remoteConnectAPI.accountLogout(); } catch (e) { log.warn('logout failed', e); } + const epoch = invalidateAccountRequests(); + syncInFlightRef.current = false; + clearSync(); + const pendingLoginId = pendingLoginIdRef.current; + if (pendingLoginId) { + try { + await cancelPendingLoginWithRetry(pendingLoginId); + if (pendingLoginIdRef.current === pendingLoginId) { + pendingLoginIdRef.current = null; + } + } catch (e) { + log.warn('pending login cancel failed', e); + if (isAccountEpochCurrent(epoch)) { + setError(e instanceof Error ? e.message : String(e)); + } + return; + } + } + if (!isAccountEpochCurrent(epoch)) return; resetState(); setView('login'); - }, [resetState]); + }, [clearSync, invalidateAccountRequests, isAccountEpochCurrent, resetState]); const handleLogout = useCallback(async () => { + const epoch = invalidateAccountRequests(); setLoading(true); + syncInFlightRef.current = false; + clearSync(); + pendingLoginIdRef.current = null; try { await remoteConnectAPI.accountLogout(); + if (!isAccountEpochCurrent(epoch)) return; resetState(); setView('login'); } catch (e: unknown) { + if (!isAccountEpochCurrent(epoch)) return; + // Logout failed before the backend changed the account; resume presence + // delivery for the still-current frontend epoch. + setActiveAccountEpoch(epoch); setError(e instanceof Error ? e.message : String(e)); - } finally { setLoading(false); } - }, [resetState]); + } finally { + if (isAccountEpochCurrent(epoch)) setLoading(false); + } + }, [clearSync, invalidateAccountRequests, isAccountEpochCurrent, resetState]); const handleDeleteDevice = useCallback(async (deviceId: string, deviceName: string) => { const isLocal = localDeviceId === deviceId; const confirmation = isLocal ? t('accountLogin.confirmRemoveCurrentDevice', { name: deviceName }) : t('accountLogin.confirmRemoveDevice', { name: deviceName }); - if (!window.confirm(confirmation)) return; + const confirmed = await confirmDanger( + isLocal + ? t('accountLogin.removeCurrentDevice') + : t('accountLogin.removeDevice'), + confirmation, + { + confirmText: isLocal + ? t('accountLogin.removeCurrentDevice') + : t('accountLogin.removeDevice'), + cancelText: t('accountLogin.cancel'), + }, + ); + if (!confirmed) return; + const previousSyncStatus = syncStatus; + const previousSyncDirection = lastSyncIsFirstLogin; setLoading(true); setError(null); + const epoch = isLocal ? invalidateAccountRequests() : accountEpochRef.current; + if (isLocal) { + // A current-device removal is also a logout. Invalidate retries and + // late progress before the backend request yields. + syncInFlightRef.current = false; + clearSync(); + } try { await remoteConnectAPI.accountDeleteDevice(deviceId); + if (!isAccountEpochCurrent(epoch)) return; if (isLocal) { success(t('accountLogin.currentDeviceRemoved')); resetState(); @@ -557,15 +804,42 @@ export const AccountPanel: React.FC = ({ void refreshDevices(); } } catch (e: unknown) { + if (!isAccountEpochCurrent(epoch)) return; if (isAccountAuthFailure(e)) { - await handleSessionExpired(e); + await handleSessionExpired(e, epoch); } else { - setError(e instanceof Error ? e.message : String(e)); + const message = e instanceof Error ? e.message : String(e); + if (isLocal) setActiveAccountEpoch(epoch); + setError(message); + if ( + isLocal + && previousSyncDirection !== null + && (previousSyncStatus === 'syncing' || previousSyncStatus === 'failed') + ) { + // Preserve the direction so Retry remains meaningful after a failed + // current-device removal invalidated the previous generation. + setSyncing(previousSyncDirection); + setSyncFailed(message); + } } } finally { - setLoading(false); + if (isAccountEpochCurrent(epoch)) setLoading(false); } - }, [handleSessionExpired, localDeviceId, refreshDevices, resetState, success, t]); + }, [ + clearSync, + handleSessionExpired, + invalidateAccountRequests, + isAccountEpochCurrent, + lastSyncIsFirstLogin, + localDeviceId, + refreshDevices, + resetState, + setSyncFailed, + setSyncing, + success, + syncStatus, + t, + ]); const selectDevice = useCallback(async (device: AccountDeviceInfo) => { if (!device.online) return; @@ -574,6 +848,9 @@ export const AccountPanel: React.FC = ({ info(t('accountLogin.syncInProgressHint')); return; } + if (syncStatus === 'failed') { + warning(t('accountLogin.syncFailedPeerHint')); + } setLoading(true); setError(null); try { @@ -585,7 +862,7 @@ export const AccountPanel: React.FC = ({ } finally { setLoading(false); } - }, [enterPeerMode, info, localDeviceId, onCloseDialog, success, syncStatus, t]); + }, [enterPeerMode, info, localDeviceId, onCloseDialog, success, syncStatus, t, warning]); return ( <> @@ -731,6 +1008,17 @@ export const AccountPanel: React.FC = ({ {syncStatus === 'done' && t('accountLogin.syncDoneShort')} {syncStatus === 'failed' && t('accountLogin.syncFailed')} + {syncStatus === 'failed' && ( + + )} {syncStatus === 'syncing' && ( {t('accountLogin.syncProgressPercent', { percent: syncProgress.percent })} @@ -781,43 +1069,48 @@ export const AccountPanel: React.FC = ({ const displayName = d.device_name || t('accountLogin.unknownDevice'); return (
isSelectable && selectDevice(d)} - onKeyDown={(event) => { - if (isSelectable && (event.key === 'Enter' || event.key === ' ')) { - event.preventDefault(); - void selectDevice(d); - } - }} - role={isSelectable ? 'button' : undefined} - tabIndex={isSelectable ? 0 : undefined}> - -
- - {displayName} - {isLocal && {t('accountLogin.thisDevice')}} - - - - {d.device_id.slice(0, 8)} + className={`account-panel__device-card ${isSelectable ? 'selectable' : ''} ${d.online ? '' : 'offline'} ${isLocal ? 'current' : ''} ${syncStatus === 'syncing' && !isLocal ? 'syncing' : ''}`}> +
- {!isLocal && d.online && syncStatus !== 'syncing' && } - {!isLocal && d.online && syncStatus === 'syncing' && ( - - )} - +
); + const handleCopyPairingUrl = useCallback(async () => { + if (!connectionResult?.qr_url) return; + const copied = await copyTextToClipboard(connectionResult.qr_url); + if (copied) { + setQrCopied(true); + window.setTimeout(() => setQrCopied(false), 2000); + } else { + notifyError(t('remoteConnect.copyUrlFailed')); + } + }, [connectionResult?.qr_url, notifyError, t]); + const renderPairingInProgress = () => { if (!connectionResult) return null; return (
{connectionResult.qr_url && ( -
{ - navigator.clipboard.writeText(connectionResult.qr_url!); - setQrCopied(true); - setTimeout(() => setQrCopied(false), 2000); - }} + title={t('remoteConnect.copyUrl')} + aria-label={t('remoteConnect.copyUrl')} + onClick={() => void handleCopyPairingUrl()} > -
+ )} {connectionResult.bot_pairing_code && (
@@ -675,13 +888,17 @@ export const RemoteConnectDialog: React.FC = ({ {qrCopied ? t('remoteConnect.urlCopied') - : activeGroup === 'bot' + : connectionOwner === 'bot' ? t('remoteConnect.stateWaitingBot') : t('remoteConnect.stateWaiting')}

- {activeGroup === 'bot' ? t('remoteConnect.botHint') : t('remoteConnect.scanHint')} + {connectionOwner === 'bot' + ? t('remoteConnect.botHint') + : connectionResult.qr_url + ? t('remoteConnect.scanHint') + : t('remoteConnect.stateWaiting')}

); } - if (connectionResult && activeGroup === 'bot') { + if (connectionResult && connectionOwner === 'bot') { return renderPairingInProgress(); } return ( @@ -939,14 +1156,14 @@ export const RemoteConnectDialog: React.FC = ({ {isWeixinRasterQrSrc(weixinQrImageUrl) ? ( WeChat QR ) : (
= ({ // ── Layout ─────────────────────────────────────────────────────── - const isNetworkConnecting = !!connectionResult && activeGroup === 'network' && !isRelayConnected; - const isBotConnecting = !!connectionResult && activeGroup === 'bot' && !isBotConnected; + const isNetworkConnecting = !!connectionResult && connectionOwner === 'network' && !isRelayConnected; + const isBotConnecting = !!connectionResult && connectionOwner === 'bot' && !isBotConnected; const handleAgreeDisclaimer = useCallback(() => { setRemoteConnectDisclaimerAgreed(); setHasAgreedDisclaimer(true); setShowDisclaimer(false); }, []); + useEffect(() => { + if (isOpen && hasAgreedDisclaimer && !hasWorkspace && activeGroup !== 'account') { + handleGroupChange('account'); + } + }, [activeGroup, handleGroupChange, hasAgreedDisclaimer, hasWorkspace, isOpen]); + + const disclaimerIsGate = isOpen && !hasAgreedDisclaimer; + const handleDisclaimerClose = disclaimerIsGate + ? handleDialogClose + : () => setShowDisclaimer(false); + return ( <> @@ -1043,7 +1271,7 @@ export const RemoteConnectDialog: React.FC = ({ >
{/* ── Group tabs ── */} -
+
= ({ >
) : ( -
+
{BOT_TABS.map((tab, i) => ( {i > 0 && }
setShowDisclaimer(false)} + isOpen={isOpen && (disclaimerIsGate || showDisclaimer)} + onClose={handleDisclaimerClose} title={t('remoteConnect.disclaimerTitle')} showCloseButton size="large" @@ -1148,7 +1444,7 @@ export const RemoteConnectDialog: React.FC = ({ > setShowDisclaimer(false)} + onClose={handleDisclaimerClose} onAgree={hasAgreedDisclaimer ? undefined : handleAgreeDisclaimer} /> diff --git a/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.scss b/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.scss index 8390bf8a93..0d4ff570c7 100644 --- a/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.scss +++ b/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.scss @@ -20,6 +20,13 @@ color: var(--color-text-muted); } +.bitfun-remote-disclaimer__section-title { + margin: 2px 0 0; + color: var(--color-text-primary); + font-size: 13px; + font-weight: 600; +} + .bitfun-remote-disclaimer__list { margin: 0; padding-left: 20px; @@ -34,6 +41,34 @@ } } +.bitfun-remote-disclaimer__list--key { + gap: 7px; +} + +.bitfun-remote-disclaimer__details { + border: 1px solid var(--border-subtle); + border-radius: 6px; + background: var(--element-bg-subtle); + + summary { + padding: 9px 12px; + color: var(--color-text-primary); + font-size: 12px; + font-weight: 500; + cursor: pointer; + } + + &[open] summary { + border-bottom: 1px solid var(--border-subtle); + } + + .bitfun-remote-disclaimer__list { + max-height: 240px; + padding: 10px 14px 12px 32px; + overflow-y: auto; + } +} + .bitfun-remote-disclaimer__actions { display: flex; justify-content: center; diff --git a/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.tsx b/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.tsx index 2f205bee3d..a49f661cae 100644 --- a/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.tsx +++ b/src/web-ui/src/app/components/RemoteConnectDialog/RemoteConnectDisclaimer.tsx @@ -27,25 +27,34 @@ export const RemoteConnectDisclaimerContent: React.FC{t('remoteConnect.disclaimerIntro')}

-
    +

    + {t('remoteConnect.disclaimerKeyRisks')} +

    +
    1. {t('remoteConnect.disclaimerItemGeneralRisk')}
    2. {t('remoteConnect.disclaimerItemSecurity')}
    3. {t('remoteConnect.disclaimerItemEncryption')}
      4. -
    4. {t('remoteConnect.disclaimerItemOpenSource')}
    5. {t('remoteConnect.disclaimerItemPrivacy')}
      6. -
    6. {t('remoteConnect.disclaimerItemDataUsage')}
      7. -
    7. {t('remoteConnect.disclaimerItemCredentials')}
      8. -
    8. {t('remoteConnect.disclaimerItemQrCode')}
      9. -
    9. {t('remoteConnect.disclaimerItemNgrok')}
      10. -
    10. {t('remoteConnect.disclaimerItemSelfHosted')}
      11. -
    11. {t('remoteConnect.disclaimerItemNetwork')}
      12. -
    12. {t('remoteConnect.disclaimerItemBot')}
      13. -
    13. {t('remoteConnect.disclaimerItemBotPersistence')}
      14. -
    14. {t('remoteConnect.disclaimerItemMobileBrowser')}
      15. -
    15. {t('remoteConnect.disclaimerItemCompliance')}
      16. -
    16. {t('remoteConnect.disclaimerItemLiability')}
    +
    + {t('remoteConnect.disclaimerFullDetails')} +
      +
    1. {t('remoteConnect.disclaimerItemOpenSource')}
      2. +
    2. {t('remoteConnect.disclaimerItemDataUsage')}
      3. +
    3. {t('remoteConnect.disclaimerItemCredentials')}
      4. +
    4. {t('remoteConnect.disclaimerItemQrCode')}
      5. +
    5. {t('remoteConnect.disclaimerItemNgrok')}
      6. +
    6. {t('remoteConnect.disclaimerItemSelfHosted')}
      7. +
    7. {t('remoteConnect.disclaimerItemNetwork')}
      8. +
    8. {t('remoteConnect.disclaimerItemBot')}
      9. +
    9. {t('remoteConnect.disclaimerItemBotPersistence')}
      10. +
    10. {t('remoteConnect.disclaimerItemMobileBrowser')}
      11. +
    11. {t('remoteConnect.disclaimerItemCompliance')}
      12. +
    12. {t('remoteConnect.disclaimerItemLiability')}
      13. +
    +
    +
    + ), + Input: (props: React.InputHTMLAttributes) => , + Select: () =>
    , + confirmDanger: mocks.confirmDanger, + confirmWarning: vi.fn(), +})); + +vi.mock('@/app/components', () => ({ + GalleryLayout: ({ children }: { children: React.ReactNode }) =>
    {children}
    , + GalleryPageHeader: ({ + title, + actions, + }: { + title: React.ReactNode; + actions?: React.ReactNode; + }) =>
    {title}{actions}
    , + GalleryEmpty: ({ message, action, testId }: { message: React.ReactNode; action?: React.ReactNode; testId?: string }) => ( +
    {message}{action}
    + ), +})); + +describe('PagesScene initial loading', () => { + let container: HTMLDivElement; + let root: Root; + + beforeEach(() => { + container = document.createElement('div'); + document.body.appendChild(container); + root = createRoot(container); + mocks.accountStatus.mockReset().mockResolvedValue({ logged_in: true, user_id: 'u1' }); + mocks.accountGetCredentialHint.mockReset().mockResolvedValue({ relay_url: 'https://relay.test' }); + mocks.listPages.mockReset().mockRejectedValue(new Error('relay unavailable')); + mocks.listVersions.mockReset().mockResolvedValue([]); + mocks.createOpenLink.mockReset(); + mocks.update.mockReset(); + mocks.deletePage.mockReset().mockResolvedValue(undefined); + mocks.confirmDanger.mockReset().mockResolvedValue(true); + mocks.listen.mockReset().mockImplementation(() => vi.fn()); + mocks.openExternal.mockReset().mockResolvedValue(undefined); + }); + + afterEach(() => { + act(() => root.unmount()); + container.remove(); + }); + + it('attempts a failed initial load only once until the user retries', async () => { + await act(async () => { + root.render(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + await act(async () => { + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + + // The failed relay call triggers one bounded auth re-check so an expired + // session can switch to the sign-in state; it must not retry listPages. + expect(mocks.accountStatus).toHaveBeenCalledTimes(2); + expect(mocks.listPages).toHaveBeenCalledTimes(1); + expect(container.querySelector('[data-testid="pages-error"]')).not.toBeNull(); + }); + + it('locks every action on one Page while an operation is pending and exposes title editing', async () => { + mocks.listPages.mockResolvedValue([{ + slug: 'demo', + generation: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + visibility: 'public', + title: 'Demo', + file_count: 1, + total_bytes: 20, + created_at: 1, + updated_at: 1, + url_path: '/p/alice/demo', + preview_url_path: '/p/alice/demo/@v/v1', + deployed_version_id: 'v1', + }]); + let resolveOpenLink: ((value: { open_url: string; expires_in_seconds: number }) => void) | undefined; + mocks.createOpenLink.mockImplementation(() => new Promise((resolve) => { + resolveOpenLink = resolve; + })); + + await act(async () => { + root.render(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + + expect(container.querySelector('input[aria-label="titleField.inputAria"]')).not.toBeNull(); + const buttons = [...container.querySelectorAll('button')]; + const open = buttons.find((button) => button.textContent?.includes('actions.openProduction')); + const remove = buttons.find((button) => button.textContent?.includes('actions.deletePage')); + expect(open).toBeDefined(); + expect(remove).toBeDefined(); + + await act(async () => { + open?.click(); + await Promise.resolve(); + }); + expect(remove?.disabled).toBe(true); + + await act(async () => { + resolveOpenLink?.({ open_url: 'https://relay.test/open', expires_in_seconds: 60 }); + await Promise.resolve(); + }); + expect(remove?.disabled).toBe(false); + }); + + it('does not restore a deleted Page from an older refresh response', async () => { + const page = { + slug: 'demo', + generation: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + visibility: 'public', + title: 'Demo', + file_count: 1, + total_bytes: 20, + created_at: 1, + updated_at: 1, + url_path: '/p/alice/demo', + preview_url_path: '/p/alice/demo/@v/v1', + deployed_version_id: 'v1', + }; + mocks.listPages.mockResolvedValueOnce([page]); + + await act(async () => { + root.render(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + + let resolveRefresh: ((pages: typeof page[]) => void) | undefined; + mocks.listPages.mockImplementationOnce(() => new Promise((resolve) => { + resolveRefresh = resolve; + })); + const refresh = [...container.querySelectorAll('button')] + .find((button) => button.textContent === 'actions.refresh'); + await act(async () => { + refresh?.click(); + await Promise.resolve(); + }); + + const remove = [...container.querySelectorAll('button')] + .find((button) => button.textContent?.includes('actions.deletePage')); + await act(async () => { + remove?.click(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(container.textContent).not.toContain('Demo'); + + await act(async () => { + resolveRefresh?.([page]); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(container.textContent).not.toContain('Demo'); + }); + + it('drops slug caches when the same account recreates a Page with a new generation', async () => { + const oldPage = { + slug: 'recreated', + generation: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + visibility: 'public' as const, + title: 'Old Page', + file_count: 1, + total_bytes: 20, + created_at: 1, + updated_at: 1, + url_path: '/p/alice/recreated', + preview_url_path: '/p/alice/recreated/@v/a1', + deployed_version_id: 'a1', + }; + const newPage = { + ...oldPage, + generation: 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb', + title: 'New Page', + preview_url_path: '/p/alice/recreated/@v/b1', + deployed_version_id: 'b1', + }; + mocks.listPages.mockResolvedValueOnce([oldPage]); + mocks.listVersions + .mockResolvedValueOnce([{ + generation: oldPage.generation, + version_id: 'a1', + title: oldPage.title, + file_count: 1, + total_bytes: 20, + has_worker: false, + note: 'old generation note', + created_at: 1, + deployed: true, + preview_url_path: oldPage.preview_url_path, + }]) + .mockResolvedValueOnce([{ + generation: newPage.generation, + version_id: 'b1', + title: newPage.title, + file_count: 1, + total_bytes: 20, + has_worker: false, + note: 'new generation note', + created_at: 2, + deployed: true, + preview_url_path: newPage.preview_url_path, + }]); + + await act(async () => { + root.render(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + const oldInput = container.querySelector('input') as HTMLInputElement; + const valueSetter = Object.getOwnPropertyDescriptor( + HTMLInputElement.prototype, + 'value', + )?.set; + await act(async () => { + valueSetter?.call(oldInput, 'old generation draft'); + oldInput.dispatchEvent(new Event('input', { bubbles: true })); + await Promise.resolve(); + }); + expect(oldInput.value).toBe('old generation draft'); + + const oldVersions = [...container.querySelectorAll('button')] + .find((button) => button.textContent?.includes('actions.versions')); + await act(async () => { + oldVersions?.click(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(container.textContent).toContain('old generation note'); + + mocks.listPages.mockResolvedValueOnce([newPage]); + const refresh = [...container.querySelectorAll('button')] + .find((button) => button.textContent === 'actions.refresh'); + await act(async () => { + refresh?.click(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(container.textContent).toContain('New Page'); + expect(container.textContent).not.toContain('old generation note'); + expect((container.querySelector('input') as HTMLInputElement).value).toBe('New Page'); + + const newVersions = [...container.querySelectorAll('button')] + .find((button) => button.textContent?.includes('actions.versions')); + await act(async () => { + newVersions?.click(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(mocks.listVersions).toHaveBeenLastCalledWith( + newPage.slug, + newPage.generation, + ); + expect(container.textContent).toContain('new generation note'); + expect(container.textContent).not.toContain('old generation note'); + }); + + it('clears account-owned state immediately and fences stale same-slug actions', async () => { + const pageA = { + slug: 'shared', + generation: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + visibility: 'public' as const, + title: 'Account A Page', + file_count: 1, + total_bytes: 20, + created_at: 1, + updated_at: 1, + url_path: '/p/alice/shared', + preview_url_path: '/p/alice/shared/@v/a1', + deployed_version_id: 'a1', + }; + const pageB = { + ...pageA, + generation: 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb', + title: 'Account B Page', + url_path: '/p/bob/shared', + preview_url_path: '/p/bob/shared/@v/b1', + deployed_version_id: 'b1', + }; + mocks.listPages.mockResolvedValueOnce([pageA]); + mocks.listVersions.mockResolvedValueOnce([{ + generation: pageA.generation, + version_id: 'a1', + title: pageA.title, + file_count: 1, + total_bytes: 20, + has_worker: false, + note: 'A-only note', + created_at: 1, + deployed: true, + preview_url_path: pageA.preview_url_path, + }]); + + await act(async () => { + root.render(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + const versions = [...container.querySelectorAll('button')] + .find((button) => button.textContent?.includes('actions.versions')); + await act(async () => { + versions?.click(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(container.textContent).toContain('A-only note'); + + let resolveConfirmation: ((confirmed: boolean) => void) | undefined; + mocks.confirmDanger.mockImplementationOnce(() => new Promise((resolve) => { + resolveConfirmation = resolve; + })); + const staleDelete = [...container.querySelectorAll('button')] + .find((button) => button.textContent?.includes('actions.deletePage')); + await act(async () => { + staleDelete?.click(); + await Promise.resolve(); + }); + + let resolveBPages: ((pages: typeof pageB[]) => void) | undefined; + mocks.accountStatus.mockResolvedValue({ logged_in: true, user_id: 'u2' }); + mocks.listPages.mockImplementationOnce(() => new Promise((resolve) => { + resolveBPages = resolve; + })); + const loginStateListener = mocks.listen.mock.calls + .find(([event]) => event === 'account://login-state')?.[1] as + ((payload: { logged_in: boolean }) => void) | undefined; + expect(loginStateListener).toBeDefined(); + await act(async () => { + loginStateListener?.({ logged_in: true }); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + + // B ownership is adopted before its list response arrives, so no A-only + // versions, drafts, or cards remain actionable during the gap. + expect(container.textContent).not.toContain('Account A Page'); + expect(container.textContent).not.toContain('A-only note'); + + await act(async () => { + resolveConfirmation?.(true); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(mocks.deletePage).not.toHaveBeenCalled(); + + await act(async () => { + resolveBPages?.([pageB]); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(container.textContent).toContain('Account B Page'); + expect(container.textContent).not.toContain('A-only note'); + expect((container.querySelector('input') as HTMLInputElement | null)?.value) + .toBe('Account B Page'); + }); + + it('invalidates in-flight actions when the same user logs in again', async () => { + const oldPage = { + slug: 'same-user', + generation: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + visibility: 'public' as const, + title: 'Old session Page', + file_count: 1, + total_bytes: 20, + created_at: 1, + updated_at: 1, + url_path: '/p/alice/same-user', + preview_url_path: '/p/alice/same-user/@v/a1', + deployed_version_id: 'a1', + }; + const freshPage = { + ...oldPage, + generation: 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb', + title: 'Fresh session Page', + preview_url_path: '/p/alice/same-user/@v/b1', + deployed_version_id: 'b1', + }; + mocks.listPages.mockResolvedValueOnce([oldPage]); + let resolveStaleOpen: ((value: { open_url: string; expires_in_seconds: number }) => void) + | undefined; + mocks.createOpenLink.mockImplementationOnce(() => new Promise((resolve) => { + resolveStaleOpen = resolve; + })); + + await act(async () => { + root.render(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + const open = [...container.querySelectorAll('button')] + .find((button) => button.textContent?.includes('actions.openProduction')); + await act(async () => { + open?.click(); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(mocks.createOpenLink).toHaveBeenCalledWith( + oldPage.slug, + oldPage.generation, + undefined, + ); + + mocks.listPages.mockResolvedValueOnce([freshPage]); + const loginStateListener = mocks.listen.mock.calls + .find(([event]) => event === 'account://login-state')?.[1] as + ((payload: { logged_in: boolean }) => void) | undefined; + expect(loginStateListener).toBeDefined(); + await act(async () => { + loginStateListener?.({ logged_in: true }); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(container.textContent).toContain('Fresh session Page'); + expect(container.textContent).not.toContain('Old session Page'); + + await act(async () => { + resolveStaleOpen?.({ + open_url: 'https://relay.test/stale-open', + expires_in_seconds: 60, + }); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(mocks.openExternal).not.toHaveBeenCalled(); + expect(container.textContent).toContain('Fresh session Page'); + }); +}); diff --git a/src/web-ui/src/app/scenes/pages/PagesScene.tsx b/src/web-ui/src/app/scenes/pages/PagesScene.tsx new file mode 100644 index 0000000000..709d92b0a0 --- /dev/null +++ b/src/web-ui/src/app/scenes/pages/PagesScene.tsx @@ -0,0 +1,933 @@ +import React, { Suspense, lazy, useCallback, useEffect, useMemo, useRef, useState } from 'react'; +import { + ChevronDown, + ChevronUp, + Copy, + ExternalLink, + FileClock, + PanelsTopLeft, + RefreshCw, + Rocket, + Save, + Trash2, +} from 'lucide-react'; +import { + Button, + Input, + Select, + confirmDanger, + confirmWarning, + type SelectOption, +} from '@/component-library'; +import { GalleryEmpty, GalleryLayout, GalleryPageHeader } from '@/app/components'; +import { + pageAPI, + type PageInfo, + type PageVersionInfo, + type PageVisibility, +} from '@/infrastructure/api/service-api/PageAPI'; +import { api } from '@/infrastructure/api/service-api/ApiClient'; +import { remoteConnectAPI } from '@/infrastructure/api/service-api/RemoteConnectAPI'; +import { systemAPI } from '@/infrastructure/api/service-api/SystemAPI'; +import { useI18n } from '@/infrastructure/i18n'; +import { useNotification } from '@/shared/notification-system'; +import { createLogger } from '@/shared/utils/logger'; +import './PagesScene.scss'; + +const log = createLogger('PagesScene'); +const RemoteConnectDialog = lazy(() => import('@/app/components/RemoteConnectDialog')); + +interface PagesSceneProps { + isActive?: boolean; +} + +interface PageOwner { + userId: string; + epoch: number; +} + +interface PageActionLease { + slug: string; + key: string; + token: string; + userId: string; + ownerEpoch: number; +} + +function errorText(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function replacePage(pages: PageInfo[], updated: PageInfo): PageInfo[] { + return pages.map((page) => (page.slug === updated.slug ? updated : page)); +} + +const PagesScene: React.FC = ({ isActive = true }) => { + const { t, formatDate, formatNumber } = useI18n('scenes/pages'); + const notification = useNotification(); + const attemptedLoadRef = useRef(false); + const pageLoadEpochRef = useRef(0); + const pageOwnerRef = useRef(null); + const pageOwnerEpochCounterRef = useRef(0); + const nextActionTokenRef = useRef(0); + const [pageOwnerEpoch, setPageOwnerEpoch] = useState(0); + const pagesRef = useRef([]); + const [pages, setPages] = useState([]); + const [relayBaseUrl, setRelayBaseUrl] = useState(''); + const [versionsBySlug, setVersionsBySlug] = useState>({}); + const [expandedSlugs, setExpandedSlugs] = useState>(() => new Set()); + const [loading, setLoading] = useState(false); + const [loadError, setLoadError] = useState(''); + const [loginRequired, setLoginRequired] = useState(false); + const [showAccountDialog, setShowAccountDialog] = useState(false); + const [pendingBySlug, setPendingBySlug] = useState>({}); + const busySlugsRef = useRef>(new Map()); + const [titleDrafts, setTitleDrafts] = useState>({}); + + const cancelPendingPageLoad = useCallback(() => { + pageLoadEpochRef.current += 1; + setLoading(false); + }, []); + + const adoptPageOwner = useCallback(( + userId: string | null, + cancelLoads: boolean, + forceNewEpoch = false, + ): number => { + const current = pageOwnerRef.current; + if (!forceNewEpoch && (current?.userId ?? null) === userId) { + return current?.epoch ?? pageOwnerEpochCounterRef.current; + } + const epoch = pageOwnerEpochCounterRef.current + 1; + pageOwnerEpochCounterRef.current = epoch; + pageOwnerRef.current = userId ? { userId, epoch } : null; + setPageOwnerEpoch(epoch); + if (cancelLoads) pageLoadEpochRef.current += 1; + busySlugsRef.current.clear(); + pagesRef.current = []; + setPages([]); + setRelayBaseUrl(''); + setVersionsBySlug({}); + setExpandedSlugs(new Set()); + setTitleDrafts({}); + setPendingBySlug({}); + setLoadError(''); + setLoading(false); + setLoginRequired(userId === null); + return epoch; + }, []); + + const updateOwnedPages = useCallback(( + update: (current: PageInfo[]) => PageInfo[], + ) => { + const next = update(pagesRef.current); + pagesRef.current = next; + setPages(next); + }, []); + + const commitLoadedPages = useCallback((nextPages: PageInfo[]) => { + const previousGenerationBySlug = new Map( + pagesRef.current.map((page) => [page.slug, page.generation]), + ); + const nextGenerationBySlug = new Map( + nextPages.map((page) => [page.slug, page.generation]), + ); + const canRetainSlugState = (slug: string) => { + const previousGeneration = previousGenerationBySlug.get(slug); + return previousGeneration !== undefined + && previousGeneration === nextGenerationBySlug.get(slug); + }; + + for (const slug of busySlugsRef.current.keys()) { + if (!canRetainSlugState(slug)) busySlugsRef.current.delete(slug); + } + setVersionsBySlug((current) => Object.fromEntries( + Object.entries(current).filter(([slug]) => canRetainSlugState(slug)), + )); + setExpandedSlugs((current) => new Set( + [...current].filter((slug) => canRetainSlugState(slug)), + )); + setTitleDrafts((current) => Object.fromEntries( + Object.entries(current).filter(([slug]) => canRetainSlugState(slug)), + )); + setPendingBySlug((current) => Object.fromEntries( + Object.entries(current).filter(([slug]) => canRetainSlugState(slug)), + )); + pagesRef.current = nextPages; + setPages(nextPages); + }, []); + + const isPageActionCurrent = useCallback((lease: PageActionLease): boolean => { + const owner = pageOwnerRef.current; + return owner?.userId === lease.userId + && owner.epoch === lease.ownerEpoch + && busySlugsRef.current.get(lease.slug) === lease.token; + }, []); + + const endPageAction = useCallback((lease: PageActionLease) => { + if (busySlugsRef.current.get(lease.slug) !== lease.token) return; + busySlugsRef.current.delete(lease.slug); + setPendingBySlug((current) => { + if (current[lease.slug] !== lease.key) return current; + const next = { ...current }; + delete next[lease.slug]; + return next; + }); + }, []); + + const beginPageAction = useCallback(async ( + page: PageInfo, + key: string, + expectedOwnerEpoch: number, + ): Promise => { + const owner = pageOwnerRef.current; + if (!owner || owner.epoch !== expectedOwnerEpoch || busySlugsRef.current.has(page.slug)) { + return null; + } + const token = `${owner.epoch}:${nextActionTokenRef.current += 1}`; + const lease: PageActionLease = { + slug: page.slug, + key, + token, + userId: owner.userId, + ownerEpoch: owner.epoch, + }; + // A list response captured before this mutation must never overwrite the + // operation's newer result when it eventually arrives. + cancelPendingPageLoad(); + busySlugsRef.current.set(page.slug, token); + setPendingBySlug((current) => ({ ...current, [page.slug]: key })); + const status = await remoteConnectAPI.accountStatus().catch(() => null); + if (!status?.logged_in || status.user_id !== owner.userId || !isPageActionCurrent(lease)) { + if (status) { + adoptPageOwner(status.logged_in ? status.user_id : null, true); + attemptedLoadRef.current = !status.logged_in; + } + endPageAction(lease); + return null; + } + return lease; + }, [adoptPageOwner, cancelPendingPageLoad, endPageAction, isPageActionCurrent]); + + const validatePageAction = useCallback(async (lease: PageActionLease): Promise => { + if (!isPageActionCurrent(lease)) return false; + const status = await remoteConnectAPI.accountStatus().catch(() => null); + if (!status?.logged_in || status.user_id !== lease.userId || !isPageActionCurrent(lease)) { + if (status) { + adoptPageOwner(status.logged_in ? status.user_id : null, true); + attemptedLoadRef.current = !status.logged_in; + } + return false; + } + return true; + }, [adoptPageOwner, isPageActionCurrent]); + + const visibilityOptions = useMemo(() => [ + { value: 'private', label: t('visibility.private') }, + { value: 'relay', label: t('visibility.relay') }, + { value: 'public', label: t('visibility.public') }, + ], [t]); + + const visibilityLabel = useCallback((visibility: PageVisibility): string => { + switch (visibility) { + case 'private': return t('visibility.private'); + case 'relay': return t('visibility.relay'); + case 'public': return t('visibility.public'); + } + }, [t]); + + const formatBytes = useCallback((bytes: number): string => { + if (bytes < 1024) return t('bytes.b', { value: formatNumber(bytes) }); + if (bytes < 1024 * 1024) { + return t('bytes.kb', { value: formatNumber(bytes / 1024, { maximumFractionDigits: 1 }) }); + } + return t('bytes.mb', { + value: formatNumber(bytes / (1024 * 1024), { maximumFractionDigits: 1 }), + }); + }, [formatNumber, t]); + + const formatTimestamp = useCallback((seconds: number): string => formatDate( + new Date(seconds * 1000), + { year: 'numeric', month: 'short', day: 'numeric', hour: '2-digit', minute: '2-digit' }, + ), [formatDate]); + + const loadPages = useCallback(async () => { + if (busySlugsRef.current.size > 0) return; + const requestEpoch = pageLoadEpochRef.current + 1; + pageLoadEpochRef.current = requestEpoch; + attemptedLoadRef.current = true; + setLoading(true); + setLoadError(''); + setLoginRequired(false); + let requestedUserId: string | null = null; + try { + const status = await remoteConnectAPI.accountStatus(); + if (pageLoadEpochRef.current !== requestEpoch) return; + requestedUserId = status.user_id; + if (!status.logged_in || !status.user_id) { + adoptPageOwner(null, false); + setLoginRequired(true); + return; + } + const ownerEpoch = adoptPageOwner(status.user_id, false); + const [nextPages, hint] = await Promise.all([ + pageAPI.listPages(), + remoteConnectAPI.accountGetCredentialHint().catch(() => null), + ]); + const latestStatus = await remoteConnectAPI.accountStatus().catch(() => null); + if (pageLoadEpochRef.current !== requestEpoch) return; + if (!latestStatus?.logged_in || latestStatus.user_id !== status.user_id) { + // The account changed while this relay request was in flight. Leave + // ownership of the UI to a fresh request for the new account. + adoptPageOwner(latestStatus?.logged_in ? latestStatus.user_id : null, true); + attemptedLoadRef.current = !latestStatus?.logged_in; + return; + } + const currentOwner = pageOwnerRef.current; + if (currentOwner?.userId !== status.user_id || currentOwner.epoch !== ownerEpoch) return; + commitLoadedPages(nextPages); + setRelayBaseUrl(hint?.relay_url?.replace(/\/$/, '') ?? ''); + } catch (error) { + if (pageLoadEpochRef.current !== requestEpoch) return; + log.error('Failed to load published Pages', { error }); + const latestStatus = await remoteConnectAPI.accountStatus().catch(() => null); + if (pageLoadEpochRef.current !== requestEpoch) return; + if (requestedUserId !== null + && latestStatus?.logged_in + && latestStatus.user_id !== requestedUserId) { + adoptPageOwner(latestStatus.user_id, true); + attemptedLoadRef.current = false; + return; + } + if (latestStatus && !latestStatus.logged_in) { + adoptPageOwner(null, true); + setLoginRequired(true); + return; + } + setLoadError(errorText(error)); + } finally { + if (pageLoadEpochRef.current === requestEpoch) { + setLoading(false); + } + } + }, [adoptPageOwner, commitLoadedPages]); + + useEffect(() => { + if (isActive && !attemptedLoadRef.current && !loading) { + void loadPages(); + } + }, [isActive, loadPages, loading]); + + useEffect(() => { + const unlisten = api.listen<{ logged_in: boolean }>( + 'account://login-state', + (payload) => { + const loggedIn = payload?.logged_in === true; + // The event intentionally carries no user id. Treat every transition, + // including a same-user re-login, as a new ownership generation before + // doing any asynchronous status lookup. This immediately removes data + // and actions owned by the previous authenticated session. + adoptPageOwner(null, true, true); + attemptedLoadRef.current = !loggedIn; + if (loggedIn && isActive) { + void loadPages(); + } + }, + ); + return unlisten; + }, [adoptPageOwner, isActive, loadPages]); + + const loadVersions = useCallback(async (page: PageInfo, ownerEpoch: number) => { + const key = `versions:${page.slug}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + const versions = await pageAPI.listVersions(page.slug, page.generation); + if (!await validatePageAction(lease)) return; + setVersionsBySlug((current) => ({ ...current, [page.slug]: versions })); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to load Page versions', { slug: page.slug, error }); + notification.error(t('notifications.versionsLoadFailed', { error: errorText(error) })); + throw error; + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, validatePageAction]); + + const toggleVersions = useCallback(async (page: PageInfo, ownerEpoch: number) => { + if (expandedSlugs.has(page.slug)) { + setExpandedSlugs((current) => { + const next = new Set(current); + next.delete(page.slug); + return next; + }); + return; + } + if (!versionsBySlug[page.slug]) { + try { + await loadVersions(page, ownerEpoch); + } catch { + return; + } + } + if (pageOwnerRef.current?.epoch === ownerEpoch) { + setExpandedSlugs((current) => new Set(current).add(page.slug)); + } + }, [expandedSlugs, loadVersions, versionsBySlug]); + + const openPage = useCallback(async ( + page: PageInfo, + ownerEpoch: number, + versionId?: string, + ) => { + const key = `open:${page.slug}:${versionId ?? 'production'}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + const link = await pageAPI.createOpenLink(page.slug, page.generation, versionId); + if (!await validatePageAction(lease)) return; + await systemAPI.openExternal(link.open_url); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to open Page', { slug: page.slug, versionId, error }); + notification.error(t('notifications.openFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, validatePageAction]); + + const copyPageLink = useCallback(async ( + page: PageInfo, + ownerEpoch: number, + version?: PageVersionInfo, + ) => { + const key = `copy:${page.slug}:${version?.version_id ?? 'production'}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + let url = ''; + let temporary = false; + const path = version?.preview_url_path ?? page.url_path; + if (page.visibility === 'public' && relayBaseUrl && path) { + url = `${relayBaseUrl}${path.startsWith('/') ? '' : '/'}${path}`; + } else { + const link = await pageAPI.createOpenLink( + page.slug, + page.generation, + version?.version_id, + ); + url = link.open_url; + temporary = true; + } + if (!await validatePageAction(lease)) return; + await systemAPI.setClipboard(url); + notification.success(temporary + ? t('notifications.temporaryLinkCopied') + : t('notifications.linkCopied')); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to copy Page link', { + slug: page.slug, + versionId: version?.version_id, + error, + }); + notification.error(t('notifications.copyFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, relayBaseUrl, t, validatePageAction]); + + const changeVisibility = useCallback(async ( + page: PageInfo, + ownerEpoch: number, + visibility: PageVisibility, + ) => { + if (visibility === page.visibility) return; + const confirmed = await confirmWarning( + t('confirm.visibilityTitle'), + t('confirm.visibilityMessage', { + slug: page.slug, + current: visibilityLabel(page.visibility), + target: visibilityLabel(visibility), + }), + { confirmText: t('actions.changeVisibility') }, + ); + if (!confirmed) return; + const key = `visibility:${page.slug}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + const updated = await pageAPI.update(page.slug, page.generation, { visibility }); + if (!await validatePageAction(lease)) return; + updateOwnedPages((current) => replacePage(current, updated)); + notification.success(t('notifications.visibilityUpdated', { + visibility: visibilityLabel(visibility), + })); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to update Page visibility', { slug: page.slug, visibility, error }); + notification.error(t('notifications.visibilityFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, updateOwnedPages, validatePageAction, visibilityLabel]); + + const saveTitle = useCallback(async (page: PageInfo, ownerEpoch: number) => { + const title = (titleDrafts[page.slug] ?? page.title).trim(); + if (!title || title === page.title) return; + const key = `title:${page.slug}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + const updated = await pageAPI.update(page.slug, page.generation, { title }); + if (!await validatePageAction(lease)) return; + updateOwnedPages((current) => replacePage(current, updated)); + setTitleDrafts((current) => ({ ...current, [page.slug]: updated.title })); + notification.success(t('notifications.titleUpdated')); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to update Page title', { slug: page.slug, error }); + notification.error(t('notifications.titleFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, titleDrafts, updateOwnedPages, validatePageAction]); + + const deployVersion = useCallback(async ( + page: PageInfo, + ownerEpoch: number, + version: PageVersionInfo, + ) => { + if (version.deployed) return; + const confirmed = await confirmWarning( + t('confirm.deployTitle'), + t('confirm.deployMessage', { + slug: page.slug, + current: page.deployed_version_id ?? t('status.notDeployed'), + target: version.version_id, + }), + { confirmText: t('actions.deploy') }, + ); + if (!confirmed) return; + const key = `deploy:${page.slug}:${version.version_id}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + const updated = await pageAPI.deploy(page.slug, page.generation, version.version_id); + if (!await validatePageAction(lease)) return; + updateOwnedPages((current) => replacePage(current, updated)); + setVersionsBySlug((current) => ({ + ...current, + [page.slug]: (current[page.slug] ?? []).map((item) => ({ + ...item, + deployed: item.version_id === version.version_id, + })), + })); + notification.success(t('notifications.deployed', { version: version.version_id })); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to deploy Page version', { + slug: page.slug, + versionId: version.version_id, + error, + }); + notification.error(t('notifications.deployFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, updateOwnedPages, validatePageAction]); + + const unpublishPage = useCallback(async (page: PageInfo, ownerEpoch: number) => { + if (!page.deployed_version_id) return; + const confirmed = await confirmWarning( + t('confirm.unpublishTitle'), + t('confirm.unpublishMessage', { + slug: page.slug, + current: page.deployed_version_id, + }), + { confirmText: t('actions.unpublish') }, + ); + if (!confirmed) return; + const key = `unpublish:${page.slug}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + await pageAPI.unpublish(page.slug, page.generation); + if (!await validatePageAction(lease)) return; + updateOwnedPages((current) => current.map((item) => ( + item.slug === page.slug ? { ...item, deployed_version_id: null } : item + ))); + setVersionsBySlug((current) => ({ + ...current, + [page.slug]: (current[page.slug] ?? []).map((item) => ({ ...item, deployed: false })), + })); + notification.success(t('notifications.unpublished')); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to unpublish Page', { slug: page.slug, error }); + notification.error(t('notifications.unpublishFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, updateOwnedPages, validatePageAction]); + + const deleteVersion = useCallback(async ( + page: PageInfo, + ownerEpoch: number, + version: PageVersionInfo, + ) => { + if (version.deployed) return; + const confirmed = await confirmDanger( + t('confirm.deleteVersionTitle'), + t('confirm.deleteVersionMessage', { version: version.version_id, title: page.title }), + { confirmText: t('actions.deleteVersion') }, + ); + if (!confirmed) return; + const key = `delete-version:${page.slug}:${version.version_id}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + await pageAPI.deleteVersion(page.slug, page.generation, version.version_id); + if (!await validatePageAction(lease)) return; + setVersionsBySlug((current) => ({ + ...current, + [page.slug]: (current[page.slug] ?? []) + .filter((item) => item.version_id !== version.version_id), + })); + notification.success(t('notifications.versionDeleted')); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to delete Page version', { + slug: page.slug, + versionId: version.version_id, + error, + }); + notification.error(t('notifications.versionDeleteFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, validatePageAction]); + + const deletePage = useCallback(async (page: PageInfo, ownerEpoch: number) => { + const confirmed = await confirmDanger( + t('confirm.deletePageTitle'), + t('confirm.deletePageMessage', { title: page.title, slug: page.slug }), + { confirmText: t('actions.deletePage') }, + ); + if (!confirmed) return; + const key = `delete-page:${page.slug}`; + const lease = await beginPageAction(page, key, ownerEpoch); + if (!lease) return; + try { + await pageAPI.deletePage(page.slug, page.generation); + if (!await validatePageAction(lease)) return; + updateOwnedPages((current) => current.filter((item) => item.slug !== page.slug)); + setVersionsBySlug((current) => { + const next = { ...current }; + delete next[page.slug]; + return next; + }); + notification.success(t('notifications.pageDeleted')); + } catch (error) { + if (!await validatePageAction(lease)) return; + log.error('Failed to delete Page', { slug: page.slug, error }); + notification.error(t('notifications.pageDeleteFailed', { error: errorText(error) })); + } finally { + endPageAction(lease); + } + }, [beginPageAction, endPageAction, notification, t, updateOwnedPages, validatePageAction]); + + const refreshButton = ( + + ); + + return ( + + + + {loadError && pages.length > 0 && ( +
    + {t('loadFailed')} + {loadError} + +
    + )} + + {loading && pages.length === 0 ? ( + } + message={t('loading')} + testId="pages-loading" + /> + ) : loginRequired ? ( + } + message={<>{t('signInRequired')}{t('signInHint')}} + action={( + + )} + testId="pages-sign-in-required" + /> + ) : loadError && pages.length === 0 ? ( + } + message={<>{t('loadFailed')}{loadError}} + isError + action={} + testId="pages-error" + /> + ) : pages.length === 0 ? ( + } + message={<>{t('empty')}{t('emptyHint')}} + testId="pages-empty" + /> + ) : ( +
    + {pages.map((page) => { + const versions = versionsBySlug[page.slug] ?? []; + const expanded = expandedSlugs.has(page.slug); + const deployed = Boolean(page.deployed_version_id); + const pendingAction = pendingBySlug[page.slug]; + const pageBusy = Boolean(pendingAction); + const titleDraft = titleDrafts[page.slug] ?? page.title; + return ( +
    +
    +
    +

    {page.title || page.slug}

    + /{page.slug} +
    + + {deployed ? t('status.deployed') : t('status.savedOnly')} + +
    + +
    + {t('meta.updated', { date: formatTimestamp(page.updated_at) })} + {t('meta.size', { size: formatBytes(page.total_bytes) })} + {t('meta.files', { count: page.file_count })} +
    + +
    + {t('titleField.label')} +
    + { + const value = event.currentTarget.value; + setTitleDrafts((current) => ({ + ...current, + [page.slug]: value, + })); + }} + onKeyDown={(event) => { + if (event.key === 'Enter') void saveTitle(page, pageOwnerEpoch); + }} + aria-label={t('titleField.inputAria', { slug: page.slug })} + /> + +
    +
    + +
    + {t('visibility.label')} +