diff --git a/src/lib.rs b/src/lib.rs index 8797421..51e1824 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -79,7 +79,7 @@ pub mod pallet { }; use frame_system::pallet_prelude::*; use sp_core::crypto::AccountId32; - use sp_runtime::traits::{AccountIdConversion, Saturating, Verify}; + use sp_runtime::traits::{AccountIdConversion, CheckedAdd, Saturating, Verify}; use sp_runtime::{MultiSignature, Perbill}; use sp_std::vec::Vec; @@ -261,8 +261,7 @@ pub mod pallet { // Pallet is configured with a zero vesting period. info.total_reward - first_paid } else { - (info.total_reward - first_paid) - .saturating_mul(payable_period.into()) + (info.total_reward - first_paid).saturating_mul(payable_period.into()) / period.into() }; @@ -329,7 +328,11 @@ pub mod pallet { #[pallet::weight(0)] pub fn initialize_reward_vec( origin: OriginFor, - rewards: Vec<(T::RelayChainAccountId, Option<::AccountId>, BalanceOf)>, + rewards: Vec<( + T::RelayChainAccountId, + Option<::AccountId>, + BalanceOf, + )>, index: u32, limit: u32, ) -> DispatchResultWithPostInfo { @@ -346,22 +349,35 @@ pub mod pallet { // Total number of contributors let mut total_contributors = TotalContributors::::get(); - let incoming_rewards: BalanceOf = rewards + let incoming_rewards: Option> = rewards .iter() - .fold(0u32.into(), |acc: BalanceOf, (_, _, reward)| { - acc + *reward + .try_fold(0u32.into(), |acc: BalanceOf, (_, _, reward)| { + acc.checked_add(reward) }); + // This is where we need to ensure we do not overflow. All subsequent payments + // will be based on these checks. + + // Ensure we do not overflow. + ensure!(incoming_rewards.is_some(), Error::::AdditionOverflow); + + let total_rewards = incoming_rewards + .unwrap() + .checked_add(¤t_initialized_rewards); + + // Ensure we do not overflow + ensure!(total_rewards.is_some(), Error::::AdditionOverflow); + // Ensure we dont go over funds ensure!( - current_initialized_rewards + incoming_rewards <= Self::pot(), + total_rewards.unwrap() <= Self::pot(), Error::::BatchBeyondFundPot ); // Let's ensure we can close initialization if index + rewards.len() as u32 == limit { ensure!( - current_initialized_rewards + incoming_rewards == Self::pot(), + total_rewards.unwrap() == Self::pot(), Error::::RewardsDoNotMatchFund ); } @@ -451,6 +467,8 @@ pub mod pallet { /// User trying to associate a native identity with a relay chain identity for posterior /// reward claiming provided an already associated relay chain identity AlreadyAssociated, + /// Addition overflow when initializating reward vec + AdditionOverflow, /// Trying to introduce a batch that goes beyond the limits of the funds BatchBeyondFundPot, /// First claim already done @@ -539,12 +557,19 @@ pub mod pallet { InitialPaymentMade(::AccountId, BalanceOf), /// Someone has proven they made a contribution and associated a native identity with it. /// Data is the relay account, native account and the total amount of _rewards_ that will be paid - NativeIdentityAssociated(T::RelayChainAccountId, ::AccountId, BalanceOf), + NativeIdentityAssociated( + T::RelayChainAccountId, + ::AccountId, + BalanceOf, + ), /// A contributor has claimed some rewards. /// Data is the account getting paid and the amount of rewards paid. RewardsPaid(::AccountId, BalanceOf), /// A contributor has updated the reward address. - RewardAddressUpdated(::AccountId, ::AccountId), + RewardAddressUpdated( + ::AccountId, + ::AccountId, + ), /// When initializing the reward vec an already initialized account was found InitializedAlreadyInitializedAccount( T::RelayChainAccountId, diff --git a/src/tests.rs b/src/tests.rs index c1e1971..f912da9 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -441,7 +441,7 @@ fn initialize_new_addresses_with_batch() { 0, DispatchError::Module { index: 2, - error: 8, + error: 9, message: None, }, ), @@ -602,3 +602,42 @@ fn reward_below_vesting_period_works() { assert_eq!(events(), expected); }); } + +#[test] +fn cannot_overflow_contributions_initializing() { + empty().execute_with(|| { + assert_noop!( + Crowdloan::initialize_reward_vec( + Origin::root(), + vec![ + ([1u8; 32].into(), Some(1), u128::MAX.into()), + ([2u8; 32].into(), Some(2), u128::MAX.into()), + ], + 0, + 2 + ), + Error::::AdditionOverflow + ); + }); +} + +#[test] +fn cannot_overflow_contributions_initializing_already_stored() { + empty().execute_with(|| { + assert_ok!(Crowdloan::initialize_reward_vec( + Origin::root(), + vec![([1u8; 32].into(), Some(1), 500u32.into())], + 0, + 2 + )); + assert_noop!( + Crowdloan::initialize_reward_vec( + Origin::root(), + vec![([2u8; 32].into(), Some(1), u128::MAX.into())], + 1, + 2 + ), + Error::::AdditionOverflow + ); + }); +}