From 2e42602edaa067e54074205a3158f67532032ec9 Mon Sep 17 00:00:00 2001 From: Landyn Date: Wed, 28 Jan 2026 18:04:53 -0600 Subject: [PATCH 1/6] proxy aware chain exts draft --- chain-extensions/src/lib.rs | 461 ++++++++++++++++++++++++++++++++++ chain-extensions/src/types.rs | 16 ++ 2 files changed, 477 insertions(+) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index eaaee70d27..53aed0f1f3 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -530,6 +530,467 @@ where Ok(RetVal::Converging(Output::Success as u32)) } + // ============================================================ + // Proxy-aware staking functions + // These allow a contract to act on behalf of a real_coldkey + // if the contract has been granted appropriate proxy permissions. + // ============================================================ + FunctionId::AddStakeAsProxyV1 => { + let weight = Weight::from_parts(340_800_000, 0) + .saturating_add(T::DbWeight::get().reads(25_u64)) + .saturating_add(T::DbWeight::get().writes(15)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey, netuid, amount_staked): ( + T::AccountId, + T::AccountId, + NetUid, + TaoCurrency, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + netuid, + amount_staked, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::RemoveStakeAsProxyV1 => { + let weight = Weight::from_parts(196_800_000, 0) + .saturating_add(T::DbWeight::get().reads(20)) + .saturating_add(T::DbWeight::get().writes(10)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey, netuid, amount_unstaked): ( + T::AccountId, + T::AccountId, + NetUid, + AlphaCurrency, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::remove_stake( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + netuid, + amount_unstaked, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::UnstakeAllAsProxyV1 => { + let weight = Weight::from_parts(28_830_000, 0) + .saturating_add(T::DbWeight::get().reads(7)) + .saturating_add(T::DbWeight::get().writes(0)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey): (T::AccountId, T::AccountId) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::unstake_all( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::UnstakeAllAlphaAsProxyV1 => { + let weight = Weight::from_parts(358_500_000, 0) + .saturating_add(T::DbWeight::get().reads(37_u64)) + .saturating_add(T::DbWeight::get().writes(21_u64)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey): (T::AccountId, T::AccountId) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::unstake_all_alpha( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::MoveStakeAsProxyV1 => { + let weight = Weight::from_parts(164_300_000, 0) + .saturating_add(T::DbWeight::get().reads(16_u64)) + .saturating_add(T::DbWeight::get().writes(7_u64)); + + env.charge_weight(weight)?; + + let ( + real_coldkey, + origin_hotkey, + destination_hotkey, + origin_netuid, + destination_netuid, + alpha_amount, + ): (T::AccountId, T::AccountId, T::AccountId, NetUid, NetUid, AlphaCurrency) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::move_stake( + RawOrigin::Signed(real_coldkey).into(), + origin_hotkey, + destination_hotkey, + origin_netuid, + destination_netuid, + alpha_amount, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::TransferStakeAsProxyV1 => { + let weight = Weight::from_parts(160_300_000, 0) + .saturating_add(T::DbWeight::get().reads(14_u64)) + .saturating_add(T::DbWeight::get().writes(6_u64)); + + env.charge_weight(weight)?; + + let (real_coldkey, destination_coldkey, hotkey, origin_netuid, destination_netuid, alpha_amount): ( + T::AccountId, + T::AccountId, + T::AccountId, + NetUid, + NetUid, + AlphaCurrency, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Transfer)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::transfer_stake( + RawOrigin::Signed(real_coldkey).into(), + destination_coldkey, + hotkey, + origin_netuid, + destination_netuid, + alpha_amount, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::SwapStakeAsProxyV1 => { + let weight = Weight::from_parts(351_300_000, 0) + .saturating_add(T::DbWeight::get().reads(36_u64)) + .saturating_add(T::DbWeight::get().writes(22_u64)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey, origin_netuid, destination_netuid, alpha_amount): ( + T::AccountId, + T::AccountId, + NetUid, + NetUid, + AlphaCurrency, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::swap_stake( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + origin_netuid, + destination_netuid, + alpha_amount, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::AddStakeLimitAsProxyV1 => { + let weight = Weight::from_parts(402_900_000, 0) + .saturating_add(T::DbWeight::get().reads(25_u64)) + .saturating_add(T::DbWeight::get().writes(15)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey, netuid, amount_staked, limit_price, allow_partial): ( + T::AccountId, + T::AccountId, + NetUid, + TaoCurrency, + TaoCurrency, + bool, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::add_stake_limit( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + netuid, + amount_staked, + limit_price, + allow_partial, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::RemoveStakeLimitAsProxyV1 => { + let weight = Weight::from_parts(377_400_000, 0) + .saturating_add(T::DbWeight::get().reads(29_u64)) + .saturating_add(T::DbWeight::get().writes(14)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey, netuid, amount_unstaked, limit_price, allow_partial): ( + T::AccountId, + T::AccountId, + NetUid, + AlphaCurrency, + TaoCurrency, + bool, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::remove_stake_limit( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + netuid, + amount_unstaked, + limit_price, + allow_partial, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::SwapStakeLimitAsProxyV1 => { + let weight = Weight::from_parts(411_500_000, 0) + .saturating_add(T::DbWeight::get().reads(36_u64)) + .saturating_add(T::DbWeight::get().writes(22_u64)); + + env.charge_weight(weight)?; + + let ( + real_coldkey, + hotkey, + origin_netuid, + destination_netuid, + alpha_amount, + limit_price, + allow_partial, + ): ( + T::AccountId, + T::AccountId, + NetUid, + NetUid, + AlphaCurrency, + TaoCurrency, + bool, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::swap_stake_limit( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + origin_netuid, + destination_netuid, + alpha_amount, + limit_price, + allow_partial, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::RemoveStakeFullLimitAsProxyV1 => { + let weight = Weight::from_parts(395_300_000, 0) + .saturating_add(T::DbWeight::get().reads(29_u64)) + .saturating_add(T::DbWeight::get().writes(14_u64)); + + env.charge_weight(weight)?; + + let (real_coldkey, hotkey, netuid, limit_price): ( + T::AccountId, + T::AccountId, + NetUid, + Option, + ) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::remove_stake_full_limit( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + netuid, + limit_price, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::SetColdkeyAutoStakeHotkeyAsProxyV1 => { + let weight = Weight::from_parts(29_930_000, 0) + .saturating_add(T::DbWeight::get().reads(5_u64)) + .saturating_add(T::DbWeight::get().writes(2_u64)); + + env.charge_weight(weight)?; + + let (real_coldkey, netuid, hotkey): (T::AccountId, NetUid, T::AccountId) = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let caller = env.caller(); + if real_coldkey != caller { + pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) + .map_err(|_| DispatchError::Other("NotProxy"))?; + } + + let call_result = pallet_subtensor::Pallet::::set_coldkey_auto_stake_hotkey( + RawOrigin::Signed(real_coldkey).into(), + netuid, + hotkey, + ); + + match call_result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } } } } diff --git a/chain-extensions/src/types.rs b/chain-extensions/src/types.rs index ee6298ad5b..a32c186d5b 100644 --- a/chain-extensions/src/types.rs +++ b/chain-extensions/src/types.rs @@ -21,6 +21,19 @@ pub enum FunctionId { AddProxyV1 = 13, RemoveProxyV1 = 14, GetAlphaPriceV1 = 15, + // Proxy-aware staking functions (caller acts on behalf of real_coldkey) + AddStakeAsProxyV1 = 16, + RemoveStakeAsProxyV1 = 17, + UnstakeAllAsProxyV1 = 18, + UnstakeAllAlphaAsProxyV1 = 19, + MoveStakeAsProxyV1 = 20, + TransferStakeAsProxyV1 = 21, + SwapStakeAsProxyV1 = 22, + AddStakeLimitAsProxyV1 = 23, + RemoveStakeLimitAsProxyV1 = 24, + SwapStakeLimitAsProxyV1 = 25, + RemoveStakeFullLimitAsProxyV1 = 26, + SetColdkeyAutoStakeHotkeyAsProxyV1 = 27, } #[derive(PartialEq, Eq, Copy, Clone, Encode, Decode, Debug)] @@ -66,6 +79,8 @@ pub enum Output { ProxyNoSelfProxy = 18, /// Proxy relationship not found ProxyNotFound = 19, + /// Caller is not an authorized proxy for the specified account + NotAuthorizedProxy = 20, } impl From for Output { @@ -93,6 +108,7 @@ impl From for Output { Some("Duplicate") => Output::ProxyDuplicate, Some("NoSelfProxy") => Output::ProxyNoSelfProxy, Some("NotFound") => Output::ProxyNotFound, + Some("NotProxy") => Output::NotAuthorizedProxy, _ => Output::RuntimeError, } } From cb7470d8a6aafa3c5fc12bcb7ec59e44c773c460 Mon Sep 17 00:00:00 2001 From: Landyn Date: Wed, 28 Jan 2026 20:53:26 -0600 Subject: [PATCH 2/6] Add V2 staking extensions supporting delegated proxy calls for state-mutating staking functions --- chain-extensions/src/lib.rs | 30 +- chain-extensions/src/tests.rs | 727 ++++++++++++++++++++++++++++++++++ chain-extensions/src/types.rs | 26 +- docs/wasm-contracts.md | 25 ++ 4 files changed, 780 insertions(+), 28 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 53aed0f1f3..fddbf56a3f 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -531,11 +531,11 @@ where Ok(RetVal::Converging(Output::Success as u32)) } // ============================================================ - // Proxy-aware staking functions - // These allow a contract to act on behalf of a real_coldkey - // if the contract has been granted appropriate proxy permissions. + // V2 staking functions + // These accept an explicit coldkey parameter. If the caller differs + // from the coldkey, proxy permissions are verified before execution. // ============================================================ - FunctionId::AddStakeAsProxyV1 => { + FunctionId::AddStakeV2 => { let weight = Weight::from_parts(340_800_000, 0) .saturating_add(T::DbWeight::get().reads(25_u64)) .saturating_add(T::DbWeight::get().writes(15)); @@ -572,7 +572,7 @@ where } } } - FunctionId::RemoveStakeAsProxyV1 => { + FunctionId::RemoveStakeV2 => { let weight = Weight::from_parts(196_800_000, 0) .saturating_add(T::DbWeight::get().reads(20)) .saturating_add(T::DbWeight::get().writes(10)); @@ -609,7 +609,7 @@ where } } } - FunctionId::UnstakeAllAsProxyV1 => { + FunctionId::UnstakeAllV2 => { let weight = Weight::from_parts(28_830_000, 0) .saturating_add(T::DbWeight::get().reads(7)) .saturating_add(T::DbWeight::get().writes(0)); @@ -639,7 +639,7 @@ where } } } - FunctionId::UnstakeAllAlphaAsProxyV1 => { + FunctionId::UnstakeAllAlphaV2 => { let weight = Weight::from_parts(358_500_000, 0) .saturating_add(T::DbWeight::get().reads(37_u64)) .saturating_add(T::DbWeight::get().writes(21_u64)); @@ -669,7 +669,7 @@ where } } } - FunctionId::MoveStakeAsProxyV1 => { + FunctionId::MoveStakeV2 => { let weight = Weight::from_parts(164_300_000, 0) .saturating_add(T::DbWeight::get().reads(16_u64)) .saturating_add(T::DbWeight::get().writes(7_u64)); @@ -710,7 +710,7 @@ where } } } - FunctionId::TransferStakeAsProxyV1 => { + FunctionId::TransferStakeV2 => { let weight = Weight::from_parts(160_300_000, 0) .saturating_add(T::DbWeight::get().reads(14_u64)) .saturating_add(T::DbWeight::get().writes(6_u64)); @@ -751,7 +751,7 @@ where } } } - FunctionId::SwapStakeAsProxyV1 => { + FunctionId::SwapStakeV2 => { let weight = Weight::from_parts(351_300_000, 0) .saturating_add(T::DbWeight::get().reads(36_u64)) .saturating_add(T::DbWeight::get().writes(22_u64)); @@ -790,7 +790,7 @@ where } } } - FunctionId::AddStakeLimitAsProxyV1 => { + FunctionId::AddStakeLimitV2 => { let weight = Weight::from_parts(402_900_000, 0) .saturating_add(T::DbWeight::get().reads(25_u64)) .saturating_add(T::DbWeight::get().writes(15)); @@ -831,7 +831,7 @@ where } } } - FunctionId::RemoveStakeLimitAsProxyV1 => { + FunctionId::RemoveStakeLimitV2 => { let weight = Weight::from_parts(377_400_000, 0) .saturating_add(T::DbWeight::get().reads(29_u64)) .saturating_add(T::DbWeight::get().writes(14)); @@ -872,7 +872,7 @@ where } } } - FunctionId::SwapStakeLimitAsProxyV1 => { + FunctionId::SwapStakeLimitV2 => { let weight = Weight::from_parts(411_500_000, 0) .saturating_add(T::DbWeight::get().reads(36_u64)) .saturating_add(T::DbWeight::get().writes(22_u64)); @@ -923,7 +923,7 @@ where } } } - FunctionId::RemoveStakeFullLimitAsProxyV1 => { + FunctionId::RemoveStakeFullLimitV2 => { let weight = Weight::from_parts(395_300_000, 0) .saturating_add(T::DbWeight::get().reads(29_u64)) .saturating_add(T::DbWeight::get().writes(14_u64)); @@ -960,7 +960,7 @@ where } } } - FunctionId::SetColdkeyAutoStakeHotkeyAsProxyV1 => { + FunctionId::SetColdkeyAutoStakeHotkeyV2 => { let weight = Weight::from_parts(29_930_000, 0) .saturating_add(T::DbWeight::get().reads(5_u64)) .saturating_add(T::DbWeight::get().writes(2_u64)); diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index bd6f46c8ab..b568af1032 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -966,6 +966,733 @@ fn unstake_all_success_unstakes_balance() { }); } +// ============================================================ +// V2 function tests (proxy-aware staking functions) +// ============================================================ + +#[test] +fn add_stake_v2_self_call_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(1); + let owner_coldkey = U256::from(2); + let coldkey = U256::from(101); + let hotkey = U256::from(202); + let min_stake = DefaultMinStake::::get(); + let amount_raw = min_stake.to_u64().saturating_mul(10); + let amount: TaoCurrency = amount_raw.into(); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + (amount_raw * 1_000_000).into(), + AlphaCurrency::from(amount_raw * 10_000_000), + ); + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, amount_raw, + ); + + assert!( + pallet_subtensor::Pallet::::get_total_stake_for_hotkey(&hotkey).is_zero() + ); + + let expected_weight = Weight::from_parts(340_800_000, 0) + .saturating_add(::DbWeight::get().reads(25)) + .saturating_add(::DbWeight::get().writes(15)); + + let mut env = MockEnv::new( + FunctionId::AddStakeV2, + coldkey, + (coldkey, hotkey, netuid, amount).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + + assert_success(ret); + assert_eq!(env.charged_weight(), Some(expected_weight)); + + let total_stake = + pallet_subtensor::Pallet::::get_total_stake_for_hotkey(&hotkey); + assert!(total_stake > TaoCurrency::ZERO); + }); +} + +#[test] +fn add_stake_v2_with_proxy_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(1); + let owner_coldkey = U256::from(2); + let real_coldkey = U256::from(101); + let proxy_contract = U256::from(102); + let hotkey = U256::from(202); + let min_stake = DefaultMinStake::::get(); + let amount_raw = min_stake.to_u64().saturating_mul(10); + let amount: TaoCurrency = amount_raw.into(); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + (amount_raw * 1_000_000).into(), + AlphaCurrency::from(amount_raw * 10_000_000), + ); + mock::register_ok_neuron(netuid, hotkey, real_coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &real_coldkey, + amount_raw + 1_000_000_000, + ); + + // Add proxy relationship: real_coldkey grants Staking proxy to proxy_contract + assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( + RawOrigin::Signed(real_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Staking, + 0u64, + )); + + let expected_weight = Weight::from_parts(340_800_000, 0) + .saturating_add(::DbWeight::get().reads(25)) + .saturating_add(::DbWeight::get().writes(15)); + + // proxy_contract calls AddStakeV2 on behalf of real_coldkey + let mut env = MockEnv::new( + FunctionId::AddStakeV2, + proxy_contract, + (real_coldkey, hotkey, netuid, amount).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + + assert_success(ret); + + let total_stake = + pallet_subtensor::Pallet::::get_total_stake_for_hotkey(&hotkey); + assert!(total_stake > TaoCurrency::ZERO); + }); +} + +#[test] +fn add_stake_v2_without_proxy_fails() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(1); + let owner_coldkey = U256::from(2); + let real_coldkey = U256::from(101); + let unauthorized_caller = U256::from(102); + let hotkey = U256::from(202); + let min_stake = DefaultMinStake::::get(); + let amount_raw = min_stake.to_u64().saturating_mul(10); + let amount: TaoCurrency = amount_raw.into(); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + (amount_raw * 1_000_000).into(), + AlphaCurrency::from(amount_raw * 10_000_000), + ); + mock::register_ok_neuron(netuid, hotkey, real_coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &real_coldkey, + amount_raw + 1_000_000_000, + ); + + // No proxy relationship established + + let expected_weight = Weight::from_parts(340_800_000, 0) + .saturating_add(::DbWeight::get().reads(25)) + .saturating_add(::DbWeight::get().writes(15)); + + // unauthorized_caller tries to act on behalf of real_coldkey without proxy + let mut env = MockEnv::new( + FunctionId::AddStakeV2, + unauthorized_caller, + (real_coldkey, hotkey, netuid, amount).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env); + + // Should fail with NotProxy error + assert!(ret.is_err()); + }); +} + +#[test] +fn remove_stake_v2_with_proxy_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(4601); + let owner_coldkey = U256::from(4602); + let real_coldkey = U256::from(5601); + let proxy_contract = U256::from(5603); + let hotkey = U256::from(5602); + let stake_amount_raw: u64 = 320_000_000_000; + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoCurrency::from(120_000_000_000), + AlphaCurrency::from(100_000_000_000), + ); + + mock::register_ok_neuron(netuid, hotkey, real_coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &real_coldkey, + stake_amount_raw + 1_000_000_000, + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + mock::remove_stake_rate_limit_for_tests(&hotkey, &real_coldkey, netuid); + + // Add proxy relationship + assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( + RawOrigin::Signed(real_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Staking, + 0u64, + )); + + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &real_coldkey, netuid, + ); + + let alpha_to_unstake: AlphaCurrency = (alpha_before.to_u64() / 2).into(); + + let expected_weight = Weight::from_parts(196_800_000, 0) + .saturating_add(::DbWeight::get().reads(20)) + .saturating_add(::DbWeight::get().writes(10)); + + let mut env = MockEnv::new( + FunctionId::RemoveStakeV2, + proxy_contract, + (real_coldkey, hotkey, netuid, alpha_to_unstake).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &real_coldkey, netuid, + ); + + assert!(alpha_after < alpha_before); + }); +} + +#[test] +fn swap_stake_v2_with_proxy_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey_a = U256::from(4401); + let owner_coldkey_a = U256::from(4402); + let owner_hotkey_b = U256::from(4403); + let owner_coldkey_b = U256::from(4404); + let real_coldkey = U256::from(5401); + let proxy_contract = U256::from(5403); + let hotkey = U256::from(5402); + + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(260); + + let netuid_a = mock::add_dynamic_network(&owner_hotkey_a, &owner_coldkey_a); + let netuid_b = mock::add_dynamic_network(&owner_hotkey_b, &owner_coldkey_b); + + mock::setup_reserves( + netuid_a, + stake_amount_raw.saturating_mul(18).into(), + AlphaCurrency::from(stake_amount_raw.saturating_mul(30)), + ); + mock::setup_reserves( + netuid_b, + stake_amount_raw.saturating_mul(20).into(), + AlphaCurrency::from(stake_amount_raw.saturating_mul(28)), + ); + + mock::register_ok_neuron(netuid_a, hotkey, real_coldkey, 0); + mock::register_ok_neuron(netuid_b, hotkey, real_coldkey, 1); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &real_coldkey, + stake_amount_raw + 1_000_000_000, + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(real_coldkey).into(), + hotkey, + netuid_a, + stake_amount_raw.into(), + )); + + mock::remove_stake_rate_limit_for_tests(&hotkey, &real_coldkey, netuid_a); + + // Add proxy relationship + assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( + RawOrigin::Signed(real_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Staking, + 0u64, + )); + + let alpha_origin_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &real_coldkey, netuid_a, + ); + let alpha_destination_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &real_coldkey, netuid_b, + ); + let alpha_to_swap: AlphaCurrency = (alpha_origin_before.to_u64() / 3).into(); + + let expected_weight = Weight::from_parts(351_300_000, 0) + .saturating_add(::DbWeight::get().reads(36)) + .saturating_add(::DbWeight::get().writes(22)); + + let mut env = MockEnv::new( + FunctionId::SwapStakeV2, + proxy_contract, + (real_coldkey, hotkey, netuid_a, netuid_b, alpha_to_swap).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let alpha_origin_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &real_coldkey, netuid_a, + ); + let alpha_destination_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &real_coldkey, netuid_b, + ); + + assert!(alpha_origin_after < alpha_origin_before); + assert!(alpha_destination_after > alpha_destination_before); + }); +} + +#[test] +fn move_stake_v2_with_proxy_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(4201); + let owner_coldkey = U256::from(4202); + let real_coldkey = U256::from(5201); + let proxy_contract = U256::from(5204); + let origin_hotkey = U256::from(5202); + let destination_hotkey = U256::from(5203); + + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(240); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + stake_amount_raw.saturating_mul(15).into(), + AlphaCurrency::from(stake_amount_raw.saturating_mul(25)), + ); + + mock::register_ok_neuron(netuid, origin_hotkey, real_coldkey, 0); + mock::register_ok_neuron(netuid, destination_hotkey, real_coldkey, 1); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &real_coldkey, + stake_amount_raw + 1_000_000_000, + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(real_coldkey).into(), + origin_hotkey, + netuid, + stake_amount_raw.into(), + )); + + mock::remove_stake_rate_limit_for_tests(&origin_hotkey, &real_coldkey, netuid); + + // Add proxy relationship + assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( + RawOrigin::Signed(real_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Staking, + 0u64, + )); + + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &origin_hotkey, + &real_coldkey, + netuid, + ); + let alpha_to_move: AlphaCurrency = (alpha_before.to_u64() / 2).into(); + + let expected_weight = Weight::from_parts(164_300_000, 0) + .saturating_add(::DbWeight::get().reads(16)) + .saturating_add(::DbWeight::get().writes(7)); + + let mut env = MockEnv::new( + FunctionId::MoveStakeV2, + proxy_contract, + ( + real_coldkey, + origin_hotkey, + destination_hotkey, + netuid, + netuid, + alpha_to_move, + ) + .encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let origin_alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &origin_hotkey, + &real_coldkey, + netuid, + ); + let destination_alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &destination_hotkey, + &real_coldkey, + netuid, + ); + + assert_eq!(origin_alpha_after, alpha_before - alpha_to_move); + assert_eq!(destination_alpha_after, alpha_to_move); + }); +} + +#[test] +fn transfer_stake_v2_requires_transfer_proxy() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(4301); + let owner_coldkey = U256::from(4302); + let origin_coldkey = U256::from(5301); + let destination_coldkey = U256::from(5302); + let proxy_contract = U256::from(5304); + let hotkey = U256::from(5303); + + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(250); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + stake_amount_raw.saturating_mul(15).into(), + AlphaCurrency::from(stake_amount_raw.saturating_mul(25)), + ); + + mock::register_ok_neuron(netuid, hotkey, origin_coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &origin_coldkey, + stake_amount_raw + 1_000_000_000, + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(origin_coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + mock::remove_stake_rate_limit_for_tests(&hotkey, &origin_coldkey, netuid); + + // Add Staking proxy (wrong type for transfer_stake) + assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( + RawOrigin::Signed(origin_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Staking, + 0u64, + )); + + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &origin_coldkey, + netuid, + ); + let alpha_to_transfer: AlphaCurrency = (alpha_before.to_u64() / 3).into(); + + let expected_weight = Weight::from_parts(160_300_000, 0) + .saturating_add(::DbWeight::get().reads(14)) + .saturating_add(::DbWeight::get().writes(6)); + + // First try with Staking proxy - should fail + let mut env = MockEnv::new( + FunctionId::TransferStakeV2, + proxy_contract, + ( + origin_coldkey, + destination_coldkey, + hotkey, + netuid, + netuid, + alpha_to_transfer, + ) + .encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env); + assert!(ret.is_err(), "Staking proxy should not work for transfer_stake"); + + // Remove Staking proxy, add Transfer proxy + assert_ok!(pallet_subtensor_proxy::Pallet::::remove_proxy( + RawOrigin::Signed(origin_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Staking, + 0u64, + )); + assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( + RawOrigin::Signed(origin_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Transfer, + 0u64, + )); + + // Now try with Transfer proxy - should succeed + let mut env = MockEnv::new( + FunctionId::TransferStakeV2, + proxy_contract, + ( + origin_coldkey, + destination_coldkey, + hotkey, + netuid, + netuid, + alpha_to_transfer, + ) + .encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let origin_alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &origin_coldkey, + netuid, + ); + let destination_alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &destination_coldkey, + netuid, + ); + + assert_eq!(origin_alpha_after, alpha_before - alpha_to_transfer); + assert_eq!(destination_alpha_after, alpha_to_transfer); + }); +} + +#[test] +fn unstake_all_v2_self_call_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(4001); + let owner_coldkey = U256::from(4002); + let coldkey = U256::from(5001); + let hotkey = U256::from(5002); + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(200); + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + + mock::setup_reserves( + netuid, + stake_amount_raw.saturating_mul(10).into(), + AlphaCurrency::from(stake_amount_raw.saturating_mul(20)), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + stake_amount_raw + 1_000_000_000, + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + mock::remove_stake_rate_limit_for_tests(&hotkey, &coldkey, netuid); + + let expected_weight = Weight::from_parts(28_830_000, 0) + .saturating_add(::DbWeight::get().reads(7)) + .saturating_add(::DbWeight::get().writes(0)); + + let pre_balance = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); + + let mut env = MockEnv::new(FunctionId::UnstakeAllV2, coldkey, (coldkey, hotkey).encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + assert_eq!(env.charged_weight(), Some(expected_weight)); + + let remaining_alpha = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(remaining_alpha <= AlphaCurrency::from(1_000)); + + let post_balance = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); + assert!(post_balance > pre_balance); + }); +} + +#[test] +fn unstake_all_alpha_v2_self_call_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(4101); + let owner_coldkey = U256::from(4102); + let coldkey = U256::from(5101); + let hotkey = U256::from(5102); + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(220); + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + + mock::setup_reserves( + netuid, + stake_amount_raw.saturating_mul(20).into(), + AlphaCurrency::from(stake_amount_raw.saturating_mul(30)), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + stake_amount_raw + 1_000_000_000, + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + mock::remove_stake_rate_limit_for_tests(&hotkey, &coldkey, netuid); + + let expected_weight = Weight::from_parts(358_500_000, 0) + .saturating_add(::DbWeight::get().reads(37)) + .saturating_add(::DbWeight::get().writes(21)); + + let mut env = + MockEnv::new(FunctionId::UnstakeAllAlphaV2, coldkey, (coldkey, hotkey).encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + assert_eq!(env.charged_weight(), Some(expected_weight)); + + let subnet_alpha = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(subnet_alpha <= AlphaCurrency::from(1_000)); + + let root_alpha = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &coldkey, + NetUid::ROOT, + ); + assert!(root_alpha > AlphaCurrency::ZERO); + }); +} + +#[test] +fn set_coldkey_auto_stake_hotkey_v2_self_call_succeeds() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(4901); + let owner_coldkey = U256::from(4902); + let coldkey = U256::from(5901); + let hotkey = U256::from(5902); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + + pallet_subtensor::Owner::::insert(hotkey, coldkey); + pallet_subtensor::OwnedHotkeys::::insert(coldkey, vec![hotkey]); + pallet_subtensor::Uids::::insert(netuid, hotkey, 0u16); + + assert_eq!( + pallet_subtensor::AutoStakeDestination::::get(coldkey, netuid), + None + ); + + let expected_weight = Weight::from_parts(29_930_000, 0) + .saturating_add(::DbWeight::get().reads(5)) + .saturating_add(::DbWeight::get().writes(2)); + + let mut env = MockEnv::new( + FunctionId::SetColdkeyAutoStakeHotkeyV2, + coldkey, + (coldkey, netuid, hotkey).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + assert_eq!(env.charged_weight(), Some(expected_weight)); + + assert_eq!( + pallet_subtensor::AutoStakeDestination::::get(coldkey, netuid), + Some(hotkey) + ); + }); +} + +#[test] +fn remove_stake_v2_self_call_returns_error_with_no_stake() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(1); + let owner_coldkey = U256::from(2); + let coldkey = U256::from(301); + let hotkey = U256::from(302); + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + let min_stake = DefaultMinStake::::get(); + let amount: AlphaCurrency = AlphaCurrency::from(min_stake.to_u64()); + + let expected_weight = Weight::from_parts(196_800_000, 0) + .saturating_add(::DbWeight::get().reads(20)) + .saturating_add(::DbWeight::get().writes(10)); + + let mut env = MockEnv::new( + FunctionId::RemoveStakeV2, + coldkey, + (coldkey, hotkey, netuid, amount).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + + match ret { + RetVal::Converging(code) => { + assert_eq!(code, Output::AmountTooLow as u32, "mismatched error output") + } + _ => panic!("unexpected return value"), + } + }); +} + #[test] fn get_alpha_price_returns_encoded_price() { mock::new_test_ext(1).execute_with(|| { diff --git a/chain-extensions/src/types.rs b/chain-extensions/src/types.rs index a32c186d5b..b0fb53ca79 100644 --- a/chain-extensions/src/types.rs +++ b/chain-extensions/src/types.rs @@ -21,19 +21,19 @@ pub enum FunctionId { AddProxyV1 = 13, RemoveProxyV1 = 14, GetAlphaPriceV1 = 15, - // Proxy-aware staking functions (caller acts on behalf of real_coldkey) - AddStakeAsProxyV1 = 16, - RemoveStakeAsProxyV1 = 17, - UnstakeAllAsProxyV1 = 18, - UnstakeAllAlphaAsProxyV1 = 19, - MoveStakeAsProxyV1 = 20, - TransferStakeAsProxyV1 = 21, - SwapStakeAsProxyV1 = 22, - AddStakeLimitAsProxyV1 = 23, - RemoveStakeLimitAsProxyV1 = 24, - SwapStakeLimitAsProxyV1 = 25, - RemoveStakeFullLimitAsProxyV1 = 26, - SetColdkeyAutoStakeHotkeyAsProxyV1 = 27, + // V2: Accept explicit coldkey with proxy verification + AddStakeV2 = 16, + RemoveStakeV2 = 17, + UnstakeAllV2 = 18, + UnstakeAllAlphaV2 = 19, + MoveStakeV2 = 20, + TransferStakeV2 = 21, + SwapStakeV2 = 22, + AddStakeLimitV2 = 23, + RemoveStakeLimitV2 = 24, + SwapStakeLimitV2 = 25, + RemoveStakeFullLimitV2 = 26, + SetColdkeyAutoStakeHotkeyV2 = 27, } #[derive(PartialEq, Eq, Copy, Clone, Encode, Decode, Debug)] diff --git a/docs/wasm-contracts.md b/docs/wasm-contracts.md index ed6e9ecdd3..fd7b9a9abf 100644 --- a/docs/wasm-contracts.md +++ b/docs/wasm-contracts.md @@ -43,6 +43,19 @@ Subtensor provides a custom chain extension that allows smart contracts to inter | 12 | `set_coldkey_auto_stake_hotkey` | Configure automatic stake destination | `(NetUid, AccountId)` | Error code | | 13 | `add_proxy` | Add a staking proxy for the caller | `(AccountId)` | Error code | | 14 | `remove_proxy` | Remove a staking proxy for the caller | `(AccountId)` | Error code | +| 15 | `get_alpha_price` | Get the current alpha price for a subnet | `(NetUid)` | `u64` (price × 10⁹) | +| 16 | `add_stake_v2` | Add stake with explicit coldkey (proxy-aware) | `(AccountId, AccountId, NetUid, TaoCurrency)` | Error code | +| 17 | `remove_stake_v2` | Remove stake with explicit coldkey | `(AccountId, AccountId, NetUid, AlphaCurrency)` | Error code | +| 18 | `unstake_all_v2` | Unstake all TAO with explicit coldkey | `(AccountId, AccountId)` | Error code | +| 19 | `unstake_all_alpha_v2` | Unstake all Alpha with explicit coldkey | `(AccountId, AccountId)` | Error code | +| 20 | `move_stake_v2` | Move stake between hotkeys with explicit coldkey | `(AccountId, AccountId, AccountId, NetUid, NetUid, AlphaCurrency)` | Error code | +| 21 | `transfer_stake_v2` | Transfer stake between coldkeys (requires Transfer proxy) | `(AccountId, AccountId, AccountId, NetUid, NetUid, AlphaCurrency)` | Error code | +| 22 | `swap_stake_v2` | Swap stake between subnets with explicit coldkey | `(AccountId, AccountId, NetUid, NetUid, AlphaCurrency)` | Error code | +| 23 | `add_stake_limit_v2` | Add stake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, TaoCurrency, TaoCurrency, bool)` | Error code | +| 24 | `remove_stake_limit_v2` | Remove stake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, AlphaCurrency, TaoCurrency, bool)` | Error code | +| 25 | `swap_stake_limit_v2` | Swap stake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, NetUid, AlphaCurrency, TaoCurrency, bool)` | Error code | +| 26 | `remove_stake_full_limit_v2` | Full unstake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, Option)` | Error code | +| 27 | `set_coldkey_auto_stake_hotkey_v2` | Set auto-stake hotkey with explicit coldkey | `(AccountId, NetUid, AccountId)` | Error code | Example usage in your ink! contract: ```rust @@ -85,6 +98,18 @@ Chain extension functions that modify state return error codes as `u32` values. | 17 | `ProxyDuplicate` | Proxy already exists | | 18 | `ProxyNoSelfProxy` | Cannot add self as proxy | | 19 | `ProxyNotFound` | Proxy relationship not found | +| 20 | `NotAuthorizedProxy` | Caller is not an authorized proxy for the account | + +#### V2 Functions (Proxy-Aware) + +Functions 16-27 are V2 versions that accept an explicit `coldkey` parameter as the first argument. These functions: + +- If `coldkey == caller`: Execute directly (no proxy check needed) +- If `coldkey != caller`: Verify caller has appropriate proxy permissions for coldkey + +**Proxy Types Required:** +- Most V2 functions require `ProxyType::Staking` +- `transfer_stake_v2` (ID 21) requires `ProxyType::Transfer` ### Call Filter From 5acabf33acfbc7e962932f3f06e2b9c9508aa979 Mon Sep 17 00:00:00 2001 From: Landyn Date: Thu, 29 Jan 2026 11:54:41 -0600 Subject: [PATCH 3/6] cargo fmt --- chain-extensions/src/lib.rs | 114 ++++++++++++++++++++++++++-------- chain-extensions/src/tests.rs | 46 ++++++++++---- 2 files changed, 122 insertions(+), 38 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index fddbf56a3f..a087de6c87 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -553,8 +553,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::add_stake( @@ -590,8 +594,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::remove_stake( @@ -622,8 +630,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::unstake_all( @@ -652,8 +664,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::unstake_all_alpha( @@ -683,14 +699,25 @@ where origin_netuid, destination_netuid, alpha_amount, - ): (T::AccountId, T::AccountId, T::AccountId, NetUid, NetUid, AlphaCurrency) = env + ): ( + T::AccountId, + T::AccountId, + T::AccountId, + NetUid, + NetUid, + AlphaCurrency, + ) = env .read_as() .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::move_stake( @@ -717,7 +744,14 @@ where env.charge_weight(weight)?; - let (real_coldkey, destination_coldkey, hotkey, origin_netuid, destination_netuid, alpha_amount): ( + let ( + real_coldkey, + destination_coldkey, + hotkey, + origin_netuid, + destination_netuid, + alpha_amount, + ): ( T::AccountId, T::AccountId, T::AccountId, @@ -730,8 +764,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Transfer)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Transfer), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::transfer_stake( @@ -770,8 +808,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::swap_stake( @@ -810,8 +852,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::add_stake_limit( @@ -851,8 +897,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::remove_stake_limit( @@ -901,8 +951,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::swap_stake_limit( @@ -941,8 +995,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::remove_stake_full_limit( @@ -973,8 +1031,12 @@ where let caller = env.caller(); if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy(&real_coldkey, &caller, Some(ProxyType::Staking)) - .map_err(|_| DispatchError::Other("NotProxy"))?; + pallet_proxy::Pallet::::find_proxy( + &real_coldkey, + &caller, + Some(ProxyType::Staking), + ) + .map_err(|_| DispatchError::Other("NotProxy"))?; } let call_result = pallet_subtensor::Pallet::::set_coldkey_auto_stake_hotkey( diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index b568af1032..a7d007a1b0 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1163,7 +1163,9 @@ fn remove_stake_v2_with_proxy_succeeds() { let alpha_before = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &real_coldkey, netuid, + &hotkey, + &real_coldkey, + netuid, ); let alpha_to_unstake: AlphaCurrency = (alpha_before.to_u64() / 2).into(); @@ -1184,7 +1186,9 @@ fn remove_stake_v2_with_proxy_succeeds() { let alpha_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &real_coldkey, netuid, + &hotkey, + &real_coldkey, + netuid, ); assert!(alpha_after < alpha_before); @@ -1246,11 +1250,15 @@ fn swap_stake_v2_with_proxy_succeeds() { let alpha_origin_before = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &real_coldkey, netuid_a, + &hotkey, + &real_coldkey, + netuid_a, ); let alpha_destination_before = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &real_coldkey, netuid_b, + &hotkey, + &real_coldkey, + netuid_b, ); let alpha_to_swap: AlphaCurrency = (alpha_origin_before.to_u64() / 3).into(); @@ -1270,11 +1278,15 @@ fn swap_stake_v2_with_proxy_succeeds() { let alpha_origin_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &real_coldkey, netuid_a, + &hotkey, + &real_coldkey, + netuid_a, ); let alpha_destination_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &real_coldkey, netuid_b, + &hotkey, + &real_coldkey, + netuid_b, ); assert!(alpha_origin_after < alpha_origin_before); @@ -1448,7 +1460,10 @@ fn transfer_stake_v2_requires_transfer_proxy() { .with_expected_weight(expected_weight); let ret = SubtensorChainExtension::::dispatch(&mut env); - assert!(ret.is_err(), "Staking proxy should not work for transfer_stake"); + assert!( + ret.is_err(), + "Staking proxy should not work for transfer_stake" + ); // Remove Staking proxy, add Transfer proxy assert_ok!(pallet_subtensor_proxy::Pallet::::remove_proxy( @@ -1539,8 +1554,12 @@ fn unstake_all_v2_self_call_succeeds() { let pre_balance = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); - let mut env = MockEnv::new(FunctionId::UnstakeAllV2, coldkey, (coldkey, hotkey).encode()) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new( + FunctionId::UnstakeAllV2, + coldkey, + (coldkey, hotkey).encode(), + ) + .with_expected_weight(expected_weight); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); assert_success(ret); @@ -1593,9 +1612,12 @@ fn unstake_all_alpha_v2_self_call_succeeds() { .saturating_add(::DbWeight::get().reads(37)) .saturating_add(::DbWeight::get().writes(21)); - let mut env = - MockEnv::new(FunctionId::UnstakeAllAlphaV2, coldkey, (coldkey, hotkey).encode()) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new( + FunctionId::UnstakeAllAlphaV2, + coldkey, + (coldkey, hotkey).encode(), + ) + .with_expected_weight(expected_weight); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); assert_success(ret); From c7936f8096caf31bc312ff895f6dc734a9cdccd7 Mon Sep 17 00:00:00 2001 From: Landyn Date: Thu, 29 Jan 2026 12:19:43 -0600 Subject: [PATCH 4/6] More specific error type checking in tests --- chain-extensions/src/tests.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index a7d007a1b0..8c31aaca44 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1116,7 +1116,7 @@ fn add_stake_v2_without_proxy_fails() { let ret = SubtensorChainExtension::::dispatch(&mut env); // Should fail with NotProxy error - assert!(ret.is_err()); + assert!(matches!(ret, Err(DispatchError::Other("NotProxy")))); }); } @@ -1460,10 +1460,8 @@ fn transfer_stake_v2_requires_transfer_proxy() { .with_expected_weight(expected_weight); let ret = SubtensorChainExtension::::dispatch(&mut env); - assert!( - ret.is_err(), - "Staking proxy should not work for transfer_stake" - ); + // Staking proxy should not work for transfer_stake - requires Transfer proxy + assert!(matches!(ret, Err(DispatchError::Other("NotProxy")))); // Remove Staking proxy, add Transfer proxy assert_ok!(pallet_subtensor_proxy::Pallet::::remove_proxy( From d24d951940773da6261c16c7c389bfa210fac63c Mon Sep 17 00:00:00 2001 From: Landyn Date: Wed, 11 Feb 2026 20:30:50 -0600 Subject: [PATCH 5/6] Replace V2 chain extensions with single ProxyCall dispatcher --- chain-extensions/src/lib.rs | 576 ++++------------------------- chain-extensions/src/mock.rs | 8 +- chain-extensions/src/tests.rs | 661 +++++++++++----------------------- chain-extensions/src/types.rs | 15 +- 4 files changed, 282 insertions(+), 978 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index a087de6c87..943f534d4b 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -8,15 +8,23 @@ mod tests; pub mod types; use crate::types::{FunctionId, Output}; -use codec::{Decode, Encode, MaxEncodedLen}; -use frame_support::{DebugNoBound, traits::Get}; +use codec::{Decode, DecodeLimit, Encode, MaxEncodedLen}; +use frame_support::{ + BoundedVec, DebugNoBound, + dispatch::GetDispatchInfo, + traits::{ConstU32, Get}, +}; use frame_system::RawOrigin; use pallet_contracts::chain_extension::{ BufInBufOutState, ChainExtension, Environment, Ext, InitState, RetVal, SysConfig, }; use pallet_subtensor_proxy as pallet_proxy; use pallet_subtensor_proxy::WeightInfo; -use sp_runtime::{DispatchError, Weight, traits::StaticLookup}; +use sp_runtime::{ + DispatchError, Weight, + traits::{Dispatchable, StaticLookup}, +}; +use sp_std::boxed::Box; use sp_std::marker::PhantomData; use substrate_fixed::types::U96F32; use subtensor_runtime_common::{AlphaCurrency, NetUid, ProxyType, TaoCurrency}; @@ -530,526 +538,66 @@ where Ok(RetVal::Converging(Output::Success as u32)) } - // ============================================================ - // V2 staking functions - // These accept an explicit coldkey parameter. If the caller differs - // from the coldkey, proxy permissions are verified before execution. - // ============================================================ - FunctionId::AddStakeV2 => { - let weight = Weight::from_parts(340_800_000, 0) - .saturating_add(T::DbWeight::get().reads(25_u64)) - .saturating_add(T::DbWeight::get().writes(15)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey, netuid, amount_staked): ( - T::AccountId, - T::AccountId, - NetUid, - TaoCurrency, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::add_stake( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - netuid, - amount_staked, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::RemoveStakeV2 => { - let weight = Weight::from_parts(196_800_000, 0) - .saturating_add(T::DbWeight::get().reads(20)) - .saturating_add(T::DbWeight::get().writes(10)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey, netuid, amount_unstaked): ( - T::AccountId, - T::AccountId, - NetUid, - AlphaCurrency, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::remove_stake( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - netuid, - amount_unstaked, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::UnstakeAllV2 => { - let weight = Weight::from_parts(28_830_000, 0) - .saturating_add(T::DbWeight::get().reads(7)) - .saturating_add(T::DbWeight::get().writes(0)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey): (T::AccountId, T::AccountId) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::unstake_all( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::UnstakeAllAlphaV2 => { - let weight = Weight::from_parts(358_500_000, 0) - .saturating_add(T::DbWeight::get().reads(37_u64)) - .saturating_add(T::DbWeight::get().writes(21_u64)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey): (T::AccountId, T::AccountId) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::unstake_all_alpha( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::MoveStakeV2 => { - let weight = Weight::from_parts(164_300_000, 0) - .saturating_add(T::DbWeight::get().reads(16_u64)) - .saturating_add(T::DbWeight::get().writes(7_u64)); - - env.charge_weight(weight)?; - - let ( - real_coldkey, - origin_hotkey, - destination_hotkey, - origin_netuid, - destination_netuid, - alpha_amount, - ): ( - T::AccountId, - T::AccountId, - T::AccountId, - NetUid, - NetUid, - AlphaCurrency, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::move_stake( - RawOrigin::Signed(real_coldkey).into(), - origin_hotkey, - destination_hotkey, - origin_netuid, - destination_netuid, - alpha_amount, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::TransferStakeV2 => { - let weight = Weight::from_parts(160_300_000, 0) - .saturating_add(T::DbWeight::get().reads(14_u64)) - .saturating_add(T::DbWeight::get().writes(6_u64)); - - env.charge_weight(weight)?; - - let ( - real_coldkey, - destination_coldkey, - hotkey, - origin_netuid, - destination_netuid, - alpha_amount, - ): ( - T::AccountId, - T::AccountId, - T::AccountId, - NetUid, - NetUid, - AlphaCurrency, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Transfer), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::transfer_stake( - RawOrigin::Signed(real_coldkey).into(), - destination_coldkey, - hotkey, - origin_netuid, - destination_netuid, - alpha_amount, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::SwapStakeV2 => { - let weight = Weight::from_parts(351_300_000, 0) - .saturating_add(T::DbWeight::get().reads(36_u64)) - .saturating_add(T::DbWeight::get().writes(22_u64)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey, origin_netuid, destination_netuid, alpha_amount): ( - T::AccountId, - T::AccountId, - NetUid, - NetUid, - AlphaCurrency, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::swap_stake( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - origin_netuid, - destination_netuid, - alpha_amount, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::AddStakeLimitV2 => { - let weight = Weight::from_parts(402_900_000, 0) - .saturating_add(T::DbWeight::get().reads(25_u64)) - .saturating_add(T::DbWeight::get().writes(15)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey, netuid, amount_staked, limit_price, allow_partial): ( - T::AccountId, - T::AccountId, - NetUid, - TaoCurrency, - TaoCurrency, - bool, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::add_stake_limit( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - netuid, - amount_staked, - limit_price, - allow_partial, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::RemoveStakeLimitV2 => { - let weight = Weight::from_parts(377_400_000, 0) - .saturating_add(T::DbWeight::get().reads(29_u64)) - .saturating_add(T::DbWeight::get().writes(14)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey, netuid, amount_unstaked, limit_price, allow_partial): ( - T::AccountId, - T::AccountId, - NetUid, - AlphaCurrency, - TaoCurrency, - bool, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::remove_stake_limit( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - netuid, - amount_unstaked, - limit_price, - allow_partial, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::SwapStakeLimitV2 => { - let weight = Weight::from_parts(411_500_000, 0) - .saturating_add(T::DbWeight::get().reads(36_u64)) - .saturating_add(T::DbWeight::get().writes(22_u64)); - - env.charge_weight(weight)?; - - let ( - real_coldkey, - hotkey, - origin_netuid, - destination_netuid, - alpha_amount, - limit_price, - allow_partial, - ): ( - T::AccountId, + FunctionId::ProxyCall => { + // 1. Read params: (real_coldkey, force_proxy_type as Option, call_data) + let (real_coldkey, force_proxy_type_raw, call_data): ( T::AccountId, - NetUid, - NetUid, - AlphaCurrency, - TaoCurrency, - bool, + Option, + BoundedVec>, ) = env .read_as() .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::swap_stake_limit( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - origin_netuid, - destination_netuid, - alpha_amount, - limit_price, - allow_partial, + // 2. Parse proxy type from raw u8 + let force_proxy_type = force_proxy_type_raw + .map(ProxyType::try_from) + .transpose() + .map_err(|_| DispatchError::Other("Invalid proxy type"))?; + + // 3. Decode the inner RuntimeCall (depth-limited for safety) + let call = <::RuntimeCall>::decode_with_depth_limit( + 8, + &mut &call_data[..], + ) + .map_err(|_| DispatchError::Other("Failed to decode call"))?; + + // 4. Dynamic weight: inner call weight + proxy overhead + let call_weight = call.get_dispatch_info().call_weight; + let proxy_overhead = ::WeightInfo::proxy( + ::MaxProxies::get(), ); + env.charge_weight(call_weight.saturating_add(proxy_overhead))?; - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) - } - } - } - FunctionId::RemoveStakeFullLimitV2 => { - let weight = Weight::from_parts(395_300_000, 0) - .saturating_add(T::DbWeight::get().reads(29_u64)) - .saturating_add(T::DbWeight::get().writes(14_u64)); - - env.charge_weight(weight)?; - - let (real_coldkey, hotkey, netuid, limit_price): ( - T::AccountId, - T::AccountId, - NetUid, - Option, - ) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - + // 5. Execute let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::remove_stake_full_limit( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - netuid, - limit_price, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) + if real_coldkey == caller { + // Direct dispatch — no proxy needed + let result = call.dispatch(RawOrigin::Signed(caller).into()); + match result { + Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), + Err(e) => Ok(RetVal::Converging(Output::from(e.error) as u32)), } - } - } - FunctionId::SetColdkeyAutoStakeHotkeyV2 => { - let weight = Weight::from_parts(29_930_000, 0) - .saturating_add(T::DbWeight::get().reads(5_u64)) - .saturating_add(T::DbWeight::get().writes(2_u64)); - - env.charge_weight(weight)?; - - let (real_coldkey, netuid, hotkey): (T::AccountId, NetUid, T::AccountId) = env - .read_as() - .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; - - let caller = env.caller(); - if real_coldkey != caller { - pallet_proxy::Pallet::::find_proxy( - &real_coldkey, - &caller, - Some(ProxyType::Staking), - ) - .map_err(|_| DispatchError::Other("NotProxy"))?; - } - - let call_result = pallet_subtensor::Pallet::::set_coldkey_auto_stake_hotkey( - RawOrigin::Signed(real_coldkey).into(), - netuid, - hotkey, - ); - - match call_result { - Ok(_) => Ok(RetVal::Converging(Output::Success as u32)), - Err(e) => { - let error_code = Output::from(e) as u32; - Ok(RetVal::Converging(error_code)) + } else { + // Proxy dispatch — pallet handles all permission checking + let real_lookup = <::Lookup as StaticLookup>::Source::from( + real_coldkey.clone(), + ); + let result = pallet_proxy::Pallet::::proxy( + RawOrigin::Signed(caller).into(), + real_lookup, + force_proxy_type, + Box::new(call), + ); + match result { + Ok(()) => { + // Inner call result is in LastCallResult storage + match pallet_proxy::LastCallResult::::get(&real_coldkey) { + Some(Ok(())) => Ok(RetVal::Converging(Output::Success as u32)), + Some(Err(e)) => Ok(RetVal::Converging(Output::from(e) as u32)), + None => Ok(RetVal::Converging(Output::RuntimeError as u32)), + } + } + Err(e) => Ok(RetVal::Converging(Output::from(e) as u32)), } } } diff --git a/chain-extensions/src/mock.rs b/chain-extensions/src/mock.rs index 27ac5bc06e..69154ceb29 100644 --- a/chain-extensions/src/mock.rs +++ b/chain-extensions/src/mock.rs @@ -156,7 +156,13 @@ impl frame_support::traits::InstanceFilter for subtensor_runtime_co | RuntimeCall::SubtensorModule(pallet_subtensor::Call::swap_stake { .. }) | RuntimeCall::SubtensorModule(pallet_subtensor::Call::swap_stake_limit { .. }) | RuntimeCall::SubtensorModule(pallet_subtensor::Call::move_stake { .. }) - | RuntimeCall::SubtensorModule(pallet_subtensor::Call::transfer_stake { .. }) + | RuntimeCall::SubtensorModule( + pallet_subtensor::Call::set_coldkey_auto_stake_hotkey { .. } + ) + ), + subtensor_runtime_common::ProxyType::Transfer => matches!( + c, + RuntimeCall::SubtensorModule(pallet_subtensor::Call::transfer_stake { .. }) ), _ => false, } diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index 8c31aaca44..edf991d0f0 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -967,11 +967,22 @@ fn unstake_all_success_unstakes_balance() { } // ============================================================ -// V2 function tests (proxy-aware staking functions) +// ProxyCall tests (proxy-aware generic call dispatcher) // ============================================================ +/// Helper: encode an inner RuntimeCall into ProxyCall input bytes. +fn encode_proxy_call_input( + real_coldkey: AccountId, + force_proxy_type: Option, + inner_call: mock::RuntimeCall, +) -> Vec { + use frame_support::{BoundedVec, traits::ConstU32}; + let call_data: BoundedVec> = inner_call.encode().try_into().unwrap(); + (real_coldkey, force_proxy_type, call_data).encode() +} + #[test] -fn add_stake_v2_self_call_succeeds() { +fn proxy_call_self_call_add_stake_succeeds() { mock::new_test_ext(1).execute_with(|| { let owner_hotkey = U256::from(1); let owner_coldkey = U256::from(2); @@ -997,21 +1008,20 @@ fn add_stake_v2_self_call_succeeds() { pallet_subtensor::Pallet::::get_total_stake_for_hotkey(&hotkey).is_zero() ); - let expected_weight = Weight::from_parts(340_800_000, 0) - .saturating_add(::DbWeight::get().reads(25)) - .saturating_add(::DbWeight::get().writes(15)); + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::add_stake { + hotkey, + netuid, + amount_staked: amount, + } + .into(); + let input = encode_proxy_call_input(coldkey, None, inner_call); - let mut env = MockEnv::new( - FunctionId::AddStakeV2, - coldkey, - (coldkey, hotkey, netuid, amount).encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, coldkey, input); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); assert_success(ret); - assert_eq!(env.charged_weight(), Some(expected_weight)); + assert!(env.charged_weight().is_some()); let total_stake = pallet_subtensor::Pallet::::get_total_stake_for_hotkey(&hotkey); @@ -1020,7 +1030,7 @@ fn add_stake_v2_self_call_succeeds() { } #[test] -fn add_stake_v2_with_proxy_succeeds() { +fn proxy_call_with_staking_proxy_add_stake_succeeds() { mock::new_test_ext(1).execute_with(|| { let owner_hotkey = U256::from(1); let owner_coldkey = U256::from(2); @@ -1052,17 +1062,16 @@ fn add_stake_v2_with_proxy_succeeds() { 0u64, )); - let expected_weight = Weight::from_parts(340_800_000, 0) - .saturating_add(::DbWeight::get().reads(25)) - .saturating_add(::DbWeight::get().writes(15)); + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::add_stake { + hotkey, + netuid, + amount_staked: amount, + } + .into(); + let input = encode_proxy_call_input(real_coldkey, None, inner_call); - // proxy_contract calls AddStakeV2 on behalf of real_coldkey - let mut env = MockEnv::new( - FunctionId::AddStakeV2, - proxy_contract, - (real_coldkey, hotkey, netuid, amount).encode(), - ) - .with_expected_weight(expected_weight); + // proxy_contract calls ProxyCall on behalf of real_coldkey + let mut env = MockEnv::new(FunctionId::ProxyCall, proxy_contract, input); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); @@ -1075,12 +1084,12 @@ fn add_stake_v2_with_proxy_succeeds() { } #[test] -fn add_stake_v2_without_proxy_fails() { +fn proxy_call_with_any_proxy_add_stake_succeeds() { mock::new_test_ext(1).execute_with(|| { let owner_hotkey = U256::from(1); let owner_coldkey = U256::from(2); let real_coldkey = U256::from(101); - let unauthorized_caller = U256::from(102); + let proxy_contract = U256::from(102); let hotkey = U256::from(202); let min_stake = DefaultMinStake::::get(); let amount_raw = min_stake.to_u64().saturating_mul(10); @@ -1099,213 +1108,155 @@ fn add_stake_v2_without_proxy_fails() { amount_raw + 1_000_000_000, ); - // No proxy relationship established + // Add proxy relationship with ProxyType::Any (was broken in V2!) + assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( + RawOrigin::Signed(real_coldkey).into(), + proxy_contract, + subtensor_runtime_common::ProxyType::Any, + 0u64, + )); - let expected_weight = Weight::from_parts(340_800_000, 0) - .saturating_add(::DbWeight::get().reads(25)) - .saturating_add(::DbWeight::get().writes(15)); + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::add_stake { + hotkey, + netuid, + amount_staked: amount, + } + .into(); + let input = encode_proxy_call_input(real_coldkey, None, inner_call); - // unauthorized_caller tries to act on behalf of real_coldkey without proxy - let mut env = MockEnv::new( - FunctionId::AddStakeV2, - unauthorized_caller, - (real_coldkey, hotkey, netuid, amount).encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, proxy_contract, input); - let ret = SubtensorChainExtension::::dispatch(&mut env); + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + + // Any proxy should allow staking calls — this was the primary bug in V2 + assert_success(ret); - // Should fail with NotProxy error - assert!(matches!(ret, Err(DispatchError::Other("NotProxy")))); + let total_stake = + pallet_subtensor::Pallet::::get_total_stake_for_hotkey(&hotkey); + assert!(total_stake > TaoCurrency::ZERO); }); } #[test] -fn remove_stake_v2_with_proxy_succeeds() { +fn proxy_call_without_proxy_fails() { mock::new_test_ext(1).execute_with(|| { - let owner_hotkey = U256::from(4601); - let owner_coldkey = U256::from(4602); - let real_coldkey = U256::from(5601); - let proxy_contract = U256::from(5603); - let hotkey = U256::from(5602); - let stake_amount_raw: u64 = 320_000_000_000; + let owner_hotkey = U256::from(1); + let owner_coldkey = U256::from(2); + let real_coldkey = U256::from(101); + let unauthorized_caller = U256::from(102); + let hotkey = U256::from(202); + let min_stake = DefaultMinStake::::get(); + let amount_raw = min_stake.to_u64().saturating_mul(10); + let amount: TaoCurrency = amount_raw.into(); let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); mock::setup_reserves( netuid, - TaoCurrency::from(120_000_000_000), - AlphaCurrency::from(100_000_000_000), + (amount_raw * 1_000_000).into(), + AlphaCurrency::from(amount_raw * 10_000_000), ); - mock::register_ok_neuron(netuid, hotkey, real_coldkey, 0); pallet_subtensor::Pallet::::add_balance_to_coldkey_account( &real_coldkey, - stake_amount_raw + 1_000_000_000, + amount_raw + 1_000_000_000, ); - assert_ok!(pallet_subtensor::Pallet::::add_stake( - RawOrigin::Signed(real_coldkey).into(), + // No proxy relationship established + + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::add_stake { hotkey, netuid, - stake_amount_raw.into(), - )); - - mock::remove_stake_rate_limit_for_tests(&hotkey, &real_coldkey, netuid); - - // Add proxy relationship - assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( - RawOrigin::Signed(real_coldkey).into(), - proxy_contract, - subtensor_runtime_common::ProxyType::Staking, - 0u64, - )); - - let alpha_before = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &real_coldkey, - netuid, - ); + amount_staked: amount, + } + .into(); + let input = encode_proxy_call_input(real_coldkey, None, inner_call); - let alpha_to_unstake: AlphaCurrency = (alpha_before.to_u64() / 2).into(); - - let expected_weight = Weight::from_parts(196_800_000, 0) - .saturating_add(::DbWeight::get().reads(20)) - .saturating_add(::DbWeight::get().writes(10)); - - let mut env = MockEnv::new( - FunctionId::RemoveStakeV2, - proxy_contract, - (real_coldkey, hotkey, netuid, alpha_to_unstake).encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, unauthorized_caller, input); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); - assert_success(ret); - let alpha_after = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &real_coldkey, - netuid, - ); - - assert!(alpha_after < alpha_before); + // Should return NotAuthorizedProxy error code (not a hard DispatchError) + match ret { + RetVal::Converging(code) => { + assert_eq!( + code, + Output::NotAuthorizedProxy as u32, + "expected NotAuthorizedProxy error" + ); + } + _ => panic!("unexpected return value"), + } }); } #[test] -fn swap_stake_v2_with_proxy_succeeds() { +fn proxy_call_with_wrong_proxy_type_fails() { mock::new_test_ext(1).execute_with(|| { - let owner_hotkey_a = U256::from(4401); - let owner_coldkey_a = U256::from(4402); - let owner_hotkey_b = U256::from(4403); - let owner_coldkey_b = U256::from(4404); - let real_coldkey = U256::from(5401); - let proxy_contract = U256::from(5403); - let hotkey = U256::from(5402); - + let owner_hotkey = U256::from(1); + let owner_coldkey = U256::from(2); + let real_coldkey = U256::from(101); + let proxy_contract = U256::from(102); + let hotkey = U256::from(202); let min_stake = DefaultMinStake::::get(); - let stake_amount_raw = min_stake.to_u64().saturating_mul(260); - - let netuid_a = mock::add_dynamic_network(&owner_hotkey_a, &owner_coldkey_a); - let netuid_b = mock::add_dynamic_network(&owner_hotkey_b, &owner_coldkey_b); + let amount_raw = min_stake.to_u64().saturating_mul(10); + let amount: TaoCurrency = amount_raw.into(); + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); mock::setup_reserves( - netuid_a, - stake_amount_raw.saturating_mul(18).into(), - AlphaCurrency::from(stake_amount_raw.saturating_mul(30)), - ); - mock::setup_reserves( - netuid_b, - stake_amount_raw.saturating_mul(20).into(), - AlphaCurrency::from(stake_amount_raw.saturating_mul(28)), + netuid, + (amount_raw * 1_000_000).into(), + AlphaCurrency::from(amount_raw * 10_000_000), ); - - mock::register_ok_neuron(netuid_a, hotkey, real_coldkey, 0); - mock::register_ok_neuron(netuid_b, hotkey, real_coldkey, 1); + mock::register_ok_neuron(netuid, hotkey, real_coldkey, 0); pallet_subtensor::Pallet::::add_balance_to_coldkey_account( &real_coldkey, - stake_amount_raw + 1_000_000_000, + amount_raw + 1_000_000_000, ); - assert_ok!(pallet_subtensor::Pallet::::add_stake( - RawOrigin::Signed(real_coldkey).into(), - hotkey, - netuid_a, - stake_amount_raw.into(), - )); - - mock::remove_stake_rate_limit_for_tests(&hotkey, &real_coldkey, netuid_a); - - // Add proxy relationship + // Add Senate proxy — can't do staking (hits _ => false in InstanceFilter) assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( RawOrigin::Signed(real_coldkey).into(), proxy_contract, - subtensor_runtime_common::ProxyType::Staking, + subtensor_runtime_common::ProxyType::Senate, 0u64, )); - let alpha_origin_before = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &real_coldkey, - netuid_a, - ); - let alpha_destination_before = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &real_coldkey, - netuid_b, - ); - let alpha_to_swap: AlphaCurrency = (alpha_origin_before.to_u64() / 3).into(); - - let expected_weight = Weight::from_parts(351_300_000, 0) - .saturating_add(::DbWeight::get().reads(36)) - .saturating_add(::DbWeight::get().writes(22)); + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::add_stake { + hotkey, + netuid, + amount_staked: amount, + } + .into(); + let input = encode_proxy_call_input(real_coldkey, None, inner_call); - let mut env = MockEnv::new( - FunctionId::SwapStakeV2, - proxy_contract, - (real_coldkey, hotkey, netuid_a, netuid_b, alpha_to_swap).encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, proxy_contract, input); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); - assert_success(ret); - let alpha_origin_after = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &real_coldkey, - netuid_a, - ); - let alpha_destination_after = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &real_coldkey, - netuid_b, - ); - - assert!(alpha_origin_after < alpha_origin_before); - assert!(alpha_destination_after > alpha_destination_before); + // Senate proxy cannot do staking calls + match ret { + RetVal::Converging(code) => { + assert_ne!(code, Output::Success as u32, "should not succeed"); + } + _ => panic!("unexpected return value"), + } }); } #[test] -fn move_stake_v2_with_proxy_succeeds() { +fn proxy_call_with_transfer_proxy_transfer_stake_succeeds() { mock::new_test_ext(1).execute_with(|| { - let owner_hotkey = U256::from(4201); - let owner_coldkey = U256::from(4202); - let real_coldkey = U256::from(5201); - let proxy_contract = U256::from(5204); - let origin_hotkey = U256::from(5202); - let destination_hotkey = U256::from(5203); + let owner_hotkey = U256::from(4301); + let owner_coldkey = U256::from(4302); + let origin_coldkey = U256::from(5301); + let destination_coldkey = U256::from(5302); + let proxy_contract = U256::from(5304); + let hotkey = U256::from(5303); let min_stake = DefaultMinStake::::get(); - let stake_amount_raw = min_stake.to_u64().saturating_mul(240); + let stake_amount_raw = min_stake.to_u64().saturating_mul(250); let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); mock::setup_reserves( @@ -1314,81 +1265,73 @@ fn move_stake_v2_with_proxy_succeeds() { AlphaCurrency::from(stake_amount_raw.saturating_mul(25)), ); - mock::register_ok_neuron(netuid, origin_hotkey, real_coldkey, 0); - mock::register_ok_neuron(netuid, destination_hotkey, real_coldkey, 1); + mock::register_ok_neuron(netuid, hotkey, origin_coldkey, 0); pallet_subtensor::Pallet::::add_balance_to_coldkey_account( - &real_coldkey, + &origin_coldkey, stake_amount_raw + 1_000_000_000, ); assert_ok!(pallet_subtensor::Pallet::::add_stake( - RawOrigin::Signed(real_coldkey).into(), - origin_hotkey, + RawOrigin::Signed(origin_coldkey).into(), + hotkey, netuid, stake_amount_raw.into(), )); - mock::remove_stake_rate_limit_for_tests(&origin_hotkey, &real_coldkey, netuid); + mock::remove_stake_rate_limit_for_tests(&hotkey, &origin_coldkey, netuid); - // Add proxy relationship + // Add Transfer proxy assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( - RawOrigin::Signed(real_coldkey).into(), + RawOrigin::Signed(origin_coldkey).into(), proxy_contract, - subtensor_runtime_common::ProxyType::Staking, + subtensor_runtime_common::ProxyType::Transfer, 0u64, )); let alpha_before = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &origin_hotkey, - &real_coldkey, + &hotkey, + &origin_coldkey, netuid, ); - let alpha_to_move: AlphaCurrency = (alpha_before.to_u64() / 2).into(); + let alpha_to_transfer: AlphaCurrency = (alpha_before.to_u64() / 3).into(); - let expected_weight = Weight::from_parts(164_300_000, 0) - .saturating_add(::DbWeight::get().reads(16)) - .saturating_add(::DbWeight::get().writes(7)); + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::transfer_stake { + destination_coldkey, + hotkey, + origin_netuid: netuid, + destination_netuid: netuid, + alpha_amount: alpha_to_transfer, + } + .into(); + let input = encode_proxy_call_input(origin_coldkey, None, inner_call); - let mut env = MockEnv::new( - FunctionId::MoveStakeV2, - proxy_contract, - ( - real_coldkey, - origin_hotkey, - destination_hotkey, - netuid, - netuid, - alpha_to_move, - ) - .encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, proxy_contract, input); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); assert_success(ret); let origin_alpha_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &origin_hotkey, - &real_coldkey, + &hotkey, + &origin_coldkey, netuid, ); let destination_alpha_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &destination_hotkey, - &real_coldkey, + &hotkey, + &destination_coldkey, netuid, ); - assert_eq!(origin_alpha_after, alpha_before - alpha_to_move); - assert_eq!(destination_alpha_after, alpha_to_move); + assert_eq!(origin_alpha_after, alpha_before - alpha_to_transfer); + assert_eq!(destination_alpha_after, alpha_to_transfer); }); } #[test] -fn transfer_stake_v2_requires_transfer_proxy() { +fn proxy_call_with_staking_proxy_transfer_stake_fails() { mock::new_test_ext(1).execute_with(|| { let owner_hotkey = U256::from(4301); let owner_coldkey = U256::from(4302); @@ -1423,7 +1366,7 @@ fn transfer_stake_v2_requires_transfer_proxy() { mock::remove_stake_rate_limit_for_tests(&hotkey, &origin_coldkey, netuid); - // Add Staking proxy (wrong type for transfer_stake) + // Add Staking proxy — cannot do transfer_stake assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( RawOrigin::Signed(origin_coldkey).into(), proxy_contract, @@ -1439,277 +1382,95 @@ fn transfer_stake_v2_requires_transfer_proxy() { ); let alpha_to_transfer: AlphaCurrency = (alpha_before.to_u64() / 3).into(); - let expected_weight = Weight::from_parts(160_300_000, 0) - .saturating_add(::DbWeight::get().reads(14)) - .saturating_add(::DbWeight::get().writes(6)); - - // First try with Staking proxy - should fail - let mut env = MockEnv::new( - FunctionId::TransferStakeV2, - proxy_contract, - ( - origin_coldkey, - destination_coldkey, - hotkey, - netuid, - netuid, - alpha_to_transfer, - ) - .encode(), - ) - .with_expected_weight(expected_weight); - - let ret = SubtensorChainExtension::::dispatch(&mut env); - // Staking proxy should not work for transfer_stake - requires Transfer proxy - assert!(matches!(ret, Err(DispatchError::Other("NotProxy")))); - - // Remove Staking proxy, add Transfer proxy - assert_ok!(pallet_subtensor_proxy::Pallet::::remove_proxy( - RawOrigin::Signed(origin_coldkey).into(), - proxy_contract, - subtensor_runtime_common::ProxyType::Staking, - 0u64, - )); - assert_ok!(pallet_subtensor_proxy::Pallet::::add_proxy( - RawOrigin::Signed(origin_coldkey).into(), - proxy_contract, - subtensor_runtime_common::ProxyType::Transfer, - 0u64, - )); + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::transfer_stake { + destination_coldkey, + hotkey, + origin_netuid: netuid, + destination_netuid: netuid, + alpha_amount: alpha_to_transfer, + } + .into(); + let input = encode_proxy_call_input(origin_coldkey, None, inner_call); - // Now try with Transfer proxy - should succeed - let mut env = MockEnv::new( - FunctionId::TransferStakeV2, - proxy_contract, - ( - origin_coldkey, - destination_coldkey, - hotkey, - netuid, - netuid, - alpha_to_transfer, - ) - .encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, proxy_contract, input); let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); - assert_success(ret); + // Staking proxy cannot perform transfer_stake (requires Transfer proxy) + match ret { + RetVal::Converging(code) => { + assert_ne!(code, Output::Success as u32, "should not succeed"); + } + _ => panic!("unexpected return value"), + } + + // Verify stake was NOT transferred let origin_alpha_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( &hotkey, &origin_coldkey, netuid, ); - let destination_alpha_after = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &destination_coldkey, - netuid, - ); - - assert_eq!(origin_alpha_after, alpha_before - alpha_to_transfer); - assert_eq!(destination_alpha_after, alpha_to_transfer); - }); -} - -#[test] -fn unstake_all_v2_self_call_succeeds() { - mock::new_test_ext(1).execute_with(|| { - let owner_hotkey = U256::from(4001); - let owner_coldkey = U256::from(4002); - let coldkey = U256::from(5001); - let hotkey = U256::from(5002); - let min_stake = DefaultMinStake::::get(); - let stake_amount_raw = min_stake.to_u64().saturating_mul(200); - let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); - - mock::setup_reserves( - netuid, - stake_amount_raw.saturating_mul(10).into(), - AlphaCurrency::from(stake_amount_raw.saturating_mul(20)), - ); - - mock::register_ok_neuron(netuid, hotkey, coldkey, 0); - pallet_subtensor::Pallet::::add_balance_to_coldkey_account( - &coldkey, - stake_amount_raw + 1_000_000_000, - ); - - assert_ok!(pallet_subtensor::Pallet::::add_stake( - RawOrigin::Signed(coldkey).into(), - hotkey, - netuid, - stake_amount_raw.into(), - )); - - mock::remove_stake_rate_limit_for_tests(&hotkey, &coldkey, netuid); - - let expected_weight = Weight::from_parts(28_830_000, 0) - .saturating_add(::DbWeight::get().reads(7)) - .saturating_add(::DbWeight::get().writes(0)); - - let pre_balance = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); - - let mut env = MockEnv::new( - FunctionId::UnstakeAllV2, - coldkey, - (coldkey, hotkey).encode(), - ) - .with_expected_weight(expected_weight); - - let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); - assert_success(ret); - assert_eq!(env.charged_weight(), Some(expected_weight)); - - let remaining_alpha = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &coldkey, netuid, - ); - assert!(remaining_alpha <= AlphaCurrency::from(1_000)); - - let post_balance = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); - assert!(post_balance > pre_balance); - }); -} - -#[test] -fn unstake_all_alpha_v2_self_call_succeeds() { - mock::new_test_ext(1).execute_with(|| { - let owner_hotkey = U256::from(4101); - let owner_coldkey = U256::from(4102); - let coldkey = U256::from(5101); - let hotkey = U256::from(5102); - let min_stake = DefaultMinStake::::get(); - let stake_amount_raw = min_stake.to_u64().saturating_mul(220); - let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); - - mock::setup_reserves( - netuid, - stake_amount_raw.saturating_mul(20).into(), - AlphaCurrency::from(stake_amount_raw.saturating_mul(30)), - ); - - mock::register_ok_neuron(netuid, hotkey, coldkey, 0); - pallet_subtensor::Pallet::::add_balance_to_coldkey_account( - &coldkey, - stake_amount_raw + 1_000_000_000, + assert_eq!( + origin_alpha_after, alpha_before, + "stake should be unchanged" ); - - assert_ok!(pallet_subtensor::Pallet::::add_stake( - RawOrigin::Signed(coldkey).into(), - hotkey, - netuid, - stake_amount_raw.into(), - )); - - mock::remove_stake_rate_limit_for_tests(&hotkey, &coldkey, netuid); - - let expected_weight = Weight::from_parts(358_500_000, 0) - .saturating_add(::DbWeight::get().reads(37)) - .saturating_add(::DbWeight::get().writes(21)); - - let mut env = MockEnv::new( - FunctionId::UnstakeAllAlphaV2, - coldkey, - (coldkey, hotkey).encode(), - ) - .with_expected_weight(expected_weight); - - let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); - assert_success(ret); - assert_eq!(env.charged_weight(), Some(expected_weight)); - - let subnet_alpha = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &coldkey, netuid, - ); - assert!(subnet_alpha <= AlphaCurrency::from(1_000)); - - let root_alpha = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, - &coldkey, - NetUid::ROOT, - ); - assert!(root_alpha > AlphaCurrency::ZERO); }); } #[test] -fn set_coldkey_auto_stake_hotkey_v2_self_call_succeeds() { +fn proxy_call_with_invalid_call_data_fails() { mock::new_test_ext(1).execute_with(|| { - let owner_hotkey = U256::from(4901); - let owner_coldkey = U256::from(4902); - let coldkey = U256::from(5901); - let hotkey = U256::from(5902); - - let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); - - pallet_subtensor::Owner::::insert(hotkey, coldkey); - pallet_subtensor::OwnedHotkeys::::insert(coldkey, vec![hotkey]); - pallet_subtensor::Uids::::insert(netuid, hotkey, 0u16); - - assert_eq!( - pallet_subtensor::AutoStakeDestination::::get(coldkey, netuid), - None - ); + let coldkey = U256::from(101); - let expected_weight = Weight::from_parts(29_930_000, 0) - .saturating_add(::DbWeight::get().reads(5)) - .saturating_add(::DbWeight::get().writes(2)); + // Use garbage bytes that won't decode as a RuntimeCall + let garbage_data: frame_support::BoundedVec> = + vec![0xFF, 0xFE, 0xFD, 0xFC, 0xFB].try_into().unwrap(); + let input = (coldkey, None::, garbage_data).encode(); - let mut env = MockEnv::new( - FunctionId::SetColdkeyAutoStakeHotkeyV2, - coldkey, - (coldkey, netuid, hotkey).encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, coldkey, input); - let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); - assert_success(ret); - assert_eq!(env.charged_weight(), Some(expected_weight)); + let ret = SubtensorChainExtension::::dispatch(&mut env); - assert_eq!( - pallet_subtensor::AutoStakeDestination::::get(coldkey, netuid), - Some(hotkey) - ); + // Should fail with decode error + assert!(matches!( + ret, + Err(DispatchError::Other("Failed to decode call")) + )); }); } #[test] -fn remove_stake_v2_self_call_returns_error_with_no_stake() { +fn proxy_call_with_invalid_proxy_type_byte_fails() { mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(101); + let hotkey = U256::from(202); let owner_hotkey = U256::from(1); let owner_coldkey = U256::from(2); - let coldkey = U256::from(301); - let hotkey = U256::from(302); + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); - mock::register_ok_neuron(netuid, hotkey, coldkey, 0); - let min_stake = DefaultMinStake::::get(); - let amount: AlphaCurrency = AlphaCurrency::from(min_stake.to_u64()); + let inner_call: mock::RuntimeCall = pallet_subtensor::Call::::add_stake { + hotkey, + netuid, + amount_staked: 1_000u64.into(), + } + .into(); + let call_data: frame_support::BoundedVec> = + inner_call.encode().try_into().unwrap(); - let expected_weight = Weight::from_parts(196_800_000, 0) - .saturating_add(::DbWeight::get().reads(20)) - .saturating_add(::DbWeight::get().writes(10)); + // Use an invalid proxy type byte (255 is not a valid ProxyType) + let input = (coldkey, Some(255u8), call_data).encode(); - let mut env = MockEnv::new( - FunctionId::RemoveStakeV2, - coldkey, - (coldkey, hotkey, netuid, amount).encode(), - ) - .with_expected_weight(expected_weight); + let mut env = MockEnv::new(FunctionId::ProxyCall, coldkey, input); - let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + let ret = SubtensorChainExtension::::dispatch(&mut env); - match ret { - RetVal::Converging(code) => { - assert_eq!(code, Output::AmountTooLow as u32, "mismatched error output") - } - _ => panic!("unexpected return value"), - } + // Should fail with invalid proxy type error + assert!(matches!( + ret, + Err(DispatchError::Other("Invalid proxy type")) + )); }); } diff --git a/chain-extensions/src/types.rs b/chain-extensions/src/types.rs index b0fb53ca79..83384d2744 100644 --- a/chain-extensions/src/types.rs +++ b/chain-extensions/src/types.rs @@ -21,19 +21,8 @@ pub enum FunctionId { AddProxyV1 = 13, RemoveProxyV1 = 14, GetAlphaPriceV1 = 15, - // V2: Accept explicit coldkey with proxy verification - AddStakeV2 = 16, - RemoveStakeV2 = 17, - UnstakeAllV2 = 18, - UnstakeAllAlphaV2 = 19, - MoveStakeV2 = 20, - TransferStakeV2 = 21, - SwapStakeV2 = 22, - AddStakeLimitV2 = 23, - RemoveStakeLimitV2 = 24, - SwapStakeLimitV2 = 25, - RemoveStakeFullLimitV2 = 26, - SetColdkeyAutoStakeHotkeyV2 = 27, + // Proxy-aware generic call dispatcher: wraps any RuntimeCall through pallet_proxy + ProxyCall = 16, } #[derive(PartialEq, Eq, Copy, Clone, Encode, Decode, Debug)] From 2c60e9333886934aa0aea50d428689c92fb45ead Mon Sep 17 00:00:00 2001 From: Landyn Date: Wed, 11 Feb 2026 21:09:30 -0600 Subject: [PATCH 6/6] Update docs: replace V2 functions with ProxyCall dispatcher --- docs/wasm-contracts.md | 28 +++++++++------------------- 1 file changed, 9 insertions(+), 19 deletions(-) diff --git a/docs/wasm-contracts.md b/docs/wasm-contracts.md index fd7b9a9abf..c9370ab14b 100644 --- a/docs/wasm-contracts.md +++ b/docs/wasm-contracts.md @@ -44,18 +44,7 @@ Subtensor provides a custom chain extension that allows smart contracts to inter | 13 | `add_proxy` | Add a staking proxy for the caller | `(AccountId)` | Error code | | 14 | `remove_proxy` | Remove a staking proxy for the caller | `(AccountId)` | Error code | | 15 | `get_alpha_price` | Get the current alpha price for a subnet | `(NetUid)` | `u64` (price × 10⁹) | -| 16 | `add_stake_v2` | Add stake with explicit coldkey (proxy-aware) | `(AccountId, AccountId, NetUid, TaoCurrency)` | Error code | -| 17 | `remove_stake_v2` | Remove stake with explicit coldkey | `(AccountId, AccountId, NetUid, AlphaCurrency)` | Error code | -| 18 | `unstake_all_v2` | Unstake all TAO with explicit coldkey | `(AccountId, AccountId)` | Error code | -| 19 | `unstake_all_alpha_v2` | Unstake all Alpha with explicit coldkey | `(AccountId, AccountId)` | Error code | -| 20 | `move_stake_v2` | Move stake between hotkeys with explicit coldkey | `(AccountId, AccountId, AccountId, NetUid, NetUid, AlphaCurrency)` | Error code | -| 21 | `transfer_stake_v2` | Transfer stake between coldkeys (requires Transfer proxy) | `(AccountId, AccountId, AccountId, NetUid, NetUid, AlphaCurrency)` | Error code | -| 22 | `swap_stake_v2` | Swap stake between subnets with explicit coldkey | `(AccountId, AccountId, NetUid, NetUid, AlphaCurrency)` | Error code | -| 23 | `add_stake_limit_v2` | Add stake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, TaoCurrency, TaoCurrency, bool)` | Error code | -| 24 | `remove_stake_limit_v2` | Remove stake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, AlphaCurrency, TaoCurrency, bool)` | Error code | -| 25 | `swap_stake_limit_v2` | Swap stake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, NetUid, AlphaCurrency, TaoCurrency, bool)` | Error code | -| 26 | `remove_stake_full_limit_v2` | Full unstake with price limit and explicit coldkey | `(AccountId, AccountId, NetUid, Option)` | Error code | -| 27 | `set_coldkey_auto_stake_hotkey_v2` | Set auto-stake hotkey with explicit coldkey | `(AccountId, NetUid, AccountId)` | Error code | +| 16 | `proxy_call` | Dispatch any RuntimeCall through pallet_proxy | `(AccountId, Option, BoundedVec)` | Error code | Example usage in your ink! contract: ```rust @@ -100,16 +89,17 @@ Chain extension functions that modify state return error codes as `u32` values. | 19 | `ProxyNotFound` | Proxy relationship not found | | 20 | `NotAuthorizedProxy` | Caller is not an authorized proxy for the account | -#### V2 Functions (Proxy-Aware) +#### ProxyCall (ID 16) — Proxy-Aware Generic Dispatcher -Functions 16-27 are V2 versions that accept an explicit `coldkey` parameter as the first argument. These functions: +Instead of per-function proxy variants, a single `proxy_call` extension dispatches any SCALE-encoded `RuntimeCall` through `pallet_proxy`. Parameters: -- If `coldkey == caller`: Execute directly (no proxy check needed) -- If `coldkey != caller`: Verify caller has appropriate proxy permissions for coldkey +- `real_coldkey: AccountId` — the account to act on behalf of +- `force_proxy_type: Option` — optional proxy type filter (e.g., `5` for Staking), or `None` to match any +- `call_data: BoundedVec` — SCALE-encoded `RuntimeCall` -**Proxy Types Required:** -- Most V2 functions require `ProxyType::Staking` -- `transfer_stake_v2` (ID 21) requires `ProxyType::Transfer` +Behavior: +- If `real_coldkey == caller`: dispatches the call directly (no proxy needed) +- If `real_coldkey != caller`: routes through `pallet_proxy::proxy()`, which checks proxy permissions via `InstanceFilter` — correctly supporting `Any`, `Staking`, `Transfer`, and all other proxy types ### Call Filter