1use std::fmt;
2use std::future::Future;
3use std::path::{Path, PathBuf};
4use std::pin::Pin;
5
6pub trait StorageBackend: Send + Sync {
11 fn write_file_str(
13 &self,
14 path: &str,
15 data: &[u8],
16 ) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>>;
17
18 fn read_file_str(
20 &self,
21 path: &str,
22 ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, StorageError>> + Send + '_>>;
23
24 fn exists_str(&self, path: &str) -> bool;
26
27 fn remove_str(
29 &self,
30 path: &str,
31 ) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>>;
32}
33
34#[derive(Debug)]
36pub enum StorageError {
37 Io(std::io::Error),
39 Config(String),
41 Path(String),
43 Encryption(String),
45 KeyGeneration(String),
47 KeyStorage(String),
49}
50
51impl fmt::Display for StorageError {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 match self {
54 StorageError::Io(e) => write!(f, "IO error: {}", e),
55 StorageError::Config(msg) => write!(f, "Configuration error: {}", msg),
56 StorageError::Path(msg) => write!(f, "Path error: {}", msg),
57 StorageError::Encryption(msg) => write!(f, "Encryption error: {}", msg),
58 StorageError::KeyGeneration(msg) => write!(f, "Key generation error: {}", msg),
59 StorageError::KeyStorage(msg) => write!(f, "Key storage error: {}", msg),
60 }
61 }
62}
63
64impl std::error::Error for StorageError {}
65
66impl From<std::io::Error> for StorageError {
67 fn from(err: std::io::Error) -> Self {
68 StorageError::Io(err)
69 }
70}
71
72pub struct FilesystemStorage {
88 base_path: PathBuf,
89}
90
91impl FilesystemStorage {
92 pub fn new(base_path: impl AsRef<Path>) -> Result<Self, StorageError> {
94 let base_path = base_path.as_ref().to_path_buf();
95
96 if !base_path.exists() {
97 std::fs::create_dir_all(&base_path)?;
98 }
99
100 Ok(Self { base_path })
101 }
102
103 fn resolve_path(&self, path: &str) -> PathBuf {
105 self.base_path.join(path)
106 }
107}
108
109impl StorageBackend for FilesystemStorage {
110 fn write_file_str(
111 &self,
112 path: &str,
113 data: &[u8],
114 ) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
115 let full_path = self.resolve_path(path);
116 let data = data.to_vec();
117
118 Box::pin(async move {
119 if let Some(parent) = full_path.parent() {
121 tokio::fs::create_dir_all(parent).await?;
122 }
123
124 tokio::fs::write(&full_path, data).await?;
125 tracing::debug!("Wrote data to filesystem: {:?}", full_path);
126 Ok(())
127 })
128 }
129
130 fn read_file_str(
131 &self,
132 path: &str,
133 ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, StorageError>> + Send + '_>> {
134 let full_path = self.resolve_path(path);
135
136 Box::pin(async move {
137 let data = tokio::fs::read(&full_path).await?;
138 tracing::debug!("Read data from filesystem: {:?}", full_path);
139 Ok(data)
140 })
141 }
142
143 fn exists_str(&self, path: &str) -> bool {
144 let full_path = self.resolve_path(path);
145 full_path.exists()
146 }
147
148 fn remove_str(
149 &self,
150 path: &str,
151 ) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
152 let full_path = self.resolve_path(path);
153
154 Box::pin(async move {
155 tokio::fs::remove_file(&full_path).await?;
156 tracing::debug!("Removed file from filesystem: {:?}", full_path);
157 Ok(())
158 })
159 }
160}
161
162pub struct EncryptedFilesystemStorage {
187 base_path: PathBuf,
188 recipient: age::x25519::Recipient,
189 identity: age::x25519::Identity,
190}
191
192impl EncryptedFilesystemStorage {
193 pub async fn new_with_instance(instance_id: &str) -> Result<Self, StorageError> {
209 let home = dirs::home_dir().ok_or_else(|| {
210 StorageError::KeyStorage("Cannot determine home directory".to_string())
211 })?;
212
213 let base_path = home.join(".runbeam").join(instance_id);
214 Self::new_with_key_path(base_path.clone(), base_path.join("encryption.key")).await
215 }
216
217 pub async fn new_with_instance_and_key(
233 instance_id: &str,
234 encryption_key: &str,
235 ) -> Result<Self, StorageError> {
236 let home = dirs::home_dir().ok_or_else(|| {
237 StorageError::KeyStorage("Cannot determine home directory".to_string())
238 })?;
239
240 let base_path = home.join(".runbeam").join(instance_id);
241
242 if !base_path.exists() {
244 tokio::fs::create_dir_all(&base_path).await?;
245 }
246
247 let (recipient, identity) = Self::load_key_from_string(encryption_key)?;
249
250 Ok(Self {
251 base_path,
252 recipient,
253 identity,
254 })
255 }
256
257 pub async fn new(base_path: impl AsRef<Path>) -> Result<Self, StorageError> {
270 let base_path = base_path.as_ref().to_path_buf();
271 let key_path = Self::get_key_path()?;
272 Self::new_with_key_path(base_path, key_path).await
273 }
274
275 async fn new_with_key_path(
277 base_path: PathBuf,
278 key_path: PathBuf,
279 ) -> Result<Self, StorageError> {
280 if !base_path.exists() {
282 tokio::fs::create_dir_all(&base_path).await?;
283 }
284
285 let (recipient, identity) = Self::setup_encryption_with_path(&key_path).await?;
287
288 Ok(Self {
289 base_path,
290 recipient,
291 identity,
292 })
293 }
294
295 async fn setup_encryption_with_path(
297 key_path: &Path,
298 ) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
299 if let Ok(key_base64) = std::env::var("RUNBEAM_ENCRYPTION_KEY") {
301 tracing::debug!(
302 "Using encryption key from RUNBEAM_ENCRYPTION_KEY environment variable"
303 );
304 return Self::load_key_from_string(&key_base64);
305 }
306
307 if key_path.exists() {
309 tracing::debug!("Loading existing encryption key from {:?}", key_path);
310 Self::load_key_from_file(key_path).await
311 } else {
312 tracing::info!(
313 "Generating new encryption key and storing at {:?}",
314 key_path
315 );
316 Self::generate_and_store_key(key_path).await
317 }
318 }
319
320 fn get_key_path() -> Result<PathBuf, StorageError> {
322 let home = dirs::home_dir().ok_or_else(|| {
323 StorageError::KeyStorage("Cannot determine home directory".to_string())
324 })?;
325
326 let key_dir = home.join(".runbeam");
327 Ok(key_dir.join("encryption.key"))
328 }
329
330 fn load_key_from_string(
339 key_input: &str,
340 ) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
341 use base64::{engine::general_purpose, Engine as _};
342
343 let key_str = key_input.trim();
344
345 if let Ok(identity) = key_str.parse::<age::x25519::Identity>() {
347 tracing::debug!("Loaded age key directly (raw format)");
348 let recipient = identity.to_public();
349 return Ok((recipient, identity));
350 }
351
352 match general_purpose::STANDARD.decode(key_str) {
354 Ok(key_bytes) => {
355 let decoded_str = String::from_utf8(key_bytes).map_err(|e| {
356 StorageError::KeyStorage(format!("Invalid UTF-8 in base64-decoded key: {}", e))
357 })?;
358
359 let identity = decoded_str.parse::<age::x25519::Identity>().map_err(|e| {
360 StorageError::KeyStorage(format!(
361 "Invalid age identity after base64 decode: {}",
362 e
363 ))
364 })?;
365
366 tracing::debug!("Loaded age key from base64-encoded format");
367 let recipient = identity.to_public();
368 Ok((recipient, identity))
369 }
370 Err(_) => Err(StorageError::KeyStorage(
371 "Key is neither a valid age identity nor valid base64-encoded age identity"
372 .to_string(),
373 )),
374 }
375 }
376
377 async fn load_key_from_file(
379 key_path: &Path,
380 ) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
381 let key_contents = tokio::fs::read_to_string(key_path)
382 .await
383 .map_err(|e| StorageError::KeyStorage(format!("Failed to read key file: {}", e)))?;
384
385 Self::load_key_from_string(&key_contents)
386 }
387
388 async fn generate_and_store_key(
390 key_path: &Path,
391 ) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
392 use base64::{engine::general_purpose, Engine as _};
393 use secrecy::ExposeSecret;
394
395 let identity = age::x25519::Identity::generate();
397 let identity_str = identity.to_string();
398
399 let identity_str_exposed = identity_str.expose_secret();
401 let key_base64 = general_purpose::STANDARD.encode(identity_str_exposed.as_bytes());
402
403 if let Some(parent) = key_path.parent() {
405 tokio::fs::create_dir_all(parent).await.map_err(|e| {
406 StorageError::KeyStorage(format!("Failed to create key directory: {}", e))
407 })?;
408 }
409
410 tokio::fs::write(key_path, &key_base64)
412 .await
413 .map_err(|e| StorageError::KeyStorage(format!("Failed to write key file: {}", e)))?;
414
415 #[cfg(unix)]
417 {
418 use std::os::unix::fs::PermissionsExt;
419 let permissions = std::fs::Permissions::from_mode(0o600);
420 std::fs::set_permissions(key_path, permissions).map_err(|e| {
421 StorageError::KeyStorage(format!("Failed to set key file permissions: {}", e))
422 })?;
423 tracing::debug!("Set encryption key file permissions to 0600");
424 }
425
426 tracing::info!("Generated and stored new encryption key");
427
428 let recipient = identity.to_public();
430
431 Ok((recipient, identity))
432 }
433
434 fn resolve_path(&self, path: &str) -> PathBuf {
436 self.base_path.join(path)
437 }
438
439 fn encrypt_data(&self, data: &[u8]) -> Result<Vec<u8>, StorageError> {
441 use std::io::Write;
442
443 let encryptor = age::Encryptor::with_recipients(vec![Box::new(self.recipient.clone())])
444 .expect("Failed to create encryptor with recipient");
445
446 let mut encrypted = Vec::new();
447 let mut writer = encryptor
448 .wrap_output(&mut encrypted)
449 .map_err(|e| StorageError::Encryption(format!("Failed to wrap output: {}", e)))?;
450
451 writer
452 .write_all(data)
453 .map_err(|e| StorageError::Encryption(format!("Failed to encrypt data: {}", e)))?;
454
455 writer.finish().map_err(|e| {
456 StorageError::Encryption(format!("Failed to finalize encryption: {}", e))
457 })?;
458
459 Ok(encrypted)
460 }
461
462 fn decrypt_data(&self, encrypted: &[u8]) -> Result<Vec<u8>, StorageError> {
464 use std::io::Read;
465
466 let decryptor = match age::Decryptor::new(encrypted)
467 .map_err(|e| StorageError::Encryption(format!("Failed to create decryptor: {}", e)))?
468 {
469 age::Decryptor::Recipients(d) => d,
470 _ => {
471 return Err(StorageError::Encryption(
472 "Unexpected decryptor type".to_string(),
473 ))
474 }
475 };
476
477 let mut decrypted = Vec::new();
478 let mut reader = decryptor
479 .decrypt(std::iter::once(&self.identity as &dyn age::Identity))
480 .map_err(|e| StorageError::Encryption(format!("Failed to decrypt data: {}", e)))?;
481
482 reader.read_to_end(&mut decrypted).map_err(|e| {
483 StorageError::Encryption(format!("Failed to read decrypted data: {}", e))
484 })?;
485
486 Ok(decrypted)
487 }
488}
489
490impl StorageBackend for EncryptedFilesystemStorage {
491 fn write_file_str(
492 &self,
493 path: &str,
494 data: &[u8],
495 ) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
496 let full_path = self.resolve_path(path);
497 let data = data.to_vec();
498
499 Box::pin(async move {
500 let encrypted = self.encrypt_data(&data)?;
502
503 if let Some(parent) = full_path.parent() {
505 tokio::fs::create_dir_all(parent).await?;
506 }
507
508 tokio::fs::write(&full_path, encrypted).await?;
510 tracing::debug!("Wrote encrypted data to filesystem: {:?}", full_path);
511 Ok(())
512 })
513 }
514
515 fn read_file_str(
516 &self,
517 path: &str,
518 ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, StorageError>> + Send + '_>> {
519 let full_path = self.resolve_path(path);
520
521 Box::pin(async move {
522 let encrypted = tokio::fs::read(&full_path).await?;
524
525 let decrypted = self.decrypt_data(&encrypted)?;
527 tracing::debug!("Read and decrypted data from filesystem: {:?}", full_path);
528 Ok(decrypted)
529 })
530 }
531
532 fn exists_str(&self, path: &str) -> bool {
533 let full_path = self.resolve_path(path);
534 full_path.exists()
535 }
536
537 fn remove_str(
538 &self,
539 path: &str,
540 ) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
541 let full_path = self.resolve_path(path);
542
543 Box::pin(async move {
544 tokio::fs::remove_file(&full_path).await?;
545 tracing::debug!("Removed encrypted file from filesystem: {:?}", full_path);
546 Ok(())
547 })
548 }
549}
550
551#[cfg(test)]
552mod tests {
553 use super::*;
554 use serial_test::serial;
555 use tempfile::TempDir;
556
557 #[tokio::test]
558 async fn test_filesystem_storage_write_and_read() {
559 let temp_dir = TempDir::new().unwrap();
560 let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
561
562 let test_data = b"test data";
563 storage.write_file_str("test.txt", test_data).await.unwrap();
564
565 assert!(storage.exists_str("test.txt"));
566
567 let read_data = storage.read_file_str("test.txt").await.unwrap();
568 assert_eq!(read_data, test_data);
569 }
570
571 #[tokio::test]
572 async fn test_filesystem_storage_nested_paths() {
573 let temp_dir = TempDir::new().unwrap();
574 let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
575
576 let test_data = b"nested data";
577 storage
578 .write_file_str("nested/path/test.txt", test_data)
579 .await
580 .unwrap();
581
582 assert!(storage.exists_str("nested/path/test.txt"));
583
584 let read_data = storage.read_file_str("nested/path/test.txt").await.unwrap();
585 assert_eq!(read_data, test_data);
586 }
587
588 #[tokio::test]
589 async fn test_filesystem_storage_remove() {
590 let temp_dir = TempDir::new().unwrap();
591 let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
592
593 let test_data = b"test data";
594 storage.write_file_str("test.txt", test_data).await.unwrap();
595 assert!(storage.exists_str("test.txt"));
596
597 storage.remove_str("test.txt").await.unwrap();
598 assert!(!storage.exists_str("test.txt"));
599 }
600
601 #[tokio::test]
602 async fn test_filesystem_storage_exists_nonexistent() {
603 let temp_dir = TempDir::new().unwrap();
604 let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
605
606 assert!(!storage.exists_str("nonexistent.txt"));
607 }
608
609 #[tokio::test]
610 async fn test_encrypted_storage_write_and_read() {
611 let temp_dir = TempDir::new().unwrap();
612 let storage = EncryptedFilesystemStorage::new(temp_dir.path())
613 .await
614 .unwrap();
615
616 let test_data = b"sensitive data";
617 storage
618 .write_file_str("secret.txt", test_data)
619 .await
620 .unwrap();
621
622 assert!(storage.exists_str("secret.txt"));
623
624 let read_data = storage.read_file_str("secret.txt").await.unwrap();
625 assert_eq!(read_data, test_data);
626 }
627
628 #[tokio::test]
629 async fn test_encrypted_storage_encryption_roundtrip() {
630 let temp_dir = TempDir::new().unwrap();
631 let storage = EncryptedFilesystemStorage::new(temp_dir.path())
632 .await
633 .unwrap();
634
635 let test_data = b"This should be encrypted";
636 storage.write_file_str("data.bin", test_data).await.unwrap();
637
638 let file_path = temp_dir.path().join("data.bin");
640 let raw_contents = std::fs::read(&file_path).unwrap();
641
642 assert_ne!(raw_contents.as_slice(), test_data);
644 assert!(raw_contents.len() > test_data.len());
646
647 let decrypted = storage.read_file_str("data.bin").await.unwrap();
649 assert_eq!(decrypted, test_data);
650 }
651
652 #[tokio::test]
653 async fn test_encrypted_storage_remove() {
654 let temp_dir = TempDir::new().unwrap();
655 let storage = EncryptedFilesystemStorage::new(temp_dir.path())
656 .await
657 .unwrap();
658
659 let test_data = b"test";
660 storage.write_file_str("file.txt", test_data).await.unwrap();
661 assert!(storage.exists_str("file.txt"));
662
663 storage.remove_str("file.txt").await.unwrap();
664 assert!(!storage.exists_str("file.txt"));
665 }
666
667 #[tokio::test]
668 #[serial]
669 async fn test_encrypted_storage_key_persistence() {
670 use std::env;
671
672 let identity = age::x25519::Identity::generate();
674 let identity_str = identity.to_string();
675 use base64::Engine;
676 use secrecy::ExposeSecret;
677 let key_base64 = base64::engine::general_purpose::STANDARD
678 .encode(identity_str.expose_secret().as_bytes());
679
680 env::set_var("RUNBEAM_ENCRYPTION_KEY", &key_base64);
682
683 let temp_dir = TempDir::new().unwrap();
684
685 let storage1 = EncryptedFilesystemStorage::new(temp_dir.path())
687 .await
688 .unwrap();
689
690 let test_data = b"persistent test";
691 storage1
692 .write_file_str("data.txt", test_data)
693 .await
694 .unwrap();
695
696 drop(storage1);
698
699 let storage2 = EncryptedFilesystemStorage::new(temp_dir.path())
701 .await
702 .unwrap();
703
704 let read_data = storage2.read_file_str("data.txt").await.unwrap();
706 assert_eq!(read_data, test_data);
707
708 env::remove_var("RUNBEAM_ENCRYPTION_KEY");
710 }
711
712 #[tokio::test]
713 #[serial]
714 async fn test_encrypted_storage_env_var_key() {
715 use base64::Engine;
716 use std::env;
717
718 let identity = age::x25519::Identity::generate();
720 let identity_str = identity.to_string();
721 use secrecy::ExposeSecret;
722 let key_base64 = base64::engine::general_purpose::STANDARD
723 .encode(identity_str.expose_secret().as_bytes());
724
725 env::set_var("RUNBEAM_ENCRYPTION_KEY", &key_base64);
727
728 let temp_dir = TempDir::new().unwrap();
729 let storage = EncryptedFilesystemStorage::new(temp_dir.path())
730 .await
731 .unwrap();
732
733 let test_data = b"env var test";
734 storage.write_file_str("test.bin", test_data).await.unwrap();
735 let read_data = storage.read_file_str("test.bin").await.unwrap();
736
737 assert_eq!(read_data, test_data);
738
739 env::remove_var("RUNBEAM_ENCRYPTION_KEY");
741 }
742
743 #[tokio::test]
744 #[serial]
745 #[cfg(unix)]
746 async fn test_encrypted_storage_key_file_permissions() {
747 use std::os::unix::fs::PermissionsExt;
748
749 let temp_dir = TempDir::new().unwrap();
750
751 std::env::remove_var("RUNBEAM_ENCRYPTION_KEY");
753
754 let _storage = EncryptedFilesystemStorage::new(temp_dir.path())
755 .await
756 .unwrap();
757
758 let key_path = dirs::home_dir().unwrap().join(".runbeam/encryption.key");
760
761 if key_path.exists() {
762 let metadata = std::fs::metadata(&key_path).unwrap();
763 let permissions = metadata.permissions();
764 let mode = permissions.mode();
765
766 assert_eq!(mode & 0o777, 0o600, "Key file should have 0600 permissions");
768 }
769 }
770}