Ease parachain candidate code fetching (#2593)

* code stored in para + modify CandidateDescriptor.

* WIP: digest + some more impl

* validation_code_hash in payload + check in inclusion

* check in client + refator

* tests

* fix encoding indices

* remove old todos

* fix test

* fix test

* add test

* fetch validation code inside collation-generation from the relay-chain

* HashMismatch -> PoVHashMismatch + miscompilation

* refactor, store hash when needed

* storage rename: more specific but slightly too verbose

* do not hash on candidate validation, fetch hash instead

* better test

* fix test

* guide updates

* don't panic in runtime

Co-authored-by: Robert Habermeier <rphmeier@gmail.com>

polkadot/node/collation-generation/src/lib.rs

@@ -33,7 +33,7 @@ use polkadot_node_subsystem::{
 };
 use polkadot_node_subsystem_util::{
 	request_availability_cores_ctx, request_persisted_validation_data_ctx,
-	request_validators_ctx,
+	request_validators_ctx, request_validation_code_ctx,
 	metrics::{self, prometheus},
 };
 use polkadot_primitives::v1::{
@@ -244,9 +244,10 @@ async fn handle_new_activations<Context: SubsystemContext>(
 				continue;
 			}
 
-			// we get validation data synchronously for each core instead of
+			// we get validation data and validation code synchronously for each core instead of
 			// within the subtask loop, because we have only a single mutable handle to the
 			// context, so the work can't really be distributed
+
 			let validation_data = match request_persisted_validation_data_ctx(
 				relay_parent,
 				scheduled_core.para_id,
@@ -270,6 +271,30 @@ async fn handle_new_activations<Context: SubsystemContext>(
 				}
 			};
 
+			let validation_code = match request_validation_code_ctx(
+				relay_parent,
+				scheduled_core.para_id,
+				assumption,
+				ctx,
+			)
+			.await?
+			.await??
+			{
+				Some(v) => v,
+				None => {
+					tracing::trace!(
+						target: LOG_TARGET,
+						core_idx = %core_idx,
+						relay_parent = ?relay_parent,
+						our_para = %config.para_id,
+						their_para = %scheduled_core.para_id,
+						"validation code is not available",
+					);
+					continue
+				}
+			};
+			let validation_code_hash = validation_code.hash();
+
 			let task_config = config.clone();
 			let mut task_sender = sender.clone();
 			let metrics = metrics.clone();
@@ -295,6 +320,7 @@ async fn handle_new_activations<Context: SubsystemContext>(
 					&scheduled_core.para_id,
 					&persisted_validation_data_hash,
 					&pov_hash,
+					&validation_code_hash,
 				);
 
 				let erasure_root = match erasure_root(
@@ -334,6 +360,7 @@ async fn handle_new_activations<Context: SubsystemContext>(
 						pov_hash,
 						erasure_root,
 						para_head: commitments.head_data.hash(),
+						validation_code_hash: validation_code_hash,
 					},
 				};
 
@@ -474,7 +501,7 @@ mod tests {
 		};
 		use polkadot_primitives::v1::{
 			BlockNumber, CollatorPair, Id as ParaId,
-			PersistedValidationData, ScheduledCore,
+			PersistedValidationData, ScheduledCore, ValidationCode,
 		};
 		use std::pin::Pin;
 
@@ -696,6 +723,16 @@ mod tests {
 						))) => {
 							tx.send(Ok(vec![Default::default(); 3])).unwrap();
 						}
+						Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+							_hash,
+							RuntimeApiRequest::ValidationCode(
+								_para_id,
+								OccupiedCoreAssumption::Free,
+								tx,
+							),
+						))) => {
+							tx.send(Ok(Some(ValidationCode(vec![1, 2, 3])))).unwrap();
+						}
 						Some(msg) => {
 							panic!("didn't expect any other overseer requests; got {:?}", msg)
 						}
@@ -733,11 +770,13 @@ mod tests {
 			let expect_validation_data_hash
 				= PersistedValidationData::<Hash, BlockNumber>::default().hash();
 			let expect_relay_parent = Hash::repeat_byte(4);
+			let expect_validation_code_hash = ValidationCode(vec![1, 2, 3]).hash();
 			let expect_payload = collator_signature_payload(
 				&expect_relay_parent,
 				&config.para_id,
 				&expect_validation_data_hash,
 				&expect_pov_hash,
+				&expect_validation_code_hash,
 			);
 			let expect_descriptor = CandidateDescriptor {
 				signature: config.key.sign(&expect_payload),
@@ -748,6 +787,7 @@ mod tests {
 				pov_hash: expect_pov_hash,
 				erasure_root: Default::default(), // this isn't something we're checking right now
 				para_head: test_collation().head_data.hash(),
+				validation_code_hash: expect_validation_code_hash,
 			};
 
 			assert_eq!(sent_messages.len(), 1);
@@ -768,6 +808,7 @@ mod tests {
 							&descriptor.para_id,
 							&descriptor.persisted_validation_data_hash,
 							&descriptor.pov_hash,
+							&descriptor.validation_code_hash,
 						)
 						.as_ref(),
 						&descriptor.collator,

polkadot/node/core/candidate-validation/src/lib.rs

@@ -361,21 +361,27 @@ async fn spawn_validate_exhaustive(
 
 /// Does basic checks of a candidate. Provide the encoded PoV-block. Returns `Ok` if basic checks
 /// are passed, `Err` otherwise.
-#[tracing::instrument(level = "trace", skip(pov), fields(subsystem = LOG_TARGET))]
+#[tracing::instrument(level = "trace", skip(pov, validation_code), fields(subsystem = LOG_TARGET))]
 fn perform_basic_checks(
 	candidate: &CandidateDescriptor,
 	max_pov_size: u32,
 	pov: &PoV,
+	validation_code: &ValidationCode,
 ) -> Result<(), InvalidCandidate> {
 	let encoded_pov = pov.encode();
-	let hash = pov.hash();
+	let pov_hash = pov.hash();
+	let validation_code_hash = validation_code.hash();
 
 	if encoded_pov.len() > max_pov_size as usize {
 		return Err(InvalidCandidate::ParamsTooLarge(encoded_pov.len() as u64));
 	}
 
-	if hash != candidate.pov_hash {
-		return Err(InvalidCandidate::HashMismatch);
+	if pov_hash != candidate.pov_hash {
+		return Err(InvalidCandidate::PoVHashMismatch);
+	}
+
+	if validation_code_hash != candidate.validation_code_hash {
+		return Err(InvalidCandidate::CodeHashMismatch);
 	}
 
 	if let Err(()) = candidate.check_collator_signature() {
@@ -431,7 +437,12 @@ fn validate_candidate_exhaustive<B: ValidationBackend, S: SpawnNamed + 'static>(
 ) -> Result<ValidationResult, ValidationFailed> {
 	let _timer = metrics.time_validate_candidate_exhaustive();
 
-	if let Err(e) = perform_basic_checks(&descriptor, persisted_validation_data.max_pov_size, &*pov) {
+	if let Err(e) = perform_basic_checks(
+		&descriptor,
+		persisted_validation_data.max_pov_size,
+		&*pov,
+		&validation_code,
+	) {
 		return Ok(ValidationResult::Invalid(e))
 	}
 
@@ -601,6 +612,7 @@ mod tests {
 			&descriptor.para_id,
 			&descriptor.persisted_validation_data_hash,
 			&descriptor.pov_hash,
+			&descriptor.validation_code_hash,
 		);
 
 		descriptor.signature = collator.sign(&payload[..]).into();
@@ -891,13 +903,21 @@ mod tests {
 
 		let pov = PoV { block_data: BlockData(vec![1; 32]) };
 		let head_data = HeadData(vec![1, 1, 1]);
+		let validation_code = ValidationCode(vec![2; 16]);
 
 		let mut descriptor = CandidateDescriptor::default();
 		descriptor.pov_hash = pov.hash();
 		descriptor.para_head = head_data.hash();
+		descriptor.validation_code_hash = validation_code.hash();
 		collator_sign(&mut descriptor, Sr25519Keyring::Alice);
 
-		assert!(perform_basic_checks(&descriptor, validation_data.max_pov_size, &pov).is_ok());
+		let check = perform_basic_checks(
+			&descriptor,
+			validation_data.max_pov_size,
+			&pov,
+			&validation_code,
+		);
+		assert!(check.is_ok());
 
 		let validation_result = WasmValidationResult {
 			head_data,
@@ -911,7 +931,7 @@ mod tests {
 		let v = validate_candidate_exhaustive::<MockValidationBackend, _>(
 			MockValidationArg { result: Ok(validation_result) },
 			validation_data.clone(),
-			vec![1, 2, 3].into(),
+			validation_code,
 			descriptor,
 			Arc::new(pov),
 			TaskExecutor::new(),
@@ -933,12 +953,20 @@ mod tests {
 		let validation_data = PersistedValidationData { max_pov_size: 1024, ..Default::default() };
 
 		let pov = PoV { block_data: BlockData(vec![1; 32]) };
+		let validation_code = ValidationCode(vec![2; 16]);
 
 		let mut descriptor = CandidateDescriptor::default();
 		descriptor.pov_hash = pov.hash();
+		descriptor.validation_code_hash = validation_code.hash();
 		collator_sign(&mut descriptor, Sr25519Keyring::Alice);
 
-		assert!(perform_basic_checks(&descriptor, validation_data.max_pov_size, &pov).is_ok());
+		let check = perform_basic_checks(
+			&descriptor,
+			validation_data.max_pov_size,
+			&pov,
+			&validation_code,
+		);
+		assert!(check.is_ok());
 
 		let v = validate_candidate_exhaustive::<MockValidationBackend, _>(
 			MockValidationArg {
@@ -947,7 +975,7 @@ mod tests {
 				))
 			},
 			validation_data,
-			vec![1, 2, 3].into(),
+			validation_code,
 			descriptor,
 			Arc::new(pov),
 			TaskExecutor::new(),
@@ -962,12 +990,20 @@ mod tests {
 		let validation_data = PersistedValidationData { max_pov_size: 1024, ..Default::default() };
 
 		let pov = PoV { block_data: BlockData(vec![1; 32]) };
+		let validation_code = ValidationCode(vec![2; 16]);
 
 		let mut descriptor = CandidateDescriptor::default();
 		descriptor.pov_hash = pov.hash();
+		descriptor.validation_code_hash = validation_code.hash();
 		collator_sign(&mut descriptor, Sr25519Keyring::Alice);
 
-		assert!(perform_basic_checks(&descriptor, validation_data.max_pov_size, &pov).is_ok());
+		let check = perform_basic_checks(
+			&descriptor,
+			validation_data.max_pov_size,
+			&pov,
+			&validation_code,
+		);
+		assert!(check.is_ok());
 
 		let v = validate_candidate_exhaustive::<MockValidationBackend, _>(
 			MockValidationArg {
@@ -976,7 +1012,7 @@ mod tests {
 				))
 			},
 			validation_data,
-			vec![1, 2, 3].into(),
+			validation_code,
 			descriptor,
 			Arc::new(pov),
 			TaskExecutor::new(),
@@ -985,4 +1021,42 @@ mod tests {
 
 		assert_matches!(v, Ok(ValidationResult::Invalid(InvalidCandidate::Timeout)));
 	}
+
+	#[test]
+	fn candidate_validation_code_mismatch_is_invalid() {
+		let validation_data = PersistedValidationData { max_pov_size: 1024, ..Default::default() };
+
+		let pov = PoV { block_data: BlockData(vec![1; 32]) };
+		let validation_code = ValidationCode(vec![2; 16]);
+
+		let mut descriptor = CandidateDescriptor::default();
+		descriptor.pov_hash = pov.hash();
+		descriptor.validation_code_hash = ValidationCode(vec![1; 16]).hash();
+		collator_sign(&mut descriptor, Sr25519Keyring::Alice);
+
+		let check = perform_basic_checks(
+			&descriptor,
+			validation_data.max_pov_size,
+			&pov,
+			&validation_code,
+		);
+		assert_matches!(check, Err(InvalidCandidate::CodeHashMismatch));
+
+		let v = validate_candidate_exhaustive::<MockValidationBackend, _>(
+			MockValidationArg {
+				result: Err(ValidationError::InvalidCandidate(
+					WasmInvalidCandidate::BadReturn
+				))
+			},
+			validation_data,
+			validation_code,
+			descriptor,
+			Arc::new(pov),
+			TaskExecutor::new(),
+			&Default::default(),
+		).unwrap();
+
+		assert_matches!(v, ValidationResult::Invalid(InvalidCandidate::CodeHashMismatch));
+	}
+
 }

polkadot/node/core/runtime-api/src/lib.rs

@@ -556,6 +556,13 @@ mod tests {
 			) -> BTreeMap<ParaId, Vec<InboundHrmpMessage>> {
 				self.hrmp_channels.get(&recipient).map(|q| q.clone()).unwrap_or_default()
 			}
+
+			fn validation_code_by_hash(
+				&self,
+				_hash: Hash,
+			) -> Option<ValidationCode> {
+				unreachable!("not used in tests");
+			}
 		}
 
 		impl BabeApi<Block> for MockRuntimeApi {

polkadot/node/primitives/src/lib.rs

@@ -124,11 +124,13 @@ pub enum InvalidCandidate {
 	/// Invalid relay chain parent.
 	BadParent,
 	/// POV hash does not match.
-	HashMismatch,
+	PoVHashMismatch,
 	/// Bad collator signature.
 	BadSignature,
 	/// Para head hash does not match.
 	ParaHeadHashMismatch,
+	/// Validation code hash does not match.
+	CodeHashMismatch,
 }
 
 /// Result of the validation of the candidate.

polkadot/parachain/src/primitives.rs

@@ -21,6 +21,7 @@ use sp_std::vec::Vec;
 
 use parity_scale_codec::{Encode, Decode, CompactAs};
 use sp_core::{RuntimeDebug, TypeId};
+use sp_runtime::traits::Hash as _;
 
 #[cfg(feature = "std")]
 use serde::{Serialize, Deserialize};
@@ -45,7 +46,6 @@ pub struct HeadData(#[cfg_attr(feature = "std", serde(with="bytes"))] pub Vec<u8
 impl HeadData {
 	/// Returns the hash of this head data.
 	pub fn hash(&self) -> Hash {
-		use sp_runtime::traits::Hash;
 		sp_runtime::traits::BlakeTwo256::hash(&self.0)
 	}
 }
@@ -55,6 +55,13 @@ impl HeadData {
 #[cfg_attr(feature = "std", derive(Serialize, Deserialize, Hash, MallocSizeOf))]
 pub struct ValidationCode(#[cfg_attr(feature = "std", serde(with="bytes"))] pub Vec<u8>);
 
+impl ValidationCode {
+	/// Get the blake2-256 hash of the validation code bytes.
+	pub fn hash(&self) -> Hash {
+		sp_runtime::traits::BlakeTwo256::hash(&self.0[..])
+	}
+}
+
 /// Parachain block data.
 ///
 /// Contains everything required to validate para-block, may contain block and witness data.

polkadot/primitives/src/v1.rs

@@ -211,14 +211,16 @@ pub fn collator_signature_payload<H: AsRef<[u8]>>(
 	para_id: &Id,
 	persisted_validation_data_hash: &Hash,
 	pov_hash: &Hash,
-) -> [u8; 100] {
+	validation_code_hash: &Hash,
+) -> [u8; 132] {
 	// 32-byte hash length is protected in a test below.
-	let mut payload = [0u8; 100];
+	let mut payload = [0u8; 132];
 
 	payload[0..32].copy_from_slice(relay_parent.as_ref());
 	u32::from(*para_id).using_encoded(|s| payload[32..32 + s.len()].copy_from_slice(s));
 	payload[36..68].copy_from_slice(persisted_validation_data_hash.as_ref());
 	payload[68..100].copy_from_slice(pov_hash.as_ref());
+	payload[100..132].copy_from_slice(validation_code_hash.as_ref());
 
 	payload
 }
@@ -228,6 +230,7 @@ fn check_collator_signature<H: AsRef<[u8]>>(
 	para_id: &Id,
 	persisted_validation_data_hash: &Hash,
 	pov_hash: &Hash,
+	validation_code_hash: &Hash,
 	collator: &CollatorId,
 	signature: &CollatorSignature,
 ) -> Result<(),()> {
@@ -236,6 +239,7 @@ fn check_collator_signature<H: AsRef<[u8]>>(
 		para_id,
 		persisted_validation_data_hash,
 		pov_hash,
+		validation_code_hash,
 	);
 
 	if signature.verify(&payload[..], collator) {
@@ -268,6 +272,8 @@ pub struct CandidateDescriptor<H = Hash> {
 	pub signature: CollatorSignature,
 	/// Hash of the para header that is being generated by this candidate.
 	pub para_head: Hash,
+	/// The blake2-256 hash of the validation code bytes.
+	pub validation_code_hash: Hash,
 }
 
 impl<H: AsRef<[u8]>> CandidateDescriptor<H> {
@@ -278,6 +284,7 @@ impl<H: AsRef<[u8]>> CandidateDescriptor<H> {
 			&self.para_id,
 			&self.persisted_validation_data_hash,
 			&self.pov_hash,
+			&self.validation_code_hash,
 			&self.collator,
 			&self.signature,
 		)
@@ -872,6 +879,10 @@ sp_api::decl_runtime_apis! {
 		/// messages in them are also included.
 		#[skip_initialize_block]
 		fn inbound_hrmp_channels_contents(recipient: Id) -> BTreeMap<Id, Vec<InboundHrmpMessage<N>>>;
+
+		/// Get the validation code from its hash.
+		#[skip_initialize_block]
+		fn validation_code_by_hash(hash: Hash) -> Option<ValidationCode>;
 	}
 }
 
@@ -1107,6 +1118,7 @@ mod tests {
 			&5u32.into(),
 			&Hash::repeat_byte(2),
 			&Hash::repeat_byte(3),
+			&Hash::repeat_byte(4),
 		);
 	}
 }

polkadot/roadmap/implementers-guide/src/runtime/paras.md

@@ -116,10 +116,10 @@ Parachains: Vec<ParaId>,
 ParaLifecycle: map ParaId => Option<ParaLifecycle>,
 /// The head-data of every registered para.
 Heads: map ParaId => Option<HeadData>;
-/// The validation code of every live para.
-ValidationCode: map ParaId => Option<ValidationCode>;
-/// Actual past code, indicated by the para id as well as the block number at which it became outdated.
-PastCode: map (ParaId, BlockNumber) => Option<ValidationCode>;
+/// The validation code hash of every live para.
+CurrentCodeHash: map ParaId => Option<Hash>;
+/// Actual past code hash, indicated by the para id as well as the block number at which it became outdated.
+PastCodeHash: map (ParaId, BlockNumber) => Option<Hash>;
 /// Past code of parachains. The parachains themselves may not be registered anymore,
 /// but we also keep their code on-chain for the same amount of time as outdated code
 /// to keep it available for secondary checkers.
@@ -136,24 +136,28 @@ PastCodePruning: Vec<(ParaId, BlockNumber)>;
 /// in the context of a relay chain block with a number >= `expected_at`.
 FutureCodeUpgrades: map ParaId => Option<BlockNumber>;
 /// The actual future code of a para.
-FutureCode: map ParaId => Option<ValidationCode>;
+FutureCodeHash: map ParaId => Option<Hash>;
 /// The actions to perform during the start of a specific session index.
 ActionsQueue: map SessionIndex => Vec<ParaId>;
 /// Upcoming paras instantiation arguments.
 UpcomingParasGenesis: map ParaId => Option<ParaGenesisArgs>;
+/// The number of references on the validation code in `CodeByHash` storage.
+CodeByHashRefs: map Hash => u32;
+/// Validation code stored by its hash.
+CoeByHash: map Hash => Option<ValidationCode>
 ```
 
 ## Session Change
 
 1. Execute all queued actions for paralifecycle changes:
   1. Clean up outgoing paras.
-     1. This means removing the entries under `Heads`, `ValidationCode`, `FutureCodeUpgrades`, and
+     1. This means removing the entries under `Heads`, `CurrentCode`, `FutureCodeUpgrades`, and
         `FutureCode`. An according entry should be added to `PastCode`, `PastCodeMeta`, and
-        `PastCodePruning` using the outgoing `ParaId` and removed `ValidationCode` value. This is
+        `PastCodePruning` using the outgoing `ParaId` and removed `CurrentCode` value. This is
         because any outdated validation code must remain available on-chain for a determined amount
         of blocks, and validation code outdated by de-registering the para is still subject to that
         invariant.
-  1. Apply all incoming paras by initializing the `Heads` and `ValidationCode` using the genesis
+  1. Apply all incoming paras by initializing the `Heads` and `CurrentCode` using the genesis
      parameters.
   1. Amend the `Parachains` list and `ParaLifecycle` to reflect changes in registered parachains.
   1. Amend the `ParaLifecycle` set to reflect changes in registered parathreads.
@@ -175,7 +179,7 @@ UpcomingParasGenesis: map ParaId => Option<ParaGenesisArgs>;
 * `schedule_para_cleanup(ParaId)`: Schedule a para to be cleaned up after the next full session.
 * `schedule_parathread_upgrade(ParaId)`: Schedule a parathread to be upgraded to a parachain.
 * `schedule_parachain_downgrade(ParaId)`: Schedule a parachain to be downgraded to a parathread.
-* `schedule_code_upgrade(ParaId, ValidationCode, expected_at: BlockNumber)`: Schedule a future code
+* `schedule_code_upgrade(ParaId, CurrentCode, expected_at: BlockNumber)`: Schedule a future code
   upgrade of the given parachain, to be applied after inclusion of a block of the same parachain
   executed in the context of a relay-chain block with number >= `expected_at`.
 * `note_new_head(ParaId, HeadData, BlockNumber)`: note that a para has progressed to a new head,
@@ -187,6 +191,7 @@ UpcomingParasGenesis: map ParaId => Option<ParaGenesisArgs>;
   intermediate parablock has been included at the given relay-chain height. This may return past,
   current, or (with certain choices of `assume_intermediate`) future code. `assume_intermediate`, if
   provided, must be before `at`. If the validation code has been pruned, this will return `None`.
+* `validation_code_hash_at(ParaId, at: BlockNumber, assume_intermediate: Option<BlockNumber>)`: Just like `validation_code_at`, but returns the code hash.
 * `lifecycle(ParaId) -> Option<ParaLifecycle>`: Return the `ParaLifecycle` of a para.
 * `is_parachain(ParaId) -> bool`: Returns true if the para ID references any live parachain,
   including those which may be transitioning to a parathread in the future.
@@ -197,9 +202,6 @@ UpcomingParasGenesis: map ParaId => Option<ParaGenesisArgs>;
 * `last_code_upgrade(id: ParaId, include_future: bool) -> Option<BlockNumber>`: The block number of
   the last scheduled upgrade of the requested para. Includes future upgrades if the flag is set.
   This is the `expected_at` number, not the `activated_at` number.
-* `persisted_validation_data(id: ParaId) -> Option<PersistedValidationData>`: Get the
-  PersistedValidationData of the given para, assuming the context is the parent block. Returns
-  `None` if the para is not known.
 
 ## Finalization
 

polkadot/runtime/kusama/src/lib.rs

@@ -1218,6 +1218,10 @@ sp_api::impl_runtime_apis! {
 		) -> BTreeMap<Id, Vec<InboundHrmpMessage<BlockNumber>>> {
 			BTreeMap::new()
 		}
+
+		fn validation_code_by_hash(_hash: Hash) -> Option<ValidationCode> {
+			None
+		}
 	}
 
 	impl beefy_primitives::BeefyApi<Block, BeefyId> for Runtime {

polkadot/runtime/parachains/src/inclusion.rs

@@ -185,6 +185,8 @@ decl_error! {
 		HrmpWatermarkMishandling,
 		/// The HRMP messages sent by the candidate is not valid.
 		InvalidOutboundHrmp,
+		/// The validation code hash of the candidate is not valid.
+		InvalidValidationCodeHash,
 	}
 }
 
@@ -444,6 +446,15 @@ impl<T: Config> Module<T> {
 					Error::<T>::NotCollatorSigned,
 				);
 
+				let validation_code_hash =
+					<paras::Module<T>>::validation_code_hash_at(para_id, now, None)
+					// A candidate for a parachain without current validation code is not scheduled.
+					.ok_or_else(|| Error::<T>::UnscheduledCandidate)?;
+				ensure!(
+					candidate.descriptor().validation_code_hash == validation_code_hash,
+					Error::<T>::InvalidValidationCodeHash,
+				);
+
 				if let Err(err) = check_cx
 					.check_validation_outputs(
 						para_id,
@@ -954,6 +965,7 @@ mod tests {
 			&candidate.descriptor.para_id,
 			&candidate.descriptor.persisted_validation_data_hash,
 			&candidate.descriptor.pov_hash,
+			&candidate.descriptor.validation_code_hash,
 		);
 
 		candidate.descriptor.signature = collator.sign(&payload[..]).into();
@@ -1109,6 +1121,7 @@ mod tests {
 		relay_parent: Hash,
 		persisted_validation_data_hash: Hash,
 		new_validation_code: Option<ValidationCode>,
+		validation_code: ValidationCode,
 		hrmp_watermark: BlockNumber,
 	}
 
@@ -1120,6 +1133,7 @@ mod tests {
 					pov_hash: self.pov_hash,
 					relay_parent: self.relay_parent,
 					persisted_validation_data_hash: self.persisted_validation_data_hash,
+					validation_code_hash: self.validation_code.hash(),
 					..Default::default()
 				},
 				commitments: CandidateCommitments {
@@ -2071,6 +2085,43 @@ mod tests {
 					Err(Error::<Test>::ValidationDataHashMismatch.into()),
 				);
 			}
+
+			// bad validation code hash
+			{
+				let mut candidate = TestCandidateBuilder {
+					para_id: chain_a,
+					relay_parent: System::parent_hash(),
+					pov_hash: Hash::repeat_byte(1),
+					persisted_validation_data_hash: make_vdata_hash(chain_a).unwrap(),
+					hrmp_watermark: RELAY_PARENT_NUM,
+					validation_code: ValidationCode(vec![1]),
+					..Default::default()
+				}.build();
+
+				collator_sign_candidate(
+					Sr25519Keyring::One,
+					&mut candidate,
+				);
+
+				let backed = block_on(back_candidate(
+					candidate,
+					&validators,
+					group_validators(GroupIndex::from(0)).unwrap().as_ref(),
+					&keystore,
+					&signing_context,
+					BackingKind::Threshold,
+				));
+
+				assert_eq!(
+					Inclusion::process_candidates(
+						Default::default(),
+						vec![backed],
+						vec![chain_a_assignment.clone()],
+						&group_validators,
+					),
+					Err(Error::<Test>::InvalidValidationCodeHash.into()),
+				);
+			}
 		});
 	}
 

polkadot/runtime/parachains/src/paras.rs

@@ -28,9 +28,9 @@ use sp_std::result;
 #[cfg(feature = "std")]
 use sp_std::marker::PhantomData;
 use primitives::v1::{
-	Id as ParaId, ValidationCode, HeadData, SessionIndex,
+	Id as ParaId, ValidationCode, HeadData, SessionIndex, Hash, ConsensusLog,
 };
-use sp_runtime::{traits::One, DispatchResult};
+use sp_runtime::{traits::One, DispatchResult, SaturatedConversion};
 use frame_system::ensure_root;
 use frame_support::{
 	decl_storage, decl_module, decl_error, decl_event, ensure,
@@ -82,7 +82,7 @@ pub struct ParaPastCodeMeta<N> {
 	/// was actually replaced, respectively. The first is used to do accurate lookups
 	/// of historic code in historic contexts, whereas the second is used to do
 	/// pruning on an accurate timeframe. These can be used as indices
-	/// into the `PastCode` map along with the `ParaId` to fetch the code itself.
+	/// into the `PastCodeHash` map along with the `ParaId` to fetch the code itself.
 	upgrade_times: Vec<ReplacementTimes<N>>,
 	/// Tracks the highest pruned code-replacement, if any. This is the `expected_at` value,
 	/// not the `activated_at` value.
@@ -261,10 +261,15 @@ decl_storage! {
 		ParaLifecycles: map hasher(twox_64_concat) ParaId => Option<ParaLifecycle>;
 		/// The head-data of every registered para.
 		Heads get(fn para_head): map hasher(twox_64_concat) ParaId => Option<HeadData>;
-		/// The validation code of every live para.
-		CurrentCode get(fn current_code): map hasher(twox_64_concat) ParaId => Option<ValidationCode>;
-		/// Actual past code, indicated by the para id as well as the block number at which it became outdated.
-		PastCode: map hasher(twox_64_concat) (ParaId, T::BlockNumber) => Option<ValidationCode>;
+		/// The validation code hash of every live para.
+		///
+		/// Corresponding code can be retrieved with [`CodeByHash`].
+		CurrentCodeHash: map hasher(twox_64_concat) ParaId => Option<Hash>;
+		/// Actual past code hash, indicated by the para id as well as the block number at which it
+		/// became outdated.
+		///
+		/// Corresponding code can be retrieved with [`CodeByHash`].
+		PastCodeHash: map hasher(twox_64_concat) (ParaId, T::BlockNumber) => Option<Hash>;
 		/// Past code of parachains. The parachains themselves may not be registered anymore,
 		/// but we also keep their code on-chain for the same amount of time as outdated code
 		/// to keep it available for secondary checkers.
@@ -281,12 +286,21 @@ decl_storage! {
 		/// The change will be applied after the first parablock for this ID included which executes
 		/// in the context of a relay chain block with a number >= `expected_at`.
 		FutureCodeUpgrades get(fn future_code_upgrade_at): map hasher(twox_64_concat) ParaId => Option<T::BlockNumber>;
-		/// The actual future code of a para.
-		FutureCode: map hasher(twox_64_concat) ParaId => Option<ValidationCode>;
+		/// The actual future code hash of a para.
+		///
+		/// Corresponding code can be retrieved with [`CodeByHash`].
+		FutureCodeHash: map hasher(twox_64_concat) ParaId => Option<Hash>;
 		/// The actions to perform during the start of a specific session index.
 		ActionsQueue get(fn actions_queue): map hasher(twox_64_concat) SessionIndex => Vec<ParaId>;
 		/// Upcoming paras instantiation arguments.
 		UpcomingParasGenesis: map hasher(twox_64_concat) ParaId => Option<ParaGenesisArgs>;
+		/// The number of reference on the validation code in [`CodeByHash`] storage.
+		CodeByHashRefs: map hasher(identity) Hash => u32;
+		/// Validation code stored by its hash.
+		///
+		/// This storage is consistent with [`FutureCodeHash`], [`CurrentCodeHash`] and
+		/// [`PastCodeHash`].
+		CodeByHash get(fn code_by_hash): map hasher(identity) Hash => Option<ValidationCode>;
 	}
 	add_extra_genesis {
 		config(paras): Vec<(ParaId, ParaGenesisArgs)>;
@@ -310,7 +324,9 @@ fn build<T: Config>(config: &GenesisConfig<T>) {
 	Parachains::put(&parachains);
 
 	for (id, genesis_args) in &config.paras {
-		<Module<T> as Store>::CurrentCode::insert(&id, &genesis_args.validation_code);
+		let code_hash = genesis_args.validation_code.hash();
+		<Module<T>>::increase_code_ref(&code_hash, &genesis_args.validation_code);
+		<Module<T> as Store>::CurrentCodeHash::insert(&id, &code_hash);
 		<Module<T> as Store>::Heads::insert(&id, &genesis_args.genesis_head);
 		if genesis_args.parachain {
 			ParaLifecycles::insert(&id, ParaLifecycle::Parachain);
@@ -361,11 +377,13 @@ decl_module! {
 		#[weight = 0]
 		fn force_set_current_code(origin, para: ParaId, new_code: ValidationCode) {
 			ensure_root(origin)?;
-			let prior_code = <Self as Store>::CurrentCode::get(&para).unwrap_or_default();
-			<Self as Store>::CurrentCode::insert(&para, new_code);
+			let prior_code_hash = <Self as Store>::CurrentCodeHash::get(&para).unwrap_or_default();
+			let new_code_hash = new_code.hash();
+			Self::increase_code_ref(&new_code_hash, &new_code);
+			<Self as Store>::CurrentCodeHash::insert(&para, new_code_hash);
 
 			let now = frame_system::Pallet::<T>::block_number();
-			Self::note_past_code(para, now, now, prior_code);
+			Self::note_past_code(para, now, now, prior_code_hash);
 			Self::deposit_event(Event::CurrentCodeUpdated(para));
 		}
 
@@ -428,6 +446,21 @@ impl<T: Config> Module<T> {
 		outgoing_paras
 	}
 
+	/// The validation code of live para.
+	pub(crate) fn current_code(para_id: &ParaId) -> Option<ValidationCode> {
+		CurrentCodeHash::get(para_id).and_then(|code_hash| {
+			let code = CodeByHash::get(&code_hash);
+			if code.is_none() {
+				log::error!(
+					"Pallet paras storage is inconsistent, code not found for hash {}",
+					code_hash,
+				);
+				debug_assert!(false, "inconsistent paras storages");
+			}
+			code
+		})
+	}
+
 	// Apply all para actions queued for the given session index.
 	//
 	// The actions to take are based on the lifecycle of of the paras.
@@ -458,8 +491,10 @@ impl<T: Config> Module<T> {
 							ParaLifecycles::insert(&para, ParaLifecycle::Parathread);
 						}
 
+						let code_hash = genesis_data.validation_code.hash();
 						<Self as Store>::Heads::insert(&para, genesis_data.genesis_head);
-						<Self as Store>::CurrentCode::insert(&para, genesis_data.validation_code);
+						Self::increase_code_ref(&code_hash, &genesis_data.validation_code);
+						<Self as Store>::CurrentCodeHash::insert(&para, code_hash);
 					}
 				},
 				// Upgrade a parathread to a parachain
@@ -484,12 +519,15 @@ impl<T: Config> Module<T> {
 
 					<Self as Store>::Heads::remove(&para);
 					<Self as Store>::FutureCodeUpgrades::remove(&para);
-					<Self as Store>::FutureCode::remove(&para);
 					ParaLifecycles::remove(&para);
+					let removed_future_code_hash = <Self as Store>::FutureCodeHash::take(&para);
+					if let Some(removed_future_code_hash) = removed_future_code_hash {
+						Self::decrease_code_ref(&removed_future_code_hash);
+					}
 
-					let removed_code = <Self as Store>::CurrentCode::take(&para);
-					if let Some(removed_code) = removed_code {
-						Self::note_past_code(para, now, now, removed_code);
+					let removed_code_hash = <Self as Store>::CurrentCodeHash::take(&para);
+					if let Some(removed_code_hash) = removed_code_hash {
+						Self::note_past_code(para, now, now, removed_code_hash);
 					}
 
 					outgoing.push(para);
@@ -513,14 +551,14 @@ impl<T: Config> Module<T> {
 		id: ParaId,
 		at: T::BlockNumber,
 		now: T::BlockNumber,
-		old_code: ValidationCode,
+		old_code_hash: Hash,
 	) -> Weight {
 
 		<Self as Store>::PastCodeMeta::mutate(&id, |past_meta| {
 			past_meta.note_replacement(at, now);
 		});
 
-		<Self as Store>::PastCode::insert(&(id, at), old_code);
+		<Self as Store>::PastCodeHash::insert(&(id, at), old_code_hash);
 
 		// Schedule pruning for this past-code to be removed as soon as it
 		// exits the slashing window.
@@ -559,7 +597,18 @@ impl<T: Config> Module<T> {
 				for (para_id, _) in pruning_tasks_to_do {
 					let full_deactivate = <Self as Store>::PastCodeMeta::mutate(&para_id, |meta| {
 						for pruned_repl_at in meta.prune_up_to(pruning_height) {
-							<Self as Store>::PastCode::remove(&(para_id, pruned_repl_at));
+							let removed_code_hash =
+								<Self as Store>::PastCodeHash::take(&(para_id, pruned_repl_at));
+
+							if let Some(removed_code_hash) = removed_code_hash {
+								Self::decrease_code_ref(&removed_code_hash);
+							} else {
+								log::warn!(
+									target: "runtime::paras",
+									"Missing code for removed hash {:?}",
+									removed_code_hash,
+								);
+							}
 						}
 
 						meta.most_recent_change().is_none() && Self::para_head(&para_id).is_none()
@@ -689,8 +738,15 @@ impl<T: Config> Module<T> {
 				T::DbWeight::get().reads_writes(1, 0)
 			} else {
 				*up = Some(expected_at);
-				FutureCode::insert(&id, new_code);
-				T::DbWeight::get().reads_writes(1, 2)
+
+				let new_code_hash = new_code.hash();
+				let expected_at_u32 = expected_at.saturated_into();
+				let log = ConsensusLog::ParaScheduleUpgradeCode(id, new_code_hash, expected_at_u32);
+				<frame_system::Pallet<T>>::deposit_log(log.into());
+
+				let (reads, writes) = Self::increase_code_ref(&new_code_hash, &new_code);
+				FutureCodeHash::insert(&id, new_code_hash);
+				T::DbWeight::get().reads_writes(1 + reads, 2 + writes)
 			}
 		})
 	}
@@ -710,9 +766,12 @@ impl<T: Config> Module<T> {
 				<Self as Store>::FutureCodeUpgrades::remove(&id);
 
 				// Both should always be `Some` in this case, since a code upgrade is scheduled.
-				let new_code = FutureCode::take(&id).unwrap_or_default();
-				let prior_code = CurrentCode::get(&id).unwrap_or_default();
-				CurrentCode::insert(&id, &new_code);
+				let new_code_hash = FutureCodeHash::take(&id).unwrap_or_default();
+				let prior_code_hash = CurrentCodeHash::get(&id).unwrap_or_default();
+				CurrentCodeHash::insert(&id, &new_code_hash);
+
+				let log = ConsensusLog::ParaUpgradeCode(id, new_code_hash);
+				<frame_system::Pallet<T>>::deposit_log(log.into());
 
 				// `now` is only used for registering pruning as part of `fn note_past_code`
 				let now = <frame_system::Pallet<T>>::block_number();
@@ -721,7 +780,7 @@ impl<T: Config> Module<T> {
 					id,
 					expected_at,
 					now,
-					prior_code,
+					prior_code_hash,
 				);
 
 				// add 1 to writes due to heads update.
@@ -734,19 +793,22 @@ impl<T: Config> Module<T> {
 		}
 	}
 
-	/// Fetches the validation code to be used when validating a block in the context of the given
-	/// relay-chain height. A second block number parameter may be used to tell the lookup to proceed
-	/// as if an intermediate parablock has been with the given relay-chain height as its context.
-	/// This may return past, current, or (with certain choices of `assume_intermediate`) future code.
+	/// Fetches the validation code hash for the validation code to be used when validating a block
+	/// in the context of the given relay-chain height. A second block number parameter may be used
+	/// to tell the lookup to proceed as if an intermediate parablock has been with the given
+	/// relay-chain height as its context. This may return the hash for the past, current, or
+	/// (with certain choices of `assume_intermediate`) future code.
 	///
 	/// `assume_intermediate`, if provided, must be before `at`. This will return `None` if the validation
 	/// code has been pruned.
+	///
+	/// To get associated code see [`Self::validation_code_at`].
 	#[allow(unused)]
-	pub(crate) fn validation_code_at(
+	pub(crate) fn validation_code_hash_at(
 		id: ParaId,
 		at: T::BlockNumber,
 		assume_intermediate: Option<T::BlockNumber>,
-	) -> Option<ValidationCode> {
+	) -> Option<Hash> {
 		let now = <frame_system::Pallet<T>>::block_number();
 		let config = <configuration::Module<T>>::config();
 
@@ -761,16 +823,36 @@ impl<T: Config> Module<T> {
 		};
 
 		if upgrade_applied_intermediate {
-			FutureCode::get(&id)
+			FutureCodeHash::get(&id)
 		} else {
 			match Self::past_code_meta(&id).code_at(at) {
 				None => None,
-				Some(UseCodeAt::Current) => CurrentCode::get(&id),
-				Some(UseCodeAt::ReplacedAt(replaced)) => <Self as Store>::PastCode::get(&(id, replaced))
+				Some(UseCodeAt::Current) => CurrentCodeHash::get(&id),
+				Some(UseCodeAt::ReplacedAt(replaced)) => <Self as Store>::PastCodeHash::get(&(id, replaced)),
 			}
 		}
 	}
 
+	/// Fetch validation code of para in specific context, see [`Self::validation_code_hash_at`].
+	#[allow(unused)]
+	pub(crate) fn validation_code_at(
+		id: ParaId,
+		at: T::BlockNumber,
+		assume_intermediate: Option<T::BlockNumber>,
+	) -> Option<ValidationCode> {
+		Self::validation_code_hash_at(id, at, assume_intermediate).and_then(|code_hash| {
+			let code = CodeByHash::get(&code_hash);
+			if code.is_none() {
+				log::error!(
+					"Pallet paras storage is inconsistent, code not found for hash {}",
+					code_hash,
+				);
+				debug_assert!(false, "inconsistent paras storages");
+			}
+			code
+		})
+	}
+
 	/// Returns the current lifecycle state of the para.
 	pub fn lifecycle(id: ParaId) -> Option<ParaLifecycle> {
 		ParaLifecycles::get(&id)
@@ -826,6 +908,34 @@ impl<T: Config> Module<T> {
 		shared::Module::<T>::scheduled_session()
 	}
 
+	/// Store the validation code if not already stored, and increase the number of reference.
+	///
+	/// Returns the number of storage reads and number of storage writes.
+	fn increase_code_ref(code_hash: &Hash, code: &ValidationCode) -> (u64, u64) {
+		let reads = 1;
+		let mut writes = 1;
+		<Self as Store>::CodeByHashRefs::mutate(code_hash, |refs| {
+			if *refs == 0 {
+				writes += 1;
+				<Self as Store>::CodeByHash::insert(code_hash, code);
+			}
+			*refs += 1;
+		});
+		(reads, writes)
+	}
+
+	/// Decrease the number of reference ofthe validation code and remove it from storage if zero
+	/// is reached.
+	fn decrease_code_ref(code_hash: &Hash) {
+		let refs = <Self as Store>::CodeByHashRefs::get(code_hash);
+		if refs <= 1 {
+			<Self as Store>::CodeByHash::remove(code_hash);
+			<Self as Store>::CodeByHashRefs::remove(code_hash);
+		} else {
+			<Self as Store>::CodeByHashRefs::insert(code_hash, refs - 1);
+		}
+	}
+
 	/// Test function for triggering a new session in this pallet.
 	#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
 	pub fn test_on_new_session() {
@@ -878,6 +988,16 @@ mod tests {
 		ReplacementTimes { expected_at, activated_at }
 	}
 
+	fn check_code_is_stored(validation_code: &ValidationCode) {
+		assert!(<Paras as Store>::CodeByHashRefs::get(validation_code.hash()) != 0);
+		assert!(<Paras as Store>::CodeByHash::contains_key(validation_code.hash()));
+	}
+
+	fn check_code_is_not_stored(validation_code: &ValidationCode) {
+		assert!(!<Paras as Store>::CodeByHashRefs::contains_key(validation_code.hash()));
+		assert!(!<Paras as Store>::CodeByHash::contains_key(validation_code.hash()));
+	}
+
 	#[test]
 	fn para_past_code_meta_gives_right_code() {
 		let mut past_code = ParaPastCodeMeta::default();
@@ -981,8 +1101,10 @@ mod tests {
 			let id = ParaId::from(0u32);
 			let at_block: BlockNumber = 10;
 			let included_block: BlockNumber = 12;
+			let validation_code = ValidationCode(vec![1, 2, 3]);
 
-			<Paras as Store>::PastCode::insert(&(id, at_block), &ValidationCode(vec![1, 2, 3]));
+			Paras::increase_code_ref(&validation_code.hash(), &validation_code);
+			<Paras as Store>::PastCodeHash::insert(&(id, at_block), &validation_code.hash());
 			<Paras as Store>::PastCodePruning::put(&vec![(id, included_block)]);
 
 			{
@@ -992,15 +1114,18 @@ mod tests {
 			}
 
 			let pruned_at: BlockNumber = included_block + acceptance_period + 1;
-			assert_eq!(<Paras as Store>::PastCode::get(&(id, at_block)), Some(vec![1, 2, 3].into()));
+			assert_eq!(<Paras as Store>::PastCodeHash::get(&(id, at_block)), Some(validation_code.hash()));
+			check_code_is_stored(&validation_code);
 
 			run_to_block(pruned_at - 1, None);
-			assert_eq!(<Paras as Store>::PastCode::get(&(id, at_block)), Some(vec![1, 2, 3].into()));
+			assert_eq!(<Paras as Store>::PastCodeHash::get(&(id, at_block)), Some(validation_code.hash()));
 			assert_eq!(Paras::past_code_meta(&id).most_recent_change(), Some(at_block));
+			check_code_is_stored(&validation_code);
 
 			run_to_block(pruned_at, None);
-			assert!(<Paras as Store>::PastCode::get(&(id, at_block)).is_none());
+			assert!(<Paras as Store>::PastCodeHash::get(&(id, at_block)).is_none());
 			assert!(Paras::past_code_meta(&id).most_recent_change().is_none());
+			check_code_is_not_stored(&validation_code);
 		});
 	}
 
@@ -1070,8 +1195,8 @@ mod tests {
 			let id_a = ParaId::from(0u32);
 			let id_b = ParaId::from(1u32);
 
-			Paras::note_past_code(id_a, 10, 12, vec![1, 2, 3].into());
-			Paras::note_past_code(id_b, 20, 23, vec![4, 5, 6].into());
+			Paras::note_past_code(id_a, 10, 12, ValidationCode(vec![1, 2, 3]).hash());
+			Paras::note_past_code(id_b, 20, 23, ValidationCode(vec![4, 5, 6]).hash());
 
 			assert_eq!(<Paras as Store>::PastCodePruning::get(), vec![(id_a, 12), (id_b, 23)]);
 			assert_eq!(
@@ -1096,11 +1221,12 @@ mod tests {
 		let acceptance_period = 10;
 		let validation_upgrade_delay = 5;
 
+		let original_code = ValidationCode(vec![1, 2, 3]);
 		let paras = vec![
 			(0u32.into(), ParaGenesisArgs {
 				parachain: true,
 				genesis_head: Default::default(),
-				validation_code: vec![1, 2, 3].into(),
+				validation_code: original_code.clone(),
 			}),
 		];
 
@@ -1118,11 +1244,13 @@ mod tests {
 		};
 
 		new_test_ext(genesis_config).execute_with(|| {
+			check_code_is_stored(&original_code);
+
 			let para_id = ParaId::from(0);
 			let new_code = ValidationCode(vec![4, 5, 6]);
 
 			run_to_block(2, None);
-			assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
 
 			let expected_at = {
 				// this parablock is in the context of block 1.
@@ -1132,8 +1260,10 @@ mod tests {
 
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
-				assert_eq!(<Paras as Store>::FutureCode::get(&para_id), Some(new_code.clone()));
-				assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
+				check_code_is_stored(&original_code);
+				check_code_is_stored(&new_code);
 
 				expected_at
 			};
@@ -1147,8 +1277,10 @@ mod tests {
 
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
-				assert_eq!(<Paras as Store>::FutureCode::get(&para_id), Some(new_code.clone()));
-				assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
+				check_code_is_stored(&original_code);
+				check_code_is_stored(&new_code);
 			}
 
 			run_to_block(expected_at + 1, None);
@@ -1163,12 +1295,14 @@ mod tests {
 					Some(expected_at),
 				);
 				assert_eq!(
-					<Paras as Store>::PastCode::get(&(para_id, expected_at)),
-					Some(vec![1, 2, 3,].into()),
+					<Paras as Store>::PastCodeHash::get(&(para_id, expected_at)),
+					Some(original_code.hash()),
 				);
 				assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
-				assert!(<Paras as Store>::FutureCode::get(&para_id).is_none());
-				assert_eq!(Paras::current_code(&para_id), Some(new_code));
+				assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
+				assert_eq!(Paras::current_code(&para_id), Some(new_code.clone()));
+				check_code_is_stored(&original_code);
+				check_code_is_stored(&new_code);
 			}
 		});
 	}
@@ -1178,11 +1312,12 @@ mod tests {
 		let acceptance_period = 10;
 		let validation_upgrade_delay = 5;
 
+		let original_code = ValidationCode(vec![1, 2, 3]);
 		let paras = vec![
 			(0u32.into(), ParaGenesisArgs {
 				parachain: true,
 				genesis_head: Default::default(),
-				validation_code: vec![1, 2, 3].into(),
+				validation_code: original_code.clone(),
 			}),
 		];
 
@@ -1204,7 +1339,7 @@ mod tests {
 			let new_code = ValidationCode(vec![4, 5, 6]);
 
 			run_to_block(2, None);
-			assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
 
 			let expected_at = {
 				// this parablock is in the context of block 1.
@@ -1214,8 +1349,8 @@ mod tests {
 
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
-				assert_eq!(<Paras as Store>::FutureCode::get(&para_id), Some(new_code.clone()));
-				assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
 
 				expected_at
 			};
@@ -1232,12 +1367,12 @@ mod tests {
 					Some(expected_at),
 				);
 				assert_eq!(
-					<Paras as Store>::PastCode::get(&(para_id, expected_at)),
-					Some(vec![1, 2, 3,].into()),
+					<Paras as Store>::PastCodeHash::get(&(para_id, expected_at)),
+					Some(original_code.hash()),
 				);
 				assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
-				assert!(<Paras as Store>::FutureCode::get(&para_id).is_none());
-				assert_eq!(Paras::current_code(&para_id), Some(new_code));
+				assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
+				assert_eq!(Paras::current_code(&para_id), Some(new_code.clone()));
 			}
 		});
 	}
@@ -1275,11 +1410,13 @@ mod tests {
 
 			Paras::schedule_code_upgrade(para_id, new_code.clone(), 8);
 			assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(8));
-			assert_eq!(<Paras as Store>::FutureCode::get(&para_id), Some(new_code.clone()));
+			assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+			check_code_is_stored(&new_code);
 
 			Paras::schedule_code_upgrade(para_id, newer_code.clone(), 10);
 			assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(8));
-			assert_eq!(<Paras as Store>::FutureCode::get(&para_id), Some(new_code.clone()));
+			assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+			check_code_is_not_stored(&newer_code);
 		});
 	}
 
@@ -1287,11 +1424,12 @@ mod tests {
 	fn full_parachain_cleanup_storage() {
 		let acceptance_period = 10;
 
+		let original_code = ValidationCode(vec![1, 2, 3]);
 		let paras = vec![
 			(0u32.into(), ParaGenesisArgs {
 				parachain: true,
 				genesis_head: Default::default(),
-				validation_code: vec![1, 2, 3].into(),
+				validation_code: original_code.clone(),
 			}),
 		];
 
@@ -1308,11 +1446,14 @@ mod tests {
 		};
 
 		new_test_ext(genesis_config).execute_with(|| {
+			check_code_is_stored(&original_code);
+
 			let para_id = ParaId::from(0);
 			let new_code = ValidationCode(vec![4, 5, 6]);
 
 			run_to_block(2, None);
-			assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
+			check_code_is_stored(&original_code);
 
 			let expected_at = {
 				// this parablock is in the context of block 1.
@@ -1322,8 +1463,10 @@ mod tests {
 
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
-				assert_eq!(<Paras as Store>::FutureCode::get(&para_id), Some(new_code.clone()));
-				assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
+				check_code_is_stored(&original_code);
+				check_code_is_stored(&new_code);
 
 				expected_at
 			};
@@ -1340,8 +1483,10 @@ mod tests {
 
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
-				assert_eq!(<Paras as Store>::FutureCode::get(&para_id), Some(new_code.clone()));
-				assert_eq!(Paras::current_code(&para_id), Some(vec![1, 2, 3].into()));
+				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
+				check_code_is_stored(&original_code);
+				check_code_is_stored(&new_code);
 
 				assert_eq!(<Paras as Store>::Heads::get(&para_id), Some(Default::default()));
 			}
@@ -1352,14 +1497,16 @@ mod tests {
 			// cleaning up the parachain should place the current parachain code
 			// into the past code buffer & schedule cleanup.
 			assert_eq!(Paras::past_code_meta(&para_id).most_recent_change(), Some(3));
-			assert_eq!(<Paras as Store>::PastCode::get(&(para_id, 3)), Some(vec![1, 2, 3].into()));
+			assert_eq!(<Paras as Store>::PastCodeHash::get(&(para_id, 3)), Some(original_code.hash()));
 			assert_eq!(<Paras as Store>::PastCodePruning::get(), vec![(para_id, 3)]);
+			check_code_is_stored(&original_code);
 
 			// any future upgrades haven't been used to validate yet, so those
 			// are cleaned up immediately.
 			assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
-			assert!(<Paras as Store>::FutureCode::get(&para_id).is_none());
+			assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
 			assert!(Paras::current_code(&para_id).is_none());
+			check_code_is_not_stored(&new_code);
 
 			// run to do the final cleanup
 			let cleaned_up_at = 3 + acceptance_period + 1;
@@ -1367,8 +1514,9 @@ mod tests {
 
 			// now the final cleanup: last past code cleaned up, and this triggers meta cleanup.
 			assert_eq!(Paras::past_code_meta(&para_id), Default::default());
-			assert!(<Paras as Store>::PastCode::get(&(para_id, 3)).is_none());
+			assert!(<Paras as Store>::PastCodeHash::get(&(para_id, 3)).is_none());
 			assert!(<Paras as Store>::PastCodePruning::get().is_empty());
+			check_code_is_not_stored(&original_code);
 		});
 	}
 
@@ -1564,4 +1712,26 @@ mod tests {
 			assert_eq!(Paras::validation_code_at(para_id, 3, None), Some(new_code.clone()));
 		});
 	}
+
+	#[test]
+	fn code_ref_is_cleaned_correctly() {
+		new_test_ext(Default::default()).execute_with(|| {
+			let code: ValidationCode = vec![1, 2, 3].into();
+			Paras::increase_code_ref(&code.hash(), &code);
+			Paras::increase_code_ref(&code.hash(), &code);
+
+			assert!(CodeByHash::contains_key(code.hash()));
+			assert_eq!(CodeByHashRefs::get(code.hash()), 2);
+
+			Paras::decrease_code_ref(&code.hash());
+
+			assert!(CodeByHash::contains_key(code.hash()));
+			assert_eq!(CodeByHashRefs::get(code.hash()), 1);
+
+			Paras::decrease_code_ref(&code.hash());
+
+			assert!(!CodeByHash::contains_key(code.hash()));
+			assert!(!CodeByHashRefs::contains_key(code.hash()));
+		});
+	}
 }

polkadot/runtime/parachains/src/runtime_api_impl/v1.rs

@@ -25,7 +25,7 @@ use primitives::v1::{
 	Id as ParaId, OccupiedCoreAssumption, SessionIndex, ValidationCode,
 	CommittedCandidateReceipt, ScheduledCore, OccupiedCore, CoreOccupied, CoreIndex,
 	GroupIndex, CandidateEvent, PersistedValidationData, SessionInfo,
-	InboundDownwardMessage, InboundHrmpMessage, AuthorityDiscoveryId
+	InboundDownwardMessage, InboundHrmpMessage, AuthorityDiscoveryId, Hash
 };
 use crate::{initializer, inclusion, scheduler, configuration, paras, session_info, dmp, hrmp, shared};
 
@@ -330,3 +330,10 @@ pub fn inbound_hrmp_channels_contents<T: hrmp::Config>(
 ) -> BTreeMap<ParaId, Vec<InboundHrmpMessage<T::BlockNumber>>> {
 	<hrmp::Module<T>>::inbound_hrmp_channels_contents(recipient)
 }
+
+/// Implementation for the `validation_code_by_hash` function of the runtime API.
+pub fn validation_code_by_hash<T: paras::Config>(
+	hash: Hash,
+) -> Option<ValidationCode> {
+	<paras::Module<T>>::code_by_hash(hash)
+}

polkadot/runtime/polkadot/src/lib.rs

@@ -1244,6 +1244,9 @@ sp_api::impl_runtime_apis! {
 			BTreeMap::new()
 		}
 
+		fn validation_code_by_hash(_hash: Hash) -> Option<ValidationCode> {
+			None
+		}
 	}
 
 	impl beefy_primitives::BeefyApi<Block, BeefyId> for Runtime {

polkadot/runtime/rococo/src/lib.rs

@@ -951,6 +951,10 @@ sp_api::impl_runtime_apis! {
 		) -> BTreeMap<Id, Vec<InboundHrmpMessage<BlockNumber>>> {
 			runtime_api_impl::inbound_hrmp_channels_contents::<Runtime>(recipient)
 		}
+
+		fn validation_code_by_hash(hash: Hash) -> Option<ValidationCode> {
+			runtime_api_impl::validation_code_by_hash::<Runtime>(hash)
+		}
 	}
 
 	impl fg_primitives::GrandpaApi<Block> for Runtime {

polkadot/runtime/test-runtime/src/lib.rs

@@ -702,6 +702,10 @@ sp_api::impl_runtime_apis! {
 		) -> BTreeMap<ParaId, Vec<InboundHrmpMessage<BlockNumber>>> {
 			runtime_impl::inbound_hrmp_channels_contents::<Runtime>(recipient)
 		}
+
+		fn validation_code_by_hash(hash: Hash) -> Option<ValidationCode> {
+			runtime_impl::validation_code_by_hash::<Runtime>(hash)
+		}
 	}
 
 	impl beefy_primitives::BeefyApi<Block, BeefyId> for Runtime {

polkadot/runtime/westend/src/lib.rs

@@ -983,6 +983,10 @@ sp_api::impl_runtime_apis! {
 		) -> BTreeMap<Id, Vec<InboundHrmpMessage<BlockNumber>>> {
 			BTreeMap::new()
 		}
+
+		fn validation_code_by_hash(_hash: Hash) -> Option<ValidationCode> {
+			None
+		}
 	}
 
 	impl beefy_primitives::BeefyApi<Block, BeefyId> for Runtime {