From 7614190e3b9aab12eccdd11e661aa8afc850d030 Mon Sep 17 00:00:00 2001
From: joe petrowski <25483142+joepetrowski@users.noreply.github.com>
Date: Thu, 2 May 2024 12:45:35 +0200
Subject: [PATCH] Use `authorize_upgrade` on Relay Chains (#29)

* authorize on relay

* fix tests
---
 src/build_upgrade.rs | 53 ++++++++++++++++++++++++++++++++++++++------
 src/tests.rs         |  3 +++
 src/types.rs         |  2 ++
 3 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/src/build_upgrade.rs b/src/build_upgrade.rs
index 04c9d93..319179c 100644
--- a/src/build_upgrade.rs
+++ b/src/build_upgrade.rs
@@ -14,6 +14,13 @@ pub(crate) struct UpgradeArgs {
 	#[clap(long = "only")]
 	pub(crate) only: bool,
 
+	/// Construct a call that will call `set_code` directly on the Relay Chain. This is generally
+	/// not recommended, as it involves submitting a large preimage (and therefore paying a large
+	/// fee). The default (false) uses `authorize_upgrade` instead, which only requires submitting
+	/// the hash. Anyone can then submit the actual runtime after it has been authorized.
+	#[clap(long = "set-relay-directly")]
+	pub(crate) set_relay_directly: bool,
+
 	/// The Fellowship release version. Should be semver and correspond to the release published.
 	#[clap(long = "relay-version")]
 	pub(crate) relay_version: Option<String>,
@@ -162,12 +169,16 @@ pub(crate) fn parse_inputs(prefs: UpgradeArgs) -> UpgradeDetails {
 		None => None,
 	};
 
+	let set_relay_directly = prefs.set_relay_directly;
+
+	// Get a version from one of the args. (This still feels dirty.)
 	let version = relay_version.clone().unwrap_or(asset_hub_version.unwrap_or(
 		bridge_hub_version.unwrap_or(encointer_version.unwrap_or(
 			collectives_version.unwrap_or(coretime_version.unwrap_or(String::from("no-version"))),
 		)),
 	));
 
+	// Set up a directory to store information fetched/written during this program.
 	let directory = format!("./upgrade-{}-{}/", &prefs.network, &version);
 	let output_file = if let Some(user_filename) = prefs.filename {
 		format!("{}{}", directory, user_filename)
@@ -177,7 +188,15 @@ pub(crate) fn parse_inputs(prefs: UpgradeArgs) -> UpgradeDetails {
 
 	make_version_directory(directory.as_str());
 
-	UpgradeDetails { relay, relay_version, networks, directory, output_file, additional }
+	UpgradeDetails {
+		relay,
+		relay_version,
+		networks,
+		directory,
+		output_file,
+		additional,
+		set_relay_directly,
+	}
 }
 
 // Create a directory into which to place runtime blobs and the final call data.
@@ -400,9 +419,19 @@ fn generate_relay_upgrade_call(upgrade_details: &UpgradeDetails) -> Option<CallI
 			let runtime_hash = blake2_256(&runtime);
 			println!("Kusama Relay Chain Runtime Hash: 0x{}", hex::encode(runtime_hash));
 
-			Some(CallInfo::from_runtime_call(NetworkRuntimeCall::Kusama(
-				KusamaRuntimeCall::System(SystemCall::set_code { code: runtime }),
-			)))
+			if !upgrade_details.set_relay_directly {
+				// authorize upgrade
+				Some(CallInfo::from_runtime_call(NetworkRuntimeCall::Kusama(
+					KusamaRuntimeCall::System(SystemCall::authorize_upgrade {
+						code_hash: H256(runtime_hash),
+					}),
+				)))
+			} else {
+				// set code directly
+				Some(CallInfo::from_runtime_call(NetworkRuntimeCall::Kusama(
+					KusamaRuntimeCall::System(SystemCall::set_code { code: runtime }),
+				)))
+			}
 		},
 		Network::Polkadot => {
 			use polkadot_relay::runtime_types::frame_system::pallet::Call as SystemCall;
@@ -415,9 +444,19 @@ fn generate_relay_upgrade_call(upgrade_details: &UpgradeDetails) -> Option<CallI
 			let runtime_hash = blake2_256(&runtime);
 			println!("Polkadot Relay Chain Runtime Hash: 0x{}", hex::encode(runtime_hash));
 
-			Some(CallInfo::from_runtime_call(NetworkRuntimeCall::Polkadot(
-				PolkadotRuntimeCall::System(SystemCall::set_code { code: runtime }),
-			)))
+			if !upgrade_details.set_relay_directly {
+				// authorize upgrade
+				Some(CallInfo::from_runtime_call(NetworkRuntimeCall::Polkadot(
+					PolkadotRuntimeCall::System(SystemCall::authorize_upgrade {
+						code_hash: H256(runtime_hash),
+					}),
+				)))
+			} else {
+				// set code directly
+				Some(CallInfo::from_runtime_call(NetworkRuntimeCall::Polkadot(
+					PolkadotRuntimeCall::System(SystemCall::set_code { code: runtime }),
+				)))
+			}
 		},
 		_ => panic!("Not a Relay Chain"),
 	}
diff --git a/src/tests.rs b/src/tests.rs
index 783af01..ddc173a 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -122,6 +122,7 @@ fn upgrade_args_for_only_relay() -> UpgradeArgs {
 	UpgradeArgs {
 		network: String::from("polkadot"),
 		only: true,
+		set_relay_directly: true,
 		relay_version: Some(String::from("v1.2.0")),
 		asset_hub: None,
 		bridge_hub: None,
@@ -137,6 +138,7 @@ fn upgrade_args_for_only_asset_hub() -> UpgradeArgs {
 	UpgradeArgs {
 		network: String::from("polkadot"),
 		only: true,
+		set_relay_directly: true,
 		relay_version: None,
 		asset_hub: Some(String::from("v1.2.0")),
 		bridge_hub: None,
@@ -152,6 +154,7 @@ fn upgrade_args_for_all() -> UpgradeArgs {
 	UpgradeArgs {
 		network: String::from("polkadot"),
 		only: false,
+		set_relay_directly: true,
 		relay_version: Some(String::from("v1.2.0")),
 		asset_hub: None,
 		bridge_hub: None,
diff --git a/src/types.rs b/src/types.rs
index 1de6763..8bf8457 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -140,6 +140,8 @@ pub(super) struct UpgradeDetails {
 	pub(super) output_file: String,
 	// An additional call to be enacted in the same batch as the system upgrade.
 	pub(super) additional: Option<CallInfo>,
+	// Call `set_code` directly on the Relay Chain.
+	pub(super) set_relay_directly: bool,
 }
 
 // A network and the version to which it will upgrade.