@@ -392,6 +392,8 @@ pub trait WeightInfo {
392392 fn force_schedule_code_upgrade ( c : u32 ) -> Weight ;
393393 fn force_note_new_head ( s : u32 ) -> Weight ;
394394 fn force_queue_action ( ) -> Weight ;
395+ fn add_trusted_validation_code ( c : u32 ) -> Weight ;
396+ fn poke_unused_validation_code ( ) -> Weight ;
395397}
396398
397399pub struct TestWeightInfo ;
@@ -411,6 +413,12 @@ impl WeightInfo for TestWeightInfo {
411413 fn force_queue_action ( ) -> Weight {
412414 Weight :: MAX
413415 }
416+ fn add_trusted_validation_code ( _c : u32 ) -> Weight {
417+ Weight :: MAX
418+ }
419+ fn poke_unused_validation_code ( ) -> Weight {
420+ Weight :: MAX
421+ }
414422}
415423
416424#[ frame_support:: pallet]
@@ -763,6 +771,79 @@ pub mod pallet {
763771 Ok ( ( ) )
764772 }
765773
774+ /// Adds the validation code to the storage.
775+ ///
776+ /// The code will not be added if it is already present. Additionally, if PVF pre-checking
777+ /// is running for that code, it will be instantly accepted.
778+ ///
779+ /// Otherwise, the code will be added into the storage. Note that the code will be added
780+ /// into storage with reference count 0. This is to account the fact that there are no users
781+ /// for this code yet. The caller will have to make sure that this code eventually gets
782+ /// used by some parachain or removed from the storage to avoid storage leaks. For the latter
783+ /// prefer to use the `poke_unused_validation_code` dispatchable to raw storage manipulation.
784+ ///
785+ /// This function is mainly meant to be used for upgrading parachains that do not follow
786+ /// the go-ahead signal while the PVF pre-checking feature is enabled.
787+ #[ pallet:: weight( <T as Config >:: WeightInfo :: add_trusted_validation_code( validation_code. 0 . len( ) as u32 ) ) ]
788+ pub fn add_trusted_validation_code (
789+ origin : OriginFor < T > ,
790+ validation_code : ValidationCode ,
791+ ) -> DispatchResult {
792+ ensure_root ( origin) ?;
793+ let code_hash = validation_code. hash ( ) ;
794+
795+ if let Some ( vote) = <Self as Store >:: PvfActiveVoteMap :: get ( & code_hash) {
796+ // Remove the existing vote.
797+ PvfActiveVoteMap :: < T > :: remove ( & code_hash) ;
798+ PvfActiveVoteList :: < T > :: mutate ( |l| {
799+ if let Ok ( i) = l. binary_search ( & code_hash) {
800+ l. remove ( i) ;
801+ }
802+ } ) ;
803+
804+ let cfg = configuration:: Pallet :: < T > :: config ( ) ;
805+ Self :: enact_pvf_accepted (
806+ <frame_system:: Pallet < T > >:: block_number ( ) ,
807+ & code_hash,
808+ & vote. causes ,
809+ vote. age ,
810+ & cfg,
811+ ) ;
812+ return Ok ( ( ) )
813+ }
814+
815+ if <Self as Store >:: CodeByHash :: contains_key ( & code_hash) {
816+ // There is no vote, but the code exists. Nothing to do here.
817+ return Ok ( ( ) )
818+ }
819+
820+ // At this point the code is unknown and there is no PVF pre-checking vote for it, so we
821+ // can just add the code into the storage.
822+ //
823+ // NOTE That we do not use `increase_code_ref` here, because the code is not yet used
824+ // by any parachain.
825+ <Self as Store >:: CodeByHash :: insert ( code_hash, & validation_code) ;
826+
827+ Ok ( ( ) )
828+ }
829+
830+ /// Remove the validation code from the storage iff the reference count is 0.
831+ ///
832+ /// This is better than removing the storage directly, because it will not remove the code
833+ /// that was suddenly got used by some parachain while this dispatchable was pending
834+ /// dispatching.
835+ #[ pallet:: weight( <T as Config >:: WeightInfo :: poke_unused_validation_code( ) ) ]
836+ pub fn poke_unused_validation_code (
837+ origin : OriginFor < T > ,
838+ validation_code_hash : ValidationCodeHash ,
839+ ) -> DispatchResult {
840+ ensure_root ( origin) ?;
841+ if <Self as Store >:: CodeByHashRefs :: get ( & validation_code_hash) == 0 {
842+ <Self as Store >:: CodeByHash :: remove ( & validation_code_hash) ;
843+ }
844+ Ok ( ( ) )
845+ }
846+
766847 /// Includes a statement for a PVF pre-checking vote. Potentially, finalizes the vote and
767848 /// enacts the results if that was the last vote before achieving the supermajority.
768849 #[ pallet:: weight( Weight :: MAX ) ]
@@ -1609,6 +1690,8 @@ impl<T: Config> Pallet<T> {
16091690 weight += T :: DbWeight :: get ( ) . reads ( 1 ) ;
16101691 match PvfActiveVoteMap :: < T > :: get ( & code_hash) {
16111692 None => {
1693+ // We deliberately are using `CodeByHash` here instead of the `CodeByHashRefs`. This
1694+ // is because the code may have been added by `add_trusted_validation_code`.
16121695 let known_code = CodeByHash :: < T > :: contains_key ( & code_hash) ;
16131696 weight += T :: DbWeight :: get ( ) . reads ( 1 ) ;
16141697
@@ -1812,7 +1895,10 @@ impl<T: Config> Pallet<T> {
18121895 fn decrease_code_ref ( code_hash : & ValidationCodeHash ) -> Weight {
18131896 let mut weight = T :: DbWeight :: get ( ) . reads ( 1 ) ;
18141897 let refs = <Self as Store >:: CodeByHashRefs :: get ( code_hash) ;
1815- debug_assert ! ( refs != 0 ) ;
1898+ if refs == 0 {
1899+ log:: error!( target: LOG_TARGET , "Code refs is already zero for {:?}" , code_hash) ;
1900+ return weight
1901+ }
18161902 if refs <= 1 {
18171903 weight += T :: DbWeight :: get ( ) . writes ( 2 ) ;
18181904 <Self as Store >:: CodeByHash :: remove ( code_hash) ;
@@ -1842,7 +1928,7 @@ impl<T: Config> Pallet<T> {
18421928#[ cfg( test) ]
18431929mod tests {
18441930 use super :: * ;
1845- use frame_support:: { assert_err, assert_ok} ;
1931+ use frame_support:: { assert_err, assert_ok, assert_storage_noop } ;
18461932 use keyring:: Sr25519Keyring ;
18471933 use primitives:: {
18481934 v0:: PARACHAIN_KEY_TYPE_ID ,
@@ -1855,7 +1941,10 @@ mod tests {
18551941
18561942 use crate :: {
18571943 configuration:: HostConfiguration ,
1858- mock:: { new_test_ext, Configuration , MockGenesisConfig , Paras , ParasShared , System , Test } ,
1944+ mock:: {
1945+ new_test_ext, Configuration , MockGenesisConfig , Origin , Paras , ParasShared , System ,
1946+ Test ,
1947+ } ,
18591948 } ;
18601949
18611950 static VALIDATORS : & [ Sr25519Keyring ] = & [
@@ -3066,6 +3155,186 @@ mod tests {
30663155 } ) ;
30673156 }
30683157
3158+ #[ test]
3159+ fn add_trusted_validation_code_inserts_with_no_users ( ) {
3160+ // This test is to ensure that trusted validation code is inserted into the storage
3161+ // with the reference count equal to 0.
3162+ let validation_code = ValidationCode ( vec ! [ 1 , 2 , 3 ] ) ;
3163+ new_test_ext ( Default :: default ( ) ) . execute_with ( || {
3164+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , validation_code. clone( ) ) ) ;
3165+ assert_eq ! ( <Paras as Store >:: CodeByHashRefs :: get( & validation_code. hash( ) ) , 0 , ) ;
3166+ } ) ;
3167+ }
3168+
3169+ #[ test]
3170+ fn add_trusted_validation_code_idempotent ( ) {
3171+ // This test makes sure that calling add_trusted_validation_code twice with the same
3172+ // parameters is a no-op.
3173+ let validation_code = ValidationCode ( vec ! [ 1 , 2 , 3 ] ) ;
3174+ new_test_ext ( Default :: default ( ) ) . execute_with ( || {
3175+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , validation_code. clone( ) ) ) ;
3176+ assert_storage_noop ! ( {
3177+ assert_ok!( Paras :: add_trusted_validation_code(
3178+ Origin :: root( ) ,
3179+ validation_code. clone( )
3180+ ) ) ;
3181+ } ) ;
3182+ } ) ;
3183+ }
3184+
3185+ #[ test]
3186+ fn poke_unused_validation_code_removes_code_cleanly ( ) {
3187+ // This test makes sure that calling poke_unused_validation_code with a code that is currently
3188+ // in the storage but has no users will remove it cleanly from the storage.
3189+ let validation_code = ValidationCode ( vec ! [ 1 , 2 , 3 ] ) ;
3190+ new_test_ext ( Default :: default ( ) ) . execute_with ( || {
3191+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , validation_code. clone( ) ) ) ;
3192+ assert_ok ! ( Paras :: poke_unused_validation_code( Origin :: root( ) , validation_code. hash( ) ) ) ;
3193+
3194+ assert_eq ! ( <Paras as Store >:: CodeByHashRefs :: get( & validation_code. hash( ) ) , 0 ) ;
3195+ assert ! ( !<Paras as Store >:: CodeByHash :: contains_key( & validation_code. hash( ) ) ) ;
3196+ } ) ;
3197+ }
3198+
3199+ #[ test]
3200+ fn poke_unused_validation_code_doesnt_remove_code_with_users ( ) {
3201+ let para_id = 100 . into ( ) ;
3202+ let validation_code = ValidationCode ( vec ! [ 1 , 2 , 3 ] ) ;
3203+ new_test_ext ( Default :: default ( ) ) . execute_with ( || {
3204+ // First we add the code to the storage.
3205+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , validation_code. clone( ) ) ) ;
3206+
3207+ // Then we add a user to the code, say by upgrading.
3208+ run_to_block ( 2 , None ) ;
3209+ Paras :: schedule_code_upgrade (
3210+ para_id,
3211+ validation_code. clone ( ) ,
3212+ 1 ,
3213+ & Configuration :: config ( ) ,
3214+ ) ;
3215+ Paras :: note_new_head ( para_id, HeadData :: default ( ) , 1 ) ;
3216+
3217+ // Finally we poke the code, which should not remove it from the storage.
3218+ assert_storage_noop ! ( {
3219+ assert_ok!( Paras :: poke_unused_validation_code(
3220+ Origin :: root( ) ,
3221+ validation_code. hash( )
3222+ ) ) ;
3223+ } ) ;
3224+ check_code_is_stored ( & validation_code) ;
3225+ } ) ;
3226+ }
3227+
3228+ #[ test]
3229+ fn increase_code_ref_doesnt_have_allergy_on_add_trusted_validation_code ( ) {
3230+ // Verify that accidential calling of increase_code_ref or decrease_code_ref does not lead
3231+ // to a disaster.
3232+ // NOTE that this test is extra paranoid, as it is not really possible to hit
3233+ // `decrease_code_ref` without calling `increase_code_ref` first.
3234+ let code = ValidationCode ( vec ! [ 1 , 2 , 3 ] ) ;
3235+
3236+ new_test_ext ( Default :: default ( ) ) . execute_with ( || {
3237+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , code. clone( ) ) ) ;
3238+ Paras :: increase_code_ref ( & code. hash ( ) , & code) ;
3239+ Paras :: increase_code_ref ( & code. hash ( ) , & code) ;
3240+ assert ! ( <Paras as Store >:: CodeByHash :: contains_key( code. hash( ) ) ) ;
3241+ assert_eq ! ( <Paras as Store >:: CodeByHashRefs :: get( code. hash( ) ) , 2 ) ;
3242+ } ) ;
3243+
3244+ new_test_ext ( Default :: default ( ) ) . execute_with ( || {
3245+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , code. clone( ) ) ) ;
3246+ Paras :: decrease_code_ref ( & code. hash ( ) ) ;
3247+ assert ! ( <Paras as Store >:: CodeByHash :: contains_key( code. hash( ) ) ) ;
3248+ assert_eq ! ( <Paras as Store >:: CodeByHashRefs :: get( code. hash( ) ) , 0 ) ;
3249+ } ) ;
3250+ }
3251+
3252+ #[ test]
3253+ fn add_trusted_validation_code_insta_approval ( ) {
3254+ // In particular, this tests that `kick_off_pvf_check` reacts to the `add_trusted_validation_code`
3255+ // and uses the `CodeByHash::contains_key` which is what `add_trusted_validation_code` uses.
3256+ let para_id = 100 . into ( ) ;
3257+ let validation_code = ValidationCode ( vec ! [ 1 , 2 , 3 ] ) ;
3258+ let validation_upgrade_delay = 25 ;
3259+ let minimum_validation_upgrade_delay = 2 ;
3260+ let genesis_config = MockGenesisConfig {
3261+ configuration : crate :: configuration:: GenesisConfig {
3262+ config : HostConfiguration {
3263+ pvf_checking_enabled : true ,
3264+ validation_upgrade_delay,
3265+ minimum_validation_upgrade_delay,
3266+ ..Default :: default ( )
3267+ } ,
3268+ ..Default :: default ( )
3269+ } ,
3270+ ..Default :: default ( )
3271+ } ;
3272+ new_test_ext ( genesis_config) . execute_with ( || {
3273+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , validation_code. clone( ) ) ) ;
3274+
3275+ // Then some parachain upgrades it's code with the relay-parent 1.
3276+ run_to_block ( 2 , None ) ;
3277+ Paras :: schedule_code_upgrade (
3278+ para_id,
3279+ validation_code. clone ( ) ,
3280+ 1 ,
3281+ & Configuration :: config ( ) ,
3282+ ) ;
3283+ Paras :: note_new_head ( para_id, HeadData :: default ( ) , 1 ) ;
3284+
3285+ // Verify that the code upgrade has `expected_at` set to `26`. This is the behavior
3286+ // equal to that of `pvf_checking_enabled: false`.
3287+ assert_eq ! (
3288+ <Paras as Store >:: FutureCodeUpgrades :: get( & para_id) ,
3289+ Some ( 1 + validation_upgrade_delay)
3290+ ) ;
3291+ } ) ;
3292+ }
3293+
3294+ #[ test]
3295+ fn add_trusted_validation_code_enacts_existing_pvf_vote ( ) {
3296+ // This test makes sure that calling `add_trusted_validation_code` with a code that is
3297+ // already going through PVF pre-checking voting will conclude the voting and enact the
3298+ // code upgrade.
3299+ let para_id = 100 . into ( ) ;
3300+ let validation_code = ValidationCode ( vec ! [ 1 , 2 , 3 ] ) ;
3301+ let validation_upgrade_delay = 25 ;
3302+ let minimum_validation_upgrade_delay = 2 ;
3303+ let genesis_config = MockGenesisConfig {
3304+ configuration : crate :: configuration:: GenesisConfig {
3305+ config : HostConfiguration {
3306+ pvf_checking_enabled : true ,
3307+ validation_upgrade_delay,
3308+ minimum_validation_upgrade_delay,
3309+ ..Default :: default ( )
3310+ } ,
3311+ ..Default :: default ( )
3312+ } ,
3313+ ..Default :: default ( )
3314+ } ;
3315+ new_test_ext ( genesis_config) . execute_with ( || {
3316+ // First, some parachain upgrades it's code with the relay-parent 1.
3317+ run_to_block ( 2 , None ) ;
3318+ Paras :: schedule_code_upgrade (
3319+ para_id,
3320+ validation_code. clone ( ) ,
3321+ 1 ,
3322+ & Configuration :: config ( ) ,
3323+ ) ;
3324+ Paras :: note_new_head ( para_id, HeadData :: default ( ) , 1 ) ;
3325+
3326+ // No upgrade should be scheduled at this point. PVF pre-checking vote should run for
3327+ // that PVF.
3328+ assert ! ( <Paras as Store >:: FutureCodeUpgrades :: get( & para_id) . is_none( ) ) ;
3329+ assert ! ( <Paras as Store >:: PvfActiveVoteMap :: contains_key( & validation_code. hash( ) ) ) ;
3330+
3331+ // Then we add a trusted validation code. That should conclude the vote.
3332+ assert_ok ! ( Paras :: add_trusted_validation_code( Origin :: root( ) , validation_code. clone( ) ) ) ;
3333+ assert ! ( <Paras as Store >:: FutureCodeUpgrades :: get( & para_id) . is_some( ) ) ;
3334+ assert ! ( !<Paras as Store >:: PvfActiveVoteMap :: contains_key( & validation_code. hash( ) ) ) ;
3335+ } ) ;
3336+ }
3337+
30693338 #[ test]
30703339 fn verify_upgrade_go_ahead_signal_is_externally_accessible ( ) {
30713340 use primitives:: v1:: well_known_keys;
0 commit comments