From 8654c94835891c33d968b1b055a661309e07bbd5 Mon Sep 17 00:00:00 2001 From: David Longman Date: Fri, 6 Mar 2026 11:37:50 -0700 Subject: [PATCH] Use atomic writes for account storage --- src-tauri/src/auth/fs_utils.rs | 106 ++++++++++++++++ src-tauri/src/auth/mod.rs | 2 + src-tauri/src/auth/storage.rs | 218 +++++++++++++++++---------------- src-tauri/src/auth/switcher.rs | 5 +- 4 files changed, 225 insertions(+), 106 deletions(-) create mode 100644 src-tauri/src/auth/fs_utils.rs diff --git a/src-tauri/src/auth/fs_utils.rs b/src-tauri/src/auth/fs_utils.rs new file mode 100644 index 0000000..d79f1aa --- /dev/null +++ b/src-tauri/src/auth/fs_utils.rs @@ -0,0 +1,106 @@ +use std::ffi::OsString; +use std::fs::{self, File, OpenOptions}; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::thread; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use anyhow::{Context, Result}; + +pub struct FileLock { + path: PathBuf, +} + +impl FileLock { + pub fn acquire(target: &Path) -> Result { + let lock_path = sibling_with_suffix(target, ".lock"); + + if let Some(parent) = lock_path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("Failed to create lock directory: {}", parent.display()))?; + } + + for _ in 0..100 { + match OpenOptions::new() + .write(true) + .create_new(true) + .open(&lock_path) + { + Ok(mut file) => { + writeln!(file, "{}", std::process::id()).ok(); + return Ok(Self { path: lock_path }); + } + Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => { + thread::sleep(Duration::from_millis(50)); + } + Err(err) => { + return Err(err) + .with_context(|| format!("Failed to acquire lock: {}", lock_path.display())); + } + } + } + + anyhow::bail!("Timed out waiting for file lock: {}", lock_path.display()); + } +} + +impl Drop for FileLock { + fn drop(&mut self) { + let _ = fs::remove_file(&self.path); + } +} + +pub fn sibling_with_suffix(path: &Path, suffix: &str) -> PathBuf { + let mut file_name: OsString = path + .file_name() + .map(|name| name.to_os_string()) + .unwrap_or_else(|| OsString::from("file")); + file_name.push(suffix); + path.with_file_name(file_name) +} + +pub fn write_bytes_atomic(path: &Path, bytes: &[u8], backup_existing: bool) -> Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("Failed to create directory: {}", parent.display()))?; + } + + if backup_existing && path.exists() { + let backup_path = sibling_with_suffix(path, ".bak"); + fs::copy(path, &backup_path).with_context(|| { + format!( + "Failed to create backup {} from {}", + backup_path.display(), + path.display() + ) + })?; + } + + let temp_path = temp_path_for(path); + { + let mut file = File::create(&temp_path) + .with_context(|| format!("Failed to create temp file: {}", temp_path.display()))?; + file.write_all(bytes) + .with_context(|| format!("Failed to write temp file: {}", temp_path.display()))?; + file.sync_all() + .with_context(|| format!("Failed to sync temp file: {}", temp_path.display()))?; + } + + fs::rename(&temp_path, path).with_context(|| { + format!( + "Failed to replace {} with {}", + path.display(), + temp_path.display() + ) + })?; + + Ok(()) +} + +fn temp_path_for(path: &Path) -> PathBuf { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + sibling_with_suffix(path, &format!(".tmp-{}-{nanos}", std::process::id())) +} diff --git a/src-tauri/src/auth/mod.rs b/src-tauri/src/auth/mod.rs index 8827172..6062d2f 100644 --- a/src-tauri/src/auth/mod.rs +++ b/src-tauri/src/auth/mod.rs @@ -1,10 +1,12 @@ //! Authentication module +pub mod fs_utils; pub mod oauth_server; pub mod storage; pub mod switcher; pub mod token_refresh; +pub use fs_utils::*; pub use oauth_server::*; pub use storage::*; pub use switcher::*; diff --git a/src-tauri/src/auth/storage.rs b/src-tauri/src/auth/storage.rs index 282f5d3..50073d4 100644 --- a/src-tauri/src/auth/storage.rs +++ b/src-tauri/src/auth/storage.rs @@ -7,6 +7,8 @@ use anyhow::{Context, Result}; use crate::types::{AccountsStore, AuthData, StoredAccount}; +use super::fs_utils::{FileLock, write_bytes_atomic}; + /// Get the path to the codex-switcher config directory pub fn get_config_dir() -> Result { let home = dirs::home_dir().context("Could not find home directory")?; @@ -21,7 +23,10 @@ pub fn get_accounts_file() -> Result { /// Load the accounts store from disk pub fn load_accounts() -> Result { let path = get_accounts_file()?; + load_accounts_from_path(&path) +} +fn load_accounts_from_path(path: &PathBuf) -> Result { if !path.exists() { return Ok(AccountsStore::default()); } @@ -38,7 +43,11 @@ pub fn load_accounts() -> Result { /// Save the accounts store to disk pub fn save_accounts(store: &AccountsStore) -> Result<()> { let path = get_accounts_file()?; + let _lock = FileLock::acquire(&path)?; + save_accounts_to_path(&path, store) +} +fn save_accounts_to_path(path: &PathBuf, store: &AccountsStore) -> Result<()> { // Ensure the config directory exists if let Some(parent) = path.parent() { fs::create_dir_all(parent) @@ -48,7 +57,7 @@ pub fn save_accounts(store: &AccountsStore) -> Result<()> { let content = serde_json::to_string_pretty(store).context("Failed to serialize accounts store")?; - fs::write(&path, content) + write_bytes_atomic(path, content.as_bytes(), true) .with_context(|| format!("Failed to write accounts file: {}", path.display()))?; // Set restrictive permissions on Unix @@ -62,59 +71,64 @@ pub fn save_accounts(store: &AccountsStore) -> Result<()> { Ok(()) } +fn update_accounts_store(mutate: F) -> Result +where + F: FnOnce(&mut AccountsStore) -> Result, +{ + let path = get_accounts_file()?; + let _lock = FileLock::acquire(&path)?; + let mut store = load_accounts_from_path(&path)?; + let result = mutate(&mut store)?; + save_accounts_to_path(&path, &store)?; + Ok(result) +} + /// Add a new account to the store pub fn add_account(account: StoredAccount) -> Result { - let mut store = load_accounts()?; - - // Check for duplicate names - if store.accounts.iter().any(|a| a.name == account.name) { - anyhow::bail!("An account with name '{}' already exists", account.name); - } + update_accounts_store(|store| { + if store.accounts.iter().any(|a| a.name == account.name) { + anyhow::bail!("An account with name '{}' already exists", account.name); + } - let account_clone = account.clone(); - store.accounts.push(account); + let account_clone = account.clone(); + store.accounts.push(account); - // If this is the first account, make it active - if store.accounts.len() == 1 { - store.active_account_id = Some(account_clone.id.clone()); - } + if store.accounts.len() == 1 { + store.active_account_id = Some(account_clone.id.clone()); + } - save_accounts(&store)?; - Ok(account_clone) + Ok(account_clone) + }) } /// Remove an account by ID pub fn remove_account(account_id: &str) -> Result<()> { - let mut store = load_accounts()?; - - let initial_len = store.accounts.len(); - store.accounts.retain(|a| a.id != account_id); + update_accounts_store(|store| { + let initial_len = store.accounts.len(); + store.accounts.retain(|a| a.id != account_id); - if store.accounts.len() == initial_len { - anyhow::bail!("Account not found: {account_id}"); - } + if store.accounts.len() == initial_len { + anyhow::bail!("Account not found: {account_id}"); + } - // If we removed the active account, clear it or set to first available - if store.active_account_id.as_deref() == Some(account_id) { - store.active_account_id = store.accounts.first().map(|a| a.id.clone()); - } + if store.active_account_id.as_deref() == Some(account_id) { + store.active_account_id = store.accounts.first().map(|a| a.id.clone()); + } - save_accounts(&store)?; - Ok(()) + Ok(()) + }) } /// Update the active account ID pub fn set_active_account(account_id: &str) -> Result<()> { - let mut store = load_accounts()?; - - // Verify the account exists - if !store.accounts.iter().any(|a| a.id == account_id) { - anyhow::bail!("Account not found: {account_id}"); - } + update_accounts_store(|store| { + if !store.accounts.iter().any(|a| a.id == account_id) { + anyhow::bail!("Account not found: {account_id}"); + } - store.active_account_id = Some(account_id.to_string()); - save_accounts(&store)?; - Ok(()) + store.active_account_id = Some(account_id.to_string()); + Ok(()) + }) } /// Get an account by ID @@ -135,14 +149,13 @@ pub fn get_active_account() -> Result> { /// Update an account's last_used_at timestamp pub fn touch_account(account_id: &str) -> Result<()> { - let mut store = load_accounts()?; - - if let Some(account) = store.accounts.iter_mut().find(|a| a.id == account_id) { - account.last_used_at = Some(chrono::Utc::now()); - save_accounts(&store)?; - } + update_accounts_store(|store| { + if let Some(account) = store.accounts.iter_mut().find(|a| a.id == account_id) { + account.last_used_at = Some(chrono::Utc::now()); + } - Ok(()) + Ok(()) + }) } /// Update an account's metadata (name, email, plan_type) @@ -152,40 +165,37 @@ pub fn update_account_metadata( email: Option, plan_type: Option, ) -> Result<()> { - let mut store = load_accounts()?; - - // Check for duplicate names first (if renaming) - if let Some(ref new_name) = name { - if store - .accounts - .iter() - .any(|a| a.id != account_id && a.name == *new_name) - { - anyhow::bail!("An account with name '{new_name}' already exists"); + update_accounts_store(|store| { + if let Some(ref new_name) = name { + if store + .accounts + .iter() + .any(|a| a.id != account_id && a.name == *new_name) + { + anyhow::bail!("An account with name '{new_name}' already exists"); + } } - } - // Now find and update the account - let account = store - .accounts - .iter_mut() - .find(|a| a.id == account_id) - .context("Account not found")?; + let account = store + .accounts + .iter_mut() + .find(|a| a.id == account_id) + .context("Account not found")?; - if let Some(new_name) = name { - account.name = new_name; - } + if let Some(new_name) = name { + account.name = new_name; + } - if email.is_some() { - account.email = email; - } + if email.is_some() { + account.email = email; + } - if plan_type.is_some() { - account.plan_type = plan_type; - } + if plan_type.is_some() { + account.plan_type = plan_type; + } - save_accounts(&store)?; - Ok(()) + Ok(()) + }) } /// Update ChatGPT OAuth tokens for an account and return the updated account. @@ -198,42 +208,40 @@ pub fn update_account_chatgpt_tokens( email: Option, plan_type: Option, ) -> Result { - let mut store = load_accounts()?; - - let account = store - .accounts - .iter_mut() - .find(|a| a.id == account_id) - .context("Account not found")?; - - match &mut account.auth_data { - AuthData::ChatGPT { - id_token: stored_id_token, - access_token: stored_access_token, - refresh_token: stored_refresh_token, - account_id: stored_account_id, - } => { - *stored_id_token = id_token; - *stored_access_token = access_token; - *stored_refresh_token = refresh_token; - if let Some(new_account_id) = chatgpt_account_id { - *stored_account_id = Some(new_account_id); + update_accounts_store(|store| { + let account = store + .accounts + .iter_mut() + .find(|a| a.id == account_id) + .context("Account not found")?; + + match &mut account.auth_data { + AuthData::ChatGPT { + id_token: stored_id_token, + access_token: stored_access_token, + refresh_token: stored_refresh_token, + account_id: stored_account_id, + } => { + *stored_id_token = id_token; + *stored_access_token = access_token; + *stored_refresh_token = refresh_token; + if let Some(new_account_id) = chatgpt_account_id { + *stored_account_id = Some(new_account_id); + } + } + AuthData::ApiKey { .. } => { + anyhow::bail!("Cannot update OAuth tokens for an API key account"); } } - AuthData::ApiKey { .. } => { - anyhow::bail!("Cannot update OAuth tokens for an API key account"); - } - } - if let Some(new_email) = email { - account.email = Some(new_email); - } + if let Some(new_email) = email { + account.email = Some(new_email); + } - if let Some(new_plan_type) = plan_type { - account.plan_type = Some(new_plan_type); - } + if let Some(new_plan_type) = plan_type { + account.plan_type = Some(new_plan_type); + } - let updated = account.clone(); - save_accounts(&store)?; - Ok(updated) + Ok(account.clone()) + }) } diff --git a/src-tauri/src/auth/switcher.rs b/src-tauri/src/auth/switcher.rs index 8fb4959..664d5f9 100644 --- a/src-tauri/src/auth/switcher.rs +++ b/src-tauri/src/auth/switcher.rs @@ -8,6 +8,8 @@ use chrono::Utc; use crate::types::{AuthData, AuthDotJson, StoredAccount, TokenData}; +use super::fs_utils::{FileLock, write_bytes_atomic}; + /// Get the official Codex home directory pub fn get_codex_home() -> Result { // Check for CODEX_HOME environment variable first @@ -35,10 +37,11 @@ pub fn switch_to_account(account: &StoredAccount) -> Result<()> { let auth_json = create_auth_json(account)?; let auth_path = codex_home.join("auth.json"); + let _lock = FileLock::acquire(&auth_path)?; let content = serde_json::to_string_pretty(&auth_json).context("Failed to serialize auth.json")?; - fs::write(&auth_path, content) + write_bytes_atomic(&auth_path, content.as_bytes(), true) .with_context(|| format!("Failed to write auth.json: {}", auth_path.display()))?; // Set restrictive permissions on Unix