Introduce BypassDesicion in pallet-rate-limiting

ales-otf · ales-otf · commit 09ba8786c1bf · 2025-12-04T20:03:58.000+03:00
- fix bypassing in pallet-rate-limiting
- fix conditional bypassing for legacy rate-limited calls
diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs
@@ -97,8 +97,12 @@
 //!         }
 //!     }
 //!
-//!     fn should_bypass(origin: &RuntimeOrigin, _call: &RuntimeCall) -> bool {
-//!         matches!(origin, RuntimeOrigin::Root)
+//!     fn should_bypass(origin: &RuntimeOrigin, _call: &RuntimeCall) -> BypassDecision {
+//!         if matches!(origin, RuntimeOrigin::Root) {
+//!             BypassDecision::bypass_and_skip()
+//!         } else {
+//!             BypassDecision::enforce_and_record()
+//!         }
 //!     }
 //!
 //!     fn adjust_span(_origin: &RuntimeOrigin, _call: &RuntimeCall, span: BlockNumber) -> BlockNumber {
@@ -140,7 +144,7 @@ pub use benchmarking::BenchmarkHelper;
 pub use pallet::*;
 pub use tx_extension::RateLimitTransactionExtension;
 pub use types::{
-    GroupSharing, RateLimit, RateLimitGroup, RateLimitKind, RateLimitScopeResolver,
+    BypassDecision, GroupSharing, RateLimit, RateLimitGroup, RateLimitKind, RateLimitScopeResolver,
     RateLimitTarget, RateLimitUsageResolver, TransactionIdentifier,
 };
 
@@ -172,8 +176,8 @@ pub mod pallet {
     #[cfg(feature = "runtime-benchmarks")]
     use crate::benchmarking::BenchmarkHelper as BenchmarkHelperTrait;
     use crate::types::{
-        GroupSharing, RateLimit, RateLimitGroup, RateLimitKind, RateLimitScopeResolver,
-        RateLimitTarget, RateLimitUsageResolver, TransactionIdentifier,
+        BypassDecision, GroupSharing, RateLimit, RateLimitGroup, RateLimitKind,
+        RateLimitScopeResolver, RateLimitTarget, RateLimitUsageResolver, TransactionIdentifier,
     };
 
     type GroupNameOf<T, I> = BoundedVec<u8, <T as Config<I>>::MaxGroupNameLength>;
@@ -542,7 +546,9 @@ pub mod pallet {
             scope: &Option<<T as Config<I>>::LimitScope>,
             usage_key: &Option<<T as Config<I>>::UsageKey>,
         ) -> Result<bool, DispatchError> {
-            if <T as Config<I>>::LimitScopeResolver::should_bypass(origin, call) {
+            let bypass: BypassDecision =
+                <T as Config<I>>::LimitScopeResolver::should_bypass(origin, call);
+            if bypass.bypass_enforcement {
                 return Ok(true);
             }
 
@@ -571,8 +577,8 @@ pub mod pallet {
             })
         }
 
-        /// Resolves the span for a target/scope and applies the configured span adjustment
-        /// (e.g., tempo scaling) using the pallet's scope resolver.
+        /// Resolves the span for a target/scope and applies the configured span adjustment (e.g.,
+        /// tempo scaling) using the pallet's scope resolver.
         pub fn effective_span(
             origin: &DispatchOriginOf<<T as Config<I>>::RuntimeCall>,
             call: &<T as Config<I>>::RuntimeCall,
@@ -594,7 +600,7 @@ pub mod pallet {
                 return true;
             }
 
-            if let Some(last) = LastSeen::<T, I>::get(target, usage_key.clone()) {
+            if let Some(last) = LastSeen::<T, I>::get(target, usage_key) {
                 let current = frame_system::Pallet::<T>::block_number();
                 let delta = current.saturating_sub(last);
                 if delta < block_span {
@@ -753,6 +759,12 @@ pub mod pallet {
             Ok(())
         }
 
+        /// Returns true when the call has been registered (either directly or via a group).
+        pub fn is_registered(identifier: &TransactionIdentifier) -> bool {
+            let tx_target = RateLimitTarget::Transaction(*identifier);
+            Limits::<T, I>::contains_key(tx_target) || CallGroups::<T, I>::contains_key(identifier)
+        }
+
         fn call_metadata(
             identifier: &TransactionIdentifier,
         ) -> Result<(Vec<u8>, Vec<u8>), DispatchError> {
diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs
@@ -75,11 +75,19 @@ impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeOrigin, RuntimeCall, Li
         }
     }
 
-    fn should_bypass(_origin: &RuntimeOrigin, call: &RuntimeCall) -> bool {
-        matches!(
-            call,
-            RuntimeCall::RateLimiting(RateLimitingCall::remove_call_from_group { .. })
-        )
+    fn should_bypass(
+        _origin: &RuntimeOrigin,
+        call: &RuntimeCall,
+    ) -> pallet_rate_limiting::types::BypassDecision {
+        match call {
+            RuntimeCall::RateLimiting(RateLimitingCall::remove_call_from_group { .. }) => {
+                pallet_rate_limiting::types::BypassDecision::bypass_and_skip()
+            }
+            RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { .. }) => {
+                pallet_rate_limiting::types::BypassDecision::bypass_and_record()
+            }
+            _ => pallet_rate_limiting::types::BypassDecision::enforce_and_record(),
+        }
     }
 
     fn adjust_span(_origin: &RuntimeOrigin, call: &RuntimeCall, span: u64) -> u64 {
diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs
@@ -16,7 +16,11 @@ fn remark_call() -> RuntimeCall {
 }
 
 fn scoped_call() -> RuntimeCall {
-    RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 1 })
+    RuntimeCall::RateLimiting(RateLimitingCall::set_rate_limit {
+        target: RateLimitTarget::Transaction(TransactionIdentifier::new(0, 0)),
+        scope: Some(1),
+        limit: RateLimitKind::Default,
+    })
 }
 
 fn register(call: RuntimeCall, group: Option<GroupId>) -> TransactionIdentifier {
@@ -577,15 +581,15 @@ fn is_within_limit_detects_rate_limited_scope() {
         let tx_target = target(identifier);
         Limits::<Test, ()>::insert(
             tx_target,
-            RateLimit::scoped_single(7u16, RateLimitKind::Exact(3)),
+            RateLimit::scoped_single(1u16, RateLimitKind::Exact(3)),
         );
         LastSeen::<Test, ()>::insert(tx_target, Some(1u16), 9);
         System::set_block_number(11);
         let result = RateLimiting::is_within_limit(
             &RuntimeOrigin::signed(1),
             &call,
             &identifier,
-            &Some(7u16),
+            &Some(1u16),
             &Some(1u16),
         )
         .expect("ok");
@@ -714,12 +718,9 @@ fn limit_for_call_names_prefers_scoped_value() {
             target(identifier),
             RateLimit::scoped_single(9u16, RateLimitKind::Exact(8)),
         );
-        let fetched = RateLimiting::limit_for_call_names(
-            "RateLimiting",
-            "set_default_rate_limit",
-            Some(9u16),
-        )
-        .expect("limit");
+        let fetched =
+            RateLimiting::limit_for_call_names("RateLimiting", "set_rate_limit", Some(9u16))
+                .expect("limit");
         assert_eq!(fetched, RateLimitKind::Exact(8));
     });
 }
diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs
@@ -117,30 +117,40 @@ where
         _inherited_implication: &impl Implication,
         _source: TransactionSource,
     ) -> ValidateResult<Self::Val, <T as Config<I>>::RuntimeCall> {
-        if <T as Config<I>>::LimitScopeResolver::should_bypass(&origin, call) {
-            return Ok((ValidTransaction::default(), None, origin));
-        }
-
         let identifier = match TransactionIdentifier::from_call::<T, I>(call) {
             Ok(identifier) => identifier,
             Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)),
         };
 
+        if !Pallet::<T, I>::is_registered(&identifier) {
+            return Ok((ValidTransaction::default(), None, origin));
+        }
+
         let scope = <T as Config<I>>::LimitScopeResolver::context(&origin, call);
         let usage = <T as Config<I>>::UsageResolver::context(&origin, call);
 
         let config_target = Pallet::<T, I>::config_target(&identifier)
             .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?;
         let usage_target = Pallet::<T, I>::usage_target(&identifier)
             .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?;
-        let should_record = Pallet::<T, I>::should_record_usage(&identifier, &usage_target);
+        let bypass = <T as Config<I>>::LimitScopeResolver::should_bypass(&origin, call);
+        let should_record =
+            bypass.record_usage && Pallet::<T, I>::should_record_usage(&identifier, &usage_target);
 
         let Some(block_span) =
             Pallet::<T, I>::effective_span(&origin, call, &config_target, &scope)
         else {
             return Ok((ValidTransaction::default(), None, origin));
         };
 
+        if bypass.bypass_enforcement {
+            return Ok((
+                ValidTransaction::default(),
+                should_record.then_some((usage_target, usage, true)),
+                origin,
+            ));
+        }
+
         if block_span.is_zero() {
             return Ok((ValidTransaction::default(), None, origin));
         }
@@ -340,6 +350,52 @@ mod tests {
         });
     }
 
+    #[test]
+    fn tx_extension_records_usage_on_bypass() {
+        new_test_ext().execute_with(|| {
+            let extension = new_tx_extension();
+            let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit {
+                block_span: 2,
+            });
+            let identifier = identifier_for(&call);
+            let target = RateLimitTarget::Transaction(identifier);
+
+            assert_ok!(RateLimiting::register_call(
+                RuntimeOrigin::root(),
+                Box::new(call.clone()),
+                None,
+            ));
+
+            System::set_block_number(5);
+
+            let (_valid, val, origin) =
+                validate_with_tx_extension(&extension, &call).expect("bypass should succeed");
+            assert!(val.is_some(), "bypass decision should still record usage");
+
+            let info = call.get_dispatch_info();
+            let len = call.encode().len();
+            let pre = extension
+                .clone()
+                .prepare(val.clone(), &origin, &call, &info, len)
+                .expect("prepare succeeds");
+
+            let mut post = PostDispatchInfo::default();
+            RateLimitTransactionExtension::<Test>::post_dispatch(
+                pre,
+                &info,
+                &mut post,
+                len,
+                &Ok(()),
+            )
+            .expect("post_dispatch succeeds");
+
+            assert_eq!(
+                LastSeen::<Test, ()>::get(target, Some(2u16)),
+                Some(5u64.into())
+            );
+        });
+    }
+
     #[test]
     fn tx_extension_records_last_seen_for_successful_call() {
         new_test_ext().execute_with(|| {
diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs
@@ -11,10 +11,9 @@ pub trait RateLimitScopeResolver<Origin, Call, Scope, Span> {
     /// limits.
     fn context(origin: &Origin, call: &Call) -> Option<Scope>;
 
-    /// Returns `true` when the rate limit should be bypassed for the provided origin/call pair.
-    /// Defaults to `false`.
-    fn should_bypass(_origin: &Origin, _call: &Call) -> bool {
-        false
+    /// Returns how the call should interact with enforcement and usage tracking.
+    fn should_bypass(_origin: &Origin, _call: &Call) -> BypassDecision {
+        BypassDecision::enforce_and_record()
     }
 
     /// Optionally adjusts the effective span used during enforcement. Defaults to the original
@@ -24,6 +23,34 @@ pub trait RateLimitScopeResolver<Origin, Call, Scope, Span> {
     }
 }
 
+/// Controls whether enforcement should run and whether usage should be recorded for a call.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct BypassDecision {
+    pub bypass_enforcement: bool,
+    pub record_usage: bool,
+}
+
+impl BypassDecision {
+    pub const fn new(bypass_enforcement: bool, record_usage: bool) -> Self {
+        Self {
+            bypass_enforcement,
+            record_usage,
+        }
+    }
+
+    pub const fn enforce_and_record() -> Self {
+        Self::new(false, true)
+    }
+
+    pub const fn bypass_and_record() -> Self {
+        Self::new(true, true)
+    }
+
+    pub const fn bypass_and_skip() -> Self {
+        Self::new(true, false)
+    }
+}
+
 /// Resolves the optional usage tracking key applied when enforcing limits.
 pub trait RateLimitUsageResolver<Origin, Call, Usage> {
     /// Returns `Some(usage)` when usage should be tracked per-key, or `None` for global usage
diff --git a/runtime/src/rate_limiting/migration.rs b/runtime/src/rate_limiting/migration.rs
@@ -404,10 +404,6 @@ pub fn rate_limited_calls() -> Vec<RateLimitedCall> {
                 subtensor_identifier(3),   // remove_stake
                 subtensor_identifier(89),  // remove_stake_limit
                 subtensor_identifier(103), // remove_stake_full_limit
-                subtensor_identifier(86),  // transfer_stake
-                subtensor_identifier(85),  // move_stake
-                subtensor_identifier(87),  // swap_stake
-                subtensor_identifier(90),  // swap_stake_limit
             ],
             legacy: sources(&["TxRateLimit"], &[]),
         },
@@ -423,7 +419,10 @@ pub fn rate_limited_calls() -> Vec<RateLimitedCall> {
             scope: LimitScopeKind::Global,
             usage: UsageKind::AccountSubnet,
             read_only: Vec::new(),
-            legacy: sources(&["TxChildkeyTakeRateLimit"], &["TransactionKeyLastBlock::SetChildkeyTake"]),
+            legacy: sources(
+                &["TxChildkeyTakeRateLimit"],
+                &["TransactionKeyLastBlock::SetChildkeyTake"],
+            ),
         },
         RateLimitedCall {
             target: TargetKind::Standalone(subtensor_identifier(67)), // set_children
@@ -462,10 +461,7 @@ pub fn rate_limited_calls() -> Vec<RateLimitedCall> {
             scope: LimitScopeKind::Global,
             usage: UsageKind::SubnetNeuron,
             read_only: Vec::new(),
-            legacy: sources(
-                &["EvmKeyAssociateRateLimit"],
-                &["AssociatedEvmAddress"],
-            ),
+            legacy: sources(&["EvmKeyAssociateRateLimit"], &["AssociatedEvmAddress"]),
         },
         RateLimitedCall {
             target: TargetKind::Standalone(admin_utils_identifier(76)), // sudo_set_mechanism_count
@@ -548,10 +544,10 @@ type LimitImporter<T> = fn(&Grouping, &mut LimitEntries<T>) -> u64;
 
 fn limit_importers<T: SubtensorConfig>() -> [LimitImporter<T>; 4] {
     [
-        import_simple_limits::<T>,      // Tx/childkey/delegate/staking lock, register, sudo, evm, children
+        import_simple_limits::<T>, // Tx/childkey/delegate/staking lock, register, sudo, evm, children
         import_owner_hparam_limits::<T>, // Owner hyperparams
-        import_serving_limits::<T>,     // Axon/prometheus serving rate limit per subnet
-        import_weight_limits::<T>,      // Weight/commit/reveal per subnet and mechanism
+        import_serving_limits::<T>, // Axon/prometheus serving rate limit per subnet
+        import_weight_limits::<T>, // Weight/commit/reveal per subnet and mechanism
     ]
 }
 
@@ -581,9 +577,9 @@ fn import_simple_limits<T: SubtensorConfig>(
         );
     }
 
-    // Share the TxRateLimit span across staking operations; add_* are marker-only via span tweaks.
-    if let Some(span) = block_number::<T>(TxRateLimit::<T>::get()) {
-        if let Some(members) = grouping.members(GROUP_STAKING_OPS) {
+    // Staking ops are gated to one operation per block in legacy (marker cleared each block).
+    if let Some(members) = grouping.members(GROUP_STAKING_OPS) {
+        if let Some(span) = block_number::<T>(1) {
             for call in members {
                 set_global_limit::<T>(limits, grouping.config_target(*call), span);
             }
@@ -761,9 +757,9 @@ fn last_seen_importers<T: SubtensorConfig>() -> [LastSeenImporter<T>; 5] {
     [
         import_last_rate_limited_blocks::<T>, // LastRateLimitedBlock (tx, delegate, owner hyperparams, sn owner)
         import_transaction_key_last_blocks::<T>, // TransactionKeyLastBlock (children, version key, mechanisms)
-        import_last_update_entries::<T>, // LastUpdate (weights/mechanism weights)
-        import_serving_entries::<T>,     // Axons/Prometheus
-        import_evm_entries::<T>,         // AssociatedEvmAddress
+        import_last_update_entries::<T>,         // LastUpdate (weights/mechanism weights)
+        import_serving_entries::<T>,             // Axons/Prometheus
+        import_evm_entries::<T>,                 // AssociatedEvmAddress
     ]
 }
 
@@ -866,8 +862,8 @@ fn import_last_update_entries<T: SubtensorConfig>(
     let mut reads: u64 = 0;
     for (index, blocks) in LastUpdate::<T>::iter() {
         reads += 1;
-        let (netuid, mecid) = Pallet::<T>::get_netuid_and_subid(index)
-            .unwrap_or((NetUid::ROOT, 0.into()));
+        let (netuid, mecid) =
+            Pallet::<T>::get_netuid_and_subid(index).unwrap_or((NetUid::ROOT, 0.into()));
         let subnet_calls = if mecid == 0.into() {
             weight_calls_subnet(grouping)
         } else {
@@ -882,7 +878,10 @@ fn import_last_update_entries<T: SubtensorConfig>(
                 continue;
             };
             let usage = if mecid == 0.into() {
-                RateLimitUsageKey::SubnetNeuron { netuid, uid: uid_u16 }
+                RateLimitUsageKey::SubnetNeuron {
+                    netuid,
+                    uid: uid_u16,
+                }
             } else {
                 RateLimitUsageKey::SubnetMechanismNeuron {
                     netuid,
diff --git a/runtime/src/rate_limiting/mod.rs b/runtime/src/rate_limiting/mod.rs
