From 09de121eea7a194a29f469c8b61c0dcb86ab4023 Mon Sep 17 00:00:00 2001
From: salman01zp <pathansalman555@gmail.com>
Date: Tue, 14 Nov 2023 22:25:50 +0530
Subject: [PATCH] block chill and withdraw-unbond calls

---
 pallets/roles/src/lib.rs      | 52 +++++++++++++++++++++++++++++++++--
 pallets/roles/src/mock.rs     | 13 +++++++++
 standalone/runtime/src/lib.rs | 18 +++++++++++-
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/pallets/roles/src/lib.rs b/pallets/roles/src/lib.rs
index 10f4007c0..30327705b 100644
--- a/pallets/roles/src/lib.rs
+++ b/pallets/roles/src/lib.rs
@@ -120,7 +120,7 @@ pub mod pallet {
 	pub enum Error<T> {
 		/// Not a validator.
 		NotValidator,
-		/// Role has already been assigned to provided validator.
+		/// Validator has active role assigned
 		HasRoleAssigned,
 		/// No role assigned to provided validator.
 		NoRoleAssigned,
@@ -251,6 +251,30 @@ pub mod pallet {
 			Ok(())
 		}
 
+		/// Declare no desire to either validate or nominate.
+		///
+		/// If you have opted for any of the roles, please submit `clear_role` extrinsic to opt out
+		/// of all the services. Once your role is cleared, your request will be processed.
+		///
+		/// # Parameters
+		///
+		/// - `origin`: Origin of the transaction.
+		///
+		/// This function will return error if
+		/// - Account is not a validator account.
+		/// - Role is assigned to the validator.
+		#[pallet::weight({0})]
+		#[pallet::call_index(2)]
+		pub fn chill(origin: OriginFor<T>) -> DispatchResult {
+			let account = ensure_signed(origin.clone())?;
+			// Ensure no role is assigned to the account before chilling.
+			ensure!(!AccountRolesMapping::<T>::contains_key(&account), Error::<T>::HasRoleAssigned);
+
+			// chill
+			let _ = pallet_staking::Pallet::<T>::chill(origin);
+			Ok(())
+		}
+
 		/// Unbound funds from the stash account.
 		/// This will allow user to unbound and later withdraw funds.
 		/// If you have opted for any of the roles, please submit `clear_role` extrinsic to opt out
@@ -266,7 +290,7 @@ pub mod pallet {
 		/// - If there is any active role assigned to the user.
 		///  
 		#[pallet::weight({0})]
-		#[pallet::call_index(2)]
+		#[pallet::call_index(3)]
 		pub fn unbound_funds(
 			origin: OriginFor<T>,
 			#[pallet::compact] amount: BalanceOf<T>,
@@ -280,5 +304,29 @@ pub mod pallet {
 
 			Ok(())
 		}
+
+		/// Withdraw unbound funds after un-bonding period has passed.
+		///
+		/// # Parameters
+		///
+		/// - `origin`: Origin of the transaction.
+		///
+		/// This function will return error if
+		/// - If there is any active role assigned to the user.
+		#[pallet::weight({0})]
+		#[pallet::call_index(4)]
+		pub fn withdraw_unbonded(origin: OriginFor<T>) -> DispatchResult {
+			let stash_account = ensure_signed(origin.clone())?;
+			// Ensure no role is assigned to the account and is eligible to withdraw.
+			ensure!(
+				!AccountRolesMapping::<T>::contains_key(&stash_account),
+				Error::<T>::HasRoleAssigned
+			);
+
+			// Withdraw unbound funds.
+			let _ = pallet_staking::Pallet::<T>::withdraw_unbonded(origin, 0);
+
+			Ok(())
+		}
 	}
 }
diff --git a/pallets/roles/src/mock.rs b/pallets/roles/src/mock.rs
index c2d1fe877..ae4141095 100644
--- a/pallets/roles/src/mock.rs
+++ b/pallets/roles/src/mock.rs
@@ -97,6 +97,19 @@ impl Contains<RuntimeCall> for BaseFilter {
 			return false
 		}
 
+		// no chill call
+		if matches!(call, RuntimeCall::Staking(pallet_staking::Call::chill { .. })) {
+			return false
+		}
+
+		// no withdraw_unbonded call
+		let is_stake_withdraw_call =
+			matches!(call, RuntimeCall::Staking(pallet_staking::Call::withdraw_unbonded { .. }));
+
+		if is_stake_withdraw_call {
+			return false
+		}
+
 		true
 	}
 }
diff --git a/standalone/runtime/src/lib.rs b/standalone/runtime/src/lib.rs
index 27e9ff404..82616bf14 100644
--- a/standalone/runtime/src/lib.rs
+++ b/standalone/runtime/src/lib.rs
@@ -1097,7 +1097,8 @@ impl Contains<RuntimeCall> for BaseFilter {
 			// no paused call
 			return false
 		}
-
+		// Following staking pallet calls will be blocked and will be allowed to execute
+		// through role pallet.
 		let is_stake_unbound_call =
 			matches!(call, RuntimeCall::Staking(pallet_staking::Call::unbond { .. }));
 
@@ -1106,6 +1107,21 @@ impl Contains<RuntimeCall> for BaseFilter {
 			return false
 		}
 
+		// no chill call
+		if matches!(call, RuntimeCall::Staking(pallet_staking::Call::chill { .. })) {
+			return false
+		}
+
+		// no withdraw_unbonded call
+		let is_stake_withdraw_call =  matches!(
+			call,
+			RuntimeCall::Staking(pallet_staking::Call::withdraw_unbonded { .. })
+		);
+
+		if is_stake_withdraw_call{
+			return false
+		}
+
 		let democracy_related = matches!(
 			call,
 			// Filter democracy proposals creation