use std::{ io::ErrorKind, ops::DerefMut, path::{Path, PathBuf}, }; use libc::statvfs64; use log::{debug, error, info}; use rand::{ distributions::{Alphanumeric, DistString}, thread_rng, Rng, }; use time::OffsetDateTime; use tokio::fs::{metadata, 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, StoredFileWithPassword, 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.file, &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, entry: StoredFileWithPassword, ) -> Result<(), FileAddError> { if self.full().await? { return Err(FileAddError::Full); } let available_size = self.available_size().await?; if entry.file.size > available_size { return Err(FileAddError::TooBig(available_size)); } let mut store = self.state.write().await; store.0.insert(key, entry); 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, value) in old.0.into_iter() { if value.file.expiry > now { store.0.insert(key, value); } else { info!("Deleting expired file {}", key); delete_file_if_exists(&self.config.storage_dir.join(&key)).await?; } } Ok(()) } fn disk_free(&self) -> std::io::Result { let storage_path_c = std::ffi::CString::new( self.config .storage_dir .to_str() .expect("storage path is not valid unicode"), ) .expect("storage path contains a null byte"); unsafe { let mut stat: statvfs64 = std::mem::zeroed(); loop { if statvfs64(storage_path_c.as_ptr(), std::ptr::addr_of_mut!(stat)) == 0 { break; } let os_error = std::io::Error::last_os_error(); if os_error.kind() != ErrorKind::Interrupted { return Err(os_error); } } Ok(stat.f_bsize.saturating_mul(stat.f_bavail)) } } async fn pending_data_size(&self) -> std::io::Result { let store = self.state.read().await; let mut total_pending_size = 0u64; for (key, value) in store.0.iter() { let file_path = self.config.storage_dir.join(key); let pending_size = value .file .size .saturating_sub(metadata(&file_path).await?.len()); total_pending_size = total_pending_size.saturating_add(pending_size); } Ok(total_pending_size) } async fn total_stored_size(&self) -> u64 { let store = self.state.read().await; store.0.iter().fold(0, |acc, (_, v)| acc + v.file.size) } pub async fn available_size(&self) -> std::io::Result { let available_policy_size = self .config .max_storage_size .saturating_sub(self.total_stored_size().await); let available_disk_size = self .disk_free()? .saturating_sub(self.config.min_disk_free) .saturating_sub(self.pending_data_size().await?); Ok(std::cmp::min(available_policy_size, available_disk_size)) } pub async fn full(&self) -> std::io::Result { Ok(self.available_size().await? == 0 || self.state.read().await.0.len() >= MAX_STORAGE_FILES) } }