staking-async: allow session keys handling on AssetHub

cumulus/parachains/runtimes/assets/asset-hub-westend/src/staking.rs

@@ -321,9 +321,15 @@ pub enum RelayChainRuntimePallets {
 
 #[derive(Encode, Decode)]
 pub enum AhClientCalls {
-	// index of `fn validator_set` in `staking-async-ah-client`. It has only one call.
+	// index of `fn validator_set` in `staking-async-ah-client`.
 	#[codec(index = 0)]
 	ValidatorSet(rc_client::ValidatorSetReport<AccountId>),
+	// index of `fn set_keys_from_ah` in `staking-async-ah-client`.
+	#[codec(index = 3)]
+	SetKeys { stash: AccountId, keys: Vec<u8>, proof: Vec<u8> },
+	// index of `fn purge_keys_from_ah` in `staking-async-ah-client`.
+	#[codec(index = 4)]
+	PurgeKeys { stash: AccountId },
 }
 
 pub struct ValidatorSetToXcm;
@@ -347,6 +353,64 @@ impl sp_runtime::traits::Convert<rc_client::ValidatorSetReport<AccountId>, Xcm<(
 	}
 }
 
+/// Message to set session keys on the Relay Chain.
+#[derive(Encode, Decode, Clone)]
+pub struct SetKeysMessage {
+	pub stash: AccountId,
+	pub keys: Vec<u8>,
+	pub proof: Vec<u8>,
+}
+
+pub struct SetKeysToXcm;
+impl sp_runtime::traits::Convert<SetKeysMessage, Xcm<()>> for SetKeysToXcm {
+	fn convert(msg: SetKeysMessage) -> Xcm<()> {
+		Xcm(vec![
+			Instruction::UnpaidExecution {
+				weight_limit: WeightLimit::Unlimited,
+				check_origin: None,
+			},
+			Instruction::Transact {
+				origin_kind: OriginKind::Native,
+				fallback_max_weight: None,
+				call: RelayChainRuntimePallets::AhClient(AhClientCalls::SetKeys {
+					stash: msg.stash,
+					keys: msg.keys,
+					proof: msg.proof,
+				})
+				.encode()
+				.into(),
+			},
+		])
+	}
+}
+
+/// Message to purge session keys on the Relay Chain.
+#[derive(Encode, Decode, Clone)]
+pub struct PurgeKeysMessage {
+	pub stash: AccountId,
+}
+
+pub struct PurgeKeysToXcm;
+impl sp_runtime::traits::Convert<PurgeKeysMessage, Xcm<()>> for PurgeKeysToXcm {
+	fn convert(msg: PurgeKeysMessage) -> Xcm<()> {
+		Xcm(vec![
+			Instruction::UnpaidExecution {
+				weight_limit: WeightLimit::Unlimited,
+				check_origin: None,
+			},
+			Instruction::Transact {
+				origin_kind: OriginKind::Native,
+				fallback_max_weight: None,
+				call: RelayChainRuntimePallets::AhClient(AhClientCalls::PurgeKeys {
+					stash: msg.stash,
+				})
+				.encode()
+				.into(),
+			},
+		])
+	}
+}
+
 parameter_types! {
 	pub RelayLocation: Location = Location::parent();
 }
@@ -363,6 +427,24 @@ impl rc_client::SendToRelayChain for StakingXcmToRelayChain {
 			ValidatorSetToXcm,
 		>::send(report)
 	}
+
+	fn set_keys(stash: Self::AccountId, keys: Vec<u8>, proof: Vec<u8>) -> Result<(), ()> {
+		rc_client::XCMSender::<
+			xcm_config::XcmRouter,
+			RelayLocation,
+			SetKeysMessage,
+			SetKeysToXcm,
+		>::send(SetKeysMessage { stash, keys, proof })
+	}
+
+	fn purge_keys(stash: Self::AccountId) -> Result<(), ()> {
+		rc_client::XCMSender::<
+			xcm_config::XcmRouter,
+			RelayLocation,
+			PurgeKeysMessage,
+			PurgeKeysToXcm,
+		>::send(PurgeKeysMessage { stash })
+	}
 }
 
 parameter_types! {

polkadot/runtime/westend/src/lib.rs

@@ -509,7 +509,6 @@ impl pallet_authorship::Config for Runtime {
 parameter_types! {
 	pub const Period: BlockNumber = 10 * MINUTES;
 	pub const Offset: BlockNumber = 0;
-	pub const KeyDeposit: Balance = deposit(1, 5 * 32 + 33);
 }
 
 impl_opaque_keys! {
@@ -535,7 +534,7 @@ impl pallet_session::Config for Runtime {
 	type DisablingStrategy = pallet_session::disabling::UpToLimitWithReEnablingDisablingStrategy;
 	type WeightInfo = weights::pallet_session::WeightInfo<Runtime>;
 	type Currency = Balances;
-	type KeyDeposit = KeyDeposit;
+	type KeyDeposit = ();
 }
 
 impl pallet_session::historical::Config for Runtime {

prdoc/pr_10666.prdoc

@@ -0,0 +1,20 @@
+title: 'staking-async: allow  session keys handling on AssetHub'
+doc:
+- audience: Runtime Dev
+  description: |-
+    - Added  session keys handling on AssetHub for staking-async: validators can now call set_keys and purge_keys on AssetHub, which forwards the request to the RelayChain via XCM.
+    - Validators are still allowed to call `set_keys` and `purge_keys` via relay-chain pallet-session's related extrinsics. This option will be deprecated in the future.
+    - No key deposit is set on AssetHub yet.
+crates:
+- name: asset-hub-westend-runtime
+  bump: major
+- name: westend-runtime
+  bump: major
+- name: pallet-session
+  bump: minor
+- name: pallet-staking-async-ah-client
+  bump: major
+- name: pallet-staking-async-rc-client
+  bump: major
+- name: pallet-staking-async
+  bump: major

substrate/frame/session/src/lib.rs

@@ -860,7 +860,10 @@ impl<T: Config> Pallet<T> {
 	///
 	/// This ensures that the reference counter in system is incremented appropriately and as such
 	/// must accept an account ID, rather than a validator ID.
-	fn do_set_keys(account: &T::AccountId, keys: T::Keys) -> DispatchResult {
+	///
+	/// This function is public to allow cross-chain session key management (e.g., from AssetHub
+	/// via `pallet-staking-async-ah-client`).
+	pub fn do_set_keys(account: &T::AccountId, keys: T::Keys) -> DispatchResult {
 		let who = T::ValidatorIdOf::convert(account.clone())
 			.ok_or(Error::<T>::NoAssociatedValidatorId)?;
 
@@ -923,7 +926,11 @@ impl<T: Config> Pallet<T> {
 		Ok(old_keys)
 	}
 
-	fn do_purge_keys(account: &T::AccountId) -> DispatchResult {
+	/// Purge session keys for an account.
+	///
+	/// This function is public to allow cross-chain session key management (e.g., from AssetHub
+	/// via `pallet-staking-async-ah-client`).
+	pub fn do_purge_keys(account: &T::AccountId) -> DispatchResult {
 		let who = T::ValidatorIdOf::convert(account.clone())
 			// `purge_keys` may not have a controller-stash pair any more. If so then we expect the
 			// stash account to be passed in directly and convert that to a `ValidatorId` using the

substrate/frame/staking-async/ah-client/src/lib.rs

@@ -65,6 +65,7 @@ use frame_support::{
 	pallet_prelude::*,
 	traits::{Defensive, DefensiveSaturating, RewardsReporter},
 };
+use pallet_session::WeightInfo as SessionWeightInfo;
 pub use pallet_staking_async_rc_client::SendToAssetHub;
 use pallet_staking_async_rc_client::{self as rc_client};
 use sp_runtime::SaturatedConversion;
@@ -103,6 +104,12 @@ pub trait SessionInterface {
 	/// The validator id type of the session pallet
 	type ValidatorId: Clone;
 
+	/// The account id type
+	type AccountId;
+
+	/// The session keys type
+	type Keys: sp_runtime::traits::OpaqueKeys + codec::Decode;
+
 	fn validators() -> Vec<Self::ValidatorId>;
 
 	/// prune up to the given session index.
@@ -112,12 +119,30 @@ pub trait SessionInterface {
 	///
 	/// This is used to disable validators directly on the RC, until the next validator set.
 	fn report_offence(offender: Self::ValidatorId, severity: OffenceSeverity);
+
+	/// Set session keys for an account.
+	///
+	/// This is called when AssetHub forwards a session key registration via XCM.
+	fn set_keys(account: &Self::AccountId, keys: Self::Keys) -> DispatchResult;
+
+	/// Purge session keys for an account.
+	///
+	/// This is called when AssetHub forwards a session key purge request via XCM.
+	fn purge_keys(account: &Self::AccountId) -> DispatchResult;
+
+	/// Weight for setting session keys.
+	fn set_keys_weight() -> Weight;
+
+	/// Weight for purging session keys.
+	fn purge_keys_weight() -> Weight;
 }
 
 impl<T: Config + pallet_session::Config + pallet_session::historical::Config> SessionInterface
 	for T
 {
 	type ValidatorId = <T as pallet_session::Config>::ValidatorId;
+	type AccountId = <T as frame_system::Config>::AccountId;
+	type Keys = <T as pallet_session::Config>::Keys;
 
 	fn validators() -> Vec<Self::ValidatorId> {
 		pallet_session::Pallet::<T>::validators()
@@ -129,6 +154,22 @@ impl<T: Config + pallet_session::Config + pallet_session::historical::Config> Se
 	fn report_offence(offender: Self::ValidatorId, severity: OffenceSeverity) {
 		pallet_session::Pallet::<T>::report_offence(offender, severity)
 	}
+
+	fn set_keys(account: &Self::AccountId, keys: Self::Keys) -> DispatchResult {
+		pallet_session::Pallet::<T>::do_set_keys(account, keys)
+	}
+
+	fn purge_keys(account: &Self::AccountId) -> DispatchResult {
+		pallet_session::Pallet::<T>::do_purge_keys(account)
+	}
+
+	fn set_keys_weight() -> Weight {
+		<T as pallet_session::Config>::WeightInfo::set_keys()
+	}
+
+	fn purge_keys_weight() -> Weight {
+		<T as pallet_session::Config>::WeightInfo::purge_keys()
+	}
 }
 
 /// Represents the operating mode of the pallet.
@@ -208,7 +249,7 @@ pub mod pallet {
 	use frame_system::pallet_prelude::*;
 	use pallet_session::{historical, SessionManager};
 	use pallet_staking_async_rc_client::SessionReport;
-	use sp_runtime::{Perbill, Saturating};
+	use sp_runtime::{traits::OpaqueKeys, Perbill, Saturating};
 	use sp_staking::{
 		offence::{OffenceSeverity, OnOffenceHandler},
 		SessionIndex,
@@ -275,7 +316,10 @@ pub mod pallet {
 		type MaxOffenceBatchSize: Get<u32>;
 
 		/// Interface to talk to the local Session pallet.
-		type SessionInterface: SessionInterface<ValidatorId = Self::AccountId>;
+		type SessionInterface: SessionInterface<
+			ValidatorId = Self::AccountId,
+			AccountId = Self::AccountId,
+		>;
 
 		/// A fallback implementation to delegate logic to when the pallet is in
 		/// [`OperatingMode::Passive`].
@@ -463,6 +507,10 @@ pub mod pallet {
 	pub enum Error<T> {
 		/// Could not process incoming message because incoming messages are blocked.
 		Blocked,
+		/// The session keys could not be decoded.
+		InvalidKeys,
+		/// Invalid ownership proof for the session keys.
+		InvalidProof,
 	}
 
 	#[pallet::event]
@@ -633,6 +681,55 @@ pub mod pallet {
 			Self::on_migration_end();
 			Ok(())
 		}
+
+		/// Set session keys for a validator, forwarded from AssetHub.
+		///
+		/// This is called when a validator sets their session keys on AssetHub, which forwards
+		/// the request to the RelayChain via XCM.
+		///
+		/// The `keys` parameter contains the encoded session keys that will be decoded and
+		/// registered with the session pallet.
+		/// The `proof` parameter is validated using the `OpaqueKeys::ownership_proof_is_valid`
+		/// method to verify key ownership.
+		#[pallet::call_index(3)]
+		#[pallet::weight(T::SessionInterface::set_keys_weight())]
+		pub fn set_keys_from_ah(
+			origin: OriginFor<T>,
+			stash: T::AccountId,
+			keys: Vec<u8>,
+			proof: Vec<u8>,
+		) -> DispatchResult {
+			T::AssetHubOrigin::ensure_origin_or_root(origin)?;
+			log::info!(target: LOG_TARGET, "Received set_keys request from AssetHub for {stash:?}");
+
+			// Decode the keys from bytes
+			let session_keys =
+				<<T as Config>::SessionInterface as SessionInterface>::Keys::decode(&mut &keys[..])
+					.map_err(|_| Error::<T>::InvalidKeys)?;
+
+			ensure!(session_keys.ownership_proof_is_valid(&proof), Error::<T>::InvalidProof);
+
+			// Forward to session pallet via SessionInterface
+			T::SessionInterface::set_keys(&stash, session_keys)?;
+
+			Ok(())
+		}
+
+		/// Purge session keys for a validator, forwarded from AssetHub.
+		///
+		/// This is called when a validator purges their session keys on AssetHub, which forwards
+		/// the request to the RelayChain via XCM.
+		#[pallet::call_index(4)]
+		#[pallet::weight(T::SessionInterface::purge_keys_weight())]
+		pub fn purge_keys_from_ah(origin: OriginFor<T>, stash: T::AccountId) -> DispatchResult {
+			T::AssetHubOrigin::ensure_origin_or_root(origin)?;
+			log::info!(target: LOG_TARGET, "Received purge_keys request from AssetHub for {stash:?}");
+
+			// Forward to session pallet via SessionInterface
+			T::SessionInterface::purge_keys(&stash)?;
+
+			Ok(())
+		}
 	}
 
 	#[pallet::hooks]

substrate/frame/staking-async/ah-client/src/mock.rs

@@ -19,9 +19,29 @@
 
 use crate::*;
 use frame_support::{derive_impl, parameter_types, weights::Weight};
-use sp_runtime::{BuildStorage, Perbill};
+use sp_runtime::{traits::OpaqueKeys, BuildStorage, KeyTypeId, Perbill};
 use sp_staking::offence::{OffenceSeverity, OnOffenceHandler};
 
+/// Mock session keys for testing.
+#[derive(Clone, PartialEq, Eq, Debug, codec::Encode, codec::Decode, scale_info::TypeInfo)]
+pub struct MockSessionKeys {
+	pub dummy: [u8; 32],
+}
+
+const MOCK_KEY_TYPE: KeyTypeId = KeyTypeId(*b"mock");
+
+impl OpaqueKeys for MockSessionKeys {
+	type KeyTypeIdProviders = ();
+
+	fn key_ids() -> &'static [KeyTypeId] {
+		&[MOCK_KEY_TYPE]
+	}
+
+	fn get_raw(&self, _: KeyTypeId) -> &[u8] {
+		&self.dummy
+	}
+}
+
 type Block = frame_system::mocking::MockBlock<Test>;
 
 frame_support::construct_runtime!(
@@ -41,11 +61,26 @@ impl frame_system::Config for Test {
 pub struct MockSessionInterface;
 impl SessionInterface for MockSessionInterface {
 	type ValidatorId = u64;
+	type AccountId = u64;
+	type Keys = MockSessionKeys;
+
 	fn validators() -> Vec<Self::ValidatorId> {
 		vec![1, 2, 3]
 	}
 	fn prune_up_to(_up_to: u32) {}
 	fn report_offence(_offender: Self::ValidatorId, _severity: OffenceSeverity) {}
+	fn set_keys(_account: &Self::AccountId, _keys: Self::Keys) -> DispatchResult {
+		Ok(())
+	}
+	fn purge_keys(_account: &Self::AccountId) -> DispatchResult {
+		Ok(())
+	}
+	fn set_keys_weight() -> Weight {
+		Weight::zero()
+	}
+	fn purge_keys_weight() -> Weight {
+		Weight::zero()
+	}
 }
 
 pub struct MockFallback;