refactor: optimise solo signer in regards to multisig

- Optimization included.
  - Block calculation updated. (for cleaning up the multisig storage)
- Included new tests.
- Updated existing tests with additional checks.
- Fixed one tiny bug that allowed users to fill out the multisig storage
  when the wrong user tried to sign the single signer script.

pallet/src/balance.rs

@@ -14,9 +14,7 @@ use move_vm_backend::balance::BalanceHandler;
 use sp_runtime::traits::Zero;
 use sp_std::{
     cell::{Ref, RefCell},
-    default::Default,
     rc::Rc,
-    vec::Vec,
 };
 
 use crate::{Config, Error, Pallet};

pallet/src/lib.rs

@@ -45,6 +45,7 @@ pub mod pallet {
         traits::{
             tokens::currency::LockIdentifier, Currency, Get, LockableCurrency, ReservableCurrency,
         },
+        BoundedBTreeSet,
     };
     use frame_system::pallet_prelude::*;
     pub use move_core_types::language_storage::TypeTag;
@@ -252,45 +253,65 @@ pub mod pallet {
             // Make sure the scripts are not maliciously trying to use forged signatures.
             let signer_count =
                 verify_script_integrity_and_check_signers(&bytecode).map_err(Error::<T>::from)?;
-            let accounts = Self::extract_account_ids_from_args(&args, signer_count)?;
-            let block_height = <frame_system::Pallet<T>>::block_number();
-
-            let (mut signature_handler, call_hash) = if signer_count > 1 {
-                // Generate the call hash to identify this multi-sign call request.
-                let call_hash = Self::transaction_bc_call_hash(&transaction_bc[..]);
+            let unique_signers = Self::extract_account_ids_from_args(&args, signer_count)?;
+
+            // Based on the number of unique signers, decide the following:
+            let (is_signature_required, contains_multisig) = match unique_signers.len() {
+                0 => (false, None),
+                1 => (true, None),
+                _ => (
+                    true,
+                    // Generate the unique script tx hash to identify this multi-sign call request.
+                    Some(Self::transaction_bc_call_hash(&transaction_bc[..])),
+                ),
+            };
 
-                let multisig = MultisigStorage::<T>::get(call_hash).unwrap_or(
-                    MultisigOf::<T>::new(accounts, block_height).map_err(Into::<Error<T>>::into)?,
+            let mut signature_handler = if let Some(script_hash) = contains_multisig {
+                let multisig_data = MultisigStorage::<T>::get(script_hash).unwrap_or(
+                    MultisigOf::<T>::new(unique_signers).map_err(Into::<Error<T>>::into)?,
                 );
 
-                (ScriptSignatureHandler::<T>::from(multisig), call_hash)
+                ScriptSignatureHandler::<T>::from(multisig_data)
             } else {
-                (
-                    ScriptSignatureHandler::<T>::new(accounts, block_height)?,
-                    [0u8; 32],
-                )
+                ScriptSignatureHandler::<T>::new(unique_signers)?
             };
-            if signer_count > 0 {
+
+            if is_signature_required {
                 let lock_id = Self::multi_signer_lock_id(&who, &transaction_bc[..], &cheque_limit);
                 signature_handler.sign_script(&who, &cheque_limit, lock_id)?;
             }
 
-            // If the script is signed correctly, we can execute it in MoveVM and update the
-            // blockchain storage or the balance sheet.
-            // If we have only one signer, it will skip this; if not, we have to wait for more signers, so we store it as a multi-signer request.
+            // The script needs to be signed by all signers before we can execute it in MoveVM and
+            // update the blockchain storage or the balance sheet.
             if !signature_handler.all_signers_approved() {
-                // Use the initial block height always when the first signer initiates the
-                // multi-signer execution request. This will prevent the lifetime of such requests
-                // from being extended with each signing.
-                let multisig = signature_handler.into_inner();
-                Self::new_multi_sign_request_chore(*multisig.block_number(), call_hash)?;
-                MultisigStorage::<T>::insert(call_hash, multisig);
+                // Everything below this block, applies to multisig scenario only.
+                let Some(script_hash) = contains_multisig else {
+                    // THIS PART SHOULD NOT BE REACHABLE - WE CAN RECONSIDER IT
+                    // If the solely required signer hasn't signed the script, reject the script entirely.
+                    return Err(Error::<T>::ScriptSignatureMissing.into());
+                };
+
+                let mut multisig = signature_handler.into_inner();
+
+                // The deadline for collecting all signatures is set by the first signer in the
+                // multisig scenario. There's no way to extend the initally set deadline.
+                if multisig.stored_block_height().is_none() {
+                    let block_height = <frame_system::Pallet<T>>::block_number();
+
+                    multisig.set_block_height(block_height);
+                    Self::new_multi_sign_request_chore(block_height, script_hash)?;
+                }
+
+                MultisigStorage::<T>::insert(script_hash, multisig);
+
                 Self::deposit_event(Event::SignedMultisigScript { who });
                 return result::execute_only_signing();
             }
+
             // If we have multiple signers and they all have signed, we have to remove the multi-signer request from the MultisigStorage.
-            if signer_count > 1 {
-                MultisigStorage::<T>::remove(call_hash);
+            if let Some(script_hash) = contains_multisig {
+                // TODO: Clear the ChoreOnIdleStorage as well here maybe?
+                MultisigStorage::<T>::remove(script_hash);
             }
 
             // We need to provide MoveVM read only access to balance sheet - MoveVM is allowed to
@@ -308,7 +329,11 @@ pub mod pallet {
             let result = result::from_vm_result::<T>(vm_result)?;
 
             // Emit an event.
-            let signers = signature_handler.into_signer_accounts()?;
+            let mut signers = signature_handler.into_signer_accounts()?;
+            if signers.is_empty() {
+                // Signer list can be empty in zero-signer scripts, so append here the user at least.
+                signers.push(who);
+            }
             Self::deposit_event(Event::ExecuteCalled { who: signers });
 
             Ok(result)
@@ -518,17 +543,20 @@ pub mod pallet {
         pub(crate) fn extract_account_ids_from_args(
             script_args: &[&[u8]],
             signer_count: usize,
-        ) -> Result<Vec<T::AccountId>, Error<T>> {
+        ) -> Result<BoundedBTreeSet<T::AccountId, T::MaxScriptSigners>, Error<T>> {
             if signer_count > script_args.len() {
                 return Err(Error::<T>::ScriptSignatureFailure);
             }
 
-            let mut accounts = Vec::<T::AccountId>::new();
+            let mut accounts = BoundedBTreeSet::<T::AccountId, T::MaxScriptSigners>::new();
             for signer in &script_args[..signer_count] {
                 let account_address =
                     bcs::from_bytes(signer).map_err(|_| Error::<T>::UnableToDeserializeAccount)?;
                 let account = Self::to_native_account(&account_address)?;
-                accounts.push(account);
+
+                accounts
+                    .try_insert(account)
+                    .map_err(|_| Error::<T>::NumberOfSignerArgumentsMismatch)?;
             }
 
             Ok(accounts)
@@ -570,7 +598,7 @@ pub mod pallet {
             let Ok(hashes_ready_for_cleanup) = ChoreOnIdleStorage::<T>::try_get(block) else {
                 return T::DbWeight::get().reads(1);
             };
-            // We don't need this entry in storage anymore, remove it!
+            // We don't need this entry in storage anymore, remove it.
             ChoreOnIdleStorage::<T>::remove(block);
 
             // Remove all that entries from MultisigStorage.
@@ -622,6 +650,10 @@ pub mod pallet {
         InsufficientBalance,
         /// Script signature failure.
         ScriptSignatureFailure,
+        /// If the script expects a list of signers, only users from that list can sign the transaction.
+        UnexpectedUserSignature,
+        /// Proper signature missing.
+        ScriptSignatureMissing,
         /// Script transaction cannot be deserialized.
         InvalidScriptTransaction,
         /// User tried to publish module in a protected memory area.

pallet/src/mock.rs

@@ -172,7 +172,7 @@ pub(crate) fn addrs_from_ss58(ss58: &str) -> Result<(AccountId32, AccountAddress
     Ok((addr_32, addr_mv))
 }
 
-/// Rolls forward in history to the given block height.
+/// Rolls forward in future to the given block height.
 pub(crate) fn roll_to(n: BlockNumberFor<Test>) {
     let weight = Weight::from_parts(100_000_000_000, 1);
     while System::block_number() < n {

pallet/src/signer.rs

@@ -10,12 +10,11 @@ use frame_support::{
         },
         Get,
     },
-    BoundedBTreeMap, Parameter,
+    BoundedBTreeMap, BoundedBTreeSet, Parameter,
 };
 use frame_system::{pallet_prelude::BlockNumberFor, Config as SysConfig};
 use scale_info::TypeInfo;
 use sp_runtime::traits::MaybeSerializeDeserialize;
-use sp_std::vec::Vec;
 
 use crate::{
     balance::{AccountIdOf, BalanceAdapter, BalanceOf},
@@ -80,8 +79,11 @@ where
 {
     /// The signers of a script transaction.
     signers: BoundedBTreeMap<AccountId, SignerData<Balance>, Size>,
-    /// The block number when this `Multisig` was created and stored.
-    block_height: BlockNumber,
+
+    /// The block height at which this `Multisig` was initally stored.
+    ///
+    /// Used only for multisigure purposes.
+    stored_block_height: Option<BlockNumber>,
 }
 
 impl<AccountId, Balance, BlockNumber, Size> Multisig<AccountId, Balance, BlockNumber, Size>
@@ -92,13 +94,13 @@ where
     Size: Get<u32>,
 {
     /// Creates a new [`Multisig`] with all blank signatures for the provided script.
-    pub fn new(signers: Vec<AccountId>, block_height: BlockNumber) -> Result<Self, MultisigError> {
+    pub fn new(signers: BoundedBTreeSet<AccountId, Size>) -> Result<Self, MultisigError> {
         if signers.len() > (Size::get() as usize) {
             return Err(MultisigError::MaxSignersExceeded);
         }
 
         let mut multisig_info = Multisig::<AccountId, Balance, BlockNumber, Size> {
-            block_height,
+            stored_block_height: None,
             ..Default::default()
         };
         for account in signers.iter() {
@@ -110,8 +112,12 @@ where
         Ok(multisig_info)
     }
 
-    pub fn block_number(&self) -> &BlockNumber {
-        &self.block_height
+    pub fn stored_block_height(&self) -> Option<&BlockNumber> {
+        self.stored_block_height.as_ref()
+    }
+
+    pub fn set_block_height(&mut self, block_height: BlockNumber) {
+        self.stored_block_height = Some(block_height);
     }
 }
 
@@ -155,7 +161,7 @@ where
     fn default() -> Self {
         Multisig {
             signers: BoundedBTreeMap::<AccountId, SignerData<Balance>, Size>::new(),
-            block_height: BlockNumber::default(),
+            stored_block_height: None,
         }
     }
 }
@@ -170,13 +176,11 @@ pub(crate) struct ScriptSignatureHandler<T: Config + SysConfig> {
 impl<T: Config + SysConfig> ScriptSignatureHandler<T> {
     /// Creates a new [`ScriptSignatureHandler`] with all blank signatures for the provided script.
     pub(crate) fn new(
-        accounts: Vec<T::AccountId>,
-        block_height: BlockNumberFor<T>,
+        accounts: BoundedBTreeSet<T::AccountId, T::MaxScriptSigners>,
     ) -> Result<Self, Error<T>> {
         Ok(Self {
             _pd_config: PhantomData,
-            multisig_info: MultisigOf::<T>::new(accounts, block_height)
-                .map_err(Into::<Error<T>>::into)?,
+            multisig_info: MultisigOf::<T>::new(accounts).map_err(Into::<Error<T>>::into)?,
         })
     }
 
@@ -193,19 +197,21 @@ impl<T: Config + SysConfig> ScriptSignatureHandler<T> {
         cheque_limit: &BalanceOf<T>,
         lock_id: LockIdentifier,
     ) -> Result<(), Error<T>> {
-        if let Some(ms_data) = self.multisig_info.get_mut(account) {
-            if matches!(ms_data.signature, Signature::Approved) {
-                Err(Error::<T>::UserHasAlreadySigned)
-            } else {
-                ms_data.signature = Signature::Approved;
-                ms_data.cheque_limit = *cheque_limit;
-                ms_data.lock_id = lock_id;
-                T::Currency::set_lock(lock_id, account, *cheque_limit, WithdrawReasons::all());
-                Ok(())
-            }
-        } else {
-            Err(Error::<T>::ScriptSignatureFailure)
+        // Only users that are on the signer list can sign this script.
+        let Some(ms_data) = self.multisig_info.get_mut(account) else {
+            return Err(Error::<T>::UnexpectedUserSignature);
+        };
+
+        // User can sign only once.
+        if matches!(ms_data.signature, Signature::Approved) {
+            return Err(Error::<T>::UserHasAlreadySigned);
         }
+
+        ms_data.signature = Signature::Approved;
+        ms_data.cheque_limit = *cheque_limit;
+        ms_data.lock_id = lock_id;
+        T::Currency::set_lock(lock_id, account, *cheque_limit, WithdrawReasons::all());
+        Ok(())
     }
 
     /// Check whether the script has been approved by all required signers.