use std::fmt;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::pin::Pin;
pub trait StorageBackend: Send + Sync {
fn write_file_str(
&self,
path: &str,
data: &[u8],
) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>>;
fn read_file_str(
&self,
path: &str,
) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, StorageError>> + Send + '_>>;
fn exists_str(&self, path: &str) -> bool;
fn remove_str(
&self,
path: &str,
) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>>;
}
#[derive(Debug)]
pub enum StorageError {
Io(std::io::Error),
Config(String),
Path(String),
Encryption(String),
KeyGeneration(String),
KeyStorage(String),
}
impl fmt::Display for StorageError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
StorageError::Io(e) => write!(f, "IO error: {}", e),
StorageError::Config(msg) => write!(f, "Configuration error: {}", msg),
StorageError::Path(msg) => write!(f, "Path error: {}", msg),
StorageError::Encryption(msg) => write!(f, "Encryption error: {}", msg),
StorageError::KeyGeneration(msg) => write!(f, "Key generation error: {}", msg),
StorageError::KeyStorage(msg) => write!(f, "Key storage error: {}", msg),
}
}
}
impl std::error::Error for StorageError {}
impl From<std::io::Error> for StorageError {
fn from(err: std::io::Error) -> Self {
StorageError::Io(err)
}
}
pub struct FilesystemStorage {
base_path: PathBuf,
}
impl FilesystemStorage {
pub fn new(base_path: impl AsRef<Path>) -> Result<Self, StorageError> {
let base_path = base_path.as_ref().to_path_buf();
if !base_path.exists() {
std::fs::create_dir_all(&base_path)?;
}
Ok(Self { base_path })
}
fn resolve_path(&self, path: &str) -> PathBuf {
self.base_path.join(path)
}
}
impl StorageBackend for FilesystemStorage {
fn write_file_str(
&self,
path: &str,
data: &[u8],
) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
let full_path = self.resolve_path(path);
let data = data.to_vec();
Box::pin(async move {
if let Some(parent) = full_path.parent() {
tokio::fs::create_dir_all(parent).await?;
}
tokio::fs::write(&full_path, data).await?;
tracing::debug!("Wrote data to filesystem: {:?}", full_path);
Ok(())
})
}
fn read_file_str(
&self,
path: &str,
) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, StorageError>> + Send + '_>> {
let full_path = self.resolve_path(path);
Box::pin(async move {
let data = tokio::fs::read(&full_path).await?;
tracing::debug!("Read data from filesystem: {:?}", full_path);
Ok(data)
})
}
fn exists_str(&self, path: &str) -> bool {
let full_path = self.resolve_path(path);
full_path.exists()
}
fn remove_str(
&self,
path: &str,
) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
let full_path = self.resolve_path(path);
Box::pin(async move {
tokio::fs::remove_file(&full_path).await?;
tracing::debug!("Removed file from filesystem: {:?}", full_path);
Ok(())
})
}
}
pub struct EncryptedFilesystemStorage {
base_path: PathBuf,
recipient: age::x25519::Recipient,
identity: age::x25519::Identity,
}
impl EncryptedFilesystemStorage {
pub async fn new_with_instance(instance_id: &str) -> Result<Self, StorageError> {
let home = dirs::home_dir().ok_or_else(|| {
StorageError::KeyStorage("Cannot determine home directory".to_string())
})?;
let base_path = home.join(".runbeam").join(instance_id);
Self::new_with_key_path(base_path.clone(), base_path.join("encryption.key")).await
}
pub async fn new_with_instance_and_key(
instance_id: &str,
encryption_key: &str,
) -> Result<Self, StorageError> {
let home = dirs::home_dir().ok_or_else(|| {
StorageError::KeyStorage("Cannot determine home directory".to_string())
})?;
let base_path = home.join(".runbeam").join(instance_id);
if !base_path.exists() {
tokio::fs::create_dir_all(&base_path).await?;
}
let (recipient, identity) = Self::load_key_from_string(encryption_key)?;
Ok(Self {
base_path,
recipient,
identity,
})
}
pub async fn new(base_path: impl AsRef<Path>) -> Result<Self, StorageError> {
let base_path = base_path.as_ref().to_path_buf();
let key_path = Self::get_key_path()?;
Self::new_with_key_path(base_path, key_path).await
}
async fn new_with_key_path(
base_path: PathBuf,
key_path: PathBuf,
) -> Result<Self, StorageError> {
if !base_path.exists() {
tokio::fs::create_dir_all(&base_path).await?;
}
let (recipient, identity) = Self::setup_encryption_with_path(&key_path).await?;
Ok(Self {
base_path,
recipient,
identity,
})
}
async fn setup_encryption_with_path(
key_path: &Path,
) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
if let Ok(key_base64) = std::env::var("RUNBEAM_ENCRYPTION_KEY") {
tracing::debug!(
"Using encryption key from RUNBEAM_ENCRYPTION_KEY environment variable"
);
return Self::load_key_from_string(&key_base64);
}
if key_path.exists() {
tracing::debug!("Loading existing encryption key from {:?}", key_path);
Self::load_key_from_file(key_path).await
} else {
tracing::info!(
"Generating new encryption key and storing at {:?}",
key_path
);
Self::generate_and_store_key(key_path).await
}
}
fn get_key_path() -> Result<PathBuf, StorageError> {
let home = dirs::home_dir().ok_or_else(|| {
StorageError::KeyStorage("Cannot determine home directory".to_string())
})?;
let key_dir = home.join(".runbeam");
Ok(key_dir.join("encryption.key"))
}
fn load_key_from_string(
key_input: &str,
) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
use base64::{engine::general_purpose, Engine as _};
let key_str = key_input.trim();
if let Ok(identity) = key_str.parse::<age::x25519::Identity>() {
tracing::debug!("Loaded age key directly (raw format)");
let recipient = identity.to_public();
return Ok((recipient, identity));
}
match general_purpose::STANDARD.decode(key_str) {
Ok(key_bytes) => {
let decoded_str = String::from_utf8(key_bytes).map_err(|e| {
StorageError::KeyStorage(format!("Invalid UTF-8 in base64-decoded key: {}", e))
})?;
let identity = decoded_str.parse::<age::x25519::Identity>().map_err(|e| {
StorageError::KeyStorage(format!(
"Invalid age identity after base64 decode: {}",
e
))
})?;
tracing::debug!("Loaded age key from base64-encoded format");
let recipient = identity.to_public();
Ok((recipient, identity))
}
Err(_) => Err(StorageError::KeyStorage(
"Key is neither a valid age identity nor valid base64-encoded age identity"
.to_string(),
)),
}
}
async fn load_key_from_file(
key_path: &Path,
) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
let key_contents = tokio::fs::read_to_string(key_path)
.await
.map_err(|e| StorageError::KeyStorage(format!("Failed to read key file: {}", e)))?;
Self::load_key_from_string(&key_contents)
}
async fn generate_and_store_key(
key_path: &Path,
) -> Result<(age::x25519::Recipient, age::x25519::Identity), StorageError> {
use base64::{engine::general_purpose, Engine as _};
use secrecy::ExposeSecret;
let identity = age::x25519::Identity::generate();
let identity_str = identity.to_string();
let identity_str_exposed = identity_str.expose_secret();
let key_base64 = general_purpose::STANDARD.encode(identity_str_exposed.as_bytes());
if let Some(parent) = key_path.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|e| {
StorageError::KeyStorage(format!("Failed to create key directory: {}", e))
})?;
}
tokio::fs::write(key_path, &key_base64)
.await
.map_err(|e| StorageError::KeyStorage(format!("Failed to write key file: {}", e)))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let permissions = std::fs::Permissions::from_mode(0o600);
std::fs::set_permissions(key_path, permissions).map_err(|e| {
StorageError::KeyStorage(format!("Failed to set key file permissions: {}", e))
})?;
tracing::debug!("Set encryption key file permissions to 0600");
}
tracing::info!("Generated and stored new encryption key");
let recipient = identity.to_public();
Ok((recipient, identity))
}
fn resolve_path(&self, path: &str) -> PathBuf {
self.base_path.join(path)
}
fn encrypt_data(&self, data: &[u8]) -> Result<Vec<u8>, StorageError> {
use std::io::Write;
let encryptor = age::Encryptor::with_recipients(vec![Box::new(self.recipient.clone())])
.expect("Failed to create encryptor with recipient");
let mut encrypted = Vec::new();
let mut writer = encryptor
.wrap_output(&mut encrypted)
.map_err(|e| StorageError::Encryption(format!("Failed to wrap output: {}", e)))?;
writer
.write_all(data)
.map_err(|e| StorageError::Encryption(format!("Failed to encrypt data: {}", e)))?;
writer.finish().map_err(|e| {
StorageError::Encryption(format!("Failed to finalize encryption: {}", e))
})?;
Ok(encrypted)
}
fn decrypt_data(&self, encrypted: &[u8]) -> Result<Vec<u8>, StorageError> {
use std::io::Read;
let decryptor = match age::Decryptor::new(encrypted)
.map_err(|e| StorageError::Encryption(format!("Failed to create decryptor: {}", e)))?
{
age::Decryptor::Recipients(d) => d,
_ => {
return Err(StorageError::Encryption(
"Unexpected decryptor type".to_string(),
))
}
};
let mut decrypted = Vec::new();
let mut reader = decryptor
.decrypt(std::iter::once(&self.identity as &dyn age::Identity))
.map_err(|e| StorageError::Encryption(format!("Failed to decrypt data: {}", e)))?;
reader.read_to_end(&mut decrypted).map_err(|e| {
StorageError::Encryption(format!("Failed to read decrypted data: {}", e))
})?;
Ok(decrypted)
}
}
impl StorageBackend for EncryptedFilesystemStorage {
fn write_file_str(
&self,
path: &str,
data: &[u8],
) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
let full_path = self.resolve_path(path);
let data = data.to_vec();
Box::pin(async move {
let encrypted = self.encrypt_data(&data)?;
if let Some(parent) = full_path.parent() {
tokio::fs::create_dir_all(parent).await?;
}
tokio::fs::write(&full_path, encrypted).await?;
tracing::debug!("Wrote encrypted data to filesystem: {:?}", full_path);
Ok(())
})
}
fn read_file_str(
&self,
path: &str,
) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, StorageError>> + Send + '_>> {
let full_path = self.resolve_path(path);
Box::pin(async move {
let encrypted = tokio::fs::read(&full_path).await?;
let decrypted = self.decrypt_data(&encrypted)?;
tracing::debug!("Read and decrypted data from filesystem: {:?}", full_path);
Ok(decrypted)
})
}
fn exists_str(&self, path: &str) -> bool {
let full_path = self.resolve_path(path);
full_path.exists()
}
fn remove_str(
&self,
path: &str,
) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
let full_path = self.resolve_path(path);
Box::pin(async move {
tokio::fs::remove_file(&full_path).await?;
tracing::debug!("Removed encrypted file from filesystem: {:?}", full_path);
Ok(())
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use serial_test::serial;
use tempfile::TempDir;
#[tokio::test]
async fn test_filesystem_storage_write_and_read() {
let temp_dir = TempDir::new().unwrap();
let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
let test_data = b"test data";
storage.write_file_str("test.txt", test_data).await.unwrap();
assert!(storage.exists_str("test.txt"));
let read_data = storage.read_file_str("test.txt").await.unwrap();
assert_eq!(read_data, test_data);
}
#[tokio::test]
async fn test_filesystem_storage_nested_paths() {
let temp_dir = TempDir::new().unwrap();
let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
let test_data = b"nested data";
storage
.write_file_str("nested/path/test.txt", test_data)
.await
.unwrap();
assert!(storage.exists_str("nested/path/test.txt"));
let read_data = storage.read_file_str("nested/path/test.txt").await.unwrap();
assert_eq!(read_data, test_data);
}
#[tokio::test]
async fn test_filesystem_storage_remove() {
let temp_dir = TempDir::new().unwrap();
let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
let test_data = b"test data";
storage.write_file_str("test.txt", test_data).await.unwrap();
assert!(storage.exists_str("test.txt"));
storage.remove_str("test.txt").await.unwrap();
assert!(!storage.exists_str("test.txt"));
}
#[tokio::test]
async fn test_filesystem_storage_exists_nonexistent() {
let temp_dir = TempDir::new().unwrap();
let storage = FilesystemStorage::new(temp_dir.path()).unwrap();
assert!(!storage.exists_str("nonexistent.txt"));
}
#[tokio::test]
async fn test_encrypted_storage_write_and_read() {
let temp_dir = TempDir::new().unwrap();
let storage = EncryptedFilesystemStorage::new(temp_dir.path())
.await
.unwrap();
let test_data = b"sensitive data";
storage
.write_file_str("secret.txt", test_data)
.await
.unwrap();
assert!(storage.exists_str("secret.txt"));
let read_data = storage.read_file_str("secret.txt").await.unwrap();
assert_eq!(read_data, test_data);
}
#[tokio::test]
async fn test_encrypted_storage_encryption_roundtrip() {
let temp_dir = TempDir::new().unwrap();
let storage = EncryptedFilesystemStorage::new(temp_dir.path())
.await
.unwrap();
let test_data = b"This should be encrypted";
storage.write_file_str("data.bin", test_data).await.unwrap();
let file_path = temp_dir.path().join("data.bin");
let raw_contents = std::fs::read(&file_path).unwrap();
assert_ne!(raw_contents.as_slice(), test_data);
assert!(raw_contents.len() > test_data.len());
let decrypted = storage.read_file_str("data.bin").await.unwrap();
assert_eq!(decrypted, test_data);
}
#[tokio::test]
async fn test_encrypted_storage_remove() {
let temp_dir = TempDir::new().unwrap();
let storage = EncryptedFilesystemStorage::new(temp_dir.path())
.await
.unwrap();
let test_data = b"test";
storage.write_file_str("file.txt", test_data).await.unwrap();
assert!(storage.exists_str("file.txt"));
storage.remove_str("file.txt").await.unwrap();
assert!(!storage.exists_str("file.txt"));
}
#[tokio::test]
#[serial]
async fn test_encrypted_storage_key_persistence() {
use std::env;
let identity = age::x25519::Identity::generate();
let identity_str = identity.to_string();
use base64::Engine;
use secrecy::ExposeSecret;
let key_base64 = base64::engine::general_purpose::STANDARD
.encode(identity_str.expose_secret().as_bytes());
env::set_var("RUNBEAM_ENCRYPTION_KEY", &key_base64);
let temp_dir = TempDir::new().unwrap();
let storage1 = EncryptedFilesystemStorage::new(temp_dir.path())
.await
.unwrap();
let test_data = b"persistent test";
storage1
.write_file_str("data.txt", test_data)
.await
.unwrap();
drop(storage1);
let storage2 = EncryptedFilesystemStorage::new(temp_dir.path())
.await
.unwrap();
let read_data = storage2.read_file_str("data.txt").await.unwrap();
assert_eq!(read_data, test_data);
env::remove_var("RUNBEAM_ENCRYPTION_KEY");
}
#[tokio::test]
#[serial]
async fn test_encrypted_storage_env_var_key() {
use base64::Engine;
use std::env;
let identity = age::x25519::Identity::generate();
let identity_str = identity.to_string();
use secrecy::ExposeSecret;
let key_base64 = base64::engine::general_purpose::STANDARD
.encode(identity_str.expose_secret().as_bytes());
env::set_var("RUNBEAM_ENCRYPTION_KEY", &key_base64);
let temp_dir = TempDir::new().unwrap();
let storage = EncryptedFilesystemStorage::new(temp_dir.path())
.await
.unwrap();
let test_data = b"env var test";
storage.write_file_str("test.bin", test_data).await.unwrap();
let read_data = storage.read_file_str("test.bin").await.unwrap();
assert_eq!(read_data, test_data);
env::remove_var("RUNBEAM_ENCRYPTION_KEY");
}
#[tokio::test]
#[serial]
#[cfg(unix)]
async fn test_encrypted_storage_key_file_permissions() {
use std::os::unix::fs::PermissionsExt;
let temp_dir = TempDir::new().unwrap();
std::env::remove_var("RUNBEAM_ENCRYPTION_KEY");
let _storage = EncryptedFilesystemStorage::new(temp_dir.path())
.await
.unwrap();
let key_path = dirs::home_dir().unwrap().join(".runbeam/encryption.key");
if key_path.exists() {
let metadata = std::fs::metadata(&key_path).unwrap();
let permissions = metadata.permissions();
let mode = permissions.mode();
assert_eq!(mode & 0o777, 0o600, "Key file should have 0600 permissions");
}
}
}