Add Additional Call to Upgrade (#19)

* allow additional call data to be appended to build upgrade

* readme

README.md

@@ -74,6 +74,7 @@ Options:
       --bridge-hub <BRIDGE_HUB>        Optional. The runtime version of Bridge Hub to which to upgrade. If not provided, it will use the Relay Chain's version
       --collectives <COLLECTIVES>      Optional. The runtime version of Collectives to which to upgrade. If not provided, it will use the Relay Chain's version
       --filename <FILENAME>            Name of the file to which to write the output. If not provided, a default will be constructed
+      --additional <ADDITIONAL>        Some additional call that you want executed on the Relay Chain along with the upgrade
   -h, --help                           Print help
 ```
 

src/build_upgrade.rs

@@ -33,6 +33,10 @@ pub(crate) struct UpgradeArgs {
 	/// constructed.
 	#[clap(long = "filename")]
 	filename: Option<String>,
+
+	/// Some additional call that you want executed on the Relay Chain along with the upgrade.
+	#[clap(long = "additional")]
+	additional: Option<String>,
 }
 
 // The sub-command's "main" function.
@@ -117,6 +121,20 @@ fn parse_inputs(prefs: UpgradeArgs) -> UpgradeDetails {
 		_ => panic!("`network` must be `polkadot` or `kusama`"),
 	};
 
+	let additional = match prefs.additional {
+		Some(c) => {
+			let additional_bytes = get_proposal_bytes(c.clone());
+			match relay.network {
+				Network::Polkadot =>
+					Some(CallInfo::from_bytes(&additional_bytes, Network::Polkadot)),
+				Network::Kusama => Some(CallInfo::from_bytes(&additional_bytes, Network::Kusama)),
+				// for now, only support additional on the relay chain
+				_ => panic!("`network` must be `polkadot` or `kusama`"),
+			}
+		},
+		None => None,
+	};
+
 	let directory = format!("./upgrade-{}-{}/", &prefs.network, &relay_version);
 	let output_file = if let Some(user_filename) = prefs.filename {
 		format!("{}{}", directory, user_filename)
@@ -126,7 +144,7 @@ fn parse_inputs(prefs: UpgradeArgs) -> UpgradeDetails {
 
 	make_version_directory(directory.as_str());
 
-	UpgradeDetails { relay, networks, directory, output_file }
+	UpgradeDetails { relay, networks, directory, output_file, additional }
 }
 
 // Create a directory into which to place runtime blobs and the final call data.
@@ -343,14 +361,21 @@ async fn construct_batch(
 ) -> CallInfo {
 	println!("\nBatching calls.");
 	match upgrade_details.relay.network {
-		Network::Kusama => construct_kusama_batch(relay_call, para_calls).await,
-		Network::Polkadot => construct_polkadot_batch(relay_call, para_calls).await,
+		Network::Kusama =>
+			construct_kusama_batch(relay_call, para_calls, upgrade_details.additional.clone()).await,
+		Network::Polkadot =>
+			construct_polkadot_batch(relay_call, para_calls, upgrade_details.additional.clone())
+				.await,
 		_ => panic!("Not a Relay Chain"),
 	}
 }
 
 // Construct the batch needed on Kusama.
-async fn construct_kusama_batch(relay_call: CallInfo, para_calls: Vec<CallInfo>) -> CallInfo {
+async fn construct_kusama_batch(
+	relay_call: CallInfo,
+	para_calls: Vec<CallInfo>,
+	additional: Option<CallInfo>,
+) -> CallInfo {
 	use kusama_relay::runtime_types::pallet_utility::pallet::Call as UtilityCall;
 
 	let mut batch_calls = Vec::new();
@@ -371,14 +396,22 @@ async fn construct_kusama_batch(relay_call: CallInfo, para_calls: Vec<CallInfo>)
 			},
 		}
 	}
+	if let Some(a) = additional {
+		batch_calls.push(a.get_kusama_call().expect("kusama call"))
+	}
+	// Relay set code goes last
 	batch_calls.push(relay_call.get_kusama_call().expect("kusama call"));
 	CallInfo::from_runtime_call(NetworkRuntimeCall::Kusama(KusamaRuntimeCall::Utility(
 		UtilityCall::force_batch { calls: batch_calls },
 	)))
 }
 
 // Construct the batch needed on Polkadot.
-async fn construct_polkadot_batch(relay_call: CallInfo, para_calls: Vec<CallInfo>) -> CallInfo {
+async fn construct_polkadot_batch(
+	relay_call: CallInfo,
+	para_calls: Vec<CallInfo>,
+	additional: Option<CallInfo>,
+) -> CallInfo {
 	use polkadot_relay::runtime_types::pallet_utility::pallet::Call as UtilityCall;
 
 	let mut batch_calls = Vec::new();
@@ -401,6 +434,10 @@ async fn construct_polkadot_batch(relay_call: CallInfo, para_calls: Vec<CallInfo
 			},
 		}
 	}
+	if let Some(a) = additional {
+		batch_calls.push(a.get_polkadot_call().expect("polkadot call"))
+	}
+	// Relay set code goes last
 	batch_calls.push(relay_call.get_polkadot_call().expect("polkadot call"));
 	CallInfo::from_runtime_call(NetworkRuntimeCall::Polkadot(PolkadotRuntimeCall::Utility(
 		UtilityCall::force_batch { calls: batch_calls },

src/functions.rs

@@ -0,0 +1,15 @@
+use std::fs;
+
+// Check what the user entered for the proposal. If it is just call data, return it back. Otherwise,
+// we expect a path to a file that contains the call data. Read that in and return it.
+pub(crate) fn get_proposal_bytes(proposal: String) -> Vec<u8> {
+	let proposal = proposal.as_str();
+	if proposal.starts_with("0x") {
+		// This is just call data
+		return hex::decode(proposal.trim_start_matches("0x")).expect("Valid proposal")
+	} else {
+		// This is a file path
+		let contents = fs::read_to_string(proposal).expect("Should give a valid file path");
+		return hex::decode(contents.as_str().trim_start_matches("0x")).expect("Valid proposal")
+	}
+}

src/main.rs

@@ -1,5 +1,7 @@
 mod types;
 use crate::types::*;
+mod functions;
+use crate::functions::*;
 mod build_upgrade;
 use crate::build_upgrade::{build_upgrade, UpgradeArgs};
 mod submit_referendum;

src/submit_referendum.rs

@@ -649,20 +649,6 @@ fn handle_batch_of_calls(output: &Output, batch: Vec<NetworkRuntimeCall>) {
 	}
 }
 
-// Check what the user entered for the proposal. If it is just call data, return it back. Otherwise,
-// we expect a path to a file that contains the call data. Read that in and return it.
-pub(crate) fn get_proposal_bytes(proposal: String) -> Vec<u8> {
-	let proposal = proposal.as_str();
-	if proposal.starts_with("0x") {
-		// This is just call data
-		return hex::decode(proposal.trim_start_matches("0x")).expect("Valid proposal")
-	} else {
-		// This is a file path
-		let contents = fs::read_to_string(proposal).expect("Should give a valid file path");
-		return hex::decode(contents.as_str().trim_start_matches("0x")).expect("Valid proposal")
-	}
-}
-
 // Format the data to print to console.
 fn print_output(output: &Output, network_call: &NetworkRuntimeCall) {
 	match network_call {

src/tests.rs

@@ -1,8 +1,8 @@
+use crate::get_proposal_bytes;
 use crate::polkadot_relay::runtime_types::frame_system::pallet::Call as PolkadotRelaySystemCall;
 use crate::{
-	submit_referendum::{generate_calls, get_proposal_bytes},
-	CallInfo, CallOrHash, KusamaOpenGovOrigin, Network, NetworkRuntimeCall, PolkadotOpenGovOrigin,
-	PolkadotRuntimeCall, ProposalDetails, Weight,
+	submit_referendum::generate_calls, CallInfo, CallOrHash, KusamaOpenGovOrigin, Network,
+	NetworkRuntimeCall, PolkadotOpenGovOrigin, PolkadotRuntimeCall, ProposalDetails, Weight,
 };
 
 fn polkadot_whitelist_remark_user_input() -> ProposalDetails {

src/types.rs

@@ -46,6 +46,7 @@ pub(super) use polkadot_collectives::runtime_types::{
 pub mod polkadot_bridge_hub {}
 pub(super) use polkadot_bridge_hub::runtime_types::bridge_hub_polkadot_runtime::RuntimeCall as PolkadotBridgeHubRuntimeCall;
 
+#[derive(Clone)]
 pub(super) enum Network {
 	Kusama,
 	KusamaAssetHub,
@@ -103,6 +104,8 @@ pub(super) struct UpgradeDetails {
 	pub(super) directory: String,
 	// The filename of the output.
 	pub(super) output_file: String,
+	// An additional call to be enacted in the same batch as the system upgrade.
+	pub(super) additional: Option<CallInfo>,
 }
 
 // A network and the version to which it will upgrade.
@@ -153,6 +156,7 @@ pub(super) enum CallOrHash {
 }
 
 // All the info associated with a call in the forms you may need it in.
+#[derive(Clone)]
 pub(super) struct CallInfo {
 	pub(super) network: Network,
 	pub(super) encoded: Vec<u8>,