pezkuwi_runtime_teyrchains/paras/
mod.rs

1// Copyright (C) Parity Technologies (UK) Ltd. and Dijital Kurdistan Tech Institute
2// This file is part of Pezkuwi.
3
4// Pezkuwi is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// Pezkuwi is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
16
17//! The paras pezpallet acts as the main registry of paras.
18//!
19//! # Tracking State of Paras
20//!
21//! The most important responsibility of this module is to track which teyrchains
22//! are active and what their current state is. The current state of a para consists of the current
23//! head data and the current validation code (AKA Teyrchain Validation Function (PVF)).
24//!
25//! A para is not considered live until it is registered and activated in this pezpallet.
26//!
27//! The set of teyrchains cannot change except at session boundaries. This is primarily to ensure
28//! that the number and meaning of bits required for the availability bitfields does not change
29//! except at session boundaries.
30//!
31//! # Validation Code Upgrades
32//!
33//! When a para signals the validation code upgrade it will be processed by this module. This can
34//! be in turn split into more fine grained items:
35//!
36//! - Part of the acceptance criteria checks if the para can indeed signal an upgrade,
37//!
38//! - When the candidate is enacted, this module schedules code upgrade, storing the prospective
39//!   validation code.
40//!
41//! - Actually assign the prospective validation code to be the current one after all conditions are
42//!   fulfilled.
43//!
44//! The conditions that must be met before the para can use the new validation code are:
45//!
46//! 1. The validation code should have been "soaked" in the storage for a given number of blocks.
47//! That    is, the validation code should have been stored in on-chain storage for some time, so
48//! that in    case of a revert with a non-extreme height difference, that validation code can still
49//! be    found on-chain.
50//!
51//! 2. The validation code was vetted by the validators and declared as non-malicious in a processes
52//!    known as PVF pre-checking.
53//!
54//! # Validation Code Management
55//!
56//! Potentially, one validation code can be used by several different paras. For example, during
57//! initial stages of deployment several paras can use the same "shell" validation code, or
58//! there can be shards of the same para that use the same validation code.
59//!
60//! In case a validation code ceases to have any users it must be pruned from the on-chain storage.
61//!
62//! # Para Lifecycle Management
63//!
64//! A para can be in one of the two stable states: it is either a lease holding teyrchain or an
65//! on-demand teyrchain.
66//!
67//! However, in order to get into one of those two states, it must first be onboarded. Onboarding
68//! can be only enacted at session boundaries. Onboarding must take at least one full session.
69//! Moreover, a brand new validation code should go through the PVF pre-checking process.
70//!
71//! Once the para is in one of the two stable states, it can switch to the other stable state or to
72//! initiate offboarding process. The result of offboarding is removal of all data related to that
73//! para.
74//!
75//! # PVF Pre-checking
76//!
77//! As was mentioned above, a brand new validation code should go through a process of approval. As
78//! part of this process, validators from the active set will take the validation code and check if
79//! it is malicious. Once they did that and have their judgement, either accept or reject, they
80//! issue a statement in a form of an unsigned extrinsic. This extrinsic is processed by this
81//! pezpallet. Once supermajority is gained for accept, then the process that initiated the check is
82//! resumed (as mentioned before this can be either upgrading of validation code or onboarding). If
83//! getting a supermajority becomes impossible (>1/3 of validators have already voted against), then
84//! we reject.
85//!
86//! Below is a state diagram that depicts states of a single PVF pre-checking vote.
87//!
88//! ```text
89//!                                            ┌──────────┐
90//!                        supermajority       │          │
91//!                    ┌────────for───────────▶│ accepted │
92//!        vote────┐   │                       │          │
93//!         │      │   │                       └──────────┘
94//!         │      │   │
95//!         │  ┌───────┐
96//!         │  │       │
97//!         └─▶│ init  │──── >1/3 against      ┌──────────┐
98//!            │       │           │           │          │
99//!            └───────┘           └──────────▶│ rejected │
100//!             ▲  │                           │          │
101//!             │  │ session                   └──────────┘
102//!             │  └──change
103//!             │     │
104//!             │     ▼
105//!             ┌─────┐
106//! start──────▶│reset│
107//!             └─────┘
108//! ```
109
110use crate::{
111	configuration,
112	inclusion::{QueueFootprinter, UmpQueueId},
113	initializer::SessionChangeNotification,
114	shared,
115};
116use alloc::{collections::btree_set::BTreeSet, vec::Vec};
117use bitvec::{order::Lsb0 as BitOrderLsb0, vec::BitVec};
118use codec::{Decode, Encode};
119use core::{cmp, mem};
120use pezframe_support::{
121	pezpallet_prelude::*,
122	traits::{EnsureOriginWithArg, EstimateNextSessionRotation},
123	DefaultNoBound,
124};
125use pezframe_system::pezpallet_prelude::*;
126use pezkuwi_primitives::{
127	ConsensusLog, HeadData, Id as ParaId, PvfCheckStatement, SessionIndex, UpgradeGoAhead,
128	UpgradeRestriction, ValidationCode, ValidationCodeHash, ValidatorSignature, MIN_CODE_SIZE,
129};
130use pezsp_core::RuntimeDebug;
131use pezsp_runtime::{
132	traits::{AppVerify, One, Saturating},
133	DispatchResult, SaturatedConversion,
134};
135use scale_info::{Type, TypeInfo};
136
137use serde::{Deserialize, Serialize};
138
139pub use crate::Origin as TeyrchainOrigin;
140
141#[cfg(feature = "runtime-benchmarks")]
142pub mod benchmarking;
143
144#[cfg(test)]
145pub(crate) mod tests;
146
147pub use pezpallet::*;
148
149const LOG_TARGET: &str = "runtime::paras";
150
151// the two key times necessary to track for every code replacement.
152#[derive(Default, Encode, Decode, TypeInfo)]
153#[cfg_attr(test, derive(Debug, Clone, PartialEq))]
154pub struct ReplacementTimes<N> {
155	/// The relay-chain block number that the code upgrade was expected to be activated.
156	/// This is when the code change occurs from the para's perspective - after the
157	/// first parablock included with a relay-parent with number >= this value.
158	expected_at: N,
159	/// The relay-chain block number at which the parablock activating the code upgrade was
160	/// actually included. This means considered included and available, so this is the time at
161	/// which that parablock enters the acceptance period in this fork of the relay-chain.
162	activated_at: N,
163}
164
165/// Metadata used to track previous teyrchain validation code that we keep in
166/// the state.
167#[derive(Default, Encode, Decode, TypeInfo)]
168#[cfg_attr(test, derive(Debug, Clone, PartialEq))]
169pub struct ParaPastCodeMeta<N> {
170	/// Block numbers where the code was expected to be replaced and where the code
171	/// was actually replaced, respectively. The first is used to do accurate look-ups
172	/// of historic code in historic contexts, whereas the second is used to do
173	/// pruning on an accurate timeframe. These can be used as indices
174	/// into the `PastCodeHash` map along with the `ParaId` to fetch the code itself.
175	upgrade_times: Vec<ReplacementTimes<N>>,
176	/// Tracks the highest pruned code-replacement, if any. This is the `activated_at` value,
177	/// not the `expected_at` value.
178	last_pruned: Option<N>,
179}
180
181/// The possible states of a para, to take into account delayed lifecycle changes.
182///
183/// If the para is in a "transition state", it is expected that the teyrchain is
184/// queued in the `ActionsQueue` to transition it into a stable state. Its lifecycle
185/// state will be used to determine the state transition to apply to the para.
186#[derive(PartialEq, Eq, Clone, Encode, Decode, RuntimeDebug, TypeInfo)]
187pub enum ParaLifecycle {
188	/// Para is new and is onboarding as an on-demand or lease holding Teyrchain.
189	Onboarding,
190	/// Para is a Parathread (on-demand teyrchain).
191	Parathread,
192	/// Para is a lease holding Teyrchain.
193	Teyrchain,
194	/// Para is a Parathread (on-demand teyrchain) which is upgrading to a lease holding Teyrchain.
195	UpgradingParathread,
196	/// Para is a lease holding Teyrchain which is downgrading to an on-demand teyrchain.
197	DowngradingTeyrchain,
198	/// Parathread (on-demand teyrchain) is queued to be offboarded.
199	OffboardingParathread,
200	/// Teyrchain is queued to be offboarded.
201	OffboardingTeyrchain,
202}
203
204impl ParaLifecycle {
205	/// Returns true if teyrchain is currently onboarding. To learn if the
206	/// teyrchain is onboarding as a lease holding or on-demand teyrchain, look at the
207	/// `UpcomingGenesis` storage item.
208	pub fn is_onboarding(&self) -> bool {
209		matches!(self, ParaLifecycle::Onboarding)
210	}
211
212	/// Returns true if para is in a stable state, i.e. it is currently
213	/// a lease holding or on-demand teyrchain, and not in any transition state.
214	pub fn is_stable(&self) -> bool {
215		matches!(self, ParaLifecycle::Parathread | ParaLifecycle::Teyrchain)
216	}
217
218	/// Returns true if para is currently treated as a teyrchain.
219	/// This also includes transitioning states, so you may want to combine
220	/// this check with `is_stable` if you specifically want `Paralifecycle::Teyrchain`.
221	pub fn is_teyrchain(&self) -> bool {
222		matches!(
223			self,
224			ParaLifecycle::Teyrchain
225				| ParaLifecycle::DowngradingTeyrchain
226				| ParaLifecycle::OffboardingTeyrchain
227		)
228	}
229
230	/// Returns true if para is currently treated as a parathread (on-demand teyrchain).
231	/// This also includes transitioning states, so you may want to combine
232	/// this check with `is_stable` if you specifically want `Paralifecycle::Parathread`.
233	pub fn is_parathread(&self) -> bool {
234		matches!(
235			self,
236			ParaLifecycle::Parathread
237				| ParaLifecycle::UpgradingParathread
238				| ParaLifecycle::OffboardingParathread
239		)
240	}
241
242	/// Returns true if para is currently offboarding.
243	pub fn is_offboarding(&self) -> bool {
244		matches!(self, ParaLifecycle::OffboardingParathread | ParaLifecycle::OffboardingTeyrchain)
245	}
246
247	/// Returns true if para is in any transitionary state.
248	pub fn is_transitioning(&self) -> bool {
249		!Self::is_stable(self)
250	}
251}
252
253impl<N: Ord + Copy + PartialEq> ParaPastCodeMeta<N> {
254	// note a replacement has occurred at a given block number.
255	pub(crate) fn note_replacement(&mut self, expected_at: N, activated_at: N) {
256		self.upgrade_times.push(ReplacementTimes { expected_at, activated_at })
257	}
258
259	/// Returns `true` if the upgrade logs list is empty.
260	fn is_empty(&self) -> bool {
261		self.upgrade_times.is_empty()
262	}
263
264	// The block at which the most recently tracked code change occurred, from the perspective
265	// of the para.
266	#[cfg(test)]
267	fn most_recent_change(&self) -> Option<N> {
268		self.upgrade_times.last().map(|x| x.expected_at)
269	}
270
271	// prunes all code upgrade logs occurring at or before `max`.
272	// note that code replaced at `x` is the code used to validate all blocks before
273	// `x`. Thus, `max` should be outside of the slashing window when this is invoked.
274	//
275	// Since we don't want to prune anything inside the acceptance period, and the parablock only
276	// enters the acceptance period after being included, we prune based on the activation height of
277	// the code change, not the expected height of the code change.
278	//
279	// returns an iterator of block numbers at which code was replaced, where the replaced
280	// code should be now pruned, in ascending order.
281	fn prune_up_to(&'_ mut self, max: N) -> impl Iterator<Item = N> + '_ {
282		let to_prune = self.upgrade_times.iter().take_while(|t| t.activated_at <= max).count();
283		let drained = if to_prune == 0 {
284			// no-op prune.
285			self.upgrade_times.drain(self.upgrade_times.len()..)
286		} else {
287			// if we are actually pruning something, update the `last_pruned` member.
288			self.last_pruned = Some(self.upgrade_times[to_prune - 1].activated_at);
289			self.upgrade_times.drain(..to_prune)
290		};
291
292		drained.map(|times| times.expected_at)
293	}
294}
295
296/// Arguments for initializing a para.
297#[derive(
298	PartialEq,
299	Eq,
300	Clone,
301	Encode,
302	Decode,
303	DecodeWithMemTracking,
304	RuntimeDebug,
305	TypeInfo,
306	Serialize,
307	Deserialize,
308)]
309pub struct ParaGenesisArgs {
310	/// The initial head data to use.
311	pub genesis_head: HeadData,
312	/// The initial validation code to use.
313	pub validation_code: ValidationCode,
314	/// Lease holding or on-demand teyrchain.
315	#[serde(rename = "teyrchain")]
316	pub para_kind: ParaKind,
317}
318
319/// Distinguishes between lease holding Teyrchain and Parathread (on-demand teyrchain)
320#[derive(DecodeWithMemTracking, PartialEq, Eq, Clone, RuntimeDebug)]
321pub enum ParaKind {
322	Parathread,
323	Teyrchain,
324}
325
326impl Serialize for ParaKind {
327	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
328	where
329		S: serde::Serializer,
330	{
331		match self {
332			ParaKind::Teyrchain => serializer.serialize_bool(true),
333			ParaKind::Parathread => serializer.serialize_bool(false),
334		}
335	}
336}
337
338impl<'de> Deserialize<'de> for ParaKind {
339	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
340	where
341		D: serde::Deserializer<'de>,
342	{
343		match serde::de::Deserialize::deserialize(deserializer) {
344			Ok(true) => Ok(ParaKind::Teyrchain),
345			Ok(false) => Ok(ParaKind::Parathread),
346			_ => Err(serde::de::Error::custom("invalid ParaKind serde representation")),
347		}
348	}
349}
350
351// Manual encoding, decoding, and TypeInfo as the parakind field in ParaGenesisArgs used to be a
352// bool
353impl Encode for ParaKind {
354	fn size_hint(&self) -> usize {
355		true.size_hint()
356	}
357
358	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
359		match self {
360			ParaKind::Teyrchain => true.using_encoded(f),
361			ParaKind::Parathread => false.using_encoded(f),
362		}
363	}
364}
365
366impl Decode for ParaKind {
367	fn decode<I: codec::Input>(input: &mut I) -> Result<Self, codec::Error> {
368		match bool::decode(input) {
369			Ok(true) => Ok(ParaKind::Teyrchain),
370			Ok(false) => Ok(ParaKind::Parathread),
371			_ => Err("Invalid ParaKind representation".into()),
372		}
373	}
374}
375
376impl TypeInfo for ParaKind {
377	type Identity = bool;
378	fn type_info() -> Type {
379		bool::type_info()
380	}
381}
382
383/// This enum describes a reason why a particular PVF pre-checking vote was initiated. When the
384/// PVF vote in question is concluded, this enum indicates what changes should be performed.
385#[derive(Debug, Encode, Decode, TypeInfo)]
386pub(crate) enum PvfCheckCause<BlockNumber> {
387	/// PVF vote was initiated by the initial onboarding process of the given para.
388	Onboarding(ParaId),
389	/// PVF vote was initiated by signalling of an upgrade by the given para.
390	Upgrade {
391		/// The ID of the teyrchain that initiated or is waiting for the conclusion of
392		/// pre-checking.
393		id: ParaId,
394		/// The relay-chain block number of **inclusion** of candidate that that initiated the
395		/// upgrade.
396		///
397		/// It's important to count upgrade enactment delay from the inclusion of this candidate
398		/// instead of its relay parent -- in order to keep PVF available in case of chain
399		/// reversions.
400		///
401		/// See https://github.com/pezkuwichain/pezkuwi-sdk/issues/294 for detailed explanation.
402		included_at: BlockNumber,
403		/// Whether or not the upgrade should be enacted directly.
404		///
405		/// If set to `Yes` it means that no `GoAheadSignal` will be set and the teyrchain code
406		/// will also be overwritten directly.
407		upgrade_strategy: UpgradeStrategy,
408	},
409}
410
411/// The strategy on how to handle a validation code upgrade.
412///
413/// When scheduling a teyrchain code upgrade the upgrade first is checked by all validators. The
414/// validators ensure that the new validation code can be compiled and instantiated. After the
415/// majority of the validators have reported their checking result the upgrade is either scheduled
416/// or aborted. This strategy then comes into play around the relay chain block this upgrade was
417/// scheduled in.
418#[derive(Debug, Copy, Clone, PartialEq, TypeInfo, Decode, Encode)]
419pub enum UpgradeStrategy {
420	/// Set the `GoAhead` signal to inform the teyrchain that it is time to upgrade.
421	///
422	/// The upgrade will then be applied after the first teyrchain block was enacted that must have
423	/// observed the `GoAhead` signal.
424	SetGoAheadSignal,
425	/// Apply the upgrade directly at the expected relay chain block.
426	///
427	/// This doesn't wait for the teyrchain to make any kind of progress.
428	ApplyAtExpectedBlock,
429}
430
431impl<BlockNumber> PvfCheckCause<BlockNumber> {
432	/// Returns the ID of the para that initiated or subscribed to the pre-checking vote.
433	fn para_id(&self) -> ParaId {
434		match *self {
435			PvfCheckCause::Onboarding(id) => id,
436			PvfCheckCause::Upgrade { id, .. } => id,
437		}
438	}
439}
440
441/// Specifies what was the outcome of a PVF pre-checking vote.
442#[derive(Copy, Clone, Encode, Decode, RuntimeDebug, TypeInfo)]
443enum PvfCheckOutcome {
444	Accepted,
445	Rejected,
446}
447
448/// This struct describes the current state of an in-progress PVF pre-checking vote.
449#[derive(Encode, Decode, TypeInfo)]
450pub(crate) struct PvfCheckActiveVoteState<BlockNumber> {
451	// The two following vectors have their length equal to the number of validators in the active
452	// set. They start with all zeroes. A 1 is set at an index when the validator at the that index
453	// makes a vote. Once a 1 is set for either of the vectors, that validator cannot vote anymore.
454	// Since the active validator set changes each session, the bit vectors are reinitialized as
455	// well: zeroed and resized so that each validator gets its own bit.
456	votes_accept: BitVec<u8, BitOrderLsb0>,
457	votes_reject: BitVec<u8, BitOrderLsb0>,
458
459	/// The number of session changes this PVF vote has observed. Therefore, this number is
460	/// increased at each session boundary. When created, it is initialized with 0.
461	age: SessionIndex,
462	/// The block number at which this PVF vote was created.
463	created_at: BlockNumber,
464	/// A list of causes for this PVF pre-checking. Has at least one.
465	causes: Vec<PvfCheckCause<BlockNumber>>,
466}
467
468impl<BlockNumber> PvfCheckActiveVoteState<BlockNumber> {
469	/// Returns a new instance of vote state, started at the specified block `now`, with the
470	/// number of validators in the current session `n_validators` and the originating `cause`.
471	fn new(now: BlockNumber, n_validators: usize, cause: PvfCheckCause<BlockNumber>) -> Self {
472		let mut causes = Vec::with_capacity(1);
473		causes.push(cause);
474		Self {
475			created_at: now,
476			votes_accept: bitvec::bitvec![u8, BitOrderLsb0; 0; n_validators],
477			votes_reject: bitvec::bitvec![u8, BitOrderLsb0; 0; n_validators],
478			age: 0,
479			causes,
480		}
481	}
482
483	/// Resets all votes and resizes the votes vectors corresponding to the number of validators
484	/// in the new session.
485	fn reinitialize_ballots(&mut self, n_validators: usize) {
486		let clear_and_resize = |v: &mut BitVec<_, _>| {
487			v.clear();
488			v.resize(n_validators, false);
489		};
490		clear_and_resize(&mut self.votes_accept);
491		clear_and_resize(&mut self.votes_reject);
492	}
493
494	/// Returns `Some(true)` if the validator at the given index has already cast their vote within
495	/// the ongoing session. Returns `None` in case the index is out of bounds.
496	fn has_vote(&self, validator_index: usize) -> Option<bool> {
497		let accept_vote = self.votes_accept.get(validator_index)?;
498		let reject_vote = self.votes_reject.get(validator_index)?;
499		Some(*accept_vote || *reject_vote)
500	}
501
502	/// Returns `None` if the quorum is not reached, or the direction of the decision.
503	fn quorum(&self, n_validators: usize) -> Option<PvfCheckOutcome> {
504		let accept_threshold = pezkuwi_primitives::supermajority_threshold(n_validators);
505		// At this threshold, a supermajority is no longer possible, so we reject.
506		let reject_threshold = n_validators - accept_threshold;
507
508		if self.votes_accept.count_ones() >= accept_threshold {
509			Some(PvfCheckOutcome::Accepted)
510		} else if self.votes_reject.count_ones() > reject_threshold {
511			Some(PvfCheckOutcome::Rejected)
512		} else {
513			None
514		}
515	}
516
517	#[cfg(test)]
518	pub(crate) fn causes(&self) -> &[PvfCheckCause<BlockNumber>] {
519		self.causes.as_slice()
520	}
521}
522
523/// Runtime hook for when a teyrchain head is updated.
524pub trait OnNewHead {
525	/// Called when a teyrchain head is updated.
526	/// Returns the weight consumed by this function.
527	fn on_new_head(id: ParaId, head: &HeadData) -> Weight;
528}
529
530#[impl_trait_for_tuples::impl_for_tuples(30)]
531impl OnNewHead for Tuple {
532	fn on_new_head(id: ParaId, head: &HeadData) -> Weight {
533		let mut weight: Weight = Default::default();
534		for_tuples!( #( weight.saturating_accrue(Tuple::on_new_head(id, head)); )* );
535		weight
536	}
537}
538
539/// Assign coretime to some teyrchain.
540///
541/// This assigns coretime to a teyrchain without using the coretime chain. Thus, this should only be
542/// used for testing purposes.
543pub trait AssignCoretime {
544	/// ONLY USE FOR TESTING OR GENESIS.
545	fn assign_coretime(id: ParaId) -> DispatchResult;
546}
547
548impl AssignCoretime for () {
549	fn assign_coretime(_: ParaId) -> DispatchResult {
550		Ok(())
551	}
552}
553
554/// Holds an authorized validation code hash along with its expiry timestamp.
555#[derive(Debug, Encode, Decode, DecodeWithMemTracking, TypeInfo)]
556#[cfg_attr(test, derive(PartialEq))]
557pub struct AuthorizedCodeHashAndExpiry<T> {
558	code_hash: ValidationCodeHash,
559	expire_at: T,
560}
561impl<T> From<(ValidationCodeHash, T)> for AuthorizedCodeHashAndExpiry<T> {
562	fn from(value: (ValidationCodeHash, T)) -> Self {
563		AuthorizedCodeHashAndExpiry { code_hash: value.0, expire_at: value.1 }
564	}
565}
566
567pub trait WeightInfo {
568	fn force_set_current_code(c: u32) -> Weight;
569	fn force_set_current_head(s: u32) -> Weight;
570	fn force_set_most_recent_context() -> Weight;
571	fn force_schedule_code_upgrade(c: u32) -> Weight;
572	fn force_note_new_head(s: u32) -> Weight;
573	fn force_queue_action() -> Weight;
574	fn add_trusted_validation_code(c: u32) -> Weight;
575	fn poke_unused_validation_code() -> Weight;
576	fn remove_upgrade_cooldown() -> Weight;
577
578	fn include_pvf_check_statement_finalize_upgrade_accept() -> Weight;
579	fn include_pvf_check_statement_finalize_upgrade_reject() -> Weight;
580	fn include_pvf_check_statement_finalize_onboarding_accept() -> Weight;
581	fn include_pvf_check_statement_finalize_onboarding_reject() -> Weight;
582	fn include_pvf_check_statement() -> Weight;
583	fn authorize_force_set_current_code_hash() -> Weight;
584	fn apply_authorized_force_set_current_code(c: u32) -> Weight;
585}
586
587pub struct TestWeightInfo;
588impl WeightInfo for TestWeightInfo {
589	fn force_set_current_code(_c: u32) -> Weight {
590		Weight::MAX
591	}
592	fn force_set_current_head(_s: u32) -> Weight {
593		Weight::MAX
594	}
595	fn force_set_most_recent_context() -> Weight {
596		Weight::MAX
597	}
598	fn force_schedule_code_upgrade(_c: u32) -> Weight {
599		Weight::MAX
600	}
601	fn force_note_new_head(_s: u32) -> Weight {
602		Weight::MAX
603	}
604	fn force_queue_action() -> Weight {
605		Weight::MAX
606	}
607	fn add_trusted_validation_code(_c: u32) -> Weight {
608		// Called during integration tests for para initialization.
609		Weight::zero()
610	}
611	fn poke_unused_validation_code() -> Weight {
612		Weight::MAX
613	}
614	fn include_pvf_check_statement_finalize_upgrade_accept() -> Weight {
615		Weight::MAX
616	}
617	fn include_pvf_check_statement_finalize_upgrade_reject() -> Weight {
618		Weight::MAX
619	}
620	fn include_pvf_check_statement_finalize_onboarding_accept() -> Weight {
621		Weight::MAX
622	}
623	fn include_pvf_check_statement_finalize_onboarding_reject() -> Weight {
624		Weight::MAX
625	}
626	fn include_pvf_check_statement() -> Weight {
627		// This special value is to distinguish from the finalizing variants above in tests.
628		Weight::MAX - Weight::from_parts(1, 1)
629	}
630	fn remove_upgrade_cooldown() -> Weight {
631		Weight::MAX
632	}
633	fn authorize_force_set_current_code_hash() -> Weight {
634		Weight::MAX
635	}
636	fn apply_authorized_force_set_current_code(_c: u32) -> Weight {
637		Weight::MAX
638	}
639}
640
641#[pezframe_support::pezpallet]
642pub mod pezpallet {
643	use super::*;
644	use pezframe_support::traits::{
645		fungible::{Inspect, Mutate},
646		tokens::{Fortitude, Precision, Preservation},
647	};
648	use pezsp_runtime::transaction_validity::{
649		InvalidTransaction, TransactionPriority, TransactionSource, TransactionValidity,
650		ValidTransaction,
651	};
652
653	type BalanceOf<T> = <<T as Config>::Fungible as Inspect<AccountIdFor<T>>>::Balance;
654
655	#[pezpallet::pezpallet]
656	#[pezpallet::without_storage_info]
657	pub struct Pezpallet<T>(_);
658
659	#[pezpallet::config]
660	pub trait Config:
661		pezframe_system::Config
662		+ configuration::Config
663		+ shared::Config
664		+ pezframe_system::offchain::CreateBare<Call<Self>>
665	{
666		#[allow(deprecated)]
667		type RuntimeEvent: From<Event<Self>>
668			+ IsType<<Self as pezframe_system::Config>::RuntimeEvent>;
669
670		#[pezpallet::constant]
671		type UnsignedPriority: Get<TransactionPriority>;
672
673		type NextSessionRotation: EstimateNextSessionRotation<BlockNumberFor<Self>>;
674
675		/// Retrieve how many UMP messages are enqueued for this para-chain.
676		///
677		/// This is used to judge whether or not a para-chain can offboard. Per default this should
678		/// be set to the `ParaInclusion` pezpallet.
679		type QueueFootprinter: QueueFootprinter<Origin = UmpQueueId>;
680
681		/// Runtime hook for when a teyrchain head is updated.
682		type OnNewHead: OnNewHead;
683
684		/// Weight information for extrinsics in this pezpallet.
685		type WeightInfo: WeightInfo;
686
687		/// Runtime hook for assigning coretime for a given teyrchain.
688		///
689		/// This is only used at genesis or by root.
690		///
691		/// TODO: Remove once coretime is the standard across all chains.
692		type AssignCoretime: AssignCoretime;
693
694		/// The fungible instance used by the runtime.
695		type Fungible: Mutate<Self::AccountId, Balance: From<BlockNumberFor<Self>>>;
696
697		/// Multiplier to determine the cost of removing upgrade cooldown.
698		///
699		/// After a teyrchain upgrades their runtime, an upgrade cooldown is applied
700		/// ([`configuration::HostConfiguration::validation_upgrade_cooldown`]). This cooldown
701		/// exists to prevent spamming the relay chain with runtime upgrades. But as life is going
702		/// on, mistakes can happen and a consequent may be required. The cooldown period can be
703		/// removed by using [`Pezpallet::remove_upgrade_cooldown`]. This dispatchable will use this
704		/// multiplier to determine the cost for removing the upgrade cooldown. Time left for the
705		/// cooldown multiplied with this multiplier determines the cost.
706		type CooldownRemovalMultiplier: Get<BalanceOf<Self>>;
707
708		/// The origin that can authorize [`Pezpallet::authorize_force_set_current_code_hash`].
709		///
710		/// In the end this allows [`Pezpallet::apply_authorized_force_set_current_code`] to force
711		/// set the current code without paying any fee. So, the origin should be chosen with
712		/// care.
713		type AuthorizeCurrentCodeOrigin: EnsureOriginWithArg<Self::RuntimeOrigin, ParaId>;
714	}
715
716	#[pezpallet::event]
717	#[pezpallet::generate_deposit(pub(super) fn deposit_event)]
718	pub enum Event<T: Config> {
719		/// Current code has been updated for a Para. `para_id`
720		CurrentCodeUpdated(ParaId),
721		/// Current head has been updated for a Para. `para_id`
722		CurrentHeadUpdated(ParaId),
723		/// A code upgrade has been scheduled for a Para. `para_id`
724		CodeUpgradeScheduled(ParaId),
725		/// A new head has been noted for a Para. `para_id`
726		NewHeadNoted(ParaId),
727		/// A para has been queued to execute pending actions. `para_id`
728		ActionQueued(ParaId, SessionIndex),
729		/// The given para either initiated or subscribed to a PVF check for the given validation
730		/// code. `code_hash` `para_id`
731		PvfCheckStarted(ValidationCodeHash, ParaId),
732		/// The given validation code was accepted by the PVF pre-checking vote.
733		/// `code_hash` `para_id`
734		PvfCheckAccepted(ValidationCodeHash, ParaId),
735		/// The given validation code was rejected by the PVF pre-checking vote.
736		/// `code_hash` `para_id`
737		PvfCheckRejected(ValidationCodeHash, ParaId),
738		/// The upgrade cooldown was removed.
739		UpgradeCooldownRemoved {
740			/// The teyrchain for which the cooldown got removed.
741			para_id: ParaId,
742		},
743		/// A new code hash has been authorized for a Para.
744		CodeAuthorized {
745			/// Para
746			para_id: ParaId,
747			/// Authorized code hash.
748			code_hash: ValidationCodeHash,
749			/// Block at which authorization expires and will be removed.
750			expire_at: BlockNumberFor<T>,
751		},
752	}
753
754	#[pezpallet::error]
755	pub enum Error<T> {
756		/// Para is not registered in our system.
757		NotRegistered,
758		/// Para cannot be onboarded because it is already tracked by our system.
759		CannotOnboard,
760		/// Para cannot be offboarded at this time.
761		CannotOffboard,
762		/// Para cannot be upgraded to a lease holding teyrchain.
763		CannotUpgrade,
764		/// Para cannot be downgraded to an on-demand teyrchain.
765		CannotDowngrade,
766		/// The statement for PVF pre-checking is stale.
767		PvfCheckStatementStale,
768		/// The statement for PVF pre-checking is for a future session.
769		PvfCheckStatementFuture,
770		/// Claimed validator index is out of bounds.
771		PvfCheckValidatorIndexOutOfBounds,
772		/// The signature for the PVF pre-checking is invalid.
773		PvfCheckInvalidSignature,
774		/// The given validator already has cast a vote.
775		PvfCheckDoubleVote,
776		/// The given PVF does not exist at the moment of process a vote.
777		PvfCheckSubjectInvalid,
778		/// Teyrchain cannot currently schedule a code upgrade.
779		CannotUpgradeCode,
780		/// Invalid validation code size.
781		InvalidCode,
782		/// No upgrade authorized.
783		NothingAuthorized,
784		/// The submitted code is not authorized.
785		Unauthorized,
786		/// Invalid block number.
787		InvalidBlockNumber,
788	}
789
790	/// All currently active PVF pre-checking votes.
791	///
792	/// Invariant:
793	/// - There are no PVF pre-checking votes that exists in list but not in the set and vice versa.
794	#[pezpallet::storage]
795	pub(super) type PvfActiveVoteMap<T: Config> = StorageMap<
796		_,
797		Twox64Concat,
798		ValidationCodeHash,
799		PvfCheckActiveVoteState<BlockNumberFor<T>>,
800		OptionQuery,
801	>;
802
803	/// The list of all currently active PVF votes. Auxiliary to `PvfActiveVoteMap`.
804	#[pezpallet::storage]
805	pub(super) type PvfActiveVoteList<T: Config> =
806		StorageValue<_, Vec<ValidationCodeHash>, ValueQuery>;
807
808	/// All lease holding teyrchains. Ordered ascending by `ParaId`. On demand teyrchains are not
809	/// included.
810	///
811	/// Consider using the [`TeyrchainsCache`] type of modifying.
812	#[pezpallet::storage]
813	pub type Teyrchains<T: Config> = StorageValue<_, Vec<ParaId>, ValueQuery>;
814
815	/// The current lifecycle of a all known Para IDs.
816	#[pezpallet::storage]
817	pub(super) type ParaLifecycles<T: Config> = StorageMap<_, Twox64Concat, ParaId, ParaLifecycle>;
818
819	/// The head-data of every registered para.
820	#[pezpallet::storage]
821	pub type Heads<T: Config> = StorageMap<_, Twox64Concat, ParaId, HeadData>;
822
823	/// The context (relay-chain block number) of the most recent teyrchain head.
824	#[pezpallet::storage]
825	pub type MostRecentContext<T: Config> = StorageMap<_, Twox64Concat, ParaId, BlockNumberFor<T>>;
826
827	/// The validation code hash of every live para.
828	///
829	/// Corresponding code can be retrieved with [`CodeByHash`].
830	#[pezpallet::storage]
831	pub type CurrentCodeHash<T: Config> = StorageMap<_, Twox64Concat, ParaId, ValidationCodeHash>;
832
833	/// Actual past code hash, indicated by the para id as well as the block number at which it
834	/// became outdated.
835	///
836	/// Corresponding code can be retrieved with [`CodeByHash`].
837	#[pezpallet::storage]
838	pub(super) type PastCodeHash<T: Config> =
839		StorageMap<_, Twox64Concat, (ParaId, BlockNumberFor<T>), ValidationCodeHash>;
840
841	/// Past code of teyrchains. The teyrchains themselves may not be registered anymore,
842	/// but we also keep their code on-chain for the same amount of time as outdated code
843	/// to keep it available for approval checkers.
844	#[pezpallet::storage]
845	pub type PastCodeMeta<T: Config> =
846		StorageMap<_, Twox64Concat, ParaId, ParaPastCodeMeta<BlockNumberFor<T>>, ValueQuery>;
847
848	/// Which paras have past code that needs pruning and the relay-chain block at which the code
849	/// was replaced. Note that this is the actual height of the included block, not the expected
850	/// height at which the code upgrade would be applied, although they may be equal.
851	/// This is to ensure the entire acceptance period is covered, not an offset acceptance period
852	/// starting from the time at which the teyrchain perceives a code upgrade as having occurred.
853	/// Multiple entries for a single para are permitted. Ordered ascending by block number.
854	#[pezpallet::storage]
855	pub(super) type PastCodePruning<T: Config> =
856		StorageValue<_, Vec<(ParaId, BlockNumberFor<T>)>, ValueQuery>;
857
858	/// The block number at which the planned code change is expected for a teyrchain.
859	///
860	/// The change will be applied after the first parablock for this ID included which executes
861	/// in the context of a relay chain block with a number >= `expected_at`.
862	#[pezpallet::storage]
863	pub type FutureCodeUpgrades<T: Config> = StorageMap<_, Twox64Concat, ParaId, BlockNumberFor<T>>;
864
865	/// The list of upcoming future code upgrades.
866	///
867	/// Each item is a pair of the teyrchain and the expected block at which the upgrade should be
868	/// applied. The upgrade will be applied at the given relay chain block. In contrast to
869	/// [`FutureCodeUpgrades`] this code upgrade will be applied regardless the teyrchain making any
870	/// progress or not.
871	///
872	/// Ordered ascending by block number.
873	#[pezpallet::storage]
874	pub(super) type FutureCodeUpgradesAt<T: Config> =
875		StorageValue<_, Vec<(ParaId, BlockNumberFor<T>)>, ValueQuery>;
876
877	/// The actual future code hash of a para.
878	///
879	/// Corresponding code can be retrieved with [`CodeByHash`].
880	#[pezpallet::storage]
881	pub type FutureCodeHash<T: Config> = StorageMap<_, Twox64Concat, ParaId, ValidationCodeHash>;
882
883	/// The code hash authorizations for a para which will expire `expire_at` `BlockNumberFor<T>`.
884	#[pezpallet::storage]
885	pub type AuthorizedCodeHash<T: Config> =
886		StorageMap<_, Twox64Concat, ParaId, AuthorizedCodeHashAndExpiry<BlockNumberFor<T>>>;
887
888	/// This is used by the relay-chain to communicate to a teyrchain a go-ahead with in the upgrade
889	/// procedure.
890	///
891	/// This value is absent when there are no upgrades scheduled or during the time the relay chain
892	/// performs the checks. It is set at the first relay-chain block when the corresponding
893	/// teyrchain can switch its upgrade function. As soon as the teyrchain's block is included, the
894	/// value gets reset to `None`.
895	///
896	/// NOTE that this field is used by teyrchains via merkle storage proofs, therefore changing
897	/// the format will require migration of teyrchains.
898	#[pezpallet::storage]
899	pub(super) type UpgradeGoAheadSignal<T: Config> =
900		StorageMap<_, Twox64Concat, ParaId, UpgradeGoAhead>;
901
902	/// This is used by the relay-chain to communicate that there are restrictions for performing
903	/// an upgrade for this teyrchain.
904	///
905	/// This may be a because the teyrchain waits for the upgrade cooldown to expire. Another
906	/// potential use case is when we want to perform some maintenance (such as storage migration)
907	/// we could restrict upgrades to make the process simpler.
908	///
909	/// NOTE that this field is used by teyrchains via merkle storage proofs, therefore changing
910	/// the format will require migration of teyrchains.
911	#[pezpallet::storage]
912	pub type UpgradeRestrictionSignal<T: Config> =
913		StorageMap<_, Twox64Concat, ParaId, UpgradeRestriction>;
914
915	/// The list of teyrchains that are awaiting for their upgrade restriction to cooldown.
916	///
917	/// Ordered ascending by block number.
918	#[pezpallet::storage]
919	pub(super) type UpgradeCooldowns<T: Config> =
920		StorageValue<_, Vec<(ParaId, BlockNumberFor<T>)>, ValueQuery>;
921
922	/// The list of upcoming code upgrades.
923	///
924	/// Each item is a pair of which para performs a code upgrade and at which relay-chain block it
925	/// is expected at.
926	///
927	/// Ordered ascending by block number.
928	#[pezpallet::storage]
929	pub(super) type UpcomingUpgrades<T: Config> =
930		StorageValue<_, Vec<(ParaId, BlockNumberFor<T>)>, ValueQuery>;
931
932	/// The actions to perform during the start of a specific session index.
933	#[pezpallet::storage]
934	pub type ActionsQueue<T: Config> =
935		StorageMap<_, Twox64Concat, SessionIndex, Vec<ParaId>, ValueQuery>;
936
937	/// Upcoming paras instantiation arguments.
938	///
939	/// NOTE that after PVF pre-checking is enabled the para genesis arg will have it's code set
940	/// to empty. Instead, the code will be saved into the storage right away via `CodeByHash`.
941	#[pezpallet::storage]
942	pub(super) type UpcomingParasGenesis<T: Config> =
943		StorageMap<_, Twox64Concat, ParaId, ParaGenesisArgs>;
944
945	/// The number of reference on the validation code in [`CodeByHash`] storage.
946	#[pezpallet::storage]
947	pub(super) type CodeByHashRefs<T: Config> =
948		StorageMap<_, Identity, ValidationCodeHash, u32, ValueQuery>;
949
950	/// Validation code stored by its hash.
951	///
952	/// This storage is consistent with [`FutureCodeHash`], [`CurrentCodeHash`] and
953	/// [`PastCodeHash`].
954	#[pezpallet::storage]
955	pub type CodeByHash<T: Config> = StorageMap<_, Identity, ValidationCodeHash, ValidationCode>;
956
957	#[pezpallet::genesis_config]
958	#[derive(DefaultNoBound)]
959	pub struct GenesisConfig<T: Config> {
960		#[serde(skip)]
961		pub _config: core::marker::PhantomData<T>,
962		pub paras: Vec<(ParaId, ParaGenesisArgs)>,
963	}
964
965	#[pezpallet::genesis_build]
966	impl<T: Config> BuildGenesisConfig for GenesisConfig<T> {
967		fn build(&self) {
968			let mut teyrchains = TeyrchainsCache::new();
969			for (id, genesis_args) in &self.paras {
970				if genesis_args.validation_code.0.is_empty() {
971					panic!("empty validation code is not allowed in genesis");
972				}
973				Pezpallet::<T>::initialize_para_now(&mut teyrchains, *id, genesis_args);
974				if genesis_args.para_kind == ParaKind::Teyrchain {
975					T::AssignCoretime::assign_coretime(*id)
976						.expect("Assigning coretime works at genesis; qed");
977				}
978			}
979			// teyrchains are flushed on drop
980		}
981	}
982
983	#[pezpallet::call]
984	impl<T: Config> Pezpallet<T> {
985		/// Set the storage for the teyrchain validation code immediately.
986		#[pezpallet::call_index(0)]
987		#[pezpallet::weight(<T as Config>::WeightInfo::force_set_current_code(new_code.0.len() as u32))]
988		pub fn force_set_current_code(
989			origin: OriginFor<T>,
990			para: ParaId,
991			new_code: ValidationCode,
992		) -> DispatchResult {
993			ensure_root(origin)?;
994			Self::do_force_set_current_code_update(para, new_code);
995			Ok(())
996		}
997
998		/// Set the storage for the current teyrchain head data immediately.
999		#[pezpallet::call_index(1)]
1000		#[pezpallet::weight(<T as Config>::WeightInfo::force_set_current_head(new_head.0.len() as u32))]
1001		pub fn force_set_current_head(
1002			origin: OriginFor<T>,
1003			para: ParaId,
1004			new_head: HeadData,
1005		) -> DispatchResult {
1006			ensure_root(origin)?;
1007			Self::set_current_head(para, new_head);
1008			Ok(())
1009		}
1010
1011		/// Schedule an upgrade as if it was scheduled in the given relay parent block.
1012		#[pezpallet::call_index(2)]
1013		#[pezpallet::weight(<T as Config>::WeightInfo::force_schedule_code_upgrade(new_code.0.len() as u32))]
1014		pub fn force_schedule_code_upgrade(
1015			origin: OriginFor<T>,
1016			para: ParaId,
1017			new_code: ValidationCode,
1018			relay_parent_number: BlockNumberFor<T>,
1019		) -> DispatchResult {
1020			ensure_root(origin)?;
1021			let config = configuration::ActiveConfig::<T>::get();
1022			Self::schedule_code_upgrade(
1023				para,
1024				new_code,
1025				relay_parent_number,
1026				&config,
1027				UpgradeStrategy::ApplyAtExpectedBlock,
1028			);
1029			Self::deposit_event(Event::CodeUpgradeScheduled(para));
1030			Ok(())
1031		}
1032
1033		/// Note a new block head for para within the context of the current block.
1034		#[pezpallet::call_index(3)]
1035		#[pezpallet::weight(<T as Config>::WeightInfo::force_note_new_head(new_head.0.len() as u32))]
1036		pub fn force_note_new_head(
1037			origin: OriginFor<T>,
1038			para: ParaId,
1039			new_head: HeadData,
1040		) -> DispatchResult {
1041			ensure_root(origin)?;
1042			let now = pezframe_system::Pezpallet::<T>::block_number();
1043			Self::note_new_head(para, new_head, now);
1044			Self::deposit_event(Event::NewHeadNoted(para));
1045			Ok(())
1046		}
1047
1048		/// Put a teyrchain directly into the next session's action queue.
1049		/// We can't queue it any sooner than this without going into the
1050		/// initializer...
1051		#[pezpallet::call_index(4)]
1052		#[pezpallet::weight(<T as Config>::WeightInfo::force_queue_action())]
1053		pub fn force_queue_action(origin: OriginFor<T>, para: ParaId) -> DispatchResult {
1054			ensure_root(origin)?;
1055			let next_session = shared::CurrentSessionIndex::<T>::get().saturating_add(One::one());
1056			ActionsQueue::<T>::mutate(next_session, |v| {
1057				if let Err(i) = v.binary_search(&para) {
1058					v.insert(i, para);
1059				}
1060			});
1061			Self::deposit_event(Event::ActionQueued(para, next_session));
1062			Ok(())
1063		}
1064
1065		/// Adds the validation code to the storage.
1066		///
1067		/// The code will not be added if it is already present. Additionally, if PVF pre-checking
1068		/// is running for that code, it will be instantly accepted.
1069		///
1070		/// Otherwise, the code will be added into the storage. Note that the code will be added
1071		/// into storage with reference count 0. This is to account the fact that there are no users
1072		/// for this code yet. The caller will have to make sure that this code eventually gets
1073		/// used by some teyrchain or removed from the storage to avoid storage leaks. For the
1074		/// latter prefer to use the `poke_unused_validation_code` dispatchable to raw storage
1075		/// manipulation.
1076		///
1077		/// This function is mainly meant to be used for upgrading teyrchains that do not follow
1078		/// the go-ahead signal while the PVF pre-checking feature is enabled.
1079		#[pezpallet::call_index(5)]
1080		#[pezpallet::weight(<T as Config>::WeightInfo::add_trusted_validation_code(validation_code.0.len() as u32))]
1081		pub fn add_trusted_validation_code(
1082			origin: OriginFor<T>,
1083			validation_code: ValidationCode,
1084		) -> DispatchResult {
1085			ensure_root(origin)?;
1086			let code_hash = validation_code.hash();
1087
1088			if let Some(vote) = PvfActiveVoteMap::<T>::get(&code_hash) {
1089				// Remove the existing vote.
1090				PvfActiveVoteMap::<T>::remove(&code_hash);
1091				PvfActiveVoteList::<T>::mutate(|l| {
1092					if let Ok(i) = l.binary_search(&code_hash) {
1093						l.remove(i);
1094					}
1095				});
1096
1097				let cfg = configuration::ActiveConfig::<T>::get();
1098				Self::enact_pvf_accepted(
1099					pezframe_system::Pezpallet::<T>::block_number(),
1100					&code_hash,
1101					&vote.causes,
1102					vote.age,
1103					&cfg,
1104				);
1105				return Ok(());
1106			}
1107
1108			if CodeByHash::<T>::contains_key(&code_hash) {
1109				// There is no vote, but the code exists. Nothing to do here.
1110				return Ok(());
1111			}
1112
1113			// At this point the code is unknown and there is no PVF pre-checking vote for it, so we
1114			// can just add the code into the storage.
1115			//
1116			// NOTE That we do not use `increase_code_ref` here, because the code is not yet used
1117			// by any teyrchain.
1118			CodeByHash::<T>::insert(code_hash, &validation_code);
1119
1120			Ok(())
1121		}
1122
1123		/// Remove the validation code from the storage iff the reference count is 0.
1124		///
1125		/// This is better than removing the storage directly, because it will not remove the code
1126		/// that was suddenly got used by some teyrchain while this dispatchable was pending
1127		/// dispatching.
1128		#[pezpallet::call_index(6)]
1129		#[pezpallet::weight(<T as Config>::WeightInfo::poke_unused_validation_code())]
1130		pub fn poke_unused_validation_code(
1131			origin: OriginFor<T>,
1132			validation_code_hash: ValidationCodeHash,
1133		) -> DispatchResult {
1134			ensure_root(origin)?;
1135			if CodeByHashRefs::<T>::get(&validation_code_hash) == 0 {
1136				CodeByHash::<T>::remove(&validation_code_hash);
1137			}
1138			Ok(())
1139		}
1140
1141		/// Includes a statement for a PVF pre-checking vote. Potentially, finalizes the vote and
1142		/// enacts the results if that was the last vote before achieving the supermajority.
1143		#[pezpallet::call_index(7)]
1144		#[pezpallet::weight(
1145			<T as Config>::WeightInfo::include_pvf_check_statement_finalize_upgrade_accept()
1146				.max(<T as Config>::WeightInfo::include_pvf_check_statement_finalize_upgrade_reject())
1147				.max(<T as Config>::WeightInfo::include_pvf_check_statement_finalize_onboarding_accept()
1148					.max(<T as Config>::WeightInfo::include_pvf_check_statement_finalize_onboarding_reject())
1149				)
1150		)]
1151		pub fn include_pvf_check_statement(
1152			origin: OriginFor<T>,
1153			stmt: PvfCheckStatement,
1154			signature: ValidatorSignature,
1155		) -> DispatchResultWithPostInfo {
1156			ensure_none(origin)?;
1157
1158			let validators = shared::ActiveValidatorKeys::<T>::get();
1159			let current_session = shared::CurrentSessionIndex::<T>::get();
1160			if stmt.session_index < current_session {
1161				return Err(Error::<T>::PvfCheckStatementStale.into());
1162			} else if stmt.session_index > current_session {
1163				return Err(Error::<T>::PvfCheckStatementFuture.into());
1164			}
1165			let validator_index = stmt.validator_index.0 as usize;
1166			let validator_public = validators
1167				.get(validator_index)
1168				.ok_or(Error::<T>::PvfCheckValidatorIndexOutOfBounds)?;
1169
1170			let signing_payload = stmt.signing_payload();
1171			ensure!(
1172				signature.verify(&signing_payload[..], &validator_public),
1173				Error::<T>::PvfCheckInvalidSignature,
1174			);
1175
1176			let mut active_vote = PvfActiveVoteMap::<T>::get(&stmt.subject)
1177				.ok_or(Error::<T>::PvfCheckSubjectInvalid)?;
1178
1179			// Ensure that the validator submitting this statement hasn't voted already.
1180			ensure!(
1181				!active_vote
1182					.has_vote(validator_index)
1183					.ok_or(Error::<T>::PvfCheckValidatorIndexOutOfBounds)?,
1184				Error::<T>::PvfCheckDoubleVote,
1185			);
1186
1187			// Finally, cast the vote and persist.
1188			if stmt.accept {
1189				active_vote.votes_accept.set(validator_index, true);
1190			} else {
1191				active_vote.votes_reject.set(validator_index, true);
1192			}
1193
1194			if let Some(outcome) = active_vote.quorum(validators.len()) {
1195				// The quorum has been achieved.
1196				//
1197				// Remove the PVF vote from the active map and finalize the PVF checking according
1198				// to the outcome.
1199				PvfActiveVoteMap::<T>::remove(&stmt.subject);
1200				PvfActiveVoteList::<T>::mutate(|l| {
1201					if let Ok(i) = l.binary_search(&stmt.subject) {
1202						l.remove(i);
1203					}
1204				});
1205				match outcome {
1206					PvfCheckOutcome::Accepted => {
1207						let cfg = configuration::ActiveConfig::<T>::get();
1208						Self::enact_pvf_accepted(
1209							pezframe_system::Pezpallet::<T>::block_number(),
1210							&stmt.subject,
1211							&active_vote.causes,
1212							active_vote.age,
1213							&cfg,
1214						);
1215					},
1216					PvfCheckOutcome::Rejected => {
1217						Self::enact_pvf_rejected(&stmt.subject, active_vote.causes);
1218					},
1219				}
1220
1221				// No weight refund since this statement was the last one and lead to finalization.
1222				Ok(().into())
1223			} else {
1224				// No quorum has been achieved.
1225				//
1226				// - So just store the updated state back into the storage.
1227				// - Only charge weight for simple vote inclusion.
1228				PvfActiveVoteMap::<T>::insert(&stmt.subject, active_vote);
1229				Ok(Some(<T as Config>::WeightInfo::include_pvf_check_statement()).into())
1230			}
1231		}
1232
1233		/// Set the storage for the current teyrchain head data immediately.
1234		#[pezpallet::call_index(8)]
1235		#[pezpallet::weight(<T as Config>::WeightInfo::force_set_most_recent_context())]
1236		pub fn force_set_most_recent_context(
1237			origin: OriginFor<T>,
1238			para: ParaId,
1239			context: BlockNumberFor<T>,
1240		) -> DispatchResult {
1241			ensure_root(origin)?;
1242			MostRecentContext::<T>::insert(&para, context);
1243			Ok(())
1244		}
1245
1246		/// Remove an upgrade cooldown for a teyrchain.
1247		///
1248		/// The cost for removing the cooldown earlier depends on the time left for the cooldown
1249		/// multiplied by [`Config::CooldownRemovalMultiplier`]. The paid tokens are burned.
1250		#[pezpallet::call_index(9)]
1251		#[pezpallet::weight(<T as Config>::WeightInfo::remove_upgrade_cooldown())]
1252		pub fn remove_upgrade_cooldown(origin: OriginFor<T>, para: ParaId) -> DispatchResult {
1253			let who = ensure_signed(origin)?;
1254
1255			let removed = UpgradeCooldowns::<T>::mutate(|cooldowns| {
1256				let Some(pos) = cooldowns.iter().position(|(p, _)| p == &para) else {
1257					return Ok::<_, DispatchError>(false);
1258				};
1259				let (_, cooldown_until) = cooldowns.remove(pos);
1260
1261				let cost = Self::calculate_remove_upgrade_cooldown_cost(cooldown_until);
1262
1263				// burn...
1264				T::Fungible::burn_from(
1265					&who,
1266					cost,
1267					Preservation::Preserve,
1268					Precision::Exact,
1269					Fortitude::Polite,
1270				)?;
1271
1272				Ok(true)
1273			})?;
1274
1275			if removed {
1276				UpgradeRestrictionSignal::<T>::remove(para);
1277
1278				Self::deposit_event(Event::UpgradeCooldownRemoved { para_id: para });
1279			}
1280
1281			Ok(())
1282		}
1283
1284		/// Sets the storage for the authorized current code hash of the teyrchain.
1285		/// If not applied, it will be removed at the `System::block_number() + valid_period` block.
1286		///
1287		/// This can be useful, when triggering `Paras::force_set_current_code(para, code)`
1288		/// from a different chain than the one where the `Paras` pezpallet is deployed.
1289		///
1290		/// The main purpose is to avoid transferring the entire `code` Wasm blob between chains.
1291		/// Instead, we authorize `code_hash` with `root`, which can later be applied by
1292		/// `Paras::apply_authorized_force_set_current_code(para, code)` by anyone.
1293		///
1294		/// Authorizations are stored in an **overwriting manner**.
1295		#[pezpallet::call_index(10)]
1296		#[pezpallet::weight(<T as Config>::WeightInfo::authorize_force_set_current_code_hash())]
1297		pub fn authorize_force_set_current_code_hash(
1298			origin: OriginFor<T>,
1299			para: ParaId,
1300			new_code_hash: ValidationCodeHash,
1301			valid_period: BlockNumberFor<T>,
1302		) -> DispatchResult {
1303			T::AuthorizeCurrentCodeOrigin::ensure_origin(origin, &para)?;
1304			// The requested para must be a valid para (neither onboarding nor offboarding).
1305			ensure!(Self::is_valid_para(para), Error::<T>::NotRegistered);
1306
1307			let now = pezframe_system::Pezpallet::<T>::block_number();
1308			let expire_at = now.saturating_add(valid_period);
1309
1310			// Insert the authorized code hash and ensure it overwrites the existing one for a para.
1311			AuthorizedCodeHash::<T>::insert(
1312				&para,
1313				AuthorizedCodeHashAndExpiry::from((new_code_hash, expire_at)),
1314			);
1315			Self::deposit_event(Event::CodeAuthorized {
1316				para_id: para,
1317				code_hash: new_code_hash,
1318				expire_at,
1319			});
1320
1321			Ok(())
1322		}
1323
1324		/// Applies the already authorized current code for the teyrchain,
1325		/// triggering the same functionality as `force_set_current_code`.
1326		#[pezpallet::call_index(11)]
1327		#[pezpallet::weight(<T as Config>::WeightInfo::apply_authorized_force_set_current_code(new_code.0.len() as u32))]
1328		pub fn apply_authorized_force_set_current_code(
1329			_origin: OriginFor<T>,
1330			para: ParaId,
1331			new_code: ValidationCode,
1332		) -> DispatchResultWithPostInfo {
1333			// no need to ensure anybody can do this
1334
1335			// Ensure `new_code` is authorized
1336			let _ = Self::validate_code_is_authorized(&new_code, &para)?;
1337			// Remove authorization
1338			AuthorizedCodeHash::<T>::remove(para);
1339
1340			// apply/dispatch
1341			Self::do_force_set_current_code_update(para, new_code);
1342
1343			Ok(Pays::No.into())
1344		}
1345	}
1346
1347	impl<T: Config> Pezpallet<T> {
1348		pub(crate) fn calculate_remove_upgrade_cooldown_cost(
1349			cooldown_until: BlockNumberFor<T>,
1350		) -> BalanceOf<T> {
1351			let time_left =
1352				cooldown_until.saturating_sub(pezframe_system::Pezpallet::<T>::block_number());
1353
1354			BalanceOf::<T>::from(time_left).saturating_mul(T::CooldownRemovalMultiplier::get())
1355		}
1356	}
1357
1358	#[pezpallet::view_functions]
1359	impl<T: Config> Pezpallet<T> {
1360		/// Returns the cost for removing an upgrade cooldown for the given `para`.
1361		pub fn remove_upgrade_cooldown_cost(para: ParaId) -> BalanceOf<T> {
1362			UpgradeCooldowns::<T>::get()
1363				.iter()
1364				.find(|(p, _)| p == &para)
1365				.map(|(_, c)| Self::calculate_remove_upgrade_cooldown_cost(*c))
1366				.unwrap_or_default()
1367		}
1368	}
1369
1370	#[pezpallet::validate_unsigned]
1371	impl<T: Config> ValidateUnsigned for Pezpallet<T> {
1372		type Call = Call<T>;
1373
1374		fn validate_unsigned(_source: TransactionSource, call: &Self::Call) -> TransactionValidity {
1375			match call {
1376				Call::include_pvf_check_statement { stmt, signature } => {
1377					let current_session = shared::CurrentSessionIndex::<T>::get();
1378					if stmt.session_index < current_session {
1379						return InvalidTransaction::Stale.into();
1380					} else if stmt.session_index > current_session {
1381						return InvalidTransaction::Future.into();
1382					}
1383
1384					let validator_index = stmt.validator_index.0 as usize;
1385					let validators = shared::ActiveValidatorKeys::<T>::get();
1386					let validator_public = match validators.get(validator_index) {
1387						Some(pk) => pk,
1388						None => {
1389							return InvalidTransaction::Custom(INVALID_TX_BAD_VALIDATOR_IDX).into()
1390						},
1391					};
1392
1393					let signing_payload = stmt.signing_payload();
1394					if !signature.verify(&signing_payload[..], &validator_public) {
1395						return InvalidTransaction::BadProof.into();
1396					}
1397
1398					let active_vote = match PvfActiveVoteMap::<T>::get(&stmt.subject) {
1399						Some(v) => v,
1400						None => return InvalidTransaction::Custom(INVALID_TX_BAD_SUBJECT).into(),
1401					};
1402
1403					match active_vote.has_vote(validator_index) {
1404						Some(false) => (),
1405						Some(true) => {
1406							return InvalidTransaction::Custom(INVALID_TX_DOUBLE_VOTE).into()
1407						},
1408						None => {
1409							return InvalidTransaction::Custom(INVALID_TX_BAD_VALIDATOR_IDX).into()
1410						},
1411					}
1412
1413					ValidTransaction::with_tag_prefix("PvfPreCheckingVote")
1414						.priority(T::UnsignedPriority::get())
1415						.longevity(
1416							TryInto::<u64>::try_into(
1417								T::NextSessionRotation::average_session_length() / 2u32.into(),
1418							)
1419							.unwrap_or(64_u64),
1420						)
1421						.and_provides((stmt.session_index, stmt.validator_index, stmt.subject))
1422						.propagate(true)
1423						.build()
1424				},
1425				Call::apply_authorized_force_set_current_code { para, new_code } => {
1426					match Self::validate_code_is_authorized(new_code, para) {
1427						Ok(authorized_code) => {
1428							let now = pezframe_system::Pezpallet::<T>::block_number();
1429							let longevity = authorized_code.expire_at.saturating_sub(now);
1430
1431							ValidTransaction::with_tag_prefix("ApplyAuthorizedForceSetCurrentCode")
1432								.priority(T::UnsignedPriority::get())
1433								.longevity(TryInto::<u64>::try_into(longevity).unwrap_or(64_u64))
1434								.and_provides((para, authorized_code.code_hash))
1435								.propagate(true)
1436								.build()
1437						},
1438						Err(_) => {
1439							return InvalidTransaction::Custom(INVALID_TX_UNAUTHORIZED_CODE).into()
1440						},
1441					}
1442				},
1443				_ => InvalidTransaction::Call.into(),
1444			}
1445		}
1446
1447		fn pre_dispatch(_call: &Self::Call) -> Result<(), TransactionValidityError> {
1448			// Return `Ok` here meaning that as soon as the transaction got into the block, it will
1449			// always dispatched. This is OK, since the `include_pvf_check_statement` dispatchable
1450			// will perform the same checks anyway, so there is no point doing it here.
1451			//
1452			// On the other hand, if we did not provide the implementation, then the default
1453			// implementation would be used. The default implementation just delegates the
1454			// pre-dispatch validation to `validate_unsigned`.
1455			Ok(())
1456		}
1457	}
1458}
1459
1460// custom transaction error codes
1461const INVALID_TX_BAD_VALIDATOR_IDX: u8 = 1;
1462const INVALID_TX_BAD_SUBJECT: u8 = 2;
1463const INVALID_TX_DOUBLE_VOTE: u8 = 3;
1464const INVALID_TX_UNAUTHORIZED_CODE: u8 = 4;
1465
1466/// This is intermediate "fix" for this issue:
1467/// <https://github.com/pezkuwichain/pezkuwi-sdk/issues/281>
1468///
1469/// It does not actually fix it, but makes the worst case better. Without that limit someone
1470/// could completely DoS the relay chain by registering a ridiculously high amount of paras.
1471/// With this limit the same attack could lead to some teyrchains ceasing to being able to
1472/// communicate via offchain XCMP. Snowbridge will still work as it only cares about `BridgeHub`.
1473pub const MAX_PARA_HEADS: usize = 1024;
1474
1475impl<T: Config> Pezpallet<T> {
1476	/// This is a call to schedule code upgrades for teyrchains which is safe to be called
1477	/// outside of this module. That means this function does all checks necessary to ensure
1478	/// that some external code is allowed to trigger a code upgrade. We do not do auth checks,
1479	/// that should be handled by whomever calls this function.
1480	pub(crate) fn schedule_code_upgrade_external(
1481		id: ParaId,
1482		new_code: ValidationCode,
1483		upgrade_strategy: UpgradeStrategy,
1484	) -> DispatchResult {
1485		// Check that we can schedule an upgrade at all.
1486		ensure!(Self::can_upgrade_validation_code(id), Error::<T>::CannotUpgradeCode);
1487		let config = configuration::ActiveConfig::<T>::get();
1488		// Validation code sanity checks:
1489		ensure!(new_code.0.len() >= MIN_CODE_SIZE as usize, Error::<T>::InvalidCode);
1490		ensure!(new_code.0.len() <= config.max_code_size as usize, Error::<T>::InvalidCode);
1491
1492		let current_block = pezframe_system::Pezpallet::<T>::block_number();
1493		// Schedule the upgrade with a delay just like if a teyrchain triggered the upgrade.
1494		let upgrade_block = current_block.saturating_add(config.validation_upgrade_delay);
1495		Self::schedule_code_upgrade(id, new_code, upgrade_block, &config, upgrade_strategy);
1496		Self::deposit_event(Event::CodeUpgradeScheduled(id));
1497		Ok(())
1498	}
1499
1500	/// Set the current head of a teyrchain.
1501	pub(crate) fn set_current_head(para: ParaId, new_head: HeadData) {
1502		Heads::<T>::insert(&para, new_head);
1503		Self::deposit_event(Event::CurrentHeadUpdated(para));
1504	}
1505
1506	/// Called by the initializer to initialize the paras pezpallet.
1507	pub(crate) fn initializer_initialize(now: BlockNumberFor<T>) -> Weight {
1508		Self::prune_old_code(now)
1509			+ Self::process_scheduled_upgrade_changes(now)
1510			+ Self::process_future_code_upgrades_at(now)
1511			+ Self::prune_expired_authorizations(now)
1512	}
1513
1514	/// Called by the initializer to finalize the paras pezpallet.
1515	pub(crate) fn initializer_finalize(now: BlockNumberFor<T>) {
1516		Self::process_scheduled_upgrade_cooldowns(now);
1517	}
1518
1519	/// Called by the initializer to note that a new session has started.
1520	///
1521	/// Returns the list of outgoing paras from the actions queue.
1522	pub(crate) fn initializer_on_new_session(
1523		notification: &SessionChangeNotification<BlockNumberFor<T>>,
1524	) -> Vec<ParaId> {
1525		let outgoing_paras = Self::apply_actions_queue(notification.session_index);
1526		Self::groom_ongoing_pvf_votes(&notification.new_config, notification.validators.len());
1527		outgoing_paras
1528	}
1529
1530	/// The validation code of live para.
1531	pub(crate) fn current_code(para_id: &ParaId) -> Option<ValidationCode> {
1532		CurrentCodeHash::<T>::get(para_id).and_then(|code_hash| {
1533			let code = CodeByHash::<T>::get(&code_hash);
1534			if code.is_none() {
1535				log::error!(
1536					"Pezpallet paras storage is inconsistent, code not found for hash {}",
1537					code_hash,
1538				);
1539				debug_assert!(false, "inconsistent paras storages");
1540			}
1541			code
1542		})
1543	}
1544
1545	/// Get a list of the first [`MAX_PARA_HEADS`] para heads sorted by para_id.
1546	/// This method is likely to be removed in the future.
1547	pub fn sorted_para_heads() -> Vec<(u32, Vec<u8>)> {
1548		let mut heads: Vec<(u32, Vec<u8>)> =
1549			Heads::<T>::iter().map(|(id, head)| (id.into(), head.0)).collect();
1550		heads.sort_by_key(|(id, _)| *id);
1551		heads.truncate(MAX_PARA_HEADS);
1552		heads
1553	}
1554
1555	// Apply all para actions queued for the given session index.
1556	//
1557	// The actions to take are based on the lifecycle of of the paras.
1558	//
1559	// The final state of any para after the actions queue should be as a
1560	// lease holding teyrchain, on-demand teyrchain, or not registered. (stable states)
1561	//
1562	// Returns the list of outgoing paras from the actions queue.
1563	fn apply_actions_queue(session: SessionIndex) -> Vec<ParaId> {
1564		let actions = ActionsQueue::<T>::take(session);
1565		let mut teyrchains = TeyrchainsCache::new();
1566		let now = pezframe_system::Pezpallet::<T>::block_number();
1567		let mut outgoing = Vec::new();
1568
1569		for para in actions {
1570			let lifecycle = ParaLifecycles::<T>::get(&para);
1571			match lifecycle {
1572				None | Some(ParaLifecycle::Parathread) | Some(ParaLifecycle::Teyrchain) => { /* Nothing to do... */
1573				},
1574				Some(ParaLifecycle::Onboarding) => {
1575					if let Some(genesis_data) = UpcomingParasGenesis::<T>::take(&para) {
1576						Self::initialize_para_now(&mut teyrchains, para, &genesis_data);
1577					}
1578				},
1579				// Upgrade an on-demand teyrchain to a lease holding teyrchain
1580				Some(ParaLifecycle::UpgradingParathread) => {
1581					teyrchains.add(para);
1582					ParaLifecycles::<T>::insert(&para, ParaLifecycle::Teyrchain);
1583				},
1584				// Downgrade a lease holding teyrchain to an on-demand teyrchain
1585				Some(ParaLifecycle::DowngradingTeyrchain) => {
1586					teyrchains.remove(para);
1587					ParaLifecycles::<T>::insert(&para, ParaLifecycle::Parathread);
1588				},
1589				// Offboard a lease holding or on-demand teyrchain from the system
1590				Some(ParaLifecycle::OffboardingTeyrchain)
1591				| Some(ParaLifecycle::OffboardingParathread) => {
1592					teyrchains.remove(para);
1593
1594					Heads::<T>::remove(&para);
1595					MostRecentContext::<T>::remove(&para);
1596					FutureCodeUpgrades::<T>::remove(&para);
1597					UpgradeGoAheadSignal::<T>::remove(&para);
1598					UpgradeRestrictionSignal::<T>::remove(&para);
1599					ParaLifecycles::<T>::remove(&para);
1600					AuthorizedCodeHash::<T>::remove(&para);
1601					let removed_future_code_hash = FutureCodeHash::<T>::take(&para);
1602					if let Some(removed_future_code_hash) = removed_future_code_hash {
1603						Self::decrease_code_ref(&removed_future_code_hash);
1604					}
1605
1606					let removed_code_hash = CurrentCodeHash::<T>::take(&para);
1607					if let Some(removed_code_hash) = removed_code_hash {
1608						Self::note_past_code(para, now, now, removed_code_hash);
1609					}
1610
1611					outgoing.push(para);
1612				},
1613			}
1614		}
1615
1616		if !outgoing.is_empty() {
1617			// Filter offboarded teyrchains from the upcoming upgrades and upgrade cooldowns list.
1618			//
1619			// We do it after the offboarding to get away with only a single read/write per list.
1620			//
1621			// NOTE both of those iterates over the list and the outgoing. We do not expect either
1622			//      of these to be large. Thus should be fine.
1623			UpcomingUpgrades::<T>::mutate(|upcoming_upgrades| {
1624				upcoming_upgrades.retain(|(para, _)| !outgoing.contains(para));
1625			});
1626			UpgradeCooldowns::<T>::mutate(|upgrade_cooldowns| {
1627				upgrade_cooldowns.retain(|(para, _)| !outgoing.contains(para));
1628			});
1629			FutureCodeUpgradesAt::<T>::mutate(|future_upgrades| {
1630				future_upgrades.retain(|(para, _)| !outgoing.contains(para));
1631			});
1632		}
1633
1634		// Persist teyrchains into the storage explicitly.
1635		drop(teyrchains);
1636
1637		outgoing
1638	}
1639
1640	// note replacement of the code of para with given `id`, which occurred in the
1641	// context of the given relay-chain block number. provide the replaced code.
1642	//
1643	// `at` for para-triggered replacement is the block number of the relay-chain
1644	// block in whose context the parablock was executed
1645	// (i.e. number of `relay_parent` in the receipt)
1646	fn note_past_code(
1647		id: ParaId,
1648		at: BlockNumberFor<T>,
1649		now: BlockNumberFor<T>,
1650		old_code_hash: ValidationCodeHash,
1651	) -> Weight {
1652		PastCodeMeta::<T>::mutate(&id, |past_meta| {
1653			past_meta.note_replacement(at, now);
1654		});
1655
1656		PastCodeHash::<T>::insert(&(id, at), old_code_hash);
1657
1658		// Schedule pruning for this past-code to be removed as soon as it
1659		// exits the slashing window.
1660		PastCodePruning::<T>::mutate(|pruning| {
1661			let insert_idx =
1662				pruning.binary_search_by_key(&now, |&(_, b)| b).unwrap_or_else(|idx| idx);
1663			pruning.insert(insert_idx, (id, now));
1664		});
1665
1666		T::DbWeight::get().reads_writes(2, 3)
1667	}
1668
1669	// looks at old code metadata, compares them to the current acceptance window, and prunes those
1670	// that are too old.
1671	fn prune_old_code(now: BlockNumberFor<T>) -> Weight {
1672		let config = configuration::ActiveConfig::<T>::get();
1673		let code_retention_period = config.code_retention_period;
1674		if now <= code_retention_period {
1675			let weight = T::DbWeight::get().reads_writes(1, 0);
1676			return weight;
1677		}
1678
1679		// The height of any changes we no longer should keep around.
1680		let pruning_height = now - (code_retention_period + One::one());
1681
1682		let pruning_tasks_done =
1683			PastCodePruning::<T>::mutate(|pruning_tasks: &mut Vec<(_, BlockNumberFor<T>)>| {
1684				let (pruning_tasks_done, pruning_tasks_to_do) = {
1685					// find all past code that has just exited the pruning window.
1686					let up_to_idx =
1687						pruning_tasks.iter().take_while(|&(_, at)| at <= &pruning_height).count();
1688					(up_to_idx, pruning_tasks.drain(..up_to_idx))
1689				};
1690
1691				for (para_id, _) in pruning_tasks_to_do {
1692					let full_deactivate = PastCodeMeta::<T>::mutate(&para_id, |meta| {
1693						for pruned_repl_at in meta.prune_up_to(pruning_height) {
1694							let removed_code_hash =
1695								PastCodeHash::<T>::take(&(para_id, pruned_repl_at));
1696
1697							if let Some(removed_code_hash) = removed_code_hash {
1698								Self::decrease_code_ref(&removed_code_hash);
1699							} else {
1700								log::warn!(
1701									target: LOG_TARGET,
1702									"Missing code for removed hash {:?}",
1703									removed_code_hash,
1704								);
1705							}
1706						}
1707
1708						meta.is_empty() && Heads::<T>::get(&para_id).is_none()
1709					});
1710
1711					// This teyrchain has been removed and now the vestigial code
1712					// has been removed from the state. clean up meta as well.
1713					if full_deactivate {
1714						PastCodeMeta::<T>::remove(&para_id);
1715					}
1716				}
1717
1718				pruning_tasks_done as u64
1719			});
1720
1721		// 1 read for the meta for each pruning task, 1 read for the config
1722		// 2 writes: updating the meta and pruning the code
1723		T::DbWeight::get().reads_writes(1 + pruning_tasks_done, 2 * pruning_tasks_done)
1724	}
1725
1726	/// This function removes authorizations that have expired,
1727	/// meaning their `expire_at` block is less than or equal to the current block (`now`).
1728	fn prune_expired_authorizations(now: BlockNumberFor<T>) -> Weight {
1729		let mut weight = T::DbWeight::get().reads(1);
1730		let to_remove = AuthorizedCodeHash::<T>::iter().filter_map(
1731			|(para, AuthorizedCodeHashAndExpiry { expire_at, .. })| {
1732				if expire_at <= now {
1733					Some(para)
1734				} else {
1735					None
1736				}
1737			},
1738		);
1739		for para in to_remove {
1740			AuthorizedCodeHash::<T>::remove(&para);
1741			weight.saturating_accrue(T::DbWeight::get().writes(1));
1742		}
1743
1744		weight
1745	}
1746
1747	/// Process the future code upgrades that should be applied directly.
1748	///
1749	/// Upgrades that should not be applied directly are being processed in
1750	/// [`Self::process_scheduled_upgrade_changes`].
1751	fn process_future_code_upgrades_at(now: BlockNumberFor<T>) -> Weight {
1752		// account weight for `FutureCodeUpgradeAt::mutate`.
1753		let mut weight = T::DbWeight::get().reads_writes(1, 1);
1754		FutureCodeUpgradesAt::<T>::mutate(
1755			|upcoming_upgrades: &mut Vec<(ParaId, BlockNumberFor<T>)>| {
1756				let num = upcoming_upgrades.iter().take_while(|&(_, at)| at <= &now).count();
1757				for (id, expected_at) in upcoming_upgrades.drain(..num) {
1758					weight += T::DbWeight::get().reads_writes(1, 1);
1759
1760					// Both should always be `Some` in this case, since a code upgrade is scheduled.
1761					let new_code_hash = if let Some(new_code_hash) = FutureCodeHash::<T>::take(&id)
1762					{
1763						new_code_hash
1764					} else {
1765						log::error!(target: LOG_TARGET, "Missing future code hash for {:?}", &id);
1766						continue;
1767					};
1768
1769					weight += Self::set_current_code(id, new_code_hash, expected_at);
1770				}
1771				num
1772			},
1773		);
1774
1775		weight
1776	}
1777
1778	/// Process the timers related to upgrades. Specifically, the upgrade go ahead signals toggle
1779	/// and the upgrade cooldown restrictions. However, this function does not actually unset
1780	/// the upgrade restriction, that will happen in the `initializer_finalize` function. However,
1781	/// this function does count the number of cooldown timers expired so that we can reserve weight
1782	/// for the `initializer_finalize` function.
1783	fn process_scheduled_upgrade_changes(now: BlockNumberFor<T>) -> Weight {
1784		// account weight for `UpcomingUpgrades::mutate`.
1785		let mut weight = T::DbWeight::get().reads_writes(1, 1);
1786		let upgrades_signaled = UpcomingUpgrades::<T>::mutate(
1787			|upcoming_upgrades: &mut Vec<(ParaId, BlockNumberFor<T>)>| {
1788				let num = upcoming_upgrades.iter().take_while(|&(_, at)| at <= &now).count();
1789				for (para, _) in upcoming_upgrades.drain(..num) {
1790					UpgradeGoAheadSignal::<T>::insert(&para, UpgradeGoAhead::GoAhead);
1791				}
1792				num
1793			},
1794		);
1795		weight += T::DbWeight::get().writes(upgrades_signaled as u64);
1796
1797		// account weight for `UpgradeCooldowns::get`.
1798		weight += T::DbWeight::get().reads(1);
1799		let cooldowns_expired =
1800			UpgradeCooldowns::<T>::get().iter().take_while(|&(_, at)| at <= &now).count();
1801
1802		// reserve weight for `initializer_finalize`:
1803		// - 1 read and 1 write for `UpgradeCooldowns::mutate`.
1804		// - 1 write per expired cooldown.
1805		weight += T::DbWeight::get().reads_writes(1, 1);
1806		weight += T::DbWeight::get().reads(cooldowns_expired as u64);
1807
1808		weight
1809	}
1810
1811	/// Actually perform unsetting the expired upgrade restrictions.
1812	///
1813	/// See `process_scheduled_upgrade_changes` for more details.
1814	fn process_scheduled_upgrade_cooldowns(now: BlockNumberFor<T>) {
1815		UpgradeCooldowns::<T>::mutate(
1816			|upgrade_cooldowns: &mut Vec<(ParaId, BlockNumberFor<T>)>| {
1817				// Remove all expired signals and also prune the cooldowns.
1818				upgrade_cooldowns.retain(|(para, at)| {
1819					if at <= &now {
1820						UpgradeRestrictionSignal::<T>::remove(&para);
1821						false
1822					} else {
1823						true
1824					}
1825				});
1826			},
1827		);
1828	}
1829
1830	/// Goes over all PVF votes in progress, reinitializes ballots, increments ages and prunes the
1831	/// active votes that reached their time-to-live.
1832	fn groom_ongoing_pvf_votes(
1833		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
1834		new_n_validators: usize,
1835	) -> Weight {
1836		let mut weight = T::DbWeight::get().reads(1);
1837
1838		let potentially_active_votes = PvfActiveVoteList::<T>::get();
1839
1840		// Initially empty list which contains all the PVF active votes that made it through this
1841		// session change.
1842		//
1843		// **Ordered** as well as `PvfActiveVoteList`.
1844		let mut actually_active_votes = Vec::with_capacity(potentially_active_votes.len());
1845
1846		for vote_subject in potentially_active_votes {
1847			let mut vote_state = match PvfActiveVoteMap::<T>::take(&vote_subject) {
1848				Some(v) => v,
1849				None => {
1850					// This branch should never be reached. This is due to the fact that the set of
1851					// `PvfActiveVoteMap`'s keys is always equal to the set of items found in
1852					// `PvfActiveVoteList`.
1853					log::warn!(
1854						target: LOG_TARGET,
1855						"The PvfActiveVoteMap is out of sync with PvfActiveVoteList!",
1856					);
1857					debug_assert!(false);
1858					continue;
1859				},
1860			};
1861
1862			vote_state.age += 1;
1863			if vote_state.age < cfg.pvf_voting_ttl {
1864				weight += T::DbWeight::get().writes(1);
1865				vote_state.reinitialize_ballots(new_n_validators);
1866				PvfActiveVoteMap::<T>::insert(&vote_subject, vote_state);
1867
1868				// push maintaining the original order.
1869				actually_active_votes.push(vote_subject);
1870			} else {
1871				// TTL is reached. Reject.
1872				weight += Self::enact_pvf_rejected(&vote_subject, vote_state.causes);
1873			}
1874		}
1875
1876		weight += T::DbWeight::get().writes(1);
1877		PvfActiveVoteList::<T>::put(actually_active_votes);
1878
1879		weight
1880	}
1881
1882	fn enact_pvf_accepted(
1883		now: BlockNumberFor<T>,
1884		code_hash: &ValidationCodeHash,
1885		causes: &[PvfCheckCause<BlockNumberFor<T>>],
1886		sessions_observed: SessionIndex,
1887		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
1888	) -> Weight {
1889		let mut weight = Weight::zero();
1890		for cause in causes {
1891			weight += T::DbWeight::get().reads_writes(3, 2);
1892			Self::deposit_event(Event::PvfCheckAccepted(*code_hash, cause.para_id()));
1893
1894			match cause {
1895				PvfCheckCause::Onboarding(id) => {
1896					weight += Self::proceed_with_onboarding(*id, sessions_observed);
1897				},
1898				PvfCheckCause::Upgrade { id, included_at, upgrade_strategy } => {
1899					weight += Self::proceed_with_upgrade(
1900						*id,
1901						code_hash,
1902						now,
1903						*included_at,
1904						cfg,
1905						*upgrade_strategy,
1906					);
1907				},
1908			}
1909		}
1910		weight
1911	}
1912
1913	fn proceed_with_onboarding(id: ParaId, sessions_observed: SessionIndex) -> Weight {
1914		let weight = T::DbWeight::get().reads_writes(2, 1);
1915
1916		// we should onboard only after `SESSION_DELAY` sessions but we should take
1917		// into account the number of sessions the PVF pre-checking occupied.
1918		//
1919		// we cannot onboard at the current session, so it must be at least one
1920		// session ahead.
1921		let onboard_at: SessionIndex = shared::CurrentSessionIndex::<T>::get()
1922			+ cmp::max(shared::SESSION_DELAY.saturating_sub(sessions_observed), 1);
1923
1924		ActionsQueue::<T>::mutate(onboard_at, |v| {
1925			if let Err(i) = v.binary_search(&id) {
1926				v.insert(i, id);
1927			}
1928		});
1929
1930		weight
1931	}
1932
1933	fn proceed_with_upgrade(
1934		id: ParaId,
1935		code_hash: &ValidationCodeHash,
1936		now: BlockNumberFor<T>,
1937		relay_parent_number: BlockNumberFor<T>,
1938		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
1939		upgrade_strategy: UpgradeStrategy,
1940	) -> Weight {
1941		let mut weight = Weight::zero();
1942
1943		// Compute the relay-chain block number starting at which the code upgrade is ready to
1944		// be applied.
1945		//
1946		// The first parablock that has a relay-parent higher or at the same height of
1947		// `expected_at` will trigger the code upgrade. The parablock that comes after that will
1948		// be validated against the new validation code.
1949		//
1950		// Here we are trying to choose the block number that will have
1951		// `validation_upgrade_delay` blocks from the relay-parent of inclusion of the the block
1952		// that scheduled code upgrade but no less than `minimum_validation_upgrade_delay`. We
1953		// want this delay out of caution so that when the last vote for pre-checking comes the
1954		// teyrchain will have some time until the upgrade finally takes place.
1955		let expected_at = cmp::max(
1956			relay_parent_number + cfg.validation_upgrade_delay,
1957			now + cfg.minimum_validation_upgrade_delay,
1958		);
1959
1960		match upgrade_strategy {
1961			UpgradeStrategy::ApplyAtExpectedBlock => {
1962				FutureCodeUpgradesAt::<T>::mutate(|future_upgrades| {
1963					let insert_idx = future_upgrades
1964						.binary_search_by_key(&expected_at, |&(_, b)| b)
1965						.unwrap_or_else(|idx| idx);
1966					future_upgrades.insert(insert_idx, (id, expected_at));
1967				});
1968
1969				weight += T::DbWeight::get().reads_writes(0, 2);
1970			},
1971			UpgradeStrategy::SetGoAheadSignal => {
1972				FutureCodeUpgrades::<T>::insert(&id, expected_at);
1973
1974				UpcomingUpgrades::<T>::mutate(|upcoming_upgrades| {
1975					let insert_idx = upcoming_upgrades
1976						.binary_search_by_key(&expected_at, |&(_, b)| b)
1977						.unwrap_or_else(|idx| idx);
1978					upcoming_upgrades.insert(insert_idx, (id, expected_at));
1979				});
1980
1981				weight += T::DbWeight::get().reads_writes(1, 3);
1982			},
1983		}
1984
1985		let expected_at = expected_at.saturated_into();
1986		let log = ConsensusLog::ParaScheduleUpgradeCode(id, *code_hash, expected_at);
1987		pezframe_system::Pezpallet::<T>::deposit_log(log.into());
1988
1989		weight
1990	}
1991
1992	fn enact_pvf_rejected(
1993		code_hash: &ValidationCodeHash,
1994		causes: Vec<PvfCheckCause<BlockNumberFor<T>>>,
1995	) -> Weight {
1996		let mut weight = Weight::zero();
1997
1998		for cause in causes {
1999			// Whenever PVF pre-checking is started or a new cause is added to it, the RC is bumped.
2000			// Now we need to unbump it.
2001			weight += Self::decrease_code_ref(code_hash);
2002
2003			weight += T::DbWeight::get().reads_writes(3, 2);
2004			Self::deposit_event(Event::PvfCheckRejected(*code_hash, cause.para_id()));
2005
2006			match cause {
2007				PvfCheckCause::Onboarding(id) => {
2008					// Here we need to undo everything that was done during
2009					// `schedule_para_initialize`. Essentially, the logic is similar to offboarding,
2010					// with exception that before actual onboarding the teyrchain did not have a
2011					// chance to reach to upgrades. Therefore we can skip all the upgrade related
2012					// storage items here.
2013					weight += T::DbWeight::get().writes(3);
2014					UpcomingParasGenesis::<T>::remove(&id);
2015					CurrentCodeHash::<T>::remove(&id);
2016					ParaLifecycles::<T>::remove(&id);
2017				},
2018				PvfCheckCause::Upgrade { id, .. } => {
2019					weight += T::DbWeight::get().writes(2);
2020					UpgradeGoAheadSignal::<T>::insert(&id, UpgradeGoAhead::Abort);
2021					FutureCodeHash::<T>::remove(&id);
2022				},
2023			}
2024		}
2025
2026		weight
2027	}
2028
2029	/// Verify that `schedule_para_initialize` can be called successfully.
2030	///
2031	/// Returns false if para is already registered in the system.
2032	pub fn can_schedule_para_initialize(id: &ParaId) -> bool {
2033		ParaLifecycles::<T>::get(id).is_none()
2034	}
2035
2036	/// Schedule a para to be initialized. If the validation code is not already stored in the
2037	/// code storage, then a PVF pre-checking process will be initiated.
2038	///
2039	/// Only after the PVF pre-checking succeeds can the para be onboarded. Note, that calling this
2040	/// does not guarantee that the teyrchain will eventually be onboarded. This can happen in case
2041	/// the PVF does not pass PVF pre-checking.
2042	///
2043	/// The Para ID should be not activated in this pezpallet. The validation code supplied in
2044	/// `genesis_data` should not be empty. If those conditions are not met, then the para cannot
2045	/// be onboarded.
2046	pub(crate) fn schedule_para_initialize(
2047		id: ParaId,
2048		mut genesis_data: ParaGenesisArgs,
2049	) -> DispatchResult {
2050		// Make sure teyrchain isn't already in our system and that the onboarding parameters are
2051		// valid.
2052		ensure!(Self::can_schedule_para_initialize(&id), Error::<T>::CannotOnboard);
2053		ensure!(!genesis_data.validation_code.0.is_empty(), Error::<T>::CannotOnboard);
2054		ParaLifecycles::<T>::insert(&id, ParaLifecycle::Onboarding);
2055
2056		// HACK: here we are doing something nasty.
2057		//
2058		// In order to fix the [soaking issue] we insert the code eagerly here. When the onboarding
2059		// is finally enacted, we do not need to insert the code anymore. Therefore, there is no
2060		// reason for the validation code to be copied into the `ParaGenesisArgs`. We also do not
2061		// want to risk it by copying the validation code needlessly to not risk adding more
2062		// memory pressure.
2063		//
2064		// That said, we also want to preserve `ParaGenesisArgs` as it is, for now. There are two
2065		// reasons:
2066		//
2067		// - Doing it within the context of the PR that introduces this change is undesirable, since
2068		//   it is already a big change, and that change would require a migration. Moreover, if we
2069		//   run the new version of the runtime, there will be less things to worry about during the
2070		//   eventual proper migration.
2071		//
2072		// - This data type already is used for generating genesis, and changing it will probably
2073		//   introduce some unnecessary burden.
2074		//
2075		// So instead of going through it right now, we will do something sneaky. Specifically:
2076		//
2077		// - Insert the `CurrentCodeHash` now, instead during the onboarding. That would allow to
2078		//   get rid of hashing of the validation code when onboarding.
2079		//
2080		// - Replace `validation_code` with a sentinel value: an empty vector. This should be fine
2081		//   as long we do not allow registering teyrchains with empty code. At the moment of
2082		//   writing this should already be the case.
2083		//
2084		// - Empty value is treated as the current code is already inserted during the onboarding.
2085		//
2086		// This is only an intermediate solution and should be fixed in foreseeable future.
2087		//
2088		// [soaking issue]: https://github.com/pezkuwichain/pezkuwi-sdk/issues/146
2089		let validation_code =
2090			mem::replace(&mut genesis_data.validation_code, ValidationCode(Vec::new()));
2091		UpcomingParasGenesis::<T>::insert(&id, genesis_data);
2092		let validation_code_hash = validation_code.hash();
2093		CurrentCodeHash::<T>::insert(&id, validation_code_hash);
2094
2095		let cfg = configuration::ActiveConfig::<T>::get();
2096		Self::kick_off_pvf_check(
2097			PvfCheckCause::Onboarding(id),
2098			validation_code_hash,
2099			validation_code,
2100			&cfg,
2101		);
2102
2103		Ok(())
2104	}
2105
2106	/// Schedule a para to be cleaned up at the start of the next session.
2107	///
2108	/// Will return error if either is true:
2109	///
2110	/// - para is not a stable teyrchain (i.e. [`ParaLifecycle::is_stable`] is `false`)
2111	/// - para has a pending upgrade.
2112	/// - para has unprocessed messages in its UMP queue.
2113	///
2114	/// No-op if para is not registered at all.
2115	pub(crate) fn schedule_para_cleanup(id: ParaId) -> DispatchResult {
2116		// Disallow offboarding in case there is a PVF pre-checking in progress.
2117		//
2118		// This is not a fundamental limitation but rather simplification: it allows us to get
2119		// away without introducing additional logic for pruning and, more importantly, enacting
2120		// ongoing PVF pre-checking votes. It also removes some nasty edge cases.
2121		//
2122		// However, an upcoming upgrade on its own imposes no restrictions. An upgrade is enacted
2123		// with a new para head, so if a para never progresses we still should be able to offboard
2124		// it.
2125		//
2126		// This implicitly assumes that the given para exists, i.e. it's lifecycle != None.
2127		if let Some(future_code_hash) = FutureCodeHash::<T>::get(&id) {
2128			let active_prechecking = PvfActiveVoteList::<T>::get();
2129			if active_prechecking.contains(&future_code_hash) {
2130				return Err(Error::<T>::CannotOffboard.into());
2131			}
2132		}
2133
2134		let lifecycle = ParaLifecycles::<T>::get(&id);
2135		match lifecycle {
2136			// If para is not registered, nothing to do!
2137			None => return Ok(()),
2138			Some(ParaLifecycle::Parathread) => {
2139				ParaLifecycles::<T>::insert(&id, ParaLifecycle::OffboardingParathread);
2140			},
2141			Some(ParaLifecycle::Teyrchain) => {
2142				ParaLifecycles::<T>::insert(&id, ParaLifecycle::OffboardingTeyrchain);
2143			},
2144			_ => return Err(Error::<T>::CannotOffboard.into()),
2145		}
2146
2147		let scheduled_session = Self::scheduled_session();
2148		ActionsQueue::<T>::mutate(scheduled_session, |v| {
2149			if let Err(i) = v.binary_search(&id) {
2150				v.insert(i, id);
2151			}
2152		});
2153
2154		if <T as Config>::QueueFootprinter::message_count(UmpQueueId::Para(id)) != 0 {
2155			return Err(Error::<T>::CannotOffboard.into());
2156		}
2157
2158		Ok(())
2159	}
2160
2161	/// Schedule a parathread (on-demand teyrchain) to be upgraded to a lease holding teyrchain.
2162	///
2163	/// Will return error if `ParaLifecycle` is not `Parathread`.
2164	pub(crate) fn schedule_parathread_upgrade(id: ParaId) -> DispatchResult {
2165		let scheduled_session = Self::scheduled_session();
2166		let lifecycle = ParaLifecycles::<T>::get(&id).ok_or(Error::<T>::NotRegistered)?;
2167
2168		ensure!(lifecycle == ParaLifecycle::Parathread, Error::<T>::CannotUpgrade);
2169
2170		ParaLifecycles::<T>::insert(&id, ParaLifecycle::UpgradingParathread);
2171		ActionsQueue::<T>::mutate(scheduled_session, |v| {
2172			if let Err(i) = v.binary_search(&id) {
2173				v.insert(i, id);
2174			}
2175		});
2176
2177		Ok(())
2178	}
2179
2180	/// Schedule a lease holding teyrchain to be downgraded to an on-demand teyrchain.
2181	///
2182	/// Noop if `ParaLifecycle` is not `Teyrchain`.
2183	pub(crate) fn schedule_teyrchain_downgrade(id: ParaId) -> DispatchResult {
2184		let scheduled_session = Self::scheduled_session();
2185		let lifecycle = ParaLifecycles::<T>::get(&id).ok_or(Error::<T>::NotRegistered)?;
2186
2187		ensure!(lifecycle == ParaLifecycle::Teyrchain, Error::<T>::CannotDowngrade);
2188
2189		ParaLifecycles::<T>::insert(&id, ParaLifecycle::DowngradingTeyrchain);
2190		ActionsQueue::<T>::mutate(scheduled_session, |v| {
2191			if let Err(i) = v.binary_search(&id) {
2192				v.insert(i, id);
2193			}
2194		});
2195
2196		Ok(())
2197	}
2198
2199	/// Schedule a future code upgrade of the given teyrchain.
2200	///
2201	/// If the new code is not known, then the PVF pre-checking will be started for that validation
2202	/// code. In case the validation code does not pass the PVF pre-checking process, the
2203	/// upgrade will be aborted.
2204	///
2205	/// Only after the code is approved by the process, the upgrade can be scheduled. Specifically,
2206	/// the relay-chain block number will be determined at which the upgrade will take place. We
2207	/// call that block `expected_at`.
2208	///
2209	/// Once the candidate with the relay-parent >= `expected_at` is enacted, the new validation
2210	/// code will be applied. Therefore, the new code will be used to validate the next candidate.
2211	///
2212	/// The new code should not be equal to the current one, otherwise the upgrade will be aborted.
2213	/// If there is already a scheduled code upgrade for the para, this is a no-op.
2214	///
2215	/// Inclusion block number specifies relay parent which enacted candidate initiating the
2216	/// upgrade.
2217	pub(crate) fn schedule_code_upgrade(
2218		id: ParaId,
2219		new_code: ValidationCode,
2220		inclusion_block_number: BlockNumberFor<T>,
2221		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
2222		upgrade_strategy: UpgradeStrategy,
2223	) {
2224		// Should be prevented by checks in `schedule_code_upgrade_external`
2225		let new_code_len = new_code.0.len();
2226		if new_code_len < MIN_CODE_SIZE as usize || new_code_len > cfg.max_code_size as usize {
2227			log::warn!(target: LOG_TARGET, "attempted to schedule an upgrade with invalid new validation code",);
2228			return;
2229		}
2230
2231		// Enacting this should be prevented by the `can_upgrade_validation_code`
2232		if FutureCodeHash::<T>::contains_key(&id) {
2233			// This branch should never be reached. Signalling an upgrade is disallowed for a para
2234			// that already has one upgrade scheduled.
2235			//
2236			// Any candidate that attempts to do that should be rejected by
2237			// `can_upgrade_validation_code`.
2238			//
2239			// NOTE: we cannot set `UpgradeGoAheadSignal` signal here since this will be reset by
2240			//       the following call `note_new_head`
2241			log::warn!(target: LOG_TARGET, "ended up scheduling an upgrade while one is pending",);
2242			return;
2243		}
2244
2245		let code_hash = new_code.hash();
2246
2247		// para signals an update to the same code? This does not make a lot of sense, so abort the
2248		// process right away.
2249		//
2250		// We do not want to allow this since it will mess with the code reference counting.
2251		if CurrentCodeHash::<T>::get(&id) == Some(code_hash) {
2252			// NOTE: we cannot set `UpgradeGoAheadSignal` signal here since this will be reset by
2253			//       the following call `note_new_head`
2254			log::warn!(
2255				target: LOG_TARGET,
2256				"para tried to upgrade to the same code. Abort the upgrade",
2257			);
2258			return;
2259		}
2260
2261		// This is the start of the upgrade process. Prevent any further attempts at upgrading.
2262		FutureCodeHash::<T>::insert(&id, &code_hash);
2263		UpgradeRestrictionSignal::<T>::insert(&id, UpgradeRestriction::Present);
2264
2265		let next_possible_upgrade_at = inclusion_block_number + cfg.validation_upgrade_cooldown;
2266		UpgradeCooldowns::<T>::mutate(|upgrade_cooldowns| {
2267			let insert_idx = upgrade_cooldowns
2268				.binary_search_by_key(&next_possible_upgrade_at, |&(_, b)| b)
2269				.unwrap_or_else(|idx| idx);
2270			upgrade_cooldowns.insert(insert_idx, (id, next_possible_upgrade_at));
2271		});
2272
2273		Self::kick_off_pvf_check(
2274			PvfCheckCause::Upgrade { id, included_at: inclusion_block_number, upgrade_strategy },
2275			code_hash,
2276			new_code,
2277			cfg,
2278		);
2279	}
2280
2281	/// Makes sure that the given code hash has passed pre-checking.
2282	///
2283	/// If the given code hash has already passed pre-checking, then the approval happens
2284	/// immediately.
2285	///
2286	/// If the code is unknown, but the pre-checking for that PVF is already running then we perform
2287	/// "coalescing". We save the cause for this PVF pre-check request and just add it to the
2288	/// existing active PVF vote.
2289	///
2290	/// And finally, if the code is unknown and pre-checking is not running, we start the
2291	/// pre-checking process anew.
2292	///
2293	/// Unconditionally increases the reference count for the passed `code`.
2294	fn kick_off_pvf_check(
2295		cause: PvfCheckCause<BlockNumberFor<T>>,
2296		code_hash: ValidationCodeHash,
2297		code: ValidationCode,
2298		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
2299	) -> Weight {
2300		let mut weight = Weight::zero();
2301
2302		weight += T::DbWeight::get().reads_writes(3, 2);
2303		Self::deposit_event(Event::PvfCheckStarted(code_hash, cause.para_id()));
2304
2305		weight += T::DbWeight::get().reads(1);
2306		match PvfActiveVoteMap::<T>::get(&code_hash) {
2307			None => {
2308				// We deliberately are using `CodeByHash` here instead of the `CodeByHashRefs`. This
2309				// is because the code may have been added by `add_trusted_validation_code`.
2310				let known_code = CodeByHash::<T>::contains_key(&code_hash);
2311				weight += T::DbWeight::get().reads(1);
2312
2313				if known_code {
2314					// The code is known and there is no active PVF vote for it meaning it is
2315					// already checked -- fast track the PVF checking into the accepted state.
2316					weight += T::DbWeight::get().reads(1);
2317					let now = pezframe_system::Pezpallet::<T>::block_number();
2318					weight += Self::enact_pvf_accepted(now, &code_hash, &[cause], 0, cfg);
2319				} else {
2320					// PVF is not being pre-checked and it is not known. Start a new pre-checking
2321					// process.
2322					weight += T::DbWeight::get().reads_writes(3, 2);
2323					let now = pezframe_system::Pezpallet::<T>::block_number();
2324					let n_validators = shared::ActiveValidatorKeys::<T>::get().len();
2325					PvfActiveVoteMap::<T>::insert(
2326						&code_hash,
2327						PvfCheckActiveVoteState::new(now, n_validators, cause),
2328					);
2329					PvfActiveVoteList::<T>::mutate(|l| {
2330						if let Err(idx) = l.binary_search(&code_hash) {
2331							l.insert(idx, code_hash);
2332						}
2333					});
2334				}
2335			},
2336			Some(mut vote_state) => {
2337				// Coalescing: the PVF is already being pre-checked so we just need to piggy back
2338				// on it.
2339				weight += T::DbWeight::get().writes(1);
2340				vote_state.causes.push(cause);
2341				PvfActiveVoteMap::<T>::insert(&code_hash, vote_state);
2342			},
2343		}
2344
2345		// We increase the code RC here in any case. Intuitively the teyrchain that requested this
2346		// action is now a user of that PVF.
2347		//
2348		// If the result of the pre-checking is reject, then we would decrease the RC for each
2349		// cause, including the current.
2350		//
2351		// If the result of the pre-checking is accept, then we do nothing to the RC because the PVF
2352		// will continue be used by the same users.
2353		//
2354		// If the PVF was fast-tracked (i.e. there is already non zero RC) and there is no
2355		// pre-checking, we also do not change the RC then.
2356		weight += Self::increase_code_ref(&code_hash, &code);
2357
2358		weight
2359	}
2360
2361	/// Note that a para has progressed to a new head, where the new head was executed in the
2362	/// context of a relay-chain block with given number. This will apply pending code upgrades
2363	/// based on the relay-parent block number provided.
2364	pub(crate) fn note_new_head(
2365		id: ParaId,
2366		new_head: HeadData,
2367		execution_context: BlockNumberFor<T>,
2368	) {
2369		Heads::<T>::insert(&id, &new_head);
2370		MostRecentContext::<T>::insert(&id, execution_context);
2371
2372		if let Some(expected_at) = FutureCodeUpgrades::<T>::get(&id) {
2373			if expected_at <= execution_context {
2374				FutureCodeUpgrades::<T>::remove(&id);
2375				UpgradeGoAheadSignal::<T>::remove(&id);
2376
2377				// Both should always be `Some` in this case, since a code upgrade is scheduled.
2378				let new_code_hash = if let Some(new_code_hash) = FutureCodeHash::<T>::take(&id) {
2379					new_code_hash
2380				} else {
2381					log::error!(target: LOG_TARGET, "Missing future code hash for {:?}", &id);
2382					return;
2383				};
2384
2385				Self::set_current_code(id, new_code_hash, expected_at);
2386			}
2387		} else {
2388			// This means there is no upgrade scheduled.
2389			//
2390			// In case the upgrade was aborted by the relay-chain we should reset
2391			// the `Abort` signal.
2392			UpgradeGoAheadSignal::<T>::remove(&id);
2393		};
2394
2395		T::OnNewHead::on_new_head(id, &new_head);
2396	}
2397
2398	/// Set the current code for the given teyrchain.
2399	// `at` for para-triggered replacement is the block number of the relay-chain
2400	// block in whose context the parablock was executed
2401	// (i.e. number of `relay_parent` in the receipt)
2402	pub(crate) fn set_current_code(
2403		id: ParaId,
2404		new_code_hash: ValidationCodeHash,
2405		at: BlockNumberFor<T>,
2406	) -> Weight {
2407		let maybe_prior_code_hash = CurrentCodeHash::<T>::get(&id);
2408		CurrentCodeHash::<T>::insert(&id, &new_code_hash);
2409
2410		let log = ConsensusLog::ParaUpgradeCode(id, new_code_hash);
2411		<pezframe_system::Pezpallet<T>>::deposit_log(log.into());
2412
2413		// `now` is only used for registering pruning as part of `fn note_past_code`
2414		let now = <pezframe_system::Pezpallet<T>>::block_number();
2415
2416		let weight = if let Some(prior_code_hash) = maybe_prior_code_hash {
2417			Self::note_past_code(id, at, now, prior_code_hash)
2418		} else {
2419			log::error!(target: LOG_TARGET, "Missing prior code hash for para {:?}", &id);
2420			Weight::zero()
2421		};
2422
2423		weight + T::DbWeight::get().writes(1)
2424	}
2425
2426	/// Force set the current code for the given teyrchain.
2427	fn do_force_set_current_code_update(para: ParaId, new_code: ValidationCode) {
2428		let new_code_hash = new_code.hash();
2429		Self::increase_code_ref(&new_code_hash, &new_code);
2430		Self::set_current_code(
2431			para,
2432			new_code_hash,
2433			pezframe_system::Pezpallet::<T>::block_number(),
2434		);
2435		Self::deposit_event(Event::CurrentCodeUpdated(para));
2436	}
2437
2438	/// Returns the list of PVFs (aka validation code) that require casting a vote by a validator in
2439	/// the active validator set.
2440	pub(crate) fn pvfs_require_precheck() -> Vec<ValidationCodeHash> {
2441		PvfActiveVoteList::<T>::get()
2442	}
2443
2444	/// Submits a given PVF check statement with corresponding signature as an unsigned transaction
2445	/// into the memory pool. Ultimately, that disseminates the transaction across the network.
2446	///
2447	/// This function expects an offchain context and cannot be callable from the on-chain logic.
2448	///
2449	/// The signature assumed to pertain to `stmt`.
2450	pub(crate) fn submit_pvf_check_statement(
2451		stmt: PvfCheckStatement,
2452		signature: ValidatorSignature,
2453	) {
2454		use pezframe_system::offchain::SubmitTransaction;
2455
2456		let xt = T::create_bare(Call::include_pvf_check_statement { stmt, signature }.into());
2457		if let Err(e) = SubmitTransaction::<T, Call<T>>::submit_transaction(xt) {
2458			log::error!(target: LOG_TARGET, "Error submitting pvf check statement: {:?}", e,);
2459		}
2460	}
2461
2462	/// Returns the current lifecycle state of the para.
2463	pub fn lifecycle(id: ParaId) -> Option<ParaLifecycle> {
2464		ParaLifecycles::<T>::get(&id)
2465	}
2466
2467	/// Returns whether the given ID refers to a valid para.
2468	///
2469	/// Paras that are onboarding or offboarding are not included.
2470	pub fn is_valid_para(id: ParaId) -> bool {
2471		if let Some(state) = ParaLifecycles::<T>::get(&id) {
2472			!state.is_onboarding() && !state.is_offboarding()
2473		} else {
2474			false
2475		}
2476	}
2477
2478	/// Returns whether the given ID refers to a para that is offboarding.
2479	///
2480	/// An invalid or non-offboarding para ID will return `false`.
2481	pub fn is_offboarding(id: ParaId) -> bool {
2482		ParaLifecycles::<T>::get(&id).map_or(false, |state| state.is_offboarding())
2483	}
2484
2485	/// Whether a para ID corresponds to any live lease holding teyrchain.
2486	///
2487	/// Includes lease holding teyrchains which will downgrade to a on-demand teyrchains in the
2488	/// future.
2489	pub fn is_teyrchain(id: ParaId) -> bool {
2490		if let Some(state) = ParaLifecycles::<T>::get(&id) {
2491			state.is_teyrchain()
2492		} else {
2493			false
2494		}
2495	}
2496
2497	/// Whether a para ID corresponds to any live parathread (on-demand teyrchain).
2498	///
2499	/// Includes on-demand teyrchains which will upgrade to lease holding teyrchains in the future.
2500	pub fn is_parathread(id: ParaId) -> bool {
2501		if let Some(state) = ParaLifecycles::<T>::get(&id) {
2502			state.is_parathread()
2503		} else {
2504			false
2505		}
2506	}
2507
2508	/// If a candidate from the specified teyrchain were submitted at the current block, this
2509	/// function returns if that candidate passes the acceptance criteria.
2510	pub(crate) fn can_upgrade_validation_code(id: ParaId) -> bool {
2511		FutureCodeHash::<T>::get(&id).is_none() && UpgradeRestrictionSignal::<T>::get(&id).is_none()
2512	}
2513
2514	/// Return the session index that should be used for any future scheduled changes.
2515	fn scheduled_session() -> SessionIndex {
2516		shared::Pezpallet::<T>::scheduled_session()
2517	}
2518
2519	/// Store the validation code if not already stored, and increase the number of reference.
2520	///
2521	/// Returns the weight consumed.
2522	fn increase_code_ref(code_hash: &ValidationCodeHash, code: &ValidationCode) -> Weight {
2523		let mut weight = T::DbWeight::get().reads_writes(1, 1);
2524		CodeByHashRefs::<T>::mutate(code_hash, |refs| {
2525			if *refs == 0 {
2526				weight += T::DbWeight::get().writes(1);
2527				CodeByHash::<T>::insert(code_hash, code);
2528			}
2529			*refs += 1;
2530		});
2531		weight
2532	}
2533
2534	/// Decrease the number of reference of the validation code and remove it from storage if zero
2535	/// is reached.
2536	///
2537	/// Returns the weight consumed.
2538	fn decrease_code_ref(code_hash: &ValidationCodeHash) -> Weight {
2539		let mut weight = T::DbWeight::get().reads(1);
2540		let refs = CodeByHashRefs::<T>::get(code_hash);
2541		if refs == 0 {
2542			log::error!(target: LOG_TARGET, "Code refs is already zero for {:?}", code_hash);
2543			return weight;
2544		}
2545		if refs <= 1 {
2546			weight += T::DbWeight::get().writes(2);
2547			CodeByHash::<T>::remove(code_hash);
2548			CodeByHashRefs::<T>::remove(code_hash);
2549		} else {
2550			weight += T::DbWeight::get().writes(1);
2551			CodeByHashRefs::<T>::insert(code_hash, refs - 1);
2552		}
2553		weight
2554	}
2555
2556	/// Test function for triggering a new session in this pezpallet.
2557	#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
2558	pub fn test_on_new_session() {
2559		Self::initializer_on_new_session(&SessionChangeNotification {
2560			session_index: shared::CurrentSessionIndex::<T>::get(),
2561			..Default::default()
2562		});
2563	}
2564
2565	#[cfg(any(feature = "runtime-benchmarks", test))]
2566	pub fn heads_insert(para_id: &ParaId, head_data: HeadData) {
2567		Heads::<T>::insert(para_id, head_data);
2568	}
2569
2570	/// A low-level function to eagerly initialize a given para.
2571	pub(crate) fn initialize_para_now(
2572		teyrchains: &mut TeyrchainsCache<T>,
2573		id: ParaId,
2574		genesis_data: &ParaGenesisArgs,
2575	) {
2576		match genesis_data.para_kind {
2577			ParaKind::Teyrchain => {
2578				teyrchains.add(id);
2579				ParaLifecycles::<T>::insert(&id, ParaLifecycle::Teyrchain);
2580			},
2581			ParaKind::Parathread => ParaLifecycles::<T>::insert(&id, ParaLifecycle::Parathread),
2582		}
2583
2584		// HACK: see the notice in `schedule_para_initialize`.
2585		//
2586		// Apparently, this is left over from a prior version of the runtime.
2587		// To handle this we just insert the code and link the current code hash
2588		// to it.
2589		if !genesis_data.validation_code.0.is_empty() {
2590			let code_hash = genesis_data.validation_code.hash();
2591			Self::increase_code_ref(&code_hash, &genesis_data.validation_code);
2592			CurrentCodeHash::<T>::insert(&id, code_hash);
2593		}
2594
2595		Heads::<T>::insert(&id, &genesis_data.genesis_head);
2596		MostRecentContext::<T>::insert(&id, BlockNumberFor::<T>::from(0u32));
2597	}
2598
2599	#[cfg(test)]
2600	pub(crate) fn active_vote_state(
2601		code_hash: &ValidationCodeHash,
2602	) -> Option<PvfCheckActiveVoteState<BlockNumberFor<T>>> {
2603		PvfActiveVoteMap::<T>::get(code_hash)
2604	}
2605
2606	/// This function checks whether the given `code.hash()` exists in the `AuthorizedCodeHash` map
2607	/// of authorized code hashes for a para. If found, it verifies that the associated code
2608	/// matches the provided `code`. If the validation is successful, it returns tuple as the
2609	/// authorized `ValidationCodeHash` with `expire_at`.
2610	pub(crate) fn validate_code_is_authorized(
2611		code: &ValidationCode,
2612		para: &ParaId,
2613	) -> Result<AuthorizedCodeHashAndExpiry<BlockNumberFor<T>>, Error<T>> {
2614		let authorized = AuthorizedCodeHash::<T>::get(para).ok_or(Error::<T>::NothingAuthorized)?;
2615		let now = pezframe_system::Pezpallet::<T>::block_number();
2616		ensure!(authorized.expire_at > now, Error::<T>::InvalidBlockNumber);
2617		ensure!(authorized.code_hash == code.hash(), Error::<T>::Unauthorized);
2618		Ok(authorized)
2619	}
2620}
2621
2622/// An overlay over the `Teyrchains` storage entry that provides a convenient interface for adding
2623/// or removing teyrchains in bulk.
2624pub(crate) struct TeyrchainsCache<T: Config> {
2625	// `None` here means the teyrchains list has not been accessed yet, nevermind modified.
2626	teyrchains: Option<BTreeSet<ParaId>>,
2627	_config: PhantomData<T>,
2628}
2629
2630impl<T: Config> TeyrchainsCache<T> {
2631	pub fn new() -> Self {
2632		Self { teyrchains: None, _config: PhantomData }
2633	}
2634
2635	fn ensure_initialized(&mut self) -> &mut BTreeSet<ParaId> {
2636		self.teyrchains
2637			.get_or_insert_with(|| Teyrchains::<T>::get().into_iter().collect())
2638	}
2639
2640	/// Adds the given para id to the list.
2641	pub fn add(&mut self, id: ParaId) {
2642		let teyrchains = self.ensure_initialized();
2643		teyrchains.insert(id);
2644	}
2645
2646	/// Removes the given para id from the list of teyrchains. Does nothing if the id is not in the
2647	/// list.
2648	pub fn remove(&mut self, id: ParaId) {
2649		let teyrchains = self.ensure_initialized();
2650		teyrchains.remove(&id);
2651	}
2652}
2653
2654impl<T: Config> Drop for TeyrchainsCache<T> {
2655	fn drop(&mut self) {
2656		if let Some(teyrchains) = self.teyrchains.take() {
2657			Teyrchains::<T>::put(teyrchains.into_iter().collect::<Vec<ParaId>>());
2658		}
2659	}
2660}