Identity staking follow up

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",

pallets/score-staking/src/lib.rs

@@ -68,6 +68,10 @@ pub trait AccountIdConvert<T: Config> {
 	fn convert(account: AccountId32) -> T::AccountId;
 }
 
+pub trait TokenStakingAuthorizer<T: Config> {
+	fn can_update_staking(sender: &T::AccountId) -> bool;
+}
+
 #[frame_support::pallet]
 pub mod pallet {
 	use super::*;
@@ -114,6 +118,8 @@ pub mod pallet {
 		type AdminOrigin: EnsureOrigin<Self::RuntimeOrigin>;
 		/// AccountId converter
 		type AccountIdConvert: AccountIdConvert<Self>;
+		/// Token staking authorizer
+		type TokenStakingAuthorizer: TokenStakingAuthorizer<Self>;
 	}
 
 	#[pallet::error]
@@ -150,20 +156,45 @@ pub mod pallet {
 		ScoreUserCountUnderflow,
 		// when the score user count would exceed `MaxScoreUserCount`
 		MaxScoreUserCountReached,
+		// the token staking amount has been updated already for the round
+		TokenStakingAmountAlreadyUpdated,
 	}
 
 	#[pallet::event]
 	#[pallet::generate_deposit(pub (crate) fn deposit_event)]
 	pub enum Event<T: Config> {
-		PoolStarted { start_block: BlockNumberFor<T> },
+		PoolStarted {
+			start_block: BlockNumberFor<T>,
+		},
 		PoolStopped {},
-		ScoreFeederSet { new_score_feeder: Option<T::AccountId> },
-		RoundConfigSet { new_config: RoundSetting },
-		ScoreUpdated { who: Identity, new_score: Score },
-		ScoreRemoved { who: Identity },
+		ScoreFeederSet {
+			new_score_feeder: Option<T::AccountId>,
+		},
+		RoundConfigSet {
+			new_config: RoundSetting,
+		},
+		ScoreUpdated {
+			who: Identity,
+			new_score: Score,
+		},
+		ScoreRemoved {
+			who: Identity,
+		},
 		ScoreCleared {},
-		RewardCalculated { total: BalanceOf<T>, distributed: BalanceOf<T> },
-		RewardClaimed { who: T::AccountId, amount: BalanceOf<T> },
+		TokenStakingAmountUpdated {
+			account: T::AccountId,
+			amount: BalanceOf<T>,
+		},
+		RewardDistributionStarted {
+			total: BalanceOf<T>,
+			distributed: BalanceOf<T>,
+			round_index: RoundIndex,
+		},
+		RewardDistributionCompleted {},
+		RewardClaimed {
+			who: T::AccountId,
+			amount: BalanceOf<T>,
+		},
 	}
 
 	#[pallet::storage]
@@ -242,7 +273,8 @@ pub mod pallet {
 
 			// We are about to start a new round
 			// 1. update round info
-			r.index = r.index.saturating_add(1);
+			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));
@@ -279,9 +311,10 @@ pub mod pallet {
 				weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1));
 			}
 
-			Self::deposit_event(Event::<T>::RewardCalculated {
+			Self::deposit_event(Event::<T>::RewardDistributionStarted {
 				total: round_reward,
 				distributed: all_user_reward,
+				round_index,
 			});
 
 			weight
@@ -445,6 +478,53 @@ 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_token_staking_amount(
+			origin: OriginFor<T>,
+			account: T::AccountId,
+			amount: BalanceOf<T>,
+			round_index: RoundIndex,
+		) -> DispatchResultWithPostInfo {
+			let sender = ensure_signed(origin.clone())?;
+			ensure!(
+				T::TokenStakingAuthorizer::can_update_staking(&sender),
+				Error::<T>::UnauthorizedOrigin
+			);
+
+			match Scores::<T>::get(&account) {
+				Some(mut s) => {
+					if round_index <= s.last_token_distributed_round {
+						return Err(Error::<T>::TokenStakingAmountAlreadyUpdated.into());
+					}
+					let last_round = Round::<T>::get();
+					s.token_staking_amount = amount;
+					s.last_token_distributed_round = last_round.index;
+					Scores::<T>::insert(account.clone(), s);
+					Self::deposit_event(Event::TokenStakingAmountUpdated {
+						account: account.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 complete_reward_distribution(origin: OriginFor<T>) -> DispatchResultWithPostInfo {
+			let sender = ensure_signed(origin.clone())?;
+			ensure!(
+				T::TokenStakingAuthorizer::can_update_staking(&sender),
+				Error::<T>::UnauthorizedOrigin
+			);
+
+			Self::deposit_event(Event::RewardDistributionCompleted {});
+
+			Ok(Pays::No.into())
+		}
 	}
 }
 

pallets/score-staking/src/mock.rs

@@ -51,6 +51,7 @@ construct_runtime!(
 		Balances: pallet_balances,
 		ParachainStaking: pallet_parachain_staking,
 		ScoreStaking: pallet_score_staking,
+		Teebag: pallet_teebag,
 	}
 );
 
@@ -166,6 +167,39 @@ impl pallet_score_staking::Config for Test {
 	type YearlyIssuance = ConstU128<{ 100_000_000 * UNIT }>;
 	type YearlyInflation = DefaultYearlyInflation;
 	type MaxScoreUserCount = ConstU32<2>;
+	type TokenStakingAuthorizer = pallet_teebag::Pallet<Test>;
+}
+
+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 = ();
+}
+
+impl pallet_score_staking::TokenStakingAuthorizer<Test> for pallet_teebag::Pallet<Test> {
+	fn can_update_staking(sender: &<Test as frame_system::Config>::AccountId) -> bool {
+		pallet_teebag::Pallet::<Test>::enclave_registry(sender).is_some()
+	}
 }
 
 pub fn alice() -> AccountId {
@@ -186,6 +220,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(),