Skip to content
This repository has been archived by the owner on Nov 15, 2023. It is now read-only.

Commit

Permalink
Disarm OnRuntimeUpgrade::pre/post_upgrade Tuple footgun (#14759)
Browse files Browse the repository at this point in the history
* 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 <kungfukeith11@gmail.com>

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

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

* address feedback

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

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

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

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

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

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

* muharem comments

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

Co-authored-by: Muharem Ismailov <ismailov.m.h@gmail.com>

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

Co-authored-by: Muharem Ismailov <ismailov.m.h@gmail.com>

* remove useless type

---------

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>
Co-authored-by: Muharem Ismailov <ismailov.m.h@gmail.com>
  • Loading branch information
3 people authored and Ank4n committed Aug 20, 2023
1 parent 3ddf325 commit 81c7234
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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";
Expand All @@ -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;
Expand Down
45 changes: 39 additions & 6 deletions frame/support/src/traits/hooks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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(())
Expand All @@ -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();
Expand Down Expand Up @@ -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`].
Expand Down Expand Up @@ -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"]);
Expand All @@ -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());
});
}

Expand Down

0 comments on commit 81c7234

Please sign in to comment.