Add rate limit adjuster

ales-otf · ales-otf · commit 5c8a8abf7466 · 2025-10-31T16:31:48.000+03:00
diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs
@@ -48,8 +48,9 @@
 //! The pallet relies on two resolvers:
 //!
 //! - [`Config::LimitScopeResolver`], which determines how limits are stored (for example by
-//!   returning a `netuid`). When this resolver returns `None`, the configuration is stored as a
-//!   global fallback.
+//!   returning a `netuid`). The resolver can also signal that a call should bypass rate limiting or
+//!   adjust the effective span at validation time. When it returns `None`, the configuration is
+//!   stored as a global fallback.
 //! - [`Config::UsageResolver`], which decides how executions are tracked in
 //!   [`LastSeen`](pallet::LastSeen). This can refine the limit scope (for example by returning a
 //!   tuple of `(netuid, hyperparameter)`).
@@ -63,7 +64,7 @@
 //!
 //! // Limits are scoped per netuid.
 //! pub struct ScopeResolver;
-//! impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, NetUid> for ScopeResolver {
+//! impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, NetUid, BlockNumber> for ScopeResolver {
 //!     fn context(call: &RuntimeCall) -> Option<NetUid> {
 //!         match call {
 //!             RuntimeCall::Subtensor(pallet_subtensor::Call::set_weights { netuid, .. }) => {
@@ -72,12 +73,15 @@
 //!             _ => None,
 //!         }
 //!     }
+//!
+//!     fn adjust_span(_call: &RuntimeCall, span: BlockNumber) -> BlockNumber {
+//!         span
+//!     }
 //! }
 //!
 //! // Usage tracking distinguishes hyperparameter + netuid.
 //! pub struct UsageResolver;
-//! impl pallet_rate_limiting::RateLimitUsageResolver<RuntimeCall, (NetUid, HyperParam)>
-//!     for UsageResolver {
+//! impl pallet_rate_limiting::RateLimitUsageResolver<RuntimeCall, (NetUid, HyperParam)> for UsageResolver {
 //!     fn context(call: &RuntimeCall) -> Option<(NetUid, HyperParam)> {
 //!         match call {
 //!             RuntimeCall::Subtensor(pallet_subtensor::Call::set_hyperparam {
@@ -156,7 +160,11 @@ pub mod pallet {
         type LimitScope: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize;
 
         /// Resolves the scope for the given runtime call when configuring limits.
-        type LimitScopeResolver: RateLimitScopeResolver<<Self as Config<I>>::RuntimeCall, Self::LimitScope>;
+        type LimitScopeResolver: RateLimitScopeResolver<
+                <Self as Config<I>>::RuntimeCall,
+                Self::LimitScope,
+                BlockNumberFor<Self>,
+            >;
 
         /// Usage key tracked in [`LastSeen`] for rate-limited calls.
         type UsageKey: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize;
@@ -306,21 +314,17 @@ pub mod pallet {
             identifier: &TransactionIdentifier,
             scope: &Option<<T as Config<I>>::LimitScope>,
             usage_key: &Option<<T as Config<I>>::UsageKey>,
+            call: &<T as Config<I>>::RuntimeCall,
         ) -> Result<bool, DispatchError> {
-            let Some(block_span) = Self::resolved_limit(identifier, scope) else {
+            if <T as Config<I>>::LimitScopeResolver::should_bypass(call) {
                 return Ok(true);
-            };
-
-            let current = frame_system::Pallet::<T>::block_number();
-
-            if let Some(last) = LastSeen::<T, I>::get(identifier, usage_key) {
-                let delta = current.saturating_sub(last);
-                if delta < block_span {
-                    return Ok(false);
-                }
             }
 
-            Ok(true)
+            let Some(block_span) = Self::effective_span(call, identifier, scope) else {
+                return Ok(true);
+            };
+
+            Ok(Self::within_span(identifier, usage_key, block_span))
         }
 
         pub(crate) fn resolved_limit(
@@ -335,6 +339,37 @@ pub mod pallet {
             })
         }
 
+        pub(crate) fn effective_span(
+            call: &<T as Config<I>>::RuntimeCall,
+            identifier: &TransactionIdentifier,
+            scope: &Option<<T as Config<I>>::LimitScope>,
+        ) -> Option<BlockNumberFor<T>> {
+            let span = Self::resolved_limit(identifier, scope)?;
+            Some(<T as Config<I>>::LimitScopeResolver::adjust_span(
+                call, span,
+            ))
+        }
+
+        pub(crate) fn within_span(
+            identifier: &TransactionIdentifier,
+            usage_key: &Option<<T as Config<I>>::UsageKey>,
+            block_span: BlockNumberFor<T>,
+        ) -> bool {
+            if block_span.is_zero() {
+                return true;
+            }
+
+            if let Some(last) = LastSeen::<T, I>::get(identifier, usage_key) {
+                let current = frame_system::Pallet::<T>::block_number();
+                let delta = current.saturating_sub(last);
+                if delta < block_span {
+                    return false;
+                }
+            }
+
+            true
+        }
+
         /// Returns the configured limit for the specified pallet/extrinsic names, if any.
         pub fn limit_for_call_names(
             pallet_name: &str,
diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs
@@ -61,7 +61,9 @@ pub type UsageKey = u16;
 pub struct TestScopeResolver;
 pub struct TestUsageResolver;
 
-impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, LimitScope> for TestScopeResolver {
+impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, LimitScope, u64>
+    for TestScopeResolver
+{
     fn context(call: &RuntimeCall) -> Option<LimitScope> {
         match call {
             RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => {
@@ -78,6 +80,17 @@ impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, LimitScope> for T
             RuntimeCall::RateLimiting(RateLimitingCall::clear_rate_limit { .. })
         )
     }
+
+    fn adjust_span(call: &RuntimeCall, span: u64) -> u64 {
+        if matches!(
+            call,
+            RuntimeCall::RateLimiting(RateLimitingCall::clear_all_rate_limits { .. })
+        ) {
+            span.saturating_mul(2)
+        } else {
+            span
+        }
+    }
 }
 
 impl pallet_rate_limiting::RateLimitUsageResolver<RuntimeCall, UsageKey> for TestUsageResolver {
diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs
@@ -121,7 +121,7 @@ fn is_within_limit_is_true_when_no_limit() {
             RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 });
         let identifier = identifier_for(&call);
 
-        let result = RateLimiting::is_within_limit(&identifier, &None, &None);
+        let result = RateLimiting::is_within_limit(&identifier, &None, &None, &call);
         assert_eq!(result.expect("no error expected"), true);
     });
 }
@@ -144,6 +144,7 @@ fn is_within_limit_false_when_rate_limited() {
             &identifier,
             &Some(1 as LimitScope),
             &Some(1 as UsageKey),
+            &call,
         )
         .expect("call succeeds");
         assert!(!within);
@@ -168,6 +169,7 @@ fn is_within_limit_true_after_required_span() {
             &identifier,
             &Some(2 as LimitScope),
             &Some(2 as UsageKey),
+            &call,
         )
         .expect("call succeeds");
         assert!(within);
diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs
@@ -109,16 +109,15 @@ where
         let scope = <T as Config<I>>::LimitScopeResolver::context(call);
         let usage = <T as Config<I>>::UsageResolver::context(call);
 
-        let Some(block_span) = Pallet::<T, I>::resolved_limit(&identifier, &scope) else {
+        let Some(block_span) = Pallet::<T, I>::effective_span(call, &identifier, &scope) else {
             return Ok((ValidTransaction::default(), None, origin));
         };
 
         if block_span.is_zero() {
             return Ok((ValidTransaction::default(), None, origin));
         }
 
-        let within_limit = Pallet::<T, I>::is_within_limit(&identifier, &scope, &usage)
-            .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?;
+        let within_limit = Pallet::<T, I>::within_span(&identifier, &usage, block_span);
 
         if !within_limit {
             return Err(TransactionValidityError::Invalid(
@@ -188,6 +187,12 @@ mod tests {
         })
     }
 
+    fn adjustable_call() -> RuntimeCall {
+        RuntimeCall::RateLimiting(RateLimitingCall::clear_all_rate_limits {
+            call: Box::new(remark_call()),
+        })
+    }
+
     fn new_tx_extension() -> RateLimitTransactionExtension<Test> {
         RateLimitTransactionExtension(Default::default())
     }
@@ -273,6 +278,29 @@ mod tests {
         });
     }
 
+    #[test]
+    fn tx_extension_applies_adjusted_span() {
+        new_test_ext().execute_with(|| {
+            let extension = new_tx_extension();
+            let call = adjustable_call();
+            let identifier = identifier_for(&call);
+            Limits::<Test, ()>::insert(identifier, RateLimit::global(RateLimitKind::Exact(4)));
+            LastSeen::<Test, ()>::insert(identifier, Some(1u16), 10);
+
+            System::set_block_number(14);
+
+            // Stored span (4) would allow the call, but adjusted span (8) should block it.
+            let err = validate_with_tx_extension(&extension, &call)
+                .expect_err("adjusted span should apply");
+            match err {
+                TransactionValidityError::Invalid(InvalidTransaction::Custom(code)) => {
+                    assert_eq!(code, RATE_LIMIT_DENIED);
+                }
+                other => panic!("unexpected error: {:?}", other),
+            }
+        });
+    }
+
     #[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
@@ -3,24 +3,30 @@ use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata};
 use scale_info::TypeInfo;
 use sp_std::collections::btree_map::BTreeMap;
 
-/// Resolves the optional identifier within which a rate limit applies and can optionally bypass
-/// enforcement.
-pub trait RateLimitScopeResolver<Call, Scope> {
-	/// Returns `Some(scope)` when the limit should be applied per-scope, or `None` for global
-	/// limits.
+/// Resolves the optional identifier within which a rate limit applies and can optionally adjust
+/// enforcement behaviour.
+pub trait RateLimitScopeResolver<Call, Scope, Span> {
+    /// Returns `Some(scope)` when the limit should be applied per-scope, or `None` for global
+    /// limits.
     fn context(call: &Call) -> Option<Scope>;
 
-	/// Returns `true` when the rate limit should be bypassed for the provided call. Defaults to
-	/// `false`.
+    /// Returns `true` when the rate limit should be bypassed for the provided call. Defaults to
+    /// `false`.
     fn should_bypass(_call: &Call) -> bool {
         false
     }
+
+    /// Optionally adjusts the effective span used during enforcement. Defaults to the original
+    /// `span`.
+    fn adjust_span(_call: &Call, span: Span) -> Span {
+        span
+    }
 }
 
 /// Resolves the optional usage tracking key applied when enforcing limits.
 pub trait RateLimitUsageResolver<Call, Usage> {
-	/// Returns `Some(usage)` when usage should be tracked per-key, or `None` for global usage
-	/// tracking.
+    /// Returns `Some(usage)` when usage should be tracked per-key, or `None` for global usage
+    /// tracking.
     fn context(call: &Call) -> Option<Usage>;
 }
 
