2022-04-30 01:38:26 -04:00
|
|
|
use std::{collections::HashMap, io::ErrorKind, path::PathBuf, str::FromStr};
|
2022-04-28 05:13:14 -04:00
|
|
|
|
2022-04-29 22:36:44 -04:00
|
|
|
use log::{debug, error, info, warn};
|
2022-04-30 01:38:26 -04:00
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
use time::OffsetDateTime;
|
2022-04-28 05:18:35 -04:00
|
|
|
use tokio::{
|
|
|
|
fs::File,
|
|
|
|
io::{AsyncReadExt, AsyncWriteExt},
|
|
|
|
};
|
2022-04-28 05:13:14 -04:00
|
|
|
|
|
|
|
const STATE_FILE_NAME: &str = "files.json";
|
2022-04-30 01:38:26 -04:00
|
|
|
const DEFAULT_STORAGE_DIR: &str = "storage";
|
|
|
|
const DEFAULT_MAX_LIFETIME: u32 = 30;
|
|
|
|
const GIGA: u64 = 1024*1024*1024;
|
|
|
|
const DEFAULT_MAX_SINGLE_SIZE: u64 = 16*GIGA;
|
|
|
|
const DEFAULT_MAX_TOTAL_SIZE: u64 = 64*GIGA;
|
|
|
|
|
|
|
|
|
|
|
|
pub(crate) fn storage_dir() -> PathBuf {
|
|
|
|
PathBuf::from(std::env::var("STORAGE_DIR").unwrap_or_else(|_| String::from(DEFAULT_STORAGE_DIR)))
|
|
|
|
}
|
|
|
|
|
|
|
|
fn parse_env_var<T: FromStr>(var: &str, default: T) -> T {
|
|
|
|
std::env::var(var).ok().and_then(|val| val.parse::<T>().ok()).unwrap_or(default)
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) fn max_lifetime() -> u32 {
|
|
|
|
parse_env_var("TRANSBEAM_MAX_LIFETIME", DEFAULT_MAX_LIFETIME)
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) fn max_single_size() -> u64 {
|
|
|
|
parse_env_var("TRANSBEAM_MAX_SINGLE_FILE_SIZE", DEFAULT_MAX_SINGLE_SIZE)
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) fn max_total_size() -> u64 {
|
|
|
|
parse_env_var("TRANSBEAM_MAX_TOTAL_FILE_SIZE", DEFAULT_MAX_TOTAL_SIZE)
|
|
|
|
}
|
|
|
|
|
|
|
|
#[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,
|
|
|
|
}
|
2022-04-28 05:13:14 -04:00
|
|
|
|
|
|
|
pub(crate) mod timestamp {
|
|
|
|
use core::fmt;
|
|
|
|
|
2022-04-28 05:18:35 -04:00
|
|
|
use serde::{de::Visitor, Deserializer, Serializer};
|
2022-04-28 05:13:14 -04:00
|
|
|
use time::OffsetDateTime;
|
|
|
|
|
2022-04-28 05:18:35 -04:00
|
|
|
pub(crate) fn serialize<S: Serializer>(
|
|
|
|
time: &OffsetDateTime,
|
|
|
|
ser: S,
|
|
|
|
) -> Result<S::Ok, S::Error> {
|
2022-04-28 05:13:14 -04:00
|
|
|
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 {
|
2022-04-28 05:25:03 -04:00
|
|
|
write!(formatter, "an integer")
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
|
|
|
|
fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E> {
|
|
|
|
Ok(v)
|
|
|
|
}
|
2022-04-28 05:25:03 -04:00
|
|
|
|
|
|
|
fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E> {
|
|
|
|
Ok(v as i64)
|
|
|
|
}
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
|
2022-04-28 05:18:35 -04:00
|
|
|
pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
|
|
|
|
de: D,
|
|
|
|
) -> Result<OffsetDateTime, D::Error> {
|
|
|
|
Ok(
|
|
|
|
OffsetDateTime::from_unix_timestamp(de.deserialize_i64(I64Visitor)?)
|
|
|
|
.unwrap_or_else(|_| OffsetDateTime::now_utc()),
|
|
|
|
)
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-04-30 01:38:26 -04:00
|
|
|
async fn is_valid_entry(key: &str, info: &StoredFile) -> bool {
|
|
|
|
if info.expiry < OffsetDateTime::now_utc() {
|
|
|
|
info!("File {} has expired", key);
|
2022-04-28 06:26:44 -04:00
|
|
|
return false;
|
|
|
|
}
|
2022-04-30 01:38:26 -04:00
|
|
|
|
2022-04-28 06:26:44 -04:00
|
|
|
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
|
|
|
|
}
|
|
|
|
|
2022-04-30 01:38:26 -04:00
|
|
|
pub(crate) struct FileStore(HashMap<String, StoredFile>);
|
|
|
|
impl FileStore {
|
2022-04-28 05:13:14 -04:00
|
|
|
pub(crate) async fn load() -> std::io::Result<Self> {
|
|
|
|
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?;
|
2022-04-30 01:38:26 -04:00
|
|
|
let map: HashMap<String, StoredFile> = serde_json::from_str(&buf)?;
|
2022-04-28 05:13:14 -04:00
|
|
|
info!("Loaded {} file entries from persistent storage", map.len());
|
2022-04-30 01:38:26 -04:00
|
|
|
let mut filtered: HashMap<String, StoredFile> = HashMap::new();
|
2022-04-28 05:13:14 -04:00
|
|
|
for (key, info) in map.into_iter() {
|
2022-04-30 01:38:26 -04:00
|
|
|
// 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 !crate::util::is_ascii_alphanumeric(&key) {
|
|
|
|
error!("Invalid key in persistent storage: {}", key);
|
|
|
|
continue;
|
|
|
|
}
|
2022-04-28 06:26:44 -04:00
|
|
|
if is_valid_entry(&key, &info).await {
|
|
|
|
filtered.insert(key, info);
|
2022-04-28 05:13:14 -04:00
|
|
|
} else {
|
2022-04-30 01:38:26 -04:00
|
|
|
info!("Deleting file {}", key);
|
2022-04-28 06:26:44 -04:00
|
|
|
if let Err(e) = tokio::fs::remove_file(storage_dir().join(&key)).await {
|
2022-04-30 01:38:26 -04:00
|
|
|
warn!("Failed to delete file {}: {}", key, e);
|
2022-04-28 06:26:44 -04:00
|
|
|
}
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
}
|
2022-04-28 06:26:44 -04:00
|
|
|
let mut loaded = Self(filtered);
|
|
|
|
loaded.save().await?;
|
|
|
|
Ok(loaded)
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
Err(e) => {
|
|
|
|
if let ErrorKind::NotFound = e.kind() {
|
|
|
|
Ok(Self(HashMap::new()))
|
|
|
|
} else {
|
|
|
|
Err(e)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-04-30 01:38:26 -04:00
|
|
|
fn total_size(&self) -> u64 {
|
|
|
|
self.0.iter().fold(0, |acc, (_, f)| acc + f.size)
|
|
|
|
}
|
|
|
|
|
2022-04-28 05:13:14 -04:00
|
|
|
async fn save(&mut self) -> std::io::Result<()> {
|
2022-04-28 06:26:44 -04:00
|
|
|
info!("saving updated state: {} entries", self.0.len());
|
2022-04-28 05:18:35 -04:00
|
|
|
File::create(storage_dir().join(STATE_FILE_NAME))
|
|
|
|
.await?
|
|
|
|
.write_all(&serde_json::to_vec_pretty(&self.0)?)
|
|
|
|
.await
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
|
2022-04-30 01:38:26 -04:00
|
|
|
/// 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.
|
2022-04-28 05:18:35 -04:00
|
|
|
pub(crate) async fn add_file(
|
|
|
|
&mut self,
|
|
|
|
key: String,
|
2022-04-30 01:38:26 -04:00
|
|
|
file: StoredFile,
|
|
|
|
) -> std::io::Result<Result<(), u64>> {
|
|
|
|
let remaining_size = max_total_size().saturating_sub(self.total_size());
|
|
|
|
let allowed_size = std::cmp::min(remaining_size, max_single_size());
|
|
|
|
if file.size > allowed_size { return Ok(Err(allowed_size)); }
|
2022-04-28 05:13:14 -04:00
|
|
|
self.0.insert(key, file);
|
2022-04-30 01:38:26 -04:00
|
|
|
self.save().await.map(Ok)
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
|
2022-04-30 01:38:26 -04:00
|
|
|
pub(crate) fn lookup_file(&self, key: &str) -> Option<StoredFile> {
|
2022-04-28 05:18:35 -04:00
|
|
|
self.0.get(key).cloned()
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) async fn remove_file(&mut self, key: &str) -> std::io::Result<()> {
|
2022-04-28 06:26:44 -04:00
|
|
|
debug!("removing entry {} from state", key);
|
2022-04-28 05:13:14 -04:00
|
|
|
self.0.remove(key);
|
|
|
|
self.save().await
|
|
|
|
}
|
2022-04-30 01:38:26 -04:00
|
|
|
|
|
|
|
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::replace(&mut self.0, HashMap::new()).into_iter() {
|
|
|
|
if file.expiry > now {
|
|
|
|
self.0.insert(key, file);
|
|
|
|
} else {
|
|
|
|
info!("Deleting expired file {}", key);
|
|
|
|
if let Err(e) = tokio::fs::remove_file(storage_dir().join(&key)).await {
|
|
|
|
warn!("Failed to delete expired file {}: {}", key, e);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
self.save().await
|
|
|
|
}
|
2022-04-28 05:13:14 -04:00
|
|
|
}
|