From 55a24f06c073e0c382a9d86e738105148c95d4f2 Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Tue, 2 Jun 2026 17:21:30 -0300 Subject: [PATCH 1/8] Allow transfer of lock to different coldkey/hotkey --- pallets/subtensor/src/macros/dispatches.rs | 2 +- pallets/subtensor/src/staking/lock.rs | 159 +++++++++++++++++---- pallets/subtensor/src/tests/locks.rs | 132 +++++++++++++++++ 3 files changed, 261 insertions(+), 32 deletions(-) diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index b471328aec..2a73a0e473 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -2575,7 +2575,7 @@ mod dispatches { netuid: NetUid, ) -> DispatchResult { let coldkey = ensure_signed(origin)?; - Self::do_move_lock(&coldkey, &destination_hotkey, netuid) + Self::do_move_lock(&coldkey, &destination_hotkey, netuid, false) } /// Sets or clears the caller's perpetual lock flag for a subnet. diff --git a/pallets/subtensor/src/staking/lock.rs b/pallets/subtensor/src/staking/lock.rs index 27c5e5d646..3cdb16687c 100644 --- a/pallets/subtensor/src/staking/lock.rs +++ b/pallets/subtensor/src/staking/lock.rs @@ -1119,9 +1119,6 @@ impl Pallet { return; } let old_owner_hotkey = SubnetOwnerHotkey::::get(netuid); - let unlock_rate = UnlockRate::::get(); - let maturity_rate = MaturityRate::::get(); - // Register new owner as a neuron if not yet registered. if Self::get_uid_for_net_and_hotkey(netuid, &king_hotkey).is_err() && Self::register_neuron(netuid, &king_hotkey).is_err() @@ -1130,6 +1127,36 @@ impl Pallet { } // Move aggregate buckets using the hotkey's new role. + Self::reassign_subnet_owner_lock_aggregates(netuid, &old_owner_hotkey, &king_hotkey); + + // Reassign subnet owner coldkey and owner hotkey. + SubnetOwner::::insert(netuid, new_owner_coldkey.clone()); + SubnetOwnerHotkey::::insert(netuid, king_hotkey.clone()); + Self::deposit_event(Event::SubnetOwnerChanged { + netuid, + old_coldkey: current_owner_coldkey, + new_coldkey: new_owner_coldkey, + }); + } + + /// Moves aggregate lock buckets when a subnet owner hotkey changes. + /// + /// Individual lock rows keep their hotkey. The previous owner hotkey's owner + /// aggregate becomes a regular hotkey aggregate, and the new owner hotkey's + /// regular aggregate becomes the owner aggregate. + pub fn reassign_subnet_owner_lock_aggregates( + netuid: NetUid, + old_owner_hotkey: &T::AccountId, + new_owner_hotkey: &T::AccountId, + ) { + if old_owner_hotkey == new_owner_hotkey { + return; + } + + let now = Self::get_current_block_as_u64(); + let unlock_rate = UnlockRate::::get(); + let maturity_rate = MaturityRate::::get(); + if let Some(owner_lock) = OwnerLock::::take(netuid) { let moved_owner_lock = ConvictionModel::roll_forward_lock( owner_lock, @@ -1139,7 +1166,7 @@ impl Pallet { true, true, ); - let current = HotkeyLock::::get(netuid, &old_owner_hotkey) + let current = HotkeyLock::::get(netuid, old_owner_hotkey) .map(|lock| { ConvictionModel::roll_forward_lock( lock, @@ -1153,7 +1180,7 @@ impl Pallet { .unwrap_or_else(|| Self::empty_lock(now)); Self::insert_hotkey_lock_state( netuid, - &old_owner_hotkey, + old_owner_hotkey, LockState { locked_mass: current .locked_mass @@ -1165,6 +1192,7 @@ impl Pallet { }, ); } + if let Some(owner_lock) = DecayingOwnerLock::::take(netuid) { let moved_owner_lock = ConvictionModel::roll_forward_lock( owner_lock, @@ -1174,7 +1202,7 @@ impl Pallet { true, false, ); - let current = DecayingHotkeyLock::::get(netuid, &old_owner_hotkey) + let current = DecayingHotkeyLock::::get(netuid, old_owner_hotkey) .map(|lock| { ConvictionModel::roll_forward_lock( lock, @@ -1188,7 +1216,7 @@ impl Pallet { .unwrap_or_else(|| Self::empty_lock(now)); Self::insert_decaying_hotkey_lock_state( netuid, - &old_owner_hotkey, + old_owner_hotkey, LockState { locked_mass: current .locked_mass @@ -1200,9 +1228,10 @@ impl Pallet { }, ); } - if let Some(king_lock) = HotkeyLock::::take(netuid, &king_hotkey) { - let moved_king_lock = ConvictionModel::roll_forward_lock( - king_lock, + + if let Some(new_owner_lock) = HotkeyLock::::take(netuid, new_owner_hotkey) { + let moved_new_owner_lock = ConvictionModel::roll_forward_lock( + new_owner_lock, now, unlock_rate, maturity_rate, @@ -1227,10 +1256,10 @@ impl Pallet { LockState { locked_mass: current .locked_mass - .saturating_add(moved_king_lock.locked_mass), + .saturating_add(moved_new_owner_lock.locked_mass), conviction: current .conviction - .saturating_add(moved_king_lock.conviction), + .saturating_add(moved_new_owner_lock.conviction), last_update: now, }, now, @@ -1241,9 +1270,10 @@ impl Pallet { ), ); } - if let Some(king_lock) = DecayingHotkeyLock::::take(netuid, &king_hotkey) { - let moved_king_lock = ConvictionModel::roll_forward_lock( - king_lock, + + if let Some(new_owner_lock) = DecayingHotkeyLock::::take(netuid, new_owner_hotkey) { + let moved_new_owner_lock = ConvictionModel::roll_forward_lock( + new_owner_lock, now, unlock_rate, maturity_rate, @@ -1268,10 +1298,10 @@ impl Pallet { LockState { locked_mass: current .locked_mass - .saturating_add(moved_king_lock.locked_mass), + .saturating_add(moved_new_owner_lock.locked_mass), conviction: current .conviction - .saturating_add(moved_king_lock.conviction), + .saturating_add(moved_new_owner_lock.conviction), last_update: now, }, now, @@ -1282,15 +1312,6 @@ impl Pallet { ), ); } - - // Reassign subnet owner coldkey and owner hotkey. - SubnetOwner::::insert(netuid, new_owner_coldkey.clone()); - SubnetOwnerHotkey::::insert(netuid, king_hotkey.clone()); - Self::deposit_event(Event::SubnetOwnerChanged { - netuid, - old_coldkey: current_owner_coldkey, - new_coldkey: new_owner_coldkey, - }); } /// Ensure the coldkey does not have an active lock on any subnets. @@ -1570,18 +1591,21 @@ impl Pallet { (reads, writes) } - /// Moves lock from one hotkey to another and clears conviction + /// Moves lock from one hotkey to another. /// /// The lock is rolled forward to the current block before switching the /// associated hotkey so that the lock stays mathematically correct and /// preserves current decayed locked mass. /// - /// The conviction is reset to zero if the destination and source hotkeys - /// are owned by different coldkeys, otherwise it is preserved. + /// Conviction is cleared when the source and destination hotkeys are owned by + /// different coldkeys, unless `preserve_conviction` is set. Lease termination + /// preserves conviction because the subnet ownership and remaining stake are + /// moving to the beneficiary-controlled hotkey as one handoff. pub fn do_move_lock( coldkey: &T::AccountId, destination_hotkey: &T::AccountId, netuid: NetUid, + preserve_conviction: bool, ) -> DispatchResult { ensure!( Self::hotkey_account_exists(destination_hotkey), @@ -1597,8 +1621,9 @@ impl Pallet { let mut lock = model.individual_lock().clone(); let removed = lock.clone(); - if Self::get_owning_coldkey_for_hotkey(&origin_hotkey) - != Self::get_owning_coldkey_for_hotkey(destination_hotkey) + if !preserve_conviction + && Self::get_owning_coldkey_for_hotkey(&origin_hotkey) + != Self::get_owning_coldkey_for_hotkey(destination_hotkey) { lock.conviction = U64F64::saturating_from_num(0); } @@ -1811,6 +1836,78 @@ impl Pallet { Ok(()) } + pub fn transfer_full_lock_to_coldkey( + origin_coldkey: &T::AccountId, + destination_coldkey: &T::AccountId, + netuid: NetUid, + destination_hotkey: &T::AccountId, + ) -> DispatchResult { + let now = Self::get_current_block_as_u64(); + let unlock_rate = UnlockRate::::get(); + let maturity_rate = MaturityRate::::get(); + + let Some((source_hotkey, mut source_model)) = + Self::read_conviction_model(origin_coldkey, netuid, now) + else { + return Ok(()); + }; + + source_model.roll_forward_individual(now, unlock_rate, maturity_rate); + let moved_lock = source_model.individual_lock().clone(); + if moved_lock.locked_mass.is_zero() + && moved_lock.conviction == U64F64::saturating_from_num(0) + { + Lock::::remove((origin_coldkey.clone(), netuid, source_hotkey)); + return Ok(()); + } + + let destination_lock = match Self::read_conviction_model(destination_coldkey, netuid, now) { + Some((existing_hotkey, mut destination_model)) => { + ensure!( + existing_hotkey == *destination_hotkey, + Error::::LockHotkeyMismatch + ); + destination_model.roll_forward_individual(now, unlock_rate, maturity_rate); + destination_model.individual_lock().clone() + } + None => Self::empty_lock(now), + }; + + let moved_lock_for_destination = ConvictionModel::roll_forward_lock( + moved_lock.clone(), + now, + unlock_rate, + maturity_rate, + Self::is_subnet_owner_hotkey(netuid, destination_hotkey), + Self::is_perpetual_lock(destination_coldkey, netuid), + ); + let destination_lock = + ConvictionModel::merge_lock(&destination_lock, &moved_lock_for_destination); + + Lock::::remove((origin_coldkey.clone(), netuid, source_hotkey.clone())); + Self::reduce_aggregate_lock( + origin_coldkey, + &source_hotkey, + netuid, + moved_lock.locked_mass, + moved_lock.conviction, + ); + Self::insert_lock_state( + destination_coldkey, + netuid, + destination_hotkey, + destination_lock, + ); + Self::add_aggregate_lock( + destination_coldkey, + destination_hotkey, + netuid, + moved_lock_for_destination, + ); + + Ok(()) + } + /// Destroys all lock maps for network dissolution pub fn destroy_lock_maps(netuid: NetUid) { // Lock: (coldkey, netuid, hotkey) diff --git a/pallets/subtensor/src/tests/locks.rs b/pallets/subtensor/src/tests/locks.rs index 4b452d639f..1c50f9ddad 100644 --- a/pallets/subtensor/src/tests/locks.rs +++ b/pallets/subtensor/src/tests/locks.rs @@ -2368,6 +2368,138 @@ fn test_change_subnet_owner_rebuilds_old_owner_hotkey_by_lock_mode() { }); } +#[test] +fn test_reassign_subnet_owner_lock_aggregates_moves_and_merges_all_buckets() { + new_test_ext(1).execute_with(|| { + let old_owner_coldkey = U256::from(1); + let old_owner_hotkey = U256::from(2); + let new_owner_hotkey = U256::from(3); + let netuid = setup_subnet_with_stake(old_owner_coldkey, old_owner_hotkey, 100_000_000_000); + let now = SubtensorModule::get_current_block_as_u64(); + + OwnerLock::::insert( + netuid, + LockState { + locked_mass: 100u64.into(), + conviction: U64F64::from_num(100), + last_update: now, + }, + ); + DecayingOwnerLock::::insert( + netuid, + LockState { + locked_mass: 200u64.into(), + conviction: U64F64::from_num(200), + last_update: now, + }, + ); + HotkeyLock::::insert( + netuid, + old_owner_hotkey, + LockState { + locked_mass: 10u64.into(), + conviction: U64F64::from_num(10), + last_update: now, + }, + ); + DecayingHotkeyLock::::insert( + netuid, + old_owner_hotkey, + LockState { + locked_mass: 20u64.into(), + conviction: U64F64::from_num(20), + last_update: now, + }, + ); + HotkeyLock::::insert( + netuid, + new_owner_hotkey, + LockState { + locked_mass: 300u64.into(), + conviction: U64F64::from_num(300), + last_update: now, + }, + ); + DecayingHotkeyLock::::insert( + netuid, + new_owner_hotkey, + LockState { + locked_mass: 400u64.into(), + conviction: U64F64::from_num(400), + last_update: now, + }, + ); + + SubtensorModule::reassign_subnet_owner_lock_aggregates( + netuid, + &old_owner_hotkey, + &new_owner_hotkey, + ); + + assert_eq!( + HotkeyLock::::get(netuid, old_owner_hotkey) + .unwrap() + .locked_mass, + 110u64.into() + ); + assert_eq!( + DecayingHotkeyLock::::get(netuid, old_owner_hotkey) + .unwrap() + .locked_mass, + 220u64.into() + ); + assert!(HotkeyLock::::get(netuid, new_owner_hotkey).is_none()); + assert!(DecayingHotkeyLock::::get(netuid, new_owner_hotkey).is_none()); + assert_eq!( + OwnerLock::::get(netuid).unwrap().locked_mass, + 300u64.into() + ); + assert_eq!( + DecayingOwnerLock::::get(netuid).unwrap().locked_mass, + 400u64.into() + ); + }); +} + +#[test] +fn test_reassign_subnet_owner_lock_aggregates_noops_for_same_hotkey() { + new_test_ext(1).execute_with(|| { + let coldkey = U256::from(1); + let hotkey = U256::from(2); + let netuid = setup_subnet_with_stake(coldkey, hotkey, 100_000_000_000); + let now = SubtensorModule::get_current_block_as_u64(); + + OwnerLock::::insert( + netuid, + LockState { + locked_mass: 100u64.into(), + conviction: U64F64::from_num(100), + last_update: now, + }, + ); + HotkeyLock::::insert( + netuid, + hotkey, + LockState { + locked_mass: 50u64.into(), + conviction: U64F64::from_num(50), + last_update: now, + }, + ); + + SubtensorModule::reassign_subnet_owner_lock_aggregates(netuid, &hotkey, &hotkey); + + assert_eq!( + OwnerLock::::get(netuid).unwrap().locked_mass, + 100u64.into() + ); + assert_eq!( + HotkeyLock::::get(netuid, hotkey).unwrap().locked_mass, + 50u64.into() + ); + }); +} + #[test] fn test_swap_hotkey_locks_moves_owner_hotkey_aggregate_to_owner_lock() { new_test_ext(1).execute_with(|| { From c375dfd9d4a99abbe9a212f3b51a87ba4917028e Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Tue, 2 Jun 2026 17:22:13 -0300 Subject: [PATCH 2/8] Allow to set enforce_min_stake to false to avoid dust when chain executed move_stake --- pallets/subtensor/src/staking/move_stake.rs | 1 + pallets/subtensor/src/staking/stake_utils.rs | 12 +++++++----- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/pallets/subtensor/src/staking/move_stake.rs b/pallets/subtensor/src/staking/move_stake.rs index aafefa28ed..961ab697a3 100644 --- a/pallets/subtensor/src/staking/move_stake.rs +++ b/pallets/subtensor/src/staking/move_stake.rs @@ -399,6 +399,7 @@ impl Pallet { destination_hotkey, origin_netuid, move_amount, + true, ) } } diff --git a/pallets/subtensor/src/staking/stake_utils.rs b/pallets/subtensor/src/staking/stake_utils.rs index 0f6a553c91..a31e1d24f4 100644 --- a/pallets/subtensor/src/staking/stake_utils.rs +++ b/pallets/subtensor/src/staking/stake_utils.rs @@ -959,6 +959,7 @@ impl Pallet { destination_hotkey: &T::AccountId, netuid: NetUid, alpha: AlphaBalance, + enforce_min_stake: bool, ) -> Result { // Transfer lock (may fail if destination coldkey has a conflicting lock) Self::transfer_lock(origin_coldkey, destination_coldkey, netuid, alpha)?; @@ -1002,11 +1003,12 @@ impl Pallet { .saturating_to_num::() .into(); - // Ensure tao_equivalent is above DefaultMinStake - ensure!( - tao_equivalent >= DefaultMinStake::::get(), - Error::::AmountTooLow - ); + if enforce_min_stake { + ensure!( + tao_equivalent >= DefaultMinStake::::get(), + Error::::AmountTooLow + ); + } // Step 3: Update StakingHotkeys if the hotkey's total alpha, across all subnets, is zero // TODO: fix. From 9a53fb6cf0c01bff1cc439495c451b9dfe128b26 Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Tue, 2 Jun 2026 17:23:01 -0300 Subject: [PATCH 3/8] Fix lease not transferring alpha/locks to beneficiary --- pallets/subtensor/src/subnets/leasing.rs | 95 ++++- pallets/subtensor/src/tests/leasing.rs | 430 +++++++++++++++++++++++ 2 files changed, 517 insertions(+), 8 deletions(-) diff --git a/pallets/subtensor/src/subnets/leasing.rs b/pallets/subtensor/src/subnets/leasing.rs index d3e1b84df5..a6a8b2b734 100644 --- a/pallets/subtensor/src/subnets/leasing.rs +++ b/pallets/subtensor/src/subnets/leasing.rs @@ -218,8 +218,39 @@ impl Pallet { Self::coldkey_owns_hotkey(&lease.beneficiary, &hotkey), Error::::BeneficiaryDoesNotOwnHotkey ); + + // Move any lease coldkey lock to the beneficiary-controlled hotkey without + // assigning ownership of the generated lease hotkey to the beneficiary. + Self::move_lease_lock_to_beneficiary_hotkey(&lease, &hotkey)?; + + // Transfer ownership to the beneficiary SubnetOwner::::insert(lease.netuid, lease.beneficiary.clone()); Self::set_subnet_owner_hotkey(lease.netuid, &hotkey)?; + Self::reassign_subnet_owner_lock_aggregates(lease.netuid, &lease.hotkey, &hotkey); + + Self::repatriate_lease_coldkey_alpha(&lease, &hotkey)?; + Self::transfer_full_lock_to_coldkey( + &lease.coldkey, + &lease.beneficiary, + lease.netuid, + &hotkey, + )?; + Self::remove_lease_coldkey_references(&lease); + + // Remove the proxy before dec_providers: its reserved deposit holds a consumer ref, + // so decrementing providers first would fail with ConsumerRemaining and leak the account. + T::ProxyInterface::remove_lease_beneficiary_proxy(&lease.coldkey, &lease.beneficiary)?; + + // Sweep the now-unreserved deposit off the keyless lease coldkey so it isn't stranded. + let remaining = ::Currency::balance(&lease.coldkey); + if !remaining.is_zero() { + ::Currency::transfer( + &lease.coldkey, + &lease.beneficiary, + remaining, + Preservation::Expendable, + )?; + } // Stop tracking the lease coldkey and hotkey let _ = frame_system::Pallet::::dec_providers(&lease.coldkey).defensive(); @@ -229,13 +260,11 @@ impl Pallet { let clear_result = SubnetLeaseShares::::clear_prefix(lease_id, T::MaxContributors::get(), None); AccumulatedLeaseDividends::::remove(lease_id); + SubnetUidToLeaseId::::remove(lease.netuid); SubnetLeases::::remove(lease_id); - // Remove the beneficiary proxy - T::ProxyInterface::remove_lease_beneficiary_proxy(&lease.coldkey, &lease.beneficiary)?; - Self::deposit_event(Event::SubnetLeaseTerminated { - beneficiary: lease.beneficiary, + beneficiary: lease.beneficiary.clone(), netuid: lease.netuid, }); @@ -251,6 +280,54 @@ impl Pallet { } } + fn move_lease_lock_to_beneficiary_hotkey( + lease: &SubnetLeaseOf, + beneficiary_hotkey: &T::AccountId, + ) -> Result<(), DispatchError> { + let Some((locked_hotkey, _)) = + Lock::::iter_prefix((&lease.coldkey, lease.netuid)).next() + else { + return Ok(()); + }; + + if locked_hotkey != *beneficiary_hotkey { + Self::do_move_lock(&lease.coldkey, beneficiary_hotkey, lease.netuid, true)?; + } + + Ok(()) + } + + fn repatriate_lease_coldkey_alpha( + lease: &SubnetLeaseOf, + beneficiary_hotkey: &T::AccountId, + ) -> Result<(), DispatchError> { + let alpha = Self::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + ); + if !alpha.is_zero() { + Self::transfer_stake_within_subnet( + &lease.coldkey, + &lease.hotkey, + &lease.beneficiary, + beneficiary_hotkey, + lease.netuid, + alpha, + false, + )?; + } + + Ok(()) + } + + fn remove_lease_coldkey_references(lease: &SubnetLeaseOf) { + OwnedHotkeys::::remove(&lease.coldkey); + StakingHotkeys::::remove(&lease.coldkey); + DecayingLock::::remove(&lease.coldkey, lease.netuid); + LastColdkeyHotkeyStakeBlock::::remove(&lease.coldkey, &lease.hotkey); + } + /// Hook used when the subnet owner's cut is distributed to split the amount into dividends /// for the contributors and the beneficiary in shares relative to their initial contributions. /// It accumulates dividends to be distributed later when the interval for distribution is reached. @@ -312,6 +389,7 @@ impl Pallet { &lease.hotkey, lease.netuid, alpha_for_contributor.into(), + true, )?; alpha_distributed = alpha_distributed.saturating_add(alpha_for_contributor.into()); @@ -332,6 +410,7 @@ impl Pallet { &lease.hotkey, lease.netuid, beneficiary_cut_alpha.into(), + true, )?; Self::deposit_event(Event::SubnetLeaseDividendsDistributed { lease_id, @@ -408,11 +487,11 @@ impl SubnetLeasingWeightInfo { } pub fn do_terminate_lease(k: u32) -> Weight { - Weight::from_parts(56_635_122, 6148) - .saturating_add(Weight::from_parts(912_993, 0).saturating_mul(k.into())) - .saturating_add(T::DbWeight::get().reads(4_u64)) + Weight::from_parts(239_842_400, 11615) + .saturating_add(Weight::from_parts(1_051_411, 0).saturating_mul(k.into())) + .saturating_add(T::DbWeight::get().reads(37_u64)) .saturating_add(T::DbWeight::get().reads((1_u64).saturating_mul(k.into()))) - .saturating_add(T::DbWeight::get().writes(6_u64)) + .saturating_add(T::DbWeight::get().writes(27_u64)) .saturating_add(T::DbWeight::get().writes((1_u64).saturating_mul(k.into()))) .saturating_add(Weight::from_parts(0, 2529).saturating_mul(k.into())) } diff --git a/pallets/subtensor/src/tests/leasing.rs b/pallets/subtensor/src/tests/leasing.rs index cc0715f451..684140b1f2 100644 --- a/pallets/subtensor/src/tests/leasing.rs +++ b/pallets/subtensor/src/tests/leasing.rs @@ -297,6 +297,436 @@ fn test_terminate_lease_works() { }); } +#[test] +fn test_terminate_lease_repatriates_residual_alpha_and_lock_state() { + new_test_ext(1).execute_with(|| { + let crowdloan_id = 0; + let beneficiary = U256::from(1); + let deposit = 10_000_000_000; + let cap = 1_000_000_000_000; + let contributions = vec![(U256::from(2), 990_000_000_000)]; + setup_crowdloan(crowdloan_id, deposit, cap, beneficiary, &contributions); + + let end_block = 500; + let residual_alpha = AlphaBalance::from(10_000_000_u64); + let lock_amount = AlphaBalance::from(4_000_000_u64); + let (lease_id, lease) = setup_leased_network( + beneficiary, + Percent::from_percent(30), + Some(end_block), + None, + ); + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + residual_alpha, + ); + LastColdkeyHotkeyStakeBlock::::insert(lease.coldkey, lease.hotkey, 42); + DecayingLock::::insert(lease.coldkey, lease.netuid, false); + assert_ok!(SubtensorModule::do_lock_stake( + &lease.coldkey, + lease.netuid, + &lease.hotkey, + lock_amount, + )); + + assert_eq!( + SubtensorModule::total_coldkey_alpha_on_subnet(&lease.coldkey, lease.netuid), + residual_alpha + ); + assert!(OwnerLock::::get(lease.netuid).is_some()); + + run_to_block(end_block); + + let beneficiary_hotkey = U256::from(3); + assert_ok!(SubtensorModule::create_account_if_non_existent( + &beneficiary, + &beneficiary_hotkey + )); + + assert_ok!(SubtensorModule::terminate_lease( + RuntimeOrigin::signed(beneficiary), + lease_id, + beneficiary_hotkey, + )); + + assert_eq!( + SubtensorModule::total_coldkey_alpha_on_subnet(&lease.coldkey, lease.netuid), + AlphaBalance::ZERO + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + ), + AlphaBalance::ZERO + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &beneficiary_hotkey, + &beneficiary, + lease.netuid, + ), + residual_alpha + ); + + assert!(Lock::::get((lease.coldkey, lease.netuid, lease.hotkey)).is_none()); + assert!(Lock::::get((lease.coldkey, lease.netuid, beneficiary_hotkey)).is_none()); + let beneficiary_lock = + Lock::::get((beneficiary, lease.netuid, beneficiary_hotkey)).unwrap(); + assert_eq!(beneficiary_lock.locked_mass, lock_amount); + assert_eq!( + beneficiary_lock.conviction, + U64F64::saturating_from_num(u64::from(lock_amount)) + ); + + assert!(HotkeyLock::::get(lease.netuid, lease.hotkey).is_none()); + assert!(DecayingHotkeyLock::::get(lease.netuid, lease.hotkey).is_none()); + assert!(OwnerLock::::get(lease.netuid).is_none()); + assert_eq!( + DecayingOwnerLock::::get(lease.netuid) + .unwrap() + .locked_mass, + lock_amount + ); + + assert_eq!(SubnetOwner::::get(lease.netuid), beneficiary); + assert_eq!( + SubnetOwnerHotkey::::get(lease.netuid), + beneficiary_hotkey + ); + assert_eq!(Owner::::get(beneficiary_hotkey), beneficiary); + assert_eq!(Owner::::get(lease.hotkey), lease.coldkey); + assert!(!OwnedHotkeys::::get(beneficiary).contains(&lease.hotkey)); + assert!(OwnedHotkeys::::get(lease.coldkey).is_empty()); + assert!(StakingHotkeys::::get(lease.coldkey).is_empty()); + assert!(DecayingLock::::get(lease.coldkey, lease.netuid).is_none()); + assert!(LastColdkeyHotkeyStakeBlock::::get(lease.coldkey, lease.hotkey).is_none()); + assert!(SubnetUidToLeaseId::::get(lease.netuid).is_none()); + }); +} + +#[test] +fn test_terminate_lease_repatriates_below_minimum_residual_alpha() { + new_test_ext(1).execute_with(|| { + let crowdloan_id = 0; + let beneficiary = U256::from(1); + let deposit = 10_000_000_000; + let cap = 1_000_000_000_000; + let contributions = vec![(U256::from(2), 990_000_000_000)]; + setup_crowdloan(crowdloan_id, deposit, cap, beneficiary, &contributions); + + let end_block = 500; + let (lease_id, lease) = setup_leased_network( + beneficiary, + Percent::from_percent(30), + Some(end_block), + None, + ); + + let dust_alpha = AlphaBalance::from(1_u64); + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + dust_alpha, + ); + + run_to_block(end_block); + + let beneficiary_hotkey = U256::from(3); + assert_ok!(SubtensorModule::create_account_if_non_existent( + &beneficiary, + &beneficiary_hotkey + )); + + assert_ok!(SubtensorModule::terminate_lease( + RuntimeOrigin::signed(beneficiary), + lease_id, + beneficiary_hotkey, + )); + + assert_eq!( + SubtensorModule::total_coldkey_alpha_on_subnet(&lease.coldkey, lease.netuid), + AlphaBalance::ZERO + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &beneficiary_hotkey, + &beneficiary, + lease.netuid, + ), + dust_alpha + ); + assert!(StakingHotkeys::::get(lease.coldkey).is_empty()); + }); +} + +#[test] +fn test_terminate_lease_repatriates_lock_without_residual_alpha() { + new_test_ext(1).execute_with(|| { + let crowdloan_id = 0; + let beneficiary = U256::from(1); + let deposit = 10_000_000_000; + let cap = 1_000_000_000_000; + let contributions = vec![(U256::from(2), 990_000_000_000)]; + setup_crowdloan(crowdloan_id, deposit, cap, beneficiary, &contributions); + + let end_block = 500; + let residual_alpha = AlphaBalance::from(10_000_000_u64); + let lock_amount = AlphaBalance::from(4_000_000_u64); + let (lease_id, lease) = setup_leased_network( + beneficiary, + Percent::from_percent(30), + Some(end_block), + None, + ); + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + residual_alpha, + ); + DecayingLock::::insert(lease.coldkey, lease.netuid, false); + assert_ok!(SubtensorModule::do_lock_stake( + &lease.coldkey, + lease.netuid, + &lease.hotkey, + lock_amount, + )); + SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + residual_alpha, + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + ), + AlphaBalance::ZERO + ); + assert!(Lock::::get((lease.coldkey, lease.netuid, lease.hotkey)).is_some()); + + run_to_block(end_block); + + let beneficiary_hotkey = U256::from(3); + assert_ok!(SubtensorModule::create_account_if_non_existent( + &beneficiary, + &beneficiary_hotkey + )); + + assert_ok!(SubtensorModule::terminate_lease( + RuntimeOrigin::signed(beneficiary), + lease_id, + beneficiary_hotkey, + )); + + assert!(Lock::::get((lease.coldkey, lease.netuid, lease.hotkey)).is_none()); + assert!(Lock::::get((lease.coldkey, lease.netuid, beneficiary_hotkey)).is_none()); + assert_eq!( + Lock::::get((beneficiary, lease.netuid, beneficiary_hotkey)) + .unwrap() + .locked_mass, + lock_amount + ); + assert_eq!( + SubtensorModule::total_coldkey_alpha_on_subnet(&lease.coldkey, lease.netuid), + AlphaBalance::ZERO + ); + assert!(OwnedHotkeys::::get(lease.coldkey).is_empty()); + assert!(StakingHotkeys::::get(lease.coldkey).is_empty()); + }); +} + +#[test] +fn test_terminate_lease_merges_partial_residual_lock_with_existing_beneficiary_lock() { + new_test_ext(1).execute_with(|| { + let crowdloan_id = 0; + let beneficiary = U256::from(1); + let deposit = 10_000_000_000; + let cap = 1_000_000_000_000; + let contributions = vec![(U256::from(2), 990_000_000_000)]; + setup_crowdloan(crowdloan_id, deposit, cap, beneficiary, &contributions); + + let end_block = 500; + let residual_alpha = AlphaBalance::from(2_000_000_u64); + let lease_lock_amount = AlphaBalance::from(4_000_000_u64); + let beneficiary_lock_amount = AlphaBalance::from(1_000_000_u64); + let (lease_id, lease) = setup_leased_network( + beneficiary, + Percent::from_percent(30), + Some(end_block), + None, + ); + + let beneficiary_hotkey = U256::from(3); + assert_ok!(SubtensorModule::create_account_if_non_existent( + &beneficiary, + &beneficiary_hotkey + )); + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &beneficiary_hotkey, + &beneficiary, + lease.netuid, + beneficiary_lock_amount, + ); + DecayingLock::::insert(beneficiary, lease.netuid, false); + assert_ok!(SubtensorModule::do_lock_stake( + &beneficiary, + lease.netuid, + &beneficiary_hotkey, + beneficiary_lock_amount, + )); + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + lease_lock_amount, + ); + DecayingLock::::insert(lease.coldkey, lease.netuid, false); + assert_ok!(SubtensorModule::do_lock_stake( + &lease.coldkey, + lease.netuid, + &lease.hotkey, + lease_lock_amount, + )); + SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + lease_lock_amount.saturating_sub(residual_alpha), + ); + + run_to_block(end_block); + + assert_ok!(SubtensorModule::terminate_lease( + RuntimeOrigin::signed(beneficiary), + lease_id, + beneficiary_hotkey, + )); + + assert!(Lock::::get((lease.coldkey, lease.netuid, lease.hotkey)).is_none()); + assert!(Lock::::get((lease.coldkey, lease.netuid, beneficiary_hotkey)).is_none()); + assert_eq!( + Lock::::get((beneficiary, lease.netuid, beneficiary_hotkey)) + .unwrap() + .locked_mass, + beneficiary_lock_amount.saturating_add(lease_lock_amount) + ); + assert_eq!( + OwnerLock::::get(lease.netuid).unwrap().locked_mass, + beneficiary_lock_amount.saturating_add(lease_lock_amount) + ); + assert!(HotkeyLock::::get(lease.netuid, beneficiary_hotkey).is_none()); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &beneficiary_hotkey, + &beneficiary, + lease.netuid, + ), + beneficiary_lock_amount.saturating_add(residual_alpha) + ); + assert_eq!( + SubtensorModule::total_coldkey_alpha_on_subnet(&lease.coldkey, lease.netuid), + AlphaBalance::ZERO + ); + assert!(StakingHotkeys::::get(lease.coldkey).is_empty()); + }); +} + +#[test] +fn test_terminate_lease_rolls_back_if_repatriated_lock_conflicts() { + new_test_ext(1).execute_with(|| { + let crowdloan_id = 0; + let beneficiary = U256::from(1); + let deposit = 10_000_000_000; + let cap = 1_000_000_000_000; + let contributions = vec![(U256::from(2), 990_000_000_000)]; + setup_crowdloan(crowdloan_id, deposit, cap, beneficiary, &contributions); + + let end_block = 500; + let residual_alpha = AlphaBalance::from(10_000_000_u64); + let lock_amount = AlphaBalance::from(4_000_000_u64); + let (lease_id, lease) = setup_leased_network( + beneficiary, + Percent::from_percent(30), + Some(end_block), + None, + ); + + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + residual_alpha, + ); + assert_ok!(SubtensorModule::do_lock_stake( + &lease.coldkey, + lease.netuid, + &lease.hotkey, + lock_amount, + )); + + let beneficiary_hotkey = U256::from(3); + let conflicting_hotkey = U256::from(4); + assert_ok!(SubtensorModule::create_account_if_non_existent( + &beneficiary, + &beneficiary_hotkey + )); + assert_ok!(SubtensorModule::create_account_if_non_existent( + &beneficiary, + &conflicting_hotkey + )); + SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet( + &conflicting_hotkey, + &beneficiary, + lease.netuid, + residual_alpha, + ); + assert_ok!(SubtensorModule::do_lock_stake( + &beneficiary, + lease.netuid, + &conflicting_hotkey, + lock_amount, + )); + + run_to_block(end_block); + + assert_err!( + SubtensorModule::terminate_lease( + RuntimeOrigin::signed(beneficiary), + lease_id, + beneficiary_hotkey, + ), + Error::::LockHotkeyMismatch, + ); + + assert!(SubnetLeases::::contains_key(lease_id)); + assert_eq!(SubnetOwner::::get(lease.netuid), lease.coldkey); + assert_eq!(SubnetOwnerHotkey::::get(lease.netuid), lease.hotkey); + assert_eq!(Owner::::get(lease.hotkey), lease.coldkey); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + ), + residual_alpha + ); + assert!(Lock::::get((lease.coldkey, lease.netuid, lease.hotkey)).is_some()); + assert!(PROXIES.with_borrow(|proxies| proxies.0 == vec![(lease.coldkey, beneficiary)])); + }); +} + #[test] fn test_terminate_lease_fails_if_bad_origin() { new_test_ext(1).execute_with(|| { From ddba12da4f4cdabaa0c53c69c127fad42649dae2 Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Tue, 2 Jun 2026 17:23:09 -0300 Subject: [PATCH 4/8] Update benchmark --- pallets/subtensor/src/benchmarks.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/pallets/subtensor/src/benchmarks.rs b/pallets/subtensor/src/benchmarks.rs index 1dd62bab0b..346d61c859 100644 --- a/pallets/subtensor/src/benchmarks.rs +++ b/pallets/subtensor/src/benchmarks.rs @@ -1768,6 +1768,22 @@ mod pallet_benchmarks { let hotkey = account::("beneficiary_hotkey", 0, 0); let _ = Subtensor::::create_account_if_non_existent(&beneficiary, &hotkey); + let residual_alpha = AlphaBalance::from(10_000_000_u64); + let lock_amount = AlphaBalance::from(4_000_000_u64); + Subtensor::::increase_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &lease.coldkey, + lease.netuid, + residual_alpha, + ); + DecayingLock::::insert(&lease.coldkey, lease.netuid, false); + assert_ok!(Subtensor::::do_lock_stake( + &lease.coldkey, + lease.netuid, + &lease.hotkey, + lock_amount, + )); + #[extrinsic_call] _( RawOrigin::Signed(beneficiary.clone()), @@ -1781,6 +1797,10 @@ mod pallet_benchmarks { assert_eq!(SubnetLeases::::get(lease_id), None); assert!(!SubnetLeaseShares::::contains_prefix(lease_id)); assert!(!AccumulatedLeaseDividends::::contains_key(lease_id)); + assert_eq!( + Subtensor::::total_coldkey_alpha_on_subnet(&lease.coldkey, lease.netuid), + AlphaBalance::ZERO + ); } #[benchmark] From 1a2c1ced5440b4a4c2de667c137c5170587bc28c Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Tue, 2 Jun 2026 17:42:05 -0300 Subject: [PATCH 5/8] Remove hardcoded leasing weights --- pallets/subtensor/src/subnets/leasing.rs | 37 ++++-------------------- 1 file changed, 6 insertions(+), 31 deletions(-) diff --git a/pallets/subtensor/src/subnets/leasing.rs b/pallets/subtensor/src/subnets/leasing.rs index a6a8b2b734..48403fe2ca 100644 --- a/pallets/subtensor/src/subnets/leasing.rs +++ b/pallets/subtensor/src/subnets/leasing.rs @@ -16,6 +16,7 @@ //! ownership will be transferred to the beneficiary. use super::*; +use crate::weights::WeightInfo; use frame_support::{ dispatch::RawOrigin, traits::{Defensive, fungible::*, tokens::Preservation}, @@ -178,12 +179,10 @@ impl Pallet { if crowdloan.contributors_count < T::MaxContributors::get() { // We have less contributors than the max allowed, so we need to refund the difference - Ok( - Some(SubnetLeasingWeightInfo::::do_register_leased_network( - crowdloan.contributors_count, - )) - .into(), - ) + Ok(Some(::WeightInfo::register_leased_network( + crowdloan.contributors_count, + )) + .into()) } else { // We have the max number of contributors, so we don't need to refund anything Ok(().into()) @@ -270,7 +269,7 @@ impl Pallet { if clear_result.unique < T::MaxContributors::get() { // We have cleared less than the max number of shareholders, so we need to refund the difference - Ok(Some(SubnetLeasingWeightInfo::::do_terminate_lease( + Ok(Some(::WeightInfo::terminate_lease( clear_result.unique, )) .into()) @@ -472,27 +471,3 @@ impl Pallet { Ok((crowdloan_id, crowdloan)) } } - -/// Weight functions needed for subnet leasing. -pub struct SubnetLeasingWeightInfo(PhantomData); -impl SubnetLeasingWeightInfo { - pub fn do_register_leased_network(k: u32) -> Weight { - Weight::from_parts(301_560_714, 10079) - .saturating_add(Weight::from_parts(26_884_006, 0).saturating_mul(k.into())) - .saturating_add(T::DbWeight::get().reads(41_u64)) - .saturating_add(T::DbWeight::get().reads(2_u64.saturating_mul(k.into()))) - .saturating_add(T::DbWeight::get().writes(55_u64)) - .saturating_add(T::DbWeight::get().writes(2_u64.saturating_mul(k.into()))) - .saturating_add(Weight::from_parts(0, 2579).saturating_mul(k.into())) - } - - pub fn do_terminate_lease(k: u32) -> Weight { - Weight::from_parts(239_842_400, 11615) - .saturating_add(Weight::from_parts(1_051_411, 0).saturating_mul(k.into())) - .saturating_add(T::DbWeight::get().reads(37_u64)) - .saturating_add(T::DbWeight::get().reads((1_u64).saturating_mul(k.into()))) - .saturating_add(T::DbWeight::get().writes(27_u64)) - .saturating_add(T::DbWeight::get().writes((1_u64).saturating_mul(k.into()))) - .saturating_add(Weight::from_parts(0, 2529).saturating_mul(k.into())) - } -} From ee4bdfed6d9f19311f2efd9d28e3d46d92f8ddfc Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Thu, 4 Jun 2026 14:17:04 -0300 Subject: [PATCH 6/8] Hotkey swap for lease hotkey to beneficiary --- pallets/subtensor/src/subnets/leasing.rs | 25 ++++++++++ pallets/subtensor/src/tests/leasing.rs | 62 +++++++++++++++++++++++- 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/pallets/subtensor/src/subnets/leasing.rs b/pallets/subtensor/src/subnets/leasing.rs index 48403fe2ca..ba2feaf34f 100644 --- a/pallets/subtensor/src/subnets/leasing.rs +++ b/pallets/subtensor/src/subnets/leasing.rs @@ -20,6 +20,7 @@ use crate::weights::WeightInfo; use frame_support::{ dispatch::RawOrigin, traits::{Defensive, fungible::*, tokens::Preservation}, + weights::Weight, }; use frame_system::pallet_prelude::OriginFor; use frame_system::pallet_prelude::*; @@ -217,6 +218,10 @@ impl Pallet { Self::coldkey_owns_hotkey(&lease.beneficiary, &hotkey), Error::::BeneficiaryDoesNotOwnHotkey ); + ensure!( + !Self::is_hotkey_registered_on_specific_network(&hotkey, lease.netuid), + Error::::HotKeyAlreadyRegisteredInSubNet + ); // Move any lease coldkey lock to the beneficiary-controlled hotkey without // assigning ownership of the generated lease hotkey to the beneficiary. @@ -224,6 +229,8 @@ impl Pallet { // Transfer ownership to the beneficiary SubnetOwner::::insert(lease.netuid, lease.beneficiary.clone()); + // Set the owner hotkey before moving locks so the destination lock is + // accounted in the subnet-owner aggregate, not the regular hotkey aggregate. Self::set_subnet_owner_hotkey(lease.netuid, &hotkey)?; Self::reassign_subnet_owner_lock_aggregates(lease.netuid, &lease.hotkey, &hotkey); @@ -234,6 +241,7 @@ impl Pallet { lease.netuid, &hotkey, )?; + Self::replace_lease_hotkey_with_beneficiary_hotkey(&lease, &hotkey)?; Self::remove_lease_coldkey_references(&lease); // Remove the proxy before dec_providers: its reserved deposit holds a consumer ref, @@ -327,6 +335,23 @@ impl Pallet { LastColdkeyHotkeyStakeBlock::::remove(&lease.coldkey, &lease.hotkey); } + fn replace_lease_hotkey_with_beneficiary_hotkey( + lease: &SubnetLeaseOf, + beneficiary_hotkey: &T::AccountId, + ) -> DispatchResult { + let mut weight = Weight::zero(); + Self::perform_hotkey_swap_on_one_subnet( + &lease.hotkey, + beneficiary_hotkey, + &mut weight, + lease.netuid, + false, + )?; + Owner::::remove(&lease.hotkey); + Delegates::::remove(&lease.hotkey); + Ok(()) + } + /// Hook used when the subnet owner's cut is distributed to split the amount into dividends /// for the contributors and the beneficiary in shares relative to their initial contributions. /// It accumulates dividends to be distributed later when the interval for distribution is reached. diff --git a/pallets/subtensor/src/tests/leasing.rs b/pallets/subtensor/src/tests/leasing.rs index 684140b1f2..2a6c75b483 100644 --- a/pallets/subtensor/src/tests/leasing.rs +++ b/pallets/subtensor/src/tests/leasing.rs @@ -399,8 +399,17 @@ fn test_terminate_lease_repatriates_residual_alpha_and_lock_state() { beneficiary_hotkey ); assert_eq!(Owner::::get(beneficiary_hotkey), beneficiary); - assert_eq!(Owner::::get(lease.hotkey), lease.coldkey); + assert!(!Owner::::contains_key(lease.hotkey)); assert!(!OwnedHotkeys::::get(beneficiary).contains(&lease.hotkey)); + assert!(!StakingHotkeys::::get(beneficiary).contains(&lease.hotkey)); + assert!(!SubtensorModule::is_hotkey_registered_on_network( + lease.netuid, + &lease.hotkey + )); + assert!(SubtensorModule::is_hotkey_registered_on_network( + lease.netuid, + &beneficiary_hotkey + )); assert!(OwnedHotkeys::::get(lease.coldkey).is_empty()); assert!(StakingHotkeys::::get(lease.coldkey).is_empty()); assert!(DecayingLock::::get(lease.coldkey, lease.netuid).is_none()); @@ -907,6 +916,57 @@ fn test_terminate_lease_fails_if_beneficiary_does_not_own_hotkey() { ); }); } + +#[test] +fn test_terminate_lease_fails_if_beneficiary_hotkey_already_registered_on_subnet() { + new_test_ext(1).execute_with(|| { + // Setup a crowdloan + let crowdloan_id = 0; + let beneficiary = U256::from(1); + let deposit = 10_000_000_000; // 10 TAO + let cap = 1_000_000_000_000; // 1000 TAO + let contributions = vec![(U256::from(2), 990_000_000_000)]; // 990 TAO + setup_crowdloan(crowdloan_id, deposit, cap, beneficiary, &contributions); + + // Setup a leased network + let end_block = 500; + let tao_to_stake = 100_000_000_000; // 100 TAO + let emissions_share = Percent::from_percent(30); + let (lease_id, lease) = setup_leased_network( + beneficiary, + emissions_share, + Some(end_block), + Some(tao_to_stake), + ); + + // Run to the end of the lease + run_to_block(end_block); + + // Create a beneficiary-owned hotkey that is already registered on the leased subnet + let beneficiary_hotkey = U256::from(3); + assert_ok!(SubtensorModule::create_account_if_non_existent( + &beneficiary, + &beneficiary_hotkey + )); + SubtensorModule::append_neuron(lease.netuid, &beneficiary_hotkey, 0); + + // Termination requires a fresh beneficiary hotkey so it can replace the lease hotkey UID. + assert_err!( + SubtensorModule::terminate_lease( + RuntimeOrigin::signed(lease.beneficiary), + lease_id, + beneficiary_hotkey, + ), + Error::::HotKeyAlreadyRegisteredInSubNet, + ); + + assert!(SubnetLeases::::contains_key(lease_id)); + assert_eq!(SubnetOwner::::get(lease.netuid), lease.coldkey); + assert_eq!(SubnetOwnerHotkey::::get(lease.netuid), lease.hotkey); + assert_eq!(Owner::::get(lease.hotkey), lease.coldkey); + }); +} + #[test] fn test_distribute_lease_network_dividends_multiple_contributors_works() { new_test_ext(1).execute_with(|| { From 4068519e24359b329ee3379e29cda76be9f0213a Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Mon, 8 Jun 2026 11:02:10 -0300 Subject: [PATCH 7/8] Charge hotkey swap weight for lease termination --- pallets/subtensor/src/macros/dispatches.rs | 5 ++++- pallets/subtensor/src/subnets/leasing.rs | 5 +++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index 2a73a0e473..836d482043 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -1950,7 +1950,10 @@ mod dispatches { /// * `hotkey` (T::AccountId): /// - The hotkey of the beneficiary to mark as subnet owner hotkey. #[pallet::call_index(111)] - #[pallet::weight(::WeightInfo::terminate_lease(T::MaxContributors::get()))] + #[pallet::weight( + ::WeightInfo::terminate_lease(T::MaxContributors::get()) + .saturating_add(::WeightInfo::swap_hotkey()) + )] pub fn terminate_lease( origin: OriginFor, lease_id: LeaseId, diff --git a/pallets/subtensor/src/subnets/leasing.rs b/pallets/subtensor/src/subnets/leasing.rs index ba2feaf34f..b60b7d49d3 100644 --- a/pallets/subtensor/src/subnets/leasing.rs +++ b/pallets/subtensor/src/subnets/leasing.rs @@ -340,7 +340,8 @@ impl Pallet { beneficiary_hotkey: &T::AccountId, ) -> DispatchResult { let mut weight = Weight::zero(); - Self::perform_hotkey_swap_on_one_subnet( + // We don't refund any weight for the hotkey swap. + let _ = Self::perform_hotkey_swap_on_one_subnet( &lease.hotkey, beneficiary_hotkey, &mut weight, @@ -349,7 +350,7 @@ impl Pallet { )?; Owner::::remove(&lease.hotkey); Delegates::::remove(&lease.hotkey); - Ok(()) + Ok(weight) } /// Hook used when the subnet owner's cut is distributed to split the amount into dividends From 8fa9a015500fd165a74f49b4366632c42918d2bd Mon Sep 17 00:00:00 2001 From: Loris Moulin Date: Mon, 8 Jun 2026 11:55:17 -0300 Subject: [PATCH 8/8] Fix reentrancy --- pallets/crowdloan/src/lib.rs | 8 +- pallets/crowdloan/src/tests.rs | 128 +++++++++++++++++++++++ pallets/subtensor/src/subnets/leasing.rs | 6 +- 3 files changed, 136 insertions(+), 6 deletions(-) diff --git a/pallets/crowdloan/src/lib.rs b/pallets/crowdloan/src/lib.rs index 8be80c0cc7..b8c45251bb 100644 --- a/pallets/crowdloan/src/lib.rs +++ b/pallets/crowdloan/src/lib.rs @@ -620,6 +620,11 @@ pub mod pallet { ensure!(crowdloan.raised == crowdloan.cap, Error::::CapNotRaised); ensure!(!crowdloan.finalized, Error::::AlreadyFinalized); + // Mark finalized *before* dispatching the creator-controlled call so re-entrant + // withdraw/refund/dissolve are rejected mid-dispatch. Reverts if the dispatch fails. + crowdloan.finalized = true; + Crowdloans::::insert(crowdloan_id, &crowdloan); + match (&crowdloan.call, &crowdloan.target_address) { (Some(call), None) => { // Set the current crowdloan id so the dispatched call @@ -659,9 +664,6 @@ pub mod pallet { } } - crowdloan.finalized = true; - Crowdloans::::insert(crowdloan_id, &crowdloan); - Self::deposit_event(Event::::Finalized { crowdloan_id }); Ok(()) diff --git a/pallets/crowdloan/src/tests.rs b/pallets/crowdloan/src/tests.rs index a23de52198..9cffa10ef5 100644 --- a/pallets/crowdloan/src/tests.rs +++ b/pallets/crowdloan/src/tests.rs @@ -1875,6 +1875,134 @@ fn test_finalize_fails_if_call_fails() { }); } +// The finalize `call` cannot re-enter `withdraw` on the same crowdloan: it is rejected and +// the extrinsic reverts, so no funds move and `raised` stays consistent with the real balance. +#[test] +fn test_finalize_blocks_reentrant_withdraw() { + TestState::default() + .with_balance(U256::from(1), 200.into()) // creator + .with_balance(U256::from(2), 200.into()) // contributor + .build_and_execute(|| { + let creator: AccountOf = U256::from(1); + let contributor: AccountOf = U256::from(2); + let deposit: BalanceOf = 50.into(); + let min_contribution: BalanceOf = 10.into(); + let cap: BalanceOf = 100.into(); + let end: BlockNumberFor = 50; + let crowdloan_id: CrowdloanId = 0; + + // The finalize call re-enters `withdraw` on the same crowdloan. + let reentrant_call = Box::new(RuntimeCall::Crowdloan( + pallet_crowdloan::Call::::withdraw { crowdloan_id }, + )); + + assert_ok!(Crowdloan::create( + RuntimeOrigin::signed(creator), + deposit, + min_contribution, + cap, + end, + Some(reentrant_call), + None, + )); + run_to_block(10); + + // Creator contributes 30 over the deposit (total 80); contributor fills the cap. + assert_ok!(Crowdloan::contribute( + RuntimeOrigin::signed(creator), + crowdloan_id, + 30.into() + )); + assert_ok!(Crowdloan::contribute( + RuntimeOrigin::signed(contributor), + crowdloan_id, + 20.into() + )); + + let funds_account = pallet_crowdloan::Pallet::::funds_account(crowdloan_id); + assert_eq!(Balances::free_balance(funds_account), cap); + let creator_balance_before = Balances::free_balance(creator); + + run_to_block(60); + + // Finalize dispatches the re-entrant withdraw, which is rejected with + // `AlreadyFinalized`. Wrap in a storage layer to model the per-extrinsic + // transaction the runtime applies in production, so the revert is observable. + let outcome = frame_support::storage::with_storage_layer(|| { + Crowdloan::finalize(RuntimeOrigin::signed(creator), crowdloan_id) + }); + assert_err!(outcome, pallet_crowdloan::Error::::AlreadyFinalized); + + // No funds were extracted and accounting is intact. + assert_eq!(Balances::free_balance(creator), creator_balance_before); + assert_eq!(Balances::free_balance(funds_account), cap); + assert_eq!(pallet_crowdloan::CurrentCrowdloanId::::get(), None); + let crowdloan = pallet_crowdloan::Crowdloans::::get(crowdloan_id).unwrap(); + assert!(!crowdloan.finalized); + assert_eq!(crowdloan.raised, cap); + + // Contributor funds are not frozen: the contributor can still withdraw. + assert_ok!(Crowdloan::withdraw( + RuntimeOrigin::signed(contributor), + crowdloan_id + )); + assert_eq!(Balances::free_balance(contributor), 200.into()); + }); +} + +// A re-entrant `refund` embedded as the finalize call is likewise rejected before moving funds. +#[test] +fn test_finalize_blocks_reentrant_refund() { + TestState::default() + .with_balance(U256::from(1), 200.into()) // creator + .with_balance(U256::from(2), 200.into()) // contributor + .build_and_execute(|| { + let creator: AccountOf = U256::from(1); + let contributor: AccountOf = U256::from(2); + let deposit: BalanceOf = 50.into(); + let min_contribution: BalanceOf = 10.into(); + let cap: BalanceOf = 100.into(); + let end: BlockNumberFor = 50; + let crowdloan_id: CrowdloanId = 0; + + let reentrant_call = Box::new(RuntimeCall::Crowdloan( + pallet_crowdloan::Call::::refund { crowdloan_id }, + )); + + assert_ok!(Crowdloan::create( + RuntimeOrigin::signed(creator), + deposit, + min_contribution, + cap, + end, + Some(reentrant_call), + None, + )); + run_to_block(10); + + assert_ok!(Crowdloan::contribute( + RuntimeOrigin::signed(creator), + crowdloan_id, + 30.into() + )); + assert_ok!(Crowdloan::contribute( + RuntimeOrigin::signed(contributor), + crowdloan_id, + 20.into() + )); + + let funds_account = pallet_crowdloan::Pallet::::funds_account(crowdloan_id); + run_to_block(60); + + // The re-entrant refund hits the `finalized` guard before transferring anything. + assert_err!( + Crowdloan::finalize(RuntimeOrigin::signed(creator), crowdloan_id), + pallet_crowdloan::Error::::AlreadyFinalized + ); + assert_eq!(Balances::free_balance(funds_account), cap); + }); +} + #[test] fn test_refund_succeeds() { TestState::default() diff --git a/pallets/subtensor/src/subnets/leasing.rs b/pallets/subtensor/src/subnets/leasing.rs index b60b7d49d3..c99b51c99a 100644 --- a/pallets/subtensor/src/subnets/leasing.rs +++ b/pallets/subtensor/src/subnets/leasing.rs @@ -339,9 +339,9 @@ impl Pallet { lease: &SubnetLeaseOf, beneficiary_hotkey: &T::AccountId, ) -> DispatchResult { - let mut weight = Weight::zero(); // We don't refund any weight for the hotkey swap. - let _ = Self::perform_hotkey_swap_on_one_subnet( + let mut weight = Weight::zero(); + Self::perform_hotkey_swap_on_one_subnet( &lease.hotkey, beneficiary_hotkey, &mut weight, @@ -350,7 +350,7 @@ impl Pallet { )?; Owner::::remove(&lease.hotkey); Delegates::::remove(&lease.hotkey); - Ok(weight) + Ok(()) } /// Hook used when the subnet owner's cut is distributed to split the amount into dividends