Skip to main content

squigit_storage/profiles/
store.rs

1// Copyright 2026 a7mddra
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::BTreeMap;
5use std::fs::{self, File, OpenOptions};
6use std::path::PathBuf;
7use std::sync::{Mutex, MutexGuard, OnceLock};
8
9use chrono::{DateTime, Utc};
10use fs2::FileExt;
11
12use super::types::{
13    EncryptedKeyRecord, KeyFile, LastLogin, Profile, ProfileAuth, ProfileIdentity, ProfileSnapshot,
14    AUTH_MODE_GOOGLE_OIDC_PKCE, AUTH_SCHEMA_VERSION, GOOGLE_PROVIDER, KEY_FILE_SCHEMA_VERSION,
15};
16use crate::error::{Result, StorageError};
17
18/// Active account state filename.
19const AUTH_FILE: &str = "auth.json";
20
21/// Consolidated profile metadata filename.
22const PROFILES_FILE: &str = "profiles.json";
23
24/// Consolidated encrypted API keys filename.
25const KEYS_FILE: &str = "keys.json";
26const KEYS_LOCK_FILE: &str = "keys.lock";
27
28type ProfileMap = BTreeMap<String, Profile>;
29static KEY_FILE_MUTEX: OnceLock<Mutex<()>> = OnceLock::new();
30
31pub struct KeyStoreTransaction<'a> {
32    store: &'a ProfileStore,
33    _process_guard: MutexGuard<'static, ()>,
34    lock_file: File,
35}
36
37impl KeyStoreTransaction<'_> {
38    pub fn load(&self) -> Result<KeyFile> {
39        self.store.load_key_file_unlocked()
40    }
41
42    pub fn save(&self, keys: &KeyFile) -> Result<()> {
43        self.store.save_key_file_unlocked(keys)
44    }
45}
46
47impl Drop for KeyStoreTransaction<'_> {
48    fn drop(&mut self) {
49        let _ = FileExt::unlock(&self.lock_file);
50    }
51}
52
53/// Manager for profile storage operations.
54///
55/// Root storage shape:
56/// - `{base_dir}/auth.json`
57/// - `{base_dir}/profiles.json`
58/// - `{base_dir}/keys.json`
59/// - `{base_dir}/threads/`
60pub struct ProfileStore {
61    /// Base directory: `{config_dir}/squigit/`
62    pub(super) base_dir: PathBuf,
63    /// Path to the active account state file.
64    pub(super) auth_path: PathBuf,
65    /// Path to the consolidated profile metadata file.
66    pub(super) profiles_path: PathBuf,
67    /// Path to the consolidated encrypted API keys file.
68    pub(super) keys_path: PathBuf,
69    /// Cross-process advisory lock for the encrypted API key store.
70    pub(super) keys_lock_path: PathBuf,
71}
72
73impl ProfileStore {
74    /// Create a new profile store.
75    ///
76    /// Uses the OS-appropriate config directory:
77    /// - Linux: `~/.config/squigit/`
78    /// - macOS: `~/Library/Application Support/squigit/`
79    /// - Windows: `%APPDATA%/squigit/`
80    pub fn new() -> Result<Self> {
81        let base_dir = crate::paths::base_config_dir().ok_or(StorageError::NoConfigDir)?;
82
83        Self::with_base_dir(base_dir)
84    }
85
86    /// Create a profile store using an explicit base directory.
87    ///
88    /// This is primarily intended for tests and future CLI integration.
89    pub fn with_base_dir(base_dir: PathBuf) -> Result<Self> {
90        let auth_path = base_dir.join(AUTH_FILE);
91        let profiles_path = base_dir.join(PROFILES_FILE);
92        let keys_path = base_dir.join(KEYS_FILE);
93        let keys_lock_path = base_dir.join(KEYS_LOCK_FILE);
94
95        fs::create_dir_all(&base_dir)?;
96        Self::ensure_private_directory(&base_dir)?;
97
98        Ok(Self {
99            base_dir,
100            auth_path,
101            profiles_path,
102            keys_path,
103            keys_lock_path,
104        })
105    }
106
107    /// Get the base storage directory path.
108    pub fn base_dir(&self) -> &PathBuf {
109        &self.base_dir
110    }
111
112    /// Get the directory path for a specific profile.
113    ///
114    /// Returns `{base_dir}/{profile_id}/`
115    pub fn get_profile_dir(&self, profile_id: &str) -> PathBuf {
116        self.base_dir.join(profile_id)
117    }
118
119    // =========================================================================
120    // Root File Operations
121    // =========================================================================
122
123    fn load_auth(&self) -> Result<ProfileAuth> {
124        if !self.auth_path.exists() {
125            return Ok(ProfileAuth::default());
126        }
127
128        let content = fs::read_to_string(&self.auth_path)?;
129        let auth: ProfileAuth = serde_json::from_str(&content)?;
130        Self::validate_auth(&auth)?;
131        Ok(auth)
132    }
133
134    fn save_auth(&self, auth: &ProfileAuth) -> Result<()> {
135        Self::validate_auth(auth)?;
136        self.write_json_atomic(&self.auth_path, auth)
137    }
138
139    fn load_profiles(&self) -> Result<ProfileMap> {
140        if !self.profiles_path.exists() {
141            return Ok(ProfileMap::default());
142        }
143
144        let content = fs::read_to_string(&self.profiles_path)?;
145        let profiles: ProfileMap = serde_json::from_str(&content)?;
146        Self::validate_profiles(&profiles)?;
147        Ok(profiles)
148    }
149
150    fn save_profiles(&self, profiles: &ProfileMap) -> Result<()> {
151        Self::validate_profiles(profiles)?;
152        self.write_json_atomic(&self.profiles_path, profiles)
153    }
154
155    fn load_key_file_unlocked(&self) -> Result<KeyFile> {
156        if !self.keys_path.exists() {
157            return Ok(KeyFile::default());
158        }
159
160        let content = fs::read_to_string(&self.keys_path)?;
161        let keys: KeyFile = serde_json::from_str(&content)
162            .map_err(|error| StorageError::KeyStore(format!("malformed-key-store: {error}")))?;
163        if keys.schema != KEY_FILE_SCHEMA_VERSION {
164            return Err(StorageError::KeyStore(format!(
165                "malformed-key-store: expected keys.json schema {KEY_FILE_SCHEMA_VERSION}"
166            )));
167        }
168        Self::validate_key_profiles(&keys)?;
169        Ok(keys)
170    }
171
172    fn save_key_file_unlocked(&self, keys: &KeyFile) -> Result<()> {
173        Self::validate_key_profiles(keys)?;
174        self.write_json_atomic(&self.keys_path, keys)
175    }
176
177    pub fn with_key_store_transaction<T, E>(
178        &self,
179        operation: impl FnOnce(&KeyStoreTransaction<'_>) -> std::result::Result<T, E>,
180    ) -> std::result::Result<T, E>
181    where
182        E: From<StorageError>,
183    {
184        let process_mutex = KEY_FILE_MUTEX.get_or_init(|| Mutex::new(()));
185        let process_guard = process_mutex
186            .lock()
187            .map_err(|_| StorageError::KeyStore("keys.lock mutex was poisoned".to_string()))?;
188
189        Self::reject_symlink(&self.keys_lock_path)?;
190        let mut options = OpenOptions::new();
191        options.read(true).write(true).create(true);
192        #[cfg(unix)]
193        {
194            use std::os::unix::fs::OpenOptionsExt;
195            options.mode(0o600);
196        }
197        let lock_file = options
198            .open(&self.keys_lock_path)
199            .map_err(StorageError::Io)?;
200        Self::set_private_file_permissions(&self.keys_lock_path)?;
201        lock_file.lock_exclusive().map_err(StorageError::Io)?;
202        let transaction = KeyStoreTransaction {
203            store: self,
204            _process_guard: process_guard,
205            lock_file,
206        };
207        operation(&transaction)
208    }
209
210    fn sorted_profiles(mut profiles: Vec<Profile>) -> Vec<Profile> {
211        profiles.sort_by_key(|profile| std::cmp::Reverse(profile.last_used_at));
212        profiles
213    }
214
215    fn newest_profile_id(profiles: &ProfileMap) -> Option<String> {
216        profiles
217            .values()
218            .max_by(|a, b| a.last_used_at.cmp(&b.last_used_at))
219            .map(|profile| profile.id.clone())
220    }
221
222    fn validate_auth(auth: &ProfileAuth) -> Result<()> {
223        if auth.schema != AUTH_SCHEMA_VERSION || auth.auth_mode != AUTH_MODE_GOOGLE_OIDC_PKCE {
224            return Err(StorageError::AuthState(format!(
225                "Unsupported auth.json schema. Delete the Squigit config folder or reinstall to start fresh with schema {}.",
226                AUTH_SCHEMA_VERSION
227            )));
228        }
229
230        if let Some(profile_id) = auth.active_profile_id.as_deref() {
231            Self::validate_profile_id(profile_id)?;
232        }
233        if let Some(last_login) = &auth.last_login {
234            let identity = ProfileIdentity::google(&last_login.issuer, &last_login.subject);
235            if last_login.provider != GOOGLE_PROVIDER
236                || last_login.profile_id != Profile::id_from_identity(&identity)
237            {
238                return Err(StorageError::InvalidProfileId(
239                    last_login.profile_id.clone(),
240                ));
241            }
242        }
243
244        Ok(())
245    }
246
247    fn validate_profile_id(profile_id: &str) -> Result<()> {
248        if Profile::is_canonical_id(profile_id) {
249            Ok(())
250        } else {
251            Err(StorageError::InvalidProfileId(profile_id.to_string()))
252        }
253    }
254
255    fn validate_profiles(profiles: &ProfileMap) -> Result<()> {
256        for (profile_id, profile) in profiles {
257            if profile_id != &profile.id || !profile.has_canonical_id() {
258                return Err(StorageError::InvalidProfileId(profile_id.clone()));
259            }
260        }
261        Ok(())
262    }
263
264    fn validate_key_profiles(keys: &KeyFile) -> Result<()> {
265        for profile_id in keys.profiles.keys() {
266            Self::validate_profile_id(profile_id)?;
267        }
268        Ok(())
269    }
270
271    pub fn load_encrypted_key_record(
272        &self,
273        profile_id: &str,
274        provider_key: &str,
275    ) -> Result<Option<EncryptedKeyRecord>> {
276        self.with_key_store_transaction(|transaction| {
277            let keys = transaction.load()?;
278            Ok(keys
279                .profiles
280                .get(profile_id)
281                .and_then(|profile_keys| profile_keys.get(provider_key))
282                .cloned())
283        })
284    }
285
286    pub fn update_last_trusted_reveal(&self) -> Result<()> {
287        self.with_key_store_transaction(|transaction| {
288            let mut keys = transaction.load()?;
289            keys.last_trusted_reveal = Some(Utc::now());
290            transaction.save(&keys)
291        })
292    }
293
294    pub fn invalidate_last_trusted_reveal(&self) -> Result<()> {
295        self.with_key_store_transaction(|transaction| {
296            let mut keys = transaction.load()?;
297            use chrono::TimeZone;
298            keys.last_trusted_reveal = Some(Utc.with_ymd_and_hms(1990, 1, 1, 0, 0, 0).unwrap());
299            transaction.save(&keys)
300        })
301    }
302
303    pub fn get_last_trusted_reveal(&self) -> Result<Option<DateTime<Utc>>> {
304        self.with_key_store_transaction(|transaction| {
305            let keys = transaction.load()?;
306            Ok(keys.last_trusted_reveal)
307        })
308    }
309
310    pub fn get_key_width(&self, profile_id: &str, provider_key: &str) -> Result<Option<u32>> {
311        self.with_key_store_transaction(|transaction| {
312            let keys = transaction.load()?;
313            Ok(keys
314                .profiles
315                .get(profile_id)
316                .and_then(|profile_keys| profile_keys.get(provider_key))
317                .map(|record| record.width))
318        })
319    }
320
321    /// Delete all encrypted key records for a profile.
322    pub fn delete_profile_key_records(&self, profile_id: &str) -> Result<bool> {
323        self.with_key_store_transaction(|transaction| {
324            let mut keys = transaction.load()?;
325            if keys.profiles.remove(profile_id).is_none() {
326                return Ok(keys.profiles.is_empty());
327            }
328            let is_empty = keys.profiles.is_empty();
329            transaction.save(&keys)?;
330            Ok(is_empty)
331        })
332    }
333
334    // =========================================================================
335    // Auth Operations
336    // =========================================================================
337
338    /// Get the ID of the currently active profile.
339    pub fn get_active_profile_id(&self) -> Result<Option<String>> {
340        let auth = self.load_auth()?;
341        let profiles = self.load_profiles()?;
342
343        Ok(auth
344            .active_profile_id
345            .filter(|profile_id| profiles.contains_key(profile_id)))
346    }
347
348    /// Set the active profile by ID.
349    ///
350    /// Returns an error if the profile doesn't exist.
351    pub fn set_active_profile_id(&self, profile_id: &str) -> Result<()> {
352        let profiles = self.load_profiles()?;
353
354        if !profiles.contains_key(profile_id) {
355            return Err(StorageError::ProfileNotFound(profile_id.to_string()));
356        }
357
358        let mut auth = self.load_auth()?;
359        auth.active_profile_id = Some(profile_id.to_string());
360        self.save_auth(&auth)?;
361        self.touch_profile(profile_id)?;
362        Ok(())
363    }
364
365    /// Record a successful provider login and activate the authenticated profile.
366    pub fn record_last_login(&self, last_login: LastLogin) -> Result<()> {
367        let profiles = self.load_profiles()?;
368
369        if !profiles.contains_key(&last_login.profile_id) {
370            return Err(StorageError::ProfileNotFound(last_login.profile_id.clone()));
371        }
372
373        self.save_auth(&ProfileAuth {
374            schema: AUTH_SCHEMA_VERSION,
375            auth_mode: AUTH_MODE_GOOGLE_OIDC_PKCE.to_string(),
376            active_profile_id: Some(last_login.profile_id.clone()),
377            last_login: Some(last_login.clone()),
378        })?;
379        self.touch_profile(&last_login.profile_id)?;
380        Ok(())
381    }
382
383    /// Clear the active profile (for Guest mode logout).
384    pub fn clear_active_profile_id(&self) -> Result<()> {
385        self.save_auth(&ProfileAuth::default())
386    }
387
388    // =========================================================================
389    // Profile CRUD
390    // =========================================================================
391
392    /// Create or update a profile.
393    ///
394    /// If the profile already exists, it will be updated with the new data.
395    /// Profile metadata is stored in the root profiles.json file.
396    pub fn upsert_profile(&self, profile: &Profile) -> Result<()> {
397        let mut profiles = self.load_profiles()?;
398        let mut stored_profile = profile.clone();
399
400        if let Some(existing_profile) = profiles.get(&profile.id) {
401            stored_profile.created_at = existing_profile.created_at;
402            if stored_profile.avatar_url.is_none() {
403                stored_profile.avatar_url = existing_profile.avatar_url.clone();
404            }
405            if stored_profile.avatar_base64.is_none()
406                && stored_profile.avatar_url == existing_profile.avatar_url
407            {
408                stored_profile.avatar_base64 = existing_profile.avatar_base64.clone();
409            }
410        }
411
412        profiles.insert(stored_profile.id.clone(), stored_profile.clone());
413        self.save_profiles(&profiles)?;
414
415        let auth = self.load_auth()?;
416        let needs_active_profile = match auth.active_profile_id.as_deref() {
417            Some(active_id) => !profiles.contains_key(active_id),
418            None => true,
419        };
420
421        if needs_active_profile {
422            let mut auth = self.load_auth()?;
423            auth.active_profile_id = Some(stored_profile.id);
424            self.save_auth(&auth)?;
425        }
426
427        Ok(())
428    }
429
430    /// Get a profile by ID.
431    pub fn get_profile(&self, profile_id: &str) -> Result<Option<Profile>> {
432        let profiles = self.load_profiles()?;
433        Ok(profiles.get(profile_id).cloned())
434    }
435
436    /// Get the currently active profile.
437    pub fn get_active_profile(&self) -> Result<Option<Profile>> {
438        let auth = self.load_auth()?;
439        let profiles = self.load_profiles()?;
440
441        Ok(auth
442            .active_profile_id
443            .and_then(|profile_id| profiles.get(&profile_id).cloned()))
444    }
445
446    /// Load active account state and all profiles from root files.
447    pub fn profile_snapshot(&self) -> Result<ProfileSnapshot> {
448        let auth = self.load_auth()?;
449        let profiles = self.load_profiles()?;
450        let active_profile_id = auth
451            .active_profile_id
452            .filter(|profile_id| profiles.contains_key(profile_id));
453        let active_profile = active_profile_id
454            .as_deref()
455            .and_then(|profile_id| profiles.get(profile_id).cloned());
456
457        Ok(ProfileSnapshot {
458            active_profile_id,
459            active_profile,
460            profiles: Self::sorted_profiles(profiles.into_values().collect()),
461        })
462    }
463
464    /// Delete a profile and all its data.
465    ///
466    /// Returns an error if trying to delete the last profile.
467    pub fn delete_profile(&self, profile_id: &str) -> Result<()> {
468        let mut profiles = self.load_profiles()?;
469
470        if profiles.len() <= 1 && profiles.contains_key(profile_id) {
471            return Err(StorageError::CannotDeleteLastProfile);
472        }
473
474        if profiles.remove(profile_id).is_none() {
475            return Err(StorageError::ProfileNotFound(profile_id.to_string()));
476        }
477
478        let profile_dir = self.get_profile_dir(profile_id);
479        if profile_dir.exists() {
480            fs::remove_dir_all(&profile_dir)?;
481        }
482
483        self.delete_profile_key_records(profile_id)?;
484        self.save_profiles(&profiles)?;
485
486        let mut auth = self.load_auth()?;
487        let active_is_missing = match auth.active_profile_id.as_deref() {
488            Some(active_id) => !profiles.contains_key(active_id),
489            None => true,
490        };
491
492        if active_is_missing {
493            auth.active_profile_id = Self::newest_profile_id(&profiles);
494        }
495
496        if auth
497            .last_login
498            .as_ref()
499            .is_some_and(|last_login| last_login.profile_id == profile_id)
500        {
501            auth.last_login = None;
502        }
503
504        self.save_auth(&auth)?;
505
506        Ok(())
507    }
508
509    fn touch_profile(&self, profile_id: &str) -> Result<()> {
510        let mut profiles = self.load_profiles()?;
511        let Some(profile) = profiles.get_mut(profile_id) else {
512            return Ok(());
513        };
514
515        profile.touch();
516        self.save_profiles(&profiles)
517    }
518}