diff --git a/src/error.rs b/src/error.rs index 02e1a31..b3f656e 100644 --- a/src/error.rs +++ b/src/error.rs @@ -28,6 +28,12 @@ pub enum UploadError { #[error("Uploaded image could not be served, extension is missing")] MissingExtension, + + #[error("Requested a file that doesn't exist")] + MissingAlias, + + #[error("Alias directed to missing file")] + MissingFile, } impl From for UploadError { @@ -54,6 +60,7 @@ impl ResponseError for UploadError { UploadError::NoFiles | UploadError::ContentType(_) | UploadError::Upload(_) => { StatusCode::BAD_REQUEST } + UploadError::MissingAlias => StatusCode::NOT_FOUND, _ => StatusCode::INTERNAL_SERVER_ERROR, } } diff --git a/src/main.rs b/src/main.rs index be23b13..158d12a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -102,10 +102,11 @@ async fn upload(value: Value) -> Result { /// Serve original files async fn serve( manager: web::Data, - filename: web::Path, + alias: web::Path, ) -> Result { + let filename = manager.from_alias(alias.into_inner()).await?; let mut path = manager.image_dir(); - path.push(filename.into_inner()); + path.push(filename); let ext = path .extension() @@ -121,13 +122,14 @@ async fn serve( /// Serve resized files async fn serve_resized( manager: web::Data, - filename: web::Path<(u32, String)>, + path_entries: web::Path<(u32, String)>, ) -> Result { use image::GenericImageView; let mut path = manager.image_dir(); - let (size, name) = filename.into_inner(); + let (size, alias) = path_entries.into_inner(); + let name = manager.from_alias(alias).await?; path.push(size.to_string()); path.push(name.clone()); diff --git a/src/upload_manager.rs b/src/upload_manager.rs index 792f9c2..90719fa 100644 --- a/src/upload_manager.rs +++ b/src/upload_manager.rs @@ -13,6 +13,7 @@ struct UploadManagerInner { hasher: sha2::Sha256, image_dir: PathBuf, db: sled::Db, + alias_tree: sled::Tree, } type UploadStream = Pin>>>; @@ -56,6 +57,7 @@ impl UploadManager { inner: Arc::new(UploadManagerInner { hasher: sha2::Sha256::new(), image_dir: root_dir, + alias_tree: db.open_tree("alias")?, db, }), }) @@ -86,21 +88,46 @@ impl UploadManager { // Cloning bytes is fine because it's actually a pointer let hash = self.hash(bytes.clone()).await?; - let (dup, path) = self.check_duplicate(hash, content_type).await?; + let alias = self.add_alias(&hash, content_type.clone()).await?; + let (dup, name) = self.check_duplicate(hash, content_type).await?; - // bail early with path to existing file if this is a duplicate + // bail early with alias to existing file if this is a duplicate if dup.exists() { + let mut path = PathBuf::new(); + path.push(alias); return Ok(Some(path)); } // TODO: validate image before saving // -- WRITE NEW FILE -- - safe_save_file(path.clone(), bytes).await?; + let mut real_path = self.image_dir(); + real_path.push(name); + safe_save_file(real_path, bytes).await?; + + // Return alias to file + let mut path = PathBuf::new(); + path.push(alias); Ok(Some(path)) } + pub(crate) async fn from_alias(&self, alias: String) -> Result { + let tree = self.inner.alias_tree.clone(); + let hash = web::block(move || tree.get(alias.as_bytes())) + .await? + .ok_or(UploadError::MissingAlias)?; + + let db = self.inner.db.clone(); + let filename = web::block(move || db.get(hash)) + .await? + .ok_or(UploadError::MissingFile)?; + + let filename = String::from_utf8(filename.to_vec())?; + + Ok(filename) + } + // produce a sh256sum of the uploaded file async fn hash(&self, bytes: bytes::Bytes) -> Result, UploadError> { let mut hasher = self.inner.hasher.clone(); @@ -118,12 +145,10 @@ impl UploadManager { &self, hash: Vec, content_type: mime::Mime, - ) -> Result<(Dup, PathBuf), UploadError> { - let mut path = self.inner.image_dir.clone(); + ) -> Result<(Dup, String), UploadError> { let db = self.inner.db.clone(); let filename = self.next_file(content_type).await?; - let filename2 = filename.clone(); let res = web::block(move || { db.compare_and_swap(hash, None as Option, Some(filename2.as_bytes())) @@ -136,19 +161,15 @@ impl UploadManager { }) = res { let name = String::from_utf8(ivec.to_vec())?; - path.push(name); - - return Ok((Dup::Exists, path)); + return Ok((Dup::Exists, name)); } - path.push(filename); - - Ok((Dup::New, path)) + Ok((Dup::New, filename)) } // generate a short filename that isn't already in-use async fn next_file(&self, content_type: mime::Mime) -> Result { - let image_dir = self.inner.image_dir.clone(); + let image_dir = self.image_dir(); use rand::distributions::{Alphanumeric, Distribution}; let mut limit: usize = 10; let rng = rand::thread_rng(); @@ -156,7 +177,7 @@ impl UploadManager { let mut path = image_dir.clone(); let s: String = Alphanumeric.sample_iter(rng).take(limit).collect(); - let filename = format!("{}{}", s, to_ext(content_type.clone())); + let filename = file_name(s, content_type.clone()); path.push(filename.clone()); @@ -170,4 +191,70 @@ impl UploadManager { limit += 1; } } + + // 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 + async fn add_alias( + &self, + hash: &[u8], + content_type: mime::Mime, + ) -> Result { + let alias = self.next_alias(hash, content_type).await?; + + loop { + 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 db = self.inner.db.clone(); + let alias2 = alias.clone(); + let res = web::block(move || { + db.compare_and_swap(key, None as Option, Some(alias2.as_bytes())) + }) + .await?; + + if res.is_ok() { + break; + } + } + + Ok(alias) + } + + // Generate an alias to the file + async fn next_alias( + &self, + hash: &[u8], + content_type: mime::Mime, + ) -> Result { + use rand::distributions::{Alphanumeric, Distribution}; + let mut limit: usize = 10; + let rng = rand::thread_rng(); + let hvec = hash.to_vec(); + loop { + let s: String = Alphanumeric.sample_iter(rng).take(limit).collect(); + let filename = file_name(s, content_type.clone()); + + let tree = self.inner.alias_tree.clone(); + let vec = hvec.clone(); + let filename2 = filename.clone(); + let res = web::block(move || { + tree.compare_and_swap(filename2.as_bytes(), None as Option, Some(vec)) + }) + .await?; + + if res.is_ok() { + return Ok(filename); + } + + limit += 1; + } + } +} + +fn file_name(name: String, content_type: mime::Mime) -> String { + format!("{}{}", name, to_ext(content_type)) }