Extend staking contract with proxy setting

opentensor · Jan 2, 2025 · 3ad54bb · 3ad54bb
1 parent d261f67
commit 3ad54bb
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 17 deletions.
diff --git a/runtime/src/precompiles/solidity/staking.abi b/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": [
             {
@@ -37,7 +63,7 @@
         ],
         "name": "removeStake",
         "outputs": [],
-        "stateMutability": "payable",
+        "stateMutability": "nonpayable",
         "type": "function"
     }
 ]
diff --git a/runtime/src/precompiles/solidity/staking.sol b/runtime/src/precompiles/solidity/staking.sol
@@ -42,4 +42,18 @@ interface IStaking {
    * - The existing stake amount must be not lower than specified amount
    */
   function removeStake(bytes32 hotkey, uint256 amount, uint16 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;
 }
diff --git a/runtime/src/precompiles/staking.rs b/runtime/src/precompiles/staking.rs
@@ -36,7 +36,10 @@ use sp_runtime::traits::Dispatchable;
 use sp_runtime::traits::{BlakeTwo256, 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,21 +55,23 @@ 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,uint16)") => {
-                Self::add_stake(handle, &method_input)
-            }
-            id if id == get_method_id("removeStake(bytes32,uint256,uint16)") => {
-                Self::remove_stake(handle, &method_input)
-            }
-            _ => Err(PrecompileFailure::Error {
+        if method_id == get_method_id("addStake(bytes32,uint16)") {
+            Self::add_stake(handle, &method_input)
+        } else if method_id == get_method_id("removeStake(bytes32,uint256,uint16)") {
+            Self::remove_stake(handle, &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 amount_sub =
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
@@ -80,8 +85,9 @@ impl StakingPrecompile {
         // Dispatch the add_stake call
         Self::dispatch(handle, call)
     }
+
     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();
 
         // We have to treat this as uint256 (because of Solidity ABI encoding rules, it pads uint64),
         // but this will never exceed 8 bytes, se we will ignore higher bytes and will only use lower
@@ -101,15 +107,37 @@ impl StakingPrecompile {
         Self::dispatch(handle, call)
     }
 
-    fn parse_hotkey(data: &[u8]) -> Result<[u8; 32], PrecompileFailure> {
+    fn add_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
+        let delegate = Self::parse_pub_key(data)?;
+        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 = Self::parse_pub_key(data)?;
+        let call = RuntimeCall::Proxy(pallet_proxy::Call::<Runtime>::remove_proxy {
+            delegate,
+            proxy_type: ProxyType::Staking,
+            delay: 0,
+        });
+
+        Self::dispatch(handle, call)
+    }
+
+    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 dispatch(handle: &mut impl PrecompileHandle, call: RuntimeCall) -> PrecompileResult {