From 9683eb284992c332ecef12e03a36f78971f13bdc Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Wed, 9 Jun 2021 16:48:15 -0700
Subject: [PATCH 01/13] Refactor remote_externalities::rpc_api

---
 .../frame/remote-externalities/src/rpc_api.rs | 69 +++++++++-----
 utils/frame/try-runtime/cli/src/lib.rs        | 91 +++++++++++++++++++
 2 files changed, 139 insertions(+), 21 deletions(-)

diff --git a/utils/frame/remote-externalities/src/rpc_api.rs b/utils/frame/remote-externalities/src/rpc_api.rs
index e7fd021bac4a8..fbd438744aa7b 100644
--- a/utils/frame/remote-externalities/src/rpc_api.rs
+++ b/utils/frame/remote-externalities/src/rpc_api.rs
@@ -18,36 +18,63 @@
 //! WS RPC API for one off RPC calls to a substrate node.
 // TODO: Consolidate one off RPC calls https://github.com/paritytech/substrate/issues/8988
 
-use super::*;
+use sp_runtime::{generic::SignedBlock, traits::Block as BlockT};
+use jsonrpsee_ws_client::{WsClientBuilder, WsClient, v2::params::JsonRpcParams, traits::Client};
 
 /// Get the header of the block identified by `at`
-pub async fn get_header<B: BlockT, S: AsRef<str>>(from: S, at: B::Hash) -> Result<B::Header, String>
+pub async fn get_header<Block, S>(from: S, at: Block::Hash) -> Result<Block::Header, String>
 where
-	B::Header: serde::de::DeserializeOwned,
+	Block: BlockT,
+	Block::Header: serde::de::DeserializeOwned,
+	S: AsRef<str>,
 {
-	use jsonrpsee_ws_client::traits::Client;
-	let at = serde_json::to_value(at)
-		.map_err(|e| format!("Block hash could not be converted to JSON due to {:?}", e))?;
-	let params = vec![at];
-	let client = WsClientBuilder::default()
-		.max_request_body_size(u32::MAX)
-		.build(from.as_ref())
-		.await
-		.map_err(|e| format!("`WsClientBuilder` failed to build do to {:?}", e))?;
-	client.request::<B::Header>("chain_getHeader", JsonRpcParams::Array(params))
+	let params = vec![hash_to_json::<Block>(at)?];
+	let client = build_client(from).await?;
+
+	client.request::<Block::Header>("chain_getHeader", JsonRpcParams::Array(params))
 		.await
 		.map_err(|e| format!("chain_getHeader request failed due to {:?}", e))
 }
 
 /// Get the finalized head
-pub async fn get_finalized_head<B: BlockT, S: AsRef<str>>(from: S) -> Result<B::Hash, String> {
-	use jsonrpsee_ws_client::traits::Client;
-	let client = WsClientBuilder::default()
-		.max_request_body_size(u32::MAX)
-		.build(from.as_ref())
-		.await
-		.map_err(|e| format!("`WsClientBuilder` failed to build do to {:?}", e))?;
-	client.request::<B::Hash>("chain_getFinalizedHead", JsonRpcParams::NoParams)
+pub async fn get_finalized_head<Block, S>(from: S) -> Result<Block::Hash, String>
+where
+	Block: BlockT,
+	S: AsRef<str>,
+{
+	let client = build_client(from).await?;
+
+	client.request::<Block::Hash>("chain_getFinalizedHead", JsonRpcParams::NoParams)
 		.await
 		.map_err(|e| format!("chain_getFinalizedHead request failed due to {:?}", e))
 }
+
+/// Get the signed block identified by `at`.
+pub async fn get_block<Block, S>(from: S, at: Block::Hash) -> Result<Block, String>
+where
+	Block: BlockT + serde::de::DeserializeOwned,
+	S: AsRef<str>,
+{
+	let params = vec![hash_to_json::<Block>(at)?];
+	let client = build_client(from).await?;
+	let signed_block = client.request::<SignedBlock<Block>>("chain_getBlock", JsonRpcParams::Array(params))
+		.await
+		.map_err(|e| format!("chain_getBlock request failed due to {:?}", e))?;
+
+	Ok(signed_block.block)
+}
+
+/// Convert a block hash to a serde json value.
+fn hash_to_json<Block: BlockT>(hash: Block::Hash) -> Result<serde_json::Value, String> {
+	serde_json::to_value(hash)
+		.map_err(|e| format!("Block hash could not be converted to JSON due to {:?}", e))
+}
+
+/// Build a website client that connects to `from`.
+async fn build_client<S: AsRef<str>>(from: S) -> Result<WsClient, String> {
+	 WsClientBuilder::default()
+		.max_request_body_size(u32::MAX)
+		.build(from.as_ref())
+		.await
+		.map_err(|e| format!("`WsClientBuilder` failed to build do to {:?}", e))
+}
\ No newline at end of file
diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index dc4cb7cd33dbd..d6b633b04ad1c 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -328,6 +328,97 @@ where
 	Ok(())
 }
 
+// async fn execute_block<Block, ExecDispatch>(
+// 	shared: SharedParams,
+// 	command: OffchainWorkerCmd,
+// 	config: Configuration,
+// )-> sc_cli::Result<()>
+// where
+// 	Block: BlockT,
+// 	Block::Hash: FromStr,
+// 	Block::Header: serde::de::DeserializeOwned,
+// 	<Block::Hash as FromStr>::Err: Debug,
+// 	NumberFor<Block>: FromStr,
+// 	<NumberFor<Block> as FromStr>::Err: Debug,
+// 	ExecDispatch: NativeExecutionDispatch + 'static,
+// {
+// 	let wasm_method = shared.wasm_method;
+// 	let execution = shared.execution;
+// 	let heap_pages = if shared.heap_pages.is_some() {
+// 		shared.heap_pages
+// 	} else {
+// 		config.default_heap_pages
+// 	};
+
+// 	let mut changes = Default::default();
+// 	let max_runtime_instances = config.max_runtime_instances;
+// 	let executor = NativeExecutor::<ExecDispatch>::new(
+// 		wasm_method.into(),
+// 		heap_pages,
+// 		max_runtime_instances,
+// 	);
+
+// 	let (mode, url) = match command.state {
+// 		State::Live {
+// 			url,
+// 			snapshot_path,
+// 			block_at,
+// 			modules
+// 		} => {
+// 			let online_config = OnlineConfig {
+// 				transport: url.to_owned().into(),
+// 				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
+// 				modules: modules.to_owned().unwrap_or_default(),
+// 				at: block_at.as_ref()
+// 					.map(|b| b.parse().map_err(|e| format!("Could not parse hash: {:?}", e))).transpose()?,
+// 				..Default::default()
+// 			};
+
+// 			(Mode::Online(online_config), url)
+// 		},
+// 		State::Snap { snapshot_path } => {
+// 			// TODO This is a temporary hack; the url is used just to get the header. We should try
+// 			// and get the header out of state, OR use an arbitrary header if thats ok, OR allow
+// 			// the user to feed in a header via file.
+// 			// https://github.com/paritytech/substrate/issues/9027
+// 			// This assumes you have a node running on local host default
+// 			let url = "ws://127.0.0.1:9944".to_string();
+// 			let mode = Mode::Offline(OfflineConfig {
+// 				state_snapshot: SnapshotConfig::new(snapshot_path),
+// 			});
+
+// 			(mode, url)
+// 		}
+// 	};
+
+// 	let ext = {
+// 		let builder = Builder::<Block>::new().mode(mode);
+// 		let mut ext = if command.overwrite_code {
+// 			let (code_key, code) = extract_code(config.chain_spec)?;
+// 			builder.inject(&[(code_key, code)]).build().await?
+// 		} else {
+// 			builder.build().await?
+// 		};
+
+// 		// register externality extensions in order to provide host interface for OCW to the runtime.
+// 		let (offchain, _offchain_state) = TestOffchainExt::new();
+// 		let (pool, _pool_state) = TestTransactionPoolExt::new();
+// 		ext.register_extension(OffchainDbExt::new(offchain.clone()));
+// 		ext.register_extension(OffchainWorkerExt::new(offchain));
+// 		ext.register_extension(KeystoreExt(Arc::new(KeyStore::new())));
+// 		ext.register_extension(TransactionPoolExt::new(pool));
+
+// 		ext
+// 	}
+
+// 	let header_hash: Block::Hash = command.header_at
+// 		.parse()
+// 		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
+// 	// let Block: Block = rpc_api::get_block::<Block, _>(url, header);
+
+// 	Ok(())
+// }
+
 impl TryRuntimeCmd {
 	pub async fn run<Block, ExecDispatch>(&self, config: Configuration) -> sc_cli::Result<()>
 	where

From 22b96f9a4bdb138f79c387ad78b71b6c11c46cf9 Mon Sep 17 00:00:00 2001
From: Zeke Mostov <32168567+emostov@users.noreply.github.com>
Date: Thu, 10 Jun 2021 17:44:02 -0700
Subject: [PATCH 02/13] try-runtime-cli: Adde `execute-block` subcommand

---
 utils/frame/remote-externalities/src/lib.rs   |   2 +-
 .../frame/remote-externalities/src/rpc_api.rs |  10 +-
 utils/frame/try-runtime/cli/src/lib.rs        | 238 ++++++++++--------
 3 files changed, 141 insertions(+), 109 deletions(-)

diff --git a/utils/frame/remote-externalities/src/lib.rs b/utils/frame/remote-externalities/src/lib.rs
index a77650d042125..0214bf263892a 100644
--- a/utils/frame/remote-externalities/src/lib.rs
+++ b/utils/frame/remote-externalities/src/lib.rs
@@ -107,7 +107,7 @@ impl From<String> for Transport {
 /// A state snapshot config may be present and will be written to in that case.
 #[derive(Clone)]
 pub struct OnlineConfig<B: BlockT> {
-	/// The block number at which to connect. Will be latest finalized head if not provided.
+	/// The block hash at which to get the runtime state. Will be latest finalized head if not provided.
 	pub at: Option<B::Hash>,
 	/// An optional state snapshot file to WRITE to, not for reading. Not written if set to `None`.
 	pub state_snapshot: Option<SnapshotConfig>,
diff --git a/utils/frame/remote-externalities/src/rpc_api.rs b/utils/frame/remote-externalities/src/rpc_api.rs
index fbd438744aa7b..a5ffa8410f1b8 100644
--- a/utils/frame/remote-externalities/src/rpc_api.rs
+++ b/utils/frame/remote-externalities/src/rpc_api.rs
@@ -18,7 +18,7 @@
 //! WS RPC API for one off RPC calls to a substrate node.
 // TODO: Consolidate one off RPC calls https://github.com/paritytech/substrate/issues/8988
 
-use sp_runtime::{generic::SignedBlock, traits::Block as BlockT};
+use sp_runtime::{generic::SignedBlock, traits::{Block as BlockT, Header as HeaderT}};
 use jsonrpsee_ws_client::{WsClientBuilder, WsClient, v2::params::JsonRpcParams, traits::Client};
 
 /// Get the header of the block identified by `at`
@@ -52,12 +52,14 @@ where
 /// Get the signed block identified by `at`.
 pub async fn get_block<Block, S>(from: S, at: Block::Hash) -> Result<Block, String>
 where
-	Block: BlockT + serde::de::DeserializeOwned,
 	S: AsRef<str>,
+	Block: BlockT + serde::de::DeserializeOwned,
+	Block::Header: HeaderT,
 {
 	let params = vec![hash_to_json::<Block>(at)?];
 	let client = build_client(from).await?;
-	let signed_block = client.request::<SignedBlock<Block>>("chain_getBlock", JsonRpcParams::Array(params))
+	let signed_block = client
+		.request::<SignedBlock<Block>>("chain_getBlock", JsonRpcParams::Array(params))
 		.await
 		.map_err(|e| format!("chain_getBlock request failed due to {:?}", e))?;
 
@@ -77,4 +79,4 @@ async fn build_client<S: AsRef<str>>(from: S) -> Result<WsClient, String> {
 		.build(from.as_ref())
 		.await
 		.map_err(|e| format!("`WsClientBuilder` failed to build do to {:?}", e))
-}
\ No newline at end of file
+}
diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index d6b633b04ad1c..3d81a6c671aca 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -25,7 +25,7 @@ use sc_executor::NativeExecutor;
 use sc_service::NativeExecutionDispatch;
 use sc_chain_spec::ChainSpec;
 use sp_state_machine::StateMachine;
-use sp_runtime::traits::{Block as BlockT, NumberFor};
+use sp_runtime::traits::{Block as BlockT, NumberFor, Header as HeaderT};
 use sp_core::{
 	offchain::{
 		OffchainWorkerExt, OffchainDbExt, TransactionPoolExt,
@@ -45,6 +45,8 @@ pub enum Command {
 	OnRuntimeUpgrade(OnRuntimeUpgradeCmd),
 	/// Execute "OffchainWorkerApi_offchain_worker" against the given runtime state.
 	OffchainWorker(OffchainWorkerCmd),
+	/// Execute "Core_execute_block" using the given block and the runtime state of the parent block.
+	ExecuteBlock(ExecuteBlockCmd),
 }
 
 #[derive(Debug, Clone, structopt::StructOpt)]
@@ -68,6 +70,16 @@ pub struct OffchainWorkerCmd {
 	pub overwrite_code: bool,
 }
 
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct ExecuteBlockCmd {
+	/// Hash of the block whose header to use to execute the offchain worker.
+	/// NOTE: `--header-at` is a required option for `snap`, but not `live`
+	#[structopt(short, long, multiple = false, parse(try_from_str = parse::hash))]
+	pub header_at: Option<String>,
+	#[structopt(subcommand)]
+	pub state: State,
+}
+
 #[derive(Debug, Clone, structopt::StructOpt)]
 pub struct SharedParams {
 	/// The shared parameters
@@ -114,12 +126,18 @@ pub struct TryRuntimeCmd {
 /// The source of runtime state to try operations against.
 #[derive(Debug, Clone, structopt::StructOpt)]
 pub enum State {
-	/// Use a state snapshot as the source of runtime state. NOTE: for the offchain-worker command this
-	/// is only partially supported at the moment and you must have a relevant archive node exposed on
-	/// localhost:9944 in order to query the block header.
-	// TODO https://github.com/paritytech/substrate/issues/9027
+	/// Use a state snapshot as the source of runtime state. NOTE: for the offchain-worker and
+	/// execute-block command this is only partially supported and requires a archive node url.
 	Snap {
+		#[structopt(short, long)]
 		snapshot_path: PathBuf,
+		/// The url to connect to for fetching data not easily retrievable from the snapshot.
+		/// NOTE: This is not necessary for `on-runtime-upgrade`
+		// TODO This is a temporary hack; the url is used just to get the header/block. We should try
+		// and get that out of state, OR allow the user to feed in a header/block via file.
+		// https://github.com/paritytech/substrate/issues/9027
+		#[structopt(default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
+		url: String,
 	},
 
 	/// Use a live chain as the source of runtime state.
@@ -174,7 +192,7 @@ where
 
 	let ext = {
 		let builder = match command.state {
-			State::Snap { snapshot_path } => {
+			State::Snap { snapshot_path, .. } => {
 				Builder::<Block>::new().mode(Mode::Offline(OfflineConfig {
 					state_snapshot: SnapshotConfig::new(snapshot_path),
 				}))
@@ -273,13 +291,7 @@ where
 
 				(Mode::Online(online_config), url)
 			},
-			State::Snap { snapshot_path } => {
-				// TODO This is a temporary hack; the url is used just to get the header. We should try
-				// and get the header out of state, OR use an arbitrary header if thats ok, OR allow
-				// the user to feed in a header via file.
-				// https://github.com/paritytech/substrate/issues/9027
-				// This assumes you have a node running on local host default
-				let url = "ws://127.0.0.1:9944".to_string();
+			State::Snap { snapshot_path, url } => {
 				let mode = Mode::Offline(OfflineConfig {
 					state_snapshot: SnapshotConfig::new(snapshot_path),
 				});
@@ -328,101 +340,116 @@ where
 	Ok(())
 }
 
-// async fn execute_block<Block, ExecDispatch>(
-// 	shared: SharedParams,
-// 	command: OffchainWorkerCmd,
-// 	config: Configuration,
-// )-> sc_cli::Result<()>
-// where
-// 	Block: BlockT,
-// 	Block::Hash: FromStr,
-// 	Block::Header: serde::de::DeserializeOwned,
-// 	<Block::Hash as FromStr>::Err: Debug,
-// 	NumberFor<Block>: FromStr,
-// 	<NumberFor<Block> as FromStr>::Err: Debug,
-// 	ExecDispatch: NativeExecutionDispatch + 'static,
-// {
-// 	let wasm_method = shared.wasm_method;
-// 	let execution = shared.execution;
-// 	let heap_pages = if shared.heap_pages.is_some() {
-// 		shared.heap_pages
-// 	} else {
-// 		config.default_heap_pages
-// 	};
-
-// 	let mut changes = Default::default();
-// 	let max_runtime_instances = config.max_runtime_instances;
-// 	let executor = NativeExecutor::<ExecDispatch>::new(
-// 		wasm_method.into(),
-// 		heap_pages,
-// 		max_runtime_instances,
-// 	);
-
-// 	let (mode, url) = match command.state {
-// 		State::Live {
-// 			url,
-// 			snapshot_path,
-// 			block_at,
-// 			modules
-// 		} => {
-// 			let online_config = OnlineConfig {
-// 				transport: url.to_owned().into(),
-// 				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
-// 				modules: modules.to_owned().unwrap_or_default(),
-// 				at: block_at.as_ref()
-// 					.map(|b| b.parse().map_err(|e| format!("Could not parse hash: {:?}", e))).transpose()?,
-// 				..Default::default()
-// 			};
-
-// 			(Mode::Online(online_config), url)
-// 		},
-// 		State::Snap { snapshot_path } => {
-// 			// TODO This is a temporary hack; the url is used just to get the header. We should try
-// 			// and get the header out of state, OR use an arbitrary header if thats ok, OR allow
-// 			// the user to feed in a header via file.
-// 			// https://github.com/paritytech/substrate/issues/9027
-// 			// This assumes you have a node running on local host default
-// 			let url = "ws://127.0.0.1:9944".to_string();
-// 			let mode = Mode::Offline(OfflineConfig {
-// 				state_snapshot: SnapshotConfig::new(snapshot_path),
-// 			});
-
-// 			(mode, url)
-// 		}
-// 	};
-
-// 	let ext = {
-// 		let builder = Builder::<Block>::new().mode(mode);
-// 		let mut ext = if command.overwrite_code {
-// 			let (code_key, code) = extract_code(config.chain_spec)?;
-// 			builder.inject(&[(code_key, code)]).build().await?
-// 		} else {
-// 			builder.build().await?
-// 		};
-
-// 		// register externality extensions in order to provide host interface for OCW to the runtime.
-// 		let (offchain, _offchain_state) = TestOffchainExt::new();
-// 		let (pool, _pool_state) = TestTransactionPoolExt::new();
-// 		ext.register_extension(OffchainDbExt::new(offchain.clone()));
-// 		ext.register_extension(OffchainWorkerExt::new(offchain));
-// 		ext.register_extension(KeystoreExt(Arc::new(KeyStore::new())));
-// 		ext.register_extension(TransactionPoolExt::new(pool));
-
-// 		ext
-// 	}
-
-// 	let header_hash: Block::Hash = command.header_at
-// 		.parse()
-// 		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
-// 	// let Block: Block = rpc_api::get_block::<Block, _>(url, header);
-
-// 	Ok(())
-// }
+async fn execute_block<Block, ExecDispatch>(
+	shared: SharedParams,
+	command: ExecuteBlockCmd,
+	config: Configuration,
+)-> sc_cli::Result<()>
+where
+	Block: BlockT + serde::de::DeserializeOwned,
+	Block::Hash: FromStr,
+	<Block::Hash as FromStr>::Err: Debug,
+	NumberFor<Block>: FromStr,
+	<NumberFor<Block> as FromStr>::Err: Debug,
+	ExecDispatch: NativeExecutionDispatch + 'static,
+{
+	let wasm_method = shared.wasm_method;
+	let execution = shared.execution;
+	let heap_pages = if shared.heap_pages.is_some() {
+		shared.heap_pages
+	} else {
+		config.default_heap_pages
+	};
+
+	let mut changes = Default::default();
+	let max_runtime_instances = config.max_runtime_instances;
+	let executor = NativeExecutor::<ExecDispatch>::new(
+		wasm_method.into(),
+		heap_pages,
+		max_runtime_instances,
+	);
+
+	let (mode, block) = match command.state {
+		State::Snap{ snapshot_path, url } => {
+			// `--header-at` is a required option for `snap`
+			let header_hash: Block::Hash = command.header_at
+				.expect("`--header-at` is a required option for `execute-block snap`.")
+				.parse()
+				.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
+			let block: Block = rpc_api::get_block::<Block, _>(url.clone(), header_hash).await?;
+
+			let mode = Mode::Offline(OfflineConfig {
+				state_snapshot: SnapshotConfig::new(snapshot_path),
+			});
+
+			(mode, block)
+		},
+		State::Live { url, snapshot_path, block_at, modules } => {
+			let block_hash: Block::Hash = block_at
+				.ok_or("execute-block requires `block_at` option")?
+				.parse()
+				.map_err(|e| format!("Could not parse hash: {:?}", e))?;
+
+			let block: Block = rpc_api::get_block::<Block, _>(url.clone(), block_hash).await?;
+			let parent_hash = block.header().parent_hash();
+
+			let mode = Mode::Online(OnlineConfig {
+				transport: url.to_owned().into(),
+				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
+				modules: modules.to_owned().unwrap_or_default(),
+				at: Some(*parent_hash),
+				..Default::default()
+			});
+
+			(mode, block)
+		}
+	};
+
+	let ext = {
+		let builder = Builder::<Block>::new().mode(mode);
+		// We do not allow overwriting code because we need exact state for storage root match.
+		let mut ext = builder.build().await?;
+
+		// register externality extensions in order to provide host interface for OCW to the runtime.
+		let (offchain, _offchain_state) = TestOffchainExt::new();
+		let (pool, _pool_state) = TestTransactionPoolExt::new();
+		ext.register_extension(OffchainDbExt::new(offchain.clone()));
+		ext.register_extension(OffchainWorkerExt::new(offchain));
+		ext.register_extension(KeystoreExt(Arc::new(KeyStore::new())));
+		ext.register_extension(TransactionPoolExt::new(pool));
+
+		ext
+	};
+
+	// A digest item gets added when the runtime is processing the block, so we need to pop
+	// the last one to be consistent with what a gossiped block would contain.
+	let (mut header, extrinsics) = block.deconstruct();
+	header.digest_mut().pop();
+	let block = Block::new(header, extrinsics);
+
+	let _ = StateMachine::<_, _, NumberFor<Block>, _>::new(
+		&ext.backend,
+		None,
+		&mut changes,
+		&executor,
+		"Core_execute_block",
+		block.encode().as_ref(),
+		ext.extensions,
+		&sp_state_machine::backend::BackendRuntimeCode::new(&ext.backend).runtime_code()?,
+		sp_core::testing::TaskExecutor::new(),
+	)
+	.execute(execution.into())
+	.map_err(|e| format!("failed to execute 'Core_execute_block' due to {:?}", e))?;
+
+	log::info!("Core_execute_block executed without errors.");
+
+	Ok(())
+}
 
 impl TryRuntimeCmd {
 	pub async fn run<Block, ExecDispatch>(&self, config: Configuration) -> sc_cli::Result<()>
 	where
-		Block: BlockT,
+		Block: BlockT + serde::de::DeserializeOwned,
 		Block::Header: serde::de::DeserializeOwned,
 		Block::Hash: FromStr,
 		<Block::Hash as FromStr>::Err: Debug,
@@ -437,6 +464,9 @@ impl TryRuntimeCmd {
 			Command::OffchainWorker(cmd) => {
 				offchain_worker::<Block, ExecDispatch>(self.shared.clone(), cmd.clone(), config).await
 			}
+			Command::ExecuteBlock(cmd) => {
+				execute_block::<Block, ExecDispatch>(self.shared.clone(), cmd.clone(), config).await
+			}
 		}
 	}
 }

From 4f380718e125e943e370e978a27e7965999ef2a2 Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Thu, 10 Jun 2021 17:53:25 -0700
Subject: [PATCH 03/13] Trivial

---
 utils/frame/try-runtime/cli/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 3d81a6c671aca..5f728f61c7292 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -397,7 +397,7 @@ where
 				transport: url.to_owned().into(),
 				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 				modules: modules.to_owned().unwrap_or_default(),
-				at: Some(*parent_hash),
+				at: Some(parent_hash.to_owned()),
 				..Default::default()
 			});
 

From ba0e5a6271bb09ec7465c99ec72cdd371d2c1725 Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Tue, 15 Jun 2021 12:08:36 -0700
Subject: [PATCH 04/13] Address some comments

---
 .../frame/remote-externalities/src/rpc_api.rs | 10 ++++----
 utils/frame/try-runtime/cli/src/lib.rs        | 25 ++++++-------------
 2 files changed, 12 insertions(+), 23 deletions(-)

diff --git a/utils/frame/remote-externalities/src/rpc_api.rs b/utils/frame/remote-externalities/src/rpc_api.rs
index a5ffa8410f1b8..6773bfd54bb19 100644
--- a/utils/frame/remote-externalities/src/rpc_api.rs
+++ b/utils/frame/remote-externalities/src/rpc_api.rs
@@ -33,7 +33,7 @@ where
 
 	client.request::<Block::Header>("chain_getHeader", JsonRpcParams::Array(params))
 		.await
-		.map_err(|e| format!("chain_getHeader request failed due to {:?}", e))
+		.map_err(|e| format!("chain_getHeader request failed: {:?}", e))
 }
 
 /// Get the finalized head
@@ -46,7 +46,7 @@ where
 
 	client.request::<Block::Hash>("chain_getFinalizedHead", JsonRpcParams::NoParams)
 		.await
-		.map_err(|e| format!("chain_getFinalizedHead request failed due to {:?}", e))
+		.map_err(|e| format!("chain_getFinalizedHead request failed: {:?}", e))
 }
 
 /// Get the signed block identified by `at`.
@@ -61,7 +61,7 @@ where
 	let signed_block = client
 		.request::<SignedBlock<Block>>("chain_getBlock", JsonRpcParams::Array(params))
 		.await
-		.map_err(|e| format!("chain_getBlock request failed due to {:?}", e))?;
+		.map_err(|e| format!("chain_getBlock request failed: {:?}", e))?;
 
 	Ok(signed_block.block)
 }
@@ -69,7 +69,7 @@ where
 /// Convert a block hash to a serde json value.
 fn hash_to_json<Block: BlockT>(hash: Block::Hash) -> Result<serde_json::Value, String> {
 	serde_json::to_value(hash)
-		.map_err(|e| format!("Block hash could not be converted to JSON due to {:?}", e))
+		.map_err(|e| format!("Block hash could not be converted to JSON: {:?}", e))
 }
 
 /// Build a website client that connects to `from`.
@@ -78,5 +78,5 @@ async fn build_client<S: AsRef<str>>(from: S) -> Result<WsClient, String> {
 		.max_request_body_size(u32::MAX)
 		.build(from.as_ref())
 		.await
-		.map_err(|e| format!("`WsClientBuilder` failed to build do to {:?}", e))
+		.map_err(|e| format!("`WsClientBuilder` failed to build: {:?}", e))
 }
diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 5f728f61c7292..06ea639bab17f 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -136,7 +136,7 @@ pub enum State {
 		// TODO This is a temporary hack; the url is used just to get the header/block. We should try
 		// and get that out of state, OR allow the user to feed in a header/block via file.
 		// https://github.com/paritytech/substrate/issues/9027
-		#[structopt(default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
+		#[structopt(short, long, default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
 		url: String,
 	},
 
@@ -156,7 +156,7 @@ pub enum State {
 		modules: Option<Vec<String>>,
 
 		/// The url to connect to.
-		#[structopt(default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
+		#[structopt(short, long, default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
 		url: String,
 	},
 }
@@ -176,11 +176,7 @@ where
 {
 	let wasm_method = shared.wasm_method;
 	let execution = shared.execution;
-	let heap_pages = if shared.heap_pages.is_some() {
-		shared.heap_pages
-	} else {
-		config.default_heap_pages
-	};
+	let heap_pages = shared.heap_pages.or(config.default_heap_pages);
 
 	let mut changes = Default::default();
 	let max_runtime_instances = config.max_runtime_instances;
@@ -259,11 +255,7 @@ where
 {
 	let wasm_method = shared.wasm_method;
 	let execution = shared.execution;
-	let heap_pages = if shared.heap_pages.is_some() {
-		shared.heap_pages
-	} else {
-		config.default_heap_pages
-	};
+	let heap_pages = shared.heap_pages.or(config.default_heap_pages);
 
 	let mut changes = Default::default();
 	let max_runtime_instances = config.max_runtime_instances;
@@ -355,11 +347,7 @@ where
 {
 	let wasm_method = shared.wasm_method;
 	let execution = shared.execution;
-	let heap_pages = if shared.heap_pages.is_some() {
-		shared.heap_pages
-	} else {
-		config.default_heap_pages
-	};
+	let heap_pages = shared.heap_pages.or(config.default_heap_pages);
 
 	let mut changes = Default::default();
 	let max_runtime_instances = config.max_runtime_instances;
@@ -427,7 +415,7 @@ where
 	header.digest_mut().pop();
 	let block = Block::new(header, extrinsics);
 
-	let _ = StateMachine::<_, _, NumberFor<Block>, _>::new(
+	let _encoded_result = StateMachine::<_, _, NumberFor<Block>, _>::new(
 		&ext.backend,
 		None,
 		&mut changes,
@@ -440,6 +428,7 @@ where
 	)
 	.execute(execution.into())
 	.map_err(|e| format!("failed to execute 'Core_execute_block' due to {:?}", e))?;
+	debug_assert!(_encoded_result == vec![1]);
 
 	log::info!("Core_execute_block executed without errors.");
 

From bada167a3a334455fe9efcfad2ea378a7c57253d Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Wed, 16 Jun 2021 14:30:26 -0700
Subject: [PATCH 05/13] Use required_if & remove header-at usage

---
 utils/frame/try-runtime/cli/src/lib.rs | 46 +++++++++++---------------
 1 file changed, 20 insertions(+), 26 deletions(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 06ea639bab17f..1316bbcd9712b 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -57,10 +57,6 @@ pub struct OnRuntimeUpgradeCmd {
 
 #[derive(Debug, Clone, structopt::StructOpt)]
 pub struct OffchainWorkerCmd {
-	/// Hash of the block whose header to use to execute the offchain worker.
-	#[structopt(short, long, multiple = false, parse(try_from_str = parse::hash))]
-	pub header_at: String,
-
 	#[structopt(subcommand)]
 	pub state: State,
 
@@ -72,10 +68,6 @@ pub struct OffchainWorkerCmd {
 
 #[derive(Debug, Clone, structopt::StructOpt)]
 pub struct ExecuteBlockCmd {
-	/// Hash of the block whose header to use to execute the offchain worker.
-	/// NOTE: `--header-at` is a required option for `snap`, but not `live`
-	#[structopt(short, long, multiple = false, parse(try_from_str = parse::hash))]
-	pub header_at: Option<String>,
 	#[structopt(subcommand)]
 	pub state: State,
 }
@@ -111,6 +103,19 @@ pub struct SharedParams {
 	/// sc_service::Configuration.default_heap_pages.
 	#[structopt(long)]
 	pub heap_pages: Option<u64>,
+
+	/// The block hash at which to read state. This is required for execute-block, offchain-worker,
+	/// or state live with any command.
+	#[structopt(
+		short,
+		long,
+		multiple = false,
+		parse(try_from_str = parse::hash),
+		required_ifs(
+			&[("command", "offchain-worker"), ("command", "execute-block"), ("subcommand", "live")]
+		)
+	)]
+	pub block_at: String,
 }
 
 /// Various commands to try out against runtime state at a specific block.
@@ -146,11 +151,6 @@ pub enum State {
 		#[structopt(short, long)]
 		snapshot_path: Option<PathBuf>,
 
-		/// The block hash at which to connect.
-		/// Will be latest finalized head if not provided.
-		#[structopt(short, long, multiple = false, parse(try_from_str = parse::hash))]
-		block_at: Option<String>,
-
 		/// The modules to scrape. If empty, entire chain state will be scraped.
 		#[structopt(short, long, require_delimiter = true)]
 		modules: Option<Vec<String>>,
@@ -196,14 +196,12 @@ where
 			State::Live {
 				url,
 				snapshot_path,
-				block_at,
 				modules
 			} => Builder::<Block>::new().mode(Mode::Online(OnlineConfig {
 				transport: url.to_owned().into(),
 				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 				modules: modules.to_owned().unwrap_or_default(),
-				at: block_at.as_ref()
-					.map(|b| b.parse().map_err(|e| format!("Could not parse hash: {:?}", e))).transpose()?,
+				at: Some(shared.block_at.parse().map_err(|e| format!("Could not parse hash: {:?}", e))?),
 				..Default::default()
 			})),
 		};
@@ -269,15 +267,14 @@ where
 			State::Live {
 				url,
 				snapshot_path,
-				block_at,
 				modules
 			} => {
+				let at = shared.block_at.parse().map_err(|e| format!("Could not parse hash: {:?}", e))?;
 				let online_config = OnlineConfig {
 					transport: url.to_owned().into(),
 					state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 					modules: modules.to_owned().unwrap_or_default(),
-					at: block_at.as_ref()
-						.map(|b| b.parse().map_err(|e| format!("Could not parse hash: {:?}", e))).transpose()?,
+					at: Some(at),
 					..Default::default()
 				};
 
@@ -307,7 +304,7 @@ where
 	ext.register_extension(KeystoreExt(Arc::new(KeyStore::new())));
 	ext.register_extension(TransactionPoolExt::new(pool));
 
-	let header_hash: Block::Hash = command.header_at
+	let header_hash: Block::Hash = shared.block_at
 		.parse()
 		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
 	let header = rpc_api::get_header::<Block, _>(url, header_hash).await?;
@@ -359,9 +356,7 @@ where
 
 	let (mode, block) = match command.state {
 		State::Snap{ snapshot_path, url } => {
-			// `--header-at` is a required option for `snap`
-			let header_hash: Block::Hash = command.header_at
-				.expect("`--header-at` is a required option for `execute-block snap`.")
+			let header_hash: Block::Hash = shared.block_at
 				.parse()
 				.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
 			let block: Block = rpc_api::get_block::<Block, _>(url.clone(), header_hash).await?;
@@ -372,9 +367,8 @@ where
 
 			(mode, block)
 		},
-		State::Live { url, snapshot_path, block_at, modules } => {
-			let block_hash: Block::Hash = block_at
-				.ok_or("execute-block requires `block_at` option")?
+		State::Live { url, snapshot_path, modules } => {
+			let block_hash: Block::Hash = shared.block_at
 				.parse()
 				.map_err(|e| format!("Could not parse hash: {:?}", e))?;
 

From 36f2043799679baeeecbb3364bcc584eb6943bf7 Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Wed, 16 Jun 2021 15:21:53 -0700
Subject: [PATCH 06/13] Improve doc

---
 utils/frame/try-runtime/cli/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 1316bbcd9712b..660c260ae0d47 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -105,7 +105,7 @@ pub struct SharedParams {
 	pub heap_pages: Option<u64>,
 
 	/// The block hash at which to read state. This is required for execute-block, offchain-worker,
-	/// or state live with any command.
+	/// or any command that used the live subcommand.
 	#[structopt(
 		short,
 		long,

From 742e9a5a3eb586dffe50907abb1e0d9609cd1d46 Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Wed, 16 Jun 2021 15:32:01 -0700
Subject: [PATCH 07/13] Update comment

---
 utils/frame/try-runtime/cli/src/lib.rs | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 660c260ae0d47..a3150a7da7e78 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -137,7 +137,7 @@ pub enum State {
 		#[structopt(short, long)]
 		snapshot_path: PathBuf,
 		/// The url to connect to for fetching data not easily retrievable from the snapshot.
-		/// NOTE: This is not necessary for `on-runtime-upgrade`
+		/// NOTE: This is not used for `on-runtime-upgrade`
 		// TODO This is a temporary hack; the url is used just to get the header/block. We should try
 		// and get that out of state, OR allow the user to feed in a header/block via file.
 		// https://github.com/paritytech/substrate/issues/9027
@@ -296,7 +296,6 @@ where
 		builder.build().await?
 	};
 
-	// register externality extensions in order to provide host interface for OCW to the runtime.
 	let (offchain, _offchain_state) = TestOffchainExt::new();
 	let (pool, _pool_state) = TestTransactionPoolExt::new();
 	ext.register_extension(OffchainDbExt::new(offchain.clone()));

From 354d88eff7bb7c98944334627df3dc3b8ad6e8c8 Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Thu, 17 Jun 2021 20:40:41 +0200
Subject: [PATCH 08/13] small tweaks

---
 Cargo.lock                                    |  1 +
 .../election-provider-multi-phase/src/lib.rs  |  2 +-
 utils/frame/remote-externalities/src/lib.rs   | 59 ++++++++++++++++---
 utils/frame/try-runtime/cli/Cargo.toml        |  1 +
 utils/frame/try-runtime/cli/src/lib.rs        | 48 +++++++++------
 5 files changed, 85 insertions(+), 26 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fb944b782abd9..79383bfccfe78 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10646,6 +10646,7 @@ dependencies = [
  "sp-blockchain",
  "sp-core",
  "sp-externalities",
+ "sp-io",
  "sp-keystore",
  "sp-runtime",
  "sp-state-machine",
diff --git a/frame/election-provider-multi-phase/src/lib.rs b/frame/election-provider-multi-phase/src/lib.rs
index 2bb47a8778074..2864ca518d068 100644
--- a/frame/election-provider-multi-phase/src/lib.rs
+++ b/frame/election-provider-multi-phase/src/lib.rs
@@ -608,7 +608,7 @@ pub mod pallet {
 		type Fallback: Get<FallbackStrategy>;
 
 		/// Origin that can control this pallet. Note that any action taken by this origin (such)
-		/// as providing an emergency solution is not checked. Thus, it must be a trusted origin. 
+		/// as providing an emergency solution is not checked. Thus, it must be a trusted origin.
 		type ForceOrigin: EnsureOrigin<Self::Origin>;
 
 		/// The configuration of benchmarking.
diff --git a/utils/frame/remote-externalities/src/lib.rs b/utils/frame/remote-externalities/src/lib.rs
index 0214bf263892a..dce9adb55fded 100644
--- a/utils/frame/remote-externalities/src/lib.rs
+++ b/utils/frame/remote-externalities/src/lib.rs
@@ -43,10 +43,12 @@ type KeyPair = (StorageKey, StorageData);
 
 const LOG_TARGET: &str = "remote-ext";
 const DEFAULT_TARGET: &str = "wss://rpc.polkadot.io";
-const BATCH_SIZE: usize = 512;
+const BATCH_SIZE: usize = 1000;
 
 jsonrpsee_proc_macros::rpc_client_api! {
 	RpcApi<B: BlockT> {
+		#[rpc(method = "state_getStorage", positional_params)]
+		fn get_storage(prefix: StorageKey, hash: Option<B::Hash>) -> StorageData;
 		#[rpc(method = "state_getKeysPaged", positional_params)]
 		fn get_keys_paged(
 			prefix: Option<StorageKey>,
@@ -159,8 +161,11 @@ impl Default for SnapshotConfig {
 pub struct Builder<B: BlockT> {
 	/// Custom key-pairs to be injected into the externalities.
 	inject: Vec<KeyPair>,
-	/// Storage entry key prefixes to be injected into the externalities. The *hashed* prefix must be given.
+	/// Storage entry key prefixes to be injected into the externalities. The *hashed* prefix must
+	/// be given.
 	hashed_prefixes: Vec<Vec<u8>>,
+	/// Storage entry keys to be injected into the externalities. The *hashed* key must be given.
+	hashed_keys: Vec<Vec<u8>>,
 	/// connectivity mode, online or offline.
 	mode: Mode<B>,
 }
@@ -169,7 +174,12 @@ pub struct Builder<B: BlockT> {
 // that.
 impl<B: BlockT> Default for Builder<B> {
 	fn default() -> Self {
-		Self { inject: Default::default(), mode: Default::default(), hashed_prefixes: Default::default() }
+		Self {
+			inject: Default::default(),
+			mode: Default::default(),
+			hashed_prefixes: Default::default(),
+			hashed_keys: Default::default(),
+		}
 	}
 }
 
@@ -192,6 +202,17 @@ impl<B: BlockT> Builder<B> {
 
 // RPC methods
 impl<B: BlockT> Builder<B> {
+	async fn rpc_get_storage(
+		&self,
+		key: StorageKey,
+		maybe_at: Option<B::Hash>,
+	) -> Result<StorageData, &'static str> {
+		trace!(target: LOG_TARGET, "rpc: get_storage");
+		RpcApi::<B>::get_storage(self.as_online().rpc_client(), key, maybe_at).await.map_err(|e| {
+			error!("Error = {:?}", e);
+			"rpc get_storage failed."
+		})
+	}
 	/// Get the latest finalized head.
 	async fn rpc_get_head(&self) -> Result<B::Hash, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: finalized_head");
@@ -356,11 +377,23 @@ impl<B: BlockT> Builder<B> {
 		};
 
 		for prefix in &self.hashed_prefixes {
-			info!(target: LOG_TARGET, "adding data for hashed prefix: {:?}", HexDisplay::from(prefix));
-			let additional_key_values = self.rpc_get_pairs_paged(StorageKey(prefix.to_vec()), at).await?;
+			debug!(
+				target: LOG_TARGET,
+				"adding data for hashed prefix: {:?}",
+				HexDisplay::from(prefix)
+			);
+			let additional_key_values =
+				self.rpc_get_pairs_paged(StorageKey(prefix.to_vec()), at).await?;
 			keys_and_values.extend(additional_key_values);
 		}
 
+		for key in &self.hashed_keys {
+			let key = StorageKey(key.to_vec());
+			debug!(target: LOG_TARGET, "adding data for hashed key: {:?}", HexDisplay::from(&key));
+			let value = self.rpc_get_storage(key.clone(), Some(at)).await?;
+			keys_and_values.push((key, value));
+		}
+
 		Ok(keys_and_values)
 	}
 
@@ -400,7 +433,7 @@ impl<B: BlockT> Builder<B> {
 
 		info!(
 			target: LOG_TARGET,
-			"extending externalities with {} manually injected keys",
+			"extending externalities with {} manually injected key-values",
 			self.inject.len()
 		);
 		base_kv.extend(self.inject.clone());
@@ -416,19 +449,29 @@ impl<B: BlockT> Builder<B> {
 	}
 
 	/// Inject a manual list of key and values to the storage.
-	pub fn inject(mut self, injections: &[KeyPair]) -> Self {
+	pub fn inject_key_value(mut self, injections: &[KeyPair]) -> Self {
 		for i in injections {
 			self.inject.push(i.clone());
 		}
 		self
 	}
 
-	/// Inject a hashed prefix. This is treated as-is, and should be pre-hashed. 
+	/// Inject a hashed prefix. This is treated as-is, and should be pre-hashed.
+	///
+	/// This should be used to inject a "PREFIX", like a storage (double) map.
 	pub fn inject_hashed_prefix(mut self, hashed: &[u8]) -> Self {
 		self.hashed_prefixes.push(hashed.to_vec());
 		self
 	}
 
+	/// Inject a hashed key to scrape. This is treated as-is, and should be pre-hashed.
+	///
+	/// This should be used to inject a "KEY", like a storage value.
+	pub fn inject_hashed_key(mut self, hashed: &[u8]) -> Self {
+		self.hashed_keys.push(hashed.to_vec());
+		self
+	}
+
 	/// Configure a state snapshot to be used.
 	pub fn mode(mut self, mode: Mode<B>) -> Self {
 		self.mode = mode;
diff --git a/utils/frame/try-runtime/cli/Cargo.toml b/utils/frame/try-runtime/cli/Cargo.toml
index f262ba4812a0e..2e2335bc5fff9 100644
--- a/utils/frame/try-runtime/cli/Cargo.toml
+++ b/utils/frame/try-runtime/cli/Cargo.toml
@@ -29,6 +29,7 @@ sp-blockchain = { version = "3.0.0", path = "../../../../primitives/blockchain"
 sp-runtime = { version = "3.0.0", path = "../../../../primitives/runtime" }
 sp-externalities = { version = "0.9.0", path = "../../../../primitives/externalities" }
 sp-core = { version = "3.0.0", path = "../../../../primitives/core" }
+sp-io = { version = "3.0.0", path = "../../../../primitives/io" }
 sp-keystore = { version = "0.9.0", path = "../../../../primitives/keystore" }
 frame-try-runtime = { version = "0.9.0", path = "../../../../frame/try-runtime" }
 
diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index a3150a7da7e78..98e001766bfd5 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -29,9 +29,10 @@ use sp_runtime::traits::{Block as BlockT, NumberFor, Header as HeaderT};
 use sp_core::{
 	offchain::{
 		OffchainWorkerExt, OffchainDbExt, TransactionPoolExt,
-		testing::{TestOffchainExt, TestTransactionPoolExt}
+		testing::{TestOffchainExt, TestTransactionPoolExt},
 	},
 	storage::{StorageData, StorageKey, well_known_keys},
+	hashing::twox_128,
 };
 use sp_keystore::{KeystoreExt, testing::KeyStore};
 use remote_externalities::{Builder, Mode, SnapshotConfig, OfflineConfig, OnlineConfig, rpc_api};
@@ -59,11 +60,6 @@ pub struct OnRuntimeUpgradeCmd {
 pub struct OffchainWorkerCmd {
 	#[structopt(subcommand)]
 	pub state: State,
-
-	/// Whether or not to overwrite the code from state with the code from
-	/// the specified chain spec.
-	#[structopt(long)]
-	pub overwrite_code: bool,
 }
 
 #[derive(Debug, Clone, structopt::StructOpt)]
@@ -207,7 +203,10 @@ where
 		};
 
 		let (code_key, code) = extract_code(config.chain_spec)?;
-		builder.inject(&[(code_key, code)]).build().await?
+		builder
+			.inject_key_value(&[(code_key, code)])
+			.inject_hashed_key(&[twox_128(b"System"), twox_128(b"LastRuntimeUpgrade")].concat())
+			.build().await?
 	};
 
 	let encoded_result = StateMachine::<_, _, NumberFor<Block>, _>::new(
@@ -241,7 +240,7 @@ async fn offchain_worker<Block, ExecDispatch>(
 	shared: SharedParams,
 	command: OffchainWorkerCmd,
 	config: Configuration,
-)-> sc_cli::Result<()>
+) -> sc_cli::Result<()>
 where
 	Block: BlockT,
 	Block::Hash: FromStr,
@@ -288,12 +287,17 @@ where
 				(mode, url)
 			}
 	};
-	let builder = Builder::<Block>::new().mode(mode);
+	let builder = Builder::<Block>::new()
+		.mode(mode)
+		.inject_hashed_key(&[twox_128(b"System"), twox_128(b"LastRuntimeUpgrade")].concat());
 	let mut ext = if command.overwrite_code {
 		let (code_key, code) = extract_code(config.chain_spec)?;
-		builder.inject(&[(code_key, code)]).build().await?
+		builder.inject_key_value(&[(code_key, code)]).build().await?
 	} else {
-		builder.build().await?
+		builder
+			.inject_hashed_key(well_known_keys::CODE)
+			.build()
+			.await?
 	};
 
 	let (offchain, _offchain_state) = TestOffchainExt::new();
@@ -354,7 +358,7 @@ where
 	);
 
 	let (mode, block) = match command.state {
-		State::Snap{ snapshot_path, url } => {
+		State::Snap { snapshot_path, url } => {
 			let header_hash: Block::Hash = shared.block_at
 				.parse()
 				.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
@@ -387,11 +391,21 @@ where
 	};
 
 	let ext = {
-		let builder = Builder::<Block>::new().mode(mode);
-		// We do not allow overwriting code because we need exact state for storage root match.
-		let mut ext = builder.build().await?;
+		let builder = Builder::<Block>::new()
+			.mode(mode)
+			.inject_hashed_key(&[twox_128(b"System"), twox_128(b"LastRuntimeUpgrade")].concat());
+		let mut ext = if command.overwrite_code {
+			let (code_key, code) = extract_code(config.chain_spec)?;
+			builder.inject_key_value(&[(code_key, code)]).build().await?
+		} else {
+			builder
+				.inject_hashed_key(well_known_keys::CODE)
+				.build()
+				.await?
+		};
 
-		// register externality extensions in order to provide host interface for OCW to the runtime.
+		// register externality extensions in order to provide host interface for OCW to the
+		// runtime.
 		let (offchain, _offchain_state) = TestOffchainExt::new();
 		let (pool, _pool_state) = TestTransactionPoolExt::new();
 		ext.register_extension(OffchainDbExt::new(offchain.clone()));
@@ -466,7 +480,7 @@ impl CliConfiguration for TryRuntimeCmd {
 	}
 }
 
-/// Extract `:code` from the given chain spec and return as `StorageData` along with the 
+/// Extract `:code` from the given chain spec and return as `StorageData` along with the
 /// corresponding `StorageKey`.
 fn extract_code(spec: Box<dyn ChainSpec>) -> sc_cli::Result<(StorageKey, StorageData)> {
 	let genesis_storage = spec.build_storage()?;

From 033dfbc02f1a703bea183b80141e65c739119a13 Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Thu, 17 Jun 2021 12:35:15 -0700
Subject: [PATCH 09/13] add overwrite-code to shared params

---
 utils/frame/try-runtime/cli/src/lib.rs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 98e001766bfd5..3f07b5a62fd74 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -112,6 +112,10 @@ pub struct SharedParams {
 		)
 	)]
 	pub block_at: String,
+	/// Whether or not to overwrite the code from state with the code from
+	/// the specified chain spec.
+	#[structopt(long)]
+	pub overwrite_code: bool,
 }
 
 /// Various commands to try out against runtime state at a specific block.
@@ -290,7 +294,7 @@ where
 	let builder = Builder::<Block>::new()
 		.mode(mode)
 		.inject_hashed_key(&[twox_128(b"System"), twox_128(b"LastRuntimeUpgrade")].concat());
-	let mut ext = if command.overwrite_code {
+	let mut ext = if shared.overwrite_code {
 		let (code_key, code) = extract_code(config.chain_spec)?;
 		builder.inject_key_value(&[(code_key, code)]).build().await?
 	} else {
@@ -394,7 +398,7 @@ where
 		let builder = Builder::<Block>::new()
 			.mode(mode)
 			.inject_hashed_key(&[twox_128(b"System"), twox_128(b"LastRuntimeUpgrade")].concat());
-		let mut ext = if command.overwrite_code {
+		let mut ext = if shared.overwrite_code {
 			let (code_key, code) = extract_code(config.chain_spec)?;
 			builder.inject_key_value(&[(code_key, code)]).build().await?
 		} else {

From 282ace2689b11254f7b9e4e2bfb15809ec5d36b4 Mon Sep 17 00:00:00 2001
From: Zeke Mostov <32168567+emostov@users.noreply.github.com>
Date: Thu, 17 Jun 2021 14:49:33 -0700
Subject: [PATCH 10/13] Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Peter Goodspeed-Niklaus <coriolinus@users.noreply.github.com>
---
 utils/frame/try-runtime/cli/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 3f07b5a62fd74..a29dd98fe5565 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -438,7 +438,7 @@ where
 		sp_core::testing::TaskExecutor::new(),
 	)
 	.execute(execution.into())
-	.map_err(|e| format!("failed to execute 'Core_execute_block' due to {:?}", e))?;
+	.map_err(|e| format!("failed to execute 'Core_execute_block': {:?}", e))?;
 	debug_assert!(_encoded_result == vec![1]);
 
 	log::info!("Core_execute_block executed without errors.");

From 59c189ca63a0a26a04761db3b187ef29eb833e7a Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Thu, 17 Jun 2021 15:20:02 -0700
Subject: [PATCH 11/13] make url a shared param

---
 utils/frame/try-runtime/cli/src/lib.rs | 72 +++++++++++---------------
 1 file changed, 31 insertions(+), 41 deletions(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 3f07b5a62fd74..493f4d6b26055 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -116,6 +116,14 @@ pub struct SharedParams {
 	/// the specified chain spec.
 	#[structopt(long)]
 	pub overwrite_code: bool,
+
+	/// The url to connect to.
+	// TODO having this a shared parm is a temporary hack; the url is used just
+	// to get the header/block. We should try and get that out of state, OR allow
+	// the user to feed in a header/block via file.
+	// https://github.com/paritytech/substrate/issues/9027
+	#[structopt(short, long, default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
+	url: String,
 }
 
 /// Various commands to try out against runtime state at a specific block.
@@ -136,13 +144,6 @@ pub enum State {
 	Snap {
 		#[structopt(short, long)]
 		snapshot_path: PathBuf,
-		/// The url to connect to for fetching data not easily retrievable from the snapshot.
-		/// NOTE: This is not used for `on-runtime-upgrade`
-		// TODO This is a temporary hack; the url is used just to get the header/block. We should try
-		// and get that out of state, OR allow the user to feed in a header/block via file.
-		// https://github.com/paritytech/substrate/issues/9027
-		#[structopt(short, long, default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
-		url: String,
 	},
 
 	/// Use a live chain as the source of runtime state.
@@ -154,11 +155,7 @@ pub enum State {
 		/// The modules to scrape. If empty, entire chain state will be scraped.
 		#[structopt(short, long, require_delimiter = true)]
 		modules: Option<Vec<String>>,
-
-		/// The url to connect to.
-		#[structopt(short, long, default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
-		url: String,
-	},
+	}
 }
 
 async fn on_runtime_upgrade<Block, ExecDispatch>(
@@ -188,17 +185,16 @@ where
 
 	let ext = {
 		let builder = match command.state {
-			State::Snap { snapshot_path, .. } => {
+			State::Snap { snapshot_path } => {
 				Builder::<Block>::new().mode(Mode::Offline(OfflineConfig {
 					state_snapshot: SnapshotConfig::new(snapshot_path),
 				}))
 			},
 			State::Live {
-				url,
 				snapshot_path,
 				modules
 			} => Builder::<Block>::new().mode(Mode::Online(OnlineConfig {
-				transport: url.to_owned().into(),
+				transport: shared.url.to_owned().into(),
 				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 				modules: modules.to_owned().unwrap_or_default(),
 				at: Some(shared.block_at.parse().map_err(|e| format!("Could not parse hash: {:?}", e))?),
@@ -226,10 +222,10 @@ where
 		sp_core::testing::TaskExecutor::new(),
 	)
 	.execute(execution.into())
-	.map_err(|e| format!("failed to execute 'TryRuntime_on_runtime_upgrade' due to {:?}", e))?;
+	.map_err(|e| format!("failed to execute 'TryRuntime_on_runtime_upgrade': {:?}", e))?;
 
 	let (weight, total_weight) = <(u64, u64) as Decode>::decode(&mut &*encoded_result)
-		.map_err(|e| format!("failed to decode output due to {:?}", e))?;
+		.map_err(|e| format!("failed to decode output: {:?}", e))?;
 	log::info!(
 		"TryRuntime_on_runtime_upgrade executed without errors. Consumed weight = {}, total weight = {} ({})",
 		weight,
@@ -266,29 +262,28 @@ where
 		max_runtime_instances,
 	);
 
-	let (mode, url) = match command.state {
+	let mode = match command.state {
 			State::Live {
-				url,
 				snapshot_path,
 				modules
 			} => {
 				let at = shared.block_at.parse().map_err(|e| format!("Could not parse hash: {:?}", e))?;
 				let online_config = OnlineConfig {
-					transport: url.to_owned().into(),
+					transport: shared.url.to_owned().into(),
 					state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 					modules: modules.to_owned().unwrap_or_default(),
 					at: Some(at),
 					..Default::default()
 				};
 
-				(Mode::Online(online_config), url)
+				Mode::Online(online_config)
 			},
-			State::Snap { snapshot_path, url } => {
+			State::Snap { snapshot_path } => {
 				let mode = Mode::Offline(OfflineConfig {
 					state_snapshot: SnapshotConfig::new(snapshot_path),
 				});
 
-				(mode, url)
+				mode
 			}
 	};
 	let builder = Builder::<Block>::new()
@@ -314,7 +309,7 @@ where
 	let header_hash: Block::Hash = shared.block_at
 		.parse()
 		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
-	let header = rpc_api::get_header::<Block, _>(url, header_hash).await?;
+	let header = rpc_api::get_header::<Block, _>(shared.url, header_hash).await?;
 
 	let _ = StateMachine::<_, _, NumberFor<Block>, _>::new(
 		&ext.backend,
@@ -329,7 +324,7 @@ where
 		sp_core::testing::TaskExecutor::new(),
 	)
 	.execute(execution.into())
-	.map_err(|e| format!("failed to execute 'OffchainWorkerApi_offchain_worker' due to {:?}", e))?;
+	.map_err(|e| format!("failed to execute 'OffchainWorkerApi_offchain_worker':  {:?}", e))?;
 
 	log::info!("OffchainWorkerApi_offchain_worker executed without errors.");
 
@@ -361,36 +356,31 @@ where
 		max_runtime_instances,
 	);
 
-	let (mode, block) = match command.state {
-		State::Snap { snapshot_path, url } => {
-			let header_hash: Block::Hash = shared.block_at
-				.parse()
-				.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
-			let block: Block = rpc_api::get_block::<Block, _>(url.clone(), header_hash).await?;
+	let block_hash: Block::Hash = shared.block_at
+		.parse()
+		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
+	let block: Block = rpc_api::get_block::<Block, _>(shared.url.clone(), block_hash).await?;
 
+	let mode = match command.state {
+		State::Snap { snapshot_path } => {
 			let mode = Mode::Offline(OfflineConfig {
 				state_snapshot: SnapshotConfig::new(snapshot_path),
 			});
 
-			(mode, block)
+			mode
 		},
-		State::Live { url, snapshot_path, modules } => {
-			let block_hash: Block::Hash = shared.block_at
-				.parse()
-				.map_err(|e| format!("Could not parse hash: {:?}", e))?;
-
-			let block: Block = rpc_api::get_block::<Block, _>(url.clone(), block_hash).await?;
+		State::Live { snapshot_path, modules } => {
 			let parent_hash = block.header().parent_hash();
 
 			let mode = Mode::Online(OnlineConfig {
-				transport: url.to_owned().into(),
+				transport: shared.url.to_owned().into(),
 				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 				modules: modules.to_owned().unwrap_or_default(),
 				at: Some(parent_hash.to_owned()),
 				..Default::default()
 			});
 
-			(mode, block)
+			mode
 		}
 	};
 
@@ -438,7 +428,7 @@ where
 		sp_core::testing::TaskExecutor::new(),
 	)
 	.execute(execution.into())
-	.map_err(|e| format!("failed to execute 'Core_execute_block' due to {:?}", e))?;
+	.map_err(|e| format!("failed to execute 'Core_execute_block': {:?}", e))?;
 	debug_assert!(_encoded_result == vec![1]);
 
 	log::info!("Core_execute_block executed without errors.");

From db29fecb90a7790886c126d5210b4772ed9fcb7d Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <coriolinus@users.noreply.github.com>
Date: Mon, 21 Jun 2021 19:11:20 +0200
Subject: [PATCH 12/13] add helper for block_at (#9153)

* add helper for block_at

* remove redundant bound

* doc for fn block_at
---
 utils/frame/try-runtime/cli/src/lib.rs | 33 +++++++++++++++++---------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/utils/frame/try-runtime/cli/src/lib.rs b/utils/frame/try-runtime/cli/src/lib.rs
index 493f4d6b26055..e0d09ff7fbcf4 100644
--- a/utils/frame/try-runtime/cli/src/lib.rs
+++ b/utils/frame/try-runtime/cli/src/lib.rs
@@ -111,7 +111,8 @@ pub struct SharedParams {
 			&[("command", "offchain-worker"), ("command", "execute-block"), ("subcommand", "live")]
 		)
 	)]
-	pub block_at: String,
+	block_at: String,
+
 	/// Whether or not to overwrite the code from state with the code from
 	/// the specified chain spec.
 	#[structopt(long)]
@@ -126,6 +127,20 @@ pub struct SharedParams {
 	url: String,
 }
 
+impl SharedParams {
+	/// Get the configured value of `block_at`, interpreted as the hash type of `Block`.
+	pub fn block_at<Block>(&self) -> sc_cli::Result<Block::Hash>
+	where
+		Block: BlockT,
+		<Block as BlockT>::Hash: FromStr,
+		<<Block as BlockT>::Hash as FromStr>::Err: Debug,
+	{
+		self.block_at
+			.parse::<<Block as BlockT>::Hash>()
+			.map_err(|e| format!("Could not parse block hash: {:?}", e).into())
+	}
+}
+
 /// Various commands to try out against runtime state at a specific block.
 #[derive(Debug, Clone, structopt::StructOpt)]
 pub struct TryRuntimeCmd {
@@ -161,7 +176,7 @@ pub enum State {
 async fn on_runtime_upgrade<Block, ExecDispatch>(
 	shared: SharedParams,
 	command: OnRuntimeUpgradeCmd,
-	config: Configuration
+	config: Configuration,
 ) -> sc_cli::Result<()>
 where
 	Block: BlockT,
@@ -197,7 +212,7 @@ where
 				transport: shared.url.to_owned().into(),
 				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 				modules: modules.to_owned().unwrap_or_default(),
-				at: Some(shared.block_at.parse().map_err(|e| format!("Could not parse hash: {:?}", e))?),
+				at: Some(shared.block_at::<Block>()?),
 				..Default::default()
 			})),
 		};
@@ -267,7 +282,7 @@ where
 				snapshot_path,
 				modules
 			} => {
-				let at = shared.block_at.parse().map_err(|e| format!("Could not parse hash: {:?}", e))?;
+				let at = shared.block_at::<Block>()?;
 				let online_config = OnlineConfig {
 					transport: shared.url.to_owned().into(),
 					state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
@@ -306,9 +321,7 @@ where
 	ext.register_extension(KeystoreExt(Arc::new(KeyStore::new())));
 	ext.register_extension(TransactionPoolExt::new(pool));
 
-	let header_hash: Block::Hash = shared.block_at
-		.parse()
-		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
+	let header_hash = shared.block_at::<Block>()?;
 	let header = rpc_api::get_header::<Block, _>(shared.url, header_hash).await?;
 
 	let _ = StateMachine::<_, _, NumberFor<Block>, _>::new(
@@ -335,7 +348,7 @@ async fn execute_block<Block, ExecDispatch>(
 	shared: SharedParams,
 	command: ExecuteBlockCmd,
 	config: Configuration,
-)-> sc_cli::Result<()>
+) -> sc_cli::Result<()>
 where
 	Block: BlockT + serde::de::DeserializeOwned,
 	Block::Hash: FromStr,
@@ -356,9 +369,7 @@ where
 		max_runtime_instances,
 	);
 
-	let block_hash: Block::Hash = shared.block_at
-		.parse()
-		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
+	let block_hash = shared.block_at::<Block>()?;
 	let block: Block = rpc_api::get_block::<Block, _>(shared.url.clone(), block_hash).await?;
 
 	let mode = match command.state {

From 2a8d77732219a486dedc3c3e81eb0b0b0e5303ce Mon Sep 17 00:00:00 2001
From: emostov <32168567+emostov@users.noreply.github.com>
Date: Mon, 21 Jun 2021 18:21:21 -0700
Subject: [PATCH 13/13] Update error message

---
 utils/frame/remote-externalities/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/frame/remote-externalities/src/lib.rs b/utils/frame/remote-externalities/src/lib.rs
index dce9adb55fded..4b6738f3b915a 100644
--- a/utils/frame/remote-externalities/src/lib.rs
+++ b/utils/frame/remote-externalities/src/lib.rs
@@ -302,7 +302,7 @@ impl<B: BlockT> Builder<B> {
 			let values = client.batch_request::<Option<StorageData>>(batch)
 				.await
 				.map_err(|e| {
-					log::error!(target: LOG_TARGET, "failed to execute batch {:?} due to {:?}", chunk_keys, e);
+					log::error!(target: LOG_TARGET, "failed to execute batch: {:?}. Error: {:?}", chunk_keys, e);
 					"batch failed."
 				})?;
 			assert_eq!(chunk_keys.len(), values.len());