109 lines
3.9 KiB
Rust
109 lines
3.9 KiB
Rust
|
use std::{collections::HashMap, io::ErrorKind};
|
||
|
|
||
|
use log::{info, error};
|
||
|
use tokio::{fs::File, io::{AsyncReadExt, AsyncWriteExt}};
|
||
|
|
||
|
use crate::{DownloadableFile, storage_dir};
|
||
|
|
||
|
const STATE_FILE_NAME: &str = "files.json";
|
||
|
|
||
|
pub(crate) mod timestamp {
|
||
|
use core::fmt;
|
||
|
|
||
|
use serde::{Serializer, Deserializer, de::Visitor};
|
||
|
use time::OffsetDateTime;
|
||
|
|
||
|
pub(crate) fn serialize<S: Serializer>(time: &OffsetDateTime, ser: S) -> Result<S::Ok, S::Error> {
|
||
|
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, "a signed integer")
|
||
|
}
|
||
|
|
||
|
fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E> {
|
||
|
Ok(v)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
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()))
|
||
|
}
|
||
|
}
|
||
|
|
||
|
pub(crate) struct PersistentState(HashMap<String, DownloadableFile>);
|
||
|
impl PersistentState {
|
||
|
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?;
|
||
|
let map: HashMap<String, DownloadableFile> = serde_json::from_str(&buf)?;
|
||
|
info!("Loaded {} file entries from persistent storage", map.len());
|
||
|
let mut filtered: HashMap<String, DownloadableFile> = HashMap::new();
|
||
|
for (key, info) in map.into_iter() {
|
||
|
if !crate::util::is_ascii_alphanumeric(&key) {
|
||
|
error!("Invalid key in persistent storage: {}", key);
|
||
|
continue;
|
||
|
}
|
||
|
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);
|
||
|
continue;
|
||
|
};
|
||
|
let metadata = if let Ok(md) = file.metadata().await {
|
||
|
md
|
||
|
} else {
|
||
|
error!("Unable to get metadata for file {} referenced in persistent storage", key);
|
||
|
continue;
|
||
|
};
|
||
|
if metadata.len() != info.size {
|
||
|
error!("Mismatched file size for file {} referenced in persistent storage: expected {}, found {}", key, info.size, metadata.len());
|
||
|
continue
|
||
|
}
|
||
|
filtered.insert(key, info);
|
||
|
}
|
||
|
Ok(Self(filtered))
|
||
|
}
|
||
|
Err(e) => {
|
||
|
if let ErrorKind::NotFound = e.kind() {
|
||
|
Ok(Self(HashMap::new()))
|
||
|
} else {
|
||
|
Err(e)
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
async fn save(&mut self) -> std::io::Result<()> {
|
||
|
File::create(storage_dir().join(STATE_FILE_NAME)).await?.write_all(&serde_json::to_vec_pretty(&self.0)?).await
|
||
|
}
|
||
|
|
||
|
pub(crate) async fn add_file(&mut self, key: String, file: DownloadableFile) -> std::io::Result<()> {
|
||
|
self.0.insert(key, file);
|
||
|
self.save().await
|
||
|
}
|
||
|
|
||
|
pub(crate) fn lookup_file(&self, key: &str) -> Option<DownloadableFile> {
|
||
|
self.0.get(key).map(|f| f.clone())
|
||
|
}
|
||
|
|
||
|
pub(crate) async fn remove_file(&mut self, key: &str) -> std::io::Result<()> {
|
||
|
self.0.remove(key);
|
||
|
self.save().await
|
||
|
}
|
||
|
|
||
|
pub(crate) fn remove_uploader(&mut self, key: &str) {
|
||
|
if let Some(f) = self.0.get_mut(key) {
|
||
|
f.uploader.take();
|
||
|
}
|
||
|
}
|
||
|
}
|