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>
2026-05-30 08:11:03 +00:00 · 2021-04-01 23:55:39 +02:00
parent 98082c5326
commit beca01f118
15 changed files with 494 additions and 102 deletions
@@ -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()),
+				);
+			}
 		});
 	}

@@ -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()));
+		});
+	}
 }
@@ -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)
+}