| 1 | //! App-owned credential storage: `modify` is the only write path. |
| 2 | //! |
| 3 | //! Ported from pi-mono `packages/ai/src/auth/types.ts` (`CredentialStore`) and |
| 4 | //! `packages/ai/src/auth/credential-store.ts` (`InMemoryCredentialStore`), |
| 5 | //! MIT, Copyright (c) 2025 Mario Zechner — full notice in the parent module. |
| 6 | //! Several doc comments below are adapted closely from pi's. |
| 7 | //! |
| 8 | //! pi's rule, kept verbatim in spirit: every mutation is a serialized |
| 9 | //! read-modify-write whose closure sees the current credential, so a refresh |
| 10 | //! and a concurrent login cannot clobber each other. |
| 11 | //! |
| 12 | //! CodeWhale already serializes the xAI OAuth refresh that way, but not |
| 13 | //! through this store: `xai_oauth` holds `with_xai_oauth_lifecycle_lock` |
| 14 | //! across the token-file read, the refresh request, and the write-back, so |
| 15 | //! two concurrent near-expiry observers share one rotated epoch rather than |
| 16 | //! overwriting each other's refresh token. That lock is process- and |
| 17 | //! file-level, not this registry's per-provider mutex, and the xAI flow is |
| 18 | //! **not** rewritten onto [`CredentialStore::modify`] here. This trait is |
| 19 | //! the shape that would put API-key slots and that OAuth flow on one write |
| 20 | //! path. |
| 21 | |
| 22 | use std::collections::HashMap; |
| 23 | use std::sync::{Arc, Mutex, OnceLock}; |
| 24 | |
| 25 | use anyhow::Result; |
| 26 | |
| 27 | use super::{Credential, CredentialKind}; |
| 28 | |
| 29 | /// Non-secret credential metadata for account/status enumeration. |
| 30 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 31 | pub(crate) struct CredentialInfo { |
| 32 | pub(crate) provider_id: String, |
| 33 | pub(crate) kind: CredentialKind, |
| 34 | } |
| 35 | |
| 36 | /// App-owned credential storage, keyed by provider id, one credential per |
| 37 | /// provider. |
| 38 | /// |
| 39 | /// `modify` is the only write path, so every mutation is a serialized |
| 40 | /// read-modify-write. Callers that need to refresh a rotated token run the |
| 41 | /// refresh *inside* `modify` so concurrent requests cannot double-refresh. |
| 42 | /// |
| 43 | /// Error semantics: `read` yields `Ok(None)` for a missing entry and `Err` |
| 44 | /// on storage failure. `list` is best-effort per slot: a single unreadable |
| 45 | /// entry is omitted so one corrupt slot cannot hide every other stored |
| 46 | /// credential from enumeration (status, `/provider`, logout). `list` fails |
| 47 | /// only when enumeration itself cannot run. `list` must not execute |
| 48 | /// configured API-key commands or open network flows. |
| 49 | pub(crate) trait CredentialStore: Send + Sync { |
| 50 | /// Read the stored credential, possibly expired. Display/status use. |
| 51 | fn read(&self, provider_id: &str) -> Result<Option<Credential>>; |
| 52 | |
| 53 | /// List stored credential metadata without exposing secrets. |
| 54 | /// |
| 55 | /// Per-slot read failures are omitted rather than propagated. |
| 56 | fn list(&self) -> Result<Vec<CredentialInfo>>; |
| 57 | |
| 58 | /// Serialized write — the only write path. `f` sees the current |
| 59 | /// credential because correct writes (refresh, login-during-refresh) |
| 60 | /// depend on it; return the new credential, or `None` to leave the entry |
| 61 | /// unchanged. Mutual exclusion is per provider id. Yields the post-write |
| 62 | /// credential. |
| 63 | /// |
| 64 | /// Not yet on any production write path: CodeWhale's two credential writes |
| 65 | /// (`save_api_key_for_identity`, logout) already interleave a config-file |
| 66 | /// mutation and a compensating rollback between their snapshot and their |
| 67 | /// store write, so they take [`with_provider_write_lock`] around the whole |
| 68 | /// sequence instead. Collapsing them onto `modify` would change their |
| 69 | /// error messages and rollback shape, which is deliberately out of this |
| 70 | /// change's scope. |
| 71 | #[cfg_attr(not(test), expect(dead_code))] |
| 72 | fn modify( |
| 73 | &self, |
| 74 | provider_id: &str, |
| 75 | f: &mut dyn FnMut(Option<Credential>) -> Result<Option<Credential>>, |
| 76 | ) -> Result<Option<Credential>>; |
| 77 | |
| 78 | /// Remove a credential (logout). Serialized against `modify`. |
| 79 | /// |
| 80 | /// Exercised by the store tests but not yet on the production logout path: |
| 81 | /// the full-wipe logout in `config.rs` still deletes slots directly rather |
| 82 | /// than through this trait. Routing it here is the remaining half of the |
| 83 | /// port and is deliberately not folded into this change — logout is |
| 84 | /// security-sensitive and deserves its own commit and its own tests. |
| 85 | #[cfg_attr(not(test), expect(dead_code))] |
| 86 | fn delete(&self, provider_id: &str) -> Result<()>; |
| 87 | } |
| 88 | |
| 89 | /// Process-wide per-provider write locks. |
| 90 | /// |
| 91 | /// pi serializes through a per-provider promise chain; the Rust equivalent is |
| 92 | /// a registry of per-id mutexes. This is process-local: it does not serialize |
| 93 | /// against another `codewhale` process writing the same slot, which is a real |
| 94 | /// remaining gap and is called out in the module docs. |
| 95 | fn provider_lock(provider_id: &str) -> Arc<Mutex<()>> { |
| 96 | static LOCKS: OnceLock<Mutex<HashMap<String, Arc<Mutex<()>>>>> = OnceLock::new(); |
| 97 | let locks = LOCKS.get_or_init(|| Mutex::new(HashMap::new())); |
| 98 | let mut guard = locks |
| 99 | .lock() |
| 100 | .unwrap_or_else(std::sync::PoisonError::into_inner); |
| 101 | Arc::clone( |
| 102 | guard |
| 103 | .entry(provider_id.to_string()) |
| 104 | .or_insert_with(|| Arc::new(Mutex::new(()))), |
| 105 | ) |
| 106 | } |
| 107 | |
| 108 | /// Run `body` holding this provider's write lock. |
| 109 | pub(crate) fn with_provider_write_lock<T>(provider_id: &str, body: impl FnOnce() -> T) -> T { |
| 110 | let lock = provider_lock(provider_id); |
| 111 | let _guard = lock |
| 112 | .lock() |
| 113 | .unwrap_or_else(std::sync::PoisonError::into_inner); |
| 114 | body() |
| 115 | } |
| 116 | |
| 117 | /// Run `body` holding every listed provider's write lock. |
| 118 | /// |
| 119 | /// Locks are acquired in sorted, deduplicated order so a full logout that |
| 120 | /// covers every slot cannot deadlock with a per-provider save (which holds |
| 121 | /// only one). Callers that also take the xAI OAuth lifecycle lock must take |
| 122 | /// that lock first, matching the documented xAI-then-config order. |
| 123 | pub(crate) fn with_provider_write_locks<T>( |
| 124 | provider_ids: impl IntoIterator<Item = impl AsRef<str>>, |
| 125 | body: impl FnOnce() -> T, |
| 126 | ) -> T { |
| 127 | let mut ids: Vec<String> = provider_ids |
| 128 | .into_iter() |
| 129 | .map(|id| id.as_ref().to_string()) |
| 130 | .collect(); |
| 131 | ids.sort(); |
| 132 | ids.dedup(); |
| 133 | let locks: Vec<Arc<Mutex<()>>> = ids.iter().map(|id| provider_lock(id)).collect(); |
| 134 | let _guards: Vec<_> = locks |
| 135 | .iter() |
| 136 | .map(|lock| { |
| 137 | lock.lock() |
| 138 | .unwrap_or_else(std::sync::PoisonError::into_inner) |
| 139 | }) |
| 140 | .collect(); |
| 141 | body() |
| 142 | } |
| 143 | |
| 144 | /// Default in-memory store. Real stores are injected; this one backs tests and |
| 145 | /// keeps the trait honest about its own contract — including the serialization |
| 146 | /// guarantee, which is only observable through `modify`. |
| 147 | #[derive(Debug, Default)] |
| 148 | pub(crate) struct InMemoryCredentialStore { |
| 149 | entries: Mutex<HashMap<String, Credential>>, |
| 150 | } |
| 151 | |
| 152 | impl InMemoryCredentialStore { |
| 153 | #[cfg_attr(not(test), expect(dead_code))] |
| 154 | pub(crate) fn new() -> Self { |
| 155 | Self::default() |
| 156 | } |
| 157 | } |
| 158 | |
| 159 | impl CredentialStore for InMemoryCredentialStore { |
| 160 | fn read(&self, provider_id: &str) -> Result<Option<Credential>> { |
| 161 | Ok(self |
| 162 | .entries |
| 163 | .lock() |
| 164 | .unwrap_or_else(std::sync::PoisonError::into_inner) |
| 165 | .get(provider_id) |
| 166 | .cloned()) |
| 167 | } |
| 168 | |
| 169 | fn list(&self) -> Result<Vec<CredentialInfo>> { |
| 170 | let entries = self |
| 171 | .entries |
| 172 | .lock() |
| 173 | .unwrap_or_else(std::sync::PoisonError::into_inner); |
| 174 | let mut infos: Vec<CredentialInfo> = entries |
| 175 | .iter() |
| 176 | .map(|(provider_id, credential)| CredentialInfo { |
| 177 | provider_id: provider_id.clone(), |
| 178 | kind: credential.kind(), |
| 179 | }) |
| 180 | .collect(); |
| 181 | infos.sort_by(|left, right| left.provider_id.cmp(&right.provider_id)); |
| 182 | Ok(infos) |
| 183 | } |
| 184 | |
| 185 | fn modify( |
| 186 | &self, |
| 187 | provider_id: &str, |
| 188 | f: &mut dyn FnMut(Option<Credential>) -> Result<Option<Credential>>, |
| 189 | ) -> Result<Option<Credential>> { |
| 190 | with_provider_write_lock(provider_id, || { |
| 191 | let current = self.read(provider_id)?; |
| 192 | let next = f(current.clone())?; |
| 193 | match next { |
| 194 | Some(credential) => { |
| 195 | self.entries |
| 196 | .lock() |
| 197 | .unwrap_or_else(std::sync::PoisonError::into_inner) |
| 198 | .insert(provider_id.to_string(), credential.clone()); |
| 199 | Ok(Some(credential)) |
| 200 | } |
| 201 | None => Ok(current), |
| 202 | } |
| 203 | }) |
| 204 | } |
| 205 | |
| 206 | fn delete(&self, provider_id: &str) -> Result<()> { |
| 207 | with_provider_write_lock(provider_id, || { |
| 208 | self.entries |
| 209 | .lock() |
| 210 | .unwrap_or_else(std::sync::PoisonError::into_inner) |
| 211 | .remove(provider_id); |
| 212 | Ok(()) |
| 213 | }) |
| 214 | } |
| 215 | } |
| 216 | |
| 217 | /// Adapter over CodeWhale's existing durable secret store. |
| 218 | /// |
| 219 | /// This changes no on-disk format: it reads and writes exactly the slots |
| 220 | /// `codewhale_secrets::Secrets` already owns. Its value is that every write now |
| 221 | /// goes through `modify` under the provider's lock, so a save racing a rotate |
| 222 | /// no longer interleaves. |
| 223 | pub(crate) struct SecretStoreCredentials { |
| 224 | secrets: codewhale_secrets::Secrets, |
| 225 | /// Slots to probe for `list`. The backing keyring exposes no key |
| 226 | /// enumeration, so the caller supplies the known slot names. |
| 227 | known_slots: Vec<String>, |
| 228 | } |
| 229 | |
| 230 | impl SecretStoreCredentials { |
| 231 | pub(crate) fn new(secrets: codewhale_secrets::Secrets, known_slots: Vec<String>) -> Self { |
| 232 | Self { |
| 233 | secrets, |
| 234 | known_slots, |
| 235 | } |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | impl std::fmt::Debug for SecretStoreCredentials { |
| 240 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 241 | f.debug_struct("SecretStoreCredentials") |
| 242 | .field("backend", &self.secrets.backend_name()) |
| 243 | .field("known_slots", &self.known_slots.len()) |
| 244 | .finish() |
| 245 | } |
| 246 | } |
| 247 | |
| 248 | impl CredentialStore for SecretStoreCredentials { |
| 249 | fn read(&self, provider_id: &str) -> Result<Option<Credential>> { |
| 250 | Ok(self |
| 251 | .secrets |
| 252 | .get(provider_id)? |
| 253 | .filter(|value| !value.trim().is_empty()) |
| 254 | .map(|key| Credential::ApiKey { key })) |
| 255 | } |
| 256 | |
| 257 | fn list(&self) -> Result<Vec<CredentialInfo>> { |
| 258 | let mut infos = Vec::new(); |
| 259 | for slot in &self.known_slots { |
| 260 | match self.read(slot) { |
| 261 | Ok(Some(_)) => infos.push(CredentialInfo { |
| 262 | provider_id: slot.clone(), |
| 263 | kind: CredentialKind::ApiKey, |
| 264 | }), |
| 265 | Ok(None) => {} |
| 266 | Err(error) => { |
| 267 | // Deliberate: skip the bad slot and keep going. Propagating |
| 268 | // `read`'s error used to fail the whole enumeration, so one |
| 269 | // unreadable entry hid every other credential from |
| 270 | // `/provider` and left logout unable to delete the rest. |
| 271 | tracing::warn!( |
| 272 | slot, |
| 273 | error = %error, |
| 274 | "skipping unreadable credential slot during enumeration" |
| 275 | ); |
| 276 | } |
| 277 | } |
| 278 | } |
| 279 | Ok(infos) |
| 280 | } |
| 281 | |
| 282 | fn modify( |
| 283 | &self, |
| 284 | provider_id: &str, |
| 285 | f: &mut dyn FnMut(Option<Credential>) -> Result<Option<Credential>>, |
| 286 | ) -> Result<Option<Credential>> { |
| 287 | with_provider_write_lock(provider_id, || { |
| 288 | let current = self.read(provider_id)?; |
| 289 | let next = f(current.clone())?; |
| 290 | match next { |
| 291 | Some(credential) => { |
| 292 | self.secrets.set(provider_id, credential.expose_secret())?; |
| 293 | Ok(Some(credential)) |
| 294 | } |
| 295 | None => Ok(current), |
| 296 | } |
| 297 | }) |
| 298 | } |
| 299 | |
| 300 | fn delete(&self, provider_id: &str) -> Result<()> { |
| 301 | with_provider_write_lock(provider_id, || { |
| 302 | self.secrets.delete(provider_id)?; |
| 303 | Ok(()) |
| 304 | }) |
| 305 | } |
| 306 | } |
| 307 |