code/pallets/subtensor/src/guards/check_coldkey_swap.rs
use super::{CallOf, DispatchableOriginOf};
use crate::weights::WeightInfo;
use crate::{Call, ColdkeySwapAnnouncements, ColdkeySwapDisputes, Config, Error};
use frame_support::{
dispatch::{DispatchErrorWithPostInfo, DispatchExtension, DispatchInfo, PostDispatchInfo},
pallet_prelude::*,
traits::{IsSubType, OriginTrait},
};
use sp_runtime::traits::Dispatchable;
use sp_std::marker::PhantomData;
/// Dispatch extension that blocks most calls when a coldkey swap is active.
///
/// When a coldkey swap has been announced for the signing account:
/// * If the swap is disputed, ALL calls are blocked.
/// * Otherwise, only swap-related calls and MEV-protected calls (`submit_encrypted`)
/// are allowed through.
///
/// Root origin bypasses this extension entirely.
/// Non-signed origins pass through.
///
/// Because this is a `DispatchExtension` (not a `TransactionExtension`), it fires at every
/// `call.dispatch(origin)` site — including inside the proxy pallet's `do_proxy()`.
/// This means nested proxies of any depth are handled automatically with the real
/// resolved origin.
pub struct CheckColdkeySwap<T: Config>(PhantomData<T>);
impl<T> CheckColdkeySwap<T>
where
T: Config + pallet_shield::Config,
CallOf<T>: IsSubType<Call<T>> + IsSubType<pallet_shield::Call<T>>,
{
pub fn check(who: &T::AccountId, call: &CallOf<T>) -> Result<(), Error<T>> {
if !ColdkeySwapAnnouncements::<T>::contains_key(who) {
return Ok(());
}
if ColdkeySwapDisputes::<T>::contains_key(who) {
return Err(Error::<T>::ColdkeySwapDisputed);
}
if Self::is_allowed_during_swap(call) {
Ok(())
} else {
Err(Error::<T>::ColdkeySwapAnnounced)
}
}
fn is_allowed_during_swap(call: &CallOf<T>) -> bool {
matches!(
call.is_sub_type(),
Some(
Call::announce_coldkey_swap { .. }
| Call::swap_coldkey_announced { .. }
| Call::dispute_coldkey_swap { .. }
| Call::clear_coldkey_swap_announcement { .. }
)
) || matches!(
IsSubType::<pallet_shield::Call<T>>::is_sub_type(call),
Some(pallet_shield::Call::submit_encrypted { .. })
)
}
}
impl<T> DispatchExtension<<T as frame_system::Config>::RuntimeCall> for CheckColdkeySwap<T>
where
T: Config + pallet_shield::Config,
<T as frame_system::Config>::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>
+ IsSubType<Call<T>>
+ IsSubType<pallet_shield::Call<T>>,
DispatchableOriginOf<T>: OriginTrait<AccountId = T::AccountId>,
{
type Pre = ();
fn weight(_call: &CallOf<T>) -> Weight {
<T as Config>::WeightInfo::check_coldkey_swap_extension()
}
fn pre_dispatch(
origin: &DispatchableOriginOf<T>,
call: &CallOf<T>,
) -> Result<Self::Pre, DispatchErrorWithPostInfo> {
// Only care about signed origins.
// Root is already bypassed by the extension before we get here.
let Some(who) = origin.as_signer() else {
return Ok(());
};
Self::check(who, call).map_err(Into::into)
}
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
use super::CheckColdkeySwap;
use crate::{
ColdkeySwapAnnouncements, ColdkeySwapDisputes, Error, tests::mock::*,
weights::WeightInfo as _,
};
use frame_support::{
BoundedVec, assert_ok,
dispatch::{DispatchExtension, DispatchResultWithPostInfo},
traits::ExtendedDispatchable,
};
use frame_system::Call as SystemCall;
use pallet_subtensor_proxy::Call as ProxyCall;
use sp_core::U256;
use sp_runtime::traits::Hash;
use subtensor_runtime_common::{ProxyType, TaoBalance};
type HashingOf<T> = <T as frame_system::Config>::Hashing;
/// Calls that should be blocked when a coldkey swap is active.
fn forbidden_calls() -> Vec<RuntimeCall> {
vec![
RuntimeCall::System(SystemCall::remark { remark: vec![] }),
RuntimeCall::SubtensorModule(crate::Call::add_stake {
hotkey: U256::from(1),
netuid: 1u16.into(),
amount_staked: 1_000u64.into(),
}),
RuntimeCall::SubtensorModule(crate::Call::remove_stake {
hotkey: U256::from(1),
netuid: 1u16.into(),
amount_unstaked: 1_000u64.into(),
}),
RuntimeCall::SubtensorModule(crate::Call::set_weights {
netuid: 1u16.into(),
dests: vec![],
weights: vec![],
version_key: 0,
}),
RuntimeCall::SubtensorModule(crate::Call::register_network {
hotkey: U256::from(1),
}),
]
}
/// Calls that should be allowed through the extension during an active (undisputed) swap.
fn authorized_calls() -> Vec<RuntimeCall> {
vec![
RuntimeCall::SubtensorModule(crate::Call::announce_coldkey_swap {
new_coldkey_hash: HashingOf::<Test>::hash_of(&U256::from(99)),
}),
RuntimeCall::SubtensorModule(crate::Call::swap_coldkey_announced {
new_coldkey: U256::from(42),
}),
RuntimeCall::SubtensorModule(crate::Call::dispute_coldkey_swap {}),
RuntimeCall::Shield(pallet_shield::Call::submit_encrypted {
ciphertext: BoundedVec::truncate_from(vec![1, 2, 3, 4]),
}),
]
}
fn setup_swap_announced(who: &U256) {
let now = System::block_number();
let hash = HashingOf::<Test>::hash_of(&U256::from(42));
ColdkeySwapAnnouncements::<Test>::insert(who, (now, hash));
}
fn setup_swap_disputed(who: &U256) {
setup_swap_announced(who);
ColdkeySwapDisputes::<Test>::insert(who, System::block_number());
}
fn remark_call() -> RuntimeCall {
RuntimeCall::System(SystemCall::remark { remark: vec![] })
}
fn add_balance_to_coldkey_account(coldkey: &U256, tao: TaoBalance) {
let credit = SubtensorModule::mint_tao(tao);
let _ = SubtensorModule::spend_tao(coldkey, credit, tao).unwrap();
}
fn dispatch_with_ext(call: RuntimeCall, origin: RuntimeOrigin) -> DispatchResultWithPostInfo {
<CheckColdkeySwap<Test> as ExtendedDispatchable<RuntimeCall>>::dispatch_with_extension(
origin, call,
)
}
#[test]
fn weight_charges_all_calls_because_swap_state_can_block_any_signed_call() {
let expected = <Test as crate::Config>::WeightInfo::check_coldkey_swap_extension();
for call in forbidden_calls().into_iter().chain(authorized_calls()) {
assert_eq!(
<CheckColdkeySwap<Test> as DispatchExtension<RuntimeCall>>::weight(&call),
expected
);
}
}
#[test]
fn no_active_swap_allows_calls() {
new_test_ext(1).execute_with(|| {
let who = U256::from(1);
assert_ok!(dispatch_with_ext(remark_call(), RuntimeOrigin::signed(who)));
});
}
#[test]
fn none_bypasses_extension() {
new_test_ext(1).execute_with(|| {
let who = U256::from(1);
setup_swap_disputed(&who);
assert_ok!(dispatch_with_ext(remark_call(), RuntimeOrigin::none()));
});
}
#[test]
fn root_bypasses_extension() {
new_test_ext(1).execute_with(|| {
let who = U256::from(1);
setup_swap_disputed(&who);
assert_ok!(dispatch_with_ext(remark_call(), RuntimeOrigin::root()));
});
}
#[test]
fn active_swap_blocks_forbidden_calls() {
new_test_ext(1).execute_with(|| {
let who = U256::from(1);
setup_swap_announced(&who);
for call in forbidden_calls() {
assert_eq!(
dispatch_with_ext(call, RuntimeOrigin::signed(who))
.unwrap_err()
.error,
Error::<Test>::ColdkeySwapAnnounced.into()
);
}
});
}
#[test]
fn active_swap_allows_authorized_calls() {
new_test_ext(1).execute_with(|| {
let who = U256::from(1);
setup_swap_announced(&who);
for call in authorized_calls() {
if let Err(err) = dispatch_with_ext(call, RuntimeOrigin::signed(who)) {
assert_ne!(
err.error,
Error::<Test>::ColdkeySwapAnnounced.into(),
"Authorized call should not be blocked by the extension"
);
}
}
});
}
#[test]
fn disputed_swap_blocks_all_calls() {
new_test_ext(1).execute_with(|| {
let who = U256::from(1);
setup_swap_disputed(&who);
// Both forbidden and authorized calls should be blocked during dispute
let all_calls = forbidden_calls()
.into_iter()
.chain(authorized_calls())
.collect::<Vec<_>>();
for call in all_calls {
assert_eq!(
dispatch_with_ext(call, RuntimeOrigin::signed(who))
.unwrap_err()
.error,
Error::<Test>::ColdkeySwapDisputed.into()
);
}
});
}
#[test]
fn proxied_forbidden_call_blocked() {
new_test_ext(1).execute_with(|| {
let real = U256::from(1);
let delegate = U256::from(2);
let now = System::block_number();
let hash = HashingOf::<Test>::hash_of(&U256::from(42));
ColdkeySwapAnnouncements::<Test>::insert(real, (now, hash));
// Give delegate enough balance for proxy deposit
add_balance_to_coldkey_account(&real, 1_000_000_000.into());
add_balance_to_coldkey_account(&delegate, 1_000_000_000.into());
// Register proxy: delegate can act on behalf of real
assert_ok!(Proxy::add_proxy(
RuntimeOrigin::signed(real),
delegate,
ProxyType::Any,
0
));
// Dispatch a proxy call as delegate
let proxy_call = RuntimeCall::Proxy(ProxyCall::proxy {
real,
force_proxy_type: None,
call: Box::new(remark_call()),
});
// The outer proxy call itself succeeds
assert_ok!(dispatch_with_ext(
proxy_call,
RuntimeOrigin::signed(delegate)
));
// The inner call was blocked — check via LastCallResult storage.
assert_eq!(
pallet_subtensor_proxy::LastCallResult::<Test>::get(real),
Some(Err(Error::<Test>::ColdkeySwapAnnounced.into()))
);
});
}
#[test]
fn nested_proxy_blocked() {
new_test_ext(1).execute_with(|| {
let real = U256::from(1);
let delegate1 = U256::from(2);
let delegate2 = U256::from(3);
let now = System::block_number();
let hash = HashingOf::<Test>::hash_of(&U256::from(42));
ColdkeySwapAnnouncements::<Test>::insert(real, (now, hash));
add_balance_to_coldkey_account(&real, 1_000_000_000.into());
add_balance_to_coldkey_account(&delegate1, 1_000_000_000.into());
add_balance_to_coldkey_account(&delegate2, 1_000_000_000.into());
// delegate1 can proxy for real, delegate2 can proxy for delegate1
assert_ok!(Proxy::add_proxy(
RuntimeOrigin::signed(real),
delegate1,
ProxyType::Any,
0
));
assert_ok!(Proxy::add_proxy(
RuntimeOrigin::signed(delegate1),
delegate2,
ProxyType::Any,
0
));
// Nested: delegate2 -> delegate1 -> proxy(real, remark)
let inner_proxy = RuntimeCall::Proxy(ProxyCall::proxy {
real,
force_proxy_type: None,
call: Box::new(remark_call()),
});
let outer_proxy = RuntimeCall::Proxy(ProxyCall::proxy {
real: delegate1,
force_proxy_type: None,
call: Box::new(inner_proxy),
});
assert_ok!(dispatch_with_ext(
outer_proxy,
RuntimeOrigin::signed(delegate2)
));
// The innermost call (remark as real) was blocked.
assert_eq!(
pallet_subtensor_proxy::LastCallResult::<Test>::get(real),
Some(Err(Error::<Test>::ColdkeySwapAnnounced.into()))
);
});
}
}