adds a doublt Tempo limit to setting childkeys

Author: unconst

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -1,5 +1,4 @@
 use super::*;
-use frame_support::storage::IterableStorageDoubleMap;
 use substrate_fixed::types::I64F64;
 use substrate_fixed::types::I96F32;
 
@@ -262,7 +261,7 @@ impl<T: Config> Pallet<T> {
         PendingdHotkeyEmission::<T>::insert(hotkey, 0);
 
         // --- 2 Retrieve the last time this hotkey's emissions were drained.
-        let last_hotkey_emission_drain: u64 = LastHotkeyEmissionDrain::<T>::get(hotkey);
+        let last_emission_drain: u64 = LastHotkeyEmissionDrain::<T>::get(hotkey);
 
         // --- 3 Update the block value to the current block number.
         LastHotkeyEmissionDrain::<T>::insert(hotkey, block_number);
@@ -282,43 +281,48 @@ impl<T: Config> Pallet<T> {
         // --- 7 Calculate the remaining emission after the hotkey's take.
         let mut remainder: u64 = emission_minus_take;
 
-        // --- 8 Iterate over each nominator.
-        for (nominator, nominator_stake) in
-            <Stake<T> as IterableStorageDoubleMap<T::AccountId, T::AccountId, u64>>::iter_prefix(
-                hotkey,
-            )
-        {
-            // --- 9 Check if the stake was manually increased by the user since the last emission drain for this hotkey.
-            // If it was, skip this nominator as they will not receive their proportion of the emission.
-            if LastAddStakeIncrease::<T>::get(hotkey, nominator.clone())
-                > last_hotkey_emission_drain
-            {
-                continue;
+        // --- 8 Iterate over each nominator and get all viable stake.
+        let mut total_viable_nominator_stake: u64 = total_hotkey_stake;
+        for (nominator, nominator_stake) in Stake::<T>::iter_prefix(hotkey) {
+            if LastAddStakeIncrease::<T>::get(hotkey, nominator) > last_emission_drain {
+                total_viable_nominator_stake =
+                    total_viable_nominator_stake.saturating_sub(nominator_stake);
             }
+        }
 
-            // --- 10 Calculate this nominator's share of the emission.
-            let nominator_emission: I64F64 = I64F64::from_num(emission_minus_take)
-                .saturating_mul(I64F64::from_num(nominator_stake))
-                .checked_div(I64F64::from_num(total_hotkey_stake))
-                .unwrap_or(I64F64::from_num(0));
-
-            // --- 11 Increase the stake for the nominator.
-            Self::increase_stake_on_coldkey_hotkey_account(
-                &nominator,
-                hotkey,
-                nominator_emission.to_num::<u64>(),
-            );
+        // --- 9 Iterate over each nominator.
+        if total_viable_nominator_stake != 0 {
+            for (nominator, nominator_stake) in Stake::<T>::iter_prefix(hotkey) {
+                // --- 10 Check if the stake was manually increased by the user since the last emission drain for this hotkey.
+                // If it was, skip this nominator as they will not receive their proportion of the emission.
+                if LastAddStakeIncrease::<T>::get(hotkey, nominator.clone()) > last_emission_drain {
+                    continue;
+                }
 
-            // --- 11* Record event and Subtract the nominator's emission from the remainder.
-            total_new_tao = total_new_tao.saturating_add(nominator_emission.to_num::<u64>());
-            remainder = remainder.saturating_sub(nominator_emission.to_num::<u64>());
+                // --- 11 Calculate this nominator's share of the emission.
+                let nominator_emission: I64F64 = I64F64::from_num(emission_minus_take)
+                    .saturating_mul(I64F64::from_num(nominator_stake))
+                    .checked_div(I64F64::from_num(total_viable_nominator_stake))
+                    .unwrap_or(I64F64::from_num(0));
+
+                // --- 12 Increase the stake for the nominator.
+                Self::increase_stake_on_coldkey_hotkey_account(
+                    &nominator,
+                    hotkey,
+                    nominator_emission.to_num::<u64>(),
+                );
+
+                // --- 13* Record event and Subtract the nominator's emission from the remainder.
+                total_new_tao = total_new_tao.saturating_add(nominator_emission.to_num::<u64>());
+                remainder = remainder.saturating_sub(nominator_emission.to_num::<u64>());
+            }
         }
 
-        // --- 13 Finally, add the stake to the hotkey itself, including its take and the remaining emission.
+        // --- 14 Finally, add the stake to the hotkey itself, including its take and the remaining emission.
         let hotkey_new_tao: u64 = hotkey_take.saturating_add(remainder);
         Self::increase_stake_on_hotkey_account(hotkey, hotkey_new_tao);
 
-        // --- 14 Record new tao creation event and return the amount created.
+        // --- 15 Record new tao creation event and return the amount created.
         total_new_tao = total_new_tao.saturating_add(hotkey_new_tao);
         total_new_tao
     }

pallets/subtensor/src/lib.rs

@@ -44,6 +44,7 @@ pub mod staking;
 pub mod subnets;
 pub mod swap;
 pub mod utils;
+use crate::utils::TransactionType;
 use macros::{config, dispatches, errors, events, genesis, hooks};
 
 // apparently this is stabilized since rust 1.36
@@ -1051,6 +1052,17 @@ pub mod pallet {
     /// =================================
     /// ==== Axon / Promo Endpoints =====
     /// =================================
+    #[pallet::storage] // --- NMAP ( hot, netuid, name ) --> last_block | Returns the last block of a transaction for a given key, netuid, and name.
+    pub type TransactionKeyLastBlock<T: Config> = StorageNMap<
+        _,
+        (
+            NMapKey<Blake2_128Concat, T::AccountId>, // hot
+            NMapKey<Identity, u16>,                  // netuid
+            NMapKey<Identity, u16>,                  // extrinsic enum.
+        ),
+        u64,
+        ValueQuery,
+    >;
     #[pallet::storage]
     /// --- MAP ( key ) --> last_block
     pub type LastTxBlock<T: Config> =

pallets/subtensor/src/macros/errors.rs

@@ -166,5 +166,7 @@ mod errors {
         ProportionOverflow,
         /// Too many children MAX 5.
         TooManyChildren,
+        /// Default transaction rate limit exceeded.
+        TxRateLimitExceeded,
     }
 }

pallets/subtensor/src/staking/add_stake.rs

@@ -70,8 +70,8 @@ impl<T: Config> Pallet<T> {
             Error::<T>::StakeRateLimitExceeded
         );
 
-        // If this is a nomination stake, check if total stake after adding will be above
-        // the minimum required stake.
+        // Set the last time the stake increased for nominator drain protection.
+        LastAddStakeIncrease::<T>::insert(&hotkey, &coldkey, Self::get_current_block_as_u64());
 
         // If coldkey is not owner of the hotkey, it's a nomination stake.
         if !Self::coldkey_owns_hotkey(&coldkey, &hotkey) {

pallets/subtensor/src/staking/set_children.rs

@@ -58,6 +58,15 @@ impl<T: Config> Pallet<T> {
             children
         );
 
+        // Ensure the hotkey passes the rate limit.
+        ensure!(
+            Self::passes_rate_limit_globally(
+                &TransactionType::SetChildren, // Set children.
+                &hotkey,                       // Specific to a hotkey.
+            ),
+            Error::<T>::TxRateLimitExceeded
+        );
+
         // --- 2. Check that this delegation is not on the root network. Child hotkeys are not valid on root.
         ensure!(
             netuid != Self::get_root_netuid(),

pallets/subtensor/src/utils.rs

@@ -7,6 +7,33 @@ use sp_core::Get;
 use sp_core::U256;
 use substrate_fixed::types::I32F32;
 
+/// Enum representing different types of transactions
+#[derive(Copy, Clone)]
+pub enum TransactionType {
+    SetChildren,
+    Unknown,
+}
+
+/// Implement conversion from TransactionType to u16
+impl From<TransactionType> for u16 {
+    fn from(tx_type: TransactionType) -> Self {
+        match tx_type {
+            TransactionType::SetChildren => 0,
+            TransactionType::Unknown => 1,
+        }
+    }
+}
+
+/// Implement conversion from u16 to TransactionType
+impl From<u16> for TransactionType {
+    fn from(value: u16) -> Self {
+        match value {
+            0 => TransactionType::SetChildren,
+            _ => TransactionType::Unknown,
+        }
+    }
+}
+
 impl<T: Config> Pallet<T> {
     pub fn ensure_subnet_owner_or_root(
         o: T::RuntimeOrigin,
@@ -278,6 +305,56 @@ impl<T: Config> Pallet<T> {
     // ========================
     // ==== Rate Limiting =====
     // ========================
+    /// Get the rate limit for a specific transaction type
+    pub fn get_rate_limit(tx_type: &TransactionType) -> u64 {
+        match tx_type {
+            TransactionType::SetChildren => (DefaultTempo::<T>::get().saturating_mul(2)).into(), // Cannot set children twice within the default tempo period.
+            TransactionType::Unknown => 0, // Default to no limit for unknown types (no limit)
+        }
+    }
+
+    /// Check if a transaction should be rate limited on a specific subnet
+    pub fn passes_rate_limit_on_subnet(
+        tx_type: &TransactionType,
+        hotkey: &T::AccountId,
+        netuid: u16,
+    ) -> bool {
+        let block: u64 = Self::get_current_block_as_u64();
+        let limit: u64 = Self::get_rate_limit(tx_type);
+        let last_block: u64 = Self::get_last_transaction_block(hotkey, netuid, tx_type);
+        block.saturating_sub(last_block) < limit
+    }
+
+    /// Check if a transaction should be rate limited globally
+    pub fn passes_rate_limit_globally(tx_type: &TransactionType, hotkey: &T::AccountId) -> bool {
+        let netuid: u16 = u16::MAX;
+        let block: u64 = Self::get_current_block_as_u64();
+        let limit: u64 = Self::get_rate_limit(tx_type);
+        let last_block: u64 = Self::get_last_transaction_block(hotkey, netuid, tx_type);
+        block.saturating_sub(last_block) >= limit
+    }
+
+    /// Get the block number of the last transaction for a specific hotkey, network, and transaction type
+    pub fn get_last_transaction_block(
+        hotkey: &T::AccountId,
+        netuid: u16,
+        tx_type: &TransactionType,
+    ) -> u64 {
+        let tx_as_u16: u16 = (*tx_type).into();
+        TransactionKeyLastBlock::<T>::get((hotkey, netuid, tx_as_u16))
+    }
+
+    /// Set the block number of the last transaction for a specific hotkey, network, and transaction type
+    pub fn set_last_transaction_block(
+        hotkey: &T::AccountId,
+        netuid: u16,
+        tx_type: &TransactionType,
+        block: u64,
+    ) {
+        let tx_as_u16: u16 = (*tx_type).into();
+        TransactionKeyLastBlock::<T>::insert((hotkey, netuid, tx_as_u16), block);
+    }
+
     pub fn set_last_tx_block(key: &T::AccountId, block: u64) {
         LastTxBlock::<T>::insert(key, block)
     }