Do not include :code in the storage proof if not required (#480)

This is basically a hack that prevents the inclusion of `:code` in the
storage proof. A proper fix requires some changes to the trie crate to
not include nodes that aren't read.

client/collator/src/lib.rs

@@ -333,6 +333,8 @@ mod tests {
 	use polkadot_overseer::{AllSubsystems, HeadSupportsParachains, Overseer};
 	use sp_consensus::BlockOrigin;
 	use sp_core::{testing::TaskExecutor, Pair};
+	use sp_runtime::traits::BlakeTwo256;
+	use sp_state_machine::Backend;
 
 	struct AlwaysSupportsParachains;
 	impl HeadSupportsParachains for AlwaysSupportsParachains {
@@ -376,7 +378,7 @@ mod tests {
 	}
 
 	#[test]
-	fn collates_produces_a_block() {
+	fn collates_produces_a_block_and_storage_proof_does_not_contains_code() {
 		let _ = env_logger::try_init();
 
 		let spawner = TaskExecutor::new();
@@ -432,8 +434,21 @@ mod tests {
 
 		let block_data = collation.proof_of_validity.block_data;
 
-		let block = Block::decode(&mut &block_data.0[..]).expect("Is a valid block");
+		let block =
+			ParachainBlockData::<Block>::decode(&mut &block_data.0[..]).expect("Is a valid block");
 
 		assert_eq!(1, *block.header().number());
+
+		// Ensure that we did not include `:code` in the proof.
+		let db = block.storage_proof().clone().into_memory_db();
+
+		let backend =
+			sp_state_machine::new_in_mem::<BlakeTwo256>().update_backend(*header.state_root(), db);
+
+		// Should return an error, as it was not included while building the proof.
+		assert!(backend
+			.storage(sp_core::storage::well_known_keys::CODE)
+			.unwrap_err()
+			.contains("Trie lookup error: Database missing expected key"));
 	}
 }

pallets/parachain-system/src/lib.rs

@@ -556,6 +556,18 @@ pub mod pallet {
 			matches!(call, Call::set_validation_data(_))
 		}
 	}
+
+	#[pallet::genesis_config]
+	#[derive(Default)]
+	pub struct GenesisConfig;
+
+	#[pallet::genesis_build]
+	impl<T: Config> GenesisBuild<T> for GenesisConfig {
+		fn build(&self) {
+			//TODO: Remove after https://github.com/paritytech/cumulus/issues/479
+			sp_io::storage::set(b":c", &[]);
+		}
+	}
 }
 
 impl<T: Config> Pallet<T> {

polkadot-parachains/rococo-runtime/src/lib.rs

@@ -442,7 +442,7 @@ construct_runtime! {
 		RandomnessCollectiveFlip: pallet_randomness_collective_flip::{Pallet, Call, Storage},
 		TransactionPayment: pallet_transaction_payment::{Pallet, Storage},
 
-		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Storage, Inherent, Event<T>, ValidateUnsigned} = 20,
+		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Config, Storage, Inherent, Event<T>, ValidateUnsigned} = 20,
 		ParachainInfo: parachain_info::{Pallet, Storage, Config} = 21,
 
 		Balances: pallet_balances::{Pallet, Call, Storage, Config<T>, Event<T>} = 30,

polkadot-parachains/shell-runtime/src/lib.rs

@@ -226,7 +226,7 @@ construct_runtime! {
 		UncheckedExtrinsic = UncheckedExtrinsic,
 	{
 		System: frame_system::{Pallet, Call, Storage, Config, Event<T>},
-		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Storage, Inherent, Event<T>, ValidateUnsigned},
+		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Config, Storage, Inherent, Event<T>, ValidateUnsigned},
 		ParachainInfo: parachain_info::{Pallet, Storage, Config},
 
 		// DMP handler.

polkadot-parachains/src/chain_spec.rs

@@ -178,6 +178,7 @@ fn testnet_genesis(
 			authorities: initial_authorities,
 		},
 		cumulus_pallet_aura_ext: Default::default(),
+		cumulus_pallet_parachain_system: Default::default(),
 	}
 }
 
@@ -190,6 +191,7 @@ fn shell_testnet_genesis(parachain_id: ParaId) -> shell_runtime::GenesisConfig {
 			changes_trie_config: Default::default(),
 		},
 		parachain_info: shell_runtime::ParachainInfoConfig { parachain_id },
+		cumulus_pallet_parachain_system: Default::default(),
 	}
 }
 
@@ -366,6 +368,7 @@ fn statemint_genesis(
 		// of this.
 		pallet_aura: Default::default(),
 		cumulus_pallet_aura_ext: Default::default(),
+		cumulus_pallet_parachain_system: Default::default(),
 	}
 }
 
@@ -540,6 +543,7 @@ fn statemine_genesis(
 		},
 		pallet_aura: Default::default(),
 		cumulus_pallet_aura_ext: Default::default(),
+		cumulus_pallet_parachain_system: Default::default(),
 	}
 }
 
@@ -722,5 +726,6 @@ fn westmint_genesis(
 		// of this.
 		pallet_aura: Default::default(),
 		cumulus_pallet_aura_ext: Default::default(),
+		cumulus_pallet_parachain_system: Default::default(),
 	}
 }

polkadot-parachains/statemine-runtime/src/lib.rs

@@ -672,7 +672,7 @@ construct_runtime!(
 	{
 		// System support stuff.
 		System: frame_system::{Pallet, Call, Config, Storage, Event<T>} = 0,
-		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Storage, Inherent, Event<T>} = 1,
+		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Config, Storage, Inherent, Event<T>} = 1,
 		RandomnessCollectiveFlip: pallet_randomness_collective_flip::{Pallet, Call, Storage} = 2,
 		Timestamp: pallet_timestamp::{Pallet, Call, Storage, Inherent} = 3,
 		ParachainInfo: parachain_info::{Pallet, Storage, Config} = 4,
@@ -734,8 +734,17 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPallets,
+	OnRuntimeUpgrade,
 >;
 
+pub struct OnRuntimeUpgrade;
+impl frame_support::traits::OnRuntimeUpgrade for OnRuntimeUpgrade {
+	fn on_runtime_upgrade() -> u64 {
+		sp_io::storage::set(b":c", &[]);
+		RocksDbWeight::get().writes(1)
+	}
+}
+
 impl_runtime_apis! {
 	impl sp_consensus_aura::AuraApi<Block, AuraId> for Runtime {
 		fn slot_duration() -> sp_consensus_aura::SlotDuration {

polkadot-parachains/statemint-runtime/src/lib.rs

@@ -615,7 +615,7 @@ construct_runtime!(
 	{
 		// System support stuff.
 		System: frame_system::{Pallet, Call, Config, Storage, Event<T>} = 0,
-		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Storage, Inherent, Event<T>} = 1,
+		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Config, Storage, Inherent, Event<T>} = 1,
 		RandomnessCollectiveFlip: pallet_randomness_collective_flip::{Pallet, Call, Storage} = 2,
 		Timestamp: pallet_timestamp::{Pallet, Call, Storage, Inherent} = 3,
 		ParachainInfo: parachain_info::{Pallet, Storage, Config} = 4,

polkadot-parachains/westmint-runtime/src/lib.rs

@@ -605,7 +605,7 @@ construct_runtime!(
 	{
 		// System support stuff;
 		System: frame_system::{Pallet, Call, Config, Storage, Event<T>},
-		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Storage, Inherent, Event<T>},
+		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Config, Storage, Inherent, Event<T>},
 		RandomnessCollectiveFlip: pallet_randomness_collective_flip::{Pallet, Call, Storage},
 		Timestamp: pallet_timestamp::{Pallet, Call, Storage, Inherent},
 		ParachainInfo: parachain_info::{Pallet, Storage, Config},
@@ -665,8 +665,17 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPallets,
+	OnRuntimeUpgrade
 >;
 
+pub struct OnRuntimeUpgrade;
+impl frame_support::traits::OnRuntimeUpgrade for OnRuntimeUpgrade {
+	fn on_runtime_upgrade() -> u64 {
+		sp_io::storage::set(b":c", &[]);
+		RocksDbWeight::get().writes(1)
+	}
+}
+
 impl_runtime_apis! {
 	impl sp_consensus_aura::AuraApi<Block, AuraId> for Runtime {
 		fn slot_duration() -> sp_consensus_aura::SlotDuration {

test/runtime/src/lib.rs

@@ -269,11 +269,11 @@ construct_runtime! {
 		UncheckedExtrinsic = UncheckedExtrinsic,
 	{
 		System: frame_system::{Pallet, Call, Storage, Config, Event<T>},
+		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Storage, Inherent, Event<T>, Config},
 		Timestamp: pallet_timestamp::{Pallet, Call, Storage, Inherent},
 		Balances: pallet_balances::{Pallet, Call, Storage, Config<T>, Event<T>},
 		Sudo: pallet_sudo::{Pallet, Call, Storage, Config<T>, Event<T>},
 		RandomnessCollectiveFlip: pallet_randomness_collective_flip::{Pallet, Call, Storage},
-		ParachainSystem: cumulus_pallet_parachain_system::{Pallet, Call, Storage, Inherent, Event<T>},
 		TransactionPayment: pallet_transaction_payment::{Pallet, Storage},
 	}
 }

test/service/src/chain_spec.rs

@@ -111,6 +111,7 @@ fn testnet_genesis(
 				.to_vec(),
 			..Default::default()
 		},
+		cumulus_pallet_parachain_system: Default::default(),
 		pallet_balances: cumulus_test_runtime::BalancesConfig {
 			balances: endowed_accounts
 				.iter()