diff --git a/vesting/src/lib.rs b/vesting/src/lib.rs index 57abc6d3c..b010d34ce 100644 --- a/vesting/src/lib.rs +++ b/vesting/src/lib.rs @@ -31,8 +31,11 @@ use codec::HasCompact; use frame_support::{ ensure, pallet_prelude::*, - traits::{Currency, EnsureOrigin, ExistenceRequirement, Get, LockIdentifier, LockableCurrency, WithdrawReasons}, - transactional, + traits::{ + Currency, EnsureOrigin, ExistenceRequirement, Get, LockIdentifier, LockableCurrency, MaxEncodedLen, + WithdrawReasons, + }, + transactional, BoundedVec, }; use frame_system::{ensure_root, ensure_signed, pallet_prelude::*}; use sp_runtime::{ @@ -41,6 +44,7 @@ use sp_runtime::{ }; use sp_std::{ cmp::{Eq, PartialEq}, + convert::TryInto, vec::Vec, }; @@ -51,16 +55,13 @@ mod weights; pub use module::*; pub use weights::WeightInfo; -/// The maximum number of vesting schedules an account can have. -pub const MAX_VESTINGS: usize = 20; - pub const VESTING_LOCK_ID: LockIdentifier = *b"ormlvest"; /// The vesting schedule. /// /// Benefits would be granted gradually, `per_period` amount every `period` /// of blocks after `start`. -#[derive(Clone, Encode, Decode, PartialEq, Eq, RuntimeDebug)] +#[derive(Clone, Encode, Decode, PartialEq, Eq, RuntimeDebug, MaxEncodedLen)] pub struct VestingSchedule { /// Vesting starting block pub start: BlockNumber, @@ -136,6 +137,9 @@ pub mod module { /// Weight information for extrinsics in this module. type WeightInfo: WeightInfo; + + /// The maximum vesting schedules + type MaxVestingSchedules: Get; } #[pallet::error] @@ -150,6 +154,8 @@ pub mod module { TooManyVestingSchedules, /// The vested transfer amount is too low AmountLow, + /// Failed because the maximum vesting schedules was exceeded + MaxVestingSchedulesExceeded, } #[pallet::event] @@ -166,8 +172,13 @@ pub mod module { /// Vesting schedules of an account. #[pallet::storage] #[pallet::getter(fn vesting_schedules)] - pub type VestingSchedules = - StorageMap<_, Blake2_128Concat, T::AccountId, Vec>, ValueQuery>; + pub type VestingSchedules = StorageMap< + _, + Blake2_128Concat, + T::AccountId, + BoundedVec, T::MaxVestingSchedules>, + ValueQuery, + >; #[pallet::genesis_config] pub struct GenesisConfig { @@ -189,21 +200,23 @@ pub mod module { .for_each(|(who, start, period, period_count, per_period)| { let total = *per_period * Into::>::into(*period_count); - assert!( - T::Currency::free_balance(who) >= total, - "Account do not have enough balance" - ); - - T::Currency::set_lock(VESTING_LOCK_ID, who, total, WithdrawReasons::all()); - VestingSchedules::::insert( - who, + let bounded_schedule: BoundedVec, T::MaxVestingSchedules> = vec![VestingSchedule { start: *start, period: *period, period_count: *period_count, per_period: *per_period, - }], + }] + .try_into() + .expect("Max vesting schedules exceeded"); + + assert!( + T::Currency::free_balance(who) >= total, + "Account do not have enough balance" ); + + T::Currency::set_lock(VESTING_LOCK_ID, who, total, WithdrawReasons::all()); + VestingSchedules::::insert(who, bounded_schedule); }); } } @@ -216,8 +229,7 @@ pub mod module { #[pallet::call] impl Pallet { - // can not get VestingSchedule count from `who`, so use `MAX_VESTINGS / 2` - #[pallet::weight(T::WeightInfo::claim((MAX_VESTINGS / 2) as u32))] + #[pallet::weight(T::WeightInfo::claim((::MaxVestingSchedules::get() / 2) as u32))] pub fn claim(origin: OriginFor) -> DispatchResultWithPostInfo { let who = ensure_signed(origin)?; let locked_amount = Self::do_claim(&who); @@ -294,36 +306,34 @@ impl Pallet { fn do_vested_transfer(from: &T::AccountId, to: &T::AccountId, schedule: VestingScheduleOf) -> DispatchResult { let schedule_amount = Self::ensure_valid_vesting_schedule(&schedule)?; - ensure!( - >::decode_len(to).unwrap_or(0) < MAX_VESTINGS, - Error::::TooManyVestingSchedules - ); - let total_amount = Self::locked_balance(to) .checked_add(&schedule_amount) .ok_or(ArithmeticError::Overflow)?; T::Currency::transfer(from, to, schedule_amount, ExistenceRequirement::AllowDeath)?; T::Currency::set_lock(VESTING_LOCK_ID, to, total_amount, WithdrawReasons::all()); - >::append(to, schedule); + >::try_append(to, schedule).map_err(|_| Error::::MaxVestingSchedulesExceeded)?; Ok(()) } fn do_update_vesting_schedules(who: &T::AccountId, schedules: Vec>) -> DispatchResult { - let total_amount = schedules.iter().try_fold::<_, _, Result, DispatchError>>( - Zero::zero(), - |acc_amount, schedule| { + let bounded_schedules: BoundedVec, T::MaxVestingSchedules> = schedules + .try_into() + .map_err(|_| Error::::MaxVestingSchedulesExceeded)?; + + let total_amount = bounded_schedules + .iter() + .try_fold::<_, _, Result, DispatchError>>(Zero::zero(), |acc_amount, schedule| { let amount = Self::ensure_valid_vesting_schedule(schedule)?; Ok(acc_amount + amount) - }, - )?; + })?; ensure!( T::Currency::free_balance(who) >= total_amount, Error::::InsufficientBalanceToLock, ); T::Currency::set_lock(VESTING_LOCK_ID, who, total_amount, WithdrawReasons::all()); - >::insert(who, schedules); + >::insert(who, bounded_schedules); Ok(()) } diff --git a/vesting/src/mock.rs b/vesting/src/mock.rs index e3b250191..c804d02b9 100644 --- a/vesting/src/mock.rs +++ b/vesting/src/mock.rs @@ -76,12 +76,17 @@ impl EnsureOrigin for EnsureAliceOrBob { } } +parameter_types! { + pub const MaxVestingSchedule: u32 = 2; +} + impl Config for Runtime { type Event = Event; type Currency = PalletBalances; type MinVestedTransfer = MinVestedTransfer; type VestedTransferOrigin = EnsureAliceOrBob; type WeightInfo = (); + type MaxVestingSchedules = MaxVestingSchedule; } type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic; diff --git a/vesting/src/tests.rs b/vesting/src/tests.rs index 79ffc8d46..35cdc049e 100644 --- a/vesting/src/tests.rs +++ b/vesting/src/tests.rs @@ -168,7 +168,7 @@ fn vested_transfer_fails_if_overflow() { start: 1u64, period: 1u64, period_count: 2u32, - per_period: u64::max_value(), + per_period: u64::MAX, }; assert_noop!( Vesting::vested_transfer(Origin::signed(ALICE), BOB, schedule), @@ -176,7 +176,7 @@ fn vested_transfer_fails_if_overflow() { ); let another_schedule = VestingSchedule { - start: u64::max_value(), + start: u64::MAX, period: 1u64, period_count: 2u32, per_period: 1u64, @@ -329,3 +329,28 @@ fn multiple_vesting_schedule_claim_works() { assert_eq!(PalletBalances::locks(&BOB), vec![]); }); } + +#[test] +fn exceeding_maximum_schedules_should_fail() { + ExtBuilder::build().execute_with(|| { + let schedule = VestingSchedule { + start: 0u64, + period: 10u64, + period_count: 2u32, + per_period: 10u64, + }; + assert_ok!(Vesting::vested_transfer(Origin::signed(ALICE), BOB, schedule.clone())); + assert_ok!(Vesting::vested_transfer(Origin::signed(ALICE), BOB, schedule.clone())); + assert_noop!( + Vesting::vested_transfer(Origin::signed(ALICE), BOB, schedule.clone()), + Error::::MaxVestingSchedulesExceeded + ); + + let schedules = vec![schedule.clone(), schedule.clone(), schedule.clone()]; + + assert_noop!( + Vesting::update_vesting_schedules(Origin::root(), BOB, schedules), + Error::::MaxVestingSchedulesExceeded + ); + }); +}