Improving Identity staking rewards distribution

parachain/pallets/score-staking/Cargo.toml

@@ -13,12 +13,14 @@ serde = { workspace = true }
 frame-benchmarking = { workspace = true, optional = true }
 frame-support = { workspace = true }
 frame-system = { workspace = true }
+pallet-timestamp = { workspace = true }
 sp-core = { workspace = true }
 sp-runtime = { workspace = true }
 sp-std = { workspace = true }
 
 core-primitives = { workspace = true }
 pallet-parachain-staking = { workspace = true }
+pallet-teebag = { workspace = true }
 
 [dev-dependencies]
 sp-io = { workspace = true }
@@ -42,12 +44,15 @@ std = [
     "pallet-balances/std",
     "core-primitives/std",
     "pallet-parachain-staking/std",
+    "pallet-teebag/std",
+    "pallet-timestamp/std",
 ]
 runtime-benchmarks = [
     "frame-benchmarking/runtime-benchmarks",
     "frame-support/runtime-benchmarks",
     "frame-system/runtime-benchmarks",
     "sp-runtime/runtime-benchmarks",
+    "pallet-timestamp/runtime-benchmarks",
 ]
 try-runtime = [
     "frame-support/try-runtime",

parachain/pallets/score-staking/src/lib.rs

@@ -47,7 +47,10 @@ use frame_support::{
 };
 use pallet_parachain_staking as ParaStaking;
 use sp_core::crypto::AccountId32;
-use sp_runtime::{traits::CheckedSub, Perbill, SaturatedConversion};
+use sp_runtime::{
+	traits::{CheckedSub, Zero},
+	Perbill, SaturatedConversion,
+};
 
 pub use pallet::*;
 
@@ -74,7 +77,6 @@ pub mod pallet {
 
 	use core_primitives::Identity;
 	use frame_system::pallet_prelude::*;
-	use sp_runtime::traits::Zero;
 
 	/// The current storage version.
 	const STORAGE_VERSION: StorageVersion = StorageVersion::new(0);
@@ -89,6 +91,7 @@ pub mod pallet {
 			if let Some(mut s) = Scores::<T>::get(delegator) {
 				let _ = Self::update_total_score(s.score, 0);
 				s.score = 0;
+				s.total_staking_amount = BalanceOf::<T>::zero();
 				Scores::<T>::insert(delegator, s);
 			}
 
@@ -114,6 +117,8 @@ pub mod pallet {
 		type AdminOrigin: EnsureOrigin<Self::RuntimeOrigin>;
 		/// AccountId converter
 		type AccountIdConvert: AccountIdConvert<Self>;
+		// For extrinsics that should only be called by origins from TEE
+		type TEECallOrigin: EnsureOrigin<Self::RuntimeOrigin>;
 	}
 
 	#[pallet::error]
@@ -150,6 +155,10 @@ pub mod pallet {
 		ScoreUserCountUnderflow,
 		// when the score user count would exceed `MaxScoreUserCount`
 		MaxScoreUserCountReached,
+		// when the token staking amount has been updated already for the round
+		TotalStakingAmountAlreadyUpdated,
+		// the token staking amount has been updated already for the round
+		RoundRewardsAlreadyDistributed,
 	}
 
 	#[pallet::event]
@@ -162,7 +171,9 @@ pub mod pallet {
 		ScoreUpdated { who: Identity, new_score: Score },
 		ScoreRemoved { who: Identity },
 		ScoreCleared {},
-		RewardCalculated { total: BalanceOf<T>, distributed: BalanceOf<T> },
+		RewardDistributionStarted { round_index: RoundIndex },
+		TotalStakingAmountUpdated { account_id: T::AccountId, amount: BalanceOf<T> },
+		RewardDistributionCompleted { round_index: RoundIndex },
 		RewardClaimed { who: T::AccountId, amount: BalanceOf<T> },
 	}
 
@@ -204,6 +215,11 @@ pub mod pallet {
 	#[pallet::getter(fn state)]
 	pub type State<T: Config> = StorageValue<_, PoolState, ValueQuery>;
 
+	/// The round index of the last token distribution
+	#[pallet::storage]
+	#[pallet::getter(fn last_rewards_distribution_round)]
+	pub type LastTokenDistributionRound<T: Config> = StorageValue<_, RoundIndex, ValueQuery>;
+
 	#[pallet::genesis_config]
 	pub struct GenesisConfig<T: Config> {
 		pub state: PoolState,
@@ -241,48 +257,14 @@ pub mod pallet {
 			}
 
 			// We are about to start a new round
-			// 1. update round info
-			r.index = r.index.saturating_add(1);
+			// update round info
+			let round_index = r.index.saturating_add(1);
+			r.index = round_index;
 			r.start_block = now;
 			Round::<T>::put(r);
 			weight = weight.saturating_add(T::DbWeight::get().reads_writes(1, 1));
 
-			// 2. calculate payout
-			let round_reward: BalanceOf<T> = (T::YearlyInflation::get() * T::YearlyIssuance::get()
-				/ YEARS.into()) * Self::round_config().interval.into();
-			let round_reward_u128 = round_reward.saturated_into::<u128>();
-
-			let total_stake_u128 = ParaStaking::Pallet::<T>::total().saturated_into::<u128>();
-			let total_score = Self::total_score();
-			let n = Self::round_config().stake_coef_n;
-			let m = Self::round_config().stake_coef_m;
-
-			let mut all_user_reward = BalanceOf::<T>::zero();
-
-			for (a, mut p) in Scores::<T>::iter() {
-				let user_stake_u128 = ParaStaking::Pallet::<T>::delegator_state(&a)
-					.map(|s| s.total)
-					.unwrap_or_default()
-					.saturated_into::<u128>();
-				let user_reward_u128 = round_reward_u128
-					.saturating_mul(p.score.into())
-					.saturating_div(total_score.into())
-					.saturating_mul(num_integer::Roots::nth_root(&user_stake_u128.pow(n), m))
-					.saturating_div(num_integer::Roots::nth_root(&total_stake_u128.pow(n), m));
-				let user_reward = user_reward_u128.saturated_into::<BalanceOf<T>>();
-
-				p.last_round_reward = user_reward;
-				p.total_reward += user_reward;
-				p.unpaid_reward += user_reward;
-				all_user_reward += user_reward;
-				Scores::<T>::insert(&a, p);
-				weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1));
-			}
-
-			Self::deposit_event(Event::<T>::RewardCalculated {
-				total: round_reward,
-				distributed: all_user_reward,
-			});
+			Self::deposit_event(Event::<T>::RewardDistributionStarted { round_index });
 
 			weight
 		}
@@ -445,6 +427,70 @@ pub mod pallet {
 			let payment = Scores::<T>::get(&account).ok_or(Error::<T>::UserNotExist)?;
 			Self::claim(origin, payment.unpaid_reward)
 		}
+
+		#[pallet::call_index(9)]
+		#[pallet::weight((195_000_000, DispatchClass::Normal))]
+		pub fn update_total_staking_amount(
+			origin: OriginFor<T>,
+			account_id: T::AccountId,
+			amount: BalanceOf<T>,
+		) -> DispatchResultWithPostInfo {
+			let _ = T::TEECallOrigin::ensure_origin(origin)?;
+
+			match Scores::<T>::get(&account_id) {
+				Some(mut s) => {
+					s.total_staking_amount = amount;
+					Scores::<T>::insert(account_id.clone(), s);
+					Self::deposit_event(Event::TotalStakingAmountUpdated {
+						account_id: account_id.clone(),
+						amount,
+					});
+					Ok(Pays::No.into())
+				},
+				None => Err(Error::<T>::UserNotExist.into()),
+			}
+		}
+
+		#[pallet::call_index(10)]
+		#[pallet::weight((195_000_000, DispatchClass::Normal))]
+		pub fn distribute_rewards(
+			origin: OriginFor<T>,
+			round_index: RoundIndex,
+		) -> DispatchResultWithPostInfo {
+			let _ = T::TEECallOrigin::ensure_origin(origin)?;
+			ensure!(
+				round_index > LastTokenDistributionRound::<T>::get(),
+				Error::<T>::RoundRewardsAlreadyDistributed
+			);
+			let round_reward: BalanceOf<T> = (T::YearlyInflation::get() * T::YearlyIssuance::get()
+				/ YEARS.into()) * Self::round_config().interval.into();
+			let round_reward_u128 = round_reward.saturated_into::<u128>();
+			let total_stake_u128 = ParaStaking::Pallet::<T>::total().saturated_into::<u128>();
+			let total_score = Self::total_score();
+			let n = Self::round_config().stake_coef_n;
+			let m = Self::round_config().stake_coef_m;
+
+			for (a, mut p) in Scores::<T>::iter() {
+				let user_stake_u128 = p.total_staking_amount.saturated_into::<u128>();
+				let user_reward_u128 = round_reward_u128
+					.saturating_mul(p.score.into())
+					.saturating_div(total_score.into())
+					.saturating_mul(num_integer::Roots::nth_root(&user_stake_u128.pow(n), m))
+					.saturating_div(num_integer::Roots::nth_root(&total_stake_u128.pow(n), m));
+				let user_reward = user_reward_u128.saturated_into::<BalanceOf<T>>();
+
+				p.last_round_reward = user_reward;
+				p.total_reward += user_reward;
+				p.unpaid_reward += user_reward;
+				Scores::<T>::insert(&a, p);
+			}
+
+			LastTokenDistributionRound::<T>::put(round_index);
+
+			Self::deposit_event(Event::RewardDistributionCompleted { round_index });
+
+			Ok(Pays::No.into())
+		}
 	}
 }
 

parachain/pallets/score-staking/src/mock.rs

@@ -18,12 +18,15 @@
 use crate::{
 	self as pallet_score_staking, AccountIdConvert, Config, Perbill, PoolState, RoundSetting,
 };
+use core::marker::PhantomData;
 use frame_support::{
-	assert_ok, construct_runtime, ord_parameter_types, parameter_types,
+	assert_ok, construct_runtime, ord_parameter_types,
+	pallet_prelude::EnsureOrigin,
+	parameter_types,
 	traits::{OnFinalize, OnInitialize},
 };
 use frame_system::{EnsureRoot, EnsureSignedBy};
-use sp_core::{ConstU128, ConstU32, H256};
+use sp_core::{ConstU128, ConstU16, ConstU32, H256};
 use sp_keyring::AccountKeyring;
 use sp_runtime::{
 	generic,
@@ -51,6 +54,7 @@ construct_runtime!(
 		Balances: pallet_balances,
 		ParachainStaking: pallet_parachain_staking,
 		ScoreStaking: pallet_score_staking,
+		Teebag: pallet_teebag,
 	}
 );
 
@@ -166,6 +170,66 @@ impl pallet_score_staking::Config for Test {
 	type YearlyIssuance = ConstU128<{ 100_000_000 * UNIT }>;
 	type YearlyInflation = DefaultYearlyInflation;
 	type MaxScoreUserCount = ConstU32<2>;
+	type TEECallOrigin = EnsureEnclaveSigner<Self>;
+}
+
+parameter_types! {
+		pub const MinimumPeriod: u64 = 6000 / 2;
+}
+
+pub type Moment = u64;
+
+impl pallet_timestamp::Config for Test {
+	type Moment = Moment;
+	type OnTimestampSet = ();
+	type MinimumPeriod = MinimumPeriod;
+	type WeightInfo = ();
+}
+
+parameter_types! {
+	pub const MomentsPerDay: u64 = 86_400_000; // [ms/d]
+}
+
+impl pallet_teebag::Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type MomentsPerDay = MomentsPerDay;
+	type SetAdminOrigin = EnsureRoot<Self::AccountId>;
+	type MaxEnclaveIdentifier = ConstU32<1>;
+	type MaxAuthorizedEnclave = ConstU32<2>;
+	type WeightInfo = ();
+}
+
+pub struct EnsureEnclaveSigner<T>(PhantomData<T>);
+impl<T> EnsureOrigin<T::RuntimeOrigin> for EnsureEnclaveSigner<T>
+where
+	T: frame_system::Config + pallet_teebag::Config + pallet_timestamp::Config<Moment = u64>,
+	<T as frame_system::Config>::AccountId: From<[u8; 32]>,
+	<T as frame_system::Config>::Hash: From<[u8; 32]>,
+{
+	type Success = T::AccountId;
+	fn try_origin(o: T::RuntimeOrigin) -> Result<Self::Success, T::RuntimeOrigin> {
+		o.into().and_then(|o| match o {
+			frame_system::RawOrigin::Signed(who)
+				if pallet_teebag::EnclaveRegistry::<T>::contains_key(&who) =>
+			{
+				Ok(who)
+			},
+			r => Err(T::RuntimeOrigin::from(r)),
+		})
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn try_successful_origin() -> Result<T::RuntimeOrigin, ()> {
+		use pallet_teebag::test_util::{get_signer, TEST8_MRENCLAVE, TEST8_SIGNER_PUB};
+		let signer: <T as frame_system::Config>::AccountId = get_signer(TEST8_SIGNER_PUB);
+		if !pallet_teebag::EnclaveRegistry::<T>::contains_key(signer.clone()) {
+			assert_ok!(pallet_teebag::Pallet::<T>::add_enclave(
+				&signer,
+				&pallet_teebag::Enclave::default().with_mrenclave(TEST8_MRENCLAVE),
+			));
+		}
+		Ok(frame_system::RawOrigin::Signed(signer).into())
+	}
 }
 
 pub fn alice() -> AccountId {
@@ -186,6 +250,14 @@ pub fn new_test_ext(fast_round: bool) -> sp_io::TestExternalities {
 		.assimilate_storage(&mut t)
 		.unwrap();
 
+	pallet_teebag::GenesisConfig::<Test> {
+		allow_sgx_debug_mode: true,
+		admin: Some(AccountKeyring::Alice.to_account_id()),
+		mode: pallet_teebag::OperationalMode::Production,
+	}
+	.assimilate_storage(&mut t)
+	.unwrap();
+
 	pallet_score_staking::GenesisConfig::<Test> {
 		state: PoolState::Stopped,
 		marker: Default::default(),