Disarm OnRuntimeUpgrade::pre/post_upgrade Tuple footgun (#14759)

* return error on incorrect tuple usage of pre_upgrade and post_upgrade

* add test

* comment lint

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Keith Yeung <[email protected]>

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Keith Yeung <[email protected]>

* address feedback

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Keith Yeung <[email protected]>

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Keith Yeung <[email protected]>

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Keith Yeung <[email protected]>

* muharem comments

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Muharem Ismailov <[email protected]>

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Muharem Ismailov <[email protected]>

* remove useless type

---------

Co-authored-by: Keith Yeung <[email protected]>
Co-authored-by: Muharem Ismailov <[email protected]>

frame/elections-phragmen/src/migrations/unlock_and_unreserve_all_funds.rs

@@ -53,8 +53,6 @@ pub trait UnlockConfig: 'static {
 	type PalletId: Get<LockIdentifier>;
 	/// The DB weight as configured in the runtime to calculate the correct weight.
 	type DbWeight: Get<RuntimeDbWeight>;
-	/// The block number as configured in the runtime.
-	type BlockNumber: Parameter + Zero + Copy + Ord;
 }
 
 #[storage_alias(dynamic)]
@@ -331,7 +329,6 @@ mod test {
 		assert_ok, parameter_types,
 		traits::{Currency, OnRuntimeUpgrade, ReservableCurrency, WithdrawReasons},
 	};
-	use frame_system::pallet_prelude::BlockNumberFor;
 
 	parameter_types! {
 		const PalletName: &'static str = "Elections";
@@ -341,7 +338,6 @@ mod test {
 	impl super::UnlockConfig for UnlockConfigImpl {
 		type Currency = Balances;
 		type AccountId = u64;
-		type BlockNumber = BlockNumberFor<Test>;
 		type DbWeight = ();
 		type PalletName = PalletName;
 		type MaxVotesPerVoter = PhragmenMaxVoters;

frame/support/src/traits/hooks.rs

@@ -106,7 +106,16 @@ pub trait OnRuntimeUpgrade {
 		Weight::zero()
 	}
 
-	/// See [`Hooks::on_runtime_upgrade`].
+	/// The expected and default behavior of this method is to handle executing `pre_upgrade` ->
+	/// `on_runtime_upgrade` -> `post_upgrade` hooks for a migration.
+	///
+	/// Internally, the default implementation
+	/// - Handles passing data from `pre_upgrade` to `post_upgrade`
+	/// - Ensure storage is not modified in `pre_upgrade` and `post_upgrade` hooks.
+	///
+	/// Combining the `pre_upgrade` -> `on_runtime_upgrade` -> `post_upgrade` logic flow into a
+	/// single method call is helpful for scenarios like testing a tuple of migrations, where the
+	/// tuple contains order-dependent migrations.
 	#[cfg(feature = "try-runtime")]
 	fn try_on_runtime_upgrade(checks: bool) -> Result<Weight, TryRuntimeError> {
 		let maybe_state = if checks {
@@ -128,13 +137,13 @@ pub trait OnRuntimeUpgrade {
 		Ok(weight)
 	}
 
-	/// See [`Hooks::on_runtime_upgrade`].
+	/// See [`Hooks::pre_upgrade`].
 	#[cfg(feature = "try-runtime")]
 	fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
 		Ok(Vec::new())
 	}
 
-	/// See [`Hooks::on_runtime_upgrade`].
+	/// See [`Hooks::post_upgrade`].
 	#[cfg(feature = "try-runtime")]
 	fn post_upgrade(_state: Vec<u8>) -> Result<(), TryRuntimeError> {
 		Ok(())
@@ -151,9 +160,8 @@ impl OnRuntimeUpgrade for Tuple {
 		weight
 	}
 
-	// We are executing pre- and post-checks sequentially in order to be able to test several
-	// consecutive migrations for the same pallet without errors. Therefore pre and post upgrade
-	// hooks for tuples are a noop.
+	/// Implements the default behavior of `try_on_runtime_upgrade` for tuples, logging any errors
+	/// that occur.
 	#[cfg(feature = "try-runtime")]
 	fn try_on_runtime_upgrade(checks: bool) -> Result<Weight, TryRuntimeError> {
 		let mut weight = Weight::zero();
@@ -188,6 +196,26 @@ impl OnRuntimeUpgrade for Tuple {
 
 		Ok(weight)
 	}
+
+	/// [`OnRuntimeUpgrade::pre_upgrade`] should not be used on a tuple.
+	///
+	/// Instead, implementors should use [`OnRuntimeUpgrade::try_on_runtime_upgrade`] which
+	/// internally calls `pre_upgrade` -> `on_runtime_upgrade` -> `post_upgrade` for each tuple
+	/// member in sequence, enabling testing of order-dependent migrations.
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
+		Err("Usage of `pre_upgrade` with Tuples is not expected. Please use `try_on_runtime_upgrade` instead, which internally calls `pre_upgrade` -> `on_runtime_upgrade` -> `post_upgrade` for each tuple member.".into())
+	}
+
+	/// [`OnRuntimeUpgrade::post_upgrade`] should not be used on a tuple.
+	///
+	/// Instead, implementors should use [`OnRuntimeUpgrade::try_on_runtime_upgrade`] which
+	/// internally calls `pre_upgrade` -> `on_runtime_upgrade` -> `post_upgrade` for each tuple
+	/// member in sequence, enabling testing of order-dependent migrations.
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(_state: Vec<u8>) -> Result<(), TryRuntimeError> {
+		Err("Usage of `post_upgrade` with Tuples is not expected. Please use `try_on_runtime_upgrade` instead, which internally calls `pre_upgrade` -> `on_runtime_upgrade` -> `post_upgrade` for each tuple member.".into())
+	}
 }
 
 /// See [`Hooks::integrity_test`].
@@ -497,6 +525,7 @@ mod tests {
 		impl_test_type!(Baz);
 
 		TestExternalities::default().execute_with(|| {
+			// try_on_runtime_upgrade works
 			Foo::try_on_runtime_upgrade(true).unwrap();
 			assert_eq!(Pre::take(), vec!["Foo"]);
 			assert_eq!(Post::take(), vec!["Foo"]);
@@ -512,6 +541,10 @@ mod tests {
 			<(Foo, (Bar, Baz))>::try_on_runtime_upgrade(true).unwrap();
 			assert_eq!(Pre::take(), vec!["Foo", "Bar", "Baz"]);
 			assert_eq!(Post::take(), vec!["Foo", "Bar", "Baz"]);
+
+			// calling pre_upgrade and post_upgrade directly on tuple of pallets fails
+			assert!(<(Foo, (Bar, Baz))>::pre_upgrade().is_err());
+			assert!(<(Foo, (Bar, Baz))>::post_upgrade(vec![]).is_err());
 		});
 	}