Fix set-controller and ref count update (#1688)

* Fix set-controller and ref count update * Improve check efficiency --------- Co-authored-by: Adam Dossa <[email protected]>
PolymeshAssociation · Jul 7, 2024 · 84065e9 · 84065e9
1 parent ca0fd5e
commit 84065e9
Show file tree

Hide file tree

Showing 7 changed files with 135 additions and 22 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pallets/identity/Cargo.toml b/pallets/identity/Cargo.toml
@@ -16,6 +16,7 @@ log = "0.4.8"
 serde = { version = "1.0.104", default-features = false }
 serde_derive = { version = "1.0.104", optional = true, default-features = false  }
 either = { version = "1.6.1", default-features = false }
+hex-literal = "0.3.0"
 
 # Cryptography
 schnorrkel = { version = "0.10.1", default-features = false, optional = true }

diff --git a/pallets/identity/src/lib.rs b/pallets/identity/src/lib.rs
@@ -116,7 +116,7 @@ use polymesh_primitives::{
 
 pub type Event<T> = polymesh_common_utilities::traits::identity::Event<T>;
 
-storage_migration_ver!(5);
+storage_migration_ver!(6);
 
 decl_storage! {
     trait Store for Module<T: Config> as Identity {
@@ -170,7 +170,7 @@ decl_storage! {
         pub CddAuthForPrimaryKeyRotation get(fn cdd_auth_for_primary_key_rotation): bool;
 
         /// Storage version.
-        StorageVersion get(fn storage_version) build(|_| Version::new(5)): Version;
+        StorageVersion get(fn storage_version) build(|_| Version::new(6)): Version;
 
         /// How many "strong" references to the account key.
         ///
@@ -264,8 +264,8 @@ decl_module! {
         fn deposit_event() = default;
 
         fn on_runtime_upgrade() -> Weight {
-            storage_migrate_on!(StorageVersion, 5, {
-                migration::migrate_to_v5::<T>();
+            storage_migrate_on!(StorageVersion, 6, {
+                migration::migrate_to_v6::<T>();
             });
             Weight::zero()
         }
@@ -784,25 +784,36 @@ fn revoke_claim_class(claim_type: ClaimType) -> frame_support::dispatch::Dispatc
 }
 
 pub mod migration {
-    use super::*;
+    use hex_literal::hex;
     use sp_runtime::runtime_logger::RuntimeLogger;
 
-    pub fn migrate_to_v5<T: Config>() {
+    use super::*;
+
+    pub fn migrate_to_v6<T: Config>() {
         RuntimeLogger::init();
-        log::info!(" >>> Initializing NextAuthId and NumberOfGivenAuths");
-        initialize_next_auth_id::<T>();
-        initialize_number_of_given_auths::<T>();
-        log::info!(" >>> NextAuthId and NumberOfGivenAuths have been initialized");
-    }
+        log::info!(" >>> Updating account ref");
 
-    fn initialize_next_auth_id<T: Config>() {
-        let next_auth_id = MultiPurposeNonce::get().saturating_add(1);
-        CurrentAuthId::put(next_auth_id);
-    }
+        // Hex for key 2EGKNqWLx2VhjgFZ6BwXZ9Tf6jQXzWjW4cNvE2Bd24z85xfq
+        if let Ok(key_1) = Decode::decode(
+            &mut hex!("50631df3b9b6a7a7d8893ae3959b3d9ca1f266ca05ea6e007cdaeb8adbae9805").as_ref(),
+        ) {
+            crate::Module::<T>::remove_account_key_ref_count(&key_1);
+        }
+
+        // Hex for key 2DR8JZWwJ7jDSHYDKoaV7Bfi6dyoQ6FkBx4KYB5owtjiFgj5
+        if let Ok(key_2) = Decode::decode(
+            &mut hex!("2aded091e73417dd618c66a10175df84cd76803ad9a87a62cb986d2e84e6d43c").as_ref(),
+        ) {
+            crate::Module::<T>::remove_account_key_ref_count(&key_2);
+        };
 
-    fn initialize_number_of_given_auths<T: Config>() {
-        for (authorizer, _) in AuthorizationsGiven::<T>::iter_keys() {
-            NumberOfGivenAuths::mutate(authorizer, |n| *n = n.saturating_add(1));
+        // Hex for key 2G3CJWkzusVqtxYoofCjGqwxgjFYJS6Rxvg8supREvQBmtvy
+        if let Ok(key_3) = Decode::decode(
+            &mut hex!("9ed994211b1b54a39a2ff9f4d94f6ca0a2da5411bfc266f8236425c2efd0876c").as_ref(),
+        ) {
+            crate::Module::<T>::remove_account_key_ref_count(&key_3);
         }
+
+        log::info!(" >>> Account ref has been updated");
     }
 }
diff --git a/pallets/runtime/tests/src/lib.rs b/pallets/runtime/tests/src/lib.rs
@@ -36,6 +36,7 @@ mod settlement_pallet;
 mod settlement_test;
 mod signed_extra;
 mod staking;
+mod staking_extra_tests;
 mod sto_test;
 mod transaction_payment_test;
 mod transfer_compliance_test;

diff --git a/pallets/runtime/tests/src/staking_extra_tests.rs b/pallets/runtime/tests/src/staking_extra_tests.rs
@@ -0,0 +1,96 @@
+use frame_support::{assert_ok, StorageMap};
+use sp_keyring::AccountKeyring;
+
+use polymesh_primitives::{AuthorizationData, Permissions};
+
+use crate::ext_builder::ExtBuilder;
+use crate::storage::{add_secondary_key, TestStorage, User};
+
+type Origin = <TestStorage as frame_system::Config>::RuntimeOrigin;
+
+#[test]
+fn updating_controller() {
+    let charlie = vec![AccountKeyring::Charlie.to_account_id()];
+    ExtBuilder::default()
+        .cdd_providers(charlie)
+        .build()
+        .execute_with(|| {
+            let alice: User = User::new(AccountKeyring::Alice);
+            let eve: User = User::new_with(alice.did, AccountKeyring::Eve);
+
+            add_secondary_key(alice.did, eve.acc());
+
+            let auth_id = pallet_identity::Module::<TestStorage>::add_auth(
+                alice.did,
+                eve.signatory_acc(),
+                AuthorizationData::RotatePrimaryKeyToSecondary(Permissions::default()),
+                None,
+            )
+            .unwrap();
+
+            assert_ok!(
+                pallet_staking::Pallet::<TestStorage>::add_permissioned_validator(
+                    Origin::root(),
+                    alice.did,
+                    None
+                )
+            );
+
+            assert_ok!(pallet_staking::Pallet::<TestStorage>::bond(
+                alice.origin(),
+                sp_runtime::MultiAddress::Id(alice.acc()),
+                10_000_000,
+                pallet_staking::RewardDestination::Controller
+            ));
+
+            assert_ok!(
+                pallet_staking::Pallet::<TestStorage>::set_min_bond_threshold(
+                    Origin::root(),
+                    50_000
+                )
+            );
+
+            assert_ok!(pallet_staking::Pallet::<TestStorage>::validate(
+                Origin::signed(alice.acc()),
+                pallet_staking::ValidatorPrefs::default()
+            ));
+
+            assert_ok!(pallet_identity::Module::<TestStorage>::revoke_claim(
+                Origin::signed(AccountKeyring::Charlie.to_account_id()),
+                alice.did,
+                polymesh_primitives::Claim::CustomerDueDiligence(Default::default())
+            ));
+
+            assert_ok!(
+                pallet_staking::Pallet::<TestStorage>::remove_permissioned_validator(
+                    Origin::root(),
+                    alice.did,
+                )
+            );
+
+            assert_eq!(
+                pallet_identity::AccountKeyRefCount::<TestStorage>::get(alice.acc()),
+                1
+            );
+
+            assert_ok!(pallet_staking::Pallet::<TestStorage>::chill(alice.origin(),));
+
+            assert_eq!(
+                pallet_identity::AccountKeyRefCount::<TestStorage>::get(alice.acc()),
+                0
+            );
+
+            assert_ok!(
+                pallet_identity::Module::<TestStorage>::rotate_primary_key_to_secondary(
+                    eve.origin(),
+                    auth_id,
+                    None
+                )
+            );
+
+            assert_ok!(pallet_staking::Pallet::<TestStorage>::set_controller(
+                alice.origin(),
+                sp_runtime::MultiAddress::Id(eve.acc()),
+            ));
+        });
+}
diff --git a/pallets/staking/src/pallet/impls.rs b/pallets/staking/src/pallet/impls.rs
@@ -702,10 +702,10 @@ impl<T: Config> Pallet<T> {
                 if let Some(p) = pref {
                     if p.running_count > 0 {
                         p.running_count -= 1;
-                        <Identity<T>>::remove_account_key_ref_count(&stash);
                     }
                 }
             });
+            <Identity<T>>::remove_account_key_ref_count(&stash);
         }
     }
 

diff --git a/pallets/staking/src/pallet/mod.rs b/pallets/staking/src/pallet/mod.rs
@@ -1428,9 +1428,12 @@ pub mod pallet {
             if <Ledger<T>>::contains_key(&controller) {
                 return Err(Error::<T>::AlreadyPaired.into());
             }
-            // Prevents stashes which are controllers from calling the extrinsic
-            if <Ledger<T>>::contains_key(&stash) {
-                return Err(Error::<T>::BadState.into());
+
+            // Prevents stashes which are controllers of another ledger from calling the extrinsic
+            if old_controller != stash {
+                if <Ledger<T>>::contains_key(&stash) {
+                    return Err(Error::<T>::BadState.into());
+                }
             }
 
             if controller != old_controller {