Skalman set keys proof with rt proof of possession function

cumulus/pallets/collator-selection/Cargo.toml

@@ -30,6 +30,7 @@ sp-runtime = { workspace = true }
 sp-staking = { workspace = true }
 
 frame-benchmarking = { optional = true, workspace = true }
+cumulus-pallet-session-benchmarking = { default-features = false, optional = true, workspace = true }
 
 [dev-dependencies]
 pallet-aura = { workspace = true, default-features = true }
@@ -42,6 +43,7 @@ sp-tracing = { workspace = true, default-features = true }
 [features]
 default = ["std"]
 runtime-benchmarks = [
+	"cumulus-pallet-session-benchmarking/runtime-benchmarks",
 	"frame-benchmarking/runtime-benchmarks",
 	"frame-support/runtime-benchmarks",
 	"frame-system/runtime-benchmarks",
@@ -53,6 +55,7 @@ runtime-benchmarks = [
 ]
 std = [
 	"codec/std",
+	"cumulus-pallet-session-benchmarking?/std",
 	"frame-benchmarking/std",
 	"frame-support/std",
 	"frame-system/std",

cumulus/pallets/collator-selection/src/benchmarking.rs

@@ -22,20 +22,20 @@ use super::*;
 #[allow(unused)]
 use crate::Pallet as CollatorSelection;
 use alloc::vec::Vec;
-use codec::Decode;
 use core::cmp;
+use cumulus_pallet_session_benchmarking as session_benchmarking;
 use frame_benchmarking::{account, v2::*, whitelisted_caller, BenchmarkError};
 use frame_support::traits::{Currency, EnsureOrigin, Get, ReservableCurrency};
 use frame_system::{pallet_prelude::BlockNumberFor, EventRecord, RawOrigin};
 use pallet_authorship::EventHandler;
 use pallet_session::{self as session, SessionManager};
 
 pub type BalanceOf<T> =
-	<<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance;
+	<<T as super::Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance;
 
 const SEED: u32 = 0;
 
-fn assert_last_event<T: Config>(generic_event: <T as Config>::RuntimeEvent) {
+fn assert_last_event<T: Config>(generic_event: <T as super::Config>::RuntimeEvent) {
 	let events = frame_system::Pallet::<T>::events();
 	let system_event: <T as frame_system::Config>::RuntimeEvent = generic_event.into();
 	// compare to the last event record
@@ -54,36 +54,24 @@ fn create_funded_user<T: Config>(
 	user
 }
 
-fn keys<T: Config + session::Config>(c: u32) -> <T as session::Config>::Keys {
-	use rand::{RngCore, SeedableRng};
+fn validator<T: Config + session_benchmarking::Config>(
+	c: u32,
+) -> (T::AccountId, <T as session::Config>::Keys, Vec<u8>) {
+	let validator = create_funded_user::<T>("candidate", c, 1000);
+	let (keys, proof) = T::generate_session_keys_and_proof(validator.clone());
 
-	let keys = {
-		let mut keys = [0u8; 128];
-
-		if c > 0 {
-			let mut rng = rand::rngs::StdRng::seed_from_u64(c as u64);
-			rng.fill_bytes(&mut keys);
-		}
-
-		keys
-	};
-
-	Decode::decode(&mut &keys[..]).unwrap()
+	(validator, keys, proof)
 }
 
-fn validator<T: Config + session::Config>(c: u32) -> (T::AccountId, <T as session::Config>::Keys) {
-	(create_funded_user::<T>("candidate", c, 1000), keys::<T>(c))
-}
-
-fn register_validators<T: Config + session::Config>(count: u32) -> Vec<T::AccountId> {
+fn register_validators<T: Config + session_benchmarking::Config>(count: u32) -> Vec<T::AccountId> {
 	let validators = (0..count).map(|c| validator::<T>(c)).collect::<Vec<_>>();
 
-	for (who, keys) in validators.clone() {
+	for (who, keys, proof) in validators.clone() {
 		<session::Pallet<T>>::ensure_can_pay_key_deposit(&who).unwrap();
-		<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new()).unwrap();
+		<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, proof).unwrap();
 	}
 
-	validators.into_iter().map(|(who, _)| who).collect()
+	validators.into_iter().map(|(who, _, _)| who).collect()
 }
 
 fn register_candidates<T: Config>(count: u32) {
@@ -114,7 +102,7 @@ fn min_invulnerables<T: Config>() -> u32 {
 	min_collators.saturating_sub(candidates_length)
 }
 
-#[benchmarks(where T: pallet_authorship::Config + session::Config)]
+#[benchmarks(where T: pallet_authorship::Config + session_benchmarking::Config)]
 mod benchmarks {
 	use super::*;
 
@@ -155,16 +143,15 @@ mod benchmarks {
 		// add one more to the list. should not be in `b` (invulnerables) because it's the account
 		// we will _add_ to invulnerables. we want it to be in `candidates` because we need the
 		// weight associated with removing it.
-		let (new_invulnerable, new_invulnerable_keys) = validator::<T>(b.max(c) + 1);
-		candidates.push((new_invulnerable.clone(), new_invulnerable_keys));
+		let (new_invulnerable, new_invulnerable_keys, new_proof) = validator::<T>(b.max(c) + 1);
+		candidates.push((new_invulnerable.clone(), new_invulnerable_keys, new_proof));
 		// set their keys ...
-		for (who, keys) in candidates.clone() {
+		for (who, keys, proof) in candidates.clone() {
 			<session::Pallet<T>>::ensure_can_pay_key_deposit(&who).unwrap();
-			<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new())
-				.unwrap();
+			<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, proof).unwrap();
 		}
 		// ... and register them.
-		for (who, _) in candidates.iter() {
+		for (who, _, _) in candidates.iter() {
 			let deposit = CandidacyBond::<T>::get();
 			<T as pallet::Config>::Currency::make_free_balance_be(who, deposit * 1000_u32.into());
 			CandidateList::<T>::try_mutate(|list| {
@@ -299,14 +286,11 @@ mod benchmarks {
 		let caller: T::AccountId = whitelisted_caller();
 		let bond: BalanceOf<T> = <T as pallet::Config>::Currency::minimum_balance() * 2u32.into();
 		<T as pallet::Config>::Currency::make_free_balance_be(&caller, bond);
+		let (keys, proof) = T::generate_session_keys_and_proof(caller.clone());
 
 		<session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
-		<session::Pallet<T>>::set_keys(
-			RawOrigin::Signed(caller.clone()).into(),
-			keys::<T>(c + 1),
-			Vec::new(),
-		)
-		.unwrap();
+		<session::Pallet<T>>::set_keys(RawOrigin::Signed(caller.clone()).into(), keys, proof)
+			.unwrap();
 
 		#[extrinsic_call]
 		_(RawOrigin::Signed(caller.clone()));
@@ -328,13 +312,11 @@ mod benchmarks {
 		let bond: BalanceOf<T> = <T as pallet::Config>::Currency::minimum_balance() * 10u32.into();
 		<T as pallet::Config>::Currency::make_free_balance_be(&caller, bond);
 
+		let (keys, proof) = T::generate_session_keys_and_proof(caller.clone());
+
 		<session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
-		<session::Pallet<T>>::set_keys(
-			RawOrigin::Signed(caller.clone()).into(),
-			keys::<T>(c + 1),
-			Vec::new(),
-		)
-		.unwrap();
+		<session::Pallet<T>>::set_keys(RawOrigin::Signed(caller.clone()).into(), keys, proof)
+			.unwrap();
 
 		let target = CandidateList::<T>::get().iter().last().unwrap().who.clone();
 

cumulus/pallets/collator-selection/src/mock.rs

@@ -174,6 +174,17 @@ impl Config for Test {
 	type WeightInfo = ();
 }
 
+#[cfg(feature = "runtime-benchmarks")]
+impl cumulus_pallet_session_benchmarking::Config for Test {
+	fn generate_session_keys_and_proof(owner: Self::AccountId) -> (Self::Keys, Vec<u8>) {
+		use codec::Encode;
+
+		let keys = MockSessionKeys::generate(&owner.encode(), None);
+
+		(keys.keys, keys.proof.encode())
+	}
+}
+
 pub fn new_test_ext() -> sp_io::TestExternalities {
 	sp_tracing::try_init_simple();
 	let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();

cumulus/pallets/collator-selection/src/tests.rs

@@ -18,6 +18,7 @@ use crate::{
 	mock::*, CandidacyBond, CandidateInfo, CandidateList, DesiredCandidates, Error, Invulnerables,
 	LastAuthoredBlock,
 };
+use codec::Encode;
 use frame_support::{
 	assert_noop, assert_ok,
 	traits::{Currency, OnInitialize},
@@ -123,8 +124,13 @@ fn invulnerable_limit_works() {
 			// only keys were registered in mock for 1 to 5
 			if ii > 5 {
 				Balances::make_free_balance_be(&ii, 100);
-				let key = MockSessionKeys { aura: UintAuthorityId(ii) };
-				Session::set_keys(RuntimeOrigin::signed(ii).into(), key, Vec::new()).unwrap();
+				let mut key = MockSessionKeys { aura: UintAuthorityId(ii) };
+				Session::set_keys(
+					RuntimeOrigin::signed(ii).into(),
+					key.clone(),
+					key.create_ownership_proof(&ii.encode()).unwrap().encode(),
+				)
+				.unwrap();
 			}
 			assert_eq!(Balances::free_balance(ii), 100);
 			if ii < 21 {
@@ -482,8 +488,13 @@ fn cannot_register_candidate_if_too_many() {
 		// candidates.
 		for i in 6..=23 {
 			Balances::make_free_balance_be(&i, 100);
-			let key = MockSessionKeys { aura: UintAuthorityId(i) };
-			Session::set_keys(RuntimeOrigin::signed(i).into(), key, Vec::new()).unwrap();
+			let mut key = MockSessionKeys { aura: UintAuthorityId(i) };
+			Session::set_keys(
+				RuntimeOrigin::signed(i).into(),
+				key.clone(),
+				key.create_ownership_proof(&i.encode()).unwrap().encode(),
+			)
+			.unwrap();
 		}
 
 		for c in 3..=22 {
@@ -765,8 +776,13 @@ fn take_candidate_slot_works() {
 		assert_eq!(CandidateList::<Test>::get().iter().count(), 3);
 
 		Balances::make_free_balance_be(&6, 100);
-		let key = MockSessionKeys { aura: UintAuthorityId(6) };
-		Session::set_keys(RuntimeOrigin::signed(6).into(), key, Vec::new()).unwrap();
+		let mut key = MockSessionKeys { aura: UintAuthorityId(6) };
+		Session::set_keys(
+			RuntimeOrigin::signed(6).into(),
+			key.clone(),
+			key.create_ownership_proof(&6u64.encode()).unwrap().encode(),
+		)
+		.unwrap();
 
 		assert_ok!(CollatorSelection::take_candidate_slot(
 			RuntimeOrigin::signed(6),

cumulus/pallets/session-benchmarking/src/inner.rs

@@ -16,14 +16,20 @@
 //! Benchmarking setup for pallet-session.
 #![cfg(feature = "runtime-benchmarks")]
 
-use alloc::{vec, vec::Vec};
+use alloc::vec::Vec;
 
-use codec::Decode;
 use frame_benchmarking::v2::*;
 use frame_system::RawOrigin;
 use pallet_session::*;
 pub struct Pallet<T: Config>(pallet_session::Pallet<T>);
-pub trait Config: pallet_session::Config {}
+pub trait Config: pallet_session::Config {
+	/// Generate a session key and a proof of ownership.
+	///
+	/// The given `owner` is the account that will call `set_keys` using the returned session keys
+	/// and proof. This means that the proof should prove the ownership of `owner` over the private
+	/// keys associated to the session keys.
+	fn generate_session_keys_and_proof(owner: Self::AccountId) -> (Self::Keys, Vec<u8>);
+}
 
 #[benchmarks]
 mod benchmarks {
@@ -33,8 +39,8 @@ mod benchmarks {
 	fn set_keys() -> Result<(), BenchmarkError> {
 		let caller: T::AccountId = whitelisted_caller();
 		frame_system::Pallet::<T>::inc_providers(&caller);
-		let keys = T::Keys::decode(&mut sp_runtime::traits::TrailingZeroInput::zeroes()).unwrap();
-		let proof: Vec<u8> = vec![0, 1, 2, 3];
+		let (keys, proof) = T::generate_session_keys_and_proof(caller.clone());
+
 		<pallet_session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
 
 		#[extrinsic_call]
@@ -47,9 +53,9 @@ mod benchmarks {
 	fn purge_keys() -> Result<(), BenchmarkError> {
 		let caller: T::AccountId = whitelisted_caller();
 		frame_system::Pallet::<T>::inc_providers(&caller);
-		let keys = T::Keys::decode(&mut sp_runtime::traits::TrailingZeroInput::zeroes()).unwrap();
-		let proof: Vec<u8> = vec![0, 1, 2, 3];
+		let (keys, proof) = T::generate_session_keys_and_proof(caller.clone());
 		<pallet_session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
+
 		let _t = pallet_session::Pallet::<T>::set_keys(
 			RawOrigin::Signed(caller.clone()).into(),
 			keys,

cumulus/parachains/runtimes/assets/asset-hub-rococo/src/lib.rs

@@ -48,6 +48,7 @@ use sp_runtime::{
 	transaction_validity::{TransactionSource, TransactionValidity},
 	ApplyExtrinsicResult, Permill,
 };
+use sp_session::OpaqueGeneratedSessionKeys;
 use testnet_parachains_constants::rococo::snowbridge::EthereumNetwork;
 
 #[cfg(feature = "std")]
@@ -1437,8 +1438,8 @@ impl_runtime_apis! {
 	}
 
 	impl sp_session::SessionKeys<Block> for Runtime {
-		fn generate_session_keys(seed: Option<Vec<u8>>) -> Vec<u8> {
-			SessionKeys::generate(seed)
+		fn generate_session_keys(owner: Vec<u8>, seed: Option<Vec<u8>>) -> OpaqueGeneratedSessionKeys {
+			SessionKeys::generate(&owner, seed).into()
 		}
 
 		fn decode_session_keys(
@@ -1695,7 +1696,12 @@ impl_runtime_apis! {
 			}
 
 			use cumulus_pallet_session_benchmarking::Pallet as SessionBench;
-			impl cumulus_pallet_session_benchmarking::Config for Runtime {}
+			impl cumulus_pallet_session_benchmarking::Config for Runtime {
+				fn generate_session_keys_and_proof(owner: Self::AccountId) -> (Self::Keys, Vec<u8>) {
+					let keys = SessionKeys::generate(&owner.encode(), None);
+					(keys.keys, keys.proof.encode())
+				}
+			}
 
 			use pallet_xcm_bridge_hub_router::benchmarking::{
 				Pallet as XcmBridgeHubRouterBench,

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

@@ -1836,8 +1836,8 @@ pallet_revive::impl_runtime_apis_plus_revive_traits!(
 	}
 
 	impl sp_session::SessionKeys<Block> for Runtime {
-		fn generate_session_keys(seed: Option<Vec<u8>>) -> Vec<u8> {
-			SessionKeys::generate(seed)
+		fn generate_session_keys(owner: Vec<u8>, seed: Option<Vec<u8>>) -> sp_session::OpaqueGeneratedSessionKeys {
+			SessionKeys::generate(&owner, seed).into()
 		}
 
 		fn decode_session_keys(
@@ -2224,9 +2224,20 @@ pallet_revive::impl_runtime_apis_plus_revive_traits!(
 			}
 
 			use cumulus_pallet_session_benchmarking::Pallet as SessionBench;
+			impl cumulus_pallet_session_benchmarking::Config for Runtime {
+				fn generate_session_keys_and_proof(owner: Self::AccountId) -> (Self::Keys, Vec<u8>) {
+					let keys = SessionKeys::generate(&owner.encode(), None);
+					(keys.keys, keys.proof.encode())
+				}
+			}
 			use xcm_config::{MaxAssetsIntoHolding, WestendLocation};
 
-			impl cumulus_pallet_session_benchmarking::Config for Runtime {}
+			impl cumulus_pallet_session_benchmarking::Config for Runtime {
+				fn generate_session_keys_and_proof(owner: Self::AccountId) -> (Self::Keys, Vec<u8>) {
+					let keys = SessionKeys::generate(&owner.encode(), None);
+					(keys.keys, keys.proof.encode())
+				}
+			}
 			use testnet_parachains_constants::westend::locations::{PeopleParaId, PeopleLocation};
 			parameter_types! {
 				pub ExistentialDepositAsset: Option<Asset> = Some((

cumulus/parachains/runtimes/bridge-hubs/bridge-hub-rococo/src/lib.rs

@@ -806,8 +806,8 @@ impl_runtime_apis! {
 	}
 
 	impl sp_session::SessionKeys<Block> for Runtime {
-		fn generate_session_keys(seed: Option<Vec<u8>>) -> Vec<u8> {
-			SessionKeys::generate(seed)
+		fn generate_session_keys(owner: Vec<u8>, seed: Option<Vec<u8>>) -> sp_session::OpaqueGeneratedSessionKeys {
+			SessionKeys::generate(&owner, seed).into()
 		}
 
 		fn decode_session_keys(
@@ -1088,6 +1088,7 @@ impl_runtime_apis! {
 		) -> Result<Vec<frame_benchmarking::BenchmarkBatch>, alloc::string::String> {
 			use frame_benchmarking::{BenchmarkBatch, BenchmarkError};
 			use sp_storage::TrackedStorageKey;
+			use codec::Encode;
 
 			use frame_system_benchmarking::Pallet as SystemBench;
 			use frame_system_benchmarking::extensions::Pallet as SystemExtensionsBench;
@@ -1103,7 +1104,12 @@ impl_runtime_apis! {
 			}
 
 			use cumulus_pallet_session_benchmarking::Pallet as SessionBench;
-			impl cumulus_pallet_session_benchmarking::Config for Runtime {}
+			impl cumulus_pallet_session_benchmarking::Config for Runtime {
+				fn generate_session_keys_and_proof(owner: Self::AccountId) -> (Self::Keys, Vec<u8>) {
+					let keys = SessionKeys::generate(&owner.encode(), None);
+					(keys.keys, keys.proof.encode())
+				}
+			}
 
 			use xcm::latest::prelude::*;
 			use xcm_config::TokenLocation;