Skip to content

Commit 232de64

Browse files
ales-otfales-otf
committed
Pass origin to rate limit context resolvers
1 parent d65baff commit 232de64

File tree

7 files changed

+88
-40
lines changed

7 files changed

+88
-40
lines changed

‎pallets/rate-limiting/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ frame-system.workspace = true
1717
scale-info = { workspace = true, features = ["derive"] }
1818
serde = { workspace = true, features = ["derive"], optional = true }
1919
sp-std.workspace = true
20+
sp-runtime.workspace = true
2021
subtensor-runtime-common.workspace = true
2122

2223
[dev-dependencies]
@@ -34,6 +35,7 @@ std = [
3435
"scale-info/std",
3536
"serde",
3637
"sp-std/std",
38+
"sp-runtime/std",
3739
"subtensor-runtime-common/std",
3840
]
3941
runtime-benchmarks = [

‎pallets/rate-limiting/src/benchmarking.rs‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
use codec::Decode;
66
use frame_benchmarking::v2::*;
77
use frame_system::{RawOrigin, pallet_prelude::BlockNumberFor};
8+
use sp_runtime::traits::DispatchOriginOf;
89

910
use super::*;
1011

@@ -36,7 +37,10 @@ mod benchmarks {
3637
fn set_rate_limit() {
3738
let call = sample_call::<T>();
3839
let limit = RateLimitKind::<BlockNumberFor<T>>::Exact(BlockNumberFor::<T>::from(10u32));
39-
let scope = <T as Config>::LimitScopeResolver::context(call.as_ref());
40+
let origin = T::RuntimeOrigin::from(RawOrigin::Root);
41+
let resolver_origin: DispatchOriginOf<<T as Config>::RuntimeCall> =
42+
Into::<DispatchOriginOf<<T as Config>::RuntimeCall>>::into(origin.clone());
43+
let scope = <T as Config>::LimitScopeResolver::context(&resolver_origin, call.as_ref());
4044
let identifier =
4145
TransactionIdentifier::from_call::<T, ()>(call.as_ref()).expect("identifier");
4246

@@ -61,7 +65,10 @@ mod benchmarks {
6165
fn clear_rate_limit() {
6266
let call = sample_call::<T>();
6367
let limit = RateLimitKind::<BlockNumberFor<T>>::Exact(BlockNumberFor::<T>::from(10u32));
64-
let scope = <T as Config>::LimitScopeResolver::context(call.as_ref());
68+
let origin = T::RuntimeOrigin::from(RawOrigin::Root);
69+
let resolver_origin: DispatchOriginOf<<T as Config>::RuntimeCall> =
70+
Into::<DispatchOriginOf<<T as Config>::RuntimeCall>>::into(origin.clone());
71+
let scope = <T as Config>::LimitScopeResolver::context(&resolver_origin, call.as_ref());
6572

6673
// Pre-populate limit for benchmark call
6774
let identifier =

‎pallets/rate-limiting/src/lib.rs‎

Lines changed: 48 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -55,17 +55,22 @@
5555
//! [`LastSeen`](pallet::LastSeen). This can refine the limit scope (for example by returning a
5656
//! tuple of `(netuid, hyperparameter)`).
5757
//!
58-
//! Each resolver receives the call and may return `Some(identifier)` when scoping is required, or
59-
//! `None` to use the global entry. Extrinsics such as
58+
//! Each resolver receives the origin and call and may return `Some(identifier)` when scoping is
59+
//! required, or `None` to use the global entry. Extrinsics such as
6060
//! [`set_rate_limit`](pallet::Pallet::set_rate_limit) automatically consult these resolvers.
6161
//!
6262
//! ```ignore
6363
//! pub struct WeightsContextResolver;
6464
//!
6565
//! // Limits are scoped per netuid.
6666
//! pub struct ScopeResolver;
67-
//! impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, NetUid, BlockNumber> for ScopeResolver {
68-
//! fn context(call: &RuntimeCall) -> Option<NetUid> {
67+
//! impl pallet_rate_limiting::RateLimitScopeResolver<
68+
//! RuntimeOrigin,
69+
//! RuntimeCall,
70+
//! NetUid,
71+
//! BlockNumber,
72+
//! > for ScopeResolver {
73+
//! fn context(origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<NetUid> {
6974
//! match call {
7075
//! RuntimeCall::Subtensor(pallet_subtensor::Call::set_weights { netuid, .. }) => {
7176
//! Some(*netuid)
@@ -74,15 +79,23 @@
7479
//! }
7580
//! }
7681
//!
77-
//! fn adjust_span(_call: &RuntimeCall, span: BlockNumber) -> BlockNumber {
82+
//! fn should_bypass(origin: &RuntimeOrigin, _call: &RuntimeCall) -> bool {
83+
//! matches!(origin, RuntimeOrigin::Root)
84+
//! }
85+
//!
86+
//! fn adjust_span(_origin: &RuntimeOrigin, _call: &RuntimeCall, span: BlockNumber) -> BlockNumber {
7887
//! span
7988
//! }
8089
//! }
8190
//!
8291
//! // Usage tracking distinguishes hyperparameter + netuid.
8392
//! pub struct UsageResolver;
84-
//! impl pallet_rate_limiting::RateLimitUsageResolver<RuntimeCall, (NetUid, HyperParam)> for UsageResolver {
85-
//! fn context(call: &RuntimeCall) -> Option<(NetUid, HyperParam)> {
93+
//! impl pallet_rate_limiting::RateLimitUsageResolver<
94+
//! RuntimeOrigin,
95+
//! RuntimeCall,
96+
//! (NetUid, HyperParam),
97+
//! > for UsageResolver {
98+
//! fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<(NetUid, HyperParam)> {
8699
//! match call {
87100
//! RuntimeCall::Subtensor(pallet_subtensor::Call::set_hyperparam {
88101
//! netuid,
@@ -128,10 +141,10 @@ pub mod pallet {
128141
use codec::Codec;
129142
use frame_support::{
130143
pallet_prelude::*,
131-
sp_runtime::traits::{Saturating, Zero},
132144
traits::{BuildGenesisConfig, EnsureOrigin, GetCallMetadata},
133145
};
134146
use frame_system::pallet_prelude::*;
147+
use sp_runtime::traits::{DispatchOriginOf, Dispatchable, Saturating, Zero};
135148
use sp_std::{convert::TryFrom, marker::PhantomData, vec::Vec};
136149

137150
#[cfg(feature = "runtime-benchmarks")]
@@ -146,11 +159,14 @@ pub mod pallet {
146159
pub trait Config<I: 'static = ()>: frame_system::Config
147160
where
148161
BlockNumberFor<Self>: MaybeSerializeDeserialize,
162+
<<Self as Config<I>>::RuntimeCall as Dispatchable>::RuntimeOrigin:
163+
From<<Self as frame_system::Config>::RuntimeOrigin>,
149164
{
150165
/// The overarching runtime call type.
151166
type RuntimeCall: Parameter
152167
+ Codec
153168
+ GetCallMetadata
169+
+ Dispatchable
154170
+ IsType<<Self as frame_system::Config>::RuntimeCall>;
155171

156172
/// Origin permitted to configure rate limits.
@@ -161,6 +177,7 @@ pub mod pallet {
161177

162178
/// Resolves the scope for the given runtime call when configuring limits.
163179
type LimitScopeResolver: RateLimitScopeResolver<
180+
DispatchOriginOf<<Self as Config<I>>::RuntimeCall>,
164181
<Self as Config<I>>::RuntimeCall,
165182
Self::LimitScope,
166183
BlockNumberFor<Self>,
@@ -170,7 +187,11 @@ pub mod pallet {
170187
type UsageKey: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize;
171188

172189
/// Resolves the usage key for the given runtime call when enforcing limits.
173-
type UsageResolver: RateLimitUsageResolver<<Self as Config<I>>::RuntimeCall, Self::UsageKey>;
190+
type UsageResolver: RateLimitUsageResolver<
191+
DispatchOriginOf<<Self as Config<I>>::RuntimeCall>,
192+
<Self as Config<I>>::RuntimeCall,
193+
Self::UsageKey,
194+
>;
174195

175196
/// Helper used to construct runtime calls for benchmarking.
176197
#[cfg(feature = "runtime-benchmarks")]
@@ -311,16 +332,17 @@ pub mod pallet {
311332
/// Returns `true` when the given transaction identifier passes its configured rate limit
312333
/// within the provided usage scope.
313334
pub fn is_within_limit(
335+
origin: &DispatchOriginOf<<T as Config<I>>::RuntimeCall>,
336+
call: &<T as Config<I>>::RuntimeCall,
314337
identifier: &TransactionIdentifier,
315338
scope: &Option<<T as Config<I>>::LimitScope>,
316339
usage_key: &Option<<T as Config<I>>::UsageKey>,
317-
call: &<T as Config<I>>::RuntimeCall,
318340
) -> Result<bool, DispatchError> {
319-
if <T as Config<I>>::LimitScopeResolver::should_bypass(call) {
341+
if <T as Config<I>>::LimitScopeResolver::should_bypass(origin, call) {
320342
return Ok(true);
321343
}
322344

323-
let Some(block_span) = Self::effective_span(call, identifier, scope) else {
345+
let Some(block_span) = Self::effective_span(origin, call, identifier, scope) else {
324346
return Ok(true);
325347
};
326348

@@ -340,13 +362,14 @@ pub mod pallet {
340362
}
341363

342364
pub(crate) fn effective_span(
365+
origin: &DispatchOriginOf<<T as Config<I>>::RuntimeCall>,
343366
call: &<T as Config<I>>::RuntimeCall,
344367
identifier: &TransactionIdentifier,
345368
scope: &Option<<T as Config<I>>::LimitScope>,
346369
) -> Option<BlockNumberFor<T>> {
347370
let span = Self::resolved_limit(identifier, scope)?;
348371
Some(<T as Config<I>>::LimitScopeResolver::adjust_span(
349-
call, span,
372+
origin, call, span,
350373
))
351374
}
352375

@@ -491,11 +514,15 @@ pub mod pallet {
491514
call: Box<<T as Config<I>>::RuntimeCall>,
492515
limit: RateLimitKind<BlockNumberFor<T>>,
493516
) -> DispatchResult {
517+
let resolver_origin: DispatchOriginOf<<T as Config<I>>::RuntimeCall> =
518+
Into::<DispatchOriginOf<<T as Config<I>>::RuntimeCall>>::into(origin.clone());
519+
let scope =
520+
<T as Config<I>>::LimitScopeResolver::context(&resolver_origin, call.as_ref());
521+
let scope_for_event = scope.clone();
522+
494523
T::AdminOrigin::ensure_origin(origin)?;
495524

496525
let identifier = TransactionIdentifier::from_call::<T, I>(call.as_ref())?;
497-
let scope = <T as Config<I>>::LimitScopeResolver::context(call.as_ref());
498-
let scope_for_event = scope.clone();
499526

500527
if let Some(ref sc) = scope {
501528
Limits::<T, I>::mutate(&identifier, |slot| match slot {
@@ -532,11 +559,15 @@ pub mod pallet {
532559
origin: OriginFor<T>,
533560
call: Box<<T as Config<I>>::RuntimeCall>,
534561
) -> DispatchResult {
562+
let resolver_origin: DispatchOriginOf<<T as Config<I>>::RuntimeCall> =
563+
Into::<DispatchOriginOf<<T as Config<I>>::RuntimeCall>>::into(origin.clone());
564+
let scope =
565+
<T as Config<I>>::LimitScopeResolver::context(&resolver_origin, call.as_ref());
566+
let usage = <T as Config<I>>::UsageResolver::context(&resolver_origin, call.as_ref());
567+
535568
T::AdminOrigin::ensure_origin(origin)?;
536569

537570
let identifier = TransactionIdentifier::from_call::<T, I>(call.as_ref())?;
538-
let scope = <T as Config<I>>::LimitScopeResolver::context(call.as_ref());
539-
let usage = <T as Config<I>>::UsageResolver::context(call.as_ref());
540571

541572
let (pallet_name, extrinsic_name) = identifier.names::<T, I>()?;
542573
let pallet = Vec::from(pallet_name.as_bytes());

‎pallets/rate-limiting/src/mock.rs‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,10 @@ pub type UsageKey = u16;
6161
pub struct TestScopeResolver;
6262
pub struct TestUsageResolver;
6363

64-
impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, LimitScope, u64>
64+
impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeOrigin, RuntimeCall, LimitScope, u64>
6565
for TestScopeResolver
6666
{
67-
fn context(call: &RuntimeCall) -> Option<LimitScope> {
67+
fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<LimitScope> {
6868
match call {
6969
RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => {
7070
(*block_span).try_into().ok()
@@ -74,14 +74,14 @@ impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, LimitScope, u64>
7474
}
7575
}
7676

77-
fn should_bypass(call: &RuntimeCall) -> bool {
77+
fn should_bypass(_origin: &RuntimeOrigin, call: &RuntimeCall) -> bool {
7878
matches!(
7979
call,
8080
RuntimeCall::RateLimiting(RateLimitingCall::clear_rate_limit { .. })
8181
)
8282
}
8383

84-
fn adjust_span(call: &RuntimeCall, span: u64) -> u64 {
84+
fn adjust_span(_origin: &RuntimeOrigin, call: &RuntimeCall, span: u64) -> u64 {
8585
if matches!(
8686
call,
8787
RuntimeCall::RateLimiting(RateLimitingCall::clear_all_rate_limits { .. })
@@ -93,8 +93,10 @@ impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeCall, LimitScope, u64>
9393
}
9494
}
9595

96-
impl pallet_rate_limiting::RateLimitUsageResolver<RuntimeCall, UsageKey> for TestUsageResolver {
97-
fn context(call: &RuntimeCall) -> Option<UsageKey> {
96+
impl pallet_rate_limiting::RateLimitUsageResolver<RuntimeOrigin, RuntimeCall, UsageKey>
97+
for TestUsageResolver
98+
{
99+
fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<UsageKey> {
98100
match call {
99101
RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => {
100102
(*block_span).try_into().ok()

‎pallets/rate-limiting/src/tests.rs‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,8 @@ fn is_within_limit_is_true_when_no_limit() {
121121
RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 });
122122
let identifier = identifier_for(&call);
123123

124-
let result = RateLimiting::is_within_limit(&identifier, &None, &None, &call);
124+
let origin = RuntimeOrigin::signed(1);
125+
let result = RateLimiting::is_within_limit(&origin, &call, &identifier, &None, &None);
125126
assert_eq!(result.expect("no error expected"), true);
126127
});
127128
}
@@ -140,11 +141,13 @@ fn is_within_limit_false_when_rate_limited() {
140141

141142
System::set_block_number(13);
142143

144+
let origin = RuntimeOrigin::signed(1);
143145
let within = RateLimiting::is_within_limit(
146+
&origin,
147+
&call,
144148
&identifier,
145149
&Some(1 as LimitScope),
146150
&Some(1 as UsageKey),
147-
&call,
148151
)
149152
.expect("call succeeds");
150153
assert!(!within);
@@ -165,11 +168,13 @@ fn is_within_limit_true_after_required_span() {
165168

166169
System::set_block_number(20);
167170

171+
let origin = RuntimeOrigin::signed(1);
168172
let within = RateLimiting::is_within_limit(
173+
&origin,
174+
&call,
169175
&identifier,
170176
&Some(2 as LimitScope),
171177
&Some(2 as UsageKey),
172-
&call,
173178
)
174179
.expect("call succeeds");
175180
assert!(within);

‎pallets/rate-limiting/src/tx_extension.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ where
9797
_inherited_implication: &impl Implication,
9898
_source: TransactionSource,
9999
) -> ValidateResult<Self::Val, <T as Config<I>>::RuntimeCall> {
100-
if <T as Config<I>>::LimitScopeResolver::should_bypass(call) {
100+
if <T as Config<I>>::LimitScopeResolver::should_bypass(&origin, call) {
101101
return Ok((ValidTransaction::default(), None, origin));
102102
}
103103

@@ -106,10 +106,11 @@ where
106106
Err(_) => return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)),
107107
};
108108

109-
let scope = <T as Config<I>>::LimitScopeResolver::context(call);
110-
let usage = <T as Config<I>>::UsageResolver::context(call);
109+
let scope = <T as Config<I>>::LimitScopeResolver::context(&origin, call);
110+
let usage = <T as Config<I>>::UsageResolver::context(&origin, call);
111111

112-
let Some(block_span) = Pallet::<T, I>::effective_span(call, &identifier, &scope) else {
112+
let Some(block_span) = Pallet::<T, I>::effective_span(&origin, call, &identifier, &scope)
113+
else {
113114
return Ok((ValidTransaction::default(), None, origin));
114115
};
115116

‎pallets/rate-limiting/src/types.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,29 +5,29 @@ use sp_std::collections::btree_map::BTreeMap;
55

66
/// Resolves the optional identifier within which a rate limit applies and can optionally adjust
77
/// enforcement behaviour.
8-
pub trait RateLimitScopeResolver<Call, Scope, Span> {
8+
pub trait RateLimitScopeResolver<Origin, Call, Scope, Span> {
99
/// Returns `Some(scope)` when the limit should be applied per-scope, or `None` for global
1010
/// limits.
11-
fn context(call: &Call) -> Option<Scope>;
11+
fn context(origin: &Origin, call: &Call) -> Option<Scope>;
1212

13-
/// Returns `true` when the rate limit should be bypassed for the provided call. Defaults to
14-
/// `false`.
15-
fn should_bypass(_call: &Call) -> bool {
13+
/// Returns `true` when the rate limit should be bypassed for the provided origin/call pair.
14+
/// Defaults to `false`.
15+
fn should_bypass(_origin: &Origin, _call: &Call) -> bool {
1616
false
1717
}
1818

1919
/// Optionally adjusts the effective span used during enforcement. Defaults to the original
2020
/// `span`.
21-
fn adjust_span(_call: &Call, span: Span) -> Span {
21+
fn adjust_span(_origin: &Origin, _call: &Call, span: Span) -> Span {
2222
span
2323
}
2424
}
2525

2626
/// Resolves the optional usage tracking key applied when enforcing limits.
27-
pub trait RateLimitUsageResolver<Call, Usage> {
27+
pub trait RateLimitUsageResolver<Origin, Call, Usage> {
2828
/// Returns `Some(usage)` when usage should be tracked per-key, or `None` for global usage
2929
/// tracking.
30-
fn context(call: &Call) -> Option<Usage>;
30+
fn context(origin: &Origin, call: &Call) -> Option<Usage>;
3131
}
3232

3333
/// Identifies a runtime call by pallet and extrinsic indices.

0 commit comments

Comments
 (0)