refactor code

Author: open-junius

runtime/src/precompiles/balance_transfer.rs

@@ -1,11 +1,12 @@
-use frame_system::RawOrigin;
 use pallet_evm::{
     BalanceConverter, ExitError, ExitSucceed, PrecompileHandle, PrecompileOutput, PrecompileResult,
 };
 use sp_runtime::traits::UniqueSaturatedInto;
 use sp_std::vec;
 
-use crate::precompiles::{get_method_id, get_pubkey, get_slice, try_dispatch_runtime_call};
+use crate::precompiles::{
+    contract_to_origin, get_method_id, get_pubkey, get_slice, try_dispatch_runtime_call,
+};
 use crate::Runtime;
 
 pub const BALANCE_TRANSFER_INDEX: u64 = 2048;
@@ -47,15 +48,13 @@ impl BalanceTransferPrecompile {
         }
 
         let address_bytes_dst = get_slice(txdata, 4, 36)?;
-        let (account_id_src, _) = get_pubkey(&CONTRACT_ADDRESS_SS58)?;
         let (account_id_dst, _) = get_pubkey(address_bytes_dst)?;
 
         let call = pallet_balances::Call::<Runtime>::transfer_allow_death {
             dest: account_id_dst.into(),
             value: amount_sub.unique_saturated_into(),
         };
-        let origin = RawOrigin::Signed(account_id_src);
 
-        try_dispatch_runtime_call(handle, call, origin)
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 }

runtime/src/precompiles/mod.rs

@@ -6,9 +6,8 @@ use core::marker::PhantomData;
 use frame_support::dispatch::{GetDispatchInfo, Pays};
 
 use pallet_evm::{
-    AddressMapping, BalanceConverter, ExitError, ExitSucceed, GasWeightMapping,
-    HashedAddressMapping, IsPrecompileResult, Precompile, PrecompileFailure, PrecompileHandle,
-    PrecompileOutput, PrecompileResult, PrecompileSet,
+    ExitError, ExitSucceed, GasWeightMapping, IsPrecompileResult, Precompile, PrecompileFailure,
+    PrecompileHandle, PrecompileOutput, PrecompileResult, PrecompileSet,
 };
 use pallet_evm_precompile_modexp::Modexp;
 use pallet_evm_precompile_sha3fips::Sha3FIPS256;
@@ -20,10 +19,6 @@ use crate::{Runtime, RuntimeCall};
 
 use frame_system::RawOrigin;
 
-use sp_core::crypto::Ss58Codec;
-use sp_core::U256;
-use sp_runtime::traits::{BlakeTwo256, UniqueSaturatedInto};
-
 use sp_std::vec;
 
 // Include custom precompiles
@@ -147,84 +142,6 @@ pub fn get_slice(data: &[u8], from: usize, to: usize) -> Result<&[u8], Precompil
     }
 }
 
-/// The function return the token to smart contract
-fn transfer_back_to_caller(
-    smart_contract_address: &str,
-    account_id: &AccountId32,
-    amount: U256,
-) -> Result<(), PrecompileFailure> {
-    // this is staking smart contract's(0x0000000000000000000000000000000000000801) sr25519 address
-    let smart_contract_account_id = match AccountId32::from_ss58check(smart_contract_address) {
-        // match AccountId32::from_ss58check("5CwnBK9Ack1mhznmCnwiibCNQc174pYQVktYW3ayRpLm4K2X") {
-        Ok(addr) => addr,
-        Err(_) => {
-            return Err(PrecompileFailure::Error {
-                exit_status: ExitError::Other("Invalid SS58 address".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::Other("Transfer back to caller failed".into()),
-        });
-    }
-
-    Ok(())
-}
-
-fn dispatch(
-    handle: &mut impl PrecompileHandle,
-    call: RuntimeCall,
-    smart_contract_address: &str,
-) -> PrecompileResult {
-    let account_id =
-        <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(
-            handle.context().caller,
-        );
-
-    // Transfer the amount back to the caller before executing the staking operation
-    // let caller = handle.context().caller;
-    let amount = handle.context().apparent_value;
-
-    if !amount.is_zero() {
-        transfer_back_to_caller(smart_contract_address, &account_id, amount)?;
-    }
-
-    let result = call.dispatch(RawOrigin::Signed(account_id.clone()).into());
-    match &result {
-        Ok(post_info) => log::info!("Dispatch succeeded. Post info: {:?}", post_info),
-        Err(dispatch_error) => log::error!("Dispatch failed. Error: {:?}", dispatch_error),
-    }
-    match result {
-        Ok(_) => Ok(PrecompileOutput {
-            exit_status: ExitSucceed::Returned,
-            output: vec![],
-        }),
-        Err(_) => Err(PrecompileFailure::Error {
-            exit_status: ExitError::Other("Subtensor call failed".into()),
-        }),
-    }
-}
-
 pub fn get_pubkey(data: &[u8]) -> Result<(AccountId32, vec::Vec<u8>), PrecompileFailure> {
     let mut pubkey = [0u8; 32];
     pubkey.copy_from_slice(get_slice(data, 0, 32)?);
@@ -235,6 +152,26 @@ pub fn get_pubkey(data: &[u8]) -> Result<(AccountId32, vec::Vec<u8>), Precompile
             .map_or_else(vec::Vec::new, |slice| slice.to_vec()),
     ))
 }
+
+fn parse_netuid(data: &[u8], offset: usize) -> Result<u16, PrecompileFailure> {
+    if data.len() < offset + 2 {
+        return Err(PrecompileFailure::Error {
+            exit_status: ExitError::InvalidRange,
+        });
+    }
+
+    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)
+}
+
+fn contract_to_origin(contract: &[u8; 32]) -> Result<RawOrigin<AccountId32>, PrecompileFailure> {
+    let (account_id, _) = get_pubkey(contract)?;
+    Ok(RawOrigin::Signed(account_id))
+}
+
 /// Dispatches a runtime call, but also checks and records the gas costs.
 fn try_dispatch_runtime_call(
     handle: &mut impl PrecompileHandle,

runtime/src/precompiles/neuron.rs

@@ -1,15 +1,21 @@
 use pallet_evm::{ExitError, PrecompileFailure, PrecompileHandle, PrecompileResult};
 
-use crate::precompiles::{dispatch, get_method_id, get_pubkey, get_slice};
+use crate::precompiles::{
+    contract_to_origin, get_method_id, get_pubkey, get_slice, parse_netuid,
+    try_dispatch_runtime_call,
+};
 use sp_runtime::AccountId32;
 use sp_std::vec;
 
 use crate::{Runtime, RuntimeCall};
 pub const NEURON_PRECOMPILE_INDEX: u64 = 2052;
 
-// this is neuron smart contract's(0x0000000000000000000000000000000000000804) sr25519 address
-pub const NEURON_CONTRACT_ADDRESS: &str = "5GKZiUUgTnWSz3BgiVBMehEKkLszsG4ZXnvgWpWFUFKqrqyn";
-
+// ss58 public key i.e., the contract sends funds it received to the destination address from the
+// method parameter.
+const CONTRACT_ADDRESS_SS58: [u8; 32] = [
+    0xbc, 0x46, 0x35, 0x79, 0xbc, 0x99, 0xf9, 0xee, 0x7c, 0x59, 0xed, 0xee, 0x20, 0x61, 0xa3, 0x09,
+    0xd2, 0x1e, 0x68, 0xd5, 0x39, 0xb6, 0x40, 0xec, 0x66, 0x46, 0x90, 0x30, 0xab, 0x74, 0xc1, 0xdb,
+];
 pub struct NeuronPrecompile;
 
 impl NeuronPrecompile {
@@ -38,7 +44,7 @@ impl NeuronPrecompile {
                 netuid,
                 hotkey,
             });
-        dispatch(handle, call, NEURON_CONTRACT_ADDRESS)
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 
     fn parse_netuid_hotkey_parameter(data: &[u8]) -> Result<(u16, AccountId32), PrecompileFailure> {
@@ -47,9 +53,7 @@ impl NeuronPrecompile {
                 exit_status: ExitError::InvalidRange,
             });
         }
-        let mut netuid_vec = [0u8; 2];
-        netuid_vec.copy_from_slice(get_slice(data, 30, 32)?);
-        let netuid = u16::from_be_bytes(netuid_vec);
+        let netuid = parse_netuid(data, 30)?;
 
         let (hotkey, _) = get_pubkey(get_slice(data, 32, 64)?)?;
 

runtime/src/precompiles/staking.rs

@@ -25,7 +25,10 @@
 //   - Precompile checks the result of do_remove_stake and, in case of a failure, reverts the transaction.
 //
 
-use crate::precompiles::{dispatch, get_method_id, get_pubkey, get_slice};
+use crate::precompiles::{
+    contract_to_origin, get_method_id, get_pubkey, get_slice, parse_netuid,
+    try_dispatch_runtime_call,
+};
 use crate::{ProxyType, Runtime, RuntimeCall};
 use pallet_evm::{
     BalanceConverter, ExitError, ExitSucceed, PrecompileFailure, PrecompileHandle,
@@ -36,8 +39,14 @@ use sp_runtime::traits::{StaticLookup, UniqueSaturatedInto};
 use sp_std::vec;
 
 pub const STAKING_PRECOMPILE_INDEX: u64 = 2049;
-// this is staking smart contract's(0x0000000000000000000000000000000000000801) sr25519 address
-pub const STAKING_CONTRACT_ADDRESS: &str = "5CwnBK9Ack1mhznmCnwiibCNQc174pYQVktYW3ayRpLm4K2X";
+
+// ss58 public key i.e., the contract sends funds it received to the destination address from the
+// method parameter.
+const CONTRACT_ADDRESS_SS58: [u8; 32] = [
+    0x26, 0xf4, 0x10, 0x1e, 0x52, 0xb7, 0x57, 0x34, 0x33, 0x24, 0x5b, 0xc3, 0x0a, 0xe1, 0x8b, 0x63,
+    0x99, 0x53, 0xd8, 0x41, 0x79, 0x33, 0x03, 0x61, 0x4d, 0xfa, 0xcf, 0xf0, 0x37, 0xf7, 0x12, 0x94,
+];
+
 pub struct StakingPrecompile;
 
 impl StakingPrecompile {
@@ -74,14 +83,16 @@ impl StakingPrecompile {
             <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
-        dispatch(handle, call, STAKING_CONTRACT_ADDRESS)
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 
     fn remove_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
@@ -104,7 +115,7 @@ impl StakingPrecompile {
             netuid,
             amount_unstaked: amount_sub.unique_saturated_into(),
         });
-        dispatch(handle, call, STAKING_CONTRACT_ADDRESS)
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 
     fn add_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
@@ -116,7 +127,7 @@ impl StakingPrecompile {
             delay: 0,
         });
 
-        dispatch(handle, call, STAKING_CONTRACT_ADDRESS)
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 
     fn remove_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
@@ -128,13 +139,13 @@ impl StakingPrecompile {
             delay: 0,
         });
 
-        dispatch(handle, call, STAKING_CONTRACT_ADDRESS)
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 
     fn get_stake(data: &[u8]) -> PrecompileResult {
         let (hotkey, left_data) = get_pubkey(data)?;
         let (coldkey, _) = get_pubkey(&left_data)?;
-        let netuid = Self::parse_netuid(data, 0x5E)?;
+        let netuid = parse_netuid(data, 0x5E)?;
 
         let stake = pallet_subtensor::Pallet::<Runtime>::get_stake_for_hotkey_and_coldkey_on_subnet(
             &hotkey, &coldkey, netuid,

runtime/src/precompiles/subnet.rs

@@ -1,4 +1,6 @@
-use crate::precompiles::{dispatch, get_method_id, get_pubkey, get_slice};
+use crate::precompiles::{
+    contract_to_origin, get_method_id, get_pubkey, get_slice, try_dispatch_runtime_call,
+};
 use crate::{Runtime, RuntimeCall};
 use pallet_evm::{ExitError, PrecompileFailure, PrecompileHandle, PrecompileResult};
 use sp_runtime::AccountId32;
@@ -8,9 +10,12 @@ pub const SUBNET_PRECOMPILE_INDEX: u64 = 2051;
 // three bytes with max lenght 1K
 pub const MAX_PARAMETER_SIZE: usize = 3 * 1024;
 
-// this is staking smart contract's(0x0000000000000000000000000000000000000803) sr25519 address
-pub const STAKING_CONTRACT_ADDRESS: &str = "5DPSUCb5mZFfizvBDSnRoAqmxV5Bmov2CS3xV773qU6VP1w2";
-
+// ss58 public key i.e., the contract sends funds it received to the destination address from the
+// method parameter.
+const CONTRACT_ADDRESS_SS58: [u8; 32] = [
+    0x3a, 0x86, 0x18, 0xfb, 0xbb, 0x1b, 0xbc, 0x47, 0x86, 0x64, 0xff, 0x53, 0x46, 0x18, 0x0c, 0x35,
+    0xd0, 0x9f, 0xac, 0x26, 0xf2, 0x02, 0x70, 0x85, 0xb3, 0x1c, 0x56, 0xc1, 0x06, 0x3c, 0x1c, 0xd3,
+];
 pub struct SubnetPrecompile;
 
 impl SubnetPrecompile {
@@ -79,7 +84,7 @@ impl SubnetPrecompile {
         };
 
         // Dispatch the register_network call
-        dispatch(handle, call, STAKING_CONTRACT_ADDRESS)
+        try_dispatch_runtime_call(handle, call, contract_to_origin(&CONTRACT_ADDRESS_SS58)?)
     }
 
     fn parse_register_network_parameters(