Forbid XVM re-entrancy.

AstarNetwork · Aug 9, 2023 · b028df4 · b028df4
1 parent da91bca
commit b028df4
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/chain-extensions/types/xvm/src/lib.rs b/chain-extensions/types/xvm/src/lib.rs
@@ -42,7 +42,8 @@ impl From<CallError> for XvmExecutionResult {
             InvalidTarget => 3,
             InputTooLarge => 4,
             ExceedMaxWeightLimit => 5,
-            ExecutionFailed(_) => 6,
+            Reentrancy => 6,
+            ExecutionFailed(_) => 7,
         };
         Self::Err(error_code)
     }

diff --git a/pallets/xvm/Cargo.toml b/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",

diff --git a/pallets/xvm/src/lib.rs b/pallets/xvm/src/lib.rs
@@ -70,6 +70,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::*;
@@ -152,7 +154,15 @@ where
             }
         );
 
-        match vm_id {
+        // Set `IN_XVM` to true & check reentrancy.
+        if IN_XVM.with(|in_evm| in_evm.replace(true)) {
+            CallErrorWithWeight {
+                error: CallError::Reentrancy,
+                used_weight: overheads,
+            };
+        }
+
+        let res = match vm_id {
             VmId::Evm => Pallet::<T>::evm_call(
                 context,
                 source,
@@ -171,7 +181,13 @@ where
                 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_evm| in_evm.take());
+
+        res
     }
 
     fn evm_call(

diff --git a/primitives/src/xvm.rs b/primitives/src/xvm.rs
@@ -68,6 +68,8 @@ pub enum CallError {
     InputTooLarge,
     /// Exceed max weight limit.
     ExceedMaxWeightLimit,
+    /// Reentrancy is not allowed.
+    Reentrancy,
     /// The call failed on EVM or WASM execution.
     ExecutionFailed(Vec<u8>),
 }