2
0
Fork 0
mirror of https://git.asonix.dog/asonix/pict-rs synced 2024-11-14 00:13:59 +00:00
pict-rs/src/upload_manager.rs

953 lines
29 KiB
Rust
Raw Normal View History

2020-07-11 20:40:40 +00:00
use crate::{
config::Format,
error::UploadError,
migrate::{alias_id_key, alias_key, alias_key_bounds, variant_key_bounds, LatestDb},
2020-07-11 20:40:40 +00:00
to_ext,
validate::validate_image,
};
2020-06-06 21:41:17 +00:00
use actix_web::web;
2021-01-19 01:54:39 +00:00
use futures::stream::{Stream, StreamExt, TryStreamExt};
2020-06-06 21:41:17 +00:00
use sha2::Digest;
use std::{path::PathBuf, pin::Pin, sync::Arc};
2020-06-14 15:07:31 +00:00
use tracing::{debug, error, info, instrument, warn, Span};
2020-06-06 21:41:17 +00:00
// TREE STRUCTURE
// - Alias Tree
// - alias -> hash
// - alias / id -> u64(id)
// - alias / delete -> delete token
// - Main Tree
// - hash -> filename
// - hash 0 u64(id) -> alias
// - hash 2 variant path -> variant path
// - Filename Tree
// - filename -> hash
2020-06-06 21:41:17 +00:00
#[derive(Clone)]
pub struct UploadManager {
inner: Arc<UploadManagerInner>,
}
struct UploadManagerInner {
2020-06-07 01:44:26 +00:00
format: Option<Format>,
2020-06-06 21:41:17 +00:00
hasher: sha2::Sha256,
image_dir: PathBuf,
2020-06-06 22:43:33 +00:00
alias_tree: sled::Tree,
filename_tree: sled::Tree,
2020-09-14 21:42:31 +00:00
main_tree: sled::Tree,
2020-06-07 00:29:15 +00:00
db: sled::Db,
2020-06-06 21:41:17 +00:00
}
impl std::fmt::Debug for UploadManager {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
2020-06-14 15:07:31 +00:00
f.debug_struct("UploadManager").finish()
}
}
2021-02-10 22:57:42 +00:00
type UploadStream<E> = Pin<Box<dyn Stream<Item = Result<web::Bytes, E>>>>;
#[derive(Clone)]
pub(crate) struct Serde<T> {
inner: T,
}
impl<T> Serde<T> {
pub(crate) fn new(inner: T) -> Self {
Serde { inner }
}
}
impl<T> serde::Serialize for Serde<T>
where
T: std::fmt::Display,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let s = self.inner.to_string();
serde::Serialize::serialize(s.as_str(), serializer)
}
}
impl<'de, T> serde::Deserialize<'de> for Serde<T>
where
T: std::str::FromStr,
<T as std::str::FromStr>::Err: std::fmt::Display,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let s: String = serde::Deserialize::deserialize(deserializer)?;
let inner = s
.parse::<T>()
.map_err(|e| serde::de::Error::custom(e.to_string()))?;
Ok(Serde { inner })
}
}
#[derive(Clone, serde::Deserialize, serde::Serialize)]
pub(crate) struct Details {
width: usize,
height: usize,
content_type: Serde<mime::Mime>,
created_at: time::OffsetDateTime,
}
fn mime_from_media_type(media_type: rexiv2::MediaType) -> mime::Mime {
match media_type {
rexiv2::MediaType::Jpeg => mime::IMAGE_JPEG,
rexiv2::MediaType::Png => mime::IMAGE_PNG,
rexiv2::MediaType::Gif => mime::IMAGE_GIF,
rexiv2::MediaType::Other(s) if s == "image/webp" => s.parse::<mime::Mime>().unwrap(),
rexiv2::MediaType::Other(s) if s == "video/mp4" || s == "video/quicktime" => {
"video/mp4".parse::<mime::Mime>().unwrap()
}
_ => mime::APPLICATION_OCTET_STREAM,
}
}
impl Details {
pub(crate) fn from_bytes(bytes: &[u8]) -> Result<Self, UploadError> {
let metadata = rexiv2::Metadata::new_from_buffer(bytes)?;
let mime_type = mime_from_media_type(metadata.get_media_type()?);
let width = metadata.get_pixel_width();
let height = metadata.get_pixel_height();
let details = Details::now(width as usize, height as usize, mime_type);
Ok(details)
}
pub(crate) async fn from_path(path: PathBuf) -> Result<Self, UploadError> {
let (mime_type, width, height) = web::block(move || {
rexiv2::Metadata::new_from_path(&path).and_then(|metadata| {
metadata
.get_media_type()
.map(mime_from_media_type)
.map(|mime_type| {
(
mime_type,
metadata.get_pixel_width(),
metadata.get_pixel_height(),
)
})
})
})
2021-02-10 22:57:42 +00:00
.await??;
Ok(Details::now(width as usize, height as usize, mime_type))
}
fn now(width: usize, height: usize, content_type: mime::Mime) -> Self {
Details {
width,
height,
content_type: Serde::new(content_type),
created_at: time::OffsetDateTime::now_utc(),
}
}
pub(crate) fn content_type(&self) -> mime::Mime {
self.content_type.inner.clone()
}
pub(crate) fn system_time(&self) -> std::time::SystemTime {
self.created_at.into()
}
}
2020-06-14 15:07:31 +00:00
struct FilenameIVec {
inner: sled::IVec,
}
2020-06-14 15:07:31 +00:00
impl FilenameIVec {
fn new(inner: sled::IVec) -> Self {
FilenameIVec { inner }
}
}
2020-06-14 15:07:31 +00:00
impl std::fmt::Debug for FilenameIVec {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "{:?}", String::from_utf8(self.inner.to_vec()))
}
}
2020-06-14 15:07:31 +00:00
struct Hash {
inner: Vec<u8>,
}
impl Hash {
fn new(inner: Vec<u8>) -> Self {
Hash { inner }
}
}
impl std::fmt::Debug for Hash {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
2020-06-14 15:07:31 +00:00
write!(f, "{}", base64::encode(&self.inner))
}
}
2020-06-06 21:41:17 +00:00
enum Dup {
Exists,
New,
}
impl Dup {
fn exists(&self) -> bool {
matches!(self, Dup::Exists)
2020-06-06 21:41:17 +00:00
}
}
impl UploadManager {
/// Get the image directory
pub(crate) fn image_dir(&self) -> PathBuf {
self.inner.image_dir.clone()
}
/// Create a new UploadManager
2020-06-07 01:44:26 +00:00
pub(crate) async fn new(
mut root_dir: PathBuf,
format: Option<Format>,
) -> Result<Self, UploadError> {
2020-07-11 20:40:40 +00:00
let root_clone = root_dir.clone();
2020-06-06 21:41:17 +00:00
// sled automatically creates it's own directories
2021-02-10 22:57:42 +00:00
let db = web::block(move || LatestDb::exists(root_clone).migrate()).await??;
2020-06-06 21:41:17 +00:00
root_dir.push("files");
// Ensure file dir exists
2021-01-19 01:54:39 +00:00
actix_fs::create_dir_all(root_dir.clone()).await?;
2020-06-06 21:41:17 +00:00
Ok(UploadManager {
inner: Arc::new(UploadManagerInner {
2020-06-07 01:44:26 +00:00
format,
2020-06-06 21:41:17 +00:00
hasher: sha2::Sha256::new(),
image_dir: root_dir,
2020-06-06 22:43:33 +00:00
alias_tree: db.open_tree("alias")?,
filename_tree: db.open_tree("filename")?,
2020-09-14 21:42:31 +00:00
main_tree: db.open_tree("main")?,
2020-06-06 21:41:17 +00:00
db,
}),
})
}
/// Store the path to a generated image variant so we can easily clean it up later
2020-06-14 15:07:31 +00:00
#[instrument(skip(self))]
2020-06-24 16:58:46 +00:00
pub(crate) async fn store_variant(
&self,
path: PathBuf,
filename: String,
) -> Result<(), UploadError> {
let path_string = path.to_str().ok_or(UploadError::Path)?.to_string();
let fname_tree = self.inner.filename_tree.clone();
2020-06-14 15:07:31 +00:00
debug!("Getting hash");
let hash: sled::IVec = web::block(move || fname_tree.get(filename.as_bytes()))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingFilename)?;
let key = variant_key(&hash, &path_string);
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2020-06-14 15:07:31 +00:00
debug!("Storing variant");
2021-02-10 22:57:42 +00:00
web::block(move || main_tree.insert(key, path_string.as_bytes())).await??;
2020-06-14 15:07:31 +00:00
debug!("Stored variant");
Ok(())
}
/// Get the image details for a given variant
pub(crate) async fn variant_details(
&self,
path: PathBuf,
filename: String,
) -> Result<Option<Details>, UploadError> {
let path_string = path.to_str().ok_or(UploadError::Path)?.to_string();
let fname_tree = self.inner.filename_tree.clone();
debug!("Getting hash");
let hash: sled::IVec = web::block(move || fname_tree.get(filename.as_bytes()))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingFilename)?;
let key = variant_details_key(&hash, &path_string);
let main_tree = self.inner.main_tree.clone();
debug!("Getting details");
2021-02-10 22:57:42 +00:00
let opt = match web::block(move || main_tree.get(key)).await?? {
Some(ivec) => match serde_json::from_slice(&ivec) {
Ok(details) => Some(details),
Err(_) => None,
},
None => None,
};
debug!("Got details");
Ok(opt)
}
pub(crate) async fn store_variant_details(
&self,
path: PathBuf,
filename: String,
details: &Details,
) -> Result<(), UploadError> {
let path_string = path.to_str().ok_or(UploadError::Path)?.to_string();
let fname_tree = self.inner.filename_tree.clone();
debug!("Getting hash");
let hash: sled::IVec = web::block(move || fname_tree.get(filename.as_bytes()))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingFilename)?;
let key = variant_details_key(&hash, &path_string);
let main_tree = self.inner.main_tree.clone();
let details_value = serde_json::to_string(details)?;
debug!("Storing details");
2021-02-10 22:57:42 +00:00
web::block(move || main_tree.insert(key, details_value.as_bytes())).await??;
debug!("Stored details");
Ok(())
}
/// Get a list of aliases for a given file
pub(crate) async fn aliases_by_filename(
&self,
filename: String,
) -> Result<Vec<String>, UploadError> {
let fname_tree = self.inner.filename_tree.clone();
let hash = web::block(move || fname_tree.get(filename.as_bytes()))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingAlias)?;
self.aliases_by_hash(&hash).await
}
/// Get a list of aliases for a given alias
pub(crate) async fn aliases_by_alias(&self, alias: String) -> Result<Vec<String>, UploadError> {
let alias_tree = self.inner.alias_tree.clone();
let hash = web::block(move || alias_tree.get(alias.as_bytes()))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingFilename)?;
self.aliases_by_hash(&hash).await
}
async fn aliases_by_hash(&self, hash: &sled::IVec) -> Result<Vec<String>, UploadError> {
let (start, end) = alias_key_bounds(hash);
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
let aliases = web::block(move || {
main_tree
.range(start..end)
.values()
.collect::<Result<Vec<_>, _>>()
})
2021-02-10 22:57:42 +00:00
.await??;
debug!("Got {} aliases for hash", aliases.len());
let aliases = aliases
.into_iter()
.filter_map(|s| String::from_utf8(s.to_vec()).ok())
.collect::<Vec<_>>();
for alias in aliases.iter() {
debug!("{}", alias);
}
Ok(aliases)
}
/// Delete an alias without a delete token
pub(crate) async fn delete_without_token(&self, alias: String) -> Result<(), UploadError> {
let token_key = delete_key(&alias);
let alias_tree = self.inner.alias_tree.clone();
let token = web::block(move || alias_tree.get(token_key.as_bytes()))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingAlias)?;
self.delete(alias, String::from_utf8(token.to_vec())?).await
}
/// Delete the alias, and the file & variants if no more aliases exist
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, alias, token))]
2020-06-07 00:29:15 +00:00
pub(crate) async fn delete(&self, alias: String, token: String) -> Result<(), UploadError> {
use sled::Transactional;
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2020-06-07 00:29:15 +00:00
let alias_tree = self.inner.alias_tree.clone();
2020-06-14 15:07:31 +00:00
let span = Span::current();
2020-06-07 00:29:15 +00:00
let alias2 = alias.clone();
let hash = web::block(move || {
2020-09-14 21:42:31 +00:00
[&main_tree, &alias_tree].transaction(|v| {
2020-06-14 15:07:31 +00:00
let entered = span.enter();
2020-09-14 21:42:31 +00:00
let main_tree = &v[0];
2020-06-07 00:29:15 +00:00
let alias_tree = &v[1];
// -- GET TOKEN --
2020-06-14 15:07:31 +00:00
debug!("Deleting alias -> delete-token mapping");
2020-06-07 00:29:15 +00:00
let existing_token = alias_tree
.remove(delete_key(&alias2).as_bytes())?
.ok_or_else(|| trans_err(UploadError::MissingAlias))?;
2020-06-07 00:29:15 +00:00
// Bail if invalid token
if existing_token != token {
warn!("Invalid delete token");
return Err(trans_err(UploadError::InvalidToken));
}
// -- GET ID FOR HASH TREE CLEANUP --
2020-06-14 15:07:31 +00:00
debug!("Deleting alias -> id mapping");
2020-06-07 00:29:15 +00:00
let id = alias_tree
.remove(alias_id_key(&alias2).as_bytes())?
.ok_or_else(|| trans_err(UploadError::MissingAlias))?;
2020-06-07 00:29:15 +00:00
let id = String::from_utf8(id.to_vec()).map_err(|e| trans_err(e.into()))?;
// -- GET HASH FOR HASH TREE CLEANUP --
2020-06-14 15:07:31 +00:00
debug!("Deleting alias -> hash mapping");
2020-06-07 00:29:15 +00:00
let hash = alias_tree
.remove(alias2.as_bytes())?
.ok_or_else(|| trans_err(UploadError::MissingAlias))?;
2020-06-07 00:29:15 +00:00
// -- REMOVE HASH TREE ELEMENT --
2020-06-14 15:07:31 +00:00
debug!("Deleting hash -> alias mapping");
2020-09-14 21:42:31 +00:00
main_tree.remove(alias_key(&hash, &id))?;
2020-06-14 15:07:31 +00:00
drop(entered);
2020-06-07 00:29:15 +00:00
Ok(hash)
})
})
2021-02-10 22:57:42 +00:00
.await??;
2020-06-07 00:29:15 +00:00
2020-06-07 03:21:42 +00:00
// -- CHECK IF ANY OTHER ALIASES EXIST --
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2020-06-07 00:29:15 +00:00
let (start, end) = alias_key_bounds(&hash);
2020-06-14 15:07:31 +00:00
debug!("Checking for additional aliases referencing hash");
2020-06-07 00:29:15 +00:00
let any_aliases = web::block(move || {
2020-09-14 21:42:31 +00:00
Ok(main_tree.range(start..end).next().is_some()) as Result<bool, UploadError>
2020-06-07 00:29:15 +00:00
})
2021-02-10 22:57:42 +00:00
.await??;
2020-06-07 00:29:15 +00:00
// Bail if there are existing aliases
if any_aliases {
2020-06-14 15:07:31 +00:00
debug!("Other aliases reference file, not removing from disk");
2020-06-07 00:29:15 +00:00
return Ok(());
}
// -- DELETE HASH ENTRY --
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
let hash2 = hash.clone();
2020-06-14 15:07:31 +00:00
debug!("Deleting hash -> filename mapping");
2020-09-14 21:42:31 +00:00
let filename = web::block(move || main_tree.remove(&hash2))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingFile)?;
2020-06-07 00:29:15 +00:00
2020-06-07 00:33:29 +00:00
// -- DELETE FILES --
let this = self.clone();
2020-06-14 15:07:31 +00:00
debug!("Spawning cleanup task");
let span = Span::current();
2020-06-07 00:33:29 +00:00
actix_rt::spawn(async move {
2020-06-14 15:07:31 +00:00
let entered = span.enter();
if let Err(e) = this
.cleanup_files(FilenameIVec::new(filename.clone()))
.await
{
2020-06-07 00:33:29 +00:00
error!("Error removing files from fs, {}", e);
}
2020-06-14 15:07:31 +00:00
info!(
"Files deleted for {:?}",
String::from_utf8(filename.to_vec())
);
drop(entered);
2020-06-07 00:33:29 +00:00
});
2020-06-07 00:29:15 +00:00
Ok(())
}
/// Generate a delete token for an alias
2020-06-14 15:07:31 +00:00
#[instrument(skip(self))]
2020-06-07 00:29:15 +00:00
pub(crate) async fn delete_token(&self, alias: String) -> Result<String, UploadError> {
debug!("Generating delete token");
2020-06-07 00:29:15 +00:00
use rand::distributions::{Alphanumeric, Distribution};
let rng = rand::thread_rng();
2021-02-10 22:57:42 +00:00
let s: String = Alphanumeric
.sample_iter(rng)
.take(10)
.map(char::from)
.collect();
2020-06-07 00:29:15 +00:00
let delete_token = s.clone();
debug!("Saving delete token");
2020-06-07 00:29:15 +00:00
let alias_tree = self.inner.alias_tree.clone();
let key = delete_key(&alias);
let res = web::block(move || {
alias_tree.compare_and_swap(
key.as_bytes(),
None as Option<sled::IVec>,
Some(s.as_bytes()),
)
})
2021-02-10 22:57:42 +00:00
.await??;
2020-06-07 00:29:15 +00:00
if let Err(sled::CompareAndSwapError {
current: Some(ivec),
..
}) = res
{
let s = String::from_utf8(ivec.to_vec())?;
2020-06-14 15:07:31 +00:00
debug!("Returning existing delete token, {}", s);
2020-06-07 00:29:15 +00:00
return Ok(s);
}
2020-06-14 15:07:31 +00:00
debug!("Returning new delete token, {}", delete_token);
2020-06-07 00:29:15 +00:00
Ok(delete_token)
}
/// Upload the file while preserving the filename, optionally validating the uploaded image
2020-06-14 15:07:31 +00:00
#[instrument(skip(self, stream))]
pub(crate) async fn import<E>(
&self,
alias: String,
content_type: mime::Mime,
validate: bool,
stream: UploadStream<E>,
) -> Result<String, UploadError>
where
UploadError: From<E>,
2020-06-15 02:41:45 +00:00
E: Unpin,
{
2020-06-15 02:41:45 +00:00
// -- READ IN BYTES FROM CLIENT --
2020-06-14 15:07:31 +00:00
debug!("Reading stream");
2020-06-24 16:58:46 +00:00
let tmpfile = crate::tmp_file();
2020-06-15 02:41:45 +00:00
safe_save_stream(tmpfile.clone(), stream).await?;
2020-06-06 21:41:17 +00:00
2020-06-15 02:41:45 +00:00
let content_type = if validate {
2020-06-14 15:07:31 +00:00
debug!("Validating image");
2020-06-11 16:46:00 +00:00
let format = self.inner.format.clone();
2020-06-15 02:41:45 +00:00
validate_image(tmpfile.clone(), format).await?
} else {
2020-06-15 02:41:45 +00:00
content_type
};
2020-06-06 21:41:17 +00:00
// -- DUPLICATE CHECKS --
2020-06-07 01:44:26 +00:00
// Cloning bytes is fine because it's actually a pointer
2020-06-14 15:07:31 +00:00
debug!("Hashing bytes");
2020-06-15 02:41:45 +00:00
let hash = self.hash(tmpfile.clone()).await?;
2020-06-07 01:44:26 +00:00
2020-06-14 15:07:31 +00:00
debug!("Storing alias");
self.add_existing_alias(&hash, &alias).await?;
2020-06-07 01:44:26 +00:00
2020-06-14 15:07:31 +00:00
debug!("Saving file");
2020-06-15 02:41:45 +00:00
self.save_upload(tmpfile, hash, content_type).await?;
// Return alias to file
Ok(alias)
}
2020-06-07 01:44:26 +00:00
/// Upload the file, discarding bytes if it's already present, or saving if it's new
2020-06-14 15:07:31 +00:00
#[instrument(skip(self, stream))]
pub(crate) async fn upload<E>(&self, stream: UploadStream<E>) -> Result<String, UploadError>
where
UploadError: From<E>,
2020-06-15 02:41:45 +00:00
E: Unpin,
{
// -- READ IN BYTES FROM CLIENT --
debug!("Reading stream");
2020-06-24 16:58:46 +00:00
let tmpfile = crate::tmp_file();
2020-06-15 02:41:45 +00:00
safe_save_stream(tmpfile.clone(), stream).await?;
// -- VALIDATE IMAGE --
debug!("Validating image");
2020-06-11 16:46:00 +00:00
let format = self.inner.format.clone();
2020-06-15 02:41:45 +00:00
let content_type = validate_image(tmpfile.clone(), format).await?;
2020-06-06 21:41:17 +00:00
// -- DUPLICATE CHECKS --
// Cloning bytes is fine because it's actually a pointer
debug!("Hashing bytes");
2020-06-15 02:41:45 +00:00
let hash = self.hash(tmpfile.clone()).await?;
2020-06-06 21:41:17 +00:00
debug!("Adding alias");
2020-06-06 22:43:33 +00:00
let alias = self.add_alias(&hash, content_type.clone()).await?;
2020-06-06 21:41:17 +00:00
debug!("Saving file");
2020-06-15 02:41:45 +00:00
self.save_upload(tmpfile, hash, content_type).await?;
2020-06-06 22:43:33 +00:00
// Return alias to file
Ok(alias)
2020-06-06 21:41:17 +00:00
}
/// Fetch the real on-disk filename given an alias
2020-06-14 15:07:31 +00:00
#[instrument(skip(self))]
2020-06-06 22:43:33 +00:00
pub(crate) async fn from_alias(&self, alias: String) -> Result<String, UploadError> {
let tree = self.inner.alias_tree.clone();
2020-06-14 15:07:31 +00:00
debug!("Getting hash from alias");
2020-06-06 22:43:33 +00:00
let hash = web::block(move || tree.get(alias.as_bytes()))
2021-02-10 22:57:42 +00:00
.await??
2020-06-06 22:43:33 +00:00
.ok_or(UploadError::MissingAlias)?;
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2020-06-14 15:07:31 +00:00
debug!("Getting filename from hash");
2020-09-14 21:42:31 +00:00
let filename = web::block(move || main_tree.get(hash))
2021-02-10 22:57:42 +00:00
.await??
2020-06-06 22:43:33 +00:00
.ok_or(UploadError::MissingFile)?;
let filename = String::from_utf8(filename.to_vec())?;
Ok(filename)
}
// Find image variants and remove them from the DB and the disk
2020-06-14 15:07:31 +00:00
#[instrument(skip(self))]
async fn cleanup_files(&self, filename: FilenameIVec) -> Result<(), UploadError> {
let filename = filename.inner;
let mut path = self.image_dir();
let fname = String::from_utf8(filename.to_vec())?;
path.push(fname);
let mut errors = Vec::new();
2020-06-14 15:07:31 +00:00
debug!("Deleting {:?}", path);
2021-01-19 01:54:39 +00:00
if let Err(e) = actix_fs::remove_file(path).await {
errors.push(e.into());
}
let fname_tree = self.inner.filename_tree.clone();
2020-06-14 15:07:31 +00:00
debug!("Deleting filename -> hash mapping");
let hash = web::block(move || fname_tree.remove(filename))
2021-02-10 22:57:42 +00:00
.await??
.ok_or(UploadError::MissingFile)?;
let (start, end) = variant_key_bounds(&hash);
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2020-06-14 15:07:31 +00:00
debug!("Fetching file variants");
let keys = web::block(move || {
let mut keys = Vec::new();
2020-09-14 21:42:31 +00:00
for key in main_tree.range(start..end).keys() {
keys.push(key?.to_owned());
}
Ok(keys) as Result<Vec<sled::IVec>, UploadError>
})
2021-02-10 22:57:42 +00:00
.await??;
2020-06-14 15:07:31 +00:00
debug!("{} files prepared for deletion", keys.len());
for key in keys {
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2021-02-10 22:57:42 +00:00
if let Some(path) = web::block(move || main_tree.remove(key)).await?? {
let s = String::from_utf8_lossy(&path);
debug!("Deleting {}", s);
// ignore json objects
if !s.starts_with('{') {
if let Err(e) = remove_path(path).await {
errors.push(e);
}
}
}
}
for error in errors {
error!("Error deleting files, {}", error);
}
Ok(())
}
// check duplicates & store image if new
async fn save_upload(
&self,
2020-06-15 02:41:45 +00:00
tmpfile: PathBuf,
2020-06-14 15:07:31 +00:00
hash: Hash,
content_type: mime::Mime,
) -> Result<(), UploadError> {
let (dup, name) = self.check_duplicate(hash, content_type).await?;
// bail early with alias to existing file if this is a duplicate
if dup.exists() {
2020-06-14 15:07:31 +00:00
debug!("Duplicate exists, not saving file");
return Ok(());
}
// -- WRITE NEW FILE --
let mut real_path = self.image_dir();
real_path.push(name);
2020-06-24 16:58:46 +00:00
crate::safe_move_file(tmpfile, real_path).await?;
Ok(())
}
2020-06-06 21:41:17 +00:00
// produce a sh256sum of the uploaded file
2020-06-15 02:41:45 +00:00
async fn hash(&self, tmpfile: PathBuf) -> Result<Hash, UploadError> {
2020-06-06 21:41:17 +00:00
let mut hasher = self.inner.hasher.clone();
2020-06-15 02:41:45 +00:00
2021-01-19 01:54:39 +00:00
let file = actix_fs::file::open(tmpfile).await?;
let mut stream = Box::pin(actix_fs::file::read_to_stream(file).await?);
2020-06-15 02:41:45 +00:00
while let Some(res) = stream.next().await {
let bytes = res?;
hasher = web::block(move || {
hasher.update(&bytes);
Ok(hasher) as Result<_, UploadError>
})
2021-02-10 22:57:42 +00:00
.await??;
2020-06-15 02:41:45 +00:00
}
let hash =
web::block(move || Ok(hasher.finalize_reset().to_vec()) as Result<_, UploadError>)
2021-02-10 22:57:42 +00:00
.await??;
2020-06-06 21:41:17 +00:00
2020-06-14 15:07:31 +00:00
Ok(Hash::new(hash))
2020-06-06 21:41:17 +00:00
}
// check for an already-uploaded image with this hash, returning the path to the target file
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, hash, content_type))]
2020-06-06 21:41:17 +00:00
async fn check_duplicate(
&self,
2020-06-14 15:07:31 +00:00
hash: Hash,
2020-06-06 21:41:17 +00:00
content_type: mime::Mime,
2020-06-06 22:43:33 +00:00
) -> Result<(Dup, String), UploadError> {
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2020-06-06 21:41:17 +00:00
let filename = self.next_file(content_type).await?;
let filename2 = filename.clone();
2020-06-14 15:07:31 +00:00
let hash2 = hash.inner.clone();
debug!("Inserting filename for hash");
2020-06-06 21:41:17 +00:00
let res = web::block(move || {
2020-09-14 21:42:31 +00:00
main_tree.compare_and_swap(
hash2,
None as Option<sled::IVec>,
Some(filename2.as_bytes()),
)
2020-06-06 21:41:17 +00:00
})
2021-02-10 22:57:42 +00:00
.await??;
2020-06-06 21:41:17 +00:00
if let Err(sled::CompareAndSwapError {
current: Some(ivec),
..
}) = res
{
let name = String::from_utf8(ivec.to_vec())?;
2020-06-14 15:07:31 +00:00
debug!("Filename exists for hash, {}", name);
2020-06-06 22:43:33 +00:00
return Ok((Dup::Exists, name));
2020-06-06 21:41:17 +00:00
}
let fname_tree = self.inner.filename_tree.clone();
let filename2 = filename.clone();
2020-06-14 15:07:31 +00:00
debug!("Saving filename -> hash relation");
2021-02-10 22:57:42 +00:00
web::block(move || fname_tree.insert(filename2, hash.inner)).await??;
2020-06-06 22:43:33 +00:00
Ok((Dup::New, filename))
2020-06-06 21:41:17 +00:00
}
// generate a short filename that isn't already in-use
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, content_type))]
2020-06-06 21:41:17 +00:00
async fn next_file(&self, content_type: mime::Mime) -> Result<String, UploadError> {
2020-06-06 22:43:33 +00:00
let image_dir = self.image_dir();
2020-06-06 21:41:17 +00:00
use rand::distributions::{Alphanumeric, Distribution};
let mut limit: usize = 10;
2021-02-10 22:57:42 +00:00
let mut rng = rand::thread_rng();
2020-06-06 21:41:17 +00:00
loop {
2020-06-14 15:07:31 +00:00
debug!("Filename generation loop");
2020-06-06 21:41:17 +00:00
let mut path = image_dir.clone();
2021-02-10 22:57:42 +00:00
let s: String = Alphanumeric
.sample_iter(&mut rng)
.take(limit)
.map(char::from)
.collect();
2020-06-06 21:41:17 +00:00
2020-06-24 16:58:46 +00:00
let filename = file_name(s, content_type.clone())?;
2020-06-06 21:41:17 +00:00
path.push(filename.clone());
2021-01-19 01:54:39 +00:00
if let Err(e) = actix_fs::metadata(path).await {
if e.kind() == Some(std::io::ErrorKind::NotFound) {
2020-06-14 15:07:31 +00:00
debug!("Generated unused filename {}", filename);
2020-06-06 21:41:17 +00:00
return Ok(filename);
}
return Err(e.into());
}
2020-06-14 15:07:31 +00:00
debug!("Filename exists, trying again");
2020-06-06 21:41:17 +00:00
limit += 1;
}
}
2020-06-06 22:43:33 +00:00
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, hash, alias))]
2020-06-14 15:07:31 +00:00
async fn add_existing_alias(&self, hash: &Hash, alias: &str) -> Result<(), UploadError> {
self.save_alias(hash, alias).await??;
self.store_alias(hash, alias).await?;
Ok(())
}
2020-06-06 22:43:33 +00:00
// Add an alias to an existing file
//
// This will help if multiple 'users' upload the same file, and one of them wants to delete it
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, hash, content_type))]
2020-06-06 22:43:33 +00:00
async fn add_alias(
&self,
2020-06-14 15:07:31 +00:00
hash: &Hash,
2020-06-06 22:43:33 +00:00
content_type: mime::Mime,
) -> Result<String, UploadError> {
let alias = self.next_alias(hash, content_type).await?;
self.store_alias(hash, &alias).await?;
Ok(alias)
}
// Add a pre-defined alias to an existin file
//
// DANGER: this can cause BAD BAD BAD conflicts if the same alias is used for multiple files
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, hash))]
2020-06-14 15:07:31 +00:00
async fn store_alias(&self, hash: &Hash, alias: &str) -> Result<(), UploadError> {
let alias = alias.to_string();
2020-06-06 22:43:33 +00:00
loop {
2020-06-14 15:07:31 +00:00
debug!("hash -> alias save loop");
2020-06-06 22:43:33 +00:00
let db = self.inner.db.clone();
2021-02-10 22:57:42 +00:00
let id = web::block(move || db.generate_id()).await??.to_string();
2020-06-06 22:43:33 +00:00
2020-06-14 15:07:31 +00:00
let key = alias_key(&hash.inner, &id);
2020-09-14 21:42:31 +00:00
let main_tree = self.inner.main_tree.clone();
2020-06-06 22:43:33 +00:00
let alias2 = alias.clone();
2020-06-14 15:07:31 +00:00
debug!("Saving hash/id -> alias mapping");
2020-06-06 22:43:33 +00:00
let res = web::block(move || {
2020-09-14 21:42:31 +00:00
main_tree.compare_and_swap(key, None as Option<sled::IVec>, Some(alias2.as_bytes()))
2020-06-06 22:43:33 +00:00
})
2021-02-10 22:57:42 +00:00
.await??;
2020-06-06 22:43:33 +00:00
if res.is_ok() {
2020-06-07 00:29:15 +00:00
let alias_tree = self.inner.alias_tree.clone();
let key = alias_id_key(&alias);
2020-06-14 15:07:31 +00:00
debug!("Saving alias -> id mapping");
2021-02-10 22:57:42 +00:00
web::block(move || alias_tree.insert(key.as_bytes(), id.as_bytes())).await??;
2020-06-07 00:29:15 +00:00
2020-06-06 22:43:33 +00:00
break;
}
2020-06-14 15:07:31 +00:00
debug!("Id exists, trying again");
2020-06-06 22:43:33 +00:00
}
Ok(())
2020-06-06 22:43:33 +00:00
}
// Generate an alias to the file
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, hash, content_type))]
2020-06-06 22:43:33 +00:00
async fn next_alias(
&self,
2020-06-14 15:07:31 +00:00
hash: &Hash,
2020-06-06 22:43:33 +00:00
content_type: mime::Mime,
) -> Result<String, UploadError> {
use rand::distributions::{Alphanumeric, Distribution};
let mut limit: usize = 10;
2021-02-10 22:57:42 +00:00
let mut rng = rand::thread_rng();
2020-06-06 22:43:33 +00:00
loop {
2020-06-14 15:07:31 +00:00
debug!("Alias gen loop");
2021-02-10 22:57:42 +00:00
let s: String = Alphanumeric
.sample_iter(&mut rng)
.take(limit)
.map(char::from)
.collect();
2020-06-24 16:58:46 +00:00
let alias = file_name(s, content_type.clone())?;
2020-06-06 22:43:33 +00:00
let res = self.save_alias(hash, &alias).await?;
2020-06-06 22:43:33 +00:00
if res.is_ok() {
return Ok(alias);
2020-06-06 22:43:33 +00:00
}
2020-06-14 15:07:31 +00:00
debug!("Alias exists, regenning");
2020-06-06 22:43:33 +00:00
limit += 1;
}
}
// Save an alias to the database
2020-06-14 18:56:42 +00:00
#[instrument(skip(self, hash))]
async fn save_alias(
&self,
2020-06-14 15:07:31 +00:00
hash: &Hash,
alias: &str,
) -> Result<Result<(), UploadError>, UploadError> {
let tree = self.inner.alias_tree.clone();
2020-06-14 15:07:31 +00:00
let vec = hash.inner.clone();
let alias = alias.to_string();
2020-06-14 15:07:31 +00:00
debug!("Saving alias");
let res = web::block(move || {
tree.compare_and_swap(alias.as_bytes(), None as Option<sled::IVec>, Some(vec))
})
2021-02-10 22:57:42 +00:00
.await??;
if res.is_err() {
2020-06-14 15:07:31 +00:00
warn!("Duplicate alias");
return Ok(Err(UploadError::DuplicateAlias));
}
Ok(Ok(()))
}
}
2020-06-14 15:07:31 +00:00
#[instrument(skip(stream))]
2021-01-19 01:54:39 +00:00
async fn safe_save_stream<E>(to: PathBuf, stream: UploadStream<E>) -> Result<(), UploadError>
where
UploadError: From<E>,
2020-06-15 02:41:45 +00:00
E: Unpin,
{
2020-06-15 02:41:45 +00:00
if let Some(path) = to.parent() {
debug!("Creating directory {:?}", path);
2021-01-19 01:54:39 +00:00
actix_fs::create_dir_all(path.to_owned()).await?;
2020-06-15 02:41:45 +00:00
}
debug!("Checking if {:?} already exists", to);
2021-01-19 01:54:39 +00:00
if let Err(e) = actix_fs::metadata(to.clone()).await {
if e.kind() != Some(std::io::ErrorKind::NotFound) {
2020-06-15 02:41:45 +00:00
return Err(e.into());
}
} else {
return Err(UploadError::FileExists);
}
2020-06-15 02:41:45 +00:00
debug!("Writing stream to {:?}", to);
2021-01-19 01:54:39 +00:00
let file = actix_fs::file::create(to).await?;
2021-01-19 01:54:39 +00:00
actix_fs::file::write_stream(file, stream.map_err(UploadError::from)).await?;
2020-06-15 02:41:45 +00:00
Ok(())
2020-06-06 22:43:33 +00:00
}
async fn remove_path(path: sled::IVec) -> Result<(), UploadError> {
let path_string = String::from_utf8(path.to_vec())?;
2021-01-19 01:54:39 +00:00
actix_fs::remove_file(path_string).await?;
2020-06-07 00:29:15 +00:00
Ok(())
}
fn trans_err(e: UploadError) -> sled::transaction::ConflictableTransactionError<UploadError> {
sled::transaction::ConflictableTransactionError::Abort(e)
}
2020-06-24 16:58:46 +00:00
fn file_name(name: String, content_type: mime::Mime) -> Result<String, UploadError> {
Ok(format!("{}{}", name, to_ext(content_type)?))
2020-06-06 21:41:17 +00:00
}
2020-06-07 00:29:15 +00:00
fn delete_key(alias: &str) -> String {
format!("{}/delete", alias)
}
fn variant_key(hash: &[u8], path: &str) -> Vec<u8> {
let mut key = hash.to_vec();
key.extend(&[2]);
key.extend(path.as_bytes());
key
}
fn variant_details_key(hash: &[u8], path: &str) -> Vec<u8> {
let mut key = hash.to_vec();
key.extend(&[2]);
key.extend(path.as_bytes());
key.extend(b"details");
key
}