From d60862a0d088bd579424cacd9aa0ba4dbb7c5adc Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 6 Jun 2020 19:29:15 -0500 Subject: [PATCH] Implement delete tokens --- src/error.rs | 16 ++++ src/main.rs | 26 ++++++- src/upload_manager.rs | 167 ++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 199 insertions(+), 10 deletions(-) diff --git a/src/error.rs b/src/error.rs index b3f656e..58c5216 100644 --- a/src/error.rs +++ b/src/error.rs @@ -20,6 +20,9 @@ pub enum UploadError { #[error("Error processing image, {0}")] Image(#[from] image::error::ImageError), + #[error("Error interacting with filesystem, {0}")] + Io(#[from] std::io::Error), + #[error("Panic in blocking operation")] Canceled, @@ -34,6 +37,18 @@ pub enum UploadError { #[error("Alias directed to missing file")] MissingFile, + + #[error("Provided token did not match expected token")] + InvalidToken, +} + +impl From> for UploadError { + fn from(e: sled::transaction::TransactionError) -> Self { + match e { + sled::transaction::TransactionError::Abort(t) => t, + sled::transaction::TransactionError::Storage(e) => e.into(), + } + } } impl From for UploadError { @@ -61,6 +76,7 @@ impl ResponseError for UploadError { StatusCode::BAD_REQUEST } UploadError::MissingAlias => StatusCode::NOT_FOUND, + UploadError::InvalidToken => StatusCode::FORBIDDEN, _ => StatusCode::INTERNAL_SERVER_ERROR, } } diff --git a/src/main.rs b/src/main.rs index 158d12a..0dbf650 100644 --- a/src/main.rs +++ b/src/main.rs @@ -76,7 +76,10 @@ fn from_ext(ext: std::ffi::OsString) -> mime::Mime { } /// Handle responding to succesful uploads -async fn upload(value: Value) -> Result { +async fn upload( + value: Value, + manager: web::Data, +) -> Result { let images = value .map() .and_then(|mut m| m.remove("images")) @@ -92,13 +95,28 @@ async fn upload(value: Value) -> Result { .and_then(|s| s.to_str()) { info!("Uploaded {} as {:?}", image.filename, saved_as); - files.push(serde_json::json!({ "file": saved_as })); + let delete_token = manager.delete_token(saved_as.to_owned()).await?; + files.push(serde_json::json!({ + "file": saved_as, + "delete_token": delete_token, + })); } } Ok(HttpResponse::Created().json(serde_json::json!({ "msg": "ok", "files": files }))) } +async fn delete( + manager: web::Data, + path_entries: web::Path<(String, String)>, +) -> Result { + let (alias, token) = path_entries.into_inner(); + + manager.delete(token, alias).await?; + + Ok(HttpResponse::NoContent().finish()) +} + /// Serve original files async fn serve( manager: web::Data, @@ -248,6 +266,10 @@ async fn main() -> Result<(), anyhow::Error> { .route(web::post().to(upload)), ) .service(web::resource("/{filename}").route(web::get().to(serve))) + .service( + web::resource("/delete/{delete_token}/{filename}") + .route(web::delete().to(delete)), + ) .service( web::resource("/{size}/{filename}").route(web::get().to(serve_resized)), ), diff --git a/src/upload_manager.rs b/src/upload_manager.rs index 90719fa..c45a243 100644 --- a/src/upload_manager.rs +++ b/src/upload_manager.rs @@ -1,6 +1,7 @@ use crate::{error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES}; use actix_web::web; use futures::stream::{Stream, StreamExt}; +use log::warn; use sha2::Digest; use std::{path::PathBuf, pin::Pin, sync::Arc}; @@ -12,8 +13,8 @@ pub struct UploadManager { struct UploadManagerInner { hasher: sha2::Sha256, image_dir: PathBuf, - db: sled::Db, alias_tree: sled::Tree, + db: sled::Db, } type UploadStream = Pin>>>; @@ -43,10 +44,7 @@ impl UploadManager { let mut sled_dir = root_dir.clone(); sled_dir.push("db"); // sled automatically creates it's own directories - // - // This is technically a blocking operation but it's fine because it happens before we - // start handling requests - let db = sled::open(sled_dir)?; + let db = web::block(move || sled::open(sled_dir)).await?; root_dir.push("files"); @@ -63,6 +61,108 @@ impl UploadManager { }) } + pub(crate) async fn delete(&self, alias: String, token: String) -> Result<(), UploadError> { + use sled::Transactional; + let db = self.inner.db.clone(); + let alias_tree = self.inner.alias_tree.clone(); + + let alias2 = alias.clone(); + let hash = web::block(move || { + [&*db, &alias_tree].transaction(|v| { + let db = &v[0]; + let alias_tree = &v[1]; + + // -- GET TOKEN -- + let existing_token = alias_tree + .remove(delete_key(&alias2).as_bytes())? + .ok_or(trans_err(UploadError::MissingAlias))?; + + // Bail if invalid token + if existing_token != token { + warn!("Invalid delete token"); + return Err(trans_err(UploadError::InvalidToken)); + } + + // -- GET ID FOR HASH TREE CLEANUP -- + let id = alias_tree + .remove(alias_id_key(&alias2).as_bytes())? + .ok_or(trans_err(UploadError::MissingAlias))?; + let id = String::from_utf8(id.to_vec()).map_err(|e| trans_err(e.into()))?; + + // -- GET HASH FOR HASH TREE CLEANUP -- + let hash = alias_tree + .remove(alias2.as_bytes())? + .ok_or(trans_err(UploadError::MissingAlias))?; + + // -- REMOVE HASH TREE ELEMENT -- + db.remove(alias_key(&hash, &id))?; + Ok(hash) + }) + }) + .await?; + + // -- CHECK IF ANY OTHER ALIASES EXIST + let db = self.inner.db.clone(); + let (start, end) = alias_key_bounds(&hash); + let any_aliases = web::block(move || { + Ok(db.range(start..end).next().is_some()) as Result + }) + .await?; + + // Bail if there are existing aliases + if any_aliases { + return Ok(()); + } + + // -- 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(); + + web::block(move || blocking_delete_all_by_filename(image_dir, &real_filename)).await?; + + Ok(()) + } + + /// Generate a delete token for an alias + pub(crate) async fn delete_token(&self, alias: String) -> Result { + use rand::distributions::{Alphanumeric, Distribution}; + let rng = rand::thread_rng(); + let s: String = Alphanumeric.sample_iter(rng).take(10).collect(); + let delete_token = s.clone(); + + 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, + Some(s.as_bytes()), + ) + }) + .await?; + + if let Err(sled::CompareAndSwapError { + current: Some(ivec), + .. + }) = res + { + let s = String::from_utf8(ivec.to_vec())?; + + return Ok(s); + } + + Ok(delete_token) + } + /// Upload the file, discarding bytes if it's already present, or saving if it's new pub(crate) async fn upload( &self, @@ -206,9 +306,7 @@ impl UploadManager { let db = self.inner.db.clone(); let id = web::block(move || db.generate_id()).await?.to_string(); - let mut key = hash.to_vec(); - key.extend(id.as_bytes()); - + let key = alias_key(hash, &id); let db = self.inner.db.clone(); let alias2 = alias.clone(); let res = web::block(move || { @@ -217,6 +315,10 @@ impl UploadManager { .await?; if res.is_ok() { + let alias_tree = self.inner.alias_tree.clone(); + let key = alias_id_key(&alias); + web::block(move || alias_tree.insert(key.as_bytes(), id.as_bytes())).await?; + break; } } @@ -255,6 +357,55 @@ 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)?; + } + + Ok(()) +} + +fn trans_err(e: UploadError) -> sled::transaction::ConflictableTransactionError { + sled::transaction::ConflictableTransactionError::Abort(e) +} + fn file_name(name: String, content_type: mime::Mime) -> String { format!("{}{}", name, to_ext(content_type)) } + +fn alias_key(hash: &[u8], id: &str) -> Vec { + let mut key = hash.to_vec(); + // add a separator to the key between the hash and the ID + key.extend(&[0]); + key.extend(id.as_bytes()); + + key +} + +fn alias_key_bounds(hash: &[u8]) -> (Vec, Vec) { + let mut start = hash.to_vec(); + start.extend(&[0]); + + let mut end = hash.to_vec(); + end.extend(&[1]); + + (start, end) +} + +fn alias_id_key(alias: &str) -> String { + format!("{}/id", alias) +} + +fn delete_key(alias: &str) -> String { + format!("{}/delete", alias) +}