Merge pull request #1123 from opentensor/feat/proxy-precompile-extension

Extend staking contract with proxy setting

Author: sam0x17

runtime/src/precompiles/solidity/staking.abi

@@ -1,4 +1,17 @@
 [
+    {
+        "inputs": [
+            {
+                "internalType": "bytes32",
+                "name": "delegate",
+                "type": "bytes32"
+            }
+        ],
+        "name": "addProxy",
+        "outputs": [],
+        "stateMutability": "nonpayable",
+        "type": "function"
+    },
     {
         "inputs": [
             {
@@ -17,6 +30,19 @@
         "stateMutability": "payable",
         "type": "function"
     },
+    {
+        "inputs": [
+            {
+                "internalType": "bytes32",
+                "name": "delegate",
+                "type": "bytes32"
+            }
+        ],
+        "name": "removeProxy",
+        "outputs": [],
+        "stateMutability": "nonpayable",
+        "type": "function"
+    },
     {
         "inputs": [
             {

runtime/src/precompiles/solidity/staking.sol

@@ -42,6 +42,20 @@ interface IStaking {
    */
   function removeStake(bytes32 hotkey, uint256 amount, uint256 netuid) external;
 
+  /**
+   * @dev Delegates staking to a proxy account.
+   *
+   * @param delegate The public key (32 bytes) of the delegate.
+   */
+  function addProxy(bytes32 delegate) external;
+
+  /**
+   * @dev Removes staking proxy account.
+   *
+   * @param delegate The public key (32 bytes) of the delegate.
+   */
+  function removeProxy(bytes32 delegate) external;
+
   /**
   * @dev Returns the stake amount associated with the specified `hotkey` and `coldkey`.
   *

runtime/src/precompiles/staking.rs

@@ -33,10 +33,13 @@ use pallet_evm::{
 use sp_core::crypto::Ss58Codec;
 use sp_core::U256;
 use sp_runtime::traits::Dispatchable;
-use sp_runtime::traits::{BlakeTwo256, UniqueSaturatedInto};
+use sp_runtime::traits::{BlakeTwo256, StaticLookup, UniqueSaturatedInto};
 use sp_runtime::AccountId32;
 
-use crate::precompiles::{get_method_id, get_slice};
+use crate::{
+    precompiles::{get_method_id, get_slice},
+    ProxyType,
+};
 use sp_std::vec;
 
 use crate::{Runtime, RuntimeCall};
@@ -52,24 +55,25 @@ impl StakingPrecompile {
             .get(4..)
             .map_or_else(vec::Vec::new, |slice| slice.to_vec()); // Avoiding borrowing conflicts
 
-        match method_id {
-            id if id == get_method_id("addStake(bytes32,uint256)") => {
-                Self::add_stake(handle, &method_input)
-            }
-            id if id == get_method_id("removeStake(bytes32,uint256,uint256)") => {
-                Self::remove_stake(handle, &method_input)
-            }
-            id if id == get_method_id("getStake(bytes32,bytes32,uint256)") => {
-                Self::get_stake(&method_input)
-            }
-            _ => Err(PrecompileFailure::Error {
+        if method_id == get_method_id("addStake(bytes32,uint256)") {
+            Self::add_stake(handle, &method_input)
+        } else if method_id == get_method_id("removeStake(bytes32,uint256,uint256)") {
+            Self::remove_stake(handle, &method_input)
+        } else if method_id == get_method_id("getStake(bytes32,bytes32,uint256)") {
+            Self::get_stake(&method_input)
+        } else if method_id == get_method_id("addProxy(bytes32)") {
+            Self::add_proxy(handle, &method_input)
+        } else if method_id == get_method_id("removeProxy(bytes32)") {
+            Self::remove_proxy(handle, &method_input)
+        } else {
+            Err(PrecompileFailure::Error {
                 exit_status: ExitError::InvalidRange,
-            }),
+            })
         }
     }
 
     fn add_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
-        let hotkey = Self::parse_hotkey(data)?.into();
+        let hotkey = Self::parse_pub_key(data)?.into();
         let amount: U256 = handle.context().apparent_value;
         let netuid = Self::parse_netuid(data, 0x3E)?;
 
@@ -88,7 +92,7 @@ impl StakingPrecompile {
     }
 
     fn remove_stake(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
-        let hotkey = Self::parse_hotkey(data)?.into();
+        let hotkey = Self::parse_pub_key(data)?.into();
         let netuid = Self::parse_netuid(data, 0x5E)?;
 
         // We have to treat this as uint256 (because of Solidity ABI encoding rules, it pads uint64),
@@ -110,6 +114,30 @@ impl StakingPrecompile {
         Self::dispatch(handle, call)
     }
 
+    fn add_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+        let delegate = AccountId32::from(Self::parse_pub_key(data)?);
+        let delegate = <Runtime as frame_system::Config>::Lookup::unlookup(delegate);
+        let call = RuntimeCall::Proxy(pallet_proxy::Call::<Runtime>::add_proxy {
+            delegate,
+            proxy_type: ProxyType::Staking,
+            delay: 0,
+        });
+
+        Self::dispatch(handle, call)
+    }
+
+    fn remove_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+        let delegate = AccountId32::from(Self::parse_pub_key(data)?);
+        let delegate = <Runtime as frame_system::Config>::Lookup::unlookup(delegate);
+        let call = RuntimeCall::Proxy(pallet_proxy::Call::<Runtime>::remove_proxy {
+            delegate,
+            proxy_type: ProxyType::Staking,
+            delay: 0,
+        });
+
+        Self::dispatch(handle, call)
+    }
+
     fn get_stake(data: &[u8]) -> PrecompileResult {
         let (hotkey, coldkey) = Self::parse_hotkey_coldkey(data)?;
         let netuid = Self::parse_netuid(data, 0x5E)?;
@@ -149,15 +177,15 @@ impl StakingPrecompile {
         Ok((hotkey, coldkey))
     }
 
-    fn parse_hotkey(data: &[u8]) -> Result<[u8; 32], PrecompileFailure> {
+    fn parse_pub_key(data: &[u8]) -> Result<[u8; 32], PrecompileFailure> {
         if data.len() < 32 {
             return Err(PrecompileFailure::Error {
                 exit_status: ExitError::InvalidRange,
             });
         }
-        let mut hotkey = [0u8; 32];
-        hotkey.copy_from_slice(get_slice(data, 0, 32)?);
-        Ok(hotkey)
+        let mut pubkey = [0u8; 32];
+        pubkey.copy_from_slice(get_slice(data, 0, 32)?);
+        Ok(pubkey)
     }
 
     fn parse_netuid(data: &[u8], offset: usize) -> Result<u16, PrecompileFailure> {