frame-session: Introduce a proper proof of key ownership

cumulus/pallets/collator-selection/Cargo.toml

@@ -30,6 +30,7 @@ pallet-balances = { workspace = true }
 pallet-session = { workspace = true }
 
 frame-benchmarking = { optional = true, workspace = true }
+cumulus-pallet-session-benchmarking = { default-features = false, optional = true, workspace = true }
 
 [dev-dependencies]
 sp-core = { workspace = true, default-features = true }
@@ -43,6 +44,7 @@ pallet-aura = { 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

@@ -24,18 +24,19 @@ 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,35 +55,23 @@ 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() {
-		<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new()).unwrap();
+	for (who, keys, proof) in validators.clone() {
+		<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) {
@@ -110,7 +99,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::*;
 
@@ -151,15 +140,14 @@ 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() {
-			<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new())
-				.unwrap();
+		for (who, keys, proof) in candidates.clone() {
+			<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::Currency::make_free_balance_be(who, deposit * 1000_u32.into());
 			CandidateList::<T>::try_mutate(|list| {
@@ -293,13 +281,10 @@ mod benchmarks {
 		let caller: T::AccountId = whitelisted_caller();
 		let bond: BalanceOf<T> = T::Currency::minimum_balance() * 2u32.into();
 		T::Currency::make_free_balance_be(&caller, bond);
+		let (keys, proof) = T::generate_session_keys_and_proof(caller.clone());
 
-		<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()));
@@ -321,12 +306,10 @@ mod benchmarks {
 		let bond: BalanceOf<T> = T::Currency::minimum_balance() * 10u32.into();
 		T::Currency::make_free_balance_be(&caller, bond);
 
-		<session::Pallet<T>>::set_keys(
-			RawOrigin::Signed(caller.clone()).into(),
-			keys::<T>(c + 1),
-			Vec::new(),
-		)
-		.unwrap();
+		let (keys, proof) = T::generate_session_keys_and_proof(caller.clone());
+
+		<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

@@ -172,6 +172,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},
@@ -124,7 +125,12 @@ fn invulnerable_limit_works() {
 			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();
+				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 {

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

@@ -22,21 +22,28 @@ use frame_benchmarking::{benchmarks, whitelisted_caller};
 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! {
 	set_keys {
 		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());
 	}: _(RawOrigin::Signed(caller), keys, proof)
 
 	purge_keys {
 		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());
 		let _t = pallet_session::Pallet::<T>::set_keys(RawOrigin::Signed(caller.clone()).into(), keys, proof);
 	}: _(RawOrigin::Signed(caller))
 }

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

@@ -46,6 +46,7 @@ use sp_runtime::{
 	transaction_validity::{TransactionSource, TransactionValidity},
 	ApplyExtrinsicResult, Permill,
 };
+use sp_session::OpaqueGeneratedSessionKeys;
 use testnet_parachains_constants::rococo::snowbridge::EthereumNetwork;
 
 #[cfg(feature = "std")]
@@ -1216,8 +1217,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(
@@ -1468,7 +1469,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

@@ -1266,8 +1266,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(
@@ -1564,7 +1564,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())
+				}
+			}
 
 			parameter_types! {
 				pub ExistentialDepositAsset: Option<Asset> = Some((

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

@@ -910,8 +910,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(
@@ -1199,6 +1199,7 @@ impl_runtime_apis! {
 		) -> Result<Vec<frame_benchmarking::BenchmarkBatch>, sp_runtime::RuntimeString> {
 			use frame_benchmarking::{Benchmarking, BenchmarkBatch, BenchmarkError};
 			use sp_storage::TrackedStorageKey;
+			use codec::Encode;
 
 			use frame_system_benchmarking::Pallet as SystemBench;
 			impl frame_system_benchmarking::Config for Runtime {
@@ -1213,7 +1214,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::benchmarking::Pallet as PalletXcmExtrinsicsBenchmark;
 			impl pallet_xcm::benchmarking::Config for Runtime {

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

@@ -53,6 +53,7 @@ use sp_runtime::{
 	ApplyExtrinsicResult,
 };
 
+use sp_session::OpaqueGeneratedSessionKeys;
 #[cfg(feature = "std")]
 use sp_version::NativeVersion;
 use sp_version::RuntimeVersion;
@@ -677,8 +678,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(
@@ -921,6 +922,7 @@ impl_runtime_apis! {
 		) -> Result<Vec<frame_benchmarking::BenchmarkBatch>, sp_runtime::RuntimeString> {
 			use frame_benchmarking::{Benchmarking, BenchmarkBatch, BenchmarkError};
 			use sp_storage::TrackedStorageKey;
+			use codec::Encode;
 
 			use frame_system_benchmarking::Pallet as SystemBench;
 			impl frame_system_benchmarking::Config for Runtime {
@@ -935,7 +937,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::benchmarking::Pallet as PalletXcmExtrinsicsBenchmark;
 			impl pallet_xcm::benchmarking::Config for Runtime {