diff --git a/Cargo.toml b/Cargo.toml index 75241ec..8e283d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "pict-rs" description = "A simple image hosting service" -version = "0.1.1" +version = "0.1.2" authors = ["asonix "] license = "AGPL-3.0" readme = "README.md" diff --git a/src/error.rs b/src/error.rs index 5f5623e..406995f 100644 --- a/src/error.rs +++ b/src/error.rs @@ -58,6 +58,9 @@ pub enum UploadError { #[error("No filename provided in request")] MissingFilename, + + #[error("Error converting Path to String")] + Path, } impl From for UploadError { diff --git a/src/main.rs b/src/main.rs index 97321d0..11aaeb3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -213,6 +213,11 @@ async fn serve( // Save the file in another task, we want to return the thumbnail now actix_rt::spawn(async move { + if let Err(e) = manager.store_variant(path2.clone()).await { + error!("Error storing variant, {}", e); + return; + } + if let Err(e) = safe_save_file(path2, img_bytes2).await { error!("Error saving file, {}", e); } diff --git a/src/upload_manager.rs b/src/upload_manager.rs index fa243ea..b471614 100644 --- a/src/upload_manager.rs +++ b/src/upload_manager.rs @@ -15,6 +15,7 @@ struct UploadManagerInner { hasher: sha2::Sha256, image_dir: PathBuf, alias_tree: sled::Tree, + filename_tree: sled::Tree, db: sled::Db, } @@ -61,11 +62,32 @@ impl UploadManager { hasher: sha2::Sha256::new(), image_dir: root_dir, alias_tree: db.open_tree("alias")?, + filename_tree: db.open_tree("filename")?, db, }), }) } + pub(crate) async fn store_variant(&self, path: PathBuf) -> Result<(), UploadError> { + let filename = path + .file_name() + .and_then(|f| f.to_str()) + .map(|s| s.to_string()) + .ok_or(UploadError::Path)?; + let path_string = path.to_str().ok_or(UploadError::Path)?.to_string(); + + let fname_tree = self.inner.filename_tree.clone(); + let hash: sled::IVec = web::block(move || fname_tree.get(filename.as_bytes())) + .await? + .ok_or(UploadError::MissingFilename)?; + + let key = variant_key(&hash, &path_string); + let db = self.inner.db.clone(); + web::block(move || db.insert(key, path_string.as_bytes())).await?; + + Ok(()) + } + pub(crate) async fn delete(&self, alias: String, token: String) -> Result<(), UploadError> { use sled::Transactional; let db = self.inner.db.clone(); @@ -121,22 +143,15 @@ impl UploadManager { // -- DELETE HASH ENTRY -- let db = self.inner.db.clone(); - let real_filename = web::block(move || { - let real_filename = db.remove(&hash)?.ok_or(UploadError::MissingFile)?; - - Ok(real_filename) as Result - }) - .await?; - - let real_filename = String::from_utf8(real_filename.to_vec())?; - - let image_dir = self.image_dir(); + let hash2 = hash.clone(); + let filename = web::block(move || db.remove(&hash2)) + .await? + .ok_or(UploadError::MissingFile)?; // -- DELETE FILES -- + let this = self.clone(); actix_rt::spawn(async move { - if let Err(e) = - web::block(move || blocking_delete_all_by_filename(image_dir, &real_filename)).await - { + if let Err(e) = this.cleanup_files(filename).await { error!("Error removing files from fs, {}", e); } }); @@ -144,6 +159,48 @@ impl UploadManager { Ok(()) } + async fn cleanup_files(&self, filename: sled::IVec) -> Result<(), UploadError> { + let mut path = self.image_dir(); + let fname = String::from_utf8(filename.to_vec())?; + path.push(fname); + + let mut errors = Vec::new(); + if let Err(e) = actix_fs::remove_file(path).await { + errors.push(e.into()); + } + + let fname_tree = self.inner.filename_tree.clone(); + let hash = web::block(move || fname_tree.remove(filename)) + .await? + .ok_or(UploadError::MissingFile)?; + + let (start, end) = variant_key_bounds(&hash); + let db = self.inner.db.clone(); + let keys = web::block(move || { + let mut keys = Vec::new(); + for key in db.range(start..end).keys() { + keys.push(key?.to_owned()); + } + + Ok(keys) as Result, UploadError> + }) + .await?; + + for key in keys { + let db = self.inner.db.clone(); + if let Some(path) = web::block(move || db.remove(key)).await? { + if let Err(e) = remove_path(path).await { + errors.push(e); + } + } + } + + for error in errors { + error!("Error deleting files, {}", error); + } + Ok(()) + } + /// Generate a delete token for an alias pub(crate) async fn delete_token(&self, alias: String) -> Result { use rand::distributions::{Alphanumeric, Distribution}; @@ -279,8 +336,13 @@ impl UploadManager { let filename = self.next_file(content_type).await?; let filename2 = filename.clone(); + let hash2 = hash.clone(); let res = web::block(move || { - db.compare_and_swap(hash, None as Option, Some(filename2.as_bytes())) + db.compare_and_swap( + hash2, + None as Option, + Some(filename2.as_bytes()), + ) }) .await?; @@ -293,6 +355,10 @@ impl UploadManager { return Ok((Dup::Exists, name)); } + let fname_tree = self.inner.filename_tree.clone(); + let filename2 = filename.clone(); + web::block(move || fname_tree.insert(filename2, hash)).await?; + Ok((Dup::New, filename)) } @@ -386,21 +452,9 @@ impl UploadManager { } } -fn blocking_delete_all_by_filename(mut dir: PathBuf, filename: &str) -> Result<(), UploadError> { - for res in std::fs::read_dir(dir.clone())? { - let entry = res?; - - if entry.path().is_dir() { - blocking_delete_all_by_filename(entry.path(), filename)?; - } - } - - dir.push(filename); - - if dir.is_file() { - std::fs::remove_file(dir)?; - } - +async fn remove_path(path: sled::IVec) -> Result<(), UploadError> { + let path_string = String::from_utf8(path.to_vec())?; + actix_fs::remove_file(path_string).await?; Ok(()) } @@ -439,6 +493,23 @@ fn delete_key(alias: &str) -> String { format!("{}/delete", alias) } +fn variant_key(hash: &[u8], path: &str) -> Vec { + let mut key = hash.to_vec(); + key.extend(&[2]); + key.extend(path.as_bytes()); + key +} + +fn variant_key_bounds(hash: &[u8]) -> (Vec, Vec) { + let mut start = hash.to_vec(); + start.extend(&[2]); + + let mut end = hash.to_vec(); + end.extend(&[3]); + + (start, end) +} + fn valid_format(format: image::ImageFormat) -> Result { match format { image::ImageFormat::Jpeg => Ok(mime::IMAGE_JPEG),