Enforce MaxCandidates limit for staking storage items

pallets/parachain-staking/src/lib.rs

@@ -80,7 +80,7 @@ pub mod pallet {
  use crate::delegation_requests::{
  CancelledScheduledRequest, DelegationAction, ScheduledRequest,
  };
- use crate::{set::OrderedSet, traits::*, types::*, InflationInfo, Range, WeightInfo};
+ use crate::{set::BoundedOrderedSet, traits::*, types::*, InflationInfo, Range, WeightInfo};
  use crate::{AutoCompoundConfig, AutoCompoundDelegations};
  use frame_support::pallet_prelude::*;
  use frame_support::traits::{
@@ -177,6 +177,9 @@ pub mod pallet {
  type OnNewRound: OnNewRound;
  /// Weight information for extrinsics in this pallet.
  type WeightInfo: WeightInfo;
+ /// Maximum candidates
+ #[pallet::constant]
+ type MaxCandidates: Get<u32>;
  }
 
  #[pallet::error]
@@ -229,8 +232,8 @@ pub mod pallet {
  TooLowCandidateAutoCompoundingDelegationCountToLeaveCandidates,
  TooLowCandidateCountWeightHint,
  TooLowCandidateCountWeightHintGoOffline,
- TooLowCandidateCountWeightHintGoOnline,
- TooLowCandidateCountWeightHintCandidateBondMore,
+ CandidateLimitReached,
+ CannotSetAboveMaxCandidates,
  }
 
  #[pallet::event]
@@ -579,7 +582,8 @@ pub mod pallet {
  #[pallet::storage]
  #[pallet::getter(fn selected_candidates)]
  /// The collator candidates selected for the current round
- type SelectedCandidates<T: Config> = StorageValue<_, Vec<T::AccountId>, ValueQuery>;
+ type SelectedCandidates<T: Config> =
+ StorageValue<_, BoundedVec<T::AccountId, T::MaxCandidates>, ValueQuery>;
 
  #[pallet::storage]
  #[pallet::getter(fn total)]
@@ -589,8 +593,11 @@ pub mod pallet {
  #[pallet::storage]
  #[pallet::getter(fn candidate_pool)]
  /// The pool of collator candidates, each with their total backing stake
- pub(crate) type CandidatePool<T: Config> =
- StorageValue<_, OrderedSet<Bond<T::AccountId, BalanceOf<T>>>, ValueQuery>;
+ pub(crate) type CandidatePool<T: Config> = StorageValue<
+ _,
+ BoundedOrderedSet<Bond<T::AccountId, BalanceOf<T>>, T::MaxCandidates>,
+ ValueQuery,
+ >;
 
  #[pallet::storage]
  #[pallet::getter(fn at_stake)]
@@ -764,6 +771,11 @@ pub mod pallet {
  "{:?}",
  Error::<T>::CannotSetBelowMin
  );
+ assert!(
+ self.num_selected_candidates <= T::MaxCandidates::get(),
+ "{:?}",
+ Error::<T>::CannotSetAboveMaxCandidates
+ );
  <TotalSelected<T>>::put(self.num_selected_candidates);
  // Choose top TotalSelected collator candidates
  let (_, v_count, _, total_staked) = <Pallet<T>>::select_top_candidates(1u32);
@@ -886,6 +898,10 @@ pub mod pallet {
  new >= T::MinSelectedCandidates::get(),
  Error::<T>::CannotSetBelowMin
  );
+ ensure!(
+ new <= T::MaxCandidates::get(),
+ Error::<T>::CannotSetAboveMaxCandidates
+ );
  let old = <TotalSelected<T>>::get();
  ensure!(old != new, Error::<T>::NoWritingSameValue);
  ensure!(
@@ -970,13 +986,14 @@ pub mod pallet {
  candidate_count >= old_count,
  Error::<T>::TooLowCandidateCountWeightHintJoinCandidates
  );
- ensure!(
- candidates.insert(Bond {
+ let maybe_inserted_candidate = candidates
+ .try_insert(Bond {
  owner: acc.clone(),
- amount: bond
- }),
- Error::<T>::CandidateExists
- );
+ amount: bond,
+ })
+ .map_err(|_| Error::<T>::CandidateLimitReached)?;
+ ensure!(maybe_inserted_candidate, Error::<T>::CandidateExists);
+
  ensure!(
  Self::get_collator_stakable_free_balance(&acc) >= bond,
  Error::<T>::InsufficientBalance,
@@ -1065,13 +1082,13 @@ pub mod pallet {
  candidates.0.len() as u32 <= candidate_count,
  Error::<T>::TooLowCandidateCountWeightHintCancelLeaveCandidates
  );
- ensure!(
- candidates.insert(Bond {
+ let maybe_inserted_candidate = candidates
+ .try_insert(Bond {
  owner: collator.clone(),
- amount: state.total_counted
- }),
- Error::<T>::AlreadyActive
- );
+ amount: state.total_counted,
+ })
+ .map_err(|_| Error::<T>::CandidateLimitReached)?;
+ ensure!(maybe_inserted_candidate, Error::<T>::AlreadyActive);
  <CandidatePool<T>>::put(candidates);
  <CandidateInfo<T>>::insert(&collator, state);
  Self::deposit_event(Event::CancelledCandidateExit {
@@ -1447,16 +1464,20 @@ pub mod pallet {
  );
  state.go_online();
 
- ensure!(
- candidates.insert(Bond {
+ let maybe_inserted_candidate = candidates
+ .try_insert(Bond {
  owner: collator.clone(),
- amount: state.total_counted
- }),
+ amount: state.total_counted,
+ })
+ .map_err(|_| Error::<T>::CandidateLimitReached)?;
+ ensure!(
+ maybe_inserted_candidate,
  DispatchErrorWithPostInfo {
  post_info: Some(actual_weight).into(),
  error: <Error<T>>::AlreadyActive.into(),
  },
  );
+
  <CandidatePool<T>>::put(candidates);
  <CandidateInfo<T>>::insert(&collator, state);
  Self::deposit_event(Event::CandidateBackOnline {
@@ -1619,10 +1640,14 @@ pub mod pallet {
  pub(crate) fn update_active(candidate: T::AccountId, total: BalanceOf<T>) {
  let mut candidates = <CandidatePool<T>>::get();
  candidates.remove(&Bond::from_owner(candidate.clone()));
- candidates.insert(Bond {
- owner: candidate,
- amount: total,
- });
+ candidates
+ .try_insert(Bond {
+ owner: candidate,
+ amount: total,
+ })
+ .expect(
+ "the candidate is removed in previous step so the length cannot increase; qed",
+ );
  <CandidatePool<T>>::put(candidates);
  }
 
@@ -1853,27 +1878,31 @@ pub mod pallet {
  return vec![];
  }
 
- let mut candidates = <CandidatePool<T>>::get().0;
+ let candidates = <CandidatePool<T>>::get().0;
 
  // If the number of candidates is greater than top_n, select the candidates with higher
  // amount. Otherwise, return all the candidates.
  if candidates.len() > top_n {
  // Partially sort candidates such that element at index `top_n - 1` is sorted, and
  // all the elements in the range 0..top_n are the top n elements.
- candidates.select_nth_unstable_by(top_n - 1, |a, b| {
- // Order by amount, then owner. The owner is needed to ensure a stable order
- // when two accounts have the same amount.
- a.amount
- .cmp(&b.amount)
- .then_with(|| a.owner.cmp(&b.owner))
- .reverse()
- });
+ let sorted_candidates = candidates
+ .try_mutate(|inner| {
+ inner.select_nth_unstable_by(top_n - 1, |a, b| {
+ // Order by amount, then owner. The owner is needed to ensure a stable order
+ // when two accounts have the same amount.
+ a.amount
+ .cmp(&b.amount)
+ .then_with(|| a.owner.cmp(&b.owner))
+ .reverse()
+ });
+ })
+ .expect("sort cannot increase item count; qed");
 
- let mut collators = candidates
+ let mut collators = sorted_candidates
  .into_iter()
  .take(top_n)
  .map(|x| x.owner)
- .collect::<Vec<T::AccountId>>();
+ .collect::<Vec<_>>();
 
  // Sort collators by AccountId
  collators.sort();
@@ -1882,10 +1911,7 @@ pub mod pallet {
  } else {
  // Return all candidates
  // The candidates are already sorted by AccountId, so no need to sort again
- candidates
- .into_iter()
- .map(|x| x.owner)
- .collect::<Vec<T::AccountId>>()
+ candidates.into_iter().map(|x| x.owner).collect::<Vec<_>>()
  }
  }
  /// Best as in most cumulatively supported in terms of stake
@@ -1967,7 +1993,10 @@ pub mod pallet {
  });
  }
  // insert canonical collator set
- <SelectedCandidates<T>>::put(collators);
+ <SelectedCandidates<T>>::put(
+ BoundedVec::try_from(collators)
+ .expect("subset of collators is always less than or equal to max candidates"),
+ );
 
  let avg_delegator_count = delegation_count.checked_div(collator_count).unwrap_or(0);
  let weight = T::WeightInfo::select_top_candidates(collator_count, avg_delegator_count);
@@ -2149,7 +2178,7 @@ pub mod pallet {
 
  impl<T: Config> Get<Vec<T::AccountId>> for Pallet<T> {
  fn get() -> Vec<T::AccountId> {
- Self::selected_candidates()
+ Self::selected_candidates().into_inner()
  }
  }
 }

pallets/parachain-staking/src/mock.rs

@@ -120,6 +120,7 @@ parameter_types! {
  pub const MinCandidateStk: u128 = 10;
  pub const MinDelegatorStk: u128 = 5;
  pub const MinDelegation: u128 = 3;
+ pub const MaxCandidates: u32 = 10;
 }
 impl Config for Test {
  type RuntimeEvent = RuntimeEvent;
@@ -144,6 +145,7 @@ impl Config for Test {
  type PayoutCollatorReward = ();
  type OnNewRound = ();
  type WeightInfo = ();
+ type MaxCandidates = MaxCandidates;
 }
 
 pub(crate) struct ExtBuilder {

pallets/parachain-staking/src/set.rs

@@ -15,11 +15,13 @@
 // along with Moonbeam. If not, see <http://www.gnu.org/licenses/>.
 
 /* TODO: use orml_utilities::OrderedSet without leaking substrate v2.0 dependencies*/
+
+use frame_support::traits::Get;
 use parity_scale_codec::{Decode, Encode};
 use scale_info::TypeInfo;
 #[cfg(feature = "std")]
 use serde::{Deserialize, Serialize};
-use sp_runtime::RuntimeDebug;
+use sp_runtime::{BoundedVec, RuntimeDebug};
 use sp_std::prelude::*;
 
 /// An ordered set backed by `Vec`
@@ -87,3 +89,78 @@ impl<T: Ord> From<Vec<T>> for OrderedSet<T> {
  Self::from(v)
  }
 }
+
+/// An ordered set backed by `BoundedVec`
+#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
+#[derive(RuntimeDebug, PartialEq, Eq, Encode, Decode, Clone, TypeInfo)]
+#[scale_info(skip_type_params(S))]
+pub struct BoundedOrderedSet<T, S: Get<u32>>(pub BoundedVec<T, S>);
+
+impl<T, S: Get<u32>> sp_std::default::Default for BoundedOrderedSet<T, S> {
+ fn default() -> Self {
+ BoundedOrderedSet(BoundedVec::default())
+ }
+}
+
+impl<T: Ord, S: Get<u32>> BoundedOrderedSet<T, S> {
+ /// Create a new empty set
+ pub fn new() -> Self {
+ Self(BoundedVec::default())
+ }
+
+ /// Create a set from a `Vec`.
+ /// `v` will be sorted and dedup first.
+ pub fn from(mut v: BoundedVec<T, S>) -> Self {
+ v.sort();
+ let v = v.try_mutate(|inner| inner.dedup()).expect(
+ "input is a valid BoundedVec and deduping can only reduce the number of entires; qed",
+ );
+ Self::from_sorted_set(v)
+ }
+
+ /// Create a set from a `Vec`.
+ /// Assume `v` is sorted and contain unique elements.
+ pub fn from_sorted_set(v: BoundedVec<T, S>) -> Self {
+ Self(v)
+ }
+
+ /// Insert an element.
+ /// Return true if insertion happened.
+ pub fn try_insert(&mut self, value: T) -> Result<bool, ()> {
+ match self.0.binary_search(&value) {
+ Ok(_) => Ok(false),
+ Err(loc) => self.0.try_insert(loc, value).map(|_| true).map_err(|_| ()),
+ }
+ }
+
+ /// Remove an element.
+ /// Return true if removal happened.
+ pub fn remove(&mut self, value: &T) -> bool {
+ match self.0.binary_search(value) {
+ Ok(loc) => {
+ self.0.remove(loc);
+ true
+ }
+ Err(_) => false,
+ }
+ }
+
+ /// Return if the set contains `value`
+ pub fn contains(&self, value: &T) -> bool {
+ self.0.binary_search(value).is_ok()
+ }
+
+ /// Clear the set
+ pub fn clear(mut self) {
+ let v = self.0.try_mutate(|inner| inner.clear()).expect(
+ "input is a valid BoundedVec and clearing can only reduce the number of entires; qed",
+ );
+ self.0 = v;
+ }
+}
+
+impl<T: Ord, S: Get<u32>> From<BoundedVec<T, S>> for BoundedOrderedSet<T, S> {
+ fn from(v: BoundedVec<T, S>) -> Self {
+ Self::from(v)
+ }
+}