chain-spec-builder: Add support for codeSubstitutes (paritytech#4685)

While working on paritytech#4600 I
found that it would be nice if `chain-spec-builder` supported
`codeSubstitutes`. After this PR is merged you can do:

```
chain-spec-builder add-code-substitute chain_spec.json my_runtime.compact.compressed.wasm 1234
```

In addition, the `chain-spec-builder` was silently removing
`relay_chain` and `para_id` fields when used on parachain chain-specs.
This is now fixed by providing a custom chain-spec extension that has
these fields marked as optional.

prdoc/pr_4685.prdoc

@@ -0,0 +1,16 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Chain-spec-builder supports `codeSubstitutes`.
+
+doc:
+  - audience: Node Operator
+    description: |
+      A new subcommand `add-code-substitute` is available for the `chain-spec-builder` binary. It allows users to provide a runtime that should be used from a given
+      block onwards. The `codeSubstitutes` field in the chain spec is used to force usage of a given runtime at a given block until the next runtime upgrade. It can be
+      used to progress chains that are stalled due to runtime bugs that prevent block-building. However, parachain usage is only possible in combination with an updated
+      validation function on the relay chain.
+
+crates:
+  - name: staging-chain-spec-builder
+    bump: minor

substrate/bin/utils/chain-spec-builder/bin/main.rs

@@ -17,16 +17,19 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 use chain_spec_builder::{
-	generate_chain_spec_for_runtime, ChainSpecBuilder, ChainSpecBuilderCmd, ConvertToRawCmd,
-	DisplayPresetCmd, ListPresetsCmd, UpdateCodeCmd, VerifyCmd,
+	generate_chain_spec_for_runtime, AddCodeSubstituteCmd, ChainSpecBuilder, ChainSpecBuilderCmd,
+	ConvertToRawCmd, DisplayPresetCmd, ListPresetsCmd, UpdateCodeCmd, VerifyCmd,
 };
 use clap::Parser;
 use sc_chain_spec::{
-	update_code_in_json_chain_spec, GenericChainSpec, GenesisConfigBuilderRuntimeCaller,
+	set_code_substitute_in_json_chain_spec, update_code_in_json_chain_spec, GenericChainSpec,
+	GenesisConfigBuilderRuntimeCaller,
 };
 use staging_chain_spec_builder as chain_spec_builder;
 use std::fs;
 
+type ChainSpec = GenericChainSpec<(), ()>;
+
 //avoid error message escaping
 fn main() {
 	match inner_main() {
@@ -50,7 +53,7 @@ fn inner_main() -> Result<(), String> {
 			ref input_chain_spec,
 			ref runtime_wasm_path,
 		}) => {
-			let chain_spec = GenericChainSpec::<()>::from_json_file(input_chain_spec.clone())?;
+			let chain_spec = ChainSpec::from_json_file(input_chain_spec.clone())?;
 
 			let mut chain_spec_json =
 				serde_json::from_str::<serde_json::Value>(&chain_spec.as_json(false)?)
@@ -65,8 +68,29 @@ fn inner_main() -> Result<(), String> {
 				.map_err(|e| format!("to pretty failed: {e}"))?;
 			fs::write(chain_spec_path, chain_spec_json).map_err(|err| err.to_string())?;
 		},
+		ChainSpecBuilderCmd::AddCodeSubstitute(AddCodeSubstituteCmd {
+			ref input_chain_spec,
+			ref runtime_wasm_path,
+			block_height,
+		}) => {
+			let chain_spec = ChainSpec::from_json_file(input_chain_spec.clone())?;
+
+			let mut chain_spec_json =
+				serde_json::from_str::<serde_json::Value>(&chain_spec.as_json(false)?)
+					.map_err(|e| format!("Conversion to json failed: {e}"))?;
+
+			set_code_substitute_in_json_chain_spec(
+				&mut chain_spec_json,
+				&fs::read(runtime_wasm_path.as_path())
+					.map_err(|e| format!("Wasm blob file could not be read: {e}"))?[..],
+				block_height,
+			);
+			let chain_spec_json = serde_json::to_string_pretty(&chain_spec_json)
+				.map_err(|e| format!("to pretty failed: {e}"))?;
+			fs::write(chain_spec_path, chain_spec_json).map_err(|err| err.to_string())?;
+		},
 		ChainSpecBuilderCmd::ConvertToRaw(ConvertToRawCmd { ref input_chain_spec }) => {
-			let chain_spec = GenericChainSpec::<()>::from_json_file(input_chain_spec.clone())?;
+			let chain_spec = ChainSpec::from_json_file(input_chain_spec.clone())?;
 
 			let chain_spec_json =
 				serde_json::from_str::<serde_json::Value>(&chain_spec.as_json(true)?)
@@ -77,7 +101,7 @@ fn inner_main() -> Result<(), String> {
 			fs::write(chain_spec_path, chain_spec_json).map_err(|err| err.to_string())?;
 		},
 		ChainSpecBuilderCmd::Verify(VerifyCmd { ref input_chain_spec }) => {
-			let chain_spec = GenericChainSpec::<()>::from_json_file(input_chain_spec.clone())?;
+			let chain_spec = ChainSpec::from_json_file(input_chain_spec.clone())?;
 			let _ = serde_json::from_str::<serde_json::Value>(&chain_spec.as_json(true)?)
 				.map_err(|e| format!("Conversion to json failed: {e}"))?;
 		},

substrate/bin/utils/chain-spec-builder/src/lib.rs

@@ -143,6 +143,7 @@ pub enum ChainSpecBuilderCmd {
 	ConvertToRaw(ConvertToRawCmd),
 	ListPresets(ListPresetsCmd),
 	DisplayPreset(DisplayPresetCmd),
+	AddCodeSubstitute(AddCodeSubstituteCmd),
 }
 
 /// Create a new chain spec by interacting with the provided runtime wasm blob.
@@ -222,6 +223,25 @@ pub struct UpdateCodeCmd {
 	pub runtime_wasm_path: PathBuf,
 }
 
+/// Add a code substitute in the chain spec.
+///
+/// The `codeSubstitute` object of the chain spec will be updated with the block height as key and
+/// runtime code as value. This operation supports both plain and raw formats. The `codeSubstitute`
+/// field instructs the node to use the provided runtime code at the given block height. This is
+/// useful when the chain can not progress on its own due to a bug that prevents block-building.
+///
+/// Note: For parachains, the validation function on the relaychain needs to be adjusted too,
+/// otherwise blocks built using the substituted parachain runtime will be rejected.
+#[derive(Parser, Debug, Clone)]
+pub struct AddCodeSubstituteCmd {
+	/// Chain spec to be updated.
+	pub input_chain_spec: PathBuf,
+	/// New runtime wasm blob that should replace the existing code.
+	pub runtime_wasm_path: PathBuf,
+	/// The block height at which the code should be substituted.
+	pub block_height: u64,
+}
+
 /// Converts the given chain spec into the raw format.
 #[derive(Parser, Debug, Clone)]
 pub struct ConvertToRawCmd {

substrate/client/chain-spec/src/chain_spec.rs

@@ -766,6 +766,16 @@ pub fn update_code_in_json_chain_spec(chain_spec: &mut json::Value, code: &[u8])
 	}
 }
 
+/// This function sets a codeSubstitute in the chain spec.
+pub fn set_code_substitute_in_json_chain_spec(
+	chain_spec: &mut json::Value,
+	code: &[u8],
+	block_height: u64,
+) {
+	let substitutes = json::json!({"codeSubstitutes":{ &block_height.to_string(): sp_core::bytes::to_hex(code, false) }});
+	crate::json_patch::merge(chain_spec, substitutes);
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;

substrate/client/chain-spec/src/lib.rs

@@ -333,8 +333,8 @@ pub mod json_patch;
 
 pub use self::{
 	chain_spec::{
-		update_code_in_json_chain_spec, ChainSpec as GenericChainSpec, ChainSpecBuilder,
-		NoExtension,
+		set_code_substitute_in_json_chain_spec, update_code_in_json_chain_spec,
+		ChainSpec as GenericChainSpec, ChainSpecBuilder, NoExtension,
 	},
 	extension::{get_extension, get_extension_mut, Extension, Fork, Forks, GetExtension, Group},
 	genesis_block::{