From c48034306bb8fdf4b6162807bccca0e69b6c785c Mon Sep 17 00:00:00 2001
From: Falco Hirschenberger <falco.hirschenberger@gmail.com>
Date: Fri, 17 Feb 2023 16:23:56 +0100
Subject: [PATCH] Rewrite set_code function

---
 Cargo.lock              |  1 +
 frame/system/Cargo.toml |  1 +
 frame/system/src/lib.rs | 53 +++++++++++++++++++++++------------------
 3 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b33f7627f230e..225bdbcdd0fbc 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2487,6 +2487,7 @@ dependencies = [
 name = "frame-system"
 version = "4.0.0-dev"
 dependencies = [
+ "cfg-if",
  "criterion",
  "frame-support",
  "log",
diff --git a/frame/system/Cargo.toml b/frame/system/Cargo.toml
index 5a6c89aae32ea..f9f4ceac4926d 100644
--- a/frame/system/Cargo.toml
+++ b/frame/system/Cargo.toml
@@ -13,6 +13,7 @@ readme = "README.md"
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
+cfg-if = "1.0"
 codec = { package = "parity-scale-codec", version = "3.2.2", default-features = false, features = ["derive"] }
 log = { version = "0.4.17", default-features = false }
 scale-info = { version = "2.1.1", default-features = false, features = ["derive"] }
diff --git a/frame/system/src/lib.rs b/frame/system/src/lib.rs
index af7e564dede16..d7e71d0818ed9 100644
--- a/frame/system/src/lib.rs
+++ b/frame/system/src/lib.rs
@@ -419,8 +419,13 @@ pub mod pallet {
 		///
 		/// ## Complexity
 		/// - `O(C)` where `C` length of `code`
+		/// - 1 storage write (codec `O(C)`).
+		/// - 1 digest item.
+		/// - 1 event.
+		/// The weight of this function is dependent on the runtime.
+		/// # </weight>
 		#[pallet::call_index(3)]
-		#[pallet::weight((T::BlockWeights::get().max_block, DispatchClass::Operational))]
+		#[pallet::weight((T::SystemWeightInfo::set_code(), DispatchClass::Operational))]
 		pub fn set_code_without_checks(
 			origin: OriginFor<T>,
 			code: Vec<u8>,
@@ -1612,29 +1617,31 @@ impl<T: Config> Pallet<T> {
 	/// it and extracting the runtime version of it. It checks that the runtime version
 	/// of the old and new runtime has the same spec name and that the spec version is increasing.
 	pub fn can_set_code(code: &[u8]) -> Result<(), sp_runtime::DispatchError> {
-		#[cfg_attr(feature = "runtime-benchmarks", allow(unused_variables))]
-		let (current_version, new_version) = (
-			T::Version::get(),
-			sp_io::misc::runtime_version(code)
-				.and_then(|v| RuntimeVersion::decode(&mut &v[..]).ok())
-				.ok_or(Error::<T>::FailedToExtractRuntimeVersion)?,
-		);
-
-		// Disable checks if we are benchmarking `set_code` to make it possible to set arbitrary
-		// runtimes without modification
-		#[cfg(not(feature = "runtime-benchmarks"))]
-		{
-			if new_version.spec_name != current_version.spec_name {
-				log::debug!("New: {new_version:?}, Current: {current_version:?}");
-				return Err(Error::<T>::InvalidSpecName.into())
-			}
-
-			if new_version.spec_version <= current_version.spec_version {
-				return Err(Error::<T>::SpecVersionNeedsToIncrease.into())
-			}
+		let current_version = T::Version::get();
+		let new_version = sp_io::misc::runtime_version(code)
+			.and_then(|v| RuntimeVersion::decode(&mut &v[..]).ok())
+			.ok_or(Error::<T>::FailedToExtractRuntimeVersion)?;
+
+		cfg_if::cfg_if! {
+			 if #[cfg(not(feature = "runtime-benchmarks"))] {
+				  if new_version.spec_name != current_version.spec_name {
+					  return Err(Error::<T>::InvalidSpecName.into())
+				   }
+
+				   if new_version.spec_version <= current_version.spec_version {
+						return Err(Error::<T>::SpecVersionNeedsToIncrease.into())
+				   }
+
+				   Ok(())
+			  } else {
+				   // Let's ensure the compiler doesn't optimize our fetching of the runtime version away.
+				   if new_version.spec_name != current_version.spec_name || new_version.spec_version != current_version.spec_version {
+					  Ok(())
+				   } else {
+					  Err(Error::<T>::InvalidSpecName.into())
+				   }
+			  }
 		}
-
-		Ok(())
 	}
 }