From 185841b7fc7e37fa09a1081305c4c0fbabe1211d Mon Sep 17 00:00:00 2001
From: Greg Zaitsev <gregz@opentensor.dev>
Date: Wed, 12 Feb 2025 18:02:34 -0500
Subject: [PATCH] Remove ownership check from transition stake validation

---
 pallets/subtensor/src/staking/stake_utils.rs |  6 ----
 pallets/subtensor/src/tests/move_stake.rs    | 38 ++++++++++++++++++--
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/pallets/subtensor/src/staking/stake_utils.rs b/pallets/subtensor/src/staking/stake_utils.rs
index 1be858529d..03c288ebef 100644
--- a/pallets/subtensor/src/staking/stake_utils.rs
+++ b/pallets/subtensor/src/staking/stake_utils.rs
@@ -960,12 +960,6 @@ impl<T: Config> Pallet<T> {
             Error::<T>::HotKeyAccountNotExists
         );
 
-        // Ensure origin coldkey owns the origin hotkey.
-        ensure!(
-            Self::coldkey_owns_hotkey(origin_coldkey, origin_hotkey),
-            Error::<T>::NonAssociatedColdKey
-        );
-
         // Ensure there is enough stake in the origin subnet.
         let origin_alpha = Self::get_stake_for_hotkey_and_coldkey_on_subnet(
             origin_hotkey,
diff --git a/pallets/subtensor/src/tests/move_stake.rs b/pallets/subtensor/src/tests/move_stake.rs
index ce0ce37fce..15832b93fe 100644
--- a/pallets/subtensor/src/tests/move_stake.rs
+++ b/pallets/subtensor/src/tests/move_stake.rs
@@ -541,7 +541,7 @@ fn test_do_move_wrong_origin() {
                 netuid,
                 alpha,
             ),
-            Error::<Test>::NonAssociatedColdKey
+            Error::<Test>::NotEnoughStakeToWithdraw
         );
 
         // Check that no stake was moved
@@ -981,7 +981,7 @@ fn test_do_transfer_wrong_origin() {
                 netuid,
                 stake_amount
             ),
-            Error::<Test>::NonAssociatedColdKey
+            Error::<Test>::NotEnoughStakeToWithdraw
         );
     });
 }
@@ -1245,7 +1245,7 @@ fn test_do_swap_wrong_origin() {
                 netuid2,
                 stake_amount
             ),
-            Error::<Test>::NonAssociatedColdKey
+            Error::<Test>::NotEnoughStakeToWithdraw
         );
     });
 }
@@ -1486,6 +1486,38 @@ fn test_do_swap_multiple_times() {
     });
 }
 
+// cargo test --package pallet-subtensor --lib -- tests::move_stake::test_do_swap_allows_non_owned_hotkey --exact --show-output
+#[test]
+fn test_do_swap_allows_non_owned_hotkey() {
+    new_test_ext(1).execute_with(|| {
+        let subnet_owner_coldkey = U256::from(1001);
+        let subnet_owner_hotkey = U256::from(1002);
+        let origin_netuid = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+        let destination_netuid = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+
+        let coldkey = U256::from(1);
+        let hotkey = U256::from(2);
+        let foreign_coldkey = U256::from(3);
+        let stake_amount = DefaultMinStake::<Test>::get() * 10;
+
+        SubtensorModule::create_account_if_non_existent(&foreign_coldkey, &hotkey);
+        SubtensorModule::stake_into_subnet(&hotkey, &coldkey, origin_netuid, stake_amount, 0);
+        let alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &coldkey,
+            origin_netuid,
+        );
+
+        assert_ok!(SubtensorModule::do_swap_stake(
+            RuntimeOrigin::signed(coldkey),
+            hotkey,
+            origin_netuid,
+            destination_netuid,
+            alpha_before,
+        ));
+    });
+}
+
 // cargo test --package pallet-subtensor --lib -- tests::move_stake::test_swap_stake_limit_validate --exact --show-output
 #[test]
 fn test_swap_stake_limit_validate() {