From a37cc79f2bc2f20e7f2a4c5074e4e7cdfbe327e5 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 17 Oct 2025 15:09:46 +0300 Subject: [PATCH 01/23] Initialize pallet-rate-limiting --- Cargo.lock | 11 ++ Cargo.toml | 1 + pallets/rate-limiting/Cargo.toml | 31 +++++ pallets/rate-limiting/src/lib.rs | 225 +++++++++++++++++++++++++++++++ 4 files changed, 268 insertions(+) create mode 100644 pallets/rate-limiting/Cargo.toml create mode 100644 pallets/rate-limiting/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 6cc892ad7c..79ed81c6a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10205,6 +10205,17 @@ dependencies = [ "sp-runtime", ] +[[package]] +name = "pallet-rate-limiting" +version = "0.1.0" +dependencies = [ + "frame-support", + "frame-system", + "parity-scale-codec", + "scale-info", + "sp-std", +] + [[package]] name = "pallet-recovery" version = "41.0.0" diff --git a/Cargo.toml b/Cargo.toml index ce2f3cf2ed..5a8d1742d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,6 +59,7 @@ pallet-subtensor = { path = "pallets/subtensor", default-features = false } pallet-subtensor-swap = { path = "pallets/swap", default-features = false } pallet-subtensor-swap-runtime-api = { path = "pallets/swap/runtime-api", default-features = false } pallet-subtensor-swap-rpc = { path = "pallets/swap/rpc", default-features = false } +pallet-rate-limiting = { path = "pallets/rate-limiting", default-features = false } procedural-fork = { path = "support/procedural-fork", default-features = false } safe-math = { path = "primitives/safe-math", default-features = false } share-pool = { path = "primitives/share-pool", default-features = false } diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml new file mode 100644 index 0000000000..559dbaf816 --- /dev/null +++ b/pallets/rate-limiting/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "pallet-rate-limiting" +version = "0.1.0" +edition.workspace = true + +[lints] +workspace = true + +[package.metadata.docs.rs] +targets = ["x86_64-unknown-linux-gnu"] + +[dependencies] +codec = { workspace = true, features = ["derive"] } +frame-support.workspace = true +frame-system.workspace = true +scale-info = { workspace = true, features = ["derive"] } +sp-std.workspace = true + +[features] +default = ["std"] +std = [ + "codec/std", + "frame-support/std", + "frame-system/std", + "scale-info/std", + "sp-std/std", +] +try-runtime = [ + "frame-support/try-runtime", + "frame-system/try-runtime", +] diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs new file mode 100644 index 0000000000..1b52782826 --- /dev/null +++ b/pallets/rate-limiting/src/lib.rs @@ -0,0 +1,225 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +//! Basic rate limiting pallet. + +pub use pallet::*; + +use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; +use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata}; +use scale_info::TypeInfo; +use sp_std::fmt; + +#[frame_support::pallet] +pub mod pallet { + use crate::TransactionIdentifier; + use codec::Codec; + use frame_support::{ + pallet_prelude::*, sp_runtime::traits::Saturating, traits::GetCallMetadata, + }; + use frame_system::pallet_prelude::*; + use sp_std::vec::Vec; + + /// Configuration trait for the rate limiting pallet. + #[pallet::config] + pub trait Config: frame_system::Config { + /// The overarching runtime call type. + type RuntimeCall: Parameter + + Codec + + GetCallMetadata + + IsType<::RuntimeCall>; + } + + /// Storage mapping from transaction identifier to its block-based rate limit. + #[pallet::storage] + #[pallet::getter(fn limits)] + pub type Limits = + StorageMap<_, Blake2_128Concat, TransactionIdentifier, BlockNumberFor, OptionQuery>; + + /// Tracks when a transaction was last observed. + #[pallet::storage] + pub type LastSeen = + StorageMap<_, Blake2_128Concat, TransactionIdentifier, BlockNumberFor, OptionQuery>; + + /// Events emitted by the rate limiting pallet. + #[pallet::event] + #[pallet::generate_deposit(pub(super) fn deposit_event)] + pub enum Event { + /// A rate limit was set or updated. + RateLimitSet { + /// Identifier of the affected transaction. + transaction: TransactionIdentifier, + /// The new limit expressed in blocks. + limit: BlockNumberFor, + /// Pallet name associated with the transaction. + pallet: Vec, + /// Extrinsic name associated with the transaction. + extrinsic: Vec, + }, + /// A rate limit was cleared. + RateLimitCleared { + /// Identifier of the affected transaction. + transaction: TransactionIdentifier, + /// Pallet name associated with the transaction. + pallet: Vec, + /// Extrinsic name associated with the transaction. + extrinsic: Vec, + }, + } + + /// Errors that can occur while configuring rate limits. + #[pallet::error] + pub enum Error { + /// Failed to extract the pallet and extrinsic indices from the call. + InvalidRuntimeCall, + /// Attempted to remove a limit that is not present. + MissingRateLimit, + } + + #[pallet::pallet] + #[pallet::without_storage_info] + pub struct Pallet(_); + + impl Pallet { + /// Returns `true` when the given transaction identifier passes its configured rate limit. + pub fn is_within_limit(identifier: &TransactionIdentifier) -> Result { + let Some(limit) = Limits::::get(identifier) else { + return Ok(true); + }; + + let current = frame_system::Pallet::::block_number(); + if let Some(last) = LastSeen::::get(identifier) { + let delta = current.saturating_sub(last); + if delta < limit { + return Ok(false); + } + } + + Ok(true) + } + } + + #[pallet::call] + impl Pallet { + /// Sets the rate limit, in blocks, for the given call. + /// + /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any + /// arguments embedded in the call are ignored**. + #[pallet::call_index(0)] + #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] + pub fn set_rate_limit( + origin: OriginFor, + call: Box<::RuntimeCall>, + limit: BlockNumberFor, + ) -> DispatchResult { + ensure_root(origin)?; + + let identifier = TransactionIdentifier::from_call::(call.as_ref())?; + Limits::::insert(&identifier, limit); + + let (pallet_name, extrinsic_name) = identifier.names::()?; + let pallet = Vec::from(pallet_name.as_bytes()); + let extrinsic = Vec::from(extrinsic_name.as_bytes()); + + Self::deposit_event(Event::RateLimitSet { + transaction: identifier, + limit, + pallet, + extrinsic, + }); + + Ok(()) + } + + /// Clears the rate limit for the given call, if present. + /// + /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any + /// arguments embedded in the call are ignored**. + #[pallet::call_index(1)] + #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] + pub fn clear_rate_limit( + origin: OriginFor, + call: Box<::RuntimeCall>, + ) -> DispatchResult { + ensure_root(origin)?; + + let identifier = TransactionIdentifier::from_call::(call.as_ref())?; + + let (pallet_name, extrinsic_name) = identifier.names::()?; + let pallet = Vec::from(pallet_name.as_bytes()); + let extrinsic = Vec::from(extrinsic_name.as_bytes()); + + ensure!( + Limits::::take(&identifier).is_some(), + Error::::MissingRateLimit + ); + + Self::deposit_event(Event::RateLimitCleared { + transaction: identifier, + pallet, + extrinsic, + }); + + Ok(()) + } + } +} + +/// Identifies a runtime call by pallet and extrinsic indices. +#[derive( + Clone, Copy, PartialEq, Eq, Encode, Decode, DecodeWithMemTracking, TypeInfo, MaxEncodedLen, +)] +pub struct TransactionIdentifier { + /// Index of the pallet containing the extrinsic. + pub pallet_index: u8, + /// Variant index of the extrinsic within the pallet. + pub extrinsic_index: u8, +} + +impl TransactionIdentifier { + /// Builds a new identifier from pallet/extrinsic indices. + const fn new(pallet_index: u8, extrinsic_index: u8) -> Self { + Self { + pallet_index, + extrinsic_index, + } + } + + /// Returns the pallet and extrinsic name associated with this identifier. + fn names(&self) -> Result<(&'static str, &'static str), DispatchError> + where + T: Config, + { + let modules = ::RuntimeCall::get_module_names(); + let pallet_name = modules + .get(self.pallet_index as usize) + .copied() + .ok_or(Error::::InvalidRuntimeCall)?; + let call_names = ::RuntimeCall::get_call_names(pallet_name); + let extrinsic_name = call_names + .get(self.extrinsic_index as usize) + .copied() + .ok_or(Error::::InvalidRuntimeCall)?; + Ok((pallet_name, extrinsic_name)) + } + + /// Builds an identifier from a runtime call by extracting its pallet/extrinsic indices. + fn from_call(call: &::RuntimeCall) -> Result + where + T: Config, + { + call.using_encoded(|encoded| { + let pallet_index = *encoded.get(0).ok_or(Error::::InvalidRuntimeCall)?; + let extrinsic_index = *encoded.get(1).ok_or(Error::::InvalidRuntimeCall)?; + Ok(TransactionIdentifier::new(pallet_index, extrinsic_index)) + }) + } +} + +impl fmt::Debug for TransactionIdentifier { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TransactionIdentifier") + .field("pallet_index", &self.pallet_index) + .field("extrinsic_index", &self.extrinsic_index) + .finish() + } +} From e903915b076c93a2b71a4ab695984bb7b2385700 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 17 Oct 2025 16:04:53 +0300 Subject: [PATCH 02/23] Add transaction extension for rate limiting --- Cargo.lock | 1 + pallets/rate-limiting/Cargo.toml | 2 + pallets/rate-limiting/src/lib.rs | 93 +++------------ .../src/transaction_extension.rs | 108 ++++++++++++++++++ pallets/rate-limiting/src/types.rs | 81 +++++++++++++ 5 files changed, 211 insertions(+), 74 deletions(-) create mode 100644 pallets/rate-limiting/src/transaction_extension.rs create mode 100644 pallets/rate-limiting/src/types.rs diff --git a/Cargo.lock b/Cargo.lock index 79ed81c6a7..4295ea61a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10214,6 +10214,7 @@ dependencies = [ "parity-scale-codec", "scale-info", "sp-std", + "subtensor-runtime-common", ] [[package]] diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml index 559dbaf816..76c6e18142 100644 --- a/pallets/rate-limiting/Cargo.toml +++ b/pallets/rate-limiting/Cargo.toml @@ -15,6 +15,7 @@ frame-support.workspace = true frame-system.workspace = true scale-info = { workspace = true, features = ["derive"] } sp-std.workspace = true +subtensor-runtime-common.workspace = true [features] default = ["std"] @@ -24,6 +25,7 @@ std = [ "frame-system/std", "scale-info/std", "sp-std/std", + "subtensor-runtime-common/std", ] try-runtime = [ "frame-support/try-runtime", diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 1b52782826..16a1130769 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -3,22 +3,24 @@ //! Basic rate limiting pallet. pub use pallet::*; +pub use transaction_extension::RateLimitTransactionExtension; +pub use types::{Scope, TransactionIdentifier}; -use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; -use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata}; -use scale_info::TypeInfo; -use sp_std::fmt; +mod transaction_extension; +mod types; #[frame_support::pallet] pub mod pallet { - use crate::TransactionIdentifier; use codec::Codec; + use core::fmt::Debug; use frame_support::{ pallet_prelude::*, sp_runtime::traits::Saturating, traits::GetCallMetadata, }; - use frame_system::pallet_prelude::*; + use frame_system::{ensure_root, pallet_prelude::*}; use sp_std::vec::Vec; + use crate::types::TransactionIdentifier; + /// Configuration trait for the rate limiting pallet. #[pallet::config] pub trait Config: frame_system::Config { @@ -27,9 +29,12 @@ pub mod pallet { + Codec + GetCallMetadata + IsType<::RuntimeCall>; + + /// Context type used for contextual (per-group) rate limits. + type ScopeContext: Parameter + Clone + PartialEq + Eq + Debug; } - /// Storage mapping from transaction identifier to its block-based rate limit. + /// Storage mapping from transaction identifier to its configured rate limit. #[pallet::storage] #[pallet::getter(fn limits)] pub type Limits = @@ -49,7 +54,7 @@ pub mod pallet { /// Identifier of the affected transaction. transaction: TransactionIdentifier, /// The new limit expressed in blocks. - limit: BlockNumberFor, + block_span: BlockNumberFor, /// Pallet name associated with the transaction. pallet: Vec, /// Extrinsic name associated with the transaction. @@ -82,14 +87,14 @@ pub mod pallet { impl Pallet { /// Returns `true` when the given transaction identifier passes its configured rate limit. pub fn is_within_limit(identifier: &TransactionIdentifier) -> Result { - let Some(limit) = Limits::::get(identifier) else { + let Some(block_span) = Limits::::get(identifier) else { return Ok(true); }; let current = frame_system::Pallet::::block_number(); if let Some(last) = LastSeen::::get(identifier) { let delta = current.saturating_sub(last); - if delta < limit { + if delta < block_span { return Ok(false); } } @@ -100,7 +105,7 @@ pub mod pallet { #[pallet::call] impl Pallet { - /// Sets the rate limit, in blocks, for the given call. + /// Sets the rate limit configuration for the given call. /// /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any /// arguments embedded in the call are ignored**. @@ -109,12 +114,12 @@ pub mod pallet { pub fn set_rate_limit( origin: OriginFor, call: Box<::RuntimeCall>, - limit: BlockNumberFor, + block_span: BlockNumberFor, ) -> DispatchResult { ensure_root(origin)?; let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - Limits::::insert(&identifier, limit); + Limits::::insert(&identifier, block_span); let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); @@ -122,7 +127,7 @@ pub mod pallet { Self::deposit_event(Event::RateLimitSet { transaction: identifier, - limit, + block_span, pallet, extrinsic, }); @@ -163,63 +168,3 @@ pub mod pallet { } } } - -/// Identifies a runtime call by pallet and extrinsic indices. -#[derive( - Clone, Copy, PartialEq, Eq, Encode, Decode, DecodeWithMemTracking, TypeInfo, MaxEncodedLen, -)] -pub struct TransactionIdentifier { - /// Index of the pallet containing the extrinsic. - pub pallet_index: u8, - /// Variant index of the extrinsic within the pallet. - pub extrinsic_index: u8, -} - -impl TransactionIdentifier { - /// Builds a new identifier from pallet/extrinsic indices. - const fn new(pallet_index: u8, extrinsic_index: u8) -> Self { - Self { - pallet_index, - extrinsic_index, - } - } - - /// Returns the pallet and extrinsic name associated with this identifier. - fn names(&self) -> Result<(&'static str, &'static str), DispatchError> - where - T: Config, - { - let modules = ::RuntimeCall::get_module_names(); - let pallet_name = modules - .get(self.pallet_index as usize) - .copied() - .ok_or(Error::::InvalidRuntimeCall)?; - let call_names = ::RuntimeCall::get_call_names(pallet_name); - let extrinsic_name = call_names - .get(self.extrinsic_index as usize) - .copied() - .ok_or(Error::::InvalidRuntimeCall)?; - Ok((pallet_name, extrinsic_name)) - } - - /// Builds an identifier from a runtime call by extracting its pallet/extrinsic indices. - fn from_call(call: &::RuntimeCall) -> Result - where - T: Config, - { - call.using_encoded(|encoded| { - let pallet_index = *encoded.get(0).ok_or(Error::::InvalidRuntimeCall)?; - let extrinsic_index = *encoded.get(1).ok_or(Error::::InvalidRuntimeCall)?; - Ok(TransactionIdentifier::new(pallet_index, extrinsic_index)) - }) - } -} - -impl fmt::Debug for TransactionIdentifier { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("TransactionIdentifier") - .field("pallet_index", &self.pallet_index) - .field("extrinsic_index", &self.extrinsic_index) - .finish() - } -} diff --git a/pallets/rate-limiting/src/transaction_extension.rs b/pallets/rate-limiting/src/transaction_extension.rs new file mode 100644 index 0000000000..deacab0874 --- /dev/null +++ b/pallets/rate-limiting/src/transaction_extension.rs @@ -0,0 +1,108 @@ +use codec::{Decode, DecodeWithMemTracking, Encode}; +use frame_support::{ + dispatch::{DispatchInfo, DispatchResult, PostDispatchInfo}, + pallet_prelude::Weight, + sp_runtime::{ + traits::{ + DispatchInfoOf, DispatchOriginOf, Dispatchable, Implication, TransactionExtension, + ValidateResult, + }, + transaction_validity::{ + InvalidTransaction, TransactionSource, TransactionValidityError, ValidTransaction, + }, + }, +}; +use scale_info::TypeInfo; +use sp_std::{marker::PhantomData, result::Result}; + +use crate::{Config, LastSeen, Limits, Pallet, types::TransactionIdentifier}; + +/// Identifier returned in the transaction metadata for the rate limiting extension. +const IDENTIFIER: &str = "RateLimitTransactionExtension"; + +/// Custom error code used to signal a rate limit violation. +const RATE_LIMIT_DENIED: u8 = 1; + +/// Transaction extension that enforces pallet rate limiting rules. +#[derive(Default, Encode, Decode, DecodeWithMemTracking, Clone, Eq, PartialEq, TypeInfo)] +pub struct RateLimitTransactionExtension(PhantomData); + +impl core::fmt::Debug for RateLimitTransactionExtension { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(IDENTIFIER) + } +} + +impl TransactionExtension<::RuntimeCall> for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + ::RuntimeCall: Dispatchable, +{ + const IDENTIFIER: &'static str = IDENTIFIER; + + type Implicit = (); + type Val = Option; + type Pre = Option; + + fn weight(&self, _call: &::RuntimeCall) -> Weight { + Weight::zero() + } + + fn validate( + &self, + origin: DispatchOriginOf<::RuntimeCall>, + call: &::RuntimeCall, + _info: &DispatchInfoOf<::RuntimeCall>, + _len: usize, + _self_implicit: Self::Implicit, + _inherited_implication: &impl Implication, + _source: TransactionSource, + ) -> ValidateResult::RuntimeCall> { + let identifier = match TransactionIdentifier::from_call::(call) { + Ok(identifier) => identifier, + Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)), + }; + + if Limits::::get(&identifier).is_none() { + return Ok((ValidTransaction::default(), None, origin)); + } + + let within_limit = Pallet::::is_within_limit(&identifier) + .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; + + if !within_limit { + return Err(TransactionValidityError::Invalid( + InvalidTransaction::Custom(RATE_LIMIT_DENIED), + )); + } + + Ok((ValidTransaction::default(), Some(identifier), origin)) + } + + fn prepare( + self, + val: Self::Val, + _origin: &DispatchOriginOf<::RuntimeCall>, + _call: &::RuntimeCall, + _info: &DispatchInfoOf<::RuntimeCall>, + _len: usize, + ) -> Result { + Ok(val) + } + + fn post_dispatch( + pre: Self::Pre, + _info: &DispatchInfoOf<::RuntimeCall>, + _post_info: &mut PostDispatchInfo, + _len: usize, + result: &DispatchResult, + ) -> Result<(), TransactionValidityError> { + if result.is_ok() { + if let Some(identifier) = pre { + let block_number = frame_system::Pallet::::block_number(); + LastSeen::::insert(&identifier, block_number); + } + } + Ok(()) + } +} diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs new file mode 100644 index 0000000000..88143c84fb --- /dev/null +++ b/pallets/rate-limiting/src/types.rs @@ -0,0 +1,81 @@ +use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; +use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata}; +use scale_info::TypeInfo; + +/// Defines the scope within which a rate limit applies. +#[derive( + Clone, PartialEq, Eq, Encode, Decode, DecodeWithMemTracking, TypeInfo, MaxEncodedLen, Debug, +)] +pub enum Scope { + /// Rate limit applies chain-wide. + Global, + /// Rate limit applies to a specific context (e.g., subnet). + Contextual(Context), +} + +/// Identifies a runtime call by pallet and extrinsic indices. +#[derive( + Clone, + Copy, + PartialEq, + Eq, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub struct TransactionIdentifier { + /// Pallet variant index. + pub pallet_index: u8, + /// Call variant index within the pallet. + pub extrinsic_index: u8, +} + +impl TransactionIdentifier { + /// Builds a new identifier from pallet/extrinsic indices. + pub const fn new(pallet_index: u8, extrinsic_index: u8) -> Self { + Self { + pallet_index, + extrinsic_index, + } + } + + /// Returns the pallet and extrinsic names associated with this identifier. + pub fn names(&self) -> Result<(&'static str, &'static str), DispatchError> + where + T: crate::pallet::Config, + ::RuntimeCall: GetCallMetadata, + { + let modules = ::RuntimeCall::get_module_names(); + let pallet_name = modules + .get(self.pallet_index as usize) + .copied() + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + let call_names = ::RuntimeCall::get_call_names(pallet_name); + let extrinsic_name = call_names + .get(self.extrinsic_index as usize) + .copied() + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + Ok((pallet_name, extrinsic_name)) + } + + /// Builds an identifier from a runtime call by extracting pallet/extrinsic indices. + pub fn from_call( + call: &::RuntimeCall, + ) -> Result + where + T: crate::pallet::Config, + { + call.using_encoded(|encoded| { + let pallet_index = *encoded + .get(0) + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + let extrinsic_index = *encoded + .get(1) + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + Ok(TransactionIdentifier::new(pallet_index, extrinsic_index)) + }) + } +} From ea289efdb4ad960be3c18597bed6c147bf32e5fb Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 17 Oct 2025 19:13:46 +0300 Subject: [PATCH 03/23] Implement contextual scope --- pallets/rate-limiting/src/lib.rs | 32 ++++++++++++++----- .../src/transaction_extension.rs | 12 ++++--- 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 16a1130769..5affba481c 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -12,14 +12,13 @@ mod types; #[frame_support::pallet] pub mod pallet { use codec::Codec; - use core::fmt::Debug; use frame_support::{ pallet_prelude::*, sp_runtime::traits::Saturating, traits::GetCallMetadata, }; use frame_system::{ensure_root, pallet_prelude::*}; use sp_std::vec::Vec; - use crate::types::TransactionIdentifier; + use crate::types::{Scope, TransactionIdentifier}; /// Configuration trait for the rate limiting pallet. #[pallet::config] @@ -31,7 +30,7 @@ pub mod pallet { + IsType<::RuntimeCall>; /// Context type used for contextual (per-group) rate limits. - type ScopeContext: Parameter + Clone + PartialEq + Eq + Debug; + type ScopeContext: Parameter + Clone + PartialEq + Eq; } /// Storage mapping from transaction identifier to its configured rate limit. @@ -41,9 +40,18 @@ pub mod pallet { StorageMap<_, Blake2_128Concat, TransactionIdentifier, BlockNumberFor, OptionQuery>; /// Tracks when a transaction was last observed. + /// + /// The second key is `None` for `Scope::Global` and `Some(context)` for `Scope::Contextual`. #[pallet::storage] - pub type LastSeen = - StorageMap<_, Blake2_128Concat, TransactionIdentifier, BlockNumberFor, OptionQuery>; + pub type LastSeen = StorageDoubleMap< + _, + Blake2_128Concat, + TransactionIdentifier, + Blake2_128Concat, + Option<::ScopeContext>, + BlockNumberFor, + OptionQuery, + >; /// Events emitted by the rate limiting pallet. #[pallet::event] @@ -85,14 +93,22 @@ pub mod pallet { pub struct Pallet(_); impl Pallet { - /// Returns `true` when the given transaction identifier passes its configured rate limit. - pub fn is_within_limit(identifier: &TransactionIdentifier) -> Result { + /// Returns `true` when the given transaction identifier passes its configured rate limit within the provided scope. + pub fn is_within_limit( + identifier: &TransactionIdentifier, + scope: Scope<::ScopeContext>, + ) -> Result { let Some(block_span) = Limits::::get(identifier) else { return Ok(true); }; let current = frame_system::Pallet::::block_number(); - if let Some(last) = LastSeen::::get(identifier) { + let context_key = match scope { + Scope::Global => None, + Scope::Contextual(ctx) => Some(ctx), + }; + + if let Some(last) = LastSeen::::get(identifier, context_key.clone()) { let delta = current.saturating_sub(last); if delta < block_span { return Ok(false); diff --git a/pallets/rate-limiting/src/transaction_extension.rs b/pallets/rate-limiting/src/transaction_extension.rs index deacab0874..a109d9c8c1 100644 --- a/pallets/rate-limiting/src/transaction_extension.rs +++ b/pallets/rate-limiting/src/transaction_extension.rs @@ -15,7 +15,10 @@ use frame_support::{ use scale_info::TypeInfo; use sp_std::{marker::PhantomData, result::Result}; -use crate::{Config, LastSeen, Limits, Pallet, types::TransactionIdentifier}; +use crate::{ + Config, LastSeen, Limits, Pallet, + types::{Scope, TransactionIdentifier}, +}; /// Identifier returned in the transaction metadata for the rate limiting extension. const IDENTIFIER: &str = "RateLimitTransactionExtension"; @@ -67,8 +70,9 @@ where return Ok((ValidTransaction::default(), None, origin)); } - let within_limit = Pallet::::is_within_limit(&identifier) - .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; + let within_limit = + Pallet::::is_within_limit(&identifier, Scope::::Global) + .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; if !within_limit { return Err(TransactionValidityError::Invalid( @@ -100,7 +104,7 @@ where if result.is_ok() { if let Some(identifier) = pre { let block_number = frame_system::Pallet::::block_number(); - LastSeen::::insert(&identifier, block_number); + LastSeen::::insert(&identifier, Option::::None, block_number); } } Ok(()) From ebcec75740e978111443fd47f64a97e67343266e Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 17 Oct 2025 19:38:16 +0300 Subject: [PATCH 04/23] Add context resolver --- pallets/rate-limiting/src/lib.rs | 28 +++++++++---------- ...ansaction_extension.rs => tx_extension.rs} | 23 +++++++++------ pallets/rate-limiting/src/types.rs | 14 ++++------ 3 files changed, 33 insertions(+), 32 deletions(-) rename pallets/rate-limiting/src/{transaction_extension.rs => tx_extension.rs} (82%) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 5affba481c..528fab3943 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -3,10 +3,10 @@ //! Basic rate limiting pallet. pub use pallet::*; -pub use transaction_extension::RateLimitTransactionExtension; -pub use types::{Scope, TransactionIdentifier}; +pub use tx_extension::RateLimitTransactionExtension; +pub use types::{RateLimitContextResolver, TransactionIdentifier}; -mod transaction_extension; +mod tx_extension; mod types; #[frame_support::pallet] @@ -18,7 +18,7 @@ pub mod pallet { use frame_system::{ensure_root, pallet_prelude::*}; use sp_std::vec::Vec; - use crate::types::{Scope, TransactionIdentifier}; + use crate::types::{RateLimitContextResolver, TransactionIdentifier}; /// Configuration trait for the rate limiting pallet. #[pallet::config] @@ -30,7 +30,10 @@ pub mod pallet { + IsType<::RuntimeCall>; /// Context type used for contextual (per-group) rate limits. - type ScopeContext: Parameter + Clone + PartialEq + Eq; + type LimitContext: Parameter + Clone + PartialEq + Eq; + + /// Resolves the context for a given runtime call. + type ContextResolver: RateLimitContextResolver<::RuntimeCall, Self::LimitContext>; } /// Storage mapping from transaction identifier to its configured rate limit. @@ -41,14 +44,14 @@ pub mod pallet { /// Tracks when a transaction was last observed. /// - /// The second key is `None` for `Scope::Global` and `Some(context)` for `Scope::Contextual`. + /// The second key is `None` for global limits and `Some(context)` for contextual limits. #[pallet::storage] pub type LastSeen = StorageDoubleMap< _, Blake2_128Concat, TransactionIdentifier, Blake2_128Concat, - Option<::ScopeContext>, + Option<::LimitContext>, BlockNumberFor, OptionQuery, >; @@ -93,22 +96,19 @@ pub mod pallet { pub struct Pallet(_); impl Pallet { - /// Returns `true` when the given transaction identifier passes its configured rate limit within the provided scope. + /// Returns `true` when the given transaction identifier passes its configured rate limit + /// within the provided context. pub fn is_within_limit( identifier: &TransactionIdentifier, - scope: Scope<::ScopeContext>, + context: Option<::LimitContext>, ) -> Result { let Some(block_span) = Limits::::get(identifier) else { return Ok(true); }; let current = frame_system::Pallet::::block_number(); - let context_key = match scope { - Scope::Global => None, - Scope::Contextual(ctx) => Some(ctx), - }; - if let Some(last) = LastSeen::::get(identifier, context_key.clone()) { + if let Some(last) = LastSeen::::get(identifier, &context) { let delta = current.saturating_sub(last); if delta < block_span { return Ok(false); diff --git a/pallets/rate-limiting/src/transaction_extension.rs b/pallets/rate-limiting/src/tx_extension.rs similarity index 82% rename from pallets/rate-limiting/src/transaction_extension.rs rename to pallets/rate-limiting/src/tx_extension.rs index a109d9c8c1..b7e55e1222 100644 --- a/pallets/rate-limiting/src/transaction_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -17,7 +17,7 @@ use sp_std::{marker::PhantomData, result::Result}; use crate::{ Config, LastSeen, Limits, Pallet, - types::{Scope, TransactionIdentifier}, + types::{RateLimitContextResolver, TransactionIdentifier}, }; /// Identifier returned in the transaction metadata for the rate limiting extension. @@ -44,8 +44,8 @@ where const IDENTIFIER: &'static str = IDENTIFIER; type Implicit = (); - type Val = Option; - type Pre = Option; + type Val = Option<(TransactionIdentifier, Option)>; + type Pre = Option<(TransactionIdentifier, Option)>; fn weight(&self, _call: &::RuntimeCall) -> Weight { Weight::zero() @@ -70,9 +70,10 @@ where return Ok((ValidTransaction::default(), None, origin)); } - let within_limit = - Pallet::::is_within_limit(&identifier, Scope::::Global) - .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; + let context = ::ContextResolver::context(call); + + let within_limit = Pallet::::is_within_limit(&identifier, context.clone()) + .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; if !within_limit { return Err(TransactionValidityError::Invalid( @@ -80,7 +81,11 @@ where )); } - Ok((ValidTransaction::default(), Some(identifier), origin)) + Ok(( + ValidTransaction::default(), + Some((identifier, context)), + origin, + )) } fn prepare( @@ -102,9 +107,9 @@ where result: &DispatchResult, ) -> Result<(), TransactionValidityError> { if result.is_ok() { - if let Some(identifier) = pre { + if let Some((identifier, context)) = pre { let block_number = frame_system::Pallet::::block_number(); - LastSeen::::insert(&identifier, Option::::None, block_number); + LastSeen::::insert(&identifier, context, block_number); } } Ok(()) diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index 88143c84fb..ee62ed894e 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -2,15 +2,11 @@ use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata}; use scale_info::TypeInfo; -/// Defines the scope within which a rate limit applies. -#[derive( - Clone, PartialEq, Eq, Encode, Decode, DecodeWithMemTracking, TypeInfo, MaxEncodedLen, Debug, -)] -pub enum Scope { - /// Rate limit applies chain-wide. - Global, - /// Rate limit applies to a specific context (e.g., subnet). - Contextual(Context), +/// Resolves the optional context within which a rate limit applies. +pub trait RateLimitContextResolver { + /// Returns `Some(context)` when the limit should be applied per-context, or `None` for global + /// limits. + fn context(call: &Call) -> Option; } /// Identifies a runtime call by pallet and extrinsic indices. From 56333c9fd4fbc98b9b618eab74cd5a2938068f86 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Tue, 21 Oct 2025 15:41:52 +0300 Subject: [PATCH 05/23] Add default rate limit --- pallets/rate-limiting/src/lib.rs | 59 +++++++++++++++++++---- pallets/rate-limiting/src/tx_extension.rs | 15 ++++-- pallets/rate-limiting/src/types.rs | 20 ++++++++ 3 files changed, 81 insertions(+), 13 deletions(-) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 528fab3943..761cbed270 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -4,7 +4,7 @@ pub use pallet::*; pub use tx_extension::RateLimitTransactionExtension; -pub use types::{RateLimitContextResolver, TransactionIdentifier}; +pub use types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; mod tx_extension; mod types; @@ -18,7 +18,7 @@ pub mod pallet { use frame_system::{ensure_root, pallet_prelude::*}; use sp_std::vec::Vec; - use crate::types::{RateLimitContextResolver, TransactionIdentifier}; + use crate::types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; /// Configuration trait for the rate limiting pallet. #[pallet::config] @@ -39,8 +39,13 @@ pub mod pallet { /// Storage mapping from transaction identifier to its configured rate limit. #[pallet::storage] #[pallet::getter(fn limits)] - pub type Limits = - StorageMap<_, Blake2_128Concat, TransactionIdentifier, BlockNumberFor, OptionQuery>; + pub type Limits = StorageMap< + _, + Blake2_128Concat, + TransactionIdentifier, + RateLimit>, + OptionQuery, + >; /// Tracks when a transaction was last observed. /// @@ -56,6 +61,11 @@ pub mod pallet { OptionQuery, >; + /// Default block span applied when an extrinsic uses the default rate limit. + #[pallet::storage] + #[pallet::getter(fn default_limit)] + pub type DefaultLimit = StorageValue<_, BlockNumberFor, ValueQuery>; + /// Events emitted by the rate limiting pallet. #[pallet::event] #[pallet::generate_deposit(pub(super) fn deposit_event)] @@ -64,8 +74,8 @@ pub mod pallet { RateLimitSet { /// Identifier of the affected transaction. transaction: TransactionIdentifier, - /// The new limit expressed in blocks. - block_span: BlockNumberFor, + /// The new limit configuration applied to the transaction. + limit: RateLimit>, /// Pallet name associated with the transaction. pallet: Vec, /// Extrinsic name associated with the transaction. @@ -80,6 +90,11 @@ pub mod pallet { /// Extrinsic name associated with the transaction. extrinsic: Vec, }, + /// The default rate limit was set or updated. + DefaultRateLimitSet { + /// The new default limit expressed in blocks. + block_span: BlockNumberFor, + }, } /// Errors that can occur while configuring rate limits. @@ -102,7 +117,7 @@ pub mod pallet { identifier: &TransactionIdentifier, context: Option<::LimitContext>, ) -> Result { - let Some(block_span) = Limits::::get(identifier) else { + let Some(block_span) = Self::resolved_limit(identifier) else { return Ok(true); }; @@ -117,6 +132,14 @@ pub mod pallet { Ok(true) } + + fn resolved_limit(identifier: &TransactionIdentifier) -> Option> { + let limit = Limits::::get(identifier)?; + Some(match limit { + RateLimit::Default => DefaultLimit::::get(), + RateLimit::Exact(block_span) => block_span, + }) + } } #[pallet::call] @@ -130,12 +153,12 @@ pub mod pallet { pub fn set_rate_limit( origin: OriginFor, call: Box<::RuntimeCall>, - block_span: BlockNumberFor, + limit: RateLimit>, ) -> DispatchResult { ensure_root(origin)?; let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - Limits::::insert(&identifier, block_span); + Limits::::insert(&identifier, limit); let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); @@ -143,7 +166,7 @@ pub mod pallet { Self::deposit_event(Event::RateLimitSet { transaction: identifier, - block_span, + limit, pallet, extrinsic, }); @@ -182,5 +205,21 @@ pub mod pallet { Ok(()) } + + /// Sets the default rate limit in blocks applied to calls configured to use it. + #[pallet::call_index(2)] + #[pallet::weight(T::DbWeight::get().writes(1))] + pub fn set_default_rate_limit( + origin: OriginFor, + block_span: BlockNumberFor, + ) -> DispatchResult { + ensure_root(origin)?; + + DefaultLimit::::put(block_span); + + Self::deposit_event(Event::DefaultRateLimitSet { block_span }); + + Ok(()) + } } } diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index b7e55e1222..097b23b961 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -5,7 +5,7 @@ use frame_support::{ sp_runtime::{ traits::{ DispatchInfoOf, DispatchOriginOf, Dispatchable, Implication, TransactionExtension, - ValidateResult, + ValidateResult, Zero, }, transaction_validity::{ InvalidTransaction, TransactionSource, TransactionValidityError, ValidTransaction, @@ -17,7 +17,7 @@ use sp_std::{marker::PhantomData, result::Result}; use crate::{ Config, LastSeen, Limits, Pallet, - types::{RateLimitContextResolver, TransactionIdentifier}, + types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}, }; /// Identifier returned in the transaction metadata for the rate limiting extension. @@ -66,7 +66,16 @@ where Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)), }; - if Limits::::get(&identifier).is_none() { + let Some(limit) = Limits::::get(&identifier) else { + return Ok((ValidTransaction::default(), None, origin)); + }; + + let block_span = match limit { + RateLimit::Default => Pallet::::default_limit(), + RateLimit::Exact(block_span) => block_span, + }; + + if block_span.is_zero() { return Ok((ValidTransaction::default(), None, origin)); } diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index ee62ed894e..124dbbe3ab 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -75,3 +75,23 @@ impl TransactionIdentifier { }) } } + +/// Configuration value for a rate limit. +#[derive( + Clone, + Copy, + PartialEq, + Eq, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub enum RateLimit { + /// Use the pallet-level default rate limit. + Default, + /// Apply an exact rate limit measured in blocks. + Exact(BlockNumber), +} From 1b6b03cbe77a551a7500f29c9f38fc69c28a8fd1 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Tue, 21 Oct 2025 16:26:09 +0300 Subject: [PATCH 06/23] Add genesis config for pallet-rate-limiting --- Cargo.lock | 1 + pallets/rate-limiting/Cargo.toml | 2 ++ pallets/rate-limiting/src/lib.rs | 31 +++++++++++++++++++++++++++++- pallets/rate-limiting/src/types.rs | 2 ++ 4 files changed, 35 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 4295ea61a8..e315f31c63 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10213,6 +10213,7 @@ dependencies = [ "frame-system", "parity-scale-codec", "scale-info", + "serde", "sp-std", "subtensor-runtime-common", ] diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml index 76c6e18142..24620e2d54 100644 --- a/pallets/rate-limiting/Cargo.toml +++ b/pallets/rate-limiting/Cargo.toml @@ -16,6 +16,7 @@ frame-system.workspace = true scale-info = { workspace = true, features = ["derive"] } sp-std.workspace = true subtensor-runtime-common.workspace = true +serde = { workspace = true, features = ["derive"], optional = true } [features] default = ["std"] @@ -26,6 +27,7 @@ std = [ "scale-info/std", "sp-std/std", "subtensor-runtime-common/std", + "serde", ] try-runtime = [ "frame-support/try-runtime", diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 761cbed270..0b132c3dda 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -13,7 +13,9 @@ mod types; pub mod pallet { use codec::Codec; use frame_support::{ - pallet_prelude::*, sp_runtime::traits::Saturating, traits::GetCallMetadata, + pallet_prelude::*, + sp_runtime::traits::{Saturating, Zero}, + traits::{BuildGenesisConfig, GetCallMetadata}, }; use frame_system::{ensure_root, pallet_prelude::*}; use sp_std::vec::Vec; @@ -106,6 +108,33 @@ pub mod pallet { MissingRateLimit, } + #[pallet::genesis_config] + pub struct GenesisConfig { + pub default_limit: BlockNumberFor, + pub limits: Vec<(TransactionIdentifier, RateLimit>)>, + } + + #[cfg(feature = "std")] + impl Default for GenesisConfig { + fn default() -> Self { + Self { + default_limit: Zero::zero(), + limits: Vec::new(), + } + } + } + + #[pallet::genesis_build] + impl BuildGenesisConfig for GenesisConfig { + fn build(&self) { + DefaultLimit::::put(self.default_limit); + + for (identifier, limit) in &self.limits { + Limits::::insert(identifier, limit.clone()); + } + } + } + #[pallet::pallet] #[pallet::without_storage_info] pub struct Pallet(_); diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index 124dbbe3ab..4e00053ec7 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -10,6 +10,7 @@ pub trait RateLimitContextResolver { } /// Identifies a runtime call by pallet and extrinsic indices. +#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] #[derive( Clone, Copy, @@ -77,6 +78,7 @@ impl TransactionIdentifier { } /// Configuration value for a rate limit. +#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] #[derive( Clone, Copy, From 220c9052ad73a8bf7906ec6aa2c9469bdeca33f8 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Tue, 21 Oct 2025 19:26:54 +0300 Subject: [PATCH 07/23] Add rpc method to fetch rate limit --- Cargo.lock | 25 ++++++ Cargo.toml | 4 + pallets/rate-limiting/rpc/Cargo.toml | 22 ++++++ pallets/rate-limiting/rpc/src/lib.rs | 82 ++++++++++++++++++++ pallets/rate-limiting/runtime-api/Cargo.toml | 26 +++++++ pallets/rate-limiting/runtime-api/src/lib.rs | 24 ++++++ pallets/rate-limiting/src/lib.rs | 37 ++++++++- pallets/rate-limiting/src/tx_extension.rs | 2 +- 8 files changed, 218 insertions(+), 4 deletions(-) create mode 100644 pallets/rate-limiting/rpc/Cargo.toml create mode 100644 pallets/rate-limiting/rpc/src/lib.rs create mode 100644 pallets/rate-limiting/runtime-api/Cargo.toml create mode 100644 pallets/rate-limiting/runtime-api/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index e315f31c63..6e9eba1a5a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10218,6 +10218,31 @@ dependencies = [ "subtensor-runtime-common", ] +[[package]] +name = "pallet-rate-limiting-rpc" +version = "0.1.0" +dependencies = [ + "jsonrpsee", + "pallet-rate-limiting-runtime-api", + "sp-api", + "sp-blockchain", + "sp-runtime", + "subtensor-runtime-common", +] + +[[package]] +name = "pallet-rate-limiting-runtime-api" +version = "0.1.0" +dependencies = [ + "pallet-rate-limiting", + "parity-scale-codec", + "scale-info", + "serde", + "sp-api", + "sp-std", + "subtensor-runtime-common", +] + [[package]] name = "pallet-recovery" version = "41.0.0" diff --git a/Cargo.toml b/Cargo.toml index 5a8d1742d9..ed71cc537e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,8 @@ members = [ "common", "node", "pallets/*", + "pallets/rate-limiting/runtime-api", + "pallets/rate-limiting/rpc", "precompiles", "primitives/*", "runtime", @@ -60,6 +62,8 @@ pallet-subtensor-swap = { path = "pallets/swap", default-features = false } pallet-subtensor-swap-runtime-api = { path = "pallets/swap/runtime-api", default-features = false } pallet-subtensor-swap-rpc = { path = "pallets/swap/rpc", default-features = false } pallet-rate-limiting = { path = "pallets/rate-limiting", default-features = false } +pallet-rate-limiting-runtime-api = { path = "pallets/rate-limiting/runtime-api", default-features = false } +pallet-rate-limiting-rpc = { path = "pallets/rate-limiting/rpc", default-features = false } procedural-fork = { path = "support/procedural-fork", default-features = false } safe-math = { path = "primitives/safe-math", default-features = false } share-pool = { path = "primitives/share-pool", default-features = false } diff --git a/pallets/rate-limiting/rpc/Cargo.toml b/pallets/rate-limiting/rpc/Cargo.toml new file mode 100644 index 0000000000..d5bf689e8b --- /dev/null +++ b/pallets/rate-limiting/rpc/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "pallet-rate-limiting-rpc" +version = "0.1.0" +description = "RPC interface for the rate limiting pallet" +edition.workspace = true + +[dependencies] +jsonrpsee = { workspace = true, features = ["client-core", "server", "macros"] } +sp-api.workspace = true +sp-blockchain.workspace = true +sp-runtime.workspace = true +pallet-rate-limiting-runtime-api.workspace = true +subtensor-runtime-common = { workspace = true, default-features = false } + +[features] +default = ["std"] +std = [ + "sp-api/std", + "sp-runtime/std", + "pallet-rate-limiting-runtime-api/std", + "subtensor-runtime-common/std", +] diff --git a/pallets/rate-limiting/rpc/src/lib.rs b/pallets/rate-limiting/rpc/src/lib.rs new file mode 100644 index 0000000000..ca7452a7a0 --- /dev/null +++ b/pallets/rate-limiting/rpc/src/lib.rs @@ -0,0 +1,82 @@ +//! RPC interface for the rate limiting pallet. + +use jsonrpsee::{ + core::RpcResult, + proc_macros::rpc, + types::{ErrorObjectOwned, error::ErrorObject}, +}; +use sp_api::ProvideRuntimeApi; +use sp_blockchain::HeaderBackend; +use sp_runtime::traits::Block as BlockT; +use std::sync::Arc; + +pub use pallet_rate_limiting_runtime_api::{RateLimitRpcResponse, RateLimitingRuntimeApi}; + +#[rpc(client, server)] +pub trait RateLimitingRpcApi { + #[method(name = "rateLimiting_getRateLimit")] + fn get_rate_limit( + &self, + pallet: Vec, + extrinsic: Vec, + at: Option, + ) -> RpcResult>; +} + +/// Error type of this RPC api. +pub enum Error { + /// The call to runtime failed. + RuntimeError(String), +} + +impl From for ErrorObjectOwned { + fn from(e: Error) -> Self { + match e { + Error::RuntimeError(e) => ErrorObject::owned(1, e, None::<()>), + } + } +} + +impl From for i32 { + fn from(e: Error) -> i32 { + match e { + Error::RuntimeError(_) => 1, + } + } +} + +/// RPC implementation for the rate limiting pallet. +pub struct RateLimiting { + client: Arc, + _marker: std::marker::PhantomData, +} + +impl RateLimiting { + /// Creates a new instance of the rate limiting RPC helper. + pub fn new(client: Arc) -> Self { + Self { + client, + _marker: Default::default(), + } + } +} + +impl RateLimitingRpcApiServer<::Hash> for RateLimiting +where + Block: BlockT, + C: ProvideRuntimeApi + HeaderBackend + Send + Sync + 'static, + C::Api: RateLimitingRuntimeApi, +{ + fn get_rate_limit( + &self, + pallet: Vec, + extrinsic: Vec, + at: Option<::Hash>, + ) -> RpcResult> { + let api = self.client.runtime_api(); + let at = at.unwrap_or_else(|| self.client.info().best_hash); + + api.get_rate_limit(at, pallet, extrinsic) + .map_err(|e| Error::RuntimeError(format!("Unable to fetch rate limit: {e:?}")).into()) + } +} diff --git a/pallets/rate-limiting/runtime-api/Cargo.toml b/pallets/rate-limiting/runtime-api/Cargo.toml new file mode 100644 index 0000000000..2847d865dd --- /dev/null +++ b/pallets/rate-limiting/runtime-api/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "pallet-rate-limiting-runtime-api" +version = "0.1.0" +description = "Runtime API for the rate limiting pallet" +edition.workspace = true + +[dependencies] +codec = { workspace = true, features = ["derive"] } +scale-info = { workspace = true, features = ["derive"] } +sp-api.workspace = true +sp-std.workspace = true +pallet-rate-limiting.workspace = true +subtensor-runtime-common = { workspace = true, default-features = false } +serde = { workspace = true, features = ["derive"], optional = true } + +[features] +default = ["std"] +std = [ + "codec/std", + "scale-info/std", + "sp-api/std", + "sp-std/std", + "pallet-rate-limiting/std", + "subtensor-runtime-common/std", + "serde", +] diff --git a/pallets/rate-limiting/runtime-api/src/lib.rs b/pallets/rate-limiting/runtime-api/src/lib.rs new file mode 100644 index 0000000000..1a32c094ea --- /dev/null +++ b/pallets/rate-limiting/runtime-api/src/lib.rs @@ -0,0 +1,24 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +use codec::{Decode, Encode}; +use pallet_rate_limiting::RateLimit; +use scale_info::TypeInfo; +use sp_std::vec::Vec; +use subtensor_runtime_common::BlockNumber; + +#[cfg(feature = "std")] +use serde::{Deserialize, Serialize}; + +#[cfg_attr(feature = "std", derive(Serialize, Deserialize))] +#[derive(Clone, Debug, Decode, Encode, Eq, PartialEq, TypeInfo)] +pub struct RateLimitRpcResponse { + pub limit: Option>, + pub default_limit: BlockNumber, + pub resolved: Option, +} + +sp_api::decl_runtime_apis! { + pub trait RateLimitingRuntimeApi { + fn get_rate_limit(pallet: Vec, extrinsic: Vec) -> Option; + } +} diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 0b132c3dda..a07d75d618 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -18,7 +18,7 @@ pub mod pallet { traits::{BuildGenesisConfig, GetCallMetadata}, }; use frame_system::{ensure_root, pallet_prelude::*}; - use sp_std::vec::Vec; + use sp_std::{convert::TryFrom, vec::Vec}; use crate::types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; @@ -144,7 +144,7 @@ pub mod pallet { /// within the provided context. pub fn is_within_limit( identifier: &TransactionIdentifier, - context: Option<::LimitContext>, + context: &Option<::LimitContext>, ) -> Result { let Some(block_span) = Self::resolved_limit(identifier) else { return Ok(true); @@ -152,7 +152,7 @@ pub mod pallet { let current = frame_system::Pallet::::block_number(); - if let Some(last) = LastSeen::::get(identifier, &context) { + if let Some(last) = LastSeen::::get(identifier, context) { let delta = current.saturating_sub(last); if delta < block_span { return Ok(false); @@ -169,6 +169,37 @@ pub mod pallet { RateLimit::Exact(block_span) => block_span, }) } + + /// Returns the configured limit for the specified pallet/extrinsic names, if any. + pub fn limit_for_call_names( + pallet_name: &str, + extrinsic_name: &str, + ) -> Option>> { + let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; + Limits::::get(&identifier) + } + + /// Returns the resolved block span for the specified pallet/extrinsic names, if any. + pub fn resolved_limit_for_call_names( + pallet_name: &str, + extrinsic_name: &str, + ) -> Option> { + let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; + Self::resolved_limit(&identifier) + } + + fn identifier_for_call_names( + pallet_name: &str, + extrinsic_name: &str, + ) -> Option { + let modules = ::RuntimeCall::get_module_names(); + let pallet_pos = modules.iter().position(|name| *name == pallet_name)?; + let call_names = ::RuntimeCall::get_call_names(pallet_name); + let extrinsic_pos = call_names.iter().position(|name| *name == extrinsic_name)?; + let pallet_index = u8::try_from(pallet_pos).ok()?; + let extrinsic_index = u8::try_from(extrinsic_pos).ok()?; + Some(TransactionIdentifier::new(pallet_index, extrinsic_index)) + } } #[pallet::call] diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index 097b23b961..6ccfe8160b 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -81,7 +81,7 @@ where let context = ::ContextResolver::context(call); - let within_limit = Pallet::::is_within_limit(&identifier, context.clone()) + let within_limit = Pallet::::is_within_limit(&identifier, &context) .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; if !within_limit { From 69151b1b9ed8cba06e03f785274cf276dd74bad8 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Mon, 27 Oct 2025 17:11:18 +0300 Subject: [PATCH 08/23] Add tests to pallet-rate-limiting --- Cargo.lock | 3 + pallets/rate-limiting/Cargo.toml | 5 + pallets/rate-limiting/src/lib.rs | 6 + pallets/rate-limiting/src/mock.rs | 90 +++++++ pallets/rate-limiting/src/tests.rs | 281 ++++++++++++++++++++++ pallets/rate-limiting/src/tx_extension.rs | 181 ++++++++++++++ 6 files changed, 566 insertions(+) create mode 100644 pallets/rate-limiting/src/mock.rs create mode 100644 pallets/rate-limiting/src/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 6e9eba1a5a..d116651d33 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10214,6 +10214,9 @@ dependencies = [ "parity-scale-codec", "scale-info", "serde", + "sp-core", + "sp-io", + "sp-runtime", "sp-std", "subtensor-runtime-common", ] diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml index 24620e2d54..b24ff40ea5 100644 --- a/pallets/rate-limiting/Cargo.toml +++ b/pallets/rate-limiting/Cargo.toml @@ -18,6 +18,11 @@ sp-std.workspace = true subtensor-runtime-common.workspace = true serde = { workspace = true, features = ["derive"], optional = true } +[dev-dependencies] +sp-core.workspace = true +sp-io.workspace = true +sp-runtime.workspace = true + [features] default = ["std"] std = [ diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index a07d75d618..2ac9c76114 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -9,6 +9,12 @@ pub use types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; mod tx_extension; mod types; +#[cfg(test)] +mod mock; + +#[cfg(test)] +mod tests; + #[frame_support::pallet] pub mod pallet { use codec::Codec; diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs new file mode 100644 index 0000000000..d7b9fde20b --- /dev/null +++ b/pallets/rate-limiting/src/mock.rs @@ -0,0 +1,90 @@ +#![allow(dead_code)] + +use frame_support::{ + derive_impl, + sp_runtime::{ + BuildStorage, + traits::{BlakeTwo256, IdentityLookup}, + }, + traits::{ConstU16, ConstU32, ConstU64, Everything}, +}; +use sp_core::H256; +use sp_io::TestExternalities; + +use crate as pallet_rate_limiting; +use crate::TransactionIdentifier; + +pub type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic; +pub type Block = frame_system::mocking::MockBlock; + +frame_support::construct_runtime!( + pub enum Test { + System: frame_system, + RateLimiting: pallet_rate_limiting, + } +); + +#[derive_impl(frame_system::config_preludes::TestDefaultConfig)] +impl frame_system::Config for Test { + type BaseCallFilter = Everything; + type BlockWeights = (); + type BlockLength = (); + type DbWeight = (); + type RuntimeOrigin = RuntimeOrigin; + type RuntimeCall = RuntimeCall; + type Nonce = u64; + type Hash = H256; + type Hashing = BlakeTwo256; + type AccountId = u64; + type Lookup = IdentityLookup; + type RuntimeEvent = RuntimeEvent; + type BlockHashCount = ConstU64<250>; + type Version = (); + type PalletInfo = PalletInfo; + type AccountData = (); + type OnNewAccount = (); + type OnKilledAccount = (); + type SystemWeightInfo = (); + type SS58Prefix = ConstU16<42>; + type OnSetCode = (); + type MaxConsumers = ConstU32<16>; + type Block = Block; +} + +pub type LimitContext = u16; + +pub struct TestContextResolver; + +impl pallet_rate_limiting::RateLimitContextResolver + for TestContextResolver +{ + fn context(_call: &RuntimeCall) -> Option { + None + } +} + +impl pallet_rate_limiting::Config for Test { + type RuntimeCall = RuntimeCall; + type LimitContext = LimitContext; + type ContextResolver = TestContextResolver; +} + +pub type RateLimitingCall = crate::Call; + +pub fn new_test_ext() -> TestExternalities { + let storage = frame_system::GenesisConfig::::default() + .build_storage() + .expect("genesis build succeeds"); + + let mut ext = TestExternalities::new(storage); + ext.execute_with(|| System::set_block_number(1)); + ext +} + +pub(crate) fn identifier_for(call: &RuntimeCall) -> TransactionIdentifier { + TransactionIdentifier::from_call::(call).expect("identifier for call") +} + +pub(crate) fn pop_last_event() -> RuntimeEvent { + System::events().pop().expect("event expected").event +} diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs new file mode 100644 index 0000000000..4eb07ad926 --- /dev/null +++ b/pallets/rate-limiting/src/tests.rs @@ -0,0 +1,281 @@ +use frame_support::{assert_noop, assert_ok, error::BadOrigin}; +use sp_runtime::DispatchError; + +use crate::{ + DefaultLimit, LastSeen, Limits, RateLimit, mock::*, pallet::Error, types::TransactionIdentifier, +}; + +#[test] +fn limit_for_call_names_returns_none_if_not_set() { + new_test_ext().execute_with(|| { + assert!( + RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit").is_none() + ); + }); +} + +#[test] +fn limit_for_call_names_returns_stored_limit() { + new_test_ext().execute_with(|| { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::Exact(7)); + + let fetched = RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit") + .expect("limit should exist"); + assert_eq!(fetched, RateLimit::Exact(7)); + }); +} + +#[test] +fn resolved_limit_for_call_names_resolves_default_value() { + new_test_ext().execute_with(|| { + DefaultLimit::::put(3); + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::Default); + + let resolved = + RateLimiting::resolved_limit_for_call_names("RateLimiting", "set_default_rate_limit") + .expect("resolved limit"); + assert_eq!(resolved, 3); + }); +} + +#[test] +fn resolved_limit_for_call_names_returns_none_when_unset() { + new_test_ext().execute_with(|| { + assert!( + RateLimiting::resolved_limit_for_call_names("RateLimiting", "set_default_rate_limit") + .is_none() + ); + }); +} + +#[test] +fn is_within_limit_is_true_when_no_limit() { + new_test_ext().execute_with(|| { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&call); + + let result = RateLimiting::is_within_limit(&identifier, &None); + assert_eq!(result.expect("no error expected"), true); + }); +} + +#[test] +fn is_within_limit_false_when_rate_limited() { + new_test_ext().execute_with(|| { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::Exact(5)); + LastSeen::::insert(identifier, Some(1 as LimitContext), 9); + + System::set_block_number(13); + + let within = RateLimiting::is_within_limit(&identifier, &Some(1 as LimitContext)) + .expect("call succeeds"); + assert!(!within); + }); +} + +#[test] +fn is_within_limit_true_after_required_span() { + new_test_ext().execute_with(|| { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::Exact(5)); + LastSeen::::insert(identifier, Some(2 as LimitContext), 10); + + System::set_block_number(20); + + let within = RateLimiting::is_within_limit(&identifier, &Some(2 as LimitContext)) + .expect("call succeeds"); + assert!(within); + }); +} + +#[test] +fn transaction_identifier_from_call_matches_expected_indices() { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + + let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); + + // System is the first pallet in the mock runtime, RateLimiting is second. + assert_eq!(identifier.pallet_index, 1); + // set_default_rate_limit has call_index 2. + assert_eq!(identifier.extrinsic_index, 2); +} + +#[test] +fn transaction_identifier_names_matches_call_metadata() { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); + + let (pallet, extrinsic) = identifier.names::().expect("call metadata"); + assert_eq!(pallet, "RateLimiting"); + assert_eq!(extrinsic, "set_default_rate_limit"); +} + +#[test] +fn transaction_identifier_names_error_for_unknown_indices() { + let identifier = TransactionIdentifier::new(99, 0); + + let err = identifier.names::().expect_err("should fail"); + let expected: DispatchError = Error::::InvalidRuntimeCall.into(); + assert_eq!(err, expected); +} + +#[test] +fn set_rate_limit_updates_storage_and_emits_event() { + new_test_ext().execute_with(|| { + System::reset_events(); + + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let limit = RateLimit::Exact(9); + + assert_ok!(RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + Box::new(target_call.clone()), + limit + )); + + let identifier = identifier_for(&target_call); + assert_eq!(Limits::::get(identifier), Some(limit)); + + match pop_last_event() { + RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitSet { + transaction, + limit: emitted_limit, + pallet, + extrinsic, + }) => { + assert_eq!(transaction, identifier); + assert_eq!(emitted_limit, limit); + assert_eq!(pallet, b"RateLimiting".to_vec()); + assert_eq!(extrinsic, b"set_default_rate_limit".to_vec()); + } + other => panic!("unexpected event: {:?}", other), + } + }); +} + +#[test] +fn set_rate_limit_requires_root() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + + assert_noop!( + RateLimiting::set_rate_limit( + RuntimeOrigin::signed(1), + Box::new(target_call), + RateLimit::Exact(1) + ), + BadOrigin + ); + }); +} + +#[test] +fn set_rate_limit_accepts_default_variant() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + + assert_ok!(RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + Box::new(target_call.clone()), + RateLimit::Default + )); + + let identifier = identifier_for(&target_call); + assert_eq!(Limits::::get(identifier), Some(RateLimit::Default)); + }); +} + +#[test] +fn clear_rate_limit_removes_entry_and_emits_event() { + new_test_ext().execute_with(|| { + System::reset_events(); + + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&target_call); + Limits::::insert(identifier, RateLimit::Exact(4)); + + assert_ok!(RateLimiting::clear_rate_limit( + RuntimeOrigin::root(), + Box::new(target_call.clone()) + )); + + assert_eq!(Limits::::get(identifier), None); + + match pop_last_event() { + RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { + transaction, + pallet, + extrinsic, + }) => { + assert_eq!(transaction, identifier); + assert_eq!(pallet, b"RateLimiting".to_vec()); + assert_eq!(extrinsic, b"set_default_rate_limit".to_vec()); + } + other => panic!("unexpected event: {:?}", other), + } + }); +} + +#[test] +fn clear_rate_limit_fails_when_missing() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + + assert_noop!( + RateLimiting::clear_rate_limit(RuntimeOrigin::root(), Box::new(target_call)), + Error::::MissingRateLimit + ); + }); +} + +#[test] +fn set_default_rate_limit_updates_storage_and_emits_event() { + new_test_ext().execute_with(|| { + System::reset_events(); + + assert_ok!(RateLimiting::set_default_rate_limit( + RuntimeOrigin::root(), + 42 + )); + + assert_eq!(DefaultLimit::::get(), 42); + + match pop_last_event() { + RuntimeEvent::RateLimiting(crate::pallet::Event::DefaultRateLimitSet { + block_span, + }) => { + assert_eq!(block_span, 42); + } + other => panic!("unexpected event: {:?}", other), + } + }); +} + +#[test] +fn set_default_rate_limit_requires_root() { + new_test_ext().execute_with(|| { + assert_noop!( + RateLimiting::set_default_rate_limit(RuntimeOrigin::signed(1), 5), + BadOrigin + ); + }); +} diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index 6ccfe8160b..f06e72975e 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -124,3 +124,184 @@ where Ok(()) } } + +#[cfg(test)] +mod tests { + use codec::Encode; + use frame_support::dispatch::{GetDispatchInfo, PostDispatchInfo}; + use sp_runtime::{ + traits::{TransactionExtension, TxBaseImplication}, + transaction_validity::{InvalidTransaction, TransactionSource, TransactionValidityError}, + }; + + use crate::{LastSeen, Limits, RateLimit, types::TransactionIdentifier}; + + use super::*; + use crate::mock::*; + + fn remark_call() -> RuntimeCall { + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }) + } + + fn new_tx_extension() -> RateLimitTransactionExtension { + RateLimitTransactionExtension(Default::default()) + } + + fn validate_with_tx_extension( + extension: &RateLimitTransactionExtension, + call: &RuntimeCall, + ) -> Result< + ( + sp_runtime::transaction_validity::ValidTransaction, + Option<(TransactionIdentifier, Option)>, + RuntimeOrigin, + ), + TransactionValidityError, + > { + let info = call.get_dispatch_info(); + let len = call.encode().len(); + extension.validate( + RuntimeOrigin::signed(42), + call, + &info, + len, + (), + &TxBaseImplication(()), + TransactionSource::External, + ) + } + + #[test] + fn tx_extension_allows_calls_without_limit() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + + let (_valid, val, _origin) = + validate_with_tx_extension(&extension, &call).expect("valid"); + assert!(val.is_none()); + + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let origin_for_prepare = RuntimeOrigin::signed(42); + let pre = extension + .clone() + .prepare(val.clone(), &origin_for_prepare, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + let identifier = identifier_for(&call); + assert_eq!( + LastSeen::::get(identifier, None::), + None + ); + }); + } + + #[test] + fn tx_extension_records_last_seen_for_successful_call() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::Exact(5)); + + System::set_block_number(10); + + let (_valid, val, _) = validate_with_tx_extension(&extension, &call).expect("valid"); + assert!(val.is_some()); + + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let origin_for_prepare = RuntimeOrigin::signed(42); + let pre = extension + .clone() + .prepare(val.clone(), &origin_for_prepare, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + assert_eq!( + LastSeen::::get(identifier, None::), + Some(10) + ); + }); + } + + #[test] + fn tx_extension_rejects_when_call_occurs_too_soon() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::Exact(5)); + LastSeen::::insert(identifier, None::, 20); + + System::set_block_number(22); + + let err = + validate_with_tx_extension(&extension, &call).expect_err("should be rate limited"); + match err { + TransactionValidityError::Invalid(InvalidTransaction::Custom(code)) => { + assert_eq!(code, 1); + } + other => panic!("unexpected error: {:?}", other), + } + }); + } + + #[test] + fn tx_extension_skips_last_seen_when_span_zero() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::Exact(0)); + + System::set_block_number(30); + + let (_valid, val, _) = validate_with_tx_extension(&extension, &call).expect("valid"); + assert!(val.is_none()); + + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let origin_for_prepare = RuntimeOrigin::signed(42); + let pre = extension + .clone() + .prepare(val.clone(), &origin_for_prepare, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + assert_eq!( + LastSeen::::get(identifier, None::), + None + ); + }); + } +} From 13b8774bdfb1bfdb6dcebbb30045fb2282da02d7 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Mon, 27 Oct 2025 18:18:27 +0300 Subject: [PATCH 09/23] Add crate-level documentation for pallet-rate-limiting --- pallets/rate-limiting/src/lib.rs | 70 +++++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 2ac9c76114..e9659c7abd 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -1,6 +1,74 @@ #![cfg_attr(not(feature = "std"), no_std)] -//! Basic rate limiting pallet. +//! Rate limiting for runtime calls with optional contextual restrictions. +//! +//! # Overview +//! +//! `pallet-rate-limiting` lets a runtime restrict how frequently particular calls can execute. +//! Limits are stored on-chain, keyed by the call's pallet/variant pair. Each entry can specify an +//! exact block span or defer to a configured default. The pallet exposes three roots-only +//! extrinsics to manage this data: +//! +//! - [`set_rate_limit`](pallet::Pallet::set_rate_limit): assign a limit to an extrinsic, either as +//! `RateLimit::Exact(blocks)` or `RateLimit::Default`. +//! - [`clear_rate_limit`](pallet::Pallet::clear_rate_limit): remove a stored limit. +//! - [`set_default_rate_limit`](pallet::Pallet::set_default_rate_limit): set the global default +//! block span used by `RateLimit::Default` entries. +//! +//! The pallet also tracks the last block in which a rate-limited call was executed, per optional +//! *context*. Context allows one limit definition (for example, “set weights”) to be enforced per +//! subnet, account, or other grouping chosen by the runtime. The storage layout is: +//! +//! - [`Limits`](pallet::Limits): `TransactionIdentifier → RateLimit` +//! - [`DefaultLimit`](pallet::DefaultLimit): `BlockNumber` +//! - [`LastSeen`](pallet::LastSeen): `(TransactionIdentifier, Option) → BlockNumber` +//! +//! # Transaction extension +//! +//! Enforcement happens via [`RateLimitTransactionExtension`], which implements +//! `sp_runtime::traits::TransactionExtension`. The extension consults `Limits`, fetches the current +//! block, and decides whether the call is eligible. If successful, it returns metadata that causes +//! [`LastSeen`](pallet::LastSeen) to update after dispatch. A rejected call yields +//! `InvalidTransaction::Custom(1)`. +//! +//! To enable the extension, add it to your runtime's transaction extension tuple. For example: +//! +//! ```rust +//! pub type TransactionExtensions = ( +//! // ... other extensions ... +//! pallet_rate_limiting::RateLimitTransactionExtension, +//! ); +//! ``` +//! +//! # Context resolver +//! +//! The extension needs to know when two invocations should share a rate limit. This is controlled +//! by implementing [`RateLimitContextResolver`] for the runtime call type (or for a helper that the +//! runtime wires into [`Config::ContextResolver`]). The resolver receives the call and returns +//! `Some(context)` if the rate should be scoped (e.g. by `netuid`), or `None` for a global limit. +//! +//! ```rust +//! pub struct WeightsContextResolver; +//! +//! impl pallet_rate_limiting::RateLimitContextResolver +//! for WeightsContextResolver +//! { +//! fn context(call: &RuntimeCall) -> Option { +//! match call { +//! RuntimeCall::Subtensor(pallet_subtensor::Call::set_weights { netuid, .. }) => { +//! Some(*netuid) +//! } +//! _ => None, +//! } +//! } +//! } +//! +//! impl pallet_rate_limiting::Config for Runtime { +//! type RuntimeCall = RuntimeCall; +//! type LimitContext = NetUid; +//! type ContextResolver = WeightsContextResolver; +//! } +//! ``` pub use pallet::*; pub use tx_extension::RateLimitTransactionExtension; From a4a1c88b1764bbe777246846f25c542a95d0460f Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Mon, 27 Oct 2025 18:24:18 +0300 Subject: [PATCH 10/23] Add benchmarks for pallet-rate-limiting --- Cargo.lock | 1 + pallets/rate-limiting/Cargo.toml | 9 ++- pallets/rate-limiting/src/benchmarking.rs | 72 +++++++++++++++++++++++ pallets/rate-limiting/src/lib.rs | 14 ++++- pallets/rate-limiting/src/mock.rs | 13 ++++ pallets/rate-limiting/src/tests.rs | 38 +----------- pallets/rate-limiting/src/types.rs | 41 +++++++++++++ 7 files changed, 147 insertions(+), 41 deletions(-) create mode 100644 pallets/rate-limiting/src/benchmarking.rs diff --git a/Cargo.lock b/Cargo.lock index d116651d33..ee22c024f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10209,6 +10209,7 @@ dependencies = [ name = "pallet-rate-limiting" version = "0.1.0" dependencies = [ + "frame-benchmarking", "frame-support", "frame-system", "parity-scale-codec", diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml index b24ff40ea5..3447145622 100644 --- a/pallets/rate-limiting/Cargo.toml +++ b/pallets/rate-limiting/Cargo.toml @@ -11,12 +11,13 @@ targets = ["x86_64-unknown-linux-gnu"] [dependencies] codec = { workspace = true, features = ["derive"] } +frame-benchmarking = { workspace = true, optional = true } frame-support.workspace = true frame-system.workspace = true scale-info = { workspace = true, features = ["derive"] } +serde = { workspace = true, features = ["derive"], optional = true } sp-std.workspace = true subtensor-runtime-common.workspace = true -serde = { workspace = true, features = ["derive"], optional = true } [dev-dependencies] sp-core.workspace = true @@ -27,12 +28,16 @@ sp-runtime.workspace = true default = ["std"] std = [ "codec/std", + "frame-benchmarking?/std", "frame-support/std", "frame-system/std", "scale-info/std", + "serde", "sp-std/std", "subtensor-runtime-common/std", - "serde", +] +runtime-benchmarks = [ + "frame-benchmarking", ] try-runtime = [ "frame-support/try-runtime", diff --git a/pallets/rate-limiting/src/benchmarking.rs b/pallets/rate-limiting/src/benchmarking.rs new file mode 100644 index 0000000000..3266674ac7 --- /dev/null +++ b/pallets/rate-limiting/src/benchmarking.rs @@ -0,0 +1,72 @@ +//! Benchmarking setup for pallet-rate-limiting +#![cfg(feature = "runtime-benchmarks")] +#![allow(clippy::arithmetic_side_effects)] + +use codec::Decode; +use frame_benchmarking::v2::*; +use frame_system::{RawOrigin, pallet_prelude::BlockNumberFor}; + +use super::*; + +pub trait BenchmarkHelper { + fn sample_call() -> Call; +} + +impl BenchmarkHelper for () +where + Call: Decode, +{ + fn sample_call() -> Call { + Decode::decode(&mut &[][..]).expect("Provide a call via BenchmarkHelper::sample_call") + } +} + +fn sample_call() -> Box<::RuntimeCall> +where + T::BenchmarkHelper: BenchmarkHelper<::RuntimeCall>, +{ + Box::new(T::BenchmarkHelper::sample_call()) +} + +#[benchmarks] +mod benchmarks { + use super::*; + + #[benchmark] + fn set_rate_limit() { + let call = sample_call::(); + let limit = RateLimit::>::Exact(BlockNumberFor::::from(10u32)); + + #[extrinsic_call] + _(RawOrigin::Root, call, limit.clone()); + + assert!(Limits::::iter().any(|(_, stored)| stored == limit)); + } + + #[benchmark] + fn clear_rate_limit() { + let call = sample_call::(); + let limit = RateLimit::>::Exact(BlockNumberFor::::from(10u32)); + + // Pre-populate limit for benchmark call + let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); + Limits::::insert(identifier, limit); + + #[extrinsic_call] + _(RawOrigin::Root, call); + + assert!(Limits::::get(identifier).is_none()); + } + + #[benchmark] + fn set_default_rate_limit() { + let block_span = BlockNumberFor::::from(10u32); + + #[extrinsic_call] + _(RawOrigin::Root, block_span); + + assert_eq!(DefaultLimit::::get(), block_span); + } + + impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Test); +} diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index e9659c7abd..637216e71b 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -33,7 +33,7 @@ //! //! To enable the extension, add it to your runtime's transaction extension tuple. For example: //! -//! ```rust +//! ```ignore //! pub type TransactionExtensions = ( //! // ... other extensions ... //! pallet_rate_limiting::RateLimitTransactionExtension, @@ -47,7 +47,7 @@ //! runtime wires into [`Config::ContextResolver`]). The resolver receives the call and returns //! `Some(context)` if the rate should be scoped (e.g. by `netuid`), or `None` for a global limit. //! -//! ```rust +//! ```ignore //! pub struct WeightsContextResolver; //! //! impl pallet_rate_limiting::RateLimitContextResolver @@ -70,10 +70,14 @@ //! } //! ``` +#[cfg(feature = "runtime-benchmarks")] +pub use benchmarking::BenchmarkHelper; pub use pallet::*; pub use tx_extension::RateLimitTransactionExtension; pub use types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; +#[cfg(feature = "runtime-benchmarks")] +mod benchmarking; mod tx_extension; mod types; @@ -94,6 +98,8 @@ pub mod pallet { use frame_system::{ensure_root, pallet_prelude::*}; use sp_std::{convert::TryFrom, vec::Vec}; + #[cfg(feature = "runtime-benchmarks")] + use crate::benchmarking::BenchmarkHelper as BenchmarkHelperTrait; use crate::types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; /// Configuration trait for the rate limiting pallet. @@ -110,6 +116,10 @@ pub mod pallet { /// Resolves the context for a given runtime call. type ContextResolver: RateLimitContextResolver<::RuntimeCall, Self::LimitContext>; + + /// Helper used to construct runtime calls for benchmarking. + #[cfg(feature = "runtime-benchmarks")] + type BenchmarkHelper: BenchmarkHelperTrait<::RuntimeCall>; } /// Storage mapping from transaction identifier to its configured rate limit. diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs index d7b9fde20b..4218e757a1 100644 --- a/pallets/rate-limiting/src/mock.rs +++ b/pallets/rate-limiting/src/mock.rs @@ -10,6 +10,7 @@ use frame_support::{ }; use sp_core::H256; use sp_io::TestExternalities; +use sp_std::vec::Vec; use crate as pallet_rate_limiting; use crate::TransactionIdentifier; @@ -67,6 +68,18 @@ impl pallet_rate_limiting::Config for Test { type RuntimeCall = RuntimeCall; type LimitContext = LimitContext; type ContextResolver = TestContextResolver; + #[cfg(feature = "runtime-benchmarks")] + type BenchmarkHelper = BenchHelper; +} + +#[cfg(feature = "runtime-benchmarks")] +pub struct BenchHelper; + +#[cfg(feature = "runtime-benchmarks")] +impl crate::BenchmarkHelper for BenchHelper { + fn sample_call() -> RuntimeCall { + RuntimeCall::System(frame_system::Call::remark { remark: Vec::new() }) + } } pub type RateLimitingCall = crate::Call; diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index 4eb07ad926..5a8d2dd933 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -1,9 +1,6 @@ use frame_support::{assert_noop, assert_ok, error::BadOrigin}; -use sp_runtime::DispatchError; -use crate::{ - DefaultLimit, LastSeen, Limits, RateLimit, mock::*, pallet::Error, types::TransactionIdentifier, -}; +use crate::{DefaultLimit, LastSeen, Limits, RateLimit, mock::*, pallet::Error}; #[test] fn limit_for_call_names_returns_none_if_not_set() { @@ -100,39 +97,6 @@ fn is_within_limit_true_after_required_span() { }); } -#[test] -fn transaction_identifier_from_call_matches_expected_indices() { - let call = - RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); - - let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); - - // System is the first pallet in the mock runtime, RateLimiting is second. - assert_eq!(identifier.pallet_index, 1); - // set_default_rate_limit has call_index 2. - assert_eq!(identifier.extrinsic_index, 2); -} - -#[test] -fn transaction_identifier_names_matches_call_metadata() { - let call = - RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); - let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); - - let (pallet, extrinsic) = identifier.names::().expect("call metadata"); - assert_eq!(pallet, "RateLimiting"); - assert_eq!(extrinsic, "set_default_rate_limit"); -} - -#[test] -fn transaction_identifier_names_error_for_unknown_indices() { - let identifier = TransactionIdentifier::new(99, 0); - - let err = identifier.names::().expect_err("should fail"); - let expected: DispatchError = Error::::InvalidRuntimeCall.into(); - assert_eq!(err, expected); -} - #[test] fn set_rate_limit_updates_storage_and_emits_event() { new_test_ext().execute_with(|| { diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index 4e00053ec7..72e2a43777 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -97,3 +97,44 @@ pub enum RateLimit { /// Apply an exact rate limit measured in blocks. Exact(BlockNumber), } + +#[cfg(test)] +mod tests { + use sp_runtime::DispatchError; + + use super::*; + use crate::{mock::*, pallet::Error}; + + #[test] + fn transaction_identifier_from_call_matches_expected_indices() { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + + let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); + + // System is the first pallet in the mock runtime, RateLimiting is second. + assert_eq!(identifier.pallet_index, 1); + // set_default_rate_limit has call_index 2. + assert_eq!(identifier.extrinsic_index, 2); + } + + #[test] + fn transaction_identifier_names_matches_call_metadata() { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); + + let (pallet, extrinsic) = identifier.names::().expect("call metadata"); + assert_eq!(pallet, "RateLimiting"); + assert_eq!(extrinsic, "set_default_rate_limit"); + } + + #[test] + fn transaction_identifier_names_error_for_unknown_indices() { + let identifier = TransactionIdentifier::new(99, 0); + + let err = identifier.names::().expect_err("should fail"); + let expected: DispatchError = Error::::InvalidRuntimeCall.into(); + assert_eq!(err, expected); + } +} From 6ccc421e94f866ed99f85499de05673107ea142d Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Wed, 29 Oct 2025 15:18:52 +0300 Subject: [PATCH 11/23] Extend rate limit setting with context --- pallets/rate-limiting/src/benchmarking.rs | 13 +- pallets/rate-limiting/src/lib.rs | 67 ++++++--- pallets/rate-limiting/src/tests.rs | 175 +++++++++++++++++++--- pallets/rate-limiting/src/tx_extension.rs | 21 +-- 4 files changed, 217 insertions(+), 59 deletions(-) diff --git a/pallets/rate-limiting/src/benchmarking.rs b/pallets/rate-limiting/src/benchmarking.rs index 3266674ac7..bf4a3b37b5 100644 --- a/pallets/rate-limiting/src/benchmarking.rs +++ b/pallets/rate-limiting/src/benchmarking.rs @@ -36,11 +36,13 @@ mod benchmarks { fn set_rate_limit() { let call = sample_call::(); let limit = RateLimit::>::Exact(BlockNumberFor::::from(10u32)); + let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); + let context = ::ContextResolver::context(call.as_ref()); #[extrinsic_call] - _(RawOrigin::Root, call, limit.clone()); + _(RawOrigin::Root, call, limit.clone(), None); - assert!(Limits::::iter().any(|(_, stored)| stored == limit)); + assert_eq!(Limits::::get(&identifier, context), Some(limit)); } #[benchmark] @@ -50,12 +52,13 @@ mod benchmarks { // Pre-populate limit for benchmark call let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); - Limits::::insert(identifier, limit); + let context = ::ContextResolver::context(call.as_ref()); + Limits::::insert(identifier, context.clone(), limit); #[extrinsic_call] - _(RawOrigin::Root, call); + _(RawOrigin::Root, call, None); - assert!(Limits::::get(identifier).is_none()); + assert!(Limits::::get(identifier, context).is_none()); } #[benchmark] diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 637216e71b..23eff3cf9a 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -10,8 +10,10 @@ //! extrinsics to manage this data: //! //! - [`set_rate_limit`](pallet::Pallet::set_rate_limit): assign a limit to an extrinsic, either as -//! `RateLimit::Exact(blocks)` or `RateLimit::Default`. -//! - [`clear_rate_limit`](pallet::Pallet::clear_rate_limit): remove a stored limit. +//! `RateLimit::Exact(blocks)` or `RateLimit::Default`. The optional context parameter lets you +//! scope the configuration to a particular subnet/key/account while keeping a global fallback. +//! - [`clear_rate_limit`](pallet::Pallet::clear_rate_limit): remove a stored limit for the provided +//! context (or for the global entry when `None` is supplied). //! - [`set_default_rate_limit`](pallet::Pallet::set_default_rate_limit): set the global default //! block span used by `RateLimit::Default` entries. //! @@ -45,7 +47,9 @@ //! The extension needs to know when two invocations should share a rate limit. This is controlled //! by implementing [`RateLimitContextResolver`] for the runtime call type (or for a helper that the //! runtime wires into [`Config::ContextResolver`]). The resolver receives the call and returns -//! `Some(context)` if the rate should be scoped (e.g. by `netuid`), or `None` for a global limit. +//! `Some(context)` if the rate should be scoped (e.g. by `netuid`), or `None` to use the global +//! entry. The resolver is only used when *tracking* executions; you still configure limits via the +//! explicit `context` argument on `set_rate_limit`/`clear_rate_limit`. //! //! ```ignore //! pub struct WeightsContextResolver; @@ -112,7 +116,7 @@ pub mod pallet { + IsType<::RuntimeCall>; /// Context type used for contextual (per-group) rate limits. - type LimitContext: Parameter + Clone + PartialEq + Eq; + type LimitContext: Parameter + Clone + PartialEq + Eq + MaybeSerializeDeserialize; /// Resolves the context for a given runtime call. type ContextResolver: RateLimitContextResolver<::RuntimeCall, Self::LimitContext>; @@ -122,13 +126,15 @@ pub mod pallet { type BenchmarkHelper: BenchmarkHelperTrait<::RuntimeCall>; } - /// Storage mapping from transaction identifier to its configured rate limit. + /// Storage mapping from transaction identifier and optional context to its configured rate limit. #[pallet::storage] #[pallet::getter(fn limits)] - pub type Limits = StorageMap< + pub type Limits = StorageDoubleMap< _, Blake2_128Concat, TransactionIdentifier, + Blake2_128Concat, + Option<::LimitContext>, RateLimit>, OptionQuery, >; @@ -160,6 +166,8 @@ pub mod pallet { RateLimitSet { /// Identifier of the affected transaction. transaction: TransactionIdentifier, + /// Context to which the limit applies, if any. + context: Option<::LimitContext>, /// The new limit configuration applied to the transaction. limit: RateLimit>, /// Pallet name associated with the transaction. @@ -171,6 +179,8 @@ pub mod pallet { RateLimitCleared { /// Identifier of the affected transaction. transaction: TransactionIdentifier, + /// Context from which the limit was cleared, if any. + context: Option<::LimitContext>, /// Pallet name associated with the transaction. pallet: Vec, /// Extrinsic name associated with the transaction. @@ -195,7 +205,11 @@ pub mod pallet { #[pallet::genesis_config] pub struct GenesisConfig { pub default_limit: BlockNumberFor, - pub limits: Vec<(TransactionIdentifier, RateLimit>)>, + pub limits: Vec<( + TransactionIdentifier, + Option<::LimitContext>, + RateLimit>, + )>, } #[cfg(feature = "std")] @@ -213,8 +227,8 @@ pub mod pallet { fn build(&self) { DefaultLimit::::put(self.default_limit); - for (identifier, limit) in &self.limits { - Limits::::insert(identifier, limit.clone()); + for (identifier, context, limit) in &self.limits { + Limits::::insert(identifier, context.clone(), limit.clone()); } } } @@ -230,7 +244,7 @@ pub mod pallet { identifier: &TransactionIdentifier, context: &Option<::LimitContext>, ) -> Result { - let Some(block_span) = Self::resolved_limit(identifier) else { + let Some(block_span) = Self::resolved_limit(identifier, context) else { return Ok(true); }; @@ -246,8 +260,13 @@ pub mod pallet { Ok(true) } - fn resolved_limit(identifier: &TransactionIdentifier) -> Option> { - let limit = Limits::::get(identifier)?; + pub(crate) fn resolved_limit( + identifier: &TransactionIdentifier, + context: &Option<::LimitContext>, + ) -> Option> { + let lookup = Limits::::get(identifier, context.clone()) + .or_else(|| Limits::::get(identifier, None::<::LimitContext>)); + let limit = lookup?; Some(match limit { RateLimit::Default => DefaultLimit::::get(), RateLimit::Exact(block_span) => block_span, @@ -258,18 +277,21 @@ pub mod pallet { pub fn limit_for_call_names( pallet_name: &str, extrinsic_name: &str, + context: Option<::LimitContext>, ) -> Option>> { let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; - Limits::::get(&identifier) + Limits::::get(&identifier, context.clone()) + .or_else(|| Limits::::get(&identifier, None::<::LimitContext>)) } /// Returns the resolved block span for the specified pallet/extrinsic names, if any. pub fn resolved_limit_for_call_names( pallet_name: &str, extrinsic_name: &str, + context: Option<::LimitContext>, ) -> Option> { let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; - Self::resolved_limit(&identifier) + Self::resolved_limit(&identifier, &context) } fn identifier_for_call_names( @@ -288,21 +310,24 @@ pub mod pallet { #[pallet::call] impl Pallet { - /// Sets the rate limit configuration for the given call. + /// Sets the rate limit configuration for the given call and optional context. /// /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any - /// arguments embedded in the call are ignored**. + /// arguments embedded in the call are ignored**. The `context` parameter determines which + /// scoped entry is updated (for example a subnet identifier). Passing `None` updates the + /// global entry, which acts as a fallback when no context-specific limit exists. #[pallet::call_index(0)] #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn set_rate_limit( origin: OriginFor, call: Box<::RuntimeCall>, limit: RateLimit>, + context: Option<::LimitContext>, ) -> DispatchResult { ensure_root(origin)?; let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - Limits::::insert(&identifier, limit); + Limits::::insert(&identifier, context.clone(), limit.clone()); let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); @@ -310,6 +335,7 @@ pub mod pallet { Self::deposit_event(Event::RateLimitSet { transaction: identifier, + context, limit, pallet, extrinsic, @@ -321,12 +347,14 @@ pub mod pallet { /// Clears the rate limit for the given call, if present. /// /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any - /// arguments embedded in the call are ignored**. + /// arguments embedded in the call are ignored**. The `context` parameter must match the + /// entry that should be removed (use `None` to remove the global configuration). #[pallet::call_index(1)] #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn clear_rate_limit( origin: OriginFor, call: Box<::RuntimeCall>, + context: Option<::LimitContext>, ) -> DispatchResult { ensure_root(origin)?; @@ -337,12 +365,13 @@ pub mod pallet { let extrinsic = Vec::from(extrinsic_name.as_bytes()); ensure!( - Limits::::take(&identifier).is_some(), + Limits::::take(&identifier, context.clone()).is_some(), Error::::MissingRateLimit ); Self::deposit_event(Event::RateLimitCleared { transaction: identifier, + context, pallet, extrinsic, }); diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index 5a8d2dd933..ce5b6d25ab 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -6,7 +6,8 @@ use crate::{DefaultLimit, LastSeen, Limits, RateLimit, mock::*, pallet::Error}; fn limit_for_call_names_returns_none_if_not_set() { new_test_ext().execute_with(|| { assert!( - RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit").is_none() + RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", None) + .is_none() ); }); } @@ -17,14 +18,36 @@ fn limit_for_call_names_returns_stored_limit() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, RateLimit::Exact(7)); + Limits::::insert(identifier, None::, RateLimit::Exact(7)); - let fetched = RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit") - .expect("limit should exist"); + let fetched = + RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", None) + .expect("limit should exist"); assert_eq!(fetched, RateLimit::Exact(7)); }); } +#[test] +fn limit_for_call_names_prefers_context_specific_limit() { + new_test_ext().execute_with(|| { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&call); + Limits::::insert(identifier, None::, RateLimit::Exact(3)); + Limits::::insert(identifier, Some(5), RateLimit::Exact(8)); + + let fetched = + RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", Some(5)) + .expect("limit should exist"); + assert_eq!(fetched, RateLimit::Exact(8)); + + let fallback = + RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", Some(1)) + .expect("limit should exist"); + assert_eq!(fallback, RateLimit::Exact(3)); + }); +} + #[test] fn resolved_limit_for_call_names_resolves_default_value() { new_test_ext().execute_with(|| { @@ -32,21 +55,55 @@ fn resolved_limit_for_call_names_resolves_default_value() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, RateLimit::Default); - - let resolved = - RateLimiting::resolved_limit_for_call_names("RateLimiting", "set_default_rate_limit") - .expect("resolved limit"); + Limits::::insert(identifier, None::, RateLimit::Default); + + let resolved = RateLimiting::resolved_limit_for_call_names( + "RateLimiting", + "set_default_rate_limit", + None, + ) + .expect("resolved limit"); assert_eq!(resolved, 3); }); } +#[test] +fn resolved_limit_for_call_names_prefers_context_specific_value() { + new_test_ext().execute_with(|| { + let call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&call); + Limits::::insert(identifier, None::, RateLimit::Exact(4)); + Limits::::insert(identifier, Some(6), RateLimit::Exact(9)); + + let resolved = RateLimiting::resolved_limit_for_call_names( + "RateLimiting", + "set_default_rate_limit", + Some(6), + ) + .expect("resolved limit"); + assert_eq!(resolved, 9); + + let fallback = RateLimiting::resolved_limit_for_call_names( + "RateLimiting", + "set_default_rate_limit", + Some(1), + ) + .expect("resolved limit"); + assert_eq!(fallback, 4); + }); +} + #[test] fn resolved_limit_for_call_names_returns_none_when_unset() { new_test_ext().execute_with(|| { assert!( - RateLimiting::resolved_limit_for_call_names("RateLimiting", "set_default_rate_limit") - .is_none() + RateLimiting::resolved_limit_for_call_names( + "RateLimiting", + "set_default_rate_limit", + None, + ) + .is_none() ); }); } @@ -69,7 +126,7 @@ fn is_within_limit_false_when_rate_limited() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, RateLimit::Exact(5)); + Limits::::insert(identifier, Some(1 as LimitContext), RateLimit::Exact(5)); LastSeen::::insert(identifier, Some(1 as LimitContext), 9); System::set_block_number(13); @@ -86,7 +143,7 @@ fn is_within_limit_true_after_required_span() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, RateLimit::Exact(5)); + Limits::::insert(identifier, Some(2 as LimitContext), RateLimit::Exact(5)); LastSeen::::insert(identifier, Some(2 as LimitContext), 10); System::set_block_number(20); @@ -109,20 +166,26 @@ fn set_rate_limit_updates_storage_and_emits_event() { assert_ok!(RateLimiting::set_rate_limit( RuntimeOrigin::root(), Box::new(target_call.clone()), - limit + limit, + None, )); let identifier = identifier_for(&target_call); - assert_eq!(Limits::::get(identifier), Some(limit)); + assert_eq!( + Limits::::get(identifier, None::), + Some(limit) + ); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitSet { transaction, + context, limit: emitted_limit, pallet, extrinsic, }) => { assert_eq!(transaction, identifier); + assert_eq!(context, None); assert_eq!(emitted_limit, limit); assert_eq!(pallet, b"RateLimiting".to_vec()); assert_eq!(extrinsic, b"set_default_rate_limit".to_vec()); @@ -132,6 +195,29 @@ fn set_rate_limit_updates_storage_and_emits_event() { }); } +#[test] +fn set_rate_limit_supports_context_specific_limit() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let context = Some(7u16); + assert_ok!(RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + Box::new(target_call.clone()), + RateLimit::Exact(11), + context, + )); + + let identifier = identifier_for(&target_call); + assert_eq!( + Limits::::get(identifier, Some(7)), + Some(RateLimit::Exact(11)) + ); + // global remains untouched + assert_eq!(Limits::::get(identifier, None::), None); + }); +} + #[test] fn set_rate_limit_requires_root() { new_test_ext().execute_with(|| { @@ -142,7 +228,8 @@ fn set_rate_limit_requires_root() { RateLimiting::set_rate_limit( RuntimeOrigin::signed(1), Box::new(target_call), - RateLimit::Exact(1) + RateLimit::Exact(1), + None, ), BadOrigin ); @@ -158,11 +245,15 @@ fn set_rate_limit_accepts_default_variant() { assert_ok!(RateLimiting::set_rate_limit( RuntimeOrigin::root(), Box::new(target_call.clone()), - RateLimit::Default + RateLimit::Default, + None, )); let identifier = identifier_for(&target_call); - assert_eq!(Limits::::get(identifier), Some(RateLimit::Default)); + assert_eq!( + Limits::::get(identifier, None::), + Some(RateLimit::Default) + ); }); } @@ -174,22 +265,25 @@ fn clear_rate_limit_removes_entry_and_emits_event() { let target_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&target_call); - Limits::::insert(identifier, RateLimit::Exact(4)); + Limits::::insert(identifier, None::, RateLimit::Exact(4)); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), - Box::new(target_call.clone()) + Box::new(target_call.clone()), + None, )); - assert_eq!(Limits::::get(identifier), None); + assert_eq!(Limits::::get(identifier, None::), None); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { transaction, + context, pallet, extrinsic, }) => { assert_eq!(transaction, identifier); + assert_eq!(context, None); assert_eq!(pallet, b"RateLimiting".to_vec()); assert_eq!(extrinsic, b"set_default_rate_limit".to_vec()); } @@ -205,12 +299,49 @@ fn clear_rate_limit_fails_when_missing() { RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); assert_noop!( - RateLimiting::clear_rate_limit(RuntimeOrigin::root(), Box::new(target_call)), + RateLimiting::clear_rate_limit(RuntimeOrigin::root(), Box::new(target_call), None), Error::::MissingRateLimit ); }); } +#[test] +fn clear_rate_limit_removes_only_selected_context() { + new_test_ext().execute_with(|| { + System::reset_events(); + + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&target_call); + Limits::::insert(identifier, None::, RateLimit::Exact(5)); + Limits::::insert(identifier, Some(9), RateLimit::Exact(7)); + + assert_ok!(RateLimiting::clear_rate_limit( + RuntimeOrigin::root(), + Box::new(target_call.clone()), + Some(9), + )); + + assert_eq!(Limits::::get(identifier, Some(9u16)), None); + assert_eq!( + Limits::::get(identifier, None::), + Some(RateLimit::Exact(5)) + ); + + match pop_last_event() { + RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { + transaction, + context, + .. + }) => { + assert_eq!(transaction, identifier); + assert_eq!(context, Some(9)); + } + other => panic!("unexpected event: {:?}", other), + } + }); +} + #[test] fn set_default_rate_limit_updates_storage_and_emits_event() { new_test_ext().execute_with(|| { diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index f06e72975e..af02bfa453 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -16,8 +16,8 @@ use scale_info::TypeInfo; use sp_std::{marker::PhantomData, result::Result}; use crate::{ - Config, LastSeen, Limits, Pallet, - types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}, + Config, LastSeen, Pallet, + types::{RateLimitContextResolver, TransactionIdentifier}, }; /// Identifier returned in the transaction metadata for the rate limiting extension. @@ -66,21 +66,16 @@ where Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)), }; - let Some(limit) = Limits::::get(&identifier) else { - return Ok((ValidTransaction::default(), None, origin)); - }; + let context = ::ContextResolver::context(call); - let block_span = match limit { - RateLimit::Default => Pallet::::default_limit(), - RateLimit::Exact(block_span) => block_span, + let Some(block_span) = Pallet::::resolved_limit(&identifier, &context) else { + return Ok((ValidTransaction::default(), None, origin)); }; if block_span.is_zero() { return Ok((ValidTransaction::default(), None, origin)); } - let context = ::ContextResolver::context(call); - let within_limit = Pallet::::is_within_limit(&identifier, &context) .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; @@ -213,7 +208,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, RateLimit::Exact(5)); + Limits::::insert(identifier, None::, RateLimit::Exact(5)); System::set_block_number(10); @@ -251,7 +246,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, RateLimit::Exact(5)); + Limits::::insert(identifier, None::, RateLimit::Exact(5)); LastSeen::::insert(identifier, None::, 20); System::set_block_number(22); @@ -273,7 +268,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, RateLimit::Exact(0)); + Limits::::insert(identifier, None::, RateLimit::Exact(0)); System::set_block_number(30); From be9d4f7ffc7a9be1c71cf652285cb4694367d636 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Wed, 29 Oct 2025 17:46:43 +0300 Subject: [PATCH 12/23] Make rate limiting pallet instantiable --- pallets/rate-limiting/src/benchmarking.rs | 19 ++-- pallets/rate-limiting/src/lib.rs | 98 +++++++++++---------- pallets/rate-limiting/src/mock.rs | 2 +- pallets/rate-limiting/src/tests.rs | 52 ++++++----- pallets/rate-limiting/src/tx_extension.rs | 102 +++++++++++++++------- pallets/rate-limiting/src/types.rs | 36 ++++---- 6 files changed, 184 insertions(+), 125 deletions(-) diff --git a/pallets/rate-limiting/src/benchmarking.rs b/pallets/rate-limiting/src/benchmarking.rs index bf4a3b37b5..8392109772 100644 --- a/pallets/rate-limiting/src/benchmarking.rs +++ b/pallets/rate-limiting/src/benchmarking.rs @@ -36,13 +36,16 @@ mod benchmarks { fn set_rate_limit() { let call = sample_call::(); let limit = RateLimit::>::Exact(BlockNumberFor::::from(10u32)); - let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); - let context = ::ContextResolver::context(call.as_ref()); + let identifier = + TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); #[extrinsic_call] _(RawOrigin::Root, call, limit.clone(), None); - assert_eq!(Limits::::get(&identifier, context), Some(limit)); + assert_eq!( + Limits::::get(&identifier, None::<::LimitContext>), + Some(limit) + ); } #[benchmark] @@ -51,14 +54,14 @@ mod benchmarks { let limit = RateLimit::>::Exact(BlockNumberFor::::from(10u32)); // Pre-populate limit for benchmark call - let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); - let context = ::ContextResolver::context(call.as_ref()); - Limits::::insert(identifier, context.clone(), limit); + let identifier = + TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); + Limits::::insert(identifier, None::<::LimitContext>, limit); #[extrinsic_call] _(RawOrigin::Root, call, None); - assert!(Limits::::get(identifier, context).is_none()); + assert!(Limits::::get(identifier, None::<::LimitContext>).is_none()); } #[benchmark] @@ -68,7 +71,7 @@ mod benchmarks { #[extrinsic_call] _(RawOrigin::Root, block_span); - assert_eq!(DefaultLimit::::get(), block_span); + assert_eq!(DefaultLimit::::get(), block_span); } impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Test); diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 23eff3cf9a..2ca3ab87ba 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -25,6 +25,9 @@ //! - [`DefaultLimit`](pallet::DefaultLimit): `BlockNumber` //! - [`LastSeen`](pallet::LastSeen): `(TransactionIdentifier, Option) → BlockNumber` //! +//! Each storage map is namespaced by pallet instance; runtimes can deploy multiple independent +//! instances to manage distinct rate-limiting scopes. +//! //! # Transaction extension //! //! Enforcement happens via [`RateLimitTransactionExtension`], which implements @@ -100,7 +103,7 @@ pub mod pallet { traits::{BuildGenesisConfig, GetCallMetadata}, }; use frame_system::{ensure_root, pallet_prelude::*}; - use sp_std::{convert::TryFrom, vec::Vec}; + use sp_std::{convert::TryFrom, marker::PhantomData, vec::Vec}; #[cfg(feature = "runtime-benchmarks")] use crate::benchmarking::BenchmarkHelper as BenchmarkHelperTrait; @@ -108,7 +111,7 @@ pub mod pallet { /// Configuration trait for the rate limiting pallet. #[pallet::config] - pub trait Config: frame_system::Config { + pub trait Config: frame_system::Config { /// The overarching runtime call type. type RuntimeCall: Parameter + Codec @@ -119,22 +122,22 @@ pub mod pallet { type LimitContext: Parameter + Clone + PartialEq + Eq + MaybeSerializeDeserialize; /// Resolves the context for a given runtime call. - type ContextResolver: RateLimitContextResolver<::RuntimeCall, Self::LimitContext>; + type ContextResolver: RateLimitContextResolver<>::RuntimeCall, Self::LimitContext>; /// Helper used to construct runtime calls for benchmarking. #[cfg(feature = "runtime-benchmarks")] - type BenchmarkHelper: BenchmarkHelperTrait<::RuntimeCall>; + type BenchmarkHelper: BenchmarkHelperTrait<>::RuntimeCall>; } /// Storage mapping from transaction identifier and optional context to its configured rate limit. #[pallet::storage] #[pallet::getter(fn limits)] - pub type Limits = StorageDoubleMap< + pub type Limits, I: 'static = ()> = StorageDoubleMap< _, Blake2_128Concat, TransactionIdentifier, Blake2_128Concat, - Option<::LimitContext>, + Option<>::LimitContext>, RateLimit>, OptionQuery, >; @@ -143,12 +146,12 @@ pub mod pallet { /// /// The second key is `None` for global limits and `Some(context)` for contextual limits. #[pallet::storage] - pub type LastSeen = StorageDoubleMap< + pub type LastSeen, I: 'static = ()> = StorageDoubleMap< _, Blake2_128Concat, TransactionIdentifier, Blake2_128Concat, - Option<::LimitContext>, + Option<>::LimitContext>, BlockNumberFor, OptionQuery, >; @@ -156,18 +159,19 @@ pub mod pallet { /// Default block span applied when an extrinsic uses the default rate limit. #[pallet::storage] #[pallet::getter(fn default_limit)] - pub type DefaultLimit = StorageValue<_, BlockNumberFor, ValueQuery>; + pub type DefaultLimit, I: 'static = ()> = + StorageValue<_, BlockNumberFor, ValueQuery>; /// Events emitted by the rate limiting pallet. #[pallet::event] #[pallet::generate_deposit(pub(super) fn deposit_event)] - pub enum Event { + pub enum Event, I: 'static = ()> { /// A rate limit was set or updated. RateLimitSet { /// Identifier of the affected transaction. transaction: TransactionIdentifier, /// Context to which the limit applies, if any. - context: Option<::LimitContext>, + context: Option<>::LimitContext>, /// The new limit configuration applied to the transaction. limit: RateLimit>, /// Pallet name associated with the transaction. @@ -180,7 +184,7 @@ pub mod pallet { /// Identifier of the affected transaction. transaction: TransactionIdentifier, /// Context from which the limit was cleared, if any. - context: Option<::LimitContext>, + context: Option<>::LimitContext>, /// Pallet name associated with the transaction. pallet: Vec, /// Extrinsic name associated with the transaction. @@ -195,7 +199,7 @@ pub mod pallet { /// Errors that can occur while configuring rate limits. #[pallet::error] - pub enum Error { + pub enum Error { /// Failed to extract the pallet and extrinsic indices from the call. InvalidRuntimeCall, /// Attempted to remove a limit that is not present. @@ -203,17 +207,17 @@ pub mod pallet { } #[pallet::genesis_config] - pub struct GenesisConfig { + pub struct GenesisConfig, I: 'static = ()> { pub default_limit: BlockNumberFor, pub limits: Vec<( TransactionIdentifier, - Option<::LimitContext>, + Option<>::LimitContext>, RateLimit>, )>, } #[cfg(feature = "std")] - impl Default for GenesisConfig { + impl, I: 'static> Default for GenesisConfig { fn default() -> Self { Self { default_limit: Zero::zero(), @@ -223,26 +227,26 @@ pub mod pallet { } #[pallet::genesis_build] - impl BuildGenesisConfig for GenesisConfig { + impl, I: 'static> BuildGenesisConfig for GenesisConfig { fn build(&self) { - DefaultLimit::::put(self.default_limit); + DefaultLimit::::put(self.default_limit); for (identifier, context, limit) in &self.limits { - Limits::::insert(identifier, context.clone(), limit.clone()); + Limits::::insert(identifier, context.clone(), limit.clone()); } } } #[pallet::pallet] #[pallet::without_storage_info] - pub struct Pallet(_); + pub struct Pallet(PhantomData<(T, I)>); - impl Pallet { + impl, I: 'static> Pallet { /// Returns `true` when the given transaction identifier passes its configured rate limit /// within the provided context. pub fn is_within_limit( identifier: &TransactionIdentifier, - context: &Option<::LimitContext>, + context: &Option<>::LimitContext>, ) -> Result { let Some(block_span) = Self::resolved_limit(identifier, context) else { return Ok(true); @@ -250,7 +254,7 @@ pub mod pallet { let current = frame_system::Pallet::::block_number(); - if let Some(last) = LastSeen::::get(identifier, context) { + if let Some(last) = LastSeen::::get(identifier, context) { let delta = current.saturating_sub(last); if delta < block_span { return Ok(false); @@ -262,13 +266,14 @@ pub mod pallet { pub(crate) fn resolved_limit( identifier: &TransactionIdentifier, - context: &Option<::LimitContext>, + context: &Option<>::LimitContext>, ) -> Option> { - let lookup = Limits::::get(identifier, context.clone()) - .or_else(|| Limits::::get(identifier, None::<::LimitContext>)); + let lookup = Limits::::get(identifier, context).or_else(|| { + Limits::::get(identifier, None::<>::LimitContext>) + }); let limit = lookup?; Some(match limit { - RateLimit::Default => DefaultLimit::::get(), + RateLimit::Default => DefaultLimit::::get(), RateLimit::Exact(block_span) => block_span, }) } @@ -277,18 +282,19 @@ pub mod pallet { pub fn limit_for_call_names( pallet_name: &str, extrinsic_name: &str, - context: Option<::LimitContext>, + context: Option<>::LimitContext>, ) -> Option>> { let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; - Limits::::get(&identifier, context.clone()) - .or_else(|| Limits::::get(&identifier, None::<::LimitContext>)) + Limits::::get(&identifier, context.clone()).or_else(|| { + Limits::::get(&identifier, None::<>::LimitContext>) + }) } /// Returns the resolved block span for the specified pallet/extrinsic names, if any. pub fn resolved_limit_for_call_names( pallet_name: &str, extrinsic_name: &str, - context: Option<::LimitContext>, + context: Option<>::LimitContext>, ) -> Option> { let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; Self::resolved_limit(&identifier, &context) @@ -298,9 +304,9 @@ pub mod pallet { pallet_name: &str, extrinsic_name: &str, ) -> Option { - let modules = ::RuntimeCall::get_module_names(); + let modules = >::RuntimeCall::get_module_names(); let pallet_pos = modules.iter().position(|name| *name == pallet_name)?; - let call_names = ::RuntimeCall::get_call_names(pallet_name); + let call_names = >::RuntimeCall::get_call_names(pallet_name); let extrinsic_pos = call_names.iter().position(|name| *name == extrinsic_name)?; let pallet_index = u8::try_from(pallet_pos).ok()?; let extrinsic_index = u8::try_from(extrinsic_pos).ok()?; @@ -309,7 +315,7 @@ pub mod pallet { } #[pallet::call] - impl Pallet { + impl, I: 'static> Pallet { /// Sets the rate limit configuration for the given call and optional context. /// /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any @@ -320,16 +326,16 @@ pub mod pallet { #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn set_rate_limit( origin: OriginFor, - call: Box<::RuntimeCall>, + call: Box<>::RuntimeCall>, limit: RateLimit>, - context: Option<::LimitContext>, + context: Option<>::LimitContext>, ) -> DispatchResult { ensure_root(origin)?; - let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - Limits::::insert(&identifier, context.clone(), limit.clone()); + let identifier = TransactionIdentifier::from_call::(call.as_ref())?; + Limits::::insert(&identifier, context.clone(), limit.clone()); - let (pallet_name, extrinsic_name) = identifier.names::()?; + let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); let extrinsic = Vec::from(extrinsic_name.as_bytes()); @@ -353,20 +359,20 @@ pub mod pallet { #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn clear_rate_limit( origin: OriginFor, - call: Box<::RuntimeCall>, - context: Option<::LimitContext>, + call: Box<>::RuntimeCall>, + context: Option<>::LimitContext>, ) -> DispatchResult { ensure_root(origin)?; - let identifier = TransactionIdentifier::from_call::(call.as_ref())?; + let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - let (pallet_name, extrinsic_name) = identifier.names::()?; + let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); let extrinsic = Vec::from(extrinsic_name.as_bytes()); ensure!( - Limits::::take(&identifier, context.clone()).is_some(), - Error::::MissingRateLimit + Limits::::take(&identifier, context.clone()).is_some(), + Error::::MissingRateLimit ); Self::deposit_event(Event::RateLimitCleared { @@ -388,7 +394,7 @@ pub mod pallet { ) -> DispatchResult { ensure_root(origin)?; - DefaultLimit::::put(block_span); + DefaultLimit::::put(block_span); Self::deposit_event(Event::DefaultRateLimitSet { block_span }); diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs index 4218e757a1..b80862edd5 100644 --- a/pallets/rate-limiting/src/mock.rs +++ b/pallets/rate-limiting/src/mock.rs @@ -95,7 +95,7 @@ pub fn new_test_ext() -> TestExternalities { } pub(crate) fn identifier_for(call: &RuntimeCall) -> TransactionIdentifier { - TransactionIdentifier::from_call::(call).expect("identifier for call") + TransactionIdentifier::from_call::(call).expect("identifier for call") } pub(crate) fn pop_last_event() -> RuntimeEvent { diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index ce5b6d25ab..62aae069db 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -18,7 +18,7 @@ fn limit_for_call_names_returns_stored_limit() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(7)); + Limits::::insert(identifier, None::, RateLimit::Exact(7)); let fetched = RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", None) @@ -33,8 +33,8 @@ fn limit_for_call_names_prefers_context_specific_limit() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(3)); - Limits::::insert(identifier, Some(5), RateLimit::Exact(8)); + Limits::::insert(identifier, None::, RateLimit::Exact(3)); + Limits::::insert(identifier, Some(5), RateLimit::Exact(8)); let fetched = RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", Some(5)) @@ -51,11 +51,11 @@ fn limit_for_call_names_prefers_context_specific_limit() { #[test] fn resolved_limit_for_call_names_resolves_default_value() { new_test_ext().execute_with(|| { - DefaultLimit::::put(3); + DefaultLimit::::put(3); let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Default); + Limits::::insert(identifier, None::, RateLimit::Default); let resolved = RateLimiting::resolved_limit_for_call_names( "RateLimiting", @@ -73,8 +73,8 @@ fn resolved_limit_for_call_names_prefers_context_specific_value() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(4)); - Limits::::insert(identifier, Some(6), RateLimit::Exact(9)); + Limits::::insert(identifier, None::, RateLimit::Exact(4)); + Limits::::insert(identifier, Some(6), RateLimit::Exact(9)); let resolved = RateLimiting::resolved_limit_for_call_names( "RateLimiting", @@ -126,8 +126,8 @@ fn is_within_limit_false_when_rate_limited() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, Some(1 as LimitContext), RateLimit::Exact(5)); - LastSeen::::insert(identifier, Some(1 as LimitContext), 9); + Limits::::insert(identifier, Some(1 as LimitContext), RateLimit::Exact(5)); + LastSeen::::insert(identifier, Some(1 as LimitContext), 9); System::set_block_number(13); @@ -143,8 +143,8 @@ fn is_within_limit_true_after_required_span() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, Some(2 as LimitContext), RateLimit::Exact(5)); - LastSeen::::insert(identifier, Some(2 as LimitContext), 10); + Limits::::insert(identifier, Some(2 as LimitContext), RateLimit::Exact(5)); + LastSeen::::insert(identifier, Some(2 as LimitContext), 10); System::set_block_number(20); @@ -172,7 +172,7 @@ fn set_rate_limit_updates_storage_and_emits_event() { let identifier = identifier_for(&target_call); assert_eq!( - Limits::::get(identifier, None::), + Limits::::get(identifier, None::), Some(limit) ); @@ -210,11 +210,14 @@ fn set_rate_limit_supports_context_specific_limit() { let identifier = identifier_for(&target_call); assert_eq!( - Limits::::get(identifier, Some(7)), + Limits::::get(identifier, Some(7)), Some(RateLimit::Exact(11)) ); // global remains untouched - assert_eq!(Limits::::get(identifier, None::), None); + assert_eq!( + Limits::::get(identifier, None::), + None + ); }); } @@ -251,7 +254,7 @@ fn set_rate_limit_accepts_default_variant() { let identifier = identifier_for(&target_call); assert_eq!( - Limits::::get(identifier, None::), + Limits::::get(identifier, None::), Some(RateLimit::Default) ); }); @@ -265,7 +268,7 @@ fn clear_rate_limit_removes_entry_and_emits_event() { let target_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&target_call); - Limits::::insert(identifier, None::, RateLimit::Exact(4)); + Limits::::insert(identifier, None::, RateLimit::Exact(4)); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), @@ -273,7 +276,10 @@ fn clear_rate_limit_removes_entry_and_emits_event() { None, )); - assert_eq!(Limits::::get(identifier, None::), None); + assert_eq!( + Limits::::get(identifier, None::), + None + ); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { @@ -300,7 +306,7 @@ fn clear_rate_limit_fails_when_missing() { assert_noop!( RateLimiting::clear_rate_limit(RuntimeOrigin::root(), Box::new(target_call), None), - Error::::MissingRateLimit + Error::::MissingRateLimit ); }); } @@ -313,8 +319,8 @@ fn clear_rate_limit_removes_only_selected_context() { let target_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&target_call); - Limits::::insert(identifier, None::, RateLimit::Exact(5)); - Limits::::insert(identifier, Some(9), RateLimit::Exact(7)); + Limits::::insert(identifier, None::, RateLimit::Exact(5)); + Limits::::insert(identifier, Some(9), RateLimit::Exact(7)); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), @@ -322,9 +328,9 @@ fn clear_rate_limit_removes_only_selected_context() { Some(9), )); - assert_eq!(Limits::::get(identifier, Some(9u16)), None); + assert_eq!(Limits::::get(identifier, Some(9u16)), None); assert_eq!( - Limits::::get(identifier, None::), + Limits::::get(identifier, None::), Some(RateLimit::Exact(5)) ); @@ -352,7 +358,7 @@ fn set_default_rate_limit_updates_storage_and_emits_event() { 42 )); - assert_eq!(DefaultLimit::::get(), 42); + assert_eq!(DefaultLimit::::get(), 42); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::DefaultRateLimitSet { diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index af02bfa453..119ad9c707 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -27,48 +27,90 @@ const IDENTIFIER: &str = "RateLimitTransactionExtension"; const RATE_LIMIT_DENIED: u8 = 1; /// Transaction extension that enforces pallet rate limiting rules. -#[derive(Default, Encode, Decode, DecodeWithMemTracking, Clone, Eq, PartialEq, TypeInfo)] -pub struct RateLimitTransactionExtension(PhantomData); +#[derive(Default, Encode, Decode, DecodeWithMemTracking, TypeInfo)] +pub struct RateLimitTransactionExtension(PhantomData<(T, I)>) +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo; + +impl Clone for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ + fn clone(&self) -> Self { + Self(PhantomData) + } +} + +impl PartialEq for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ + fn eq(&self, _other: &Self) -> bool { + true + } +} -impl core::fmt::Debug for RateLimitTransactionExtension { +impl Eq for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ +} + +impl core::fmt::Debug for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.write_str(IDENTIFIER) } } -impl TransactionExtension<::RuntimeCall> for RateLimitTransactionExtension +impl TransactionExtension<>::RuntimeCall> + for RateLimitTransactionExtension where - T: Config + Send + Sync + TypeInfo, - ::RuntimeCall: Dispatchable, + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo + Send + Sync, + >::RuntimeCall: Dispatchable, { const IDENTIFIER: &'static str = IDENTIFIER; type Implicit = (); - type Val = Option<(TransactionIdentifier, Option)>; - type Pre = Option<(TransactionIdentifier, Option)>; - - fn weight(&self, _call: &::RuntimeCall) -> Weight { + type Val = Option<( + TransactionIdentifier, + Option<>::LimitContext>, + )>; + type Pre = Option<( + TransactionIdentifier, + Option<>::LimitContext>, + )>; + + fn weight(&self, _call: &>::RuntimeCall) -> Weight { Weight::zero() } fn validate( &self, - origin: DispatchOriginOf<::RuntimeCall>, - call: &::RuntimeCall, - _info: &DispatchInfoOf<::RuntimeCall>, + origin: DispatchOriginOf<>::RuntimeCall>, + call: &>::RuntimeCall, + _info: &DispatchInfoOf<>::RuntimeCall>, _len: usize, _self_implicit: Self::Implicit, _inherited_implication: &impl Implication, _source: TransactionSource, - ) -> ValidateResult::RuntimeCall> { - let identifier = match TransactionIdentifier::from_call::(call) { + ) -> ValidateResult>::RuntimeCall> { + let identifier = match TransactionIdentifier::from_call::(call) { Ok(identifier) => identifier, Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)), }; - let context = ::ContextResolver::context(call); + let context = >::ContextResolver::context(call); - let Some(block_span) = Pallet::::resolved_limit(&identifier, &context) else { + let Some(block_span) = Pallet::::resolved_limit(&identifier, &context) else { return Ok((ValidTransaction::default(), None, origin)); }; @@ -76,7 +118,7 @@ where return Ok((ValidTransaction::default(), None, origin)); } - let within_limit = Pallet::::is_within_limit(&identifier, &context) + let within_limit = Pallet::::is_within_limit(&identifier, &context) .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; if !within_limit { @@ -95,9 +137,9 @@ where fn prepare( self, val: Self::Val, - _origin: &DispatchOriginOf<::RuntimeCall>, - _call: &::RuntimeCall, - _info: &DispatchInfoOf<::RuntimeCall>, + _origin: &DispatchOriginOf<>::RuntimeCall>, + _call: &>::RuntimeCall, + _info: &DispatchInfoOf<>::RuntimeCall>, _len: usize, ) -> Result { Ok(val) @@ -105,7 +147,7 @@ where fn post_dispatch( pre: Self::Pre, - _info: &DispatchInfoOf<::RuntimeCall>, + _info: &DispatchInfoOf<>::RuntimeCall>, _post_info: &mut PostDispatchInfo, _len: usize, result: &DispatchResult, @@ -113,7 +155,7 @@ where if result.is_ok() { if let Some((identifier, context)) = pre { let block_number = frame_system::Pallet::::block_number(); - LastSeen::::insert(&identifier, context, block_number); + LastSeen::::insert(&identifier, context, block_number); } } Ok(()) @@ -196,7 +238,7 @@ mod tests { let identifier = identifier_for(&call); assert_eq!( - LastSeen::::get(identifier, None::), + LastSeen::::get(identifier, None::), None ); }); @@ -208,7 +250,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(5)); + Limits::::insert(identifier, None::, RateLimit::Exact(5)); System::set_block_number(10); @@ -234,7 +276,7 @@ mod tests { .expect("post_dispatch succeeds"); assert_eq!( - LastSeen::::get(identifier, None::), + LastSeen::::get(identifier, None::), Some(10) ); }); @@ -246,8 +288,8 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(5)); - LastSeen::::insert(identifier, None::, 20); + Limits::::insert(identifier, None::, RateLimit::Exact(5)); + LastSeen::::insert(identifier, None::, 20); System::set_block_number(22); @@ -268,7 +310,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(0)); + Limits::::insert(identifier, None::, RateLimit::Exact(0)); System::set_block_number(30); @@ -294,7 +336,7 @@ mod tests { .expect("post_dispatch succeeds"); assert_eq!( - LastSeen::::get(identifier, None::), + LastSeen::::get(identifier, None::), None ); }); diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index 72e2a43777..7d53a34ac6 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -40,38 +40,40 @@ impl TransactionIdentifier { } /// Returns the pallet and extrinsic names associated with this identifier. - pub fn names(&self) -> Result<(&'static str, &'static str), DispatchError> + pub fn names(&self) -> Result<(&'static str, &'static str), DispatchError> where - T: crate::pallet::Config, - ::RuntimeCall: GetCallMetadata, + T: crate::pallet::Config, + I: 'static, + >::RuntimeCall: GetCallMetadata, { - let modules = ::RuntimeCall::get_module_names(); + let modules = >::RuntimeCall::get_module_names(); let pallet_name = modules .get(self.pallet_index as usize) .copied() - .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; - let call_names = ::RuntimeCall::get_call_names(pallet_name); + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + let call_names = >::RuntimeCall::get_call_names(pallet_name); let extrinsic_name = call_names .get(self.extrinsic_index as usize) .copied() - .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; Ok((pallet_name, extrinsic_name)) } /// Builds an identifier from a runtime call by extracting pallet/extrinsic indices. - pub fn from_call( - call: &::RuntimeCall, + pub fn from_call( + call: &>::RuntimeCall, ) -> Result where - T: crate::pallet::Config, + T: crate::pallet::Config, + I: 'static, { call.using_encoded(|encoded| { let pallet_index = *encoded .get(0) - .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; let extrinsic_index = *encoded .get(1) - .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; + .ok_or(crate::pallet::Error::::InvalidRuntimeCall)?; Ok(TransactionIdentifier::new(pallet_index, extrinsic_index)) }) } @@ -110,7 +112,7 @@ mod tests { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); - let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); + let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); // System is the first pallet in the mock runtime, RateLimiting is second. assert_eq!(identifier.pallet_index, 1); @@ -122,9 +124,9 @@ mod tests { fn transaction_identifier_names_matches_call_metadata() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); - let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); + let identifier = TransactionIdentifier::from_call::(&call).expect("identifier"); - let (pallet, extrinsic) = identifier.names::().expect("call metadata"); + let (pallet, extrinsic) = identifier.names::().expect("call metadata"); assert_eq!(pallet, "RateLimiting"); assert_eq!(extrinsic, "set_default_rate_limit"); } @@ -133,8 +135,8 @@ mod tests { fn transaction_identifier_names_error_for_unknown_indices() { let identifier = TransactionIdentifier::new(99, 0); - let err = identifier.names::().expect_err("should fail"); - let expected: DispatchError = Error::::InvalidRuntimeCall.into(); + let err = identifier.names::().expect_err("should fail"); + let expected: DispatchError = Error::::InvalidRuntimeCall.into(); assert_eq!(err, expected); } } From 61eb4d286d6621d21e428841f6947cc00d3d77d7 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Thu, 30 Oct 2025 15:56:19 +0300 Subject: [PATCH 13/23] Refactor RateLimit type --- pallets/rate-limiting/runtime-api/src/lib.rs | 5 +- pallets/rate-limiting/src/benchmarking.rs | 12 +- pallets/rate-limiting/src/lib.rs | 143 ++++++++++++------- pallets/rate-limiting/src/mock.rs | 9 +- pallets/rate-limiting/src/tests.rs | 120 ++++++++++------ pallets/rate-limiting/src/tx_extension.rs | 11 +- pallets/rate-limiting/src/types.rs | 76 +++++++++- 7 files changed, 263 insertions(+), 113 deletions(-) diff --git a/pallets/rate-limiting/runtime-api/src/lib.rs b/pallets/rate-limiting/runtime-api/src/lib.rs index 1a32c094ea..98b55e9a26 100644 --- a/pallets/rate-limiting/runtime-api/src/lib.rs +++ b/pallets/rate-limiting/runtime-api/src/lib.rs @@ -1,7 +1,7 @@ #![cfg_attr(not(feature = "std"), no_std)] use codec::{Decode, Encode}; -use pallet_rate_limiting::RateLimit; +use pallet_rate_limiting::RateLimitKind; use scale_info::TypeInfo; use sp_std::vec::Vec; use subtensor_runtime_common::BlockNumber; @@ -12,7 +12,8 @@ use serde::{Deserialize, Serialize}; #[cfg_attr(feature = "std", derive(Serialize, Deserialize))] #[derive(Clone, Debug, Decode, Encode, Eq, PartialEq, TypeInfo)] pub struct RateLimitRpcResponse { - pub limit: Option>, + pub global: Option>, + pub contextual: Vec<(Vec, RateLimitKind)>, pub default_limit: BlockNumber, pub resolved: Option, } diff --git a/pallets/rate-limiting/src/benchmarking.rs b/pallets/rate-limiting/src/benchmarking.rs index 8392109772..4c9ce17708 100644 --- a/pallets/rate-limiting/src/benchmarking.rs +++ b/pallets/rate-limiting/src/benchmarking.rs @@ -35,7 +35,7 @@ mod benchmarks { #[benchmark] fn set_rate_limit() { let call = sample_call::(); - let limit = RateLimit::>::Exact(BlockNumberFor::::from(10u32)); + let limit = RateLimitKind::>::Exact(BlockNumberFor::::from(10u32)); let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); @@ -43,25 +43,25 @@ mod benchmarks { _(RawOrigin::Root, call, limit.clone(), None); assert_eq!( - Limits::::get(&identifier, None::<::LimitContext>), - Some(limit) + Limits::::get(&identifier), + Some(RateLimit::global(limit)) ); } #[benchmark] fn clear_rate_limit() { let call = sample_call::(); - let limit = RateLimit::>::Exact(BlockNumberFor::::from(10u32)); + let limit = RateLimitKind::>::Exact(BlockNumberFor::::from(10u32)); // Pre-populate limit for benchmark call let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); - Limits::::insert(identifier, None::<::LimitContext>, limit); + Limits::::insert(identifier, RateLimit::global(limit)); #[extrinsic_call] _(RawOrigin::Root, call, None); - assert!(Limits::::get(identifier, None::<::LimitContext>).is_none()); + assert!(Limits::::get(identifier).is_none()); } #[benchmark] diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 2ca3ab87ba..cf0ebeefb6 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -6,24 +6,20 @@ //! //! `pallet-rate-limiting` lets a runtime restrict how frequently particular calls can execute. //! Limits are stored on-chain, keyed by the call's pallet/variant pair. Each entry can specify an -//! exact block span or defer to a configured default. The pallet exposes three roots-only -//! extrinsics to manage this data: +//! exact block span or defer to a configured default. The pallet exposes three extrinsics, +//! restricted by [`Config::AdminOrigin`], to manage this data: //! -//! - [`set_rate_limit`](pallet::Pallet::set_rate_limit): assign a limit to an extrinsic, either as -//! `RateLimit::Exact(blocks)` or `RateLimit::Default`. The optional context parameter lets you -//! scope the configuration to a particular subnet/key/account while keeping a global fallback. +//! - [`set_rate_limit`](pallet::Pallet::set_rate_limit): assign a limit to an extrinsic by +//! supplying a [`RateLimitKind`] span and optionally a contextual identifier. When a contextual +//! span is stored, any previously configured global span is replaced. //! - [`clear_rate_limit`](pallet::Pallet::clear_rate_limit): remove a stored limit for the provided -//! context (or for the global entry when `None` is supplied). +//! scope (either the global entry when `None` is supplied, or a specific context). //! - [`set_default_rate_limit`](pallet::Pallet::set_default_rate_limit): set the global default -//! block span used by `RateLimit::Default` entries. +//! block span used by `RateLimitKind::Default` entries. //! //! The pallet also tracks the last block in which a rate-limited call was executed, per optional //! *context*. Context allows one limit definition (for example, “set weights”) to be enforced per -//! subnet, account, or other grouping chosen by the runtime. The storage layout is: -//! -//! - [`Limits`](pallet::Limits): `TransactionIdentifier → RateLimit` -//! - [`DefaultLimit`](pallet::DefaultLimit): `BlockNumber` -//! - [`LastSeen`](pallet::LastSeen): `(TransactionIdentifier, Option) → BlockNumber` +//! subnet, account, or other grouping chosen by the runtime. //! //! Each storage map is namespaced by pallet instance; runtimes can deploy multiple independent //! instances to manage distinct rate-limiting scopes. @@ -74,6 +70,7 @@ //! type RuntimeCall = RuntimeCall; //! type LimitContext = NetUid; //! type ContextResolver = WeightsContextResolver; +//! type AdminOrigin = frame_system::EnsureRoot; //! } //! ``` @@ -81,7 +78,7 @@ pub use benchmarking::BenchmarkHelper; pub use pallet::*; pub use tx_extension::RateLimitTransactionExtension; -pub use types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; +pub use types::{RateLimit, RateLimitContextResolver, RateLimitKind, TransactionIdentifier}; #[cfg(feature = "runtime-benchmarks")] mod benchmarking; @@ -100,26 +97,32 @@ pub mod pallet { use frame_support::{ pallet_prelude::*, sp_runtime::traits::{Saturating, Zero}, - traits::{BuildGenesisConfig, GetCallMetadata}, + traits::{BuildGenesisConfig, EnsureOrigin, GetCallMetadata}, }; - use frame_system::{ensure_root, pallet_prelude::*}; + use frame_system::pallet_prelude::*; use sp_std::{convert::TryFrom, marker::PhantomData, vec::Vec}; #[cfg(feature = "runtime-benchmarks")] use crate::benchmarking::BenchmarkHelper as BenchmarkHelperTrait; - use crate::types::{RateLimit, RateLimitContextResolver, TransactionIdentifier}; + use crate::types::{RateLimit, RateLimitContextResolver, RateLimitKind, TransactionIdentifier}; /// Configuration trait for the rate limiting pallet. #[pallet::config] - pub trait Config: frame_system::Config { + pub trait Config: frame_system::Config + where + BlockNumberFor: MaybeSerializeDeserialize, + { /// The overarching runtime call type. type RuntimeCall: Parameter + Codec + GetCallMetadata + IsType<::RuntimeCall>; + /// Origin permitted to configure rate limits. + type AdminOrigin: EnsureOrigin>; + /// Context type used for contextual (per-group) rate limits. - type LimitContext: Parameter + Clone + PartialEq + Eq + MaybeSerializeDeserialize; + type LimitContext: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; /// Resolves the context for a given runtime call. type ContextResolver: RateLimitContextResolver<>::RuntimeCall, Self::LimitContext>; @@ -129,16 +132,14 @@ pub mod pallet { type BenchmarkHelper: BenchmarkHelperTrait<>::RuntimeCall>; } - /// Storage mapping from transaction identifier and optional context to its configured rate limit. + /// Storage mapping from transaction identifier to its configured rate limit. #[pallet::storage] #[pallet::getter(fn limits)] - pub type Limits, I: 'static = ()> = StorageDoubleMap< + pub type Limits, I: 'static = ()> = StorageMap< _, Blake2_128Concat, TransactionIdentifier, - Blake2_128Concat, - Option<>::LimitContext>, - RateLimit>, + RateLimit<>::LimitContext, BlockNumberFor>, OptionQuery, >; @@ -172,8 +173,8 @@ pub mod pallet { transaction: TransactionIdentifier, /// Context to which the limit applies, if any. context: Option<>::LimitContext>, - /// The new limit configuration applied to the transaction. - limit: RateLimit>, + /// The rate limit policy applied to the transaction. + limit: RateLimitKind>, /// Pallet name associated with the transaction. pallet: Vec, /// Extrinsic name associated with the transaction. @@ -204,6 +205,8 @@ pub mod pallet { InvalidRuntimeCall, /// Attempted to remove a limit that is not present. MissingRateLimit, + /// Contextual configuration was requested but no context can be resolved for the call. + ContextUnavailable, } #[pallet::genesis_config] @@ -212,7 +215,7 @@ pub mod pallet { pub limits: Vec<( TransactionIdentifier, Option<>::LimitContext>, - RateLimit>, + RateLimitKind>, )>, } @@ -231,8 +234,19 @@ pub mod pallet { fn build(&self) { DefaultLimit::::put(self.default_limit); - for (identifier, context, limit) in &self.limits { - Limits::::insert(identifier, context.clone(), limit.clone()); + for (identifier, context, kind) in &self.limits { + Limits::::mutate(identifier, |entry| match context { + None => { + *entry = Some(RateLimit::global(*kind)); + } + Some(ctx) => { + if let Some(config) = entry { + config.upsert_context(ctx.clone(), *kind); + } else { + *entry = Some(RateLimit::contextual_single(ctx.clone(), *kind)); + } + } + }); } } } @@ -268,13 +282,11 @@ pub mod pallet { identifier: &TransactionIdentifier, context: &Option<>::LimitContext>, ) -> Option> { - let lookup = Limits::::get(identifier, context).or_else(|| { - Limits::::get(identifier, None::<>::LimitContext>) - }); - let limit = lookup?; - Some(match limit { - RateLimit::Default => DefaultLimit::::get(), - RateLimit::Exact(block_span) => block_span, + let config = Limits::::get(identifier)?; + let kind = config.kind_for(context.as_ref())?; + Some(match *kind { + RateLimitKind::Default => DefaultLimit::::get(), + RateLimitKind::Exact(block_span) => block_span, }) } @@ -283,11 +295,10 @@ pub mod pallet { pallet_name: &str, extrinsic_name: &str, context: Option<>::LimitContext>, - ) -> Option>> { + ) -> Option>> { let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; - Limits::::get(&identifier, context.clone()).or_else(|| { - Limits::::get(&identifier, None::<>::LimitContext>) - }) + Limits::::get(&identifier) + .and_then(|config| config.kind_for(context.as_ref()).copied()) } /// Returns the resolved block span for the specified pallet/extrinsic names, if any. @@ -327,13 +338,28 @@ pub mod pallet { pub fn set_rate_limit( origin: OriginFor, call: Box<>::RuntimeCall>, - limit: RateLimit>, + limit: RateLimitKind>, context: Option<>::LimitContext>, ) -> DispatchResult { - ensure_root(origin)?; + T::AdminOrigin::ensure_origin(origin)?; + + if context.is_some() + && >::ContextResolver::context(call.as_ref()).is_none() + { + return Err(Error::::ContextUnavailable.into()); + } let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - Limits::::insert(&identifier, context.clone(), limit.clone()); + let context_for_event = context.clone(); + + if let Some(ref ctx) = context { + Limits::::mutate(&identifier, |slot| match slot { + Some(config) => config.upsert_context(ctx.clone(), limit), + None => *slot = Some(RateLimit::contextual_single(ctx.clone(), limit)), + }); + } else { + Limits::::insert(&identifier, RateLimit::global(limit)); + } let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); @@ -341,12 +367,11 @@ pub mod pallet { Self::deposit_event(Event::RateLimitSet { transaction: identifier, - context, + context: context_for_event, limit, pallet, extrinsic, }); - Ok(()) } @@ -362,7 +387,7 @@ pub mod pallet { call: Box<>::RuntimeCall>, context: Option<>::LimitContext>, ) -> DispatchResult { - ensure_root(origin)?; + T::AdminOrigin::ensure_origin(origin)?; let identifier = TransactionIdentifier::from_call::(call.as_ref())?; @@ -370,10 +395,28 @@ pub mod pallet { let pallet = Vec::from(pallet_name.as_bytes()); let extrinsic = Vec::from(extrinsic_name.as_bytes()); - ensure!( - Limits::::take(&identifier, context.clone()).is_some(), - Error::::MissingRateLimit - ); + let mut removed = false; + Limits::::mutate_exists(&identifier, |maybe_config| { + if let Some(config) = maybe_config { + match (&context, config) { + (None, _) => { + removed = true; + *maybe_config = None; + } + (Some(ctx), RateLimit::Contextual(map)) => { + if map.remove(ctx).is_some() { + removed = true; + if map.is_empty() { + *maybe_config = None; + } + } + } + (Some(_), RateLimit::Global(_)) => {} + } + } + }); + + ensure!(removed, Error::::MissingRateLimit); Self::deposit_event(Event::RateLimitCleared { transaction: identifier, @@ -392,7 +435,7 @@ pub mod pallet { origin: OriginFor, block_span: BlockNumberFor, ) -> DispatchResult { - ensure_root(origin)?; + T::AdminOrigin::ensure_origin(origin)?; DefaultLimit::::put(block_span); diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs index b80862edd5..1c792dde16 100644 --- a/pallets/rate-limiting/src/mock.rs +++ b/pallets/rate-limiting/src/mock.rs @@ -8,6 +8,7 @@ use frame_support::{ }, traits::{ConstU16, ConstU32, ConstU64, Everything}, }; +use frame_system::EnsureRoot; use sp_core::H256; use sp_io::TestExternalities; use sp_std::vec::Vec; @@ -59,8 +60,11 @@ pub struct TestContextResolver; impl pallet_rate_limiting::RateLimitContextResolver for TestContextResolver { - fn context(_call: &RuntimeCall) -> Option { - None + fn context(call: &RuntimeCall) -> Option { + match call { + RuntimeCall::RateLimiting(_) => Some(1), + _ => None, + } } } @@ -68,6 +72,7 @@ impl pallet_rate_limiting::Config for Test { type RuntimeCall = RuntimeCall; type LimitContext = LimitContext; type ContextResolver = TestContextResolver; + type AdminOrigin = EnsureRoot; #[cfg(feature = "runtime-benchmarks")] type BenchmarkHelper = BenchHelper; } diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index 62aae069db..27bf6e5472 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -1,6 +1,7 @@ use frame_support::{assert_noop, assert_ok, error::BadOrigin}; +use sp_std::{collections::btree_map::BTreeMap, vec::Vec}; -use crate::{DefaultLimit, LastSeen, Limits, RateLimit, mock::*, pallet::Error}; +use crate::{DefaultLimit, LastSeen, Limits, RateLimit, RateLimitKind, mock::*, pallet::Error}; #[test] fn limit_for_call_names_returns_none_if_not_set() { @@ -18,12 +19,12 @@ fn limit_for_call_names_returns_stored_limit() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(7)); + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(7))); let fetched = RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", None) .expect("limit should exist"); - assert_eq!(fetched, RateLimit::Exact(7)); + assert_eq!(fetched, RateLimitKind::Exact(7)); }); } @@ -33,18 +34,20 @@ fn limit_for_call_names_prefers_context_specific_limit() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(3)); - Limits::::insert(identifier, Some(5), RateLimit::Exact(8)); + Limits::::insert( + identifier, + RateLimit::contextual_single(5u16, RateLimitKind::Exact(8)), + ); let fetched = RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", Some(5)) .expect("limit should exist"); - assert_eq!(fetched, RateLimit::Exact(8)); + assert_eq!(fetched, RateLimitKind::Exact(8)); - let fallback = + assert!( RateLimiting::limit_for_call_names("RateLimiting", "set_default_rate_limit", Some(1)) - .expect("limit should exist"); - assert_eq!(fallback, RateLimit::Exact(3)); + .is_none() + ); }); } @@ -55,7 +58,7 @@ fn resolved_limit_for_call_names_resolves_default_value() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Default); + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Default)); let resolved = RateLimiting::resolved_limit_for_call_names( "RateLimiting", @@ -73,8 +76,10 @@ fn resolved_limit_for_call_names_prefers_context_specific_value() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(4)); - Limits::::insert(identifier, Some(6), RateLimit::Exact(9)); + let mut map = BTreeMap::new(); + map.insert(6u16, RateLimitKind::Exact(9)); + map.insert(2u16, RateLimitKind::Exact(4)); + Limits::::insert(identifier, RateLimit::Contextual(map)); let resolved = RateLimiting::resolved_limit_for_call_names( "RateLimiting", @@ -84,13 +89,14 @@ fn resolved_limit_for_call_names_prefers_context_specific_value() { .expect("resolved limit"); assert_eq!(resolved, 9); - let fallback = RateLimiting::resolved_limit_for_call_names( - "RateLimiting", - "set_default_rate_limit", - Some(1), - ) - .expect("resolved limit"); - assert_eq!(fallback, 4); + assert!( + RateLimiting::resolved_limit_for_call_names( + "RateLimiting", + "set_default_rate_limit", + Some(1), + ) + .is_none() + ); }); } @@ -126,7 +132,10 @@ fn is_within_limit_false_when_rate_limited() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, Some(1 as LimitContext), RateLimit::Exact(5)); + Limits::::insert( + identifier, + RateLimit::contextual_single(1 as LimitContext, RateLimitKind::Exact(5)), + ); LastSeen::::insert(identifier, Some(1 as LimitContext), 9); System::set_block_number(13); @@ -143,7 +152,10 @@ fn is_within_limit_true_after_required_span() { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); - Limits::::insert(identifier, Some(2 as LimitContext), RateLimit::Exact(5)); + Limits::::insert( + identifier, + RateLimit::contextual_single(2 as LimitContext, RateLimitKind::Exact(5)), + ); LastSeen::::insert(identifier, Some(2 as LimitContext), 10); System::set_block_number(20); @@ -161,7 +173,7 @@ fn set_rate_limit_updates_storage_and_emits_event() { let target_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); - let limit = RateLimit::Exact(9); + let limit = RateLimitKind::Exact(9); assert_ok!(RateLimiting::set_rate_limit( RuntimeOrigin::root(), @@ -172,8 +184,8 @@ fn set_rate_limit_updates_storage_and_emits_event() { let identifier = identifier_for(&target_call); assert_eq!( - Limits::::get(identifier, None::), - Some(limit) + Limits::::get(identifier), + Some(RateLimit::global(limit)) ); match pop_last_event() { @@ -204,19 +216,15 @@ fn set_rate_limit_supports_context_specific_limit() { assert_ok!(RateLimiting::set_rate_limit( RuntimeOrigin::root(), Box::new(target_call.clone()), - RateLimit::Exact(11), - context, + RateLimitKind::Exact(11), + context.clone(), )); let identifier = identifier_for(&target_call); + let config = Limits::::get(identifier).expect("config stored"); assert_eq!( - Limits::::get(identifier, Some(7)), - Some(RateLimit::Exact(11)) - ); - // global remains untouched - assert_eq!( - Limits::::get(identifier, None::), - None + config.kind_for(context.as_ref()).copied(), + Some(RateLimitKind::Exact(11)) ); }); } @@ -231,7 +239,7 @@ fn set_rate_limit_requires_root() { RateLimiting::set_rate_limit( RuntimeOrigin::signed(1), Box::new(target_call), - RateLimit::Exact(1), + RateLimitKind::Exact(1), None, ), BadOrigin @@ -248,14 +256,14 @@ fn set_rate_limit_accepts_default_variant() { assert_ok!(RateLimiting::set_rate_limit( RuntimeOrigin::root(), Box::new(target_call.clone()), - RateLimit::Default, + RateLimitKind::Default, None, )); let identifier = identifier_for(&target_call); assert_eq!( - Limits::::get(identifier, None::), - Some(RateLimit::Default) + Limits::::get(identifier), + Some(RateLimit::global(RateLimitKind::Default)) ); }); } @@ -268,7 +276,7 @@ fn clear_rate_limit_removes_entry_and_emits_event() { let target_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&target_call); - Limits::::insert(identifier, None::, RateLimit::Exact(4)); + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(4))); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), @@ -276,10 +284,7 @@ fn clear_rate_limit_removes_entry_and_emits_event() { None, )); - assert_eq!( - Limits::::get(identifier, None::), - None - ); + assert!(Limits::::get(identifier).is_none()); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { @@ -319,8 +324,10 @@ fn clear_rate_limit_removes_only_selected_context() { let target_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&target_call); - Limits::::insert(identifier, None::, RateLimit::Exact(5)); - Limits::::insert(identifier, Some(9), RateLimit::Exact(7)); + let mut map = BTreeMap::new(); + map.insert(9u16, RateLimitKind::Exact(7)); + map.insert(10u16, RateLimitKind::Exact(5)); + Limits::::insert(identifier, RateLimit::Contextual(map)); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), @@ -328,10 +335,11 @@ fn clear_rate_limit_removes_only_selected_context() { Some(9), )); - assert_eq!(Limits::::get(identifier, Some(9u16)), None); + let config = Limits::::get(identifier).expect("config remains"); + assert!(config.kind_for(Some(&9u16)).is_none()); assert_eq!( - Limits::::get(identifier, None::), - Some(RateLimit::Exact(5)) + config.kind_for(Some(&10u16)).copied(), + Some(RateLimitKind::Exact(5)) ); match pop_last_event() { @@ -348,6 +356,24 @@ fn clear_rate_limit_removes_only_selected_context() { }); } +#[test] +fn set_rate_limit_rejects_unresolvable_context() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }); + + assert_noop!( + RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + Box::new(target_call), + RateLimitKind::Exact(5), + Some(42), + ), + Error::::ContextUnavailable + ); + }); +} + #[test] fn set_default_rate_limit_updates_storage_and_emits_event() { new_test_ext().execute_with(|| { diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index 119ad9c707..e72b473fd4 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -171,7 +171,10 @@ mod tests { transaction_validity::{InvalidTransaction, TransactionSource, TransactionValidityError}, }; - use crate::{LastSeen, Limits, RateLimit, types::TransactionIdentifier}; + use crate::{ + LastSeen, Limits, + types::{RateLimit, RateLimitKind, TransactionIdentifier}, + }; use super::*; use crate::mock::*; @@ -250,7 +253,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(5)); + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(5))); System::set_block_number(10); @@ -288,7 +291,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(5)); + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(5))); LastSeen::::insert(identifier, None::, 20); System::set_block_number(22); @@ -310,7 +313,7 @@ mod tests { let extension = new_tx_extension(); let call = remark_call(); let identifier = identifier_for(&call); - Limits::::insert(identifier, None::, RateLimit::Exact(0)); + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(0))); System::set_block_number(30); diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index 7d53a34ac6..a18fff37c6 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -1,6 +1,7 @@ use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata}; use scale_info::TypeInfo; +use sp_std::collections::btree_map::BTreeMap; /// Resolves the optional context within which a rate limit applies. pub trait RateLimitContextResolver { @@ -79,7 +80,7 @@ impl TransactionIdentifier { } } -/// Configuration value for a rate limit. +/// Policy describing the block span enforced by a rate limit. #[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] #[derive( Clone, @@ -93,13 +94,84 @@ impl TransactionIdentifier { MaxEncodedLen, Debug, )] -pub enum RateLimit { +pub enum RateLimitKind { /// Use the pallet-level default rate limit. Default, /// Apply an exact rate limit measured in blocks. Exact(BlockNumber), } +/// Stored rate limit configuration for a transaction identifier. +/// +/// The configuration is mutually exclusive: either the call is globally limited or it stores a set +/// of per-context spans. +#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] +#[cfg_attr( + feature = "std", + serde( + bound = "Context: Ord + serde::Serialize + serde::de::DeserializeOwned, BlockNumber: serde::Serialize + serde::de::DeserializeOwned" + ) +)] +#[derive(Clone, PartialEq, Eq, Encode, Decode, DecodeWithMemTracking, TypeInfo, Debug)] +pub enum RateLimit { + /// Global span applied to every invocation. + Global(RateLimitKind), + /// Per-context spans keyed by `Context`. + Contextual(BTreeMap>), +} + +impl RateLimit +where + Context: Ord, +{ + /// Convenience helper to build a global configuration. + pub fn global(kind: RateLimitKind) -> Self { + Self::Global(kind) + } + + /// Convenience helper to build a contextual configuration containing a single entry. + pub fn contextual_single(context: Context, kind: RateLimitKind) -> Self { + let mut map = BTreeMap::new(); + map.insert(context, kind); + Self::Contextual(map) + } + + /// Returns the span configured for the provided context, if any. + pub fn kind_for(&self, context: Option<&Context>) -> Option<&RateLimitKind> { + match self { + RateLimit::Global(kind) => Some(kind), + RateLimit::Contextual(map) => context.and_then(|ctx| map.get(ctx)), + } + } + + /// Inserts or updates a contextual entry, converting from a global configuration if needed. + pub fn upsert_context(&mut self, context: Context, kind: RateLimitKind) { + match self { + RateLimit::Global(_) => { + let mut map = BTreeMap::new(); + map.insert(context, kind); + *self = RateLimit::Contextual(map); + } + RateLimit::Contextual(map) => { + map.insert(context, kind); + } + } + } + + /// Removes a contextual entry, returning whether one existed. + pub fn remove_context(&mut self, context: &Context) -> bool { + match self { + RateLimit::Global(_) => false, + RateLimit::Contextual(map) => map.remove(context).is_some(), + } + } + + /// Returns true when the contextual configuration contains no entries. + pub fn is_contextual_empty(&self) -> bool { + matches!(self, RateLimit::Contextual(map) if map.is_empty()) + } +} + #[cfg(test)] mod tests { use sp_runtime::DispatchError; From c84cb92ec13ccdc8f05e2140add73087a0dbb678 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Thu, 30 Oct 2025 17:17:19 +0300 Subject: [PATCH 14/23] Separate context between usage key and limit scope --- pallets/rate-limiting/src/benchmarking.rs | 29 ++-- pallets/rate-limiting/src/lib.rs | 166 +++++++++++++--------- pallets/rate-limiting/src/mock.rs | 35 +++-- pallets/rate-limiting/src/tests.rs | 128 +++++++++-------- pallets/rate-limiting/src/tx_extension.rs | 33 ++--- pallets/rate-limiting/src/types.rs | 54 +++---- 6 files changed, 252 insertions(+), 193 deletions(-) diff --git a/pallets/rate-limiting/src/benchmarking.rs b/pallets/rate-limiting/src/benchmarking.rs index 4c9ce17708..65d547ab0b 100644 --- a/pallets/rate-limiting/src/benchmarking.rs +++ b/pallets/rate-limiting/src/benchmarking.rs @@ -36,30 +36,43 @@ mod benchmarks { fn set_rate_limit() { let call = sample_call::(); let limit = RateLimitKind::>::Exact(BlockNumberFor::::from(10u32)); + let scope = ::LimitScopeResolver::context(call.as_ref()); let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); #[extrinsic_call] - _(RawOrigin::Root, call, limit.clone(), None); - - assert_eq!( - Limits::::get(&identifier), - Some(RateLimit::global(limit)) - ); + _(RawOrigin::Root, call, limit.clone()); + + let stored = Limits::::get(&identifier).expect("limit stored"); + match (scope, &stored) { + (Some(ref sc), RateLimit::Scoped(map)) => { + assert_eq!(map.get(sc), Some(&limit)); + } + (None, RateLimit::Global(kind)) | (Some(_), RateLimit::Global(kind)) => { + assert_eq!(kind, &limit); + } + (None, RateLimit::Scoped(map)) => { + assert!(map.values().any(|k| k == &limit)); + } + } } #[benchmark] fn clear_rate_limit() { let call = sample_call::(); let limit = RateLimitKind::>::Exact(BlockNumberFor::::from(10u32)); + let scope = ::LimitScopeResolver::context(call.as_ref()); // Pre-populate limit for benchmark call let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); - Limits::::insert(identifier, RateLimit::global(limit)); + match scope.clone() { + Some(sc) => Limits::::insert(identifier, RateLimit::scoped_single(sc, limit)), + None => Limits::::insert(identifier, RateLimit::global(limit)), + } #[extrinsic_call] - _(RawOrigin::Root, call, None); + _(RawOrigin::Root, call); assert!(Limits::::get(identifier).is_none()); } diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index cf0ebeefb6..0579249b36 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -10,16 +10,18 @@ //! restricted by [`Config::AdminOrigin`], to manage this data: //! //! - [`set_rate_limit`](pallet::Pallet::set_rate_limit): assign a limit to an extrinsic by -//! supplying a [`RateLimitKind`] span and optionally a contextual identifier. When a contextual -//! span is stored, any previously configured global span is replaced. -//! - [`clear_rate_limit`](pallet::Pallet::clear_rate_limit): remove a stored limit for the provided -//! scope (either the global entry when `None` is supplied, or a specific context). +//! supplying a [`RateLimitKind`] span. The pallet infers the *limit scope* (for example a +//! `netuid`) using [`Config::LimitScopeResolver`] and stores the configuration for that scope, or +//! globally when no scope is resolved. +//! - [`clear_rate_limit`](pallet::Pallet::clear_rate_limit): remove a stored limit for the scope +//! derived from the provided call (or the global entry when no scope resolves). //! - [`set_default_rate_limit`](pallet::Pallet::set_default_rate_limit): set the global default //! block span used by `RateLimitKind::Default` entries. //! //! The pallet also tracks the last block in which a rate-limited call was executed, per optional -//! *context*. Context allows one limit definition (for example, “set weights”) to be enforced per -//! subnet, account, or other grouping chosen by the runtime. +//! *usage key*. A usage key may refine tracking beyond the limit scope (for example combining a +//! `netuid` with a hyperparameter name), so the two concepts are explicitly separated in the +//! configuration. //! //! Each storage map is namespaced by pallet instance; runtimes can deploy multiple independent //! instances to manage distinct rate-limiting scopes. @@ -41,21 +43,27 @@ //! ); //! ``` //! -//! # Context resolver +//! # Context resolvers //! -//! The extension needs to know when two invocations should share a rate limit. This is controlled -//! by implementing [`RateLimitContextResolver`] for the runtime call type (or for a helper that the -//! runtime wires into [`Config::ContextResolver`]). The resolver receives the call and returns -//! `Some(context)` if the rate should be scoped (e.g. by `netuid`), or `None` to use the global -//! entry. The resolver is only used when *tracking* executions; you still configure limits via the -//! explicit `context` argument on `set_rate_limit`/`clear_rate_limit`. +//! The pallet relies on two resolvers, both implementing [`RateLimitContextResolver`]: +//! +//! - [`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. +//! - [`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)`). +//! +//! Each resolver receives the call and may return `Some(identifier)` when scoping is required, or +//! `None` to use the global entry. Extrinsics such as +//! [`set_rate_limit`](pallet::Pallet::set_rate_limit) automatically consult these resolvers. //! //! ```ignore //! pub struct WeightsContextResolver; //! -//! impl pallet_rate_limiting::RateLimitContextResolver -//! for WeightsContextResolver -//! { +//! // Limits are scoped per netuid. +//! pub struct ScopeResolver; +//! impl pallet_rate_limiting::RateLimitContextResolver for ScopeResolver { //! fn context(call: &RuntimeCall) -> Option { //! match call { //! RuntimeCall::Subtensor(pallet_subtensor::Call::set_weights { netuid, .. }) => { @@ -66,10 +74,29 @@ //! } //! } //! +//! // Usage tracking distinguishes hyperparameter + netuid. +//! pub struct UsageResolver; +//! impl pallet_rate_limiting::RateLimitContextResolver +//! for UsageResolver +//! { +//! fn context(call: &RuntimeCall) -> Option<(NetUid, HyperParam)> { +//! match call { +//! RuntimeCall::Subtensor(pallet_subtensor::Call::set_hyperparam { +//! netuid, +//! hyper, +//! .. +//! }) => Some((*netuid, *hyper)), +//! _ => None, +//! } +//! } +//! } +//! //! impl pallet_rate_limiting::Config for Runtime { //! type RuntimeCall = RuntimeCall; -//! type LimitContext = NetUid; -//! type ContextResolver = WeightsContextResolver; +//! type LimitScope = NetUid; +//! type LimitScopeResolver = ScopeResolver; +//! type UsageKey = (NetUid, HyperParam); +//! type UsageResolver = UsageResolver; //! type AdminOrigin = frame_system::EnsureRoot; //! } //! ``` @@ -121,11 +148,17 @@ pub mod pallet { /// Origin permitted to configure rate limits. type AdminOrigin: EnsureOrigin>; - /// Context type used for contextual (per-group) rate limits. - type LimitContext: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; + /// Scope identifier used to namespace stored rate limits. + type LimitScope: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; + + /// Resolves the scope for the given runtime call when configuring limits. + type LimitScopeResolver: RateLimitContextResolver<>::RuntimeCall, Self::LimitScope>; - /// Resolves the context for a given runtime call. - type ContextResolver: RateLimitContextResolver<>::RuntimeCall, Self::LimitContext>; + /// Usage key tracked in [`LastSeen`] for rate-limited calls. + type UsageKey: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; + + /// Resolves the usage key for the given runtime call when enforcing limits. + type UsageResolver: RateLimitContextResolver<>::RuntimeCall, Self::UsageKey>; /// Helper used to construct runtime calls for benchmarking. #[cfg(feature = "runtime-benchmarks")] @@ -139,20 +172,20 @@ pub mod pallet { _, Blake2_128Concat, TransactionIdentifier, - RateLimit<>::LimitContext, BlockNumberFor>, + RateLimit<>::LimitScope, BlockNumberFor>, OptionQuery, >; /// Tracks when a transaction was last observed. /// - /// The second key is `None` for global limits and `Some(context)` for contextual limits. + /// The second key is `None` for global tracking and `Some(key)` for scoped usage tracking. #[pallet::storage] pub type LastSeen, I: 'static = ()> = StorageDoubleMap< _, Blake2_128Concat, TransactionIdentifier, Blake2_128Concat, - Option<>::LimitContext>, + Option<>::UsageKey>, BlockNumberFor, OptionQuery, >; @@ -171,8 +204,8 @@ pub mod pallet { RateLimitSet { /// Identifier of the affected transaction. transaction: TransactionIdentifier, - /// Context to which the limit applies, if any. - context: Option<>::LimitContext>, + /// Limit scope to which the configuration applies, if any. + scope: Option<>::LimitScope>, /// The rate limit policy applied to the transaction. limit: RateLimitKind>, /// Pallet name associated with the transaction. @@ -184,8 +217,8 @@ pub mod pallet { RateLimitCleared { /// Identifier of the affected transaction. transaction: TransactionIdentifier, - /// Context from which the limit was cleared, if any. - context: Option<>::LimitContext>, + /// Limit scope from which the configuration was cleared, if any. + scope: Option<>::LimitScope>, /// Pallet name associated with the transaction. pallet: Vec, /// Extrinsic name associated with the transaction. @@ -205,8 +238,6 @@ pub mod pallet { InvalidRuntimeCall, /// Attempted to remove a limit that is not present. MissingRateLimit, - /// Contextual configuration was requested but no context can be resolved for the call. - ContextUnavailable, } #[pallet::genesis_config] @@ -214,7 +245,7 @@ pub mod pallet { pub default_limit: BlockNumberFor, pub limits: Vec<( TransactionIdentifier, - Option<>::LimitContext>, + Option<>::LimitScope>, RateLimitKind>, )>, } @@ -234,16 +265,16 @@ pub mod pallet { fn build(&self) { DefaultLimit::::put(self.default_limit); - for (identifier, context, kind) in &self.limits { - Limits::::mutate(identifier, |entry| match context { + for (identifier, scope, kind) in &self.limits { + Limits::::mutate(identifier, |entry| match scope { None => { *entry = Some(RateLimit::global(*kind)); } - Some(ctx) => { + Some(sc) => { if let Some(config) = entry { - config.upsert_context(ctx.clone(), *kind); + config.upsert_scope(sc.clone(), *kind); } else { - *entry = Some(RateLimit::contextual_single(ctx.clone(), *kind)); + *entry = Some(RateLimit::scoped_single(sc.clone(), *kind)); } } }); @@ -257,18 +288,19 @@ pub mod pallet { impl, I: 'static> Pallet { /// Returns `true` when the given transaction identifier passes its configured rate limit - /// within the provided context. + /// within the provided usage scope. pub fn is_within_limit( identifier: &TransactionIdentifier, - context: &Option<>::LimitContext>, + scope: &Option<>::LimitScope>, + usage_key: &Option<>::UsageKey>, ) -> Result { - let Some(block_span) = Self::resolved_limit(identifier, context) else { + let Some(block_span) = Self::resolved_limit(identifier, scope) else { return Ok(true); }; let current = frame_system::Pallet::::block_number(); - if let Some(last) = LastSeen::::get(identifier, context) { + if let Some(last) = LastSeen::::get(identifier, usage_key) { let delta = current.saturating_sub(last); if delta < block_span { return Ok(false); @@ -280,10 +312,10 @@ pub mod pallet { pub(crate) fn resolved_limit( identifier: &TransactionIdentifier, - context: &Option<>::LimitContext>, + scope: &Option<>::LimitScope>, ) -> Option> { let config = Limits::::get(identifier)?; - let kind = config.kind_for(context.as_ref())?; + let kind = config.kind_for(scope.as_ref())?; Some(match *kind { RateLimitKind::Default => DefaultLimit::::get(), RateLimitKind::Exact(block_span) => block_span, @@ -294,21 +326,21 @@ pub mod pallet { pub fn limit_for_call_names( pallet_name: &str, extrinsic_name: &str, - context: Option<>::LimitContext>, + scope: Option<>::LimitScope>, ) -> Option>> { let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; Limits::::get(&identifier) - .and_then(|config| config.kind_for(context.as_ref()).copied()) + .and_then(|config| config.kind_for(scope.as_ref()).copied()) } /// Returns the resolved block span for the specified pallet/extrinsic names, if any. pub fn resolved_limit_for_call_names( pallet_name: &str, extrinsic_name: &str, - context: Option<>::LimitContext>, + scope: Option<>::LimitScope>, ) -> Option> { let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; - Self::resolved_limit(&identifier, &context) + Self::resolved_limit(&identifier, &scope) } fn identifier_for_call_names( @@ -327,35 +359,29 @@ pub mod pallet { #[pallet::call] impl, I: 'static> Pallet { - /// Sets the rate limit configuration for the given call and optional context. + /// Sets the rate limit configuration for the given call. /// /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any - /// arguments embedded in the call are ignored**. The `context` parameter determines which - /// scoped entry is updated (for example a subnet identifier). Passing `None` updates the - /// global entry, which acts as a fallback when no context-specific limit exists. + /// arguments embedded in the call are ignored**. The applicable scope is discovered via + /// [`Config::LimitScopeResolver`]. When a scope resolves, the configuration is stored + /// against that scope; otherwise the global entry is updated. #[pallet::call_index(0)] #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn set_rate_limit( origin: OriginFor, call: Box<>::RuntimeCall>, limit: RateLimitKind>, - context: Option<>::LimitContext>, ) -> DispatchResult { T::AdminOrigin::ensure_origin(origin)?; - if context.is_some() - && >::ContextResolver::context(call.as_ref()).is_none() - { - return Err(Error::::ContextUnavailable.into()); - } - let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - let context_for_event = context.clone(); + let scope = >::LimitScopeResolver::context(call.as_ref()); + let scope_for_event = scope.clone(); - if let Some(ref ctx) = context { + if let Some(ref sc) = scope { Limits::::mutate(&identifier, |slot| match slot { - Some(config) => config.upsert_context(ctx.clone(), limit), - None => *slot = Some(RateLimit::contextual_single(ctx.clone(), limit)), + Some(config) => config.upsert_scope(sc.clone(), limit), + None => *slot = Some(RateLimit::scoped_single(sc.clone(), limit)), }); } else { Limits::::insert(&identifier, RateLimit::global(limit)); @@ -367,7 +393,7 @@ pub mod pallet { Self::deposit_event(Event::RateLimitSet { transaction: identifier, - context: context_for_event, + scope: scope_for_event, limit, pallet, extrinsic, @@ -378,18 +404,18 @@ pub mod pallet { /// Clears the rate limit for the given call, if present. /// /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any - /// arguments embedded in the call are ignored**. The `context` parameter must match the - /// entry that should be removed (use `None` to remove the global configuration). + /// arguments embedded in the call are ignored**. The configuration scope is determined via + /// [`Config::LimitScopeResolver`]. When no scope resolves, the global entry is cleared. #[pallet::call_index(1)] #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn clear_rate_limit( origin: OriginFor, call: Box<>::RuntimeCall>, - context: Option<>::LimitContext>, ) -> DispatchResult { T::AdminOrigin::ensure_origin(origin)?; let identifier = TransactionIdentifier::from_call::(call.as_ref())?; + let scope = >::LimitScopeResolver::context(call.as_ref()); let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); @@ -398,13 +424,13 @@ pub mod pallet { let mut removed = false; Limits::::mutate_exists(&identifier, |maybe_config| { if let Some(config) = maybe_config { - match (&context, config) { + match (&scope, config) { (None, _) => { removed = true; *maybe_config = None; } - (Some(ctx), RateLimit::Contextual(map)) => { - if map.remove(ctx).is_some() { + (Some(sc), RateLimit::Scoped(map)) => { + if map.remove(sc).is_some() { removed = true; if map.is_empty() { *maybe_config = None; @@ -420,7 +446,7 @@ pub mod pallet { Self::deposit_event(Event::RateLimitCleared { transaction: identifier, - context, + scope, pallet, extrinsic, }); diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs index 1c792dde16..aec00b45bb 100644 --- a/pallets/rate-limiting/src/mock.rs +++ b/pallets/rate-limiting/src/mock.rs @@ -1,5 +1,7 @@ #![allow(dead_code)] +use core::convert::TryInto; + use frame_support::{ derive_impl, sp_runtime::{ @@ -53,15 +55,30 @@ impl frame_system::Config for Test { type Block = Block; } -pub type LimitContext = u16; +pub type LimitScope = u16; +pub type UsageKey = u16; + +pub struct TestScopeResolver; +pub struct TestUsageResolver; -pub struct TestContextResolver; +impl pallet_rate_limiting::RateLimitContextResolver for TestScopeResolver { + fn context(call: &RuntimeCall) -> Option { + match call { + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { + (*block_span).try_into().ok() + } + RuntimeCall::RateLimiting(_) => Some(1), + _ => None, + } + } +} -impl pallet_rate_limiting::RateLimitContextResolver - for TestContextResolver -{ - fn context(call: &RuntimeCall) -> Option { +impl pallet_rate_limiting::RateLimitContextResolver for TestUsageResolver { + fn context(call: &RuntimeCall) -> Option { match call { + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { + (*block_span).try_into().ok() + } RuntimeCall::RateLimiting(_) => Some(1), _ => None, } @@ -70,8 +87,10 @@ impl pallet_rate_limiting::RateLimitContextResolver impl pallet_rate_limiting::Config for Test { type RuntimeCall = RuntimeCall; - type LimitContext = LimitContext; - type ContextResolver = TestContextResolver; + type LimitScope = LimitScope; + type LimitScopeResolver = TestScopeResolver; + type UsageKey = UsageKey; + type UsageResolver = TestUsageResolver; type AdminOrigin = EnsureRoot; #[cfg(feature = "runtime-benchmarks")] type BenchmarkHelper = BenchHelper; diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index 27bf6e5472..f02c2c52b0 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -29,14 +29,14 @@ fn limit_for_call_names_returns_stored_limit() { } #[test] -fn limit_for_call_names_prefers_context_specific_limit() { +fn limit_for_call_names_prefers_scope_specific_limit() { new_test_ext().execute_with(|| { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); let identifier = identifier_for(&call); Limits::::insert( identifier, - RateLimit::contextual_single(5u16, RateLimitKind::Exact(8)), + RateLimit::scoped_single(5u16, RateLimitKind::Exact(8)), ); let fetched = @@ -71,7 +71,7 @@ fn resolved_limit_for_call_names_resolves_default_value() { } #[test] -fn resolved_limit_for_call_names_prefers_context_specific_value() { +fn resolved_limit_for_call_names_prefers_scope_specific_value() { new_test_ext().execute_with(|| { let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); @@ -79,7 +79,7 @@ fn resolved_limit_for_call_names_prefers_context_specific_value() { let mut map = BTreeMap::new(); map.insert(6u16, RateLimitKind::Exact(9)); map.insert(2u16, RateLimitKind::Exact(4)); - Limits::::insert(identifier, RateLimit::Contextual(map)); + Limits::::insert(identifier, RateLimit::Scoped(map)); let resolved = RateLimiting::resolved_limit_for_call_names( "RateLimiting", @@ -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); + let result = RateLimiting::is_within_limit(&identifier, &None, &None); assert_eq!(result.expect("no error expected"), true); }); } @@ -134,14 +134,18 @@ fn is_within_limit_false_when_rate_limited() { let identifier = identifier_for(&call); Limits::::insert( identifier, - RateLimit::contextual_single(1 as LimitContext, RateLimitKind::Exact(5)), + RateLimit::scoped_single(1 as LimitScope, RateLimitKind::Exact(5)), ); - LastSeen::::insert(identifier, Some(1 as LimitContext), 9); + LastSeen::::insert(identifier, Some(1 as UsageKey), 9); System::set_block_number(13); - let within = RateLimiting::is_within_limit(&identifier, &Some(1 as LimitContext)) - .expect("call succeeds"); + let within = RateLimiting::is_within_limit( + &identifier, + &Some(1 as LimitScope), + &Some(1 as UsageKey), + ) + .expect("call succeeds"); assert!(!within); }); } @@ -154,14 +158,18 @@ fn is_within_limit_true_after_required_span() { let identifier = identifier_for(&call); Limits::::insert( identifier, - RateLimit::contextual_single(2 as LimitContext, RateLimitKind::Exact(5)), + RateLimit::scoped_single(2 as LimitScope, RateLimitKind::Exact(5)), ); - LastSeen::::insert(identifier, Some(2 as LimitContext), 10); + LastSeen::::insert(identifier, Some(2 as UsageKey), 10); System::set_block_number(20); - let within = RateLimiting::is_within_limit(&identifier, &Some(2 as LimitContext)) - .expect("call succeeds"); + let within = RateLimiting::is_within_limit( + &identifier, + &Some(2 as LimitScope), + &Some(2 as UsageKey), + ) + .expect("call succeeds"); assert!(within); }); } @@ -179,25 +187,24 @@ fn set_rate_limit_updates_storage_and_emits_event() { RuntimeOrigin::root(), Box::new(target_call.clone()), limit, - None, )); let identifier = identifier_for(&target_call); assert_eq!( Limits::::get(identifier), - Some(RateLimit::global(limit)) + Some(RateLimit::scoped_single(0, limit)) ); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitSet { transaction, - context, + scope, limit: emitted_limit, pallet, extrinsic, }) => { assert_eq!(transaction, identifier); - assert_eq!(context, None); + assert_eq!(scope, Some(0)); assert_eq!(emitted_limit, limit); assert_eq!(pallet, b"RateLimiting".to_vec()); assert_eq!(extrinsic, b"set_default_rate_limit".to_vec()); @@ -208,24 +215,42 @@ fn set_rate_limit_updates_storage_and_emits_event() { } #[test] -fn set_rate_limit_supports_context_specific_limit() { +fn set_rate_limit_stores_global_when_scope_absent() { new_test_ext().execute_with(|| { + System::reset_events(); + let target_call = - RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); - let context = Some(7u16); + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }); + let limit = RateLimitKind::Exact(11); + assert_ok!(RateLimiting::set_rate_limit( RuntimeOrigin::root(), Box::new(target_call.clone()), - RateLimitKind::Exact(11), - context.clone(), + limit, )); let identifier = identifier_for(&target_call); - let config = Limits::::get(identifier).expect("config stored"); assert_eq!( - config.kind_for(context.as_ref()).copied(), - Some(RateLimitKind::Exact(11)) + Limits::::get(identifier), + Some(RateLimit::global(limit)) ); + + match pop_last_event() { + RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitSet { + transaction, + scope, + limit: emitted_limit, + pallet, + extrinsic, + }) => { + assert_eq!(transaction, identifier); + assert_eq!(scope, None); + assert_eq!(emitted_limit, limit); + assert_eq!(pallet, b"System".to_vec()); + assert_eq!(extrinsic, b"remark".to_vec()); + } + other => panic!("unexpected event: {:?}", other), + } }); } @@ -240,7 +265,6 @@ fn set_rate_limit_requires_root() { RuntimeOrigin::signed(1), Box::new(target_call), RateLimitKind::Exact(1), - None, ), BadOrigin ); @@ -257,13 +281,12 @@ fn set_rate_limit_accepts_default_variant() { RuntimeOrigin::root(), Box::new(target_call.clone()), RateLimitKind::Default, - None, )); let identifier = identifier_for(&target_call); assert_eq!( Limits::::get(identifier), - Some(RateLimit::global(RateLimitKind::Default)) + Some(RateLimit::scoped_single(0, RateLimitKind::Default)) ); }); } @@ -274,14 +297,13 @@ fn clear_rate_limit_removes_entry_and_emits_event() { System::reset_events(); let target_call = - RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }); let identifier = identifier_for(&target_call); Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(4))); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), Box::new(target_call.clone()), - None, )); assert!(Limits::::get(identifier).is_none()); @@ -289,14 +311,14 @@ fn clear_rate_limit_removes_entry_and_emits_event() { match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { transaction, - context, + scope, pallet, extrinsic, }) => { assert_eq!(transaction, identifier); - assert_eq!(context, None); - assert_eq!(pallet, b"RateLimiting".to_vec()); - assert_eq!(extrinsic, b"set_default_rate_limit".to_vec()); + assert_eq!(scope, None); + assert_eq!(pallet, b"System".to_vec()); + assert_eq!(extrinsic, b"remark".to_vec()); } other => panic!("unexpected event: {:?}", other), } @@ -310,29 +332,31 @@ fn clear_rate_limit_fails_when_missing() { RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); assert_noop!( - RateLimiting::clear_rate_limit(RuntimeOrigin::root(), Box::new(target_call), None), + RateLimiting::clear_rate_limit(RuntimeOrigin::root(), Box::new(target_call)), Error::::MissingRateLimit ); }); } #[test] -fn clear_rate_limit_removes_only_selected_context() { +fn clear_rate_limit_removes_only_selected_scope() { new_test_ext().execute_with(|| { System::reset_events(); - let target_call = + let base_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); - let identifier = identifier_for(&target_call); + let identifier = identifier_for(&base_call); let mut map = BTreeMap::new(); map.insert(9u16, RateLimitKind::Exact(7)); map.insert(10u16, RateLimitKind::Exact(5)); - Limits::::insert(identifier, RateLimit::Contextual(map)); + Limits::::insert(identifier, RateLimit::Scoped(map)); + + let scoped_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 9 }); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), - Box::new(target_call.clone()), - Some(9), + Box::new(scoped_call.clone()), )); let config = Limits::::get(identifier).expect("config remains"); @@ -345,35 +369,17 @@ fn clear_rate_limit_removes_only_selected_context() { match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { transaction, - context, + scope, .. }) => { assert_eq!(transaction, identifier); - assert_eq!(context, Some(9)); + assert_eq!(scope, Some(9)); } other => panic!("unexpected event: {:?}", other), } }); } -#[test] -fn set_rate_limit_rejects_unresolvable_context() { - new_test_ext().execute_with(|| { - let target_call = - RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }); - - assert_noop!( - RateLimiting::set_rate_limit( - RuntimeOrigin::root(), - Box::new(target_call), - RateLimitKind::Exact(5), - Some(42), - ), - Error::::ContextUnavailable - ); - }); -} - #[test] fn set_default_rate_limit_updates_storage_and_emits_event() { new_test_ext().execute_with(|| { diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index e72b473fd4..5276f1f396 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -80,14 +80,8 @@ where const IDENTIFIER: &'static str = IDENTIFIER; type Implicit = (); - type Val = Option<( - TransactionIdentifier, - Option<>::LimitContext>, - )>; - type Pre = Option<( - TransactionIdentifier, - Option<>::LimitContext>, - )>; + type Val = Option<(TransactionIdentifier, Option<>::UsageKey>)>; + type Pre = Option<(TransactionIdentifier, Option<>::UsageKey>)>; fn weight(&self, _call: &>::RuntimeCall) -> Weight { Weight::zero() @@ -108,9 +102,10 @@ where Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)), }; - let context = >::ContextResolver::context(call); + let scope = >::LimitScopeResolver::context(call); + let usage = >::UsageResolver::context(call); - let Some(block_span) = Pallet::::resolved_limit(&identifier, &context) else { + let Some(block_span) = Pallet::::resolved_limit(&identifier, &scope) else { return Ok((ValidTransaction::default(), None, origin)); }; @@ -118,7 +113,7 @@ where return Ok((ValidTransaction::default(), None, origin)); } - let within_limit = Pallet::::is_within_limit(&identifier, &context) + let within_limit = Pallet::::is_within_limit(&identifier, &scope, &usage) .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; if !within_limit { @@ -129,7 +124,7 @@ where Ok(( ValidTransaction::default(), - Some((identifier, context)), + Some((identifier, usage)), origin, )) } @@ -153,9 +148,9 @@ where result: &DispatchResult, ) -> Result<(), TransactionValidityError> { if result.is_ok() { - if let Some((identifier, context)) = pre { + if let Some((identifier, usage)) = pre { let block_number = frame_system::Pallet::::block_number(); - LastSeen::::insert(&identifier, context, block_number); + LastSeen::::insert(&identifier, usage, block_number); } } Ok(()) @@ -193,7 +188,7 @@ mod tests { ) -> Result< ( sp_runtime::transaction_validity::ValidTransaction, - Option<(TransactionIdentifier, Option)>, + Option<(TransactionIdentifier, Option)>, RuntimeOrigin, ), TransactionValidityError, @@ -241,7 +236,7 @@ mod tests { let identifier = identifier_for(&call); assert_eq!( - LastSeen::::get(identifier, None::), + LastSeen::::get(identifier, None::), None ); }); @@ -279,7 +274,7 @@ mod tests { .expect("post_dispatch succeeds"); assert_eq!( - LastSeen::::get(identifier, None::), + LastSeen::::get(identifier, None::), Some(10) ); }); @@ -292,7 +287,7 @@ mod tests { let call = remark_call(); let identifier = identifier_for(&call); Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(5))); - LastSeen::::insert(identifier, None::, 20); + LastSeen::::insert(identifier, None::, 20); System::set_block_number(22); @@ -339,7 +334,7 @@ mod tests { .expect("post_dispatch succeeds"); assert_eq!( - LastSeen::::get(identifier, None::), + LastSeen::::get(identifier, None::), None ); }); diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index a18fff37c6..1daf2915b3 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -3,7 +3,7 @@ use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata}; use scale_info::TypeInfo; use sp_std::collections::btree_map::BTreeMap; -/// Resolves the optional context within which a rate limit applies. +/// Resolves the optional identifier within which a rate limit applies. pub trait RateLimitContextResolver { /// Returns `Some(context)` when the limit should be applied per-context, or `None` for global /// limits. @@ -104,71 +104,71 @@ pub enum RateLimitKind { /// Stored rate limit configuration for a transaction identifier. /// /// The configuration is mutually exclusive: either the call is globally limited or it stores a set -/// of per-context spans. +/// of per-scope spans. #[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] #[cfg_attr( feature = "std", serde( - bound = "Context: Ord + serde::Serialize + serde::de::DeserializeOwned, BlockNumber: serde::Serialize + serde::de::DeserializeOwned" + bound = "Scope: Ord + serde::Serialize + serde::de::DeserializeOwned, BlockNumber: serde::Serialize + serde::de::DeserializeOwned" ) )] #[derive(Clone, PartialEq, Eq, Encode, Decode, DecodeWithMemTracking, TypeInfo, Debug)] -pub enum RateLimit { +pub enum RateLimit { /// Global span applied to every invocation. Global(RateLimitKind), - /// Per-context spans keyed by `Context`. - Contextual(BTreeMap>), + /// Per-scope spans keyed by `Scope`. + Scoped(BTreeMap>), } -impl RateLimit +impl RateLimit where - Context: Ord, + Scope: Ord, { /// Convenience helper to build a global configuration. pub fn global(kind: RateLimitKind) -> Self { Self::Global(kind) } - /// Convenience helper to build a contextual configuration containing a single entry. - pub fn contextual_single(context: Context, kind: RateLimitKind) -> Self { + /// Convenience helper to build a scoped configuration containing a single entry. + pub fn scoped_single(scope: Scope, kind: RateLimitKind) -> Self { let mut map = BTreeMap::new(); - map.insert(context, kind); - Self::Contextual(map) + map.insert(scope, kind); + Self::Scoped(map) } - /// Returns the span configured for the provided context, if any. - pub fn kind_for(&self, context: Option<&Context>) -> Option<&RateLimitKind> { + /// Returns the span configured for the provided scope, if any. + pub fn kind_for(&self, scope: Option<&Scope>) -> Option<&RateLimitKind> { match self { RateLimit::Global(kind) => Some(kind), - RateLimit::Contextual(map) => context.and_then(|ctx| map.get(ctx)), + RateLimit::Scoped(map) => scope.and_then(|key| map.get(key)), } } - /// Inserts or updates a contextual entry, converting from a global configuration if needed. - pub fn upsert_context(&mut self, context: Context, kind: RateLimitKind) { + /// Inserts or updates a scoped entry, converting from a global configuration if needed. + pub fn upsert_scope(&mut self, scope: Scope, kind: RateLimitKind) { match self { RateLimit::Global(_) => { let mut map = BTreeMap::new(); - map.insert(context, kind); - *self = RateLimit::Contextual(map); + map.insert(scope, kind); + *self = RateLimit::Scoped(map); } - RateLimit::Contextual(map) => { - map.insert(context, kind); + RateLimit::Scoped(map) => { + map.insert(scope, kind); } } } - /// Removes a contextual entry, returning whether one existed. - pub fn remove_context(&mut self, context: &Context) -> bool { + /// Removes a scoped entry, returning whether one existed. + pub fn remove_scope(&mut self, scope: &Scope) -> bool { match self { RateLimit::Global(_) => false, - RateLimit::Contextual(map) => map.remove(context).is_some(), + RateLimit::Scoped(map) => map.remove(scope).is_some(), } } - /// Returns true when the contextual configuration contains no entries. - pub fn is_contextual_empty(&self) -> bool { - matches!(self, RateLimit::Contextual(map) if map.is_empty()) + /// Returns true when the scoped configuration contains no entries. + pub fn is_scoped_empty(&self) -> bool { + matches!(self, RateLimit::Scoped(map) if map.is_empty()) } } From 2df6d4b8c95e15f024dcf4e961f9d30d14391542 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 31 Oct 2025 14:48:38 +0300 Subject: [PATCH 15/23] Update docs --- pallets/rate-limiting/src/lib.rs | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 0579249b36..e1c7a1665a 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -361,10 +361,11 @@ pub mod pallet { impl, I: 'static> Pallet { /// Sets the rate limit configuration for the given call. /// - /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any - /// arguments embedded in the call are ignored**. The applicable scope is discovered via - /// [`Config::LimitScopeResolver`]. When a scope resolves, the configuration is stored - /// against that scope; otherwise the global entry is updated. + /// The supplied `call` is inspected to derive the pallet/extrinsic indices and passed to + /// [`Config::LimitScopeResolver`] to determine the applicable scope. The pallet never + /// persists the call arguments directly, but a resolver may read them in order to resolve + /// its context. When a scope resolves, the configuration is stored against that scope; + /// otherwise the global entry is updated. #[pallet::call_index(0)] #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn set_rate_limit( @@ -403,9 +404,10 @@ pub mod pallet { /// Clears the rate limit for the given call, if present. /// - /// The supplied `call` is only used to derive the pallet and extrinsic indices; **any - /// arguments embedded in the call are ignored**. The configuration scope is determined via - /// [`Config::LimitScopeResolver`]. When no scope resolves, the global entry is cleared. + /// The supplied `call` is inspected to derive the pallet/extrinsic indices and passed to + /// [`Config::LimitScopeResolver`] when determining which scoped configuration to clear. + /// The pallet does not persist the call arguments, but resolvers may read them while + /// computing the scope. When no scope resolves, the global entry is cleared. #[pallet::call_index(1)] #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] pub fn clear_rate_limit( From 15c9aa9b7072fd98618b193cccf7fe3c7744371a Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 31 Oct 2025 14:58:02 +0300 Subject: [PATCH 16/23] Add an extrinsic to clear all scoped rate limits in pallet-rate-limiting --- pallets/rate-limiting/src/lib.rs | 44 +++++++++++++++++++++++- pallets/rate-limiting/src/tests.rs | 54 ++++++++++++++++++++++++++++++ pallets/rate-limiting/src/types.rs | 2 +- 3 files changed, 98 insertions(+), 2 deletions(-) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index e1c7a1665a..7b960e566e 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -224,6 +224,15 @@ pub mod pallet { /// Extrinsic name associated with the transaction. extrinsic: Vec, }, + /// All scoped and global rate limits for a call were cleared. + AllRateLimitsCleared { + /// Identifier of the affected transaction. + transaction: TransactionIdentifier, + /// Pallet name associated with the transaction. + pallet: Vec, + /// Extrinsic name associated with the transaction. + extrinsic: Vec, + }, /// The default rate limit was set or updated. DefaultRateLimitSet { /// The new default limit expressed in blocks. @@ -456,8 +465,41 @@ pub mod pallet { Ok(()) } - /// Sets the default rate limit in blocks applied to calls configured to use it. + /// Clears every stored rate limit configuration for the given call, including scoped + /// entries. + /// + /// The supplied `call` is inspected to derive the pallet and extrinsic indices. All stored + /// scopes for that call, along with any associated usage tracking entries, are removed when + /// this extrinsic succeeds. #[pallet::call_index(2)] + #[pallet::weight(T::DbWeight::get().reads_writes(1, 1))] + pub fn clear_all_rate_limits( + origin: OriginFor, + call: Box<>::RuntimeCall>, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + let identifier = TransactionIdentifier::from_call::(call.as_ref())?; + let (pallet_name, extrinsic_name) = identifier.names::()?; + let pallet = Vec::from(pallet_name.as_bytes()); + let extrinsic = Vec::from(extrinsic_name.as_bytes()); + + let removed = Limits::::take(&identifier).is_some(); + ensure!(removed, Error::::MissingRateLimit); + + let _ = LastSeen::::clear_prefix(&identifier, u32::MAX, None); + + Self::deposit_event(Event::AllRateLimitsCleared { + transaction: identifier, + pallet, + extrinsic, + }); + + Ok(()) + } + + /// Sets the default rate limit in blocks applied to calls configured to use it. + #[pallet::call_index(3)] #[pallet::weight(T::DbWeight::get().writes(1))] pub fn set_default_rate_limit( origin: OriginFor, diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index f02c2c52b0..1a89951193 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -380,6 +380,60 @@ fn clear_rate_limit_removes_only_selected_scope() { }); } +#[test] +fn clear_all_rate_limits_removes_entire_configuration() { + new_test_ext().execute_with(|| { + System::reset_events(); + + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&target_call); + + let mut map = BTreeMap::new(); + map.insert(3u16, RateLimitKind::Exact(6)); + map.insert(4u16, RateLimitKind::Exact(7)); + Limits::::insert(identifier, RateLimit::Scoped(map)); + + LastSeen::::insert(identifier, Some(3u16), 11); + LastSeen::::insert(identifier, None::, 12); + + assert_ok!(RateLimiting::clear_all_rate_limits( + RuntimeOrigin::root(), + Box::new(target_call.clone()), + )); + + assert!(Limits::::get(identifier).is_none()); + assert!(LastSeen::::get(identifier, Some(3u16)).is_none()); + assert!(LastSeen::::get(identifier, None::).is_none()); + + match pop_last_event() { + RuntimeEvent::RateLimiting(crate::pallet::Event::AllRateLimitsCleared { + transaction, + pallet, + extrinsic, + }) => { + assert_eq!(transaction, identifier); + assert_eq!(pallet, b"RateLimiting".to_vec()); + assert_eq!(extrinsic, b"set_default_rate_limit".to_vec()); + } + other => panic!("unexpected event: {:?}", other), + } + }); +} + +#[test] +fn clear_all_rate_limits_fails_when_missing() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }); + + assert_noop!( + RateLimiting::clear_all_rate_limits(RuntimeOrigin::root(), Box::new(target_call)), + Error::::MissingRateLimit + ); + }); +} + #[test] fn set_default_rate_limit_updates_storage_and_emits_event() { new_test_ext().execute_with(|| { diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index 1daf2915b3..4e68d0205a 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -189,7 +189,7 @@ mod tests { // System is the first pallet in the mock runtime, RateLimiting is second. assert_eq!(identifier.pallet_index, 1); // set_default_rate_limit has call_index 2. - assert_eq!(identifier.extrinsic_index, 2); + assert_eq!(identifier.extrinsic_index, 3); } #[test] From e74d65015d33e2ef116c7ec1b970e25f37da9410 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 31 Oct 2025 15:15:17 +0300 Subject: [PATCH 17/23] Clear LastSeen on clear_rate_limit call --- pallets/rate-limiting/src/lib.rs | 18 ++++++++++++++++++ pallets/rate-limiting/src/tests.rs | 13 +++++++++++++ 2 files changed, 31 insertions(+) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 7b960e566e..38cfce1a5e 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -427,6 +427,7 @@ pub mod pallet { let identifier = TransactionIdentifier::from_call::(call.as_ref())?; let scope = >::LimitScopeResolver::context(call.as_ref()); + let usage = >::UsageResolver::context(call.as_ref()); let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); @@ -455,6 +456,23 @@ pub mod pallet { ensure!(removed, Error::::MissingRateLimit); + if removed { + match (scope.as_ref(), usage) { + (None, _) => { + let _ = LastSeen::::clear_prefix(&identifier, u32::MAX, None); + } + (_, Some(key)) => { + LastSeen::::remove(&identifier, Some(key)); + } + (_, None) => { + LastSeen::::remove( + &identifier, + Option::<>::UsageKey>::None, + ); + } + } + } + Self::deposit_event(Event::RateLimitCleared { transaction: identifier, scope, diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index 1a89951193..16639e2d5b 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -300,6 +300,8 @@ fn clear_rate_limit_removes_entry_and_emits_event() { RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }); let identifier = identifier_for(&target_call); Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(4))); + LastSeen::::insert(identifier, None::, 7); + LastSeen::::insert(identifier, Some(88u16), 9); assert_ok!(RateLimiting::clear_rate_limit( RuntimeOrigin::root(), @@ -307,6 +309,8 @@ fn clear_rate_limit_removes_entry_and_emits_event() { )); assert!(Limits::::get(identifier).is_none()); + assert!(LastSeen::::get(identifier, None::).is_none()); + assert!(LastSeen::::get(identifier, Some(88u16)).is_none()); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { @@ -350,6 +354,9 @@ fn clear_rate_limit_removes_only_selected_scope() { map.insert(9u16, RateLimitKind::Exact(7)); map.insert(10u16, RateLimitKind::Exact(5)); Limits::::insert(identifier, RateLimit::Scoped(map)); + LastSeen::::insert(identifier, Some(9u16), 11); + LastSeen::::insert(identifier, Some(10u16), 12); + LastSeen::::insert(identifier, None::, 13); let scoped_call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 9 }); @@ -365,6 +372,12 @@ fn clear_rate_limit_removes_only_selected_scope() { config.kind_for(Some(&10u16)).copied(), Some(RateLimitKind::Exact(5)) ); + assert!(LastSeen::::get(identifier, Some(9u16)).is_none()); + assert_eq!(LastSeen::::get(identifier, Some(10u16)), Some(12)); + assert_eq!( + LastSeen::::get(identifier, None::), + Some(13) + ); match pop_last_event() { RuntimeEvent::RateLimiting(crate::pallet::Event::RateLimitCleared { From fc60465b506b26f62aa71685ef2e0fae3dc6d99b Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 31 Oct 2025 16:02:15 +0300 Subject: [PATCH 18/23] Add rate limit bypassing --- pallets/rate-limiting/src/lib.rs | 27 ++++++++++--------- pallets/rate-limiting/src/mock.rs | 11 ++++++-- pallets/rate-limiting/src/tx_extension.rs | 33 ++++++++++++++++++++++- pallets/rate-limiting/src/types.rs | 24 +++++++++++++---- 4 files changed, 74 insertions(+), 21 deletions(-) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 38cfce1a5e..a5c5982307 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -45,7 +45,7 @@ //! //! # Context resolvers //! -//! The pallet relies on two resolvers, both implementing [`RateLimitContextResolver`]: +//! 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 @@ -63,7 +63,7 @@ //! //! // Limits are scoped per netuid. //! pub struct ScopeResolver; -//! impl pallet_rate_limiting::RateLimitContextResolver for ScopeResolver { +//! impl pallet_rate_limiting::RateLimitScopeResolver for ScopeResolver { //! fn context(call: &RuntimeCall) -> Option { //! match call { //! RuntimeCall::Subtensor(pallet_subtensor::Call::set_weights { netuid, .. }) => { @@ -76,9 +76,8 @@ //! //! // Usage tracking distinguishes hyperparameter + netuid. //! pub struct UsageResolver; -//! impl pallet_rate_limiting::RateLimitContextResolver -//! for UsageResolver -//! { +//! impl pallet_rate_limiting::RateLimitUsageResolver +//! for UsageResolver { //! fn context(call: &RuntimeCall) -> Option<(NetUid, HyperParam)> { //! match call { //! RuntimeCall::Subtensor(pallet_subtensor::Call::set_hyperparam { @@ -105,7 +104,9 @@ pub use benchmarking::BenchmarkHelper; pub use pallet::*; pub use tx_extension::RateLimitTransactionExtension; -pub use types::{RateLimit, RateLimitContextResolver, RateLimitKind, TransactionIdentifier}; +pub use types::{ + RateLimit, RateLimitKind, RateLimitScopeResolver, RateLimitUsageResolver, TransactionIdentifier, +}; #[cfg(feature = "runtime-benchmarks")] mod benchmarking; @@ -131,7 +132,10 @@ pub mod pallet { #[cfg(feature = "runtime-benchmarks")] use crate::benchmarking::BenchmarkHelper as BenchmarkHelperTrait; - use crate::types::{RateLimit, RateLimitContextResolver, RateLimitKind, TransactionIdentifier}; + use crate::types::{ + RateLimit, RateLimitKind, RateLimitScopeResolver, RateLimitUsageResolver, + TransactionIdentifier, + }; /// Configuration trait for the rate limiting pallet. #[pallet::config] @@ -152,13 +156,13 @@ pub mod pallet { type LimitScope: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; /// Resolves the scope for the given runtime call when configuring limits. - type LimitScopeResolver: RateLimitContextResolver<>::RuntimeCall, Self::LimitScope>; + type LimitScopeResolver: RateLimitScopeResolver<>::RuntimeCall, Self::LimitScope>; /// Usage key tracked in [`LastSeen`] for rate-limited calls. type UsageKey: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; /// Resolves the usage key for the given runtime call when enforcing limits. - type UsageResolver: RateLimitContextResolver<>::RuntimeCall, Self::UsageKey>; + type UsageResolver: RateLimitUsageResolver<>::RuntimeCall, Self::UsageKey>; /// Helper used to construct runtime calls for benchmarking. #[cfg(feature = "runtime-benchmarks")] @@ -465,10 +469,7 @@ pub mod pallet { LastSeen::::remove(&identifier, Some(key)); } (_, None) => { - LastSeen::::remove( - &identifier, - Option::<>::UsageKey>::None, - ); + LastSeen::::remove(&identifier, None::<>::UsageKey>); } } } diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs index aec00b45bb..5fab86e4ab 100644 --- a/pallets/rate-limiting/src/mock.rs +++ b/pallets/rate-limiting/src/mock.rs @@ -61,7 +61,7 @@ pub type UsageKey = u16; pub struct TestScopeResolver; pub struct TestUsageResolver; -impl pallet_rate_limiting::RateLimitContextResolver for TestScopeResolver { +impl pallet_rate_limiting::RateLimitScopeResolver for TestScopeResolver { fn context(call: &RuntimeCall) -> Option { match call { RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { @@ -71,9 +71,16 @@ impl pallet_rate_limiting::RateLimitContextResolver for _ => None, } } + + fn should_bypass(call: &RuntimeCall) -> bool { + matches!( + call, + RuntimeCall::RateLimiting(RateLimitingCall::clear_rate_limit { .. }) + ) + } } -impl pallet_rate_limiting::RateLimitContextResolver for TestUsageResolver { +impl pallet_rate_limiting::RateLimitUsageResolver for TestUsageResolver { fn context(call: &RuntimeCall) -> Option { match call { RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs index 5276f1f396..95696409ee 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -17,7 +17,7 @@ use sp_std::{marker::PhantomData, result::Result}; use crate::{ Config, LastSeen, Pallet, - types::{RateLimitContextResolver, TransactionIdentifier}, + types::{RateLimitScopeResolver, RateLimitUsageResolver, TransactionIdentifier}, }; /// Identifier returned in the transaction metadata for the rate limiting extension. @@ -97,6 +97,10 @@ where _inherited_implication: &impl Implication, _source: TransactionSource, ) -> ValidateResult>::RuntimeCall> { + if >::LimitScopeResolver::should_bypass(call) { + return Ok((ValidTransaction::default(), None, origin)); + } + let identifier = match TransactionIdentifier::from_call::(call) { Ok(identifier) => identifier, Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)), @@ -178,6 +182,12 @@ mod tests { RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }) } + fn bypass_call() -> RuntimeCall { + RuntimeCall::RateLimiting(RateLimitingCall::clear_rate_limit { + call: Box::new(remark_call()), + }) + } + fn new_tx_extension() -> RateLimitTransactionExtension { RateLimitTransactionExtension(Default::default()) } @@ -242,6 +252,27 @@ mod tests { }); } + #[test] + fn tx_extension_honors_bypass_signal() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = bypass_call(); + + let (valid, val, _) = + validate_with_tx_extension(&extension, &call).expect("bypass should succeed"); + assert_eq!(valid.priority, 0); + assert!(val.is_none()); + + let identifier = identifier_for(&call); + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(3))); + LastSeen::::insert(identifier, None::, 1); + + let (_valid, post_val, _) = + validate_with_tx_extension(&extension, &call).expect("still bypassed"); + assert!(post_val.is_none()); + }); + } + #[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 index 4e68d0205a..0f4d4948f1 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -3,11 +3,25 @@ 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. -pub trait RateLimitContextResolver { - /// Returns `Some(context)` when the limit should be applied per-context, or `None` for global - /// limits. - fn context(call: &Call) -> Option; +/// Resolves the optional identifier within which a rate limit applies and can optionally bypass +/// enforcement. +pub trait RateLimitScopeResolver { + /// Returns `Some(scope)` when the limit should be applied per-scope, or `None` for global + /// limits. + fn context(call: &Call) -> Option; + + /// Returns `true` when the rate limit should be bypassed for the provided call. Defaults to + /// `false`. + fn should_bypass(_call: &Call) -> bool { + false + } +} + +/// Resolves the optional usage tracking key applied when enforcing limits. +pub trait RateLimitUsageResolver { + /// Returns `Some(usage)` when usage should be tracked per-key, or `None` for global usage + /// tracking. + fn context(call: &Call) -> Option; } /// Identifies a runtime call by pallet and extrinsic indices. From 5c8a8abf746681bf18ef47c62cfa9c9aba212a04 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Fri, 31 Oct 2025 16:31:48 +0300 Subject: [PATCH 19/23] Add rate limit adjuster --- pallets/rate-limiting/src/lib.rs | 69 +++++++++++++++++------ pallets/rate-limiting/src/mock.rs | 15 ++++- pallets/rate-limiting/src/tests.rs | 4 +- pallets/rate-limiting/src/tx_extension.rs | 34 ++++++++++- pallets/rate-limiting/src/types.rs | 24 +++++--- 5 files changed, 115 insertions(+), 31 deletions(-) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index a5c5982307..e40b857ed6 100644 --- 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 for ScopeResolver { +//! impl pallet_rate_limiting::RateLimitScopeResolver for ScopeResolver { //! fn context(call: &RuntimeCall) -> Option { //! 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 -//! for UsageResolver { +//! impl pallet_rate_limiting::RateLimitUsageResolver 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<>::RuntimeCall, Self::LimitScope>; + type LimitScopeResolver: RateLimitScopeResolver< + >::RuntimeCall, + Self::LimitScope, + BlockNumberFor, + >; /// 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<>::LimitScope>, usage_key: &Option<>::UsageKey>, + call: &>::RuntimeCall, ) -> Result { - let Some(block_span) = Self::resolved_limit(identifier, scope) else { + if >::LimitScopeResolver::should_bypass(call) { return Ok(true); - }; - - let current = frame_system::Pallet::::block_number(); - - if let Some(last) = LastSeen::::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: &>::RuntimeCall, + identifier: &TransactionIdentifier, + scope: &Option<>::LimitScope>, + ) -> Option> { + let span = Self::resolved_limit(identifier, scope)?; + Some(>::LimitScopeResolver::adjust_span( + call, span, + )) + } + + pub(crate) fn within_span( + identifier: &TransactionIdentifier, + usage_key: &Option<>::UsageKey>, + block_span: BlockNumberFor, + ) -> bool { + if block_span.is_zero() { + return true; + } + + if let Some(last) = LastSeen::::get(identifier, usage_key) { + let current = frame_system::Pallet::::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 index 5fab86e4ab..67321731a1 100644 --- 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 for TestScopeResolver { +impl pallet_rate_limiting::RateLimitScopeResolver + for TestScopeResolver +{ fn context(call: &RuntimeCall) -> Option { match call { RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { @@ -78,6 +80,17 @@ impl pallet_rate_limiting::RateLimitScopeResolver 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 for TestUsageResolver { diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index 16639e2d5b..c89543ae4b 100644 --- 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 index 95696409ee..623a6af3ac 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -109,7 +109,7 @@ where let scope = >::LimitScopeResolver::context(call); let usage = >::UsageResolver::context(call); - let Some(block_span) = Pallet::::resolved_limit(&identifier, &scope) else { + let Some(block_span) = Pallet::::effective_span(call, &identifier, &scope) else { return Ok((ValidTransaction::default(), None, origin)); }; @@ -117,8 +117,7 @@ where return Ok((ValidTransaction::default(), None, origin)); } - let within_limit = Pallet::::is_within_limit(&identifier, &scope, &usage) - .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; + let within_limit = Pallet::::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 { 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::::insert(identifier, RateLimit::global(RateLimitKind::Exact(4))); + LastSeen::::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 index 0f4d4948f1..5d537bf64f 100644 --- 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 { - /// 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 { + /// Returns `Some(scope)` when the limit should be applied per-scope, or `None` for global + /// limits. fn context(call: &Call) -> Option; - /// 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 { - /// 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; } From 1e8db7db79ef51089efa1b74923e1f95bc1206db Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Mon, 3 Nov 2025 15:28:09 +0300 Subject: [PATCH 20/23] Add api to migrate scope and usage keys --- pallets/rate-limiting/src/lib.rs | 70 +++++++++++++++++++ pallets/rate-limiting/src/tests.rs | 104 +++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+) diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index e40b857ed6..f57b9f0074 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -370,6 +370,76 @@ pub mod pallet { true } + /// Migrates a stored rate limit configuration from one scope to another. + /// + /// Returns `true` when an entry was moved. Passing identical `from`/`to` scopes simply + /// checks that a configuration exists. + pub fn migrate_limit_scope( + identifier: &TransactionIdentifier, + from: Option<>::LimitScope>, + to: Option<>::LimitScope>, + ) -> bool { + if from == to { + return Limits::::contains_key(identifier); + } + + let mut migrated = false; + Limits::::mutate(identifier, |maybe_config| { + if let Some(config) = maybe_config { + match (from.as_ref(), to.as_ref()) { + (None, Some(target)) => { + if let RateLimit::Global(kind) = config { + *config = RateLimit::scoped_single(target.clone(), *kind); + migrated = true; + } + } + (Some(source), Some(target)) => { + if let RateLimit::Scoped(map) = config { + if let Some(kind) = map.remove(source) { + map.insert(target.clone(), kind); + migrated = true; + } + } + } + (Some(source), None) => { + if let RateLimit::Scoped(map) = config { + if map.len() == 1 && map.contains_key(source) { + if let Some(kind) = map.remove(source) { + *config = RateLimit::global(kind); + migrated = true; + } + } + } + } + _ => {} + } + } + }); + + migrated + } + + /// Migrates the cached usage information for a rate-limited call to a new key. + /// + /// Returns `true` when an entry was moved. Passing identical keys simply checks that an + /// entry exists. + pub fn migrate_usage_key( + identifier: &TransactionIdentifier, + from: Option<>::UsageKey>, + to: Option<>::UsageKey>, + ) -> bool { + if from == to { + return LastSeen::::contains_key(identifier, &to); + } + + let Some(block) = LastSeen::::take(identifier, from) else { + return false; + }; + + LastSeen::::insert(identifier, to, block); + 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/tests.rs b/pallets/rate-limiting/src/tests.rs index c89543ae4b..0e051e9eb0 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -176,6 +176,110 @@ fn is_within_limit_true_after_required_span() { }); } +#[test] +fn migrate_limit_scope_global_to_scoped() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }); + let identifier = identifier_for(&target_call); + + Limits::::insert(identifier, RateLimit::global(RateLimitKind::Exact(3))); + + assert!(RateLimiting::migrate_limit_scope( + &identifier, + None, + Some(9) + )); + + match RateLimiting::limits(identifier).expect("config") { + RateLimit::Scoped(map) => { + assert_eq!(map.len(), 1); + assert_eq!(map.get(&9), Some(&RateLimitKind::Exact(3))); + } + other => panic!("unexpected config: {:?}", other), + } + }); +} + +#[test] +fn migrate_limit_scope_scoped_to_scoped() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&target_call); + + let mut map = sp_std::collections::btree_map::BTreeMap::new(); + map.insert(1u16, RateLimitKind::Exact(4)); + map.insert(2u16, RateLimitKind::Exact(6)); + Limits::::insert(identifier, RateLimit::Scoped(map)); + + assert!(RateLimiting::migrate_limit_scope( + &identifier, + Some(1), + Some(3) + )); + + match RateLimiting::limits(identifier).expect("config") { + RateLimit::Scoped(map) => { + assert!(map.get(&1).is_none()); + assert_eq!(map.get(&3), Some(&RateLimitKind::Exact(4))); + assert_eq!(map.get(&2), Some(&RateLimitKind::Exact(6))); + } + other => panic!("unexpected config: {:?}", other), + } + }); +} + +#[test] +fn migrate_limit_scope_scoped_to_global() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&target_call); + + let mut map = sp_std::collections::btree_map::BTreeMap::new(); + map.insert(7u16, RateLimitKind::Exact(8)); + Limits::::insert(identifier, RateLimit::Scoped(map)); + + assert!(RateLimiting::migrate_limit_scope( + &identifier, + Some(7), + None + )); + + match RateLimiting::limits(identifier).expect("config") { + RateLimit::Global(kind) => assert_eq!(kind, RateLimitKind::Exact(8)), + other => panic!("unexpected config: {:?}", other), + } + }); +} + +#[test] +fn migrate_usage_key_moves_entry() { + new_test_ext().execute_with(|| { + let target_call = + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 }); + let identifier = identifier_for(&target_call); + + LastSeen::::insert(identifier, Some(5u16), 11); + + assert!(RateLimiting::migrate_usage_key( + &identifier, + Some(5), + Some(6) + )); + assert!(LastSeen::::get(identifier, Some(5u16)).is_none()); + assert_eq!(LastSeen::::get(identifier, Some(6u16)), Some(11)); + + assert!(RateLimiting::migrate_usage_key(&identifier, Some(6), None)); + assert!(LastSeen::::get(identifier, Some(6u16)).is_none()); + assert_eq!( + LastSeen::::get(identifier, None::), + Some(11) + ); + }); +} + #[test] fn set_rate_limit_updates_storage_and_emits_event() { new_test_ext().execute_with(|| { From 232de64664e7cc9affe9394e45fcbf68023cdf55 Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Mon, 3 Nov 2025 18:47:54 +0300 Subject: [PATCH 21/23] Pass origin to rate limit context resolvers --- pallets/rate-limiting/Cargo.toml | 2 + pallets/rate-limiting/src/benchmarking.rs | 11 +++- pallets/rate-limiting/src/lib.rs | 65 +++++++++++++++++------ pallets/rate-limiting/src/mock.rs | 14 ++--- pallets/rate-limiting/src/tests.rs | 11 ++-- pallets/rate-limiting/src/tx_extension.rs | 9 ++-- pallets/rate-limiting/src/types.rs | 16 +++--- 7 files changed, 88 insertions(+), 40 deletions(-) diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml index 3447145622..36343ec2cf 100644 --- a/pallets/rate-limiting/Cargo.toml +++ b/pallets/rate-limiting/Cargo.toml @@ -17,6 +17,7 @@ frame-system.workspace = true scale-info = { workspace = true, features = ["derive"] } serde = { workspace = true, features = ["derive"], optional = true } sp-std.workspace = true +sp-runtime.workspace = true subtensor-runtime-common.workspace = true [dev-dependencies] @@ -34,6 +35,7 @@ std = [ "scale-info/std", "serde", "sp-std/std", + "sp-runtime/std", "subtensor-runtime-common/std", ] runtime-benchmarks = [ diff --git a/pallets/rate-limiting/src/benchmarking.rs b/pallets/rate-limiting/src/benchmarking.rs index 65d547ab0b..2d700a4ef6 100644 --- a/pallets/rate-limiting/src/benchmarking.rs +++ b/pallets/rate-limiting/src/benchmarking.rs @@ -5,6 +5,7 @@ use codec::Decode; use frame_benchmarking::v2::*; use frame_system::{RawOrigin, pallet_prelude::BlockNumberFor}; +use sp_runtime::traits::DispatchOriginOf; use super::*; @@ -36,7 +37,10 @@ mod benchmarks { fn set_rate_limit() { let call = sample_call::(); let limit = RateLimitKind::>::Exact(BlockNumberFor::::from(10u32)); - let scope = ::LimitScopeResolver::context(call.as_ref()); + let origin = T::RuntimeOrigin::from(RawOrigin::Root); + let resolver_origin: DispatchOriginOf<::RuntimeCall> = + Into::::RuntimeCall>>::into(origin.clone()); + let scope = ::LimitScopeResolver::context(&resolver_origin, call.as_ref()); let identifier = TransactionIdentifier::from_call::(call.as_ref()).expect("identifier"); @@ -61,7 +65,10 @@ mod benchmarks { fn clear_rate_limit() { let call = sample_call::(); let limit = RateLimitKind::>::Exact(BlockNumberFor::::from(10u32)); - let scope = ::LimitScopeResolver::context(call.as_ref()); + let origin = T::RuntimeOrigin::from(RawOrigin::Root); + let resolver_origin: DispatchOriginOf<::RuntimeCall> = + Into::::RuntimeCall>>::into(origin.clone()); + let scope = ::LimitScopeResolver::context(&resolver_origin, call.as_ref()); // Pre-populate limit for benchmark call let identifier = diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index f57b9f0074..4389b833a8 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -55,8 +55,8 @@ //! [`LastSeen`](pallet::LastSeen). This can refine the limit scope (for example by returning a //! tuple of `(netuid, hyperparameter)`). //! -//! Each resolver receives the call and may return `Some(identifier)` when scoping is required, or -//! `None` to use the global entry. Extrinsics such as +//! Each resolver receives the origin and call and may return `Some(identifier)` when scoping is +//! required, or `None` to use the global entry. Extrinsics such as //! [`set_rate_limit`](pallet::Pallet::set_rate_limit) automatically consult these resolvers. //! //! ```ignore @@ -64,8 +64,13 @@ //! //! // Limits are scoped per netuid. //! pub struct ScopeResolver; -//! impl pallet_rate_limiting::RateLimitScopeResolver for ScopeResolver { -//! fn context(call: &RuntimeCall) -> Option { +//! impl pallet_rate_limiting::RateLimitScopeResolver< +//! RuntimeOrigin, +//! RuntimeCall, +//! NetUid, +//! BlockNumber, +//! > for ScopeResolver { +//! fn context(origin: &RuntimeOrigin, call: &RuntimeCall) -> Option { //! match call { //! RuntimeCall::Subtensor(pallet_subtensor::Call::set_weights { netuid, .. }) => { //! Some(*netuid) @@ -74,15 +79,23 @@ //! } //! } //! -//! fn adjust_span(_call: &RuntimeCall, span: BlockNumber) -> BlockNumber { +//! fn should_bypass(origin: &RuntimeOrigin, _call: &RuntimeCall) -> bool { +//! matches!(origin, RuntimeOrigin::Root) +//! } +//! +//! fn adjust_span(_origin: &RuntimeOrigin, _call: &RuntimeCall, span: BlockNumber) -> BlockNumber { //! span //! } //! } //! //! // Usage tracking distinguishes hyperparameter + netuid. //! pub struct UsageResolver; -//! impl pallet_rate_limiting::RateLimitUsageResolver for UsageResolver { -//! fn context(call: &RuntimeCall) -> Option<(NetUid, HyperParam)> { +//! impl pallet_rate_limiting::RateLimitUsageResolver< +//! RuntimeOrigin, +//! RuntimeCall, +//! (NetUid, HyperParam), +//! > for UsageResolver { +//! fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<(NetUid, HyperParam)> { //! match call { //! RuntimeCall::Subtensor(pallet_subtensor::Call::set_hyperparam { //! netuid, @@ -128,10 +141,10 @@ pub mod pallet { use codec::Codec; use frame_support::{ pallet_prelude::*, - sp_runtime::traits::{Saturating, Zero}, traits::{BuildGenesisConfig, EnsureOrigin, GetCallMetadata}, }; use frame_system::pallet_prelude::*; + use sp_runtime::traits::{DispatchOriginOf, Dispatchable, Saturating, Zero}; use sp_std::{convert::TryFrom, marker::PhantomData, vec::Vec}; #[cfg(feature = "runtime-benchmarks")] @@ -146,11 +159,14 @@ pub mod pallet { pub trait Config: frame_system::Config where BlockNumberFor: MaybeSerializeDeserialize, + <>::RuntimeCall as Dispatchable>::RuntimeOrigin: + From<::RuntimeOrigin>, { /// The overarching runtime call type. type RuntimeCall: Parameter + Codec + GetCallMetadata + + Dispatchable + IsType<::RuntimeCall>; /// Origin permitted to configure rate limits. @@ -161,6 +177,7 @@ pub mod pallet { /// Resolves the scope for the given runtime call when configuring limits. type LimitScopeResolver: RateLimitScopeResolver< + DispatchOriginOf<>::RuntimeCall>, >::RuntimeCall, Self::LimitScope, BlockNumberFor, @@ -170,7 +187,11 @@ pub mod pallet { type UsageKey: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; /// Resolves the usage key for the given runtime call when enforcing limits. - type UsageResolver: RateLimitUsageResolver<>::RuntimeCall, Self::UsageKey>; + type UsageResolver: RateLimitUsageResolver< + DispatchOriginOf<>::RuntimeCall>, + >::RuntimeCall, + Self::UsageKey, + >; /// Helper used to construct runtime calls for benchmarking. #[cfg(feature = "runtime-benchmarks")] @@ -311,16 +332,17 @@ pub mod pallet { /// Returns `true` when the given transaction identifier passes its configured rate limit /// within the provided usage scope. pub fn is_within_limit( + origin: &DispatchOriginOf<>::RuntimeCall>, + call: &>::RuntimeCall, identifier: &TransactionIdentifier, scope: &Option<>::LimitScope>, usage_key: &Option<>::UsageKey>, - call: &>::RuntimeCall, ) -> Result { - if >::LimitScopeResolver::should_bypass(call) { + if >::LimitScopeResolver::should_bypass(origin, call) { return Ok(true); } - let Some(block_span) = Self::effective_span(call, identifier, scope) else { + let Some(block_span) = Self::effective_span(origin, call, identifier, scope) else { return Ok(true); }; @@ -340,13 +362,14 @@ pub mod pallet { } pub(crate) fn effective_span( + origin: &DispatchOriginOf<>::RuntimeCall>, call: &>::RuntimeCall, identifier: &TransactionIdentifier, scope: &Option<>::LimitScope>, ) -> Option> { let span = Self::resolved_limit(identifier, scope)?; Some(>::LimitScopeResolver::adjust_span( - call, span, + origin, call, span, )) } @@ -491,11 +514,15 @@ pub mod pallet { call: Box<>::RuntimeCall>, limit: RateLimitKind>, ) -> DispatchResult { + let resolver_origin: DispatchOriginOf<>::RuntimeCall> = + Into::>::RuntimeCall>>::into(origin.clone()); + let scope = + >::LimitScopeResolver::context(&resolver_origin, call.as_ref()); + let scope_for_event = scope.clone(); + T::AdminOrigin::ensure_origin(origin)?; let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - let scope = >::LimitScopeResolver::context(call.as_ref()); - let scope_for_event = scope.clone(); if let Some(ref sc) = scope { Limits::::mutate(&identifier, |slot| match slot { @@ -532,11 +559,15 @@ pub mod pallet { origin: OriginFor, call: Box<>::RuntimeCall>, ) -> DispatchResult { + let resolver_origin: DispatchOriginOf<>::RuntimeCall> = + Into::>::RuntimeCall>>::into(origin.clone()); + let scope = + >::LimitScopeResolver::context(&resolver_origin, call.as_ref()); + let usage = >::UsageResolver::context(&resolver_origin, call.as_ref()); + T::AdminOrigin::ensure_origin(origin)?; let identifier = TransactionIdentifier::from_call::(call.as_ref())?; - let scope = >::LimitScopeResolver::context(call.as_ref()); - let usage = >::UsageResolver::context(call.as_ref()); let (pallet_name, extrinsic_name) = identifier.names::()?; let pallet = Vec::from(pallet_name.as_bytes()); diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs index 67321731a1..fb7de0a400 100644 --- a/pallets/rate-limiting/src/mock.rs +++ b/pallets/rate-limiting/src/mock.rs @@ -61,10 +61,10 @@ pub type UsageKey = u16; pub struct TestScopeResolver; pub struct TestUsageResolver; -impl pallet_rate_limiting::RateLimitScopeResolver +impl pallet_rate_limiting::RateLimitScopeResolver for TestScopeResolver { - fn context(call: &RuntimeCall) -> Option { + fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option { match call { RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { (*block_span).try_into().ok() @@ -74,14 +74,14 @@ impl pallet_rate_limiting::RateLimitScopeResolver } } - fn should_bypass(call: &RuntimeCall) -> bool { + fn should_bypass(_origin: &RuntimeOrigin, call: &RuntimeCall) -> bool { matches!( call, RuntimeCall::RateLimiting(RateLimitingCall::clear_rate_limit { .. }) ) } - fn adjust_span(call: &RuntimeCall, span: u64) -> u64 { + fn adjust_span(_origin: &RuntimeOrigin, call: &RuntimeCall, span: u64) -> u64 { if matches!( call, RuntimeCall::RateLimiting(RateLimitingCall::clear_all_rate_limits { .. }) @@ -93,8 +93,10 @@ impl pallet_rate_limiting::RateLimitScopeResolver } } -impl pallet_rate_limiting::RateLimitUsageResolver for TestUsageResolver { - fn context(call: &RuntimeCall) -> Option { +impl pallet_rate_limiting::RateLimitUsageResolver + for TestUsageResolver +{ + fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option { match call { RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { (*block_span).try_into().ok() diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs index 0e051e9eb0..a377d71656 100644 --- a/pallets/rate-limiting/src/tests.rs +++ b/pallets/rate-limiting/src/tests.rs @@ -121,7 +121,8 @@ 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, &call); + let origin = RuntimeOrigin::signed(1); + let result = RateLimiting::is_within_limit(&origin, &call, &identifier, &None, &None); assert_eq!(result.expect("no error expected"), true); }); } @@ -140,11 +141,13 @@ fn is_within_limit_false_when_rate_limited() { System::set_block_number(13); + let origin = RuntimeOrigin::signed(1); let within = RateLimiting::is_within_limit( + &origin, + &call, &identifier, &Some(1 as LimitScope), &Some(1 as UsageKey), - &call, ) .expect("call succeeds"); assert!(!within); @@ -165,11 +168,13 @@ fn is_within_limit_true_after_required_span() { System::set_block_number(20); + let origin = RuntimeOrigin::signed(1); let within = RateLimiting::is_within_limit( + &origin, + &call, &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 index 623a6af3ac..c6c3eb745c 100644 --- a/pallets/rate-limiting/src/tx_extension.rs +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -97,7 +97,7 @@ where _inherited_implication: &impl Implication, _source: TransactionSource, ) -> ValidateResult>::RuntimeCall> { - if >::LimitScopeResolver::should_bypass(call) { + if >::LimitScopeResolver::should_bypass(&origin, call) { return Ok((ValidTransaction::default(), None, origin)); } @@ -106,10 +106,11 @@ where Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)), }; - let scope = >::LimitScopeResolver::context(call); - let usage = >::UsageResolver::context(call); + let scope = >::LimitScopeResolver::context(&origin, call); + let usage = >::UsageResolver::context(&origin, call); - let Some(block_span) = Pallet::::effective_span(call, &identifier, &scope) else { + let Some(block_span) = Pallet::::effective_span(&origin, call, &identifier, &scope) + else { return Ok((ValidTransaction::default(), None, origin)); }; diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index 5d537bf64f..f2f46a6beb 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -5,29 +5,29 @@ use sp_std::collections::btree_map::BTreeMap; /// Resolves the optional identifier within which a rate limit applies and can optionally adjust /// enforcement behaviour. -pub trait RateLimitScopeResolver { +pub trait RateLimitScopeResolver { /// Returns `Some(scope)` when the limit should be applied per-scope, or `None` for global /// limits. - fn context(call: &Call) -> Option; + fn context(origin: &Origin, call: &Call) -> Option; - /// Returns `true` when the rate limit should be bypassed for the provided call. Defaults to - /// `false`. - fn should_bypass(_call: &Call) -> bool { + /// 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 } /// Optionally adjusts the effective span used during enforcement. Defaults to the original /// `span`. - fn adjust_span(_call: &Call, span: Span) -> Span { + fn adjust_span(_origin: &Origin, _call: &Call, span: Span) -> Span { span } } /// Resolves the optional usage tracking key applied when enforcing limits. -pub trait RateLimitUsageResolver { +pub trait RateLimitUsageResolver { /// Returns `Some(usage)` when usage should be tracked per-key, or `None` for global usage /// tracking. - fn context(call: &Call) -> Option; + fn context(origin: &Origin, call: &Call) -> Option; } /// Identifies a runtime call by pallet and extrinsic indices. From e00342fc612ce0ae6b1b44684d4f34f03979bc0e Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Mon, 10 Nov 2025 18:03:52 +0300 Subject: [PATCH 22/23] Implement rate-limiting resolvers for pallet-subtensor --- Cargo.lock | 1 + common/src/lib.rs | 2 + common/src/rate_limiting.rs | 66 +++++ pallets/rate-limiting/Cargo.toml | 4 +- pallets/rate-limiting/src/lib.rs | 2 +- pallets/rate-limiting/src/types.rs | 28 +- pallets/subtensor/src/utils/rate_limiting.rs | 2 + runtime/Cargo.toml | 3 + runtime/src/lib.rs | 5 + runtime/src/rate_limiting.rs | 256 +++++++++++++++++++ 10 files changed, 357 insertions(+), 12 deletions(-) create mode 100644 common/src/rate_limiting.rs create mode 100644 runtime/src/rate_limiting.rs diff --git a/Cargo.lock b/Cargo.lock index a919bc7486..df6e2f9266 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8297,6 +8297,7 @@ dependencies = [ "pallet-nomination-pools-runtime-api", "pallet-offences", "pallet-preimage", + "pallet-rate-limiting", "pallet-registry", "pallet-safe-mode", "pallet-scheduler", diff --git a/common/src/lib.rs b/common/src/lib.rs index a98a957ad8..db4c314bbe 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -15,8 +15,10 @@ use sp_runtime::{ use subtensor_macros::freeze_struct; pub use currency::*; +pub use rate_limiting::{RateLimitScope, RateLimitUsageKey}; mod currency; +mod rate_limiting; /// Balance of an account. pub type Balance = u64; diff --git a/common/src/rate_limiting.rs b/common/src/rate_limiting.rs new file mode 100644 index 0000000000..3c88758943 --- /dev/null +++ b/common/src/rate_limiting.rs @@ -0,0 +1,66 @@ +use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; +use frame_support::pallet_prelude::Parameter; +use scale_info::TypeInfo; +use serde::{Deserialize, Serialize}; + +use crate::{MechId, NetUid}; + +#[derive( + Serialize, + Deserialize, + Encode, + Decode, + DecodeWithMemTracking, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Debug, + TypeInfo, + MaxEncodedLen, +)] +pub enum RateLimitScope { + Subnet(NetUid), + SubnetMechanism { netuid: NetUid, mecid: MechId }, +} + +#[derive( + Serialize, + Deserialize, + Encode, + Decode, + DecodeWithMemTracking, + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + Debug, + TypeInfo, + MaxEncodedLen, +)] +#[scale_info(skip_type_params(AccountId))] +pub enum RateLimitUsageKey { + Account(AccountId), + Subnet(NetUid), + AccountSubnet { + account: AccountId, + netuid: NetUid, + }, + ColdkeyHotkeySubnet { + coldkey: AccountId, + hotkey: AccountId, + netuid: NetUid, + }, + SubnetNeuron { + netuid: NetUid, + uid: u16, + }, + SubnetMechanismNeuron { + netuid: NetUid, + mecid: MechId, + uid: u16, + }, +} diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml index 36343ec2cf..67e2710f4b 100644 --- a/pallets/rate-limiting/Cargo.toml +++ b/pallets/rate-limiting/Cargo.toml @@ -15,7 +15,7 @@ frame-benchmarking = { workspace = true, optional = true } frame-support.workspace = true frame-system.workspace = true scale-info = { workspace = true, features = ["derive"] } -serde = { workspace = true, features = ["derive"], optional = true } +serde = { workspace = true, features = ["derive"] } sp-std.workspace = true sp-runtime.workspace = true subtensor-runtime-common.workspace = true @@ -33,7 +33,7 @@ std = [ "frame-support/std", "frame-system/std", "scale-info/std", - "serde", + "serde/std", "sp-std/std", "sp-runtime/std", "subtensor-runtime-common/std", diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 4389b833a8..54dd54f5f2 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -145,7 +145,7 @@ pub mod pallet { }; use frame_system::pallet_prelude::*; use sp_runtime::traits::{DispatchOriginOf, Dispatchable, Saturating, Zero}; - use sp_std::{convert::TryFrom, marker::PhantomData, vec::Vec}; + use sp_std::{boxed::Box, convert::TryFrom, marker::PhantomData, vec::Vec}; #[cfg(feature = "runtime-benchmarks")] use crate::benchmarking::BenchmarkHelper as BenchmarkHelperTrait; diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs index f2f46a6beb..4748f1576e 100644 --- a/pallets/rate-limiting/src/types.rs +++ b/pallets/rate-limiting/src/types.rs @@ -1,6 +1,7 @@ use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata}; use scale_info::TypeInfo; +use serde::{Deserialize, Serialize}; use sp_std::collections::btree_map::BTreeMap; /// Resolves the optional identifier within which a rate limit applies and can optionally adjust @@ -31,8 +32,9 @@ pub trait RateLimitUsageResolver { } /// Identifies a runtime call by pallet and extrinsic indices. -#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] #[derive( + Serialize, + Deserialize, Clone, Copy, PartialEq, @@ -101,8 +103,9 @@ impl TransactionIdentifier { } /// Policy describing the block span enforced by a rate limit. -#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] #[derive( + Serialize, + Deserialize, Clone, Copy, PartialEq, @@ -125,14 +128,21 @@ pub enum RateLimitKind { /// /// The configuration is mutually exclusive: either the call is globally limited or it stores a set /// of per-scope spans. -#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))] -#[cfg_attr( - feature = "std", - serde( - bound = "Scope: Ord + serde::Serialize + serde::de::DeserializeOwned, BlockNumber: serde::Serialize + serde::de::DeserializeOwned" - ) +#[derive( + Serialize, + Deserialize, + Clone, + PartialEq, + Eq, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + Debug, +)] +#[serde( + bound = "Scope: Ord + serde::Serialize + serde::de::DeserializeOwned, BlockNumber: serde::Serialize + serde::de::DeserializeOwned" )] -#[derive(Clone, PartialEq, Eq, Encode, Decode, DecodeWithMemTracking, TypeInfo, Debug)] pub enum RateLimit { /// Global span applied to every invocation. Global(RateLimitKind), diff --git a/pallets/subtensor/src/utils/rate_limiting.rs b/pallets/subtensor/src/utils/rate_limiting.rs index 85f58cfc64..468aecd1c1 100644 --- a/pallets/subtensor/src/utils/rate_limiting.rs +++ b/pallets/subtensor/src/utils/rate_limiting.rs @@ -1,3 +1,5 @@ +use codec::{Decode, Encode}; +use scale_info::TypeInfo; use subtensor_runtime_common::NetUid; use super::*; diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml index 9760ac1b53..b363eb4f5f 100644 --- a/runtime/Cargo.toml +++ b/runtime/Cargo.toml @@ -38,6 +38,7 @@ frame-system = { workspace = true } frame-try-runtime = { workspace = true, optional = true } pallet-timestamp.workspace = true pallet-transaction-payment.workspace = true +pallet-rate-limiting.workspace = true pallet-subtensor-utility.workspace = true frame-executive.workspace = true frame-metadata-hash-extension.workspace = true @@ -187,6 +188,7 @@ std = [ "pallet-timestamp/std", "pallet-transaction-payment-rpc-runtime-api/std", "pallet-transaction-payment/std", + "pallet-rate-limiting/std", "pallet-subtensor-utility/std", "pallet-sudo/std", "pallet-multisig/std", @@ -328,6 +330,7 @@ try-runtime = [ "pallet-insecure-randomness-collective-flip/try-runtime", "pallet-timestamp/try-runtime", "pallet-transaction-payment/try-runtime", + "pallet-rate-limiting/try-runtime", "pallet-subtensor-utility/try-runtime", "pallet-safe-mode/try-runtime", "pallet-subtensor/try-runtime", diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 266a755708..9d80693ddb 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -12,6 +12,7 @@ use core::num::NonZeroU64; pub mod check_nonce; mod migrations; +mod rate_limiting; pub mod transaction_payment_wrapper; extern crate alloc; @@ -70,6 +71,10 @@ use subtensor_precompiles::Precompiles; use subtensor_runtime_common::{AlphaCurrency, TaoCurrency, time::*, *}; use subtensor_swap_interface::{Order, SwapHandler}; +pub use rate_limiting::{ + ScopeResolver as RuntimeScopeResolver, UsageResolver as RuntimeUsageResolver, +}; + // A few exports that help ease life for downstream crates. pub use frame_support::{ StorageValue, construct_runtime, parameter_types, diff --git a/runtime/src/rate_limiting.rs b/runtime/src/rate_limiting.rs new file mode 100644 index 0000000000..4f569c2daf --- /dev/null +++ b/runtime/src/rate_limiting.rs @@ -0,0 +1,256 @@ +use frame_system::RawOrigin; +use pallet_admin_utils::Call as AdminUtilsCall; +use pallet_rate_limiting::{RateLimitScopeResolver, RateLimitUsageResolver}; +use pallet_subtensor::{Call as SubtensorCall, Tempo}; +use subtensor_runtime_common::{BlockNumber, MechId, NetUid, RateLimitScope, RateLimitUsageKey}; + +use crate::{AccountId, Runtime, RuntimeCall, RuntimeOrigin}; + +fn signed_origin(origin: &RuntimeOrigin) -> Option { + match origin.clone().into() { + Ok(RawOrigin::Signed(who)) => Some(who), + _ => None, + } +} + +fn tempo_scaled(netuid: NetUid, span: BlockNumber) -> BlockNumber { + if span == 0 { + return span; + } + let tempo = BlockNumber::from(Tempo::::get(netuid) as u32); + span.saturating_mul(tempo) +} + +fn neuron_identity(origin: &RuntimeOrigin, netuid: NetUid) -> Option<(AccountId, u16)> { + let hotkey = signed_origin(origin)?; + let uid = + pallet_subtensor::Pallet::::get_uid_for_net_and_hotkey(netuid, &hotkey).ok()?; + Some((hotkey, uid)) +} + +fn owner_hparam_netuid(call: &AdminUtilsCall) -> Option { + match call { + AdminUtilsCall::sudo_set_activity_cutoff { netuid, .. } + | AdminUtilsCall::sudo_set_adjustment_alpha { netuid, .. } + | AdminUtilsCall::sudo_set_alpha_sigmoid_steepness { netuid, .. } + | AdminUtilsCall::sudo_set_alpha_values { netuid, .. } + | AdminUtilsCall::sudo_set_bonds_moving_average { netuid, .. } + | AdminUtilsCall::sudo_set_bonds_penalty { netuid, .. } + | AdminUtilsCall::sudo_set_bonds_reset_enabled { netuid, .. } + | AdminUtilsCall::sudo_set_commit_reveal_weights_enabled { netuid, .. } + | AdminUtilsCall::sudo_set_commit_reveal_weights_interval { netuid, .. } + | AdminUtilsCall::sudo_set_immunity_period { netuid, .. } + | AdminUtilsCall::sudo_set_liquid_alpha_enabled { netuid, .. } + | AdminUtilsCall::sudo_set_max_allowed_uids { netuid, .. } + | AdminUtilsCall::sudo_set_max_burn { netuid, .. } + | AdminUtilsCall::sudo_set_max_difficulty { netuid, .. } + | AdminUtilsCall::sudo_set_min_allowed_weights { netuid, .. } + | AdminUtilsCall::sudo_set_min_burn { netuid, .. } + | AdminUtilsCall::sudo_set_network_pow_registration_allowed { netuid, .. } + | AdminUtilsCall::sudo_set_owner_immune_neuron_limit { netuid, .. } + | AdminUtilsCall::sudo_set_recycle_or_burn { netuid, .. } + | AdminUtilsCall::sudo_set_rho { netuid, .. } + | AdminUtilsCall::sudo_set_serving_rate_limit { netuid, .. } + | AdminUtilsCall::sudo_set_sn_owner_hotkey { netuid, .. } + | AdminUtilsCall::sudo_set_toggle_transfer { netuid, .. } + | AdminUtilsCall::sudo_set_weights_version_key { netuid, .. } + | AdminUtilsCall::sudo_set_yuma3_enabled { netuid, .. } => Some(*netuid), + _ => None, + } +} + +fn admin_scope_netuid(call: &AdminUtilsCall) -> Option { + owner_hparam_netuid(call).or_else(|| match call { + AdminUtilsCall::sudo_set_mechanism_count { netuid, .. } + | AdminUtilsCall::sudo_set_mechanism_emission_split { netuid, .. } + | AdminUtilsCall::sudo_trim_to_max_allowed_uids { netuid, .. } => Some(*netuid), + _ => None, + }) +} + +#[derive(Default)] +pub struct UsageResolver; + +impl RateLimitUsageResolver> + for UsageResolver +{ + fn context(origin: &RuntimeOrigin, call: &RuntimeCall) -> Option> { + match call { + RuntimeCall::SubtensorModule(inner) => match inner { + SubtensorCall::swap_hotkey { .. } => { + signed_origin(origin).map(RateLimitUsageKey::::Account) + } + SubtensorCall::register_network { .. } + | SubtensorCall::register_network_with_identity { .. } => { + signed_origin(origin).map(RateLimitUsageKey::::Account) + } + SubtensorCall::increase_take { hotkey, .. } => { + Some(RateLimitUsageKey::::Account(hotkey.clone())) + } + SubtensorCall::set_childkey_take { hotkey, netuid, .. } + | SubtensorCall::set_children { hotkey, netuid, .. } => { + Some(RateLimitUsageKey::::AccountSubnet { + account: hotkey.clone(), + netuid: *netuid, + }) + } + SubtensorCall::set_weights { netuid, .. } + | SubtensorCall::commit_weights { netuid, .. } + | SubtensorCall::reveal_weights { netuid, .. } + | SubtensorCall::batch_reveal_weights { netuid, .. } + | SubtensorCall::commit_timelocked_weights { netuid, .. } => { + let (_, uid) = neuron_identity(origin, netuid)?; + Some(RateLimitUsageKey::::SubnetNeuron { netuid, uid }) + } + SubtensorCall::set_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::reveal_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_crv3_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_timelocked_mechanism_weights { netuid, mecid, .. } => { + let (_, uid) = neuron_identity(origin, netuid)?; + Some(RateLimitUsageKey::::SubnetMechanismNeuron { + netuid, + mecid, + uid, + }) + } + SubtensorCall::serve_axon { netuid, .. } + | SubtensorCall::serve_axon_tls { netuid, .. } + | SubtensorCall::serve_prometheus { netuid, .. } => { + let hotkey = signed_origin(origin)?; + Some(RateLimitUsageKey::::AccountSubnet { + account: hotkey, + netuid: *netuid, + }) + } + SubtensorCall::associate_evm_key { netuid, .. } => { + let hotkey = signed_origin(origin)?; + let uid = pallet_subtensor::Pallet::::get_uid_for_net_and_hotkey( + *netuid, &hotkey, + ) + .ok()?; + Some(RateLimitUsageKey::::SubnetNeuron { + netuid: *netuid, + uid, + }) + } + SubtensorCall::add_stake { hotkey, netuid, .. } + | SubtensorCall::add_stake_limit { hotkey, netuid, .. } + | SubtensorCall::remove_stake { hotkey, netuid, .. } + | SubtensorCall::remove_stake_limit { hotkey, netuid, .. } + | SubtensorCall::remove_stake_full_limit { hotkey, netuid, .. } + | SubtensorCall::transfer_stake { + hotkey, + origin_netuid: netuid, + .. + } + | SubtensorCall::swap_stake { + hotkey, + origin_netuid: netuid, + .. + } + | SubtensorCall::swap_stake_limit { + hotkey, + origin_netuid: netuid, + .. + } + | SubtensorCall::move_stake { + origin_hotkey: hotkey, + origin_netuid: netuid, + .. + } + | SubtensorCall::recycle_alpha { hotkey, netuid, .. } + | SubtensorCall::burn_alpha { hotkey, netuid, .. } => { + let coldkey = signed_origin(origin)?; + Some(RateLimitUsageKey::::ColdkeyHotkeySubnet { + coldkey, + hotkey: hotkey.clone(), + netuid: *netuid, + }) + } + _ => None, + }, + RuntimeCall::AdminUtils(inner) => { + if let Some(netuid) = owner_hparam_netuid(inner) { + // Hyperparameter setters share a global span but are tracked per subnet. + Some(RateLimitUsageKey::::Subnet(netuid)) + } else { + match inner { + AdminUtilsCall::sudo_set_mechanism_count { netuid, .. } + | AdminUtilsCall::sudo_set_mechanism_emission_split { netuid, .. } + | AdminUtilsCall::sudo_trim_to_max_allowed_uids { netuid, .. } => { + let who = signed_origin(origin)?; + Some(RateLimitUsageKey::::AccountSubnet { + account: who, + netuid: *netuid, + }) + } + _ => None, + } + } + } + _ => None, + } + } +} + +#[derive(Default)] +pub struct ScopeResolver; + +impl RateLimitScopeResolver + for ScopeResolver +{ + fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option { + match call { + RuntimeCall::SubtensorModule(inner) => match inner { + SubtensorCall::serve_axon { netuid, .. } + | SubtensorCall::serve_axon_tls { netuid, .. } + | SubtensorCall::serve_prometheus { netuid, .. } + | SubtensorCall::set_weights { netuid, .. } + | SubtensorCall::commit_weights { netuid, .. } + | SubtensorCall::reveal_weights { netuid, .. } + | SubtensorCall::batch_reveal_weights { netuid, .. } + | SubtensorCall::commit_timelocked_weights { netuid, .. } => { + Some(RateLimitScope::Subnet(*netuid)) + } + SubtensorCall::set_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::reveal_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_crv3_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_timelocked_mechanism_weights { netuid, mecid, .. } => { + Some(RateLimitScope::SubnetMechanism { + netuid: *netuid, + mecid: *mecid, + }) + } + _ => None, + }, + RuntimeCall::AdminUtils(inner) => { + if owner_hparam_netuid(inner).is_some() { + // Hyperparameter setters share a global limit span; usage is tracked per subnet. + None + } else { + admin_scope_netuid(inner).map(RateLimitScope::Subnet) + } + } + _ => None, + } + } + + fn should_bypass(origin: &RuntimeOrigin, _call: &RuntimeCall) -> bool { + matches!(origin.clone().into(), Ok(RawOrigin::Root)) + } + + fn adjust_span(_origin: &RuntimeOrigin, call: &RuntimeCall, span: BlockNumber) -> BlockNumber { + match call { + RuntimeCall::AdminUtils(inner) => { + if let Some(netuid) = owner_hparam_netuid(inner) { + tempo_scaled(netuid, span) + } else { + span + } + } + _ => span, + } + } +} From 69715be7c0d3b185c47812302b6563af7aedd58e Mon Sep 17 00:00:00 2001 From: Aliaksandr Tsurko Date: Wed, 12 Nov 2025 20:00:37 +0300 Subject: [PATCH 23/23] Migrate pallet-subtensor's rate limiting to pallet-rate-limiting --- Cargo.lock | 1 + pallets/rate-limiting/src/lib.rs | 12 + pallets/subtensor/Cargo.toml | 2 + runtime/Cargo.toml | 2 +- runtime/src/rate_limiting/migration.rs | 733 ++++++++++++++++++ .../mod.rs} | 2 + 6 files changed, 751 insertions(+), 1 deletion(-) create mode 100644 runtime/src/rate_limiting/migration.rs rename runtime/src/{rate_limiting.rs => rate_limiting/mod.rs} (99%) diff --git a/Cargo.lock b/Cargo.lock index df6e2f9266..04686202fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10757,6 +10757,7 @@ dependencies = [ "pallet-crowdloan", "pallet-drand", "pallet-preimage", + "pallet-rate-limiting", "pallet-scheduler", "pallet-subtensor-proxy", "pallet-subtensor-swap", diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs index 54dd54f5f2..d823a94cd5 100644 --- a/pallets/rate-limiting/src/lib.rs +++ b/pallets/rate-limiting/src/lib.rs @@ -393,6 +393,18 @@ pub mod pallet { true } + /// Inserts or updates the cached usage timestamp for a rate-limited call. + /// + /// This is primarily intended for migrations that need to hydrate the new tracking storage + /// from legacy pallets. + pub fn record_last_seen( + identifier: &TransactionIdentifier, + usage_key: Option<>::UsageKey>, + block_number: BlockNumberFor, + ) { + LastSeen::::insert(identifier, usage_key, block_number); + } + /// Migrates a stored rate limit configuration from one scope to another. /// /// Returns `true` when an entry was moved. Passing identical `from`/`to` scopes simply diff --git a/pallets/subtensor/Cargo.toml b/pallets/subtensor/Cargo.toml index fdd5e5f9ab..3cde2ee24b 100644 --- a/pallets/subtensor/Cargo.toml +++ b/pallets/subtensor/Cargo.toml @@ -55,6 +55,7 @@ sha2.workspace = true rand_chacha.workspace = true pallet-crowdloan.workspace = true pallet-subtensor-proxy.workspace = true +pallet-rate-limiting.workspace = true [dev-dependencies] pallet-balances = { workspace = true, features = ["std"] } @@ -112,6 +113,7 @@ std = [ "pallet-crowdloan/std", "pallet-drand/std", "pallet-subtensor-proxy/std", + "pallet-rate-limiting/std", "pallet-subtensor-swap/std", "subtensor-swap-interface/std", "pallet-subtensor-utility/std", diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml index b363eb4f5f..5d40215c49 100644 --- a/runtime/Cargo.toml +++ b/runtime/Cargo.toml @@ -53,6 +53,7 @@ sp-inherents.workspace = true sp-offchain.workspace = true sp-runtime.workspace = true sp-session.workspace = true +sp-io.workspace = true sp-std.workspace = true sp-transaction-pool.workspace = true sp-version.workspace = true @@ -154,7 +155,6 @@ ethereum.workspace = true [dev-dependencies] frame-metadata.workspace = true -sp-io.workspace = true sp-tracing.workspace = true [build-dependencies] diff --git a/runtime/src/rate_limiting/migration.rs b/runtime/src/rate_limiting/migration.rs new file mode 100644 index 0000000000..243c43b7ea --- /dev/null +++ b/runtime/src/rate_limiting/migration.rs @@ -0,0 +1,733 @@ +use core::convert::TryFrom; + +use codec::Encode; +use frame_support::{pallet_prelude::Parameter, traits::Get, weights::Weight}; +use frame_system::pallet_prelude::BlockNumberFor; +use log::info; +use pallet_rate_limiting::{RateLimit, RateLimitKind, TransactionIdentifier}; +use sp_io::{ + hashing::{blake2_128, twox_128}, + storage, +}; +use sp_runtime::traits::SaturatedConversion; +use sp_std::{collections::btree_map::BTreeMap, vec::Vec}; +use subtensor_runtime_common::{MechId, NetUid, RateLimitScope, RateLimitUsageKey}; + +use pallet_subtensor::{ + self, + utils::rate_limiting::{Hyperparameter, TransactionType}, + AssociatedEvmAddress, Axons, Config as SubtensorConfig, HasMigrationRun, LastRateLimitedBlock, + LastUpdate, MaxUidsTrimmingRateLimit, MechanismCountCurrent, MechanismCountSetRateLimit, + MechanismEmissionRateLimit, NetworkRateLimit, OwnerHyperparamRateLimit, Pallet, Prometheus, + RateLimitKey, TransactionKeyLastBlock, TxChildkeyTakeRateLimit, TxDelegateTakeRateLimit, + TxRateLimit, WeightsVersionKeyRateLimit, +}; + +/// Pallet index assigned to `pallet_subtensor` in `construct_runtime!`. +const SUBTENSOR_PALLET_INDEX: u8 = 7; +/// Pallet index assigned to `pallet_admin_utils` in `construct_runtime!`. +const ADMIN_UTILS_PALLET_INDEX: u8 = 19; + +/// Marker stored in `HasMigrationRun` once the migration finishes. +const MIGRATION_NAME: &[u8] = b"migrate_rate_limiting"; + +/// `set_children` is rate-limited to once every 150 blocks. +const SET_CHILDREN_RATE_LIMIT: u64 = 150; +/// `set_sn_owner_hotkey` default interval (blocks). +const DEFAULT_SET_SN_OWNER_HOTKEY_LIMIT: u64 = 50_400; + +/// Subtensor call indices that reuse the serving rate-limit configuration. +/// TODO(grouped-rate-limits): `serve_axon` (4), `serve_axon_tls` (40), and +/// `serve_prometheus` (5) share one cooldown today. The new pallet still misses +/// grouped identifiers, so we simply port the timers as-is. +const SERVE_CALLS: [u8; 3] = [4, 40, 5]; +/// Subtensor call indices that reuse the per-subnet weight limit. +/// TODO(grouped-rate-limits): Weight commits via call 100 still touch the same +/// `LastUpdate` entries but cannot be expressed here until grouping exists. +const WEIGHT_CALLS_SUBNET: [u8; 3] = [0, 96, 113]; +/// Subtensor call indices that reuse the per-mechanism weight limit. +const WEIGHT_CALLS_MECHANISM: [u8; 4] = [119, 115, 117, 118]; +/// Subtensor call indices for register-network extrinsics. +/// TODO(grouped-rate-limits): `register_network` (59) and +/// `register_network_with_identity` (79) still share the same helper and should +/// remain grouped once pallet-rate-limiting supports aliases. +const REGISTER_NETWORK_CALLS: [u8; 2] = [59, 79]; + +/// Hyperparameter extrinsics routed through owner-or-root rate limiting. +const HYPERPARAMETERS: &[Hyperparameter] = &[ + Hyperparameter::ServingRateLimit, + Hyperparameter::MaxDifficulty, + Hyperparameter::AdjustmentAlpha, + Hyperparameter::ImmunityPeriod, + Hyperparameter::MinAllowedWeights, + Hyperparameter::MaxAllowedUids, + Hyperparameter::Kappa, + Hyperparameter::Rho, + Hyperparameter::ActivityCutoff, + Hyperparameter::PowRegistrationAllowed, + Hyperparameter::MinBurn, + Hyperparameter::MaxBurn, + Hyperparameter::BondsMovingAverage, + Hyperparameter::BondsPenalty, + Hyperparameter::CommitRevealEnabled, + Hyperparameter::LiquidAlphaEnabled, + Hyperparameter::AlphaValues, + Hyperparameter::WeightCommitInterval, + Hyperparameter::TransferEnabled, + Hyperparameter::AlphaSigmoidSteepness, + Hyperparameter::Yuma3Enabled, + Hyperparameter::BondsResetEnabled, + Hyperparameter::ImmuneNeuronLimit, + Hyperparameter::RecycleOrBurn, +]; + +type RateLimitConfigOf = RateLimit>; +type LimitEntries = Vec<(TransactionIdentifier, RateLimitConfigOf)>; +type LastSeenKey = ( + TransactionIdentifier, + Option::AccountId>>, +); +type LastSeenEntries = Vec<(LastSeenKey, BlockNumberFor)>; + +pub fn migrate_rate_limiting() -> Weight { + let mut weight = T::DbWeight::get().reads(1); + if HasMigrationRun::::get(MIGRATION_NAME) { + info!("Rate-limiting migration already executed. Skipping."); + return weight; + } + + let (limits, limit_reads) = build_limits::(); + let (last_seen, seen_reads) = build_last_seen::(); + + let limit_writes = write_limits::(&limits); + let seen_writes = write_last_seen::(&last_seen); + + HasMigrationRun::::insert(MIGRATION_NAME, true); + + weight = weight + .saturating_add(T::DbWeight::get().reads(limit_reads.saturating_add(seen_reads))) + .saturating_add( + T::DbWeight::get().writes(limit_writes.saturating_add(seen_writes).saturating_add(1)), + ); + + info!( + "Migrated {} rate-limit configs and {} last-seen entries into pallet-rate-limiting", + limits.len(), + last_seen.len() + ); + + weight +} + +fn build_limits() -> (LimitEntries, u64) { + let mut limits = LimitEntries::::new(); + let mut reads: u64 = 0; + + reads += gather_simple_limits::(&mut limits); + reads += gather_owner_hparam_limits::(&mut limits); + reads += gather_serving_limits::(&mut limits); + reads += gather_weight_limits::(&mut limits); + + (limits, reads) +} + +fn gather_simple_limits(limits: &mut LimitEntries) -> u64 { + let mut reads: u64 = 0; + + reads += 1; + if let Some(span) = block_number::(TxRateLimit::::get()) { + set_global_limit::(limits, subtensor_identifier(70), span); + } + + reads += 1; + if let Some(span) = block_number::(TxDelegateTakeRateLimit::::get()) { + // TODO(grouped-rate-limits): `decrease_take` shares the same timestamp but + // does not have its own ID here yet. + set_global_limit::(limits, subtensor_identifier(66), span); + } + + reads += 1; + if let Some(span) = block_number::(TxChildkeyTakeRateLimit::::get()) { + set_global_limit::(limits, subtensor_identifier(75), span); + } + + reads += 1; + if let Some(span) = block_number::(NetworkRateLimit::::get()) { + for call in REGISTER_NETWORK_CALLS { + set_global_limit::(limits, subtensor_identifier(call), span); + } + } + + reads += 1; + if let Some(span) = block_number::(WeightsVersionKeyRateLimit::::get()) { + set_global_limit::(limits, admin_utils_identifier(6), span); + } + + if let Some(span) = block_number::(DEFAULT_SET_SN_OWNER_HOTKEY_LIMIT) { + set_global_limit::(limits, admin_utils_identifier(67), span); + } + + if let Some(span) = block_number::(::EvmKeyAssociateRateLimit::get()) { + set_global_limit::(limits, subtensor_identifier(93), span); + } + + if let Some(span) = block_number::(MechanismCountSetRateLimit::::get()) { + set_global_limit::(limits, admin_utils_identifier(76), span); + } + + if let Some(span) = block_number::(MechanismEmissionRateLimit::::get()) { + set_global_limit::(limits, admin_utils_identifier(77), span); + } + + if let Some(span) = block_number::(MaxUidsTrimmingRateLimit::::get()) { + set_global_limit::(limits, admin_utils_identifier(78), span); + } + + if let Some(span) = block_number::(SET_CHILDREN_RATE_LIMIT) { + set_global_limit::(limits, subtensor_identifier(67), span); + } + + reads +} + +fn gather_owner_hparam_limits(limits: &mut LimitEntries) -> u64 { + let mut reads: u64 = 0; + + reads += 1; + if let Some(span) = block_number::(u64::from(OwnerHyperparamRateLimit::::get())) { + for hparam in HYPERPARAMETERS { + if let Some(identifier) = identifier_for_hyperparameter(*hparam) { + set_global_limit::(limits, identifier, span); + } + } + } + + reads +} + +fn gather_serving_limits(limits: &mut LimitEntries) -> u64 { + let mut reads: u64 = 0; + let netuids = Pallet::::get_all_subnet_netuids(); + + for netuid in netuids { + reads += 1; + if let Some(span) = block_number::(Pallet::::get_serving_rate_limit(netuid)) { + for call in SERVE_CALLS { + set_scoped_limit::( + limits, + subtensor_identifier(call), + RateLimitScope::Subnet(netuid), + span, + ); + } + } + } + + reads +} + +fn gather_weight_limits(limits: &mut LimitEntries) -> u64 { + let mut reads: u64 = 0; + let netuids = Pallet::::get_all_subnet_netuids(); + + let mut subnet_limits = BTreeMap::>::new(); + for netuid in &netuids { + reads += 1; + if let Some(span) = block_number::(Pallet::::get_weights_set_rate_limit(*netuid)) { + subnet_limits.insert(*netuid, span); + for call in WEIGHT_CALLS_SUBNET { + set_scoped_limit::( + limits, + subtensor_identifier(call), + RateLimitScope::Subnet(*netuid), + span, + ); + } + } + } + + for netuid in &netuids { + reads += 1; + let mech_count: u8 = MechanismCountCurrent::::get(*netuid).into(); + if mech_count <= 1 { + continue; + } + let Some(span) = subnet_limits.get(netuid).copied() else { + continue; + }; + for mecid in 1..mech_count { + let scope = RateLimitScope::SubnetMechanism { + netuid: *netuid, + mecid: MechId::from(mecid), + }; + for call in WEIGHT_CALLS_MECHANISM { + set_scoped_limit::(limits, subtensor_identifier(call), scope.clone(), span); + } + } + } + + reads +} + +fn build_last_seen() -> (LastSeenEntries, u64) { + let mut last_seen = LastSeenEntries::::new(); + let mut reads: u64 = 0; + + reads += import_last_rate_limited_blocks::(&mut last_seen); + reads += import_transaction_key_last_blocks::(&mut last_seen); + reads += import_last_update_entries::(&mut last_seen); + reads += import_serving_entries::(&mut last_seen); + reads += import_evm_entries::(&mut last_seen); + + (last_seen, reads) +} + +fn import_last_rate_limited_blocks(entries: &mut LastSeenEntries) -> u64 { + let mut reads: u64 = 0; + for (key, block) in LastRateLimitedBlock::::iter() { + reads += 1; + if block == 0 { + continue; + } + match key { + RateLimitKey::SetSNOwnerHotkey(netuid) => { + if let Some(identifier) = + identifier_for_transaction_type(TransactionType::SetSNOwnerHotkey) + { + record_last_seen_entry::( + entries, + identifier, + Some(RateLimitUsageKey::Subnet(netuid)), + block, + ); + } + } + RateLimitKey::OwnerHyperparamUpdate(netuid, hyper) => { + if let Some(identifier) = identifier_for_hyperparameter(hyper) { + record_last_seen_entry::( + entries, + identifier, + Some(RateLimitUsageKey::Subnet(netuid)), + block, + ); + } + } + RateLimitKey::LastTxBlock(account) => { + record_last_seen_entry::( + entries, + subtensor_identifier(70), + Some(RateLimitUsageKey::Account(account.clone())), + block, + ); + } + RateLimitKey::LastTxBlockDelegateTake(account) => { + record_last_seen_entry::( + entries, + subtensor_identifier(66), + Some(RateLimitUsageKey::Account(account.clone())), + block, + ); + } + RateLimitKey::NetworkLastRegistered | RateLimitKey::LastTxBlockChildKeyTake(_) => { + // TODO(grouped-rate-limits): Global network registration lock is still outside + // pallet-rate-limiting. We will migrate it once grouped identifiers land. + } + } + } + reads +} + +fn import_transaction_key_last_blocks(entries: &mut LastSeenEntries) -> u64 { + let mut reads: u64 = 0; + for ((account, netuid, tx_kind), block) in TransactionKeyLastBlock::::iter() { + reads += 1; + if block == 0 { + continue; + } + let tx_type = TransactionType::from(tx_kind); + let Some(identifier) = identifier_for_transaction_type(tx_type) else { + continue; + }; + let Some(usage) = usage_key_from_transaction_type(tx_type, &account, netuid) else { + continue; + }; + record_last_seen_entry::(entries, identifier, Some(usage), block); + } + reads +} + +fn import_last_update_entries(entries: &mut LastSeenEntries) -> u64 { + let mut reads: u64 = 0; + for (index, blocks) in LastUpdate::::iter() { + reads += 1; + let netuid = Pallet::::get_netuid(index); + let sub_id = u16::from(index) + .checked_div(pallet_subtensor::subnets::mechanism::GLOBAL_MAX_SUBNET_COUNT) + .unwrap_or_default(); + let is_mechanism = sub_id != 0; + let Ok(sub_id) = u8::try_from(sub_id) else { + continue; + }; + let mecid = MechId::from(sub_id); + + for (uid, last_block) in blocks.into_iter().enumerate() { + if last_block == 0 { + continue; + } + let Ok(uid_u16) = u16::try_from(uid) else { + continue; + }; + let usage = if is_mechanism { + RateLimitUsageKey::SubnetMechanismNeuron { + netuid, + mecid, + uid: uid_u16, + } + } else { + RateLimitUsageKey::SubnetNeuron { + netuid, + uid: uid_u16, + } + }; + + let call_set: &[u8] = if is_mechanism { + &WEIGHT_CALLS_MECHANISM + } else { + &WEIGHT_CALLS_SUBNET + }; + + for call in call_set { + record_last_seen_entry::( + entries, + subtensor_identifier(*call), + Some(usage.clone()), + last_block, + ); + } + } + } + reads +} + +fn import_serving_entries(entries: &mut LastSeenEntries) -> u64 { + let mut reads: u64 = 0; + for (netuid, hotkey, axon) in Axons::::iter() { + reads += 1; + if axon.block == 0 { + continue; + } + let usage = RateLimitUsageKey::AccountSubnet { + account: hotkey.clone(), + netuid, + }; + for call in [4u8, 40u8] { + record_last_seen_entry::( + entries, + subtensor_identifier(call), + Some(usage.clone()), + axon.block, + ); + } + } + + for (netuid, hotkey, prom) in Prometheus::::iter() { + reads += 1; + if prom.block == 0 { + continue; + } + let usage = RateLimitUsageKey::AccountSubnet { + account: hotkey, + netuid, + }; + record_last_seen_entry::(entries, subtensor_identifier(5), Some(usage), prom.block); + } + + reads +} + +fn import_evm_entries(entries: &mut LastSeenEntries) -> u64 { + let mut reads: u64 = 0; + for (netuid, uid, (_, block)) in AssociatedEvmAddress::::iter() { + reads += 1; + if block == 0 { + continue; + } + record_last_seen_entry::( + entries, + subtensor_identifier(93), + Some(RateLimitUsageKey::SubnetNeuron { netuid, uid }), + block, + ); + } + reads +} + +/// TODO(rate-limiting-storage): Swap these manual writes for +/// `pallet_rate_limiting::Pallet` APIs once the runtime wires the pallet in. +fn write_limits(limits: &LimitEntries) -> u64 { + if limits.is_empty() { + return 0; + } + let prefix = storage_prefix("RateLimiting", "Limits"); + let mut writes = 0; + for (identifier, limit) in limits.iter() { + let key = map_storage_key(&prefix, identifier); + storage::set(&key, &limit.encode()); + writes += 1; + } + writes +} + +fn write_last_seen(entries: &LastSeenEntries) -> u64 { + if entries.is_empty() { + return 0; + } + let prefix = storage_prefix("RateLimiting", "LastSeen"); + let mut writes = 0; + for ((identifier, usage), block) in entries.iter() { + let key = double_map_storage_key(&prefix, identifier, usage); + storage::set(&key, &block.encode()); + writes += 1; + } + writes +} + +fn block_number(value: u64) -> Option> { + if value == 0 { + return None; + } + Some(value.saturated_into::>()) +} + +fn set_global_limit( + limits: &mut LimitEntries, + identifier: TransactionIdentifier, + span: BlockNumberFor, +) { + if let Some((_, config)) = limits.iter_mut().find(|(id, _)| *id == identifier) { + *config = RateLimit::global(RateLimitKind::Exact(span)); + } else { + limits.push((identifier, RateLimit::global(RateLimitKind::Exact(span)))); + } +} + +fn set_scoped_limit( + limits: &mut LimitEntries, + identifier: TransactionIdentifier, + scope: RateLimitScope, + span: BlockNumberFor, +) { + if let Some((_, config)) = limits.iter_mut().find(|(id, _)| *id == identifier) { + match config { + RateLimit::Global(_) => { + *config = RateLimit::scoped_single(scope, RateLimitKind::Exact(span)); + } + RateLimit::Scoped(map) => { + map.insert(scope, RateLimitKind::Exact(span)); + } + } + } else { + limits.push(( + identifier, + RateLimit::scoped_single(scope, RateLimitKind::Exact(span)), + )); + } +} + +fn record_last_seen_entry( + entries: &mut LastSeenEntries, + identifier: TransactionIdentifier, + usage: Option>, + block: u64, +) { + let Some(block_number) = block_number::(block) else { + return; + }; + + let key = (identifier, usage); + if let Some((_, existing)) = entries.iter_mut().find(|(entry_key, _)| *entry_key == key) { + if block_number > *existing { + *existing = block_number; + } + } else { + entries.push((key, block_number)); + } +} + +fn storage_prefix(pallet: &str, storage: &str) -> Vec { + let mut out = Vec::with_capacity(32); + out.extend_from_slice(&twox_128(pallet.as_bytes())); + out.extend_from_slice(&twox_128(storage.as_bytes())); + out +} + +fn map_storage_key(prefix: &[u8], key: impl Encode) -> Vec { + let mut final_key = Vec::with_capacity(prefix.len() + 32); + final_key.extend_from_slice(prefix); + let encoded = key.encode(); + let hash = blake2_128(&encoded); + final_key.extend_from_slice(&hash); + final_key.extend_from_slice(&encoded); + final_key +} + +fn double_map_storage_key(prefix: &[u8], key1: impl Encode, key2: impl Encode) -> Vec { + let mut final_key = Vec::with_capacity(prefix.len() + 64); + final_key.extend_from_slice(prefix); + let first = map_storage_key(&[], key1); + final_key.extend_from_slice(&first); + let second = map_storage_key(&[], key2); + final_key.extend_from_slice(&second); + final_key +} + +const fn admin_utils_identifier(call_index: u8) -> TransactionIdentifier { + TransactionIdentifier::new(ADMIN_UTILS_PALLET_INDEX, call_index) +} + +const fn subtensor_identifier(call_index: u8) -> TransactionIdentifier { + TransactionIdentifier::new(SUBTENSOR_PALLET_INDEX, call_index) +} + +/// Returns the `TransactionIdentifier` for the admin-utils extrinsic that controls `hparam`. +/// +/// Only hyperparameters that are currently rate-limited (i.e. routed through +/// `ensure_sn_owner_or_root_with_limits`) are mapped; others return `None`. +pub fn identifier_for_hyperparameter(hparam: Hyperparameter) -> Option { + use Hyperparameter::*; + + let identifier = match hparam { + Unknown | MaxWeightLimit => return None, + ServingRateLimit => admin_utils_identifier(3), + MaxDifficulty => admin_utils_identifier(5), + AdjustmentAlpha => admin_utils_identifier(9), + ImmunityPeriod => admin_utils_identifier(13), + MinAllowedWeights => admin_utils_identifier(14), + MaxAllowedUids => admin_utils_identifier(15), + Kappa => admin_utils_identifier(16), + Rho => admin_utils_identifier(17), + ActivityCutoff => admin_utils_identifier(18), + PowRegistrationAllowed => admin_utils_identifier(20), + MinBurn => admin_utils_identifier(22), + MaxBurn => admin_utils_identifier(23), + BondsMovingAverage => admin_utils_identifier(26), + BondsPenalty => admin_utils_identifier(60), + CommitRevealEnabled => admin_utils_identifier(49), + LiquidAlphaEnabled => admin_utils_identifier(50), + AlphaValues => admin_utils_identifier(51), + WeightCommitInterval => admin_utils_identifier(57), + TransferEnabled => admin_utils_identifier(61), + AlphaSigmoidSteepness => admin_utils_identifier(68), + Yuma3Enabled => admin_utils_identifier(69), + BondsResetEnabled => admin_utils_identifier(70), + ImmuneNeuronLimit => admin_utils_identifier(72), + RecycleOrBurn => admin_utils_identifier(80), + _ => return None, + }; + + Some(identifier) +} + +/// Returns the `TransactionIdentifier` for the extrinsic associated with the given transaction +/// type, mirroring current rate-limit enforcement. +pub fn identifier_for_transaction_type(tx: TransactionType) -> Option { + use TransactionType::*; + + let identifier = match tx { + SetChildren => subtensor_identifier(67), + SetChildkeyTake => subtensor_identifier(75), + RegisterNetwork => subtensor_identifier(59), + SetWeightsVersionKey => admin_utils_identifier(6), + SetSNOwnerHotkey => admin_utils_identifier(67), + OwnerHyperparamUpdate(hparam) => return identifier_for_hyperparameter(hparam), + MechanismCountUpdate => admin_utils_identifier(76), + MechanismEmission => admin_utils_identifier(77), + MaxUidsTrimming => admin_utils_identifier(78), + Unknown => return None, + _ => return None, + }; + + Some(identifier) +} + +/// Maps legacy `RateLimitKey` entries to the new usage-key representation. +pub fn usage_key_from_legacy_key( + key: &RateLimitKey, +) -> Option> +where + AccountId: Parameter + Clone, +{ + match key { + RateLimitKey::SetSNOwnerHotkey(netuid) => Some(RateLimitUsageKey::Subnet(*netuid)), + RateLimitKey::OwnerHyperparamUpdate(netuid, _) => Some(RateLimitUsageKey::Subnet(*netuid)), + RateLimitKey::NetworkLastRegistered => None, + RateLimitKey::LastTxBlock(account) + | RateLimitKey::LastTxBlockChildKeyTake(account) + | RateLimitKey::LastTxBlockDelegateTake(account) => { + Some(RateLimitUsageKey::Account(account.clone())) + } + } +} + +/// Produces the usage key for a `TransactionType` that was stored in `TransactionKeyLastBlock`. +pub fn usage_key_from_transaction_type( + tx: TransactionType, + account: &AccountId, + netuid: NetUid, +) -> Option> +where + AccountId: Parameter + Clone, +{ + match tx { + TransactionType::SetChildren | TransactionType::SetChildkeyTake => { + Some(RateLimitUsageKey::AccountSubnet { + account: account.clone(), + netuid, + }) + } + TransactionType::SetWeightsVersionKey => Some(RateLimitUsageKey::Subnet(netuid)), + TransactionType::MechanismCountUpdate + | TransactionType::MechanismEmission + | TransactionType::MaxUidsTrimming => Some(RateLimitUsageKey::AccountSubnet { + account: account.clone(), + netuid, + }), + TransactionType::OwnerHyperparamUpdate(_) => Some(RateLimitUsageKey::Subnet(netuid)), + TransactionType::RegisterNetwork => Some(RateLimitUsageKey::Account(account.clone())), + TransactionType::SetSNOwnerHotkey => Some(RateLimitUsageKey::Subnet(netuid)), + TransactionType::Unknown => None, + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn maps_hyperparameters() { + assert_eq!( + identifier_for_hyperparameter(Hyperparameter::ServingRateLimit), + Some(admin_utils_identifier(3)) + ); + assert!(identifier_for_hyperparameter(Hyperparameter::MaxWeightLimit).is_none()); + } + + #[test] + fn maps_transaction_types() { + assert_eq!( + identifier_for_transaction_type(TransactionType::SetChildren), + Some(subtensor_identifier(67)) + ); + assert!(identifier_for_transaction_type(TransactionType::Unknown).is_none()); + } + + #[test] + fn maps_usage_keys() { + let acct = 42u64; + assert!(matches!( + usage_key_from_legacy_key(&RateLimitKey::LastTxBlock(acct)), + Some(RateLimitUsageKey::Account(42)) + )); + } +} diff --git a/runtime/src/rate_limiting.rs b/runtime/src/rate_limiting/mod.rs similarity index 99% rename from runtime/src/rate_limiting.rs rename to runtime/src/rate_limiting/mod.rs index 019cdcd458..713c8bacf6 100644 --- a/runtime/src/rate_limiting.rs +++ b/runtime/src/rate_limiting/mod.rs @@ -6,6 +6,8 @@ use subtensor_runtime_common::{BlockNumber, NetUid, RateLimitScope, RateLimitUsa use crate::{AccountId, Runtime, RuntimeCall, RuntimeOrigin}; +pub(crate) mod migration; + fn signed_origin(origin: &RuntimeOrigin) -> Option { match origin.clone().into() { Ok(RawOrigin::Signed(who)) => Some(who),