fix: fly upload bug

CESSProject · Dec 16, 2023 · 0b7bc9f · 0b7bc9f
1 parent 67577ba
commit 0b7bc9f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 16 deletions.
diff --git a/pallets/file-bank/src/impls/receptionist.rs b/pallets/file-bank/src/impls/receptionist.rs
@@ -3,23 +3,32 @@ use crate::*;
 pub struct Receptionist<T: Config>(PhantomData<T>);
 
 impl<T: Config> Receptionist<T> {
-    pub fn fly_upload_file(file_hash: Hash, user_brief: UserBrief<T>, used_space: u128) -> DispatchResult {
-        T::StorageHandle::update_user_space(&user_brief.user, 1, used_space)?;
+    pub fn fly_upload_file(file_hash: Hash, user_brief: UserBrief<T>) -> DispatchResult {
 
-        if <Bucket<T>>::contains_key(&user_brief.user, &user_brief.bucket_name) {
-            Pallet::<T>::add_file_to_bucket(&user_brief.user, &user_brief.bucket_name, &file_hash)?;
-        } else {
-            Pallet::<T>::create_bucket_helper(&user_brief.user, &user_brief.bucket_name, Some(file_hash))?;
-        }
-
-        Pallet::<T>::add_user_hold_fileslice(&user_brief.user, file_hash, used_space)?;
 
         <File<T>>::try_mutate(&file_hash, |file_opt| -> DispatchResult {
             let file = file_opt.as_mut().ok_or(Error::<T>::FileNonExistent)?;
+            let needed_space = SEGMENT_SIZE
+                .checked_mul(15).ok_or(Error::<T>::Overflow)?
+                .checked_div(10).ok_or(Error::<T>::Overflow)?
+                .checked_mul(file.segment_list.len() as u128).ok_or(Error::<T>::Overflow)?;
+            ensure!(T::StorageHandle::get_user_avail_space(&user_brief.user)? > needed_space, Error::<T>::InsufficientAvailableSpace);
+            T::StorageHandle::update_user_space(&user_brief.user, 1, needed_space)?;
+
+            if <Bucket<T>>::contains_key(&user_brief.user, &user_brief.bucket_name) {
+                Pallet::<T>::add_file_to_bucket(&user_brief.user, &user_brief.bucket_name, &file_hash)?;
+            } else {
+                Pallet::<T>::create_bucket_helper(&user_brief.user, &user_brief.bucket_name, Some(file_hash))?;
+            }
+
+            Pallet::<T>::add_user_hold_fileslice(&user_brief.user, file_hash, needed_space)?;
             file.owner.try_push(user_brief.clone()).map_err(|_e| Error::<T>::BoundedVecError)?;
+
             Ok(())
         })?;
 
+
+
         Ok(())
     }
 

diff --git a/pallets/file-bank/src/lib.rs b/pallets/file-bank/src/lib.rs
@@ -435,15 +435,14 @@ pub mod pallet {
 			ensure!(user_brief.file_name.len() as u32 >= minimum, Error::<T>::SpecError);
 			ensure!(user_brief.bucket_name.len() as u32 >= minimum, Error::<T>::SpecError);
 
-			let needed_space = SEGMENT_SIZE
-				.checked_mul(15).ok_or(Error::<T>::Overflow)?
-				.checked_div(10).ok_or(Error::<T>::Overflow)?
-				.checked_mul(deal_info.len() as u128).ok_or(Error::<T>::Overflow)?;
-			ensure!(T::StorageHandle::get_user_avail_space(&user_brief.user)? > needed_space, Error::<T>::InsufficientAvailableSpace);
-
 			if <File<T>>::contains_key(&file_hash) {
-				Receptionist::<T>::fly_upload_file(file_hash, user_brief.clone(), needed_space)?;
+				Receptionist::<T>::fly_upload_file(file_hash, user_brief.clone())?;
 			} else {
+				let needed_space = SEGMENT_SIZE
+					.checked_mul(15).ok_or(Error::<T>::Overflow)?
+					.checked_div(10).ok_or(Error::<T>::Overflow)?
+					.checked_mul(deal_info.len() as u128).ok_or(Error::<T>::Overflow)?;
+            	ensure!(T::StorageHandle::get_user_avail_space(&user_brief.user)? > needed_space, Error::<T>::InsufficientAvailableSpace);
 				Receptionist::<T>::generate_deal(file_hash, deal_info, user_brief.clone(), needed_space, file_size)?;
 			}