fix wrong origin

Author: open-junius

runtime/src/precompiles/staking.rs

@@ -26,16 +26,17 @@
 //
 
 use crate::precompiles::{
-    contract_to_origin, get_method_id, get_pubkey, get_slice, parse_netuid,
-    try_dispatch_runtime_call,
+    get_method_id, get_pubkey, get_slice, parse_netuid, try_dispatch_runtime_call,
 };
 use crate::{ProxyType, Runtime, RuntimeCall};
+use frame_system::RawOrigin;
 use pallet_evm::{
-    BalanceConverter, ExitError, ExitSucceed, PrecompileFailure, PrecompileHandle,
-    PrecompileOutput, PrecompileResult,
+    AddressMapping, BalanceConverter, ExitError, ExitSucceed, HashedAddressMapping,
+    PrecompileFailure, PrecompileHandle, PrecompileOutput, PrecompileResult,
 };
 use sp_core::U256;
-use sp_runtime::traits::{StaticLookup, UniqueSaturatedInto};
+use sp_runtime::traits::{BlakeTwo256, Dispatchable, StaticLookup, UniqueSaturatedInto};
+use sp_runtime::AccountId32;
 use sp_std::vec;
 
 pub const STAKING_PRECOMPILE_INDEX: u64 = 2049;
@@ -75,27 +76,37 @@ impl StakingPrecompile {
     }
 
     fn add_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+        let account_id =
+            <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
+                handle.context().caller,
+            );
+
         let (hotkey, _) = get_pubkey(data)?;
         let amount: U256 = handle.context().apparent_value;
         let netuid = Self::parse_netuid(data, 0x3E)?;
 
+        if !amount.is_zero() {
+            Self::transfer_back_to_caller(&account_id, amount)?;
+        }
+
         let amount_sub =
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
                 .ok_or(ExitError::OutOfFund)?;
 
-        // let (account_id_src, _) = get_pubkey(&CONTRACT_ADDRESS_SS58)?;
-        // Create the add_stake call
         let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::add_stake {
             hotkey,
             netuid,
             amount_staked: amount_sub.unique_saturated_into(),
         });
-        // let origin = RawOrigin::Signed(account_id_src);
-        // Dispatch the add_stake call
-        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
+
+        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
     fn remove_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+        let account_id =
+            <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
+                handle.context().caller,
+            );
         let (hotkey, _) = get_pubkey(data)?;
         let netuid = Self::parse_netuid(data, 0x5E)?;
 
@@ -115,10 +126,14 @@ impl StakingPrecompile {
             netuid,
             amount_unstaked: amount_sub.unique_saturated_into(),
         });
-        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
+        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
     fn add_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+        let account_id =
+            <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
+                handle.context().caller,
+            );
         let (delegate, _) = get_pubkey(data)?;
         let delegate = <Runtime as frame_system::Config>::Lookup::unlookup(delegate);
         let call = RuntimeCall::Proxy(pallet_proxy::Call::<Runtime>::add_proxy {
@@ -127,10 +142,14 @@ impl StakingPrecompile {
             delay: 0,
         });
 
-        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
+        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
     fn remove_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+        let account_id =
+            <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
+                handle.context().caller,
+            );
         let (delegate, _) = get_pubkey(data)?;
         let delegate = <Runtime as frame_system::Config>::Lookup::unlookup(delegate);
         let call = RuntimeCall::Proxy(pallet_proxy::Call::<Runtime>::remove_proxy {
@@ -139,7 +158,7 @@ impl StakingPrecompile {
             delay: 0,
         });
 
-        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
+        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
     fn get_stake(data: &[u8]) -> PrecompileResult {
@@ -167,17 +186,37 @@ impl StakingPrecompile {
         })
     }
 
-    fn parse_netuid(data: &[u8], offset: usize) -> Result<u16, PrecompileFailure> {
-        if data.len() < offset + 2 {
+    fn transfer_back_to_caller(
+        account_id: &AccountId32,
+        amount: U256,
+    ) -> Result<(), PrecompileFailure> {
+        let smart_contract_account_id: AccountId32 = CONTRACT_ADDRESS_SS58.into();
+
+        let amount_sub =
+            <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
+                .ok_or(ExitError::OutOfFund)?;
+
+        // Create a transfer call from the smart contract to the caller
+        let transfer_call =
+            RuntimeCall::Balances(pallet_balances::Call::<Runtime>::transfer_allow_death {
+                dest: account_id.clone().into(),
+                value: amount_sub.unique_saturated_into(),
+            });
+
+        // Execute the transfer
+        let transfer_result =
+            transfer_call.dispatch(RawOrigin::Signed(smart_contract_account_id).into());
+
+        if let Err(dispatch_error) = transfer_result {
+            log::error!(
+                "Transfer back to caller failed. Error: {:?}",
+                dispatch_error
+            );
             return Err(PrecompileFailure::Error {
-                exit_status: ExitError::InvalidRange,
+                exit_status: ExitError::Other("Transfer back to caller failed".into()),
             });
         }
 
-        let mut netuid_bytes = [0u8; 2];
-        netuid_bytes.copy_from_slice(get_slice(data, offset, offset + 2)?);
-        let netuid: u16 = netuid_bytes[1] as u16 | ((netuid_bytes[0] as u16) << 8u16);
-
-        Ok(netuid)
+        Ok(())
     }
 }

runtime/src/precompiles/subnet.rs

@@ -1,11 +1,13 @@
-use crate::precompiles::{
-    contract_to_origin, get_method_id, get_pubkey, get_slice, try_dispatch_runtime_call,
-};
+use crate::precompiles::{get_method_id, get_pubkey, get_slice, try_dispatch_runtime_call};
 use crate::{Runtime, RuntimeCall};
-use pallet_evm::{ExitError, PrecompileFailure, PrecompileHandle, PrecompileResult};
+use frame_system::RawOrigin;
+use pallet_evm::{
+    AddressMapping, ExitError, HashedAddressMapping, PrecompileFailure, PrecompileHandle,
+    PrecompileResult,
+};
+use sp_runtime::traits::BlakeTwo256;
 use sp_runtime::AccountId32;
 use sp_std::vec;
-
 pub const SUBNET_PRECOMPILE_INDEX: u64 = 2051;
 // bytes with max lenght 1K
 pub const MAX_SINGLE_PARAMETER_SIZE: usize = 1024;
@@ -51,7 +53,6 @@ impl SubnetPrecompile {
         let call = match data.len() {
             32 => {
                 let (hotkey, _) = get_pubkey(data)?;
-
                 RuntimeCall::SubtensorModule(
                     pallet_subtensor::Call::<Runtime>::register_network_with_identity {
                         hotkey,
@@ -85,8 +86,13 @@ impl SubnetPrecompile {
             }
         };
 
+        let account_id =
+            <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
+                handle.context().caller,
+            );
+
         // Dispatch the register_network call
-        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
+        try_dispatch_runtime_call(handle, call, RawOrigin::Signed(account_id))
     }
 
     fn parse_register_network_parameters(