Merge pull request #1959 from opentensor/refactor-extension2

Refactor transaction extension

Author: sam0x17

pallets/subtensor/src/macros/dispatches.rs

@@ -2156,7 +2156,7 @@ mod dispatches {
         /// Emits a `SymbolUpdated` event on success.
         #[pallet::call_index(112)]
         #[pallet::weight((
-            Weight::from_parts(26_880_000, 0).saturating_add(T::DbWeight::get().reads_writes(4, 1)),
+            Weight::from_parts(26_200_000, 0).saturating_add(T::DbWeight::get().reads_writes(4, 1)),
             DispatchClass::Operational,
             Pays::Yes
         ))]

pallets/subtensor/src/tests/swap_coldkey.rs

@@ -2493,8 +2493,12 @@ fn test_coldkey_in_swap_schedule_prevents_funds_usage() {
             &TxBaseImplication(()),
             TransactionSource::External,
         );
-        // Should pass, not in list.
-        assert_ok!(result);
+        // Should fail
+        assert_eq!(
+            // Should get an invalid transaction error
+            result.unwrap_err(),
+            CustomTransactionError::ColdkeyInSwapSchedule.into()
+        );
 
         // Remove stake limit
         let call = RuntimeCall::SubtensorModule(SubtensorCall::remove_stake_limit {
@@ -2513,7 +2517,27 @@ fn test_coldkey_in_swap_schedule_prevents_funds_usage() {
             &TxBaseImplication(()),
             TransactionSource::External,
         );
-        // Should pass, not in list.
+        // Should fail
+        assert_eq!(
+            // Should get an invalid transaction error
+            result.unwrap_err(),
+            CustomTransactionError::ColdkeyInSwapSchedule.into()
+        );
+
+        // Schedule swap should succeed
+        let call = RuntimeCall::SubtensorModule(SubtensorCall::schedule_swap_coldkey {
+            new_coldkey: hotkey,
+        });
+        let result = extension.validate(
+            RawOrigin::Signed(who).into(),
+            &call.clone(),
+            &info,
+            10,
+            (),
+            &TxBaseImplication(()),
+            TransactionSource::External,
+        );
+        // Should be ok
         assert_ok!(result);
     });
 }

pallets/subtensor/src/transaction_extension.rs

@@ -120,6 +120,16 @@ where
             return Ok((Default::default(), None, origin));
         };
 
+        // Verify ColdkeySwapScheduled map for coldkey
+        match call.is_sub_type() {
+            // Whitelist
+            Some(Call::schedule_swap_coldkey { .. }) => {}
+            _ => {
+                if ColdkeySwapScheduled::<T>::contains_key(who) {
+                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
+                }
+            }
+        }
         match call.is_sub_type() {
             Some(Call::commit_weights { netuid, .. }) => {
                 if Self::check_weights_min_stake(who, *netuid) {
@@ -248,54 +258,7 @@ where
                     Err(CustomTransactionError::StakeAmountTooLow.into())
                 }
             }
-            Some(Call::add_stake { .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                }
-
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
-            Some(Call::add_stake_limit { .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                }
-
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
-            Some(Call::remove_stake { .. }) => Ok((Default::default(), Some(who.clone()), origin)),
-            Some(Call::remove_stake_limit { .. }) => {
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
-            Some(Call::move_stake { .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                }
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
-            Some(Call::transfer_stake { .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                }
-
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
-            Some(Call::swap_stake { .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                }
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
-            Some(Call::swap_stake_limit { .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                }
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
             Some(Call::register { netuid, .. } | Call::burned_register { netuid, .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                }
-
                 let registrations_this_interval =
                     Pallet::<T>::get_registrations_this_interval(*netuid);
                 let max_registrations_per_interval =
@@ -308,13 +271,6 @@ where
 
                 Ok((Default::default(), Some(who.clone()), origin))
             }
-            Some(Call::dissolve_network { .. }) => {
-                if ColdkeySwapScheduled::<T>::contains_key(who) {
-                    Err(CustomTransactionError::ColdkeyInSwapSchedule.into())
-                } else {
-                    Ok((Default::default(), Some(who.clone()), origin))
-                }
-            }
             Some(Call::serve_axon {
                 netuid,
                 version,
@@ -342,20 +298,7 @@ where
                 )
                 .map(|validity| (validity, Some(who.clone()), origin.clone()))
             }
-            _ => {
-                if let Some(
-                    BalancesCall::transfer_keep_alive { .. }
-                    | BalancesCall::transfer_all { .. }
-                    | BalancesCall::transfer_allow_death { .. },
-                ) = call.is_sub_type()
-                {
-                    if ColdkeySwapScheduled::<T>::contains_key(who) {
-                        return Err(CustomTransactionError::ColdkeyInSwapSchedule.into());
-                    }
-                }
-
-                Ok((Default::default(), Some(who.clone()), origin))
-            }
+            _ => Ok((Default::default(), Some(who.clone()), origin)),
         }
     }