Accept only one block in validate_block when upgrading a runtime

cumulus/pallets/parachain-system/src/lib.rs

@@ -585,7 +585,7 @@ pub mod pallet {
 			// TODO: This is more than zero, but will need benchmarking to figure out what.
 			let mut total_weight = Weight::zero();
 
-			// NOTE: the inherent data is expected to be unique, even if this block is built
+			// NOTE: the inherent data is expected to be unique, even if this block is build
 			// in the context of the same relay parent as the previous one. In particular,
 			// the inherent shouldn't contain messages that were already processed by any of the
 			// ancestors.
@@ -643,7 +643,7 @@ pub mod pallet {
 				),
 			);
 
-			// initialization logic: we know that this runs exactly once every block,
+			// Initialization logic: we know that this runs exactly once every block,
 			// which means we can put the initialization logic here to remove the
 			// sequencing problem.
 			let upgrade_go_ahead_signal = relay_state_proof

cumulus/pallets/parachain-system/src/validate_block/implementation.rs

@@ -28,7 +28,7 @@ use frame_support::{
 	BoundedVec,
 };
 use polkadot_parachain_primitives::primitives::{HeadData, ValidationResult};
-use sp_core::storage::{ChildInfo, StateVersion};
+use sp_core::storage::{well_known_keys, ChildInfo, StateVersion};
 use sp_externalities::{set_and_run_with_externalities, Externalities};
 use sp_io::{hashing::blake2_128, KillStorageResult};
 use sp_runtime::traits::{
@@ -218,6 +218,10 @@ where
 			},
 		);
 
+		if overlay.storage(well_known_keys::CODE).is_some() && num_blocks > 1 {
+			panic!("When applying a runtime upgrade, only one block per PoV is allowed. Received {num_blocks}.")
+		}
+
 		run_with_externalities_and_recorder::<B, _, _>(
 			&backend,
 			&mut Default::default(),

cumulus/pallets/parachain-system/src/validate_block/tests.rs

@@ -16,11 +16,12 @@
 
 use crate::{validate_block::MemoryOptimizedValidationParams, *};
 use codec::{Decode, DecodeAll, Encode};
-use cumulus_primitives_core::{ParachainBlockData, PersistedValidationData};
+use cumulus_primitives_core::{relay_chain, ParachainBlockData, PersistedValidationData};
 use cumulus_test_client::{
 	generate_extrinsic, generate_extrinsic_with_pair,
 	runtime::{
-		self as test_runtime, Block, Hash, Header, TestPalletCall, UncheckedExtrinsic, WASM_BINARY,
+		self as test_runtime, Block, Hash, Header, SudoCall, SystemCall, TestPalletCall,
+		UncheckedExtrinsic, WASM_BINARY,
 	},
 	seal_block, transfer, BlockData, BlockOrigin, BuildParachainBlockData, Client,
 	DefaultTestClientBuilderExt, HeadData, InitBlockBuilder,
@@ -654,3 +655,89 @@ fn ensure_we_only_like_blockchains() {
 			.contains("Not a valid chain of blocks :("));
 	}
 }
+
+#[test]
+fn rejects_multiple_blocks_per_pov_when_applying_runtime_upgrade() {
+	sp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, genesis_head) = create_elastic_scaling_test_client();
+
+		let code = test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!")
+			.to_vec();
+		let code_len = code.len() as u32;
+
+		let mut proof_builder =
+			RelayStateSproofBuilder { current_slot: 1.into(), ..Default::default() };
+		proof_builder.host_config.max_code_size = code_len * 2;
+
+		// Build the block that send the runtime upgrade.
+		let TestBlockData { block: initial_block_data, .. } = build_block_with_witness(
+			&client,
+			vec![generate_extrinsic_with_pair(
+				&client,
+				Alice.into(),
+				SudoCall::sudo {
+					call: Box::new(SystemCall::set_code_without_checks { code }.into()),
+				},
+				Some(0),
+			)],
+			genesis_head.clone(),
+			proof_builder,
+			Vec::new(),
+		);
+
+		let initial_block = initial_block_data.blocks()[0].clone();
+		let (mut header, extrinsics) = initial_block.clone().deconstruct();
+		let seal = header.digest.pop().unwrap();
+
+		let mut import = BlockImportParams::new(BlockOrigin::Own, header.clone());
+		import.body = Some(extrinsics);
+		import.post_digests.push(seal);
+		import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+
+		futures::executor::block_on(BlockImport::import_block(&client, import)).unwrap();
+		let initial_block_header = initial_block.header().clone();
+
+		let mut proof_builder = RelayStateSproofBuilder {
+			current_slot: 2.into(),
+			upgrade_go_ahead: Some(relay_chain::UpgradeGoAhead::GoAhead),
+			..Default::default()
+		};
+		proof_builder.host_config.max_code_size = code_len * 2;
+
+		// 2. Build a PoV that consists of multiple blocks.
+		let TestBlockData { block: pov_block_data, validation_data: pov_validation_data } =
+			build_multiple_blocks_with_witness(
+				&client,
+				initial_block_header.clone(), // Start building PoV from the initial block's header
+				proof_builder,
+				4,
+				|_| Vec::new(),
+			);
+
+		// 3. Validate the PoV.
+		call_validate_block_elastic_scaling(
+			initial_block_header, // The parent is the head of the initial block before the PoV
+			pov_block_data,
+			pov_validation_data.relay_parent_storage_root,
+		)
+		.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args([
+				"rejects_multiple_blocks_per_pov_when_applying_runtime_upgrade",
+				"--",
+				"--nocapture",
+			])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("only one block per PoV is allowed"));
+	}
+}

cumulus/test/runtime/src/lib.rs

@@ -99,6 +99,7 @@ pub use frame_support::{
 	},
 	StorageValue,
 };
+pub use frame_system::Call as SystemCall;
 use frame_system::{
 	limits::{BlockLength, BlockWeights},
 	EnsureRoot,

prdoc/pr_10280.prdoc

@@ -0,0 +1,8 @@
+title: Accept only one block in `validate_block` when upgrading a runtime
+doc:
+- audience: Node Dev
+  description: |-
+    As the validation is running the entire time using the same validation code, we can not accept any other blocks after a runtime upgrade was applied.
+crates:
+- name: cumulus-pallet-parachain-system
+  bump: patch