secrets.rs

crates/cortex-core/src/secrets.rs

1. //! Where API keys are kept: the OS keychain, or memory in tests, behind a
2. //! per-session cache.
4. use std::collections::HashMap;
5. use std::sync::{Mutex, MutexGuard, PoisonError};
7. pub const SERVICE: &str = "com.mattrguerra.cortex";
9. /// A panic cannot leave these maps half-written, so a poisoned lock is still safe to use.
10. fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
11. m.lock().unwrap_or_else(PoisonError::into_inner)
12. }
14. /// What people paste is not always just the key: a trailing newline, quotes,
15. /// a `Bearer ` prefix, or an invisible character copied from a web page.
16. pub fn clean_key(raw: &str) -> String {
17. let s: String = raw
18. .chars()
19. .filter(|c| !c.is_whitespace() && !matches!(c, '\u{200b}'..='\u{200d}' | '\u{2060}' | '\u{feff}'))
20. .collect();
21. let s = s.trim_matches(|c| c == '"' || c == '\'' || c == '`');
22. let s = s.strip_prefix("Bearer").or_else(|| s.strip_prefix("bearer")).unwrap_or(s);
23. s.to_string()
24. }
26. pub fn account_for(provider_id: &str) -> String {
27. format!("provider:{provider_id}")
28. }
30. pub trait SecretStore: Send + Sync {
31. fn get(&self, account: &str) -> Result<Option<String>, String>;
32. fn set(&self, account: &str, secret: &str) -> Result<(), String>;
33. fn delete(&self, account: &str) -> Result<(), String>;
34. }
36. /// For tests and the end-to-end run.
37. #[derive(Default)]
38. pub struct MemoryStore(Mutex<HashMap<String, String>>);
40. impl SecretStore for MemoryStore {
41. fn get(&self, account: &str) -> Result<Option<String>, String> {
42. Ok(lock(&self.0).get(account).cloned())
43. }
44. fn set(&self, account: &str, secret: &str) -> Result<(), String> {
45. lock(&self.0).insert(account.into(), secret.into());
46. Ok(())
47. }
48. fn delete(&self, account: &str) -> Result<(), String> {
49. lock(&self.0).remove(account);
50. Ok(())
51. }
52. }
54. /// The OS keychain.
55. #[cfg(feature = "keychain")]
56. pub struct KeychainStore;
58. #[cfg(feature = "keychain")]
59. impl SecretStore for KeychainStore {
60. fn get(&self, account: &str) -> Result<Option<String>, String> {
61. let e = keyring::Entry::new(SERVICE, account).map_err(|e| e.to_string())?;
62. match e.get_password() {
63. Ok(s) => Ok(Some(s)),
64. Err(keyring::Error::NoEntry) => Ok(None),
65. Err(e) => Err(e.to_string()),
66. }
67. }
68. fn set(&self, account: &str, secret: &str) -> Result<(), String> {
69. keyring::Entry::new(SERVICE, account).and_then(|e| e.set_password(secret)).map_err(|e| e.to_string())
70. }
71. fn delete(&self, account: &str) -> Result<(), String> {
72. let e = keyring::Entry::new(SERVICE, account).map_err(|e| e.to_string())?;
73. match e.delete_credential() {
74. Ok(()) | Err(keyring::Error::NoEntry) => Ok(()),
75. Err(e) => Err(e.to_string()),
76. }
77. }
78. }
80. /// Reads each key from the store once per session (so macOS asks once) and
81. /// keeps it in memory after that.
82. pub struct KeyCache {
83. store: Box<dyn SecretStore>,
84. cache: Mutex<HashMap<String, Option<String>>>,
85. }
87. impl KeyCache {
88. pub fn new(store: Box<dyn SecretStore>) -> Self {
89. Self { store, cache: Mutex::new(HashMap::new()) }
90. }
92. pub fn get(&self, provider_id: &str) -> Result<Option<String>, String> {
93. let acct = account_for(provider_id);
94. if let Some(v) = lock(&self.cache).get(&acct) {
95. return Ok(v.clone());
96. }
97. let v = self.store.get(&acct)?;
98. lock(&self.cache).insert(acct, v.clone());
99. Ok(v)
100. }
102. pub fn set(&self, provider_id: &str, key: &str) -> Result<(), String> {
103. let key = clean_key(key);
104. let key = key.as_str();
105. let acct = account_for(provider_id);
106. self.store.set(&acct, key)?;
107. lock(&self.cache).insert(acct, Some(key.to_string()));
108. Ok(())
109. }
111. pub fn clear(&self, provider_id: &str) -> Result<(), String> {
112. let acct = account_for(provider_id);
113. self.store.delete(&acct)?;
114. lock(&self.cache).insert(acct, None);
115. Ok(())
116. }
118. pub fn has(&self, provider_id: &str) -> Result<bool, String> {
119. Ok(self.get(provider_id)?.is_some_and(|k| !k.is_empty()))
120. }
121. }
123. #[cfg(test)]
124. mod tests {
125. use std::sync::atomic::{AtomicUsize, Ordering};
126. use std::sync::Arc;
128. use super::*;
130. #[test]
131. fn pasted_keys_are_cleaned() {
132. for raw in
133. ["abc123", " abc123\n", "\"abc123\"", "Bearer abc123", "abc\u{200b}123", "'abc123'\r\n", "\u{feff}abc123"]
134. {
135. assert_eq!(clean_key(raw), "abc123", "{raw:?}");
136. }
137. assert_eq!(clean_key(" \n"), "");
138. }
140. struct Counting(MemoryStore, Arc<AtomicUsize>);
141. impl SecretStore for Counting {
142. fn get(&self, a: &str) -> Result<Option<String>, String> {
143. self.1.fetch_add(1, Ordering::SeqCst);
144. self.0.get(a)
145. }
146. fn set(&self, a: &str, s: &str) -> Result<(), String> {
147. self.0.set(a, s)
148. }
149. fn delete(&self, a: &str) -> Result<(), String> {
150. self.0.delete(a)
151. }
152. }
154. #[test]
155. fn reads_the_store_once_per_provider() {
156. let reads = Arc::new(AtomicUsize::new(0));
157. let c = KeyCache::new(Box::new(Counting(MemoryStore::default(), reads.clone())));
158. assert!(!c.has("deepinfra").unwrap());
159. assert!(!c.has("deepinfra").unwrap());
160. assert_eq!(reads.load(Ordering::SeqCst), 1);
161. c.set("deepinfra", " key-1 \n").unwrap();
162. assert_eq!(c.get("deepinfra").unwrap().as_deref(), Some("key-1"));
163. c.clear("deepinfra").unwrap();
164. assert!(!c.has("deepinfra").unwrap());
165. assert_eq!(reads.load(Ordering::SeqCst), 1);
166. }
167. }