add netuid argument

opentensor · Nov 18, 2024 · 4f92d34 · 4f92d34
1 parent 52b5dc3
commit 4f92d34
Showing 1 changed file with 31 additions and 4 deletions.
diff --git a/runtime/src/precompiles/staking.rs b/runtime/src/precompiles/staking.rs
@@ -72,11 +72,17 @@ impl StakingPrecompile {
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
                 .ok_or(ExitError::OutOfFund)?;
 
+        let netuid =
+            Self::parse_netuid(data.get(56..64).unwrap_or(Err(PrecompileFailure::Error {
+                exit_status: ExitError::InvalidRange,
+            })?))?;
+
+        // let netuid_u16 = netuid.as_u32();
+
         // Create the add_stake call
         let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::add_stake {
             hotkey,
-            // TODO update contract to add netuid
-            netuid: 1,
+            netuid,
             amount_staked: amount_sub.unique_saturated_into(),
         });
         // Dispatch the add_stake call
@@ -96,10 +102,14 @@ impl StakingPrecompile {
             <Runtime as pallet_evm::Config>::BalanceConverter::into_substrate_balance(amount)
                 .ok_or(ExitError::OutOfFund)?;
 
+        let netuid =
+            Self::parse_netuid(data.get(64..72).unwrap_or(Err(PrecompileFailure::Error {
+                exit_status: ExitError::InvalidRange,
+            })?))?;
+
         let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::<Runtime>::remove_stake {
             hotkey,
-            // TODO update contract to add netuid
-            netuid: 1,
+            netuid,
             amount_unstaked: amount_sub.unique_saturated_into(),
         });
         Self::dispatch(handle, call)
@@ -116,6 +126,23 @@ impl StakingPrecompile {
         Ok(hotkey)
     }
 
+    fn parse_netuid(data: &[u8]) -> Result<u16, PrecompileFailure> {
+        let netuid = data
+            .get(0..8)
+            .map(U256::from_big_endian)
+            .ok_or(ExitError::InvalidRange)?;
+
+        let u16_max_u256 = U256::from(u16::MAX);
+        if netuid > u16_max_u256 {
+            // if netuid.as_u128() > u16::MAX as u128 {
+            return Err(PrecompileFailure::Error {
+                exit_status: ExitError::InvalidRange,
+            });
+        }
+
+        Ok(netuid.as_u32() as u16)
+    }
+
     fn dispatch(handle: &mut impl PrecompileHandle, call: RuntimeCall) -> PrecompileResult {
         let account_id =
             <HashedAddressMapping<BlakeTwo256> as AddressMapping<AccountId32>>::into_account_id(