Add unstake all call validation in signed ext

Author: gztensor

pallets/subtensor/src/lib.rs

@@ -2165,6 +2165,20 @@ where
                     Self::get_priority_staking(who, hotkey, *amount_unstaked),
                 )
             }
+            Some(Call::unstake_all { hotkey }) => {
+                // Fully validate the user input
+                Self::result_to_validity(
+                    Pallet::<T>::validate_unstake_all(who, hotkey, false),
+                    Self::get_priority_vanilla(),
+                )
+            }
+            Some(Call::unstake_all_alpha { hotkey }) => {
+                // Fully validate the user input
+                Self::result_to_validity(
+                    Pallet::<T>::validate_unstake_all(who, hotkey, true),
+                    Self::get_priority_vanilla(),
+                )
+            }
             Some(Call::remove_stake_limit {
                 hotkey,
                 netuid,

pallets/subtensor/src/staking/stake_utils.rs

@@ -995,6 +995,47 @@ impl<T: Config> Pallet<T> {
         Ok(())
     }
 
+    /// Validate if unstake_all can be executed
+    ///
+    pub fn validate_unstake_all(
+        coldkey: &T::AccountId,
+        hotkey: &T::AccountId,
+        only_alpha: bool,
+    ) -> Result<(), Error<T>> {
+        // Get all netuids (filter out root)
+        let subnets: Vec<u16> = Self::get_all_subnet_netuids();
+
+        // Ensure that the hotkey account exists this is only possible through registration.
+        ensure!(
+            Self::hotkey_account_exists(hotkey),
+            Error::<T>::HotKeyAccountNotExists
+        );
+
+        let mut unstaking_any = false;
+        for netuid in subnets.iter() {
+            if !SubtokenEnabled::<T>::get(netuid) {
+                continue;
+            }
+
+            if only_alpha && (*netuid == Self::get_root_netuid()) {
+                continue;
+            }
+
+            // Get user's stake in this subnet
+            let alpha = Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, *netuid);
+
+            if Self::validate_remove_stake(coldkey, hotkey, *netuid, alpha, alpha, false).is_ok()
+            {
+                unstaking_any = true;
+            }
+        }
+
+        // If no unstaking happens, return error
+        ensure!(unstaking_any, Error::<T>::AmountTooLow);
+
+        Ok(())
+    }
+
     /// Validate stake transition user input
     /// That works for move_stake, transfer_stake, and swap_stake
     ///

pallets/subtensor/src/tests/staking.rs

@@ -3841,6 +3841,57 @@ fn test_unstake_low_liquidity_validate() {
     });
 }
 
+#[test]
+fn test_unstake_all_validate() {
+    // Testing the signed extension validate function
+    // correctly filters the `unstake_all` transaction.
+
+    new_test_ext(0).execute_with(|| {
+        let subnet_owner_coldkey = U256::from(1001);
+        let subnet_owner_hotkey = U256::from(1002);
+        let hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let amount_staked = DefaultMinStake::<Test>::get() * 10 + DefaultStakingFee::<Test>::get();
+
+        let netuid = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+        SubtensorModule::create_account_if_non_existent(&coldkey, &hotkey);
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey, amount_staked);
+
+        // Simulate stake for hotkey
+        SubnetTAO::<Test>::insert(netuid, u64::MAX / 1000);
+        SubnetAlphaIn::<Test>::insert(netuid, u64::MAX / 1000);
+        SubtensorModule::stake_into_subnet(&hotkey, &coldkey, netuid, amount_staked, 0);
+
+        // Set the liquidity at lowest possible value so that all staking requests fail
+        SubnetTAO::<Test>::insert(
+            netuid,
+            DefaultMinimumPoolLiquidity::<Test>::get().to_num::<u64>(),
+        );
+        SubnetAlphaIn::<Test>::insert(
+            netuid,
+            DefaultMinimumPoolLiquidity::<Test>::get().to_num::<u64>(),
+        );
+
+        // unstake_all call
+        let call = RuntimeCall::SubtensorModule(SubtensorCall::unstake_all { hotkey });
+
+        let info: DispatchInfo =
+            DispatchInfoOf::<<Test as frame_system::Config>::RuntimeCall>::default();
+
+        let extension = SubtensorSignedExtension::<Test>::new();
+        // Submit to the signed extension validate function
+        let result_no_stake = extension.validate(&coldkey, &call.clone(), &info, 10);
+
+        // Should fail due to insufficient stake
+        assert_err!(
+            result_no_stake,
+            TransactionValidityError::Invalid(InvalidTransaction::Custom(
+                CustomTransactionError::StakeAmountTooLow.into()
+            ))
+        );
+    });
+}
+
 #[test]
 fn test_max_amount_add_root() {
     new_test_ext(0).execute_with(|| {