use std::{ io::ErrorKind, ops::DerefMut, path::{Path, PathBuf}, }; use log::{debug, error, info}; use rand::{ distributions::{Alphanumeric, DistString}, thread_rng, Rng, }; use time::OffsetDateTime; use tokio::fs::File; const MAX_STORAGE_FILES: usize = 1024; pub fn gen_storage_code(use_mnemonic: bool) -> String { if use_mnemonic { mnemonic::to_string(thread_rng().gen::<[u8; 4]>()) } else { Alphanumeric.sample_string(&mut thread_rng(), 8) } } pub fn is_valid_storage_code(s: &str) -> bool { s.as_bytes() .iter() .all(|c| c.is_ascii_alphanumeric() || c == &b'-') } pub use crate::state::v1::{StoredFile, StoredFiles}; async fn is_valid_entry(key: &str, info: &StoredFile, storage_dir: &Path) -> bool { if info.expiry < OffsetDateTime::now_utc() { info!("File {} has expired", key); return false; } let file = if let Ok(f) = File::open(storage_dir.join(&key)).await { f } else { error!( "Unable to open file {} referenced in persistent storage", key ); return false; }; let metadata = if let Ok(md) = file.metadata().await { md } else { error!( "Unable to get metadata for file {} referenced in persistent storage", key ); return false; }; if metadata.len() != info.size { error!("Mismatched file size for file {} referenced in persistent storage: expected {}, found {}", key, info.size, metadata.len()); return false; } true } async fn delete_file_if_exists(file: &PathBuf) -> std::io::Result<()> { if let Err(e) = tokio::fs::remove_file(file).await { if e.kind() != ErrorKind::NotFound { error!("Failed to delete file {}: {}", file.to_string_lossy(), e); return Err(e); } } Ok(()) } #[derive(thiserror::Error, Debug)] pub enum FileAddError { #[error("Failed to write metadata to filesystem")] FileSystem(#[from] std::io::Error), #[error("File was too large, available space is {0} bytes")] TooBig(u64), #[error("File store is full")] Full, } impl crate::AppData { // This must only be run at startup, or it will delete in-progress // uploads. pub async fn cleanup(&self) -> std::io::Result<()> { let mut store = self.state.write().await; let old = std::mem::take(store.deref_mut()); for (key, info) in old.0.into_iter() { // Handle this case separately, because we don't // want to try to delete it if it's not the sort // of path we're expecting if !is_valid_storage_code(&key) { error!("Invalid key in persistent storage: {}", key); continue; } if is_valid_entry(&key, &info, &self.config.storage_dir).await { store.0.insert(key, info); } else { info!("Deleting file {}", key); delete_file_if_exists(&self.config.storage_dir.join(&key)).await?; } } Ok(()) } /// Attempts to add a file to the store. Returns an I/O error if /// something's broken, or a u64 of the maximum allowed file size /// if the file was too big, or a unit if everything worked. pub async fn add_file(&self, key: String, file: StoredFile) -> Result<(), FileAddError> { let mut store = self.state.write().await; if store.full(self.config.max_storage_size) { return Err(FileAddError::Full); } let available_size = store.available_size(self.config.max_storage_size); if file.size > available_size { return Err(FileAddError::TooBig(available_size)); } store.0.insert(key, file); Ok(()) } pub async fn remove_file(&self, key: &str) -> std::io::Result<()> { debug!("removing entry {} from state", key); let mut store = self.state.write().await; store.0.remove(key); if is_valid_storage_code(key) { delete_file_if_exists(&self.config.storage_dir.join(key)).await?; } Ok(()) } pub async fn remove_expired_files(&self) -> std::io::Result<()> { info!("Checking for expired files"); let now = OffsetDateTime::now_utc(); let mut store = self.state.write().await; let old = std::mem::take(store.deref_mut()); for (key, file) in old.0.into_iter() { if file.expiry > now { store.0.insert(key, file); } else { info!("Deleting expired file {}", key); delete_file_if_exists(&self.config.storage_dir.join(&key)).await?; } } Ok(()) } } impl StoredFiles { fn total_size(&self) -> u64 { self.0.iter().fold(0, |acc, (_, f)| acc + f.size) } pub fn available_size(&self, max_storage_size: u64) -> u64 { max_storage_size.saturating_sub(self.total_size()) } pub fn full(&self, max_storage_size: u64) -> bool { self.available_size(max_storage_size) == 0 || self.0.len() >= MAX_STORAGE_FILES } }