try-runtime: add cli option --export-proof (#12539)

* try-runtime: add cli option --export-proof

* extract proof in raw json format

* fix build

* fix(try-runtime execute-block): wrong block parsing

* fmt

* apply suggestions

* Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Anton <anton.kalyaev@gmail.com>

* Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Anton <anton.kalyaev@gmail.com>

* Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Anton <anton.kalyaev@gmail.com>

* Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Anton <anton.kalyaev@gmail.com>

* split off external dependencies

* fmt

* fix try-runtime compilation

Co-authored-by: Anton <anton.kalyaev@gmail.com>

utils/frame/try-runtime/cli/Cargo.toml

@@ -30,11 +30,12 @@ sp-weights = { version = "4.0.0", path = "../../../../primitives/weights" }
 frame-try-runtime = { optional = true, path = "../../../../frame/try-runtime" }
 substrate-rpc-client = { path = "../../rpc/client" }
 
-parity-scale-codec = "3.0.0"
-hex = "0.4.3"
 clap = { version = "4.0.9", features = ["derive"] }
+hex = { version = "0.4.3", default-features = false }
 log = "0.4.17"
+parity-scale-codec = "3.0.0"
 serde = "1.0.136"
+serde_json = "1.0.85"
 zstd = { version = "0.11.2", default-features = false }
 
 [dev-dependencies]

utils/frame/try-runtime/cli/src/commands/execute_block.rs

@@ -134,6 +134,7 @@ where
 		"TryRuntime_execute_block",
 		&payload,
 		full_extensions(),
+		shared.export_proof,
 	)?;
 
 	Ok(())

utils/frame/try-runtime/cli/src/commands/follow_chain.rs

@@ -150,6 +150,10 @@ where
 			"TryRuntime_execute_block",
 			(block, command.state_root_check, command.try_state.clone()).encode().as_ref(),
 			full_extensions(),
+			shared
+				.export_proof
+				.as_ref()
+				.map(|path| path.as_path().join(&format!("{}.json", number))),
 		);
 
 		if let Err(why) = result {

utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs

@@ -59,6 +59,7 @@ where
 		"TryRuntime_on_runtime_upgrade",
 		command.checks.encode().as_ref(),
 		Default::default(), // we don't really need any extensions here.
+		shared.export_proof,
 	)?;
 
 	let (weight, total_weight) = <(Weight, Weight) as Decode>::decode(&mut &*encoded_result)

utils/frame/try-runtime/cli/src/lib.rs

@@ -523,6 +523,12 @@ pub struct SharedParams {
 	#[arg(long)]
 	pub heap_pages: Option<u64>,
 
+	/// Path to a file to export the storage proof into (as a JSON).
+	/// If several blocks are executed, the path is interpreted as a folder
+	/// where one file per block will be written (named `{block_number}-{block_hash}`).
+	#[clap(long)]
+	pub export_proof: Option<PathBuf>,
+
 	/// Overwrite the `state_version`.
 	///
 	/// Otherwise `remote-externalities` will automatically set the correct state version.
@@ -863,6 +869,7 @@ pub(crate) fn state_machine_call_with_proof<Block: BlockT, HostFns: HostFunction
 	method: &'static str,
 	data: &[u8],
 	extensions: Extensions,
+	maybe_export_proof: Option<PathBuf>,
 ) -> sc_cli::Result<(OverlayedChanges, Vec<u8>)> {
 	use parity_scale_codec::Encode;
 
@@ -891,6 +898,32 @@ pub(crate) fn state_machine_call_with_proof<Block: BlockT, HostFns: HostFunction
 	let proof = proving_backend
 		.extract_proof()
 		.expect("A recorder was set and thus, a storage proof can be extracted; qed");
+
+	if let Some(path) = maybe_export_proof {
+		let mut file = std::fs::File::create(&path).map_err(|e| {
+			log::error!(
+				target: LOG_TARGET,
+				"Failed to create file {}: {:?}",
+				path.to_string_lossy(),
+				e
+			);
+			e
+		})?;
+
+		log::info!(target: LOG_TARGET, "Writing storage proof to {}", path.to_string_lossy());
+
+		use std::io::Write as _;
+		file.write_all(storage_proof_to_raw_json(&proof).as_bytes()).map_err(|e| {
+			log::error!(
+				target: LOG_TARGET,
+				"Failed to write storage proof to {}: {:?}",
+				path.to_string_lossy(),
+				e
+			);
+			e
+		})?;
+	}
+
 	let proof_size = proof.encoded_size();
 	let compact_proof = proof
 		.clone()
@@ -951,3 +984,21 @@ pub(crate) fn rpc_err_handler(error: impl Debug) -> &'static str {
 	log::error!(target: LOG_TARGET, "rpc error: {:?}", error);
 	"rpc error."
 }
+
+/// Converts a [`sp_state_machine::StorageProof`] into a JSON string.
+fn storage_proof_to_raw_json(storage_proof: &sp_state_machine::StorageProof) -> String {
+	serde_json::Value::Object(
+		storage_proof
+			.to_memory_db::<sp_runtime::traits::BlakeTwo256>()
+			.drain()
+			.iter()
+			.map(|(key, (value, _n))| {
+				(
+					format!("0x{}", hex::encode(key.as_bytes())),
+					serde_json::Value::String(format!("0x{}", hex::encode(value))),
+				)
+			})
+			.collect(),
+	)
+	.to_string()
+}