diff --git a/runtime/src/precompiles/staking.rs b/runtime/src/precompiles/staking.rs index daa95ba93..8b91346ad 100644 --- a/runtime/src/precompiles/staking.rs +++ b/runtime/src/precompiles/staking.rs @@ -72,11 +72,14 @@ impl StakingPrecompile { ::BalanceConverter::into_substrate_balance(amount) .ok_or(ExitError::OutOfFund)?; + let netuid_vec = data.get(56..64).ok_or(PrecompileFailure::Error { + exit_status: ExitError::InvalidRange, + })?; + let netuid = Self::parse_netuid(netuid_vec)?; // Create the add_stake call let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::::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 +99,14 @@ impl StakingPrecompile { ::BalanceConverter::into_substrate_balance(amount) .ok_or(ExitError::OutOfFund)?; + let netuid_vec = data.get(88..96).ok_or(PrecompileFailure::Error { + exit_status: ExitError::InvalidRange, + })?; + let netuid = Self::parse_netuid(netuid_vec)?; + let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::::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 +123,26 @@ impl StakingPrecompile { Ok(hotkey) } + fn parse_netuid(data: &[u8]) -> Result { + let mut bytes = [0_u8; 8]; + bytes.copy_from_slice(data.get(0..8).ok_or(ExitError::InvalidRange)?); + + let netuid = U256::from_big_endian(&bytes); + + let u16_max_u256 = U256::from(u16::MAX); + + if netuid > u16_max_u256 { + return Err(PrecompileFailure::Error { + exit_status: ExitError::InvalidRange, + }); + } + let mut two_bytes = [0_u8; 2]; + two_bytes.copy_from_slice(&bytes[6..8]); + let result: u16 = u16::from_be_bytes(two_bytes); + + Ok(result) + } + fn dispatch(handle: &mut impl PrecompileHandle, call: RuntimeCall) -> PrecompileResult { let account_id = as AddressMapping>::into_account_id(