Skip to content

Commit 953288f

Browse files
ales-otfales-otf
committed
Fix swap-keys rate-limiting migration
1 parent 56c1af9 commit 953288f

File tree

6 files changed

+135
-74
lines changed

6 files changed

+135
-74
lines changed

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -105,12 +105,12 @@ impl pallet_rate_limiting::RateLimitScopeResolver<RuntimeOrigin, RuntimeCall, Li
105105
impl pallet_rate_limiting::RateLimitUsageResolver<RuntimeOrigin, RuntimeCall, UsageKey>
106106
for TestUsageResolver
107107
{
108-
fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<UsageKey> {
108+
fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<Vec<UsageKey>> {
109109
match call {
110110
RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => {
111-
(*block_span).try_into().ok()
111+
(*block_span).try_into().ok().map(|key| vec![key])
112112
}
113-
RuntimeCall::RateLimiting(_) => Some(1),
113+
RuntimeCall::RateLimiting(_) => Some(vec![1]),
114114
_ => None,
115115
}
116116
}

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

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use frame_support::{
1313
},
1414
};
1515
use scale_info::TypeInfo;
16-
use sp_std::{marker::PhantomData, result::Result};
16+
use sp_std::{marker::PhantomData, result::Result, vec, vec::Vec};
1717

1818
use crate::{
1919
Config, LastSeen, Pallet,
@@ -94,12 +94,12 @@ where
9494
type Implicit = ();
9595
type Val = Option<(
9696
RateLimitTarget<<T as Config<I>>::GroupId>,
97-
Option<<T as Config<I>>::UsageKey>,
97+
Option<Vec<<T as Config<I>>::UsageKey>>,
9898
bool,
9999
)>;
100100
type Pre = Option<(
101101
RateLimitTarget<<T as Config<I>>::GroupId>,
102-
Option<<T as Config<I>>::UsageKey>,
102+
Option<Vec<<T as Config<I>>::UsageKey>>,
103103
bool,
104104
)>;
105105

@@ -155,7 +155,14 @@ where
155155
return Ok((ValidTransaction::default(), None, origin));
156156
}
157157

158-
let within_limit = Pallet::<T, I>::within_span(&usage_target, &usage, block_span);
158+
let usage_keys: Vec<Option<<T as Config<I>>::UsageKey>> = match usage.clone() {
159+
None => vec![None],
160+
Some(keys) => keys.into_iter().map(Some).collect(),
161+
};
162+
163+
let within_limit = usage_keys
164+
.iter()
165+
.all(|key| Pallet::<T, I>::within_span(&usage_target, key, block_span));
159166

160167
if !within_limit {
161168
return Err(TransactionValidityError::Invalid(
@@ -194,7 +201,18 @@ where
194201
return Ok(());
195202
}
196203
let block_number = frame_system::Pallet::<T>::block_number();
197-
LastSeen::<T, I>::insert(target, usage, block_number);
204+
match usage {
205+
None => LastSeen::<T, I>::insert(
206+
target,
207+
None::<<T as Config<I>>::UsageKey>,
208+
block_number,
209+
),
210+
Some(keys) => {
211+
for key in keys {
212+
LastSeen::<T, I>::insert(target, Some(key), block_number);
213+
}
214+
}
215+
}
198216
}
199217
}
200218
Ok(())
@@ -253,7 +271,7 @@ mod tests {
253271
) -> Result<
254272
(
255273
sp_runtime::transaction_validity::ValidTransaction,
256-
Option<(RateLimitTarget<GroupId>, Option<UsageKey>, bool)>,
274+
Option<(RateLimitTarget<GroupId>, Option<Vec<UsageKey>>, bool)>,
257275
RuntimeOrigin,
258276
),
259277
TransactionValidityError,

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen};
22
use frame_support::{pallet_prelude::DispatchError, traits::GetCallMetadata};
33
use scale_info::TypeInfo;
44
use serde::{Deserialize, Serialize};
5-
use sp_std::collections::btree_map::BTreeMap;
5+
use sp_std::{collections::btree_map::BTreeMap, vec::Vec};
66

77
/// Resolves the optional identifier within which a rate limit applies and can optionally adjust
88
/// enforcement behaviour.
@@ -53,9 +53,11 @@ impl BypassDecision {
5353

5454
/// Resolves the optional usage tracking key applied when enforcing limits.
5555
pub trait RateLimitUsageResolver<Origin, Call, Usage> {
56-
/// Returns `Some(usage)` when usage should be tracked per-key, or `None` for global usage
57-
/// tracking.
58-
fn context(origin: &Origin, call: &Call) -> Option<Usage>;
56+
/// Returns `Some(keys)` to track usage per key, or `None` for global usage tracking.
57+
///
58+
/// When multiple keys are returned, the rate limit is enforced against each key and all are
59+
/// recorded on success.
60+
fn context(origin: &Origin, call: &Call) -> Option<Vec<Usage>>;
5961
}
6062

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

‎runtime/src/rate_limiting/migration.rs‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ const GROUP_WEIGHTS_SUBNET: GroupId = 2;
4343
pub const GROUP_REGISTER_NETWORK: GroupId = 3;
4444
const GROUP_OWNER_HPARAMS: GroupId = 4;
4545
const GROUP_STAKING_OPS: GroupId = 5;
46+
const GROUP_SWAP_KEYS: GroupId = 6;
4647

4748
// `set_children` is rate-limited to once every 150 blocks, it's hard-coded in the legacy code.
4849
const SET_CHILDREN_RATE_LIMIT: u64 = 150;
@@ -144,9 +145,9 @@ fn commits() -> (Vec<GroupConfig>, Vec<Commit>, u64) {
144145
reads += build_register_network(&mut groups, &mut commits);
145146
reads += build_owner_hparams(&mut groups, &mut commits);
146147
reads += build_staking_ops(&mut groups, &mut commits);
148+
reads += build_swap_keys(&mut groups, &mut commits);
147149

148150
// standalone
149-
reads += build_swap_hotkey(&mut commits);
150151
reads += build_childkey_take(&mut commits);
151152
reads += build_set_children(&mut commits);
152153
reads += build_weights_version_key(&mut commits);
@@ -571,14 +572,23 @@ fn build_staking_ops(groups: &mut Vec<GroupConfig>, commits: &mut Vec<Commit>) -
571572
0
572573
}
573574

574-
// Standalone swap_hotkey.
575-
// usage: account
575+
// Swap hotkey/coldkey share the lock and usage; swap_coldkey bypasses enforcement but records
576+
// usage.
577+
// usage: account (coldkey)
576578
// legacy sources: TxRateLimit, LastRateLimitedBlock per LastTxBlock
577-
fn build_swap_hotkey(commits: &mut Vec<Commit>) -> u64 {
579+
fn build_swap_keys(groups: &mut Vec<GroupConfig>, commits: &mut Vec<Commit>) -> u64 {
578580
let mut reads: u64 = 0;
579-
let target =
580-
RateLimitTarget::Transaction(TransactionIdentifier::new(SUBTENSOR_PALLET_INDEX, 70));
581+
groups.push(GroupConfig {
582+
id: GROUP_SWAP_KEYS,
583+
name: b"swap-keys".to_vec(),
584+
sharing: GroupSharing::ConfigAndUsage,
585+
members: vec![
586+
MigratedCall::subtensor(70, false), // swap_hotkey
587+
MigratedCall::subtensor(71, false), // swap_coldkey
588+
],
589+
});
581590

591+
let target = RateLimitTarget::Group(GROUP_SWAP_KEYS);
582592
reads += 1;
583593
push_limit_commit_if_non_zero(commits, target, TxRateLimit::<Runtime>::get(), None);
584594

@@ -1055,8 +1065,7 @@ mod tests {
10551065
MIGRATION_NAME
10561066
));
10571067

1058-
let tx_target =
1059-
RateLimitTarget::Transaction(MigratedCall::subtensor(70, false).identifier());
1068+
let tx_target = RateLimitTarget::Group(GROUP_SWAP_KEYS);
10601069
let delegate_group = RateLimitTarget::Group(DELEGATE_TAKE_GROUP_ID);
10611070

10621071
assert_eq!(
@@ -1088,7 +1097,7 @@ mod tests {
10881097
),
10891098
Some(DELEGATE_TAKE_GROUP_ID)
10901099
);
1091-
assert_eq!(pallet_rate_limiting::NextGroupId::<Runtime>::get(), 6);
1100+
assert_eq!(pallet_rate_limiting::NextGroupId::<Runtime>::get(), 7);
10921101
});
10931102
}
10941103

‎runtime/src/rate_limiting/mod.rs‎

Lines changed: 54 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use pallet_rate_limiting::{RateLimitScopeResolver, RateLimitUsageResolver};
77
use pallet_subtensor::{Call as SubtensorCall, Tempo};
88
use scale_info::TypeInfo;
99
use serde::{Deserialize, Serialize};
10+
use sp_std::{vec, vec::Vec};
1011
use subtensor_runtime_common::{BlockNumber, MechId, NetUid};
1112

1213
use crate::{AccountId, Runtime, RuntimeCall, RuntimeOrigin};
@@ -107,11 +108,15 @@ impl RateLimitScopeResolver<RuntimeOrigin, RuntimeCall, NetUid, BlockNumber> for
107108
}
108109

109110
fn should_bypass(origin: &RuntimeOrigin, call: &RuntimeCall) -> BypassDecision {
110-
if matches!(origin.clone().into(), Ok(RawOrigin::Root)) {
111-
return BypassDecision::bypass_and_skip();
112-
}
113-
114111
if let RuntimeCall::SubtensorModule(inner) = call {
112+
if matches!(origin.clone().into(), Ok(RawOrigin::Root)) {
113+
// swap_coldkey should record last-seen but never fail; other root calls skip.
114+
if matches!(inner, SubtensorCall::swap_coldkey { .. }) {
115+
return BypassDecision::bypass_and_record();
116+
}
117+
return BypassDecision::bypass_and_skip();
118+
}
119+
115120
match inner {
116121
SubtensorCall::set_childkey_take {
117122
hotkey,
@@ -129,7 +134,8 @@ impl RateLimitScopeResolver<RuntimeOrigin, RuntimeCall, NetUid, BlockNumber> for
129134
}
130135
SubtensorCall::add_stake { .. }
131136
| SubtensorCall::add_stake_limit { .. }
132-
| SubtensorCall::decrease_take { .. } => {
137+
| SubtensorCall::decrease_take { .. }
138+
| SubtensorCall::swap_coldkey { .. } => {
133139
return BypassDecision::bypass_and_record();
134140
}
135141
SubtensorCall::reveal_weights { netuid, .. }
@@ -179,32 +185,48 @@ pub struct UsageResolver;
179185
impl RateLimitUsageResolver<RuntimeOrigin, RuntimeCall, RateLimitUsageKey<AccountId>>
180186
for UsageResolver
181187
{
182-
fn context(origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<RateLimitUsageKey<AccountId>> {
188+
fn context(
189+
origin: &RuntimeOrigin,
190+
call: &RuntimeCall,
191+
) -> Option<Vec<RateLimitUsageKey<AccountId>>> {
183192
match call {
184193
RuntimeCall::SubtensorModule(inner) => match inner {
185-
SubtensorCall::swap_hotkey { .. } => {
186-
signed_origin(origin).map(RateLimitUsageKey::<AccountId>::Account)
194+
SubtensorCall::swap_coldkey { new_coldkey, .. } => {
195+
Some(vec![RateLimitUsageKey::<AccountId>::Account(
196+
new_coldkey.clone(),
197+
)])
198+
}
199+
SubtensorCall::swap_hotkey { new_hotkey, .. } => {
200+
// Record against the coldkey (enforcement) and the new hotkey to mirror legacy
201+
// writes.
202+
let coldkey = signed_origin(origin)?;
203+
Some(vec![
204+
RateLimitUsageKey::<AccountId>::Account(coldkey),
205+
RateLimitUsageKey::<AccountId>::Account(new_hotkey.clone()),
206+
])
187207
}
188208
SubtensorCall::increase_take { hotkey, .. } => {
189-
Some(RateLimitUsageKey::<AccountId>::Account(hotkey.clone()))
209+
Some(vec![RateLimitUsageKey::<AccountId>::Account(
210+
hotkey.clone(),
211+
)])
190212
}
191213
SubtensorCall::set_childkey_take { hotkey, netuid, .. }
192214
| SubtensorCall::set_children { hotkey, netuid, .. } => {
193-
Some(RateLimitUsageKey::<AccountId>::AccountSubnet {
215+
Some(vec![RateLimitUsageKey::<AccountId>::AccountSubnet {
194216
account: hotkey.clone(),
195217
netuid: *netuid,
196-
})
218+
}])
197219
}
198220
SubtensorCall::set_weights { netuid, .. }
199221
| SubtensorCall::commit_weights { netuid, .. }
200222
| SubtensorCall::reveal_weights { netuid, .. }
201223
| SubtensorCall::batch_reveal_weights { netuid, .. }
202224
| SubtensorCall::commit_timelocked_weights { netuid, .. } => {
203225
let (_, uid) = neuron_identity(origin, *netuid)?;
204-
Some(RateLimitUsageKey::<AccountId>::SubnetNeuron {
226+
Some(vec![RateLimitUsageKey::<AccountId>::SubnetNeuron {
205227
netuid: *netuid,
206228
uid,
207-
})
229+
}])
208230
}
209231
// legacy implementation still used netuid only, but it was recalculating it using
210232
// mecid, so switching to netuid AND mecid is logical here
@@ -214,39 +236,41 @@ impl RateLimitUsageResolver<RuntimeOrigin, RuntimeCall, RateLimitUsageKey<Accoun
214236
| SubtensorCall::commit_crv3_mechanism_weights { netuid, mecid, .. }
215237
| SubtensorCall::commit_timelocked_mechanism_weights { netuid, mecid, .. } => {
216238
let (_, uid) = neuron_identity(origin, *netuid)?;
217-
Some(RateLimitUsageKey::<AccountId>::SubnetMechanismNeuron {
218-
netuid: *netuid,
219-
mecid: *mecid,
220-
uid,
221-
})
239+
Some(vec![
240+
RateLimitUsageKey::<AccountId>::SubnetMechanismNeuron {
241+
netuid: *netuid,
242+
mecid: *mecid,
243+
uid,
244+
},
245+
])
222246
}
223247
SubtensorCall::serve_axon { netuid, .. }
224248
| SubtensorCall::serve_axon_tls { netuid, .. } => {
225249
let hotkey = signed_origin(origin)?;
226-
Some(RateLimitUsageKey::<AccountId>::AccountSubnetServing {
250+
Some(vec![RateLimitUsageKey::<AccountId>::AccountSubnetServing {
227251
account: hotkey,
228252
netuid: *netuid,
229253
endpoint: ServingEndpoint::Axon,
230-
})
254+
}])
231255
}
232256
SubtensorCall::serve_prometheus { netuid, .. } => {
233257
let hotkey = signed_origin(origin)?;
234-
Some(RateLimitUsageKey::<AccountId>::AccountSubnetServing {
258+
Some(vec![RateLimitUsageKey::<AccountId>::AccountSubnetServing {
235259
account: hotkey,
236260
netuid: *netuid,
237261
endpoint: ServingEndpoint::Prometheus,
238-
})
262+
}])
239263
}
240264
SubtensorCall::associate_evm_key { netuid, .. } => {
241265
let hotkey = signed_origin(origin)?;
242266
let uid = pallet_subtensor::Pallet::<Runtime>::get_uid_for_net_and_hotkey(
243267
*netuid, &hotkey,
244268
)
245269
.ok()?;
246-
Some(RateLimitUsageKey::<AccountId>::SubnetNeuron {
270+
Some(vec![RateLimitUsageKey::<AccountId>::SubnetNeuron {
247271
netuid: *netuid,
248272
uid,
249-
})
273+
}])
250274
}
251275
// Staking calls share a group lock; only add_* write usage, the rest are read-only.
252276
// Keep the usage key granular so the lock applies per (coldkey, hotkey, netuid).
@@ -276,31 +300,31 @@ impl RateLimitUsageResolver<RuntimeOrigin, RuntimeCall, RateLimitUsageKey<Accoun
276300
..
277301
} => {
278302
let coldkey = signed_origin(origin)?;
279-
Some(RateLimitUsageKey::<AccountId>::ColdkeyHotkeySubnet {
303+
Some(vec![RateLimitUsageKey::<AccountId>::ColdkeyHotkeySubnet {
280304
coldkey,
281305
hotkey: hotkey.clone(),
282306
netuid: *netuid,
283-
})
307+
}])
284308
}
285309
_ => None,
286310
},
287311
RuntimeCall::AdminUtils(inner) => {
288312
if let Some(netuid) = owner_hparam_netuid(inner) {
289-
Some(RateLimitUsageKey::<AccountId>::Subnet(netuid))
313+
Some(vec![RateLimitUsageKey::<AccountId>::Subnet(netuid)])
290314
} else {
291315
match inner {
292316
AdminUtilsCall::sudo_set_sn_owner_hotkey { netuid, .. } => {
293-
Some(RateLimitUsageKey::<AccountId>::Subnet(*netuid))
317+
Some(vec![RateLimitUsageKey::<AccountId>::Subnet(*netuid)])
294318
}
295319
AdminUtilsCall::sudo_set_weights_version_key { netuid, .. }
296320
| AdminUtilsCall::sudo_set_mechanism_count { netuid, .. }
297321
| AdminUtilsCall::sudo_set_mechanism_emission_split { netuid, .. }
298322
| AdminUtilsCall::sudo_trim_to_max_allowed_uids { netuid, .. } => {
299323
let who = signed_origin(origin)?;
300-
Some(RateLimitUsageKey::<AccountId>::AccountSubnet {
324+
Some(vec![RateLimitUsageKey::<AccountId>::AccountSubnet {
301325
account: who,
302326
netuid: *netuid,
303-
})
327+
}])
304328
}
305329
_ => None,
306330
}

0 commit comments

Comments
 (0)