edgefirst_client/
storage.rs1use directories::ProjectDirs;
40use log::debug;
41use std::{path::PathBuf, sync::RwLock};
42
43#[derive(Debug)]
45pub enum StorageError {
46 NotAvailable(String),
48 ReadError(String),
50 WriteError(String),
52 ClearError(String),
54}
55
56impl std::fmt::Display for StorageError {
57 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58 match self {
59 StorageError::NotAvailable(msg) => write!(f, "Token storage not available: {}", msg),
60 StorageError::ReadError(msg) => write!(f, "Failed to read token: {}", msg),
61 StorageError::WriteError(msg) => write!(f, "Failed to write token: {}", msg),
62 StorageError::ClearError(msg) => write!(f, "Failed to clear token: {}", msg),
63 }
64 }
65}
66
67impl std::error::Error for StorageError {}
68
69pub trait TokenStorage: Send + Sync {
109 fn store(&self, token: &str) -> Result<(), StorageError>;
111
112 fn load(&self) -> Result<Option<String>, StorageError>;
116
117 fn clear(&self) -> Result<(), StorageError>;
119}
120
121#[derive(Debug, Clone)]
140pub struct FileTokenStorage {
141 path: PathBuf,
142}
143
144impl FileTokenStorage {
145 pub fn new() -> Result<Self, StorageError> {
155 let path = ProjectDirs::from("ai", "EdgeFirst", "EdgeFirst Studio")
156 .ok_or_else(|| {
157 StorageError::NotAvailable("Could not determine user config directory".to_string())
158 })?
159 .config_dir()
160 .join("token");
161
162 debug!("FileTokenStorage using default path: {:?}", path);
163 Ok(Self { path })
164 }
165
166 pub fn with_path(path: PathBuf) -> Self {
168 debug!("FileTokenStorage using custom path: {:?}", path);
169 Self { path }
170 }
171
172 pub fn path(&self) -> &PathBuf {
174 &self.path
175 }
176}
177
178impl TokenStorage for FileTokenStorage {
179 fn store(&self, token: &str) -> Result<(), StorageError> {
180 if let Some(parent) = self.path.parent() {
182 std::fs::create_dir_all(parent).map_err(|e| {
183 StorageError::WriteError(format!("Failed to create directory {:?}: {}", parent, e))
184 })?;
185 }
186
187 std::fs::write(&self.path, token).map_err(|e| {
188 StorageError::WriteError(format!("Failed to write token to {:?}: {}", self.path, e))
189 })?;
190
191 debug!("Token stored to {:?}", self.path);
192 Ok(())
193 }
194
195 fn load(&self) -> Result<Option<String>, StorageError> {
196 if !self.path.exists() {
197 debug!("No token file found at {:?}", self.path);
198 return Ok(None);
199 }
200
201 let token = std::fs::read_to_string(&self.path).map_err(|e| {
202 StorageError::ReadError(format!("Failed to read token from {:?}: {}", self.path, e))
203 })?;
204
205 if token.is_empty() {
206 debug!("Token file at {:?} is empty", self.path);
207 return Ok(None);
208 }
209
210 debug!("Token loaded from {:?}", self.path);
211 Ok(Some(token))
212 }
213
214 fn clear(&self) -> Result<(), StorageError> {
215 if self.path.exists() {
216 std::fs::remove_file(&self.path).map_err(|e| {
217 StorageError::ClearError(format!(
218 "Failed to remove token file {:?}: {}",
219 self.path, e
220 ))
221 })?;
222 debug!("Token file removed from {:?}", self.path);
223 }
224 Ok(())
225 }
226}
227
228#[derive(Debug, Default)]
249pub struct MemoryTokenStorage {
250 token: RwLock<Option<String>>,
251}
252
253impl MemoryTokenStorage {
254 pub fn new() -> Self {
256 Self::default()
257 }
258}
259
260impl TokenStorage for MemoryTokenStorage {
261 fn store(&self, token: &str) -> Result<(), StorageError> {
262 let mut guard = self.token.write().map_err(|e| {
263 StorageError::WriteError(format!("Failed to acquire write lock: {}", e))
264 })?;
265 *guard = Some(token.to_string());
266 Ok(())
267 }
268
269 fn load(&self) -> Result<Option<String>, StorageError> {
270 let guard = self
271 .token
272 .read()
273 .map_err(|e| StorageError::ReadError(format!("Failed to acquire read lock: {}", e)))?;
274 Ok(guard.clone())
275 }
276
277 fn clear(&self) -> Result<(), StorageError> {
278 let mut guard = self.token.write().map_err(|e| {
279 StorageError::ClearError(format!("Failed to acquire write lock: {}", e))
280 })?;
281 *guard = None;
282 Ok(())
283 }
284}
285
286#[cfg(test)]
287mod tests {
288 use super::*;
289 use std::sync::Arc;
290 use tempfile::TempDir;
291
292 #[test]
293 fn test_memory_storage_store_load_clear() {
294 let storage = MemoryTokenStorage::new();
295
296 assert_eq!(storage.load().unwrap(), None);
298
299 storage.store("test-token").unwrap();
301 assert_eq!(storage.load().unwrap(), Some("test-token".to_string()));
302
303 storage.clear().unwrap();
305 assert_eq!(storage.load().unwrap(), None);
306 }
307
308 #[test]
309 fn test_memory_storage_overwrite() {
310 let storage = MemoryTokenStorage::new();
311
312 storage.store("token-1").unwrap();
313 assert_eq!(storage.load().unwrap(), Some("token-1".to_string()));
314
315 storage.store("token-2").unwrap();
316 assert_eq!(storage.load().unwrap(), Some("token-2".to_string()));
317 }
318
319 #[test]
320 fn test_memory_storage_thread_safety() {
321 let storage = Arc::new(MemoryTokenStorage::new());
322 let storage_clone = Arc::clone(&storage);
323
324 let handle = std::thread::spawn(move || {
325 storage_clone.store("thread-token").unwrap();
326 });
327
328 handle.join().unwrap();
329 assert_eq!(storage.load().unwrap(), Some("thread-token".to_string()));
330 }
331
332 #[test]
333 fn test_file_storage_store_load_clear() {
334 let temp_dir = TempDir::new().unwrap();
335 let token_path = temp_dir.path().join("token");
336 let storage = FileTokenStorage::with_path(token_path.clone());
337
338 assert_eq!(storage.load().unwrap(), None);
340
341 storage.store("file-test-token").unwrap();
343 assert!(token_path.exists());
344 assert_eq!(storage.load().unwrap(), Some("file-test-token".to_string()));
345
346 storage.clear().unwrap();
348 assert!(!token_path.exists());
349 assert_eq!(storage.load().unwrap(), None);
350 }
351
352 #[test]
353 fn test_file_storage_creates_parent_dirs() {
354 let temp_dir = TempDir::new().unwrap();
355 let token_path = temp_dir.path().join("nested").join("dirs").join("token");
356 let storage = FileTokenStorage::with_path(token_path.clone());
357
358 storage.store("nested-token").unwrap();
359 assert!(token_path.exists());
360 assert_eq!(storage.load().unwrap(), Some("nested-token".to_string()));
361 }
362
363 #[test]
364 fn test_file_storage_overwrite() {
365 let temp_dir = TempDir::new().unwrap();
366 let token_path = temp_dir.path().join("token");
367 let storage = FileTokenStorage::with_path(token_path);
368
369 storage.store("token-1").unwrap();
370 assert_eq!(storage.load().unwrap(), Some("token-1".to_string()));
371
372 storage.store("token-2").unwrap();
373 assert_eq!(storage.load().unwrap(), Some("token-2".to_string()));
374 }
375
376 #[test]
377 fn test_file_storage_clear_nonexistent() {
378 let temp_dir = TempDir::new().unwrap();
379 let token_path = temp_dir.path().join("nonexistent_token");
380 let storage = FileTokenStorage::with_path(token_path);
381
382 assert!(storage.clear().is_ok());
384 }
385
386 #[test]
387 fn test_file_storage_path() {
388 let path = PathBuf::from("/custom/path/token");
389 let storage = FileTokenStorage::with_path(path.clone());
390 assert_eq!(storage.path(), &path);
391 }
392
393 #[test]
394 fn test_storage_error_display() {
395 let err = StorageError::NotAvailable("test".to_string());
396 assert!(err.to_string().contains("test"));
397 assert!(err.to_string().contains("not available"));
398
399 let err = StorageError::ReadError("read failed".to_string());
400 assert!(err.to_string().contains("read failed"));
401
402 let err = StorageError::WriteError("write failed".to_string());
403 assert!(err.to_string().contains("write failed"));
404
405 let err = StorageError::ClearError("clear failed".to_string());
406 assert!(err.to_string().contains("clear failed"));
407 }
408}