diff --git a/pallets/subtensor/src/coinbase/run_coinbase.rs b/pallets/subtensor/src/coinbase/run_coinbase.rs index abe44d55a..fcf76728f 100644 --- a/pallets/subtensor/src/coinbase/run_coinbase.rs +++ b/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 Pallet { PendingdHotkeyEmission::::insert(hotkey, 0); // --- 2 Retrieve the last time this hotkey's emissions were drained. - let last_hotkey_emission_drain: u64 = LastHotkeyEmissionDrain::::get(hotkey); + let last_emission_drain: u64 = LastHotkeyEmissionDrain::::get(hotkey); // --- 3 Update the block value to the current block number. LastHotkeyEmissionDrain::::insert(hotkey, block_number); @@ -282,43 +281,48 @@ impl Pallet { // --- 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 - as IterableStorageDoubleMap>::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::::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::::iter_prefix(hotkey) { + if LastAddStakeIncrease::::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::(), - ); + // --- 9 Iterate over each nominator. + if total_viable_nominator_stake != 0 { + for (nominator, nominator_stake) in Stake::::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::::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::()); - remainder = remainder.saturating_sub(nominator_emission.to_num::()); + // --- 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::(), + ); + + // --- 13* Record event and Subtract the nominator's emission from the remainder. + total_new_tao = total_new_tao.saturating_add(nominator_emission.to_num::()); + remainder = remainder.saturating_sub(nominator_emission.to_num::()); + } } - // --- 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 } diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 6c16902d9..abf6a8613 100644 --- a/pallets/subtensor/src/lib.rs +++ b/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 = StorageNMap< + _, + ( + NMapKey, // hot + NMapKey, // netuid + NMapKey, // extrinsic enum. + ), + u64, + ValueQuery, + >; #[pallet::storage] /// --- MAP ( key ) --> last_block pub type LastTxBlock = diff --git a/pallets/subtensor/src/macros/errors.rs b/pallets/subtensor/src/macros/errors.rs index 7e244834d..156cbea56 100644 --- a/pallets/subtensor/src/macros/errors.rs +++ b/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, } } diff --git a/pallets/subtensor/src/staking/add_stake.rs b/pallets/subtensor/src/staking/add_stake.rs index eb7762bec..f62fd7cd2 100644 --- a/pallets/subtensor/src/staking/add_stake.rs +++ b/pallets/subtensor/src/staking/add_stake.rs @@ -70,8 +70,8 @@ impl Pallet { Error::::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::::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) { diff --git a/pallets/subtensor/src/staking/set_children.rs b/pallets/subtensor/src/staking/set_children.rs index 071764734..f413db23f 100644 --- a/pallets/subtensor/src/staking/set_children.rs +++ b/pallets/subtensor/src/staking/set_children.rs @@ -58,6 +58,15 @@ impl Pallet { children ); + // Ensure the hotkey passes the rate limit. + ensure!( + Self::passes_rate_limit_globally( + &TransactionType::SetChildren, // Set children. + &hotkey, // Specific to a hotkey. + ), + Error::::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(), diff --git a/pallets/subtensor/src/utils.rs b/pallets/subtensor/src/utils.rs index 6f5dbeaff..2cd49e198 100644 --- a/pallets/subtensor/src/utils.rs +++ b/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 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 for TransactionType { + fn from(value: u16) -> Self { + match value { + 0 => TransactionType::SetChildren, + _ => TransactionType::Unknown, + } + } +} + impl Pallet { pub fn ensure_subnet_owner_or_root( o: T::RuntimeOrigin, @@ -278,6 +305,56 @@ impl Pallet { // ======================== // ==== 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::::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::::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::::insert((hotkey, netuid, tx_as_u16), block); + } + pub fn set_last_tx_block(key: &T::AccountId, block: u64) { LastTxBlock::::insert(key, block) }