From 0d153ca239450e07cef5ec6b781e8d01c323072e Mon Sep 17 00:00:00 2001 From: "Aode (lion)" Date: Tue, 16 Nov 2021 09:45:16 -0600 Subject: [PATCH] Make stores take filename into account when saving files --- src/main.rs | 2 +- src/store.rs | 7 +++++- src/store/file_store.rs | 16 ++++++++------ src/store/object_store.rs | 15 ++++++++----- src/upload_manager.rs | 4 ++-- src/upload_manager/session.rs | 41 +++++++++++++++++++---------------- 6 files changed, 49 insertions(+), 36 deletions(-) diff --git a/src/main.rs b/src/main.rs index 1a4fb6b..604ca04 100644 --- a/src/main.rs +++ b/src/main.rs @@ -411,7 +411,7 @@ where let bytes2 = bytes.clone(); actix_rt::spawn( async move { - let identifier = match store.save_bytes(bytes2).await { + let identifier = match store.save_bytes(bytes2, &name).await { Ok(identifier) => identifier, Err(e) => { tracing::warn!("Failed to generate directory path: {}", e); diff --git a/src/store.rs b/src/store.rs index 9bb2afe..d9af237 100644 --- a/src/store.rs +++ b/src/store.rs @@ -27,11 +27,16 @@ pub(crate) trait Store: Send + Sync + Clone + Debug + 'static { async fn save_async_read( &self, reader: &mut Reader, + filename: &str, ) -> Result where Reader: AsyncRead + Unpin; - async fn save_bytes(&self, bytes: Bytes) -> Result; + async fn save_bytes( + &self, + bytes: Bytes, + filename: &str, + ) -> Result; async fn to_stream( &self, diff --git a/src/store/file_store.rs b/src/store/file_store.rs index 6519c21..d473764 100644 --- a/src/store/file_store.rs +++ b/src/store/file_store.rs @@ -8,7 +8,6 @@ use std::{ use storage_path_generator::Generator; use tokio::io::{AsyncRead, AsyncWrite}; use tracing::{debug, error, instrument}; -use uuid::Uuid; mod file_id; mod restructure; @@ -58,11 +57,12 @@ impl Store for FileStore { async fn save_async_read( &self, reader: &mut Reader, + filename: &str, ) -> Result where Reader: AsyncRead + Unpin, { - let path = self.next_file()?; + let path = self.next_file(filename)?; if let Err(e) = self.safe_save_reader(&path, reader).await { self.safe_remove_file(&path).await?; @@ -73,8 +73,12 @@ impl Store for FileStore { } #[tracing::instrument(skip(bytes))] - async fn save_bytes(&self, bytes: Bytes) -> Result { - let path = self.next_file()?; + async fn save_bytes( + &self, + bytes: Bytes, + filename: &str, + ) -> Result { + let path = self.next_file(filename)?; if let Err(e) = self.safe_save_bytes(&path, bytes).await { self.safe_remove_file(&path).await?; @@ -163,11 +167,9 @@ impl FileStore { Ok(target_path) } - fn next_file(&self) -> Result { + fn next_file(&self, filename: &str) -> Result { let target_path = self.next_directory()?; - let filename = Uuid::new_v4().to_string(); - Ok(target_path.join(filename)) } diff --git a/src/store/object_store.rs b/src/store/object_store.rs index 5239257..d4d9ffa 100644 --- a/src/store/object_store.rs +++ b/src/store/object_store.rs @@ -11,7 +11,6 @@ use std::{ }; use storage_path_generator::{Generator, Path}; use tokio::io::{AsyncRead, AsyncWrite}; -use uuid::Uuid; mod object_id; pub(crate) use object_id::ObjectId; @@ -64,11 +63,12 @@ impl Store for ObjectStore { async fn save_async_read( &self, reader: &mut Reader, + filename: &str, ) -> Result where Reader: AsyncRead + Unpin, { - let path = self.next_file()?; + let path = self.next_file(filename)?; self.bucket .put_object_stream(&self.client, reader, &path) @@ -78,8 +78,12 @@ impl Store for ObjectStore { } #[tracing::instrument(skip(bytes))] - async fn save_bytes(&self, bytes: Bytes) -> Result { - let path = self.next_file()?; + async fn save_bytes( + &self, + bytes: Bytes, + filename: &str, + ) -> Result { + let path = self.next_file(filename)?; self.bucket.put_object(&self.client, &path, &bytes).await?; @@ -195,8 +199,7 @@ impl ObjectStore { Ok(path) } - fn next_file(&self) -> Result { - let filename = Uuid::new_v4().to_string(); + fn next_file(&self, filename: &str) -> Result { let path = self.next_directory()?.to_strings().join("/"); Ok(format!("{}/{}", path, filename)) diff --git a/src/upload_manager.rs b/src/upload_manager.rs index 1de0848..e934d31 100644 --- a/src/upload_manager.rs +++ b/src/upload_manager.rs @@ -120,7 +120,7 @@ impl UploadManager { futures_util::pin_mut!(stream); let mut reader = tokio_util::io::StreamReader::new(stream); - let new_identifier = to.save_async_read(&mut reader).await?; + let new_identifier = to.save_async_read(&mut reader, &filename).await?; let details_key = self.details_key(&identifier, &filename)?; @@ -194,7 +194,7 @@ impl UploadManager { ThumbnailFormat::Jpeg, ) .await?; - let motion_identifier = store.save_async_read(&mut reader).await?; + let motion_identifier = store.save_async_read(&mut reader, &filename).await?; drop(permit); self.store_motion_path(&filename, &motion_identifier) diff --git a/src/upload_manager/session.rs b/src/upload_manager/session.rs index 24415b1..9274697 100644 --- a/src/upload_manager/session.rs +++ b/src/upload_manager/session.rs @@ -167,7 +167,12 @@ where let mut hasher_reader = Hasher::new(validated_reader, self.manager.inner.hasher.clone()); - let identifier = self.store.save_async_read(&mut hasher_reader).await?; + let filename = self.next_file(content_type).await?; + + let identifier = self + .store + .save_async_read(&mut hasher_reader, &filename) + .await?; let hash = hasher_reader.finalize_reset().await?; debug!("Storing alias"); @@ -175,7 +180,7 @@ where self.add_existing_alias(&hash, &alias).await?; debug!("Saving file"); - self.save_upload(&identifier, hash, content_type).await?; + self.save_upload(&identifier, hash, filename).await?; // Return alias to file Ok(self) @@ -205,14 +210,19 @@ where let mut hasher_reader = Hasher::new(validated_reader, self.manager.inner.hasher.clone()); - let identifier = self.store.save_async_read(&mut hasher_reader).await?; + let filename = self.next_file(input_type).await?; + + let identifier = self + .store + .save_async_read(&mut hasher_reader, &filename) + .await?; let hash = hasher_reader.finalize_reset().await?; debug!("Adding alias"); self.add_alias(&hash, input_type).await?; debug!("Saving file"); - self.save_upload(&identifier, hash, input_type).await?; + self.save_upload(&identifier, hash, filename).await?; // Return alias to file Ok(self) @@ -223,9 +233,9 @@ where &self, identifier: &S::Identifier, hash: Hash, - input_type: ValidInputType, + filename: String, ) -> Result<(), Error> { - let (dup, name) = self.check_duplicate(hash, input_type).await?; + let dup = self.check_duplicate(hash, filename.clone()).await?; // bail early with alias to existing file if this is a duplicate if dup.exists() { @@ -235,22 +245,16 @@ where return Ok(()); } - self.manager.store_identifier(name, identifier).await?; + self.manager.store_identifier(filename, identifier).await?; Ok(()) } // check for an already-uploaded image with this hash, returning the path to the target file - #[instrument(skip(self, hash, input_type))] - async fn check_duplicate( - &self, - hash: Hash, - input_type: ValidInputType, - ) -> Result<(Dup, String), Error> { + #[instrument(skip(self, hash))] + async fn check_duplicate(&self, hash: Hash, filename: String) -> Result { let main_tree = self.manager.inner.main_tree.clone(); - let filename = self.next_file(input_type).await?; - let filename2 = filename.clone(); let hash2 = hash.as_slice().to_vec(); debug!("Inserting filename for hash"); @@ -270,15 +274,14 @@ where { let name = String::from_utf8(ivec.to_vec())?; debug!("Filename exists for hash, {}", name); - return Ok((Dup::Exists, name)); + return Ok(Dup::Exists); } let fname_tree = self.manager.inner.filename_tree.clone(); - let filename2 = filename.clone(); debug!("Saving filename -> hash relation"); - web::block(move || fname_tree.insert(filename2, hash.into_inner())).await??; + web::block(move || fname_tree.insert(filename, hash.into_inner())).await??; - Ok((Dup::New, filename)) + Ok(Dup::New) } // generate a short filename that isn't already in-use