pallet-xvm refactor follow-up

chain-extensions/types/xvm/src/lib.rs

@@ -38,10 +38,11 @@ impl From<CallError> for XvmExecutionResult {
         // `0` is reserved for `Ok`
         let error_code = match input {
             InvalidVmId => 1,
-            SameVmCallNotAllowed => 2,
+            SameVmCallDenied => 2,
             InvalidTarget => 3,
             InputTooLarge => 4,
-            ExecutionFailed(_) => 5,
+            ReentranceDenied => 5,
+            ExecutionFailed(_) => 6,
         };
         Self::Err(error_code)
     }

pallets/xvm/Cargo.toml

@@ -7,6 +7,7 @@ homepage.workspace = true
 repository.workspace = true
 
 [dependencies]
+environmental = { workspace = true }
 log = { workspace = true }
 serde = { workspace = true, optional = true }
 
@@ -42,6 +43,7 @@ sp-io = { workspace = true }
 [features]
 default = ["std"]
 std = [
+	"environmental/std",
 	"log/std",
 	"parity-scale-codec/std",
 	"frame-support/std",

pallets/xvm/src/lib.rs

@@ -37,7 +37,7 @@
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
-use frame_support::{ensure, traits::Currency};
+use frame_support::{ensure, traits::Currency, weights::Weight};
 use pallet_contracts::{CollectEvents, DebugInfo, Determinism};
 use pallet_evm::GasWeightMapping;
 use parity_scale_codec::Decode;
@@ -66,6 +66,8 @@ pub use pallet::*;
 
 pub type WeightInfoOf<T> = <T as Config>::WeightInfo;
 
+environmental::thread_local_impl!(static IN_XVM: environmental::RefCell<bool> = environmental::RefCell::new(false));
+
 #[frame_support::pallet]
 pub mod pallet {
     use super::*;
@@ -120,25 +122,53 @@ where
         value: Balance,
         skip_execution: bool,
     ) -> CallResult {
+        let overheads = match vm_id {
+            VmId::Evm => WeightInfoOf::<T>::evm_call_overheads(),
+            VmId::Wasm => WeightInfoOf::<T>::wasm_call_overheads(),
+        };
+
         ensure!(
             context.source_vm_id != vm_id,
             CallErrorWithWeight {
-                error: CallError::SameVmCallNotAllowed,
-                used_weight: match vm_id {
-                    VmId::Evm => WeightInfoOf::<T>::evm_call_overheads(),
-                    VmId::Wasm => WeightInfoOf::<T>::wasm_call_overheads(),
-                },
+                error: CallError::SameVmCallDenied,
+                used_weight: overheads,
             }
         );
 
-        match vm_id {
-            VmId::Evm => {
-                Pallet::<T>::evm_call(context, source, target, input, value, skip_execution)
-            }
-            VmId::Wasm => {
-                Pallet::<T>::wasm_call(context, source, target, input, value, skip_execution)
-            }
+        // Set `IN_XVM` to true & check reentrance.
+        if IN_XVM.with(|in_xvm| in_xvm.replace(true)) {
+            return Err(CallErrorWithWeight {
+                error: CallError::ReentranceDenied,
+                used_weight: overheads,
+            });
         }
+
+        let res = match vm_id {
+            VmId::Evm => Pallet::<T>::evm_call(
+                context,
+                source,
+                target,
+                input,
+                value,
+                overheads,
+                skip_execution,
+            ),
+            VmId::Wasm => Pallet::<T>::wasm_call(
+                context,
+                source,
+                target,
+                input,
+                value,
+                overheads,
+                skip_execution,
+            ),
+        };
+
+        // Set `IN_XVM` to false.
+        // We should make sure that this line is executed whatever the execution path.
+        let _ = IN_XVM.with(|in_xvm| in_xvm.take());
+
+        res
     }
 
     fn evm_call(
@@ -147,6 +177,7 @@ where
         target: Vec<u8>,
         input: Vec<u8>,
         value: Balance,
+        overheads: Weight,
         skip_execution: bool,
     ) -> CallResult {
         log::trace!(
@@ -159,24 +190,22 @@ where
             target.len() == H160::len_bytes(),
             CallErrorWithWeight {
                 error: CallError::InvalidTarget,
-                used_weight: WeightInfoOf::<T>::evm_call_overheads(),
+                used_weight: overheads,
             }
         );
         let target_decoded =
             Decode::decode(&mut target.as_ref()).map_err(|_| CallErrorWithWeight {
                 error: CallError::InvalidTarget,
-                used_weight: WeightInfoOf::<T>::evm_call_overheads(),
+                used_weight: overheads,
             })?;
         let bounded_input = EthereumTxInput::try_from(input).map_err(|_| CallErrorWithWeight {
             error: CallError::InputTooLarge,
-            used_weight: WeightInfoOf::<T>::evm_call_overheads(),
+            used_weight: overheads,
         })?;
 
         let value_u256 = U256::from(value);
         // With overheads, less weight is available.
-        let weight_limit = context
-            .weight_limit
-            .saturating_sub(WeightInfoOf::<T>::evm_call_overheads());
+        let weight_limit = context.weight_limit.saturating_sub(overheads);
         let gas_limit = U256::from(T::GasWeightMapping::weight_to_gas(weight_limit));
 
         let source = T::AccountMapping::into_h160(source);
@@ -193,7 +222,7 @@ where
         if skip_execution {
             return Ok(CallInfo {
                 output: vec![],
-                used_weight: WeightInfoOf::<T>::evm_call_overheads(),
+                used_weight: overheads,
             });
         }
 
@@ -208,7 +237,7 @@ where
                 let used_weight = post_dispatch_info
                     .actual_weight
                     .unwrap_or_default()
-                    .saturating_add(WeightInfoOf::<T>::evm_call_overheads());
+                    .saturating_add(overheads);
                 CallInfo {
                     output: call_info.value,
                     used_weight,
@@ -219,7 +248,7 @@ where
                     .post_info
                     .actual_weight
                     .unwrap_or_default()
-                    .saturating_add(WeightInfoOf::<T>::evm_call_overheads());
+                    .saturating_add(overheads);
                 CallErrorWithWeight {
                     error: CallError::ExecutionFailed(Into::<&str>::into(e.error).into()),
                     used_weight,
@@ -233,6 +262,7 @@ where
         target: Vec<u8>,
         input: Vec<u8>,
         value: Balance,
+        overheads: Weight,
         skip_execution: bool,
     ) -> CallResult {
         log::trace!(
@@ -244,23 +274,21 @@ where
         let dest = {
             let error = CallErrorWithWeight {
                 error: CallError::InvalidTarget,
-                used_weight: WeightInfoOf::<T>::wasm_call_overheads(),
+                used_weight: overheads,
             };
             let decoded = Decode::decode(&mut target.as_ref()).map_err(|_| error.clone())?;
             T::Lookup::lookup(decoded).map_err(|_| error)
         }?;
 
         // With overheads, less weight is available.
-        let weight_limit = context
-            .weight_limit
-            .saturating_sub(WeightInfoOf::<T>::wasm_call_overheads());
+        let weight_limit = context.weight_limit.saturating_sub(overheads);
 
         // Note the skip execution check should be exactly before `pallet_contracts::bare_call`
         // to benchmark the correct overheads.
         if skip_execution {
             return Ok(CallInfo {
                 output: vec![],
-                used_weight: WeightInfoOf::<T>::wasm_call_overheads(),
+                used_weight: overheads,
             });
         }
 
@@ -277,9 +305,7 @@ where
         );
         log::trace!(target: "xvm::wasm_call", "WASM call result: {:?}", call_result);
 
-        let used_weight = call_result
-            .gas_consumed
-            .saturating_add(WeightInfoOf::<T>::wasm_call_overheads());
+        let used_weight = call_result.gas_consumed.saturating_add(overheads);
         match call_result.result {
             Ok(success) => Ok(CallInfo {
                 output: success.data,

pallets/xvm/src/tests.rs

@@ -49,7 +49,7 @@ fn calling_into_same_vm_is_not_allowed() {
                 value
             ),
             CallErrorWithWeight {
-                error: CallError::SameVmCallNotAllowed,
+                error: CallError::SameVmCallDenied,
                 used_weight: evm_used_weight
             },
         );
@@ -66,7 +66,7 @@ fn calling_into_same_vm_is_not_allowed() {
         assert_noop!(
             Xvm::call(wasm_context, wasm_vm_id, ALICE, wasm_target, input, value),
             CallErrorWithWeight {
-                error: CallError::SameVmCallNotAllowed,
+                error: CallError::SameVmCallDenied,
                 used_weight: wasm_used_weight
             },
         );

precompiles/utils/src/testing.rs

@@ -303,6 +303,11 @@ impl<'p, P: PrecompileSet> PrecompilesTester<'p, P> {
         self
     }
 
+    pub fn with_gas_limit(mut self, gas_limit: u64) -> Self {
+        self.handle.gas_limit = gas_limit;
+        self
+    }
+
     pub fn expect_cost(mut self, cost: u64) -> Self {
         self.expected_cost = Some(cost);
         self

precompiles/xvm/Cargo.toml

@@ -8,6 +8,7 @@ homepage.workspace = true
 repository.workspace = true
 
 [dependencies]
+hex = { workspace = true }
 log = { workspace = true }
 num_enum = { workspace = true }
 precompile-utils = { workspace = true }

precompiles/xvm/src/lib.rs

@@ -83,8 +83,12 @@ where
             id.try_into().map_err(|_| revert("invalid vm id"))
         }?;
 
-        let remaining_gas = handle.remaining_gas();
-        let weight_limit = R::GasWeightMapping::gas_to_weight(remaining_gas, true);
+        let mut gas_limit = handle.remaining_gas();
+        // If user specified a gas limit, make sure it's not exceeded.
+        if let Some(user_limit) = handle.gas_limit() {
+            gas_limit = gas_limit.min(user_limit);
+        }
+        let weight_limit = R::GasWeightMapping::gas_to_weight(gas_limit, true);
         let xvm_context = Context {
             source_vm_id: VmId::Evm,
             weight_limit,

precompiles/xvm/src/mock.rs

@@ -38,6 +38,7 @@ use sp_runtime::{
     testing::Header,
     traits::{BlakeTwo256, IdentityLookup},
 };
+use sp_std::cell::RefCell;
 
 use astar_primitives::xvm::{CallError::*, CallErrorWithWeight, CallInfo, CallResult};
 
@@ -236,10 +237,29 @@ impl pallet_evm::Config for Runtime {
     type GasLimitPovSizeRatio = ConstU64<4>;
 }
 
+thread_local! {
+    static WEIGHT_LIMIT: RefCell<Weight> = RefCell::new(Weight::zero());
+}
+
+pub(crate) struct WeightLimitCalledWith;
+impl WeightLimitCalledWith {
+    pub(crate) fn get() -> Weight {
+        WEIGHT_LIMIT.with(|gas_limit| *gas_limit.borrow())
+    }
+
+    pub(crate) fn set(value: Weight) {
+        WEIGHT_LIMIT.with(|gas_limit| *gas_limit.borrow_mut() = value)
+    }
+
+    pub(crate) fn reset() {
+        Self::set(Weight::zero());
+    }
+}
+
 struct MockXvmWithArgsCheck;
 impl XvmCall<AccountId> for MockXvmWithArgsCheck {
     fn call(
-        _context: Context,
+        context: Context,
         vm_id: VmId,
         _source: AccountId,
         target: Vec<u8>,
@@ -249,7 +269,7 @@ impl XvmCall<AccountId> for MockXvmWithArgsCheck {
         ensure!(
             vm_id != VmId::Evm,
             CallErrorWithWeight {
-                error: SameVmCallNotAllowed,
+                error: SameVmCallDenied,
                 used_weight: Weight::zero()
             }
         );
@@ -267,6 +287,9 @@ impl XvmCall<AccountId> for MockXvmWithArgsCheck {
                 used_weight: Weight::zero()
             }
         );
+
+        WeightLimitCalledWith::set(context.weight_limit);
+
         Ok(CallInfo {
             output: vec![],
             used_weight: Weight::zero(),
@@ -297,6 +320,8 @@ impl ExtBuilder {
             .build_storage::<Runtime>()
             .expect("Frame system builds valid default genesis config");
 
+        WeightLimitCalledWith::reset();
+
         let mut ext = sp_io::TestExternalities::new(t);
         ext.execute_with(|| System::set_block_number(1));
         ext