use std::{ collections::HashMap, io::ErrorKind, path::{Path, PathBuf}, }; use log::{debug, error, info}; use rand::{ distributions::{Alphanumeric, DistString}, thread_rng, Rng, }; use serde::{Deserialize, Serialize}; use time::OffsetDateTime; use tokio::{ fs::File, io::{AsyncReadExt, AsyncWriteExt}, }; const STATE_FILE_NAME: &str = "files.json"; 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'-') } #[derive(Clone, Deserialize, Serialize)] pub struct StoredFile { pub name: String, pub size: u64, #[serde(with = "timestamp")] pub modtime: OffsetDateTime, #[serde(with = "timestamp")] pub expiry: OffsetDateTime, } pub(crate) mod timestamp { use core::fmt; use serde::{de::Visitor, Deserializer, Serializer}; use time::OffsetDateTime; pub(crate) fn serialize( time: &OffsetDateTime, ser: S, ) -> Result { ser.serialize_i64(time.unix_timestamp()) } struct I64Visitor; impl<'de> Visitor<'de> for I64Visitor { type Value = i64; fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { write!(formatter, "an integer") } fn visit_i64(self, v: i64) -> Result { Ok(v) } fn visit_u64(self, v: u64) -> Result { Ok(v as i64) } } pub(crate) fn deserialize<'de, D: Deserializer<'de>>( de: D, ) -> Result { Ok( OffsetDateTime::from_unix_timestamp(de.deserialize_i64(I64Visitor)?) .unwrap_or_else(|_| OffsetDateTime::now_utc()), ) } } 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, } pub struct FileStore { files: HashMap, storage_dir: PathBuf, max_storage_size: u64, } impl FileStore { pub(crate) async fn load(storage_dir: PathBuf, max_storage_size: u64) -> std::io::Result { let open_result = File::open(storage_dir.join(STATE_FILE_NAME)).await; match open_result { Ok(mut f) => { let mut buf = String::new(); f.read_to_string(&mut buf).await?; let map: HashMap = serde_json::from_str(&buf)?; info!("Loaded {} file entries from persistent storage", map.len()); let mut filtered: HashMap = HashMap::new(); for (key, info) in map.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, &storage_dir).await { filtered.insert(key, info); } else { info!("Deleting file {}", key); delete_file_if_exists(&storage_dir.join(&key)).await?; } } let mut loaded = Self { files: filtered, storage_dir, max_storage_size, }; loaded.save().await?; Ok(loaded) } Err(e) => { if let ErrorKind::NotFound = e.kind() { Ok(Self { files: HashMap::new(), storage_dir, max_storage_size, }) } else { Err(e) } } } } fn total_size(&self) -> u64 { self.files.iter().fold(0, |acc, (_, f)| acc + f.size) } pub fn available_size(&self) -> u64 { self.max_storage_size.saturating_sub(self.total_size()) } async fn save(&mut self) -> std::io::Result<()> { info!("saving updated state: {} entries", self.files.len()); File::create(self.storage_dir.join(STATE_FILE_NAME)) .await? .write_all(&serde_json::to_vec_pretty(&self.files)?) .await } pub fn full(&self) -> bool { self.available_size() == 0 || self.files.len() >= MAX_STORAGE_FILES } /// 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(crate) async fn add_file( &mut self, key: String, file: StoredFile, ) -> Result<(), FileAddError> { if self.full() { return Err(FileAddError::Full); } let available_size = self.available_size(); if file.size > available_size { return Err(FileAddError::TooBig(available_size)); } self.files.insert(key, file); self.save().await?; Ok(()) } pub(crate) fn lookup_file(&self, key: &str) -> Option { self.files.get(key).cloned() } pub(crate) async fn remove_file(&mut self, key: &str) -> std::io::Result<()> { debug!("removing entry {} from state", key); self.files.remove(key); self.save().await?; if is_valid_storage_code(key) { delete_file_if_exists(&self.storage_dir.join(key)).await?; } Ok(()) } pub(crate) async fn remove_expired_files(&mut self) -> std::io::Result<()> { info!("Checking for expired files"); let now = OffsetDateTime::now_utc(); for (key, file) in std::mem::take(&mut self.files).into_iter() { if file.expiry > now { self.files.insert(key, file); } else { info!("Deleting expired file {}", key); delete_file_if_exists(&self.storage_dir.join(&key)).await?; } } self.save().await } }