transbeam/src/store.rs

215 lines
6.9 KiB
Rust
Raw Normal View History

use std::{
io::ErrorKind,
ops::DerefMut,
path::{Path, PathBuf},
};
2022-04-28 05:13:14 -04:00
2024-03-27 16:53:02 -04:00
use libc::statvfs64;
use log::{debug, error, info};
2022-05-01 05:28:50 -04:00
use rand::{
distributions::{Alphanumeric, DistString},
thread_rng, Rng,
};
use time::OffsetDateTime;
2024-03-27 16:53:02 -04:00
use tokio::fs::{metadata, File};
2022-04-28 05:13:14 -04:00
const MAX_STORAGE_FILES: usize = 1024;
pub fn gen_storage_code(use_mnemonic: bool) -> String {
if use_mnemonic {
2022-05-01 03:28:25 -04:00
mnemonic::to_string(thread_rng().gen::<[u8; 4]>())
} else {
Alphanumeric.sample_string(&mut thread_rng(), 8)
2022-05-01 03:28:25 -04:00
}
}
pub fn is_valid_storage_code(s: &str) -> bool {
2022-05-01 05:28:50 -04:00
s.as_bytes()
.iter()
.all(|c| c.is_ascii_alphanumeric() || c == &b'-')
2022-05-01 03:28:25 -04:00
}
pub use crate::state::v1::{StoredFile, StoredFileWithPassword, StoredFiles};
2022-08-16 04:54:18 -04:00
async fn is_valid_entry(key: &str, info: &StoredFile, storage_dir: &Path) -> bool {
if info.expiry < OffsetDateTime::now_utc() {
info!("File {} has expired", key);
2022-04-28 06:26:44 -04:00
return false;
}
2023-09-09 17:04:45 -04:00
let file = if let Ok(f) = File::open(storage_dir.join(key)).await {
2022-04-28 06:26:44 -04:00
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,
}
2022-08-16 04:54:18 -04:00
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;
2022-04-28 05:13:14 -04:00
}
if is_valid_entry(&key, &info.file, &self.config.storage_dir).await {
2022-08-16 04:54:18 -04:00
store.0.insert(key, info);
} else {
info!("Deleting file {}", key);
delete_file_if_exists(&self.config.storage_dir.join(&key)).await?;
2022-04-28 05:13:14 -04:00
}
}
2022-08-16 04:54:18 -04:00
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.
2023-09-09 17:04:45 -04:00
pub async fn add_file(
&self,
key: String,
entry: StoredFileWithPassword,
) -> Result<(), FileAddError> {
2024-03-27 16:53:02 -04:00
if self.full().await? {
return Err(FileAddError::Full);
2022-05-01 05:28:50 -04:00
}
2024-03-27 16:53:02 -04:00
let available_size = self.available_size().await?;
if entry.file.size > available_size {
return Err(FileAddError::TooBig(available_size));
}
2024-03-27 16:53:02 -04:00
let mut store = self.state.write().await;
store.0.insert(key, entry);
Ok(())
2022-04-28 05:13:14 -04:00
}
2022-08-16 04:54:18 -04:00
pub async fn remove_file(&self, key: &str) -> std::io::Result<()> {
2022-04-28 06:26:44 -04:00
debug!("removing entry {} from state", key);
2022-08-16 04:54:18 -04:00
let mut store = self.state.write().await;
store.0.remove(key);
if is_valid_storage_code(key) {
2022-08-16 04:54:18 -04:00
delete_file_if_exists(&self.config.storage_dir.join(key)).await?;
}
Ok(())
2022-04-28 05:13:14 -04:00
}
2022-08-16 04:54:18 -04:00
pub async fn remove_expired_files(&self) -> std::io::Result<()> {
info!("Checking for expired files");
let now = OffsetDateTime::now_utc();
2022-08-16 04:54:18 -04:00
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);
2022-08-16 04:54:18 -04:00
delete_file_if_exists(&self.config.storage_dir.join(&key)).await?;
}
}
2022-08-16 04:54:18 -04:00
Ok(())
}
2024-03-27 16:53:02 -04:00
fn disk_free(&self) -> std::io::Result<u64> {
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<u64> {
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)
2022-08-16 04:54:18 -04:00
}
2024-03-27 16:53:02 -04:00
pub async fn available_size(&self) -> std::io::Result<u64> {
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))
2022-08-16 04:54:18 -04:00
}
2024-03-27 16:53:02 -04:00
pub async fn full(&self) -> std::io::Result<bool> {
Ok(self.available_size().await? == 0
|| self.state.read().await.0.len() >= MAX_STORAGE_FILES)
}
2022-04-28 05:13:14 -04:00
}