[Pools] fix derivation of pool account (#4999)

closes paritytech-secops/srlabs_findings#408.
This fixes how ProxyDelegator accounts are derived but may cause issues
in Westend since it would use the old derivative accounts. Does not
affect Polkadot/Kusama as this pallet is not deployed to them yet.

---------

Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

polkadot/runtime/westend/src/lib.rs

@@ -1760,11 +1760,8 @@ pub type SignedExtra = (
 );
 
 parameter_types! {
-	// This is the max pools that will be migrated in the runtime upgrade. Westend has more pools
-	// than this, but we want to emulate some non migrated pools. In prod runtimes, if weight is not
-	// a concern, it is recommended to set to (existing pools + 10) to also account for any new
-	// pools getting created before the migration is actually executed.
-	pub const MaxPoolsToMigrate: u32 = 250;
+	/// Bounding number of agent pot accounts to be migrated in a single block.
+	pub const MaxAgentsToMigrate: u32 = 300;
 }
 
 /// All migrations that will run on the next runtime upgrade.
@@ -1798,13 +1795,11 @@ pub mod migrations {
 
 	/// Unreleased migrations. Add new ones here:
 	pub type Unreleased = (
-		// Migrate NominationPools to `DelegateStake` adapter. This is unversioned upgrade and
-		// should not be applied yet in Kusama/Polkadot.
-		pallet_nomination_pools::migration::unversioned::DelegationStakeMigration<
+		// This is only needed for Westend.
+		pallet_delegated_staking::migration::unversioned::ProxyDelegatorMigration<
 			Runtime,
-			MaxPoolsToMigrate,
+			MaxAgentsToMigrate,
 		>,
-		pallet_staking::migrations::v15::MigrateV14ToV15<Runtime>,
 	);
 }
 

prdoc/pr_4999.prdoc

@@ -0,0 +1,18 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Fixes entropy for derivation of proxy delegator account.
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      This fixes how ProxyDelegator accounts are derived but may cause issues in Westend since it would use the old
+      derivative accounts. Does not affect Polkadot/Kusama as this pallet is not deployed to them yet.
+
+crates:
+  - name: westend-runtime
+    bump: patch
+  - name: pallet-delegated-staking
+    bump: patch
+  - name: pallet-nomination-pools
+    bump: patch

substrate/frame/delegated-staking/Cargo.toml

@@ -18,6 +18,8 @@ frame-system = { workspace = true }
 scale-info = { features = ["derive"], workspace = true }
 sp-runtime = { workspace = true }
 sp-staking = { workspace = true }
+sp-io = { workspace = true }
+log = { workspace = true }
 
 [dev-dependencies]
 sp-core = { workspace = true, default-features = true }
@@ -38,6 +40,7 @@ std = [
 	"frame-election-provider-support/std",
 	"frame-support/std",
 	"frame-system/std",
+	"log/std",
 	"pallet-balances/std",
 	"pallet-nomination-pools/std",
 	"pallet-staking/std",

substrate/frame/delegated-staking/src/lib.rs

@@ -126,6 +126,7 @@
 #![deny(rustdoc::broken_intra_doc_links)]
 
 mod impls;
+pub mod migration;
 #[cfg(test)]
 mod mock;
 #[cfg(test)]
@@ -152,12 +153,25 @@ use frame_support::{
 		Defensive, DefensiveOption, Imbalance, OnUnbalanced,
 	},
 };
+use sp_io::hashing::blake2_256;
 use sp_runtime::{
-	traits::{AccountIdConversion, CheckedAdd, CheckedSub, Zero},
+	traits::{CheckedAdd, CheckedSub, TrailingZeroInput, Zero},
 	ArithmeticError, DispatchResult, Perbill, RuntimeDebug, Saturating,
 };
 use sp_staking::{Agent, Delegator, EraIndex, StakingInterface, StakingUnchecked};
 
+/// The log target of this pallet.
+pub const LOG_TARGET: &str = "runtime::delegated-staking";
+// syntactic sugar for logging.
+#[macro_export]
+macro_rules! log {
+	($level:tt, $patter:expr $(, $values:expr)* $(,)?) => {
+		log::$level!(
+			target: $crate::LOG_TARGET,
+			concat!("[{:?}] 🏊‍♂️ ", $patter), <frame_system::Pallet<T>>::block_number() $(, $values)*
+		)
+	};
+}
 pub type BalanceOf<T> =
 	<<T as Config>::Currency as FunInspect<<T as frame_system::Config>::AccountId>>::Balance;
 
@@ -463,7 +477,9 @@ impl<T: Config> Pallet<T> {
 
 	/// Derive a (keyless) pot account from the given agent account and account type.
 	fn sub_account(account_type: AccountType, acc: T::AccountId) -> T::AccountId {
-		T::PalletId::get().into_sub_account_truncating((account_type, acc.clone()))
+		let entropy = (T::PalletId::get(), acc, account_type).using_encoded(blake2_256);
+		Decode::decode(&mut TrailingZeroInput::new(entropy.as_ref()))
+			.expect("infinite length input; no invalid inputs for type; qed")
 	}
 
 	/// Held balance of a delegator.

substrate/frame/delegated-staking/src/migration.rs

@@ -0,0 +1,107 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::*;
+use frame_support::traits::OnRuntimeUpgrade;
+
+#[cfg(feature = "try-runtime")]
+use sp_runtime::TryRuntimeError;
+
+pub mod unversioned {
+	use super::*;
+	#[cfg(feature = "try-runtime")]
+	use alloc::vec::Vec;
+	use sp_runtime::traits::AccountIdConversion;
+
+	/// Migrates `ProxyDelegator` accounts with better entropy than the old logic which didn't take
+	/// into account all the bytes of the agent account ID.
+	pub struct ProxyDelegatorMigration<T, MaxAgents>(PhantomData<(T, MaxAgents)>);
+
+	impl<T: Config, MaxAgents: Get<u32>> OnRuntimeUpgrade for ProxyDelegatorMigration<T, MaxAgents> {
+		fn on_runtime_upgrade() -> Weight {
+			let mut weight = Weight::zero();
+			let old_proxy_delegator = |agent: T::AccountId| {
+				T::PalletId::get()
+					.into_sub_account_truncating((AccountType::ProxyDelegator, agent.clone()))
+			};
+
+			Agents::<T>::iter_keys().take(MaxAgents::get() as usize).for_each(|agent| {
+				weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0));
+				let old_proxy = old_proxy_delegator(agent.clone());
+
+				// if delegation does not exist, it does not need to be migrated.
+				if let Some(delegation) = Delegation::<T>::get(&old_proxy) {
+					weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0));
+
+					let new_proxy =
+						Pallet::<T>::generate_proxy_delegator(Agent::from(agent.clone()));
+
+					// accrue read writes for `do_migrate_delegation`
+					weight.saturating_accrue(T::DbWeight::get().reads_writes(8, 8));
+					let _ = Pallet::<T>::do_migrate_delegation(
+						Delegator::from(old_proxy.clone()),
+						new_proxy.clone(),
+						delegation.amount,
+					)
+					.map_err(|e| {
+						log!(
+							error,
+							"Failed to migrate old proxy delegator {:?} to new proxy {:?} for agent {:?} with error: {:?}",
+								old_proxy,
+								new_proxy,
+								agent,
+								e,
+						);
+					});
+				};
+			});
+
+			log!(info, "Finished migrating old proxy delegator accounts to new ones");
+			weight
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_data: Vec<u8>) -> Result<(), TryRuntimeError> {
+			let mut unmigrated_count = 0;
+			let old_proxy_delegator = |agent: T::AccountId| {
+				T::PalletId::get()
+					.into_sub_account_truncating((AccountType::ProxyDelegator, agent.clone()))
+			};
+
+			Agents::<T>::iter_keys().take(MaxAgents::get() as usize).for_each(|agent| {
+				let old_proxy: T::AccountId = old_proxy_delegator(agent.clone());
+				let held_balance = Pallet::<T>::held_balance_of(Delegator::from(old_proxy.clone()));
+				let delegation = Delegation::<T>::get(&old_proxy);
+				if delegation.is_some() || !held_balance.is_zero() {
+					log!(
+						error,
+						"Old proxy delegator {:?} for agent {:?} is not migrated.",
+						old_proxy,
+						agent,
+					);
+					unmigrated_count += 1;
+				}
+			});
+
+			if unmigrated_count > 0 {
+				Err(TryRuntimeError::Other("Some old proxy delegator accounts are not migrated."))
+			} else {
+				Ok(())
+			}
+		}
+	}
+}