Improve handling of unset StorageVersion (paritytech#13417)

* Improve handling of unset `StorageVersion`

When a user is forgetting to set the storage version in a pallet and calls
`current_storage_version` to compare it against the `on_chain_storage_version` it will now fail to
compile the code. Before the pallet macro just returned `StorageVersion::default()` for
`current_storage_version` leading to potential issues with migrations. Besides that it also checks
in `post_upgrade` that the pallet storage version was upgraded and thus, no migration was missed.

* Use correct `Cargo.lock`

* Fixes

* Fix test

* Update frame/support/test/tests/pallet.rs

* Ensure we don't set a storage version when the pallet is missing the attribute

* Fix merge conflict

* Update frame/support/procedural/src/pallet/expand/hooks.rs

Co-authored-by: Roman Useinov <roman.useinov@gmail.com>

* Update frame/support/procedural/src/pallet/expand/hooks.rs

Co-authored-by: Roman Useinov <roman.useinov@gmail.com>

* Fix compilation

* Do not run everything with `try-runtime`

* Fix test

* Apply suggestions from code review

Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* Fix `no-metadata-docs`

---------

Co-authored-by: Roman Useinov <roman.useinov@gmail.com>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: parity-processbot <>

frame/support/procedural/src/pallet/expand/hooks.rs

@@ -82,6 +82,57 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 		proc_macro2::TokenStream::new()
 	};
 
+	// If a storage version is set, we should ensure that the storage version on chain matches the
+	// current storage version. This assumes that `Executive` is running custom migrations before
+	// the pallets are called.
+	let post_storage_version_check = if def.pallet_struct.storage_version.is_some() {
+		quote::quote! {
+			let on_chain_version = <Self as #frame_support::traits::GetStorageVersion>::on_chain_storage_version();
+			let current_version = <Self as #frame_support::traits::GetStorageVersion>::current_storage_version();
+
+			if on_chain_version != current_version {
+				let pallet_name = <
+					<T as #frame_system::Config>::PalletInfo
+					as
+					#frame_support::traits::PalletInfo
+				>::name::<Self>().unwrap_or("<unknown pallet name>");
+
+				#frame_support::log::error!(
+					target: #frame_support::LOG_TARGET,
+					"{}: On chain storage version {:?} doesn't match current storage version {:?}.",
+					pallet_name,
+					on_chain_version,
+					current_version,
+				);
+
+				return Err("On chain and current storage version do not match. Missing runtime upgrade?");
+			}
+		}
+	} else {
+		quote::quote! {
+			let on_chain_version = <Self as #frame_support::traits::GetStorageVersion>::on_chain_storage_version();
+
+			if on_chain_version != #frame_support::traits::StorageVersion::new(0) {
+				let pallet_name = <
+					<T as #frame_system::Config>::PalletInfo
+					as
+					#frame_support::traits::PalletInfo
+				>::name::<Self>().unwrap_or("<unknown pallet name>");
+
+				#frame_support::log::error!(
+					target: #frame_support::LOG_TARGET,
+					"{}: On chain storage version {:?} is set to non zero,\
+					 while the pallet is missing the `#[pallet::storage_version(VERSION)]` attribute.",
+					pallet_name,
+					on_chain_version,
+				);
+
+				return Err("On chain storage version set, while the pallet doesn't \
+							have the `#[pallet::storage_version(VERSION)]` attribute.");
+			}
+		}
+	};
+
 	quote::quote_spanned!(span =>
 		#hooks_impl
 
@@ -170,6 +221,8 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 
 			#[cfg(feature = "try-runtime")]
 			fn post_upgrade(state: #frame_support::sp_std::vec::Vec<u8>) -> Result<(), &'static str> {
+				#post_storage_version_check
+
 				<
 					Self
 					as

frame/support/procedural/src/pallet/expand/pallet_struct.rs

@@ -160,11 +160,15 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 		}
 	);
 
-	let storage_version = if let Some(v) = def.pallet_struct.storage_version.as_ref() {
-		quote::quote! { #v }
-	} else {
-		quote::quote! { #frame_support::traits::StorageVersion::default() }
-	};
+	let (storage_version, current_storage_version_ty) =
+		if let Some(v) = def.pallet_struct.storage_version.as_ref() {
+			(quote::quote! { #v }, quote::quote! { #frame_support::traits::StorageVersion })
+		} else {
+			(
+				quote::quote! { core::default::Default::default() },
+				quote::quote! { #frame_support::traits::NoStorageVersionSet },
+			)
+		};
 
 	let whitelisted_storage_idents: Vec<syn::Ident> = def
 		.storages
@@ -199,7 +203,9 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 			for #pallet_ident<#type_use_gen>
 			#config_where_clause
 		{
-			fn current_storage_version() -> #frame_support::traits::StorageVersion {
+			type CurrentStorageVersion = #current_storage_version_ty;
+
+			fn current_storage_version() -> Self::CurrentStorageVersion {
 				#storage_version
 			}
 
@@ -214,7 +220,7 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 			#config_where_clause
 		{
 			fn on_genesis() {
-				let storage_version = #storage_version;
+				let storage_version: #frame_support::traits::StorageVersion = #storage_version;
 				storage_version.put::<Self>();
 			}
 		}

frame/support/src/dispatch.rs

@@ -2500,14 +2500,26 @@ macro_rules! decl_module {
 		impl<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?> $crate::traits::GetStorageVersion
 			for $module<$trait_instance $(, $instance)?> where $( $other_where_bounds )*
 		{
-			fn current_storage_version() -> $crate::traits::StorageVersion {
+			type CurrentStorageVersion = $crate::traits::StorageVersion;
+
+			fn current_storage_version() -> Self::CurrentStorageVersion {
 				$( $storage_version )*
 			}
 
 			fn on_chain_storage_version() -> $crate::traits::StorageVersion {
 				$crate::traits::StorageVersion::get::<Self>()
 			}
 		}
+
+		// Implement `OnGenesis` for `Module`
+		impl<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?> $crate::traits::OnGenesis
+			for $module<$trait_instance $(, $instance)?> where $( $other_where_bounds )*
+		{
+			fn on_genesis() {
+				let storage_version = <Self as $crate::traits::GetStorageVersion>::current_storage_version();
+				storage_version.put::<Self>();
+			}
+		}
 	};
 
 	// Implementation for `GetStorageVersion` when no storage version is passed.
@@ -2519,14 +2531,26 @@ macro_rules! decl_module {
 		impl<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?> $crate::traits::GetStorageVersion
 			for $module<$trait_instance $(, $instance)?> where $( $other_where_bounds )*
 		{
-			fn current_storage_version() -> $crate::traits::StorageVersion {
+			type CurrentStorageVersion = $crate::traits::NoStorageVersionSet;
+
+			fn current_storage_version() -> Self::CurrentStorageVersion {
 				Default::default()
 			}
 
 			fn on_chain_storage_version() -> $crate::traits::StorageVersion {
 				$crate::traits::StorageVersion::get::<Self>()
 			}
 		}
+
+		// Implement `OnGenesis` for `Module`
+		impl<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?> $crate::traits::OnGenesis
+			for $module<$trait_instance $(, $instance)?> where $( $other_where_bounds )*
+		{
+			fn on_genesis() {
+				let storage_version = $crate::traits::StorageVersion::default();
+				storage_version.put::<Self>();
+			}
+		}
 	};
 
 	// The main macro expansion that actually renders the module code.
@@ -2814,16 +2838,6 @@ macro_rules! decl_module {
 			}
 		}
 
-		// Implement `OnGenesis` for `Module`
-		impl<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?> $crate::traits::OnGenesis
-			for $mod_type<$trait_instance $(, $instance)?> where $( $other_where_bounds )*
-		{
-			fn on_genesis() {
-				let storage_version = <Self as $crate::traits::GetStorageVersion>::current_storage_version();
-				storage_version.put::<Self>();
-			}
-		}
-
 		// manual implementation of clone/eq/partialeq because using derive erroneously requires
 		// clone/eq/partialeq from T.
 		impl<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?> $crate::dispatch::Clone

frame/support/src/migrations.rs

@@ -18,7 +18,7 @@
 #[cfg(feature = "try-runtime")]
 use crate::storage::unhashed::contains_prefixed_key;
 use crate::{
-	traits::{GetStorageVersion, PalletInfoAccess},
+	traits::{GetStorageVersion, NoStorageVersionSet, PalletInfoAccess, StorageVersion},
 	weights::{RuntimeDbWeight, Weight},
 };
 use impl_trait_for_tuples::impl_for_tuples;
@@ -28,12 +28,38 @@ use sp_std::marker::PhantomData;
 #[cfg(feature = "try-runtime")]
 use sp_std::vec::Vec;
 
+/// Can store the current pallet version in storage.
+pub trait StoreCurrentStorageVersion<T: GetStorageVersion + PalletInfoAccess> {
+	/// Write the current storage version to the storage.
+	fn store_current_storage_version();
+}
+
+impl<T: GetStorageVersion<CurrentStorageVersion = StorageVersion> + PalletInfoAccess>
+	StoreCurrentStorageVersion<T> for StorageVersion
+{
+	fn store_current_storage_version() {
+		let version = <T as GetStorageVersion>::current_storage_version();
+		version.put::<T>();
+	}
+}
+
+impl<T: GetStorageVersion<CurrentStorageVersion = NoStorageVersionSet> + PalletInfoAccess>
+	StoreCurrentStorageVersion<T> for NoStorageVersionSet
+{
+	fn store_current_storage_version() {
+		StorageVersion::default().put::<T>();
+	}
+}
+
 /// Trait used by [`migrate_from_pallet_version_to_storage_version`] to do the actual migration.
 pub trait PalletVersionToStorageVersionHelper {
 	fn migrate(db_weight: &RuntimeDbWeight) -> Weight;
 }
 
-impl<T: GetStorageVersion + PalletInfoAccess> PalletVersionToStorageVersionHelper for T {
+impl<T: GetStorageVersion + PalletInfoAccess> PalletVersionToStorageVersionHelper for T
+where
+	T::CurrentStorageVersion: StoreCurrentStorageVersion<T>,
+{
 	fn migrate(db_weight: &RuntimeDbWeight) -> Weight {
 		const PALLET_VERSION_STORAGE_KEY_POSTFIX: &[u8] = b":__PALLET_VERSION__:";
 
@@ -43,8 +69,8 @@ impl<T: GetStorageVersion + PalletInfoAccess> PalletVersionToStorageVersionHelpe
 
 		sp_io::storage::clear(&pallet_version_key(<T as PalletInfoAccess>::name()));
 
-		let version = <T as GetStorageVersion>::current_storage_version();
-		version.put::<T>();
+		<T::CurrentStorageVersion as StoreCurrentStorageVersion<T>>::store_current_storage_version(
+		);
 
 		db_weight.writes(2)
 	}

frame/support/src/traits.rs

@@ -75,8 +75,8 @@ pub use randomness::Randomness;
 mod metadata;
 pub use metadata::{
 	CallMetadata, CrateVersion, GetCallIndex, GetCallMetadata, GetCallName, GetStorageVersion,
-	PalletInfo, PalletInfoAccess, PalletInfoData, PalletsInfoAccess, StorageVersion,
-	STORAGE_VERSION_STORAGE_KEY_POSTFIX,
+	NoStorageVersionSet, PalletInfo, PalletInfoAccess, PalletInfoData, PalletsInfoAccess,
+	StorageVersion, STORAGE_VERSION_STORAGE_KEY_POSTFIX,
 };
 
 mod hooks;

frame/support/src/traits/metadata.rs

@@ -232,6 +232,16 @@ impl PartialOrd<u16> for StorageVersion {
 	}
 }
 
+/// Special marker struct if no storage version is set for a pallet.
+///
+/// If you (the reader) end up here, it probably means that you tried to compare
+/// [`GetStorageVersion::on_chain_storage_version`] against
+/// [`GetStorageVersion::current_storage_version`]. This basically means that the
+/// [`storage_version`](crate::pallet_macros::storage_version) is missing in the pallet where the
+/// mentioned functions are being called.
+#[derive(Debug, Default)]
+pub struct NoStorageVersionSet;
+
 /// Provides information about the storage version of a pallet.
 ///
 /// It differentiates between current and on-chain storage version. Both should be only out of sync
@@ -244,8 +254,18 @@ impl PartialOrd<u16> for StorageVersion {
 ///
 /// It is required to update the on-chain storage version manually when a migration was applied.
 pub trait GetStorageVersion {
+	/// This will be filled out by the [`pallet`](crate::pallet) macro.
+	///
+	/// If the [`storage_version`](crate::pallet_macros::storage_version) attribute isn't given
+	/// this is set to [`NoStorageVersionSet`] to inform the user that the attribute is missing.
+	/// This should prevent that the user forgets to set a storage version when required. However,
+	/// this will only work when the user actually tries to call [`Self::current_storage_version`]
+	/// to compare it against the [`Self::on_chain_storage_version`]. If the attribute is given,
+	/// this will be set to [`StorageVersion`].
+	type CurrentStorageVersion;
+
 	/// Returns the current storage version as supported by the pallet.
-	fn current_storage_version() -> StorageVersion;
+	fn current_storage_version() -> Self::CurrentStorageVersion;
 	/// Returns the on-chain storage version of the pallet as stored in the storage.
 	fn on_chain_storage_version() -> StorageVersion;
 }

frame/support/test/Cargo.toml

@@ -26,9 +26,10 @@ sp-core = { version = "7.0.0", default-features = false, path = "../../../primit
 sp-std = { version = "5.0.0", default-features = false, path = "../../../primitives/std" }
 sp-version = { version = "5.0.0", default-features = false, path = "../../../primitives/version" }
 trybuild = { version = "1.0.74", features = [ "diff" ] }
-pretty_assertions = "1.2.1"
+pretty_assertions = "1.3.0"
 rustversion = "1.0.6"
 frame-system = { version = "4.0.0-dev", default-features = false, path = "../../system" }
+frame-executive = { version = "4.0.0-dev", default-features = false, path = "../../executive" }
 # The "std" feature for this pallet is never activated on purpose, in order to test construct_runtime error message
 test-pallet = { package = "frame-support-test-pallet", default-features = false, path = "pallet" }
 
@@ -39,6 +40,7 @@ std = [
 	"codec/std",
 	"scale-info/std",
 	"frame-benchmarking/std",
+	"frame-executive/std",
 	"frame-support/std",
 	"frame-system/std",
 	"sp-core/std",
@@ -50,7 +52,11 @@ std = [
 	"sp-version/std",
 	"sp-api/std",
 ]
-try-runtime = ["frame-support/try-runtime"]
+try-runtime = [
+	"frame-support/try-runtime",
+	"frame-system/try-runtime",
+	"frame-executive/try-runtime",
+]
 # WARNING:
 # Only CI runs with this feature enabled. This feature is for testing stuff related to the FRAME macros
 # in conjunction with rust features.