Feat/rate limit set weights version key

pallets/admin-utils/src/lib.rs

@@ -25,9 +25,13 @@ pub mod pallet {
     use super::*;
     use frame_support::pallet_prelude::*;
     use frame_support::traits::tokens::Balance;
-    use frame_support::{dispatch::DispatchResult, pallet_prelude::StorageMap};
+    use frame_support::{
+        dispatch::{DispatchResult, RawOrigin},
+        pallet_prelude::StorageMap,
+    };
     use frame_system::pallet_prelude::*;
     use pallet_evm_chain_id::{self, ChainId};
+    use pallet_subtensor::utils::rate_limiting::TransactionType;
     use sp_runtime::BoundedVec;
     use substrate_fixed::types::I96F32;
 
@@ -249,12 +253,35 @@ pub mod pallet {
             netuid: u16,
             weights_version_key: u64,
         ) -> DispatchResult {
-            pallet_subtensor::Pallet::<T>::ensure_subnet_owner_or_root(origin, netuid)?;
+            pallet_subtensor::Pallet::<T>::ensure_subnet_owner_or_root(origin.clone(), netuid)?;
 
             ensure!(
                 pallet_subtensor::Pallet::<T>::if_subnet_exist(netuid),
                 Error::<T>::SubnetDoesNotExist
             );
+
+            if let Ok(RawOrigin::Signed(who)) = origin.into() {
+                // SN Owner
+                // Ensure the origin passes the rate limit.
+                ensure!(
+                    pallet_subtensor::Pallet::<T>::passes_rate_limit_on_subnet(
+                        &TransactionType::SetWeightsVersionKey,
+                        &who,
+                        netuid,
+                    ),
+                    pallet_subtensor::Error::<T>::TxRateLimitExceeded
+                );
+
+                // Set last transaction block
+                let current_block = pallet_subtensor::Pallet::<T>::get_current_block_as_u64();
+                pallet_subtensor::Pallet::<T>::set_last_transaction_block_on_subnet(
+                    &who,
+                    netuid,
+                    &TransactionType::SetWeightsVersionKey,
+                    current_block,
+                );
+            }
+
             pallet_subtensor::Pallet::<T>::set_weights_version_key(netuid, weights_version_key);
             log::debug!(
                 "WeightsVersionKeySet( netuid: {:?} weights_version_key: {:?} ) ",

pallets/admin-utils/src/tests/mod.rs

@@ -5,7 +5,7 @@ use frame_support::{
     traits::Hooks,
 };
 use frame_system::Config;
-use pallet_subtensor::Error as SubtensorError;
+use pallet_subtensor::{Error as SubtensorError, SubnetOwner, Tempo, WeightsVersionKeyRateLimit};
 // use pallet_subtensor::{migrations, Event};
 use pallet_subtensor::Event;
 use sp_consensus_grandpa::AuthorityId as GrandpaId;
@@ -162,6 +162,107 @@ fn test_sudo_set_weights_version_key() {
     });
 }
 
+#[test]
+fn test_sudo_set_weights_version_key_rate_limit() {
+    new_test_ext().execute_with(|| {
+        let netuid: u16 = 1;
+        let to_be_set: u64 = 10;
+
+        let sn_owner = U256::from(1);
+        add_network(netuid, 10);
+        // Set the Subnet Owner
+        SubnetOwner::<Test>::insert(netuid, sn_owner);
+
+        let rate_limit = WeightsVersionKeyRateLimit::<Test>::get();
+        let tempo: u16 = Tempo::<Test>::get(netuid);
+
+        let rate_limit_period = rate_limit * (tempo as u64);
+
+        assert_ok!(AdminUtils::sudo_set_weights_version_key(
+            <<Test as Config>::RuntimeOrigin>::signed(sn_owner),
+            netuid,
+            to_be_set
+        ));
+        assert_eq!(SubtensorModule::get_weights_version_key(netuid), to_be_set);
+
+        // Try to set again with
+        // Assert rate limit not passed
+        assert!(!SubtensorModule::passes_rate_limit_on_subnet(
+            &pallet_subtensor::utils::rate_limiting::TransactionType::SetWeightsVersionKey,
+            &sn_owner,
+            netuid
+        ));
+
+        // Try transaction
+        assert_noop!(
+            AdminUtils::sudo_set_weights_version_key(
+                <<Test as Config>::RuntimeOrigin>::signed(sn_owner),
+                netuid,
+                to_be_set + 1
+            ),
+            pallet_subtensor::Error::<Test>::TxRateLimitExceeded
+        );
+
+        // Wait for rate limit to pass
+        run_to_block(rate_limit_period + 2);
+        assert!(SubtensorModule::passes_rate_limit_on_subnet(
+            &pallet_subtensor::utils::rate_limiting::TransactionType::SetWeightsVersionKey,
+            &sn_owner,
+            netuid
+        ));
+
+        // Try transaction
+        assert_ok!(AdminUtils::sudo_set_weights_version_key(
+            <<Test as Config>::RuntimeOrigin>::signed(sn_owner),
+            netuid,
+            to_be_set + 1
+        ));
+        assert_eq!(
+            SubtensorModule::get_weights_version_key(netuid),
+            to_be_set + 1
+        );
+    });
+}
+
+#[test]
+fn test_sudo_set_weights_version_key_rate_limit_root() {
+    // root should not be effected by rate limit
+    new_test_ext().execute_with(|| {
+        let netuid: u16 = 1;
+        let to_be_set: u64 = 10;
+
+        let sn_owner = U256::from(1);
+        add_network(netuid, 10);
+        // Set the Subnet Owner
+        SubnetOwner::<Test>::insert(netuid, sn_owner);
+
+        let rate_limit = WeightsVersionKeyRateLimit::<Test>::get();
+        let tempo: u16 = Tempo::<Test>::get(netuid);
+
+        let rate_limit_period = rate_limit * (tempo as u64);
+        // Verify the rate limit is more than 0 blocks
+        assert!(rate_limit_period > 0);
+
+        assert_ok!(AdminUtils::sudo_set_weights_version_key(
+            <<Test as Config>::RuntimeOrigin>::root(),
+            netuid,
+            to_be_set
+        ));
+        assert_eq!(SubtensorModule::get_weights_version_key(netuid), to_be_set);
+
+        // Try transaction
+        assert_ok!(AdminUtils::sudo_set_weights_version_key(
+            <<Test as Config>::RuntimeOrigin>::signed(sn_owner),
+            netuid,
+            to_be_set + 1
+        ));
+        assert_eq!(
+            SubtensorModule::get_weights_version_key(netuid),
+            to_be_set + 1
+        );
+    });
+}
+
 #[test]
 fn test_sudo_set_weights_set_rate_limit() {
     new_test_ext().execute_with(|| {

pallets/subtensor/src/lib.rs

@@ -509,6 +509,12 @@ pub mod pallet {
         T::InitialNetworkRateLimit::get()
     }
     #[pallet::type_value]
+    /// Default value for weights version key rate limit.
+    /// In units of tempos.
+    pub fn DefaultWeightsVersionKeyRateLimit<T: Config>() -> u64 {
+        5 // 5 tempos
+    }
+    #[pallet::type_value]
     /// Default value for pending emission.
     pub fn DefaultPendingEmission<T: Config>() -> u64 {
         0
@@ -1077,6 +1083,10 @@ pub mod pallet {
     pub type NetworkRateLimit<T> = StorageValue<_, u64, ValueQuery, DefaultNetworkRateLimit<T>>;
     #[pallet::storage] // --- ITEM( nominator_min_required_stake )
     pub type NominatorMinRequiredStake<T> = StorageValue<_, u64, ValueQuery, DefaultZeroU64<T>>;
+    #[pallet::storage]
+    /// ITEM( weights_version_key_rate_limit ) --- Rate limit in tempos.
+    pub type WeightsVersionKeyRateLimit<T> =
+        StorageValue<_, u64, ValueQuery, DefaultWeightsVersionKeyRateLimit<T>>;
 
     /// ============================
     /// ==== Subnet Locks =====

pallets/subtensor/src/utils/rate_limiting.rs

@@ -7,6 +7,7 @@ pub enum TransactionType {
     SetChildkeyTake,
     Unknown,
     RegisterNetwork,
+    SetWeightsVersionKey,
 }
 
 /// Implement conversion from TransactionType to u16
@@ -17,6 +18,7 @@ impl From<TransactionType> for u16 {
             TransactionType::SetChildkeyTake => 1,
             TransactionType::Unknown => 2,
             TransactionType::RegisterNetwork => 3,
+            TransactionType::SetWeightsVersionKey => 4,
         }
     }
 }
@@ -28,6 +30,7 @@ impl From<u16> for TransactionType {
             0 => TransactionType::SetChildren,
             1 => TransactionType::SetChildkeyTake,
             3 => TransactionType::RegisterNetwork,
+            4 => TransactionType::SetWeightsVersionKey,
             _ => TransactionType::Unknown,
         }
     }
@@ -41,14 +44,18 @@ impl<T: Config> Pallet<T> {
         match tx_type {
             TransactionType::SetChildren => 150, // 30 minutes
             TransactionType::SetChildkeyTake => TxChildkeyTakeRateLimit::<T>::get(),
-            TransactionType::Unknown => 0, // Default to no limit for unknown types (no limit)
             TransactionType::RegisterNetwork => NetworkRateLimit::<T>::get(),
+
+            TransactionType::Unknown => 0, // Default to no limit for unknown types (no limit)
+            _ => 0,
         }
     }
 
-    pub fn get_rate_limit_on_subnet(tx_type: &TransactionType, _netuid: u16) -> u64 {
+    pub fn get_rate_limit_on_subnet(tx_type: &TransactionType, netuid: u16) -> u64 {
         #[allow(clippy::match_single_binding)]
         match tx_type {
+            TransactionType::SetWeightsVersionKey => (Tempo::<T>::get(netuid) as u64)
+                .saturating_mul(WeightsVersionKeyRateLimit::<T>::get()),
             _ => Self::get_rate_limit(tx_type),
         }
     }