Add counter to wasmv1 and test interface

massalabs · Jul 16, 2024 · ce91892 · ce91892
1 parent 3cb1caf
commit ce91892
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/src/tests/mod.rs b/src/tests/mod.rs
@@ -18,6 +18,16 @@ impl InterfaceClone for TestInterface {
 }
 
 impl Interface for TestInterface {
+    fn increment_recursion_counter(&self) -> Result<()> {
+        println!("Increment recursion counter");
+        Ok(())
+    }
+
+    fn decrement_recursion_counter(&self) -> Result<()> {
+        println!("Decrement recursion counter");
+        Ok(())
+    }
+
     fn init_call(&self, address: &str, raw_coins: u64) -> Result<Vec<u8>> {
         println!("Init call to {}, with {} coins", address, raw_coins);
         Ok(vec![])

diff --git a/src/wasmv1_execution/abi/abis.rs b/src/wasmv1_execution/abi/abis.rs
@@ -152,6 +152,9 @@ fn abi_call(store_env: FunctionEnvMut<ABIEnv>, arg_offset: i32) -> Result<i32, W
             let remaining_gas = handler.get_remaining_gas();
             let interface = handler.exec_env.get_interface();
             let module = helper_get_module(interface, bytecode, remaining_gas)?;
+            interface.increment_recursion_counter().map_err(|e| {
+                WasmV1Error::RuntimeError(format!("Could not increment recursion counter: {}", e))
+            })?;
             let response = crate::execution::run_function(
                 interface,
                 module,
@@ -161,6 +164,9 @@ fn abi_call(store_env: FunctionEnvMut<ABIEnv>, arg_offset: i32) -> Result<i32, W
                 handler.get_gas_costs().clone(),
             )
             .map_err(|err| WasmV1Error::RuntimeError(format!("Could not run function: {}", err)))?;
+            interface.decrement_recursion_counter().map_err(|e| {
+                WasmV1Error::RuntimeError(format!("Could not decrement recursion counter: {}", e))
+            })?;
             handler.set_remaining_gas(response.remaining_gas);
             let interface = handler.exec_env.get_interface();
             interface.finish_call().map_err(|err| {
@@ -201,7 +207,9 @@ fn abi_local_call(store_env: FunctionEnvMut<ABIEnv>, arg_offset: i32) -> Result<
             let remaining_gas = handler.get_remaining_gas();
             let interface = handler.exec_env.get_interface();
             let module = helper_get_module(interface, bytecode.clone(), remaining_gas)?;
-
+            interface.increment_recursion_counter().map_err(|e| {
+                WasmV1Error::RuntimeError(format!("Could not increment recursion counter: {}", e))
+            })?;
             let response = crate::execution::run_function(
                 interface,
                 module,
@@ -211,6 +219,9 @@ fn abi_local_call(store_env: FunctionEnvMut<ABIEnv>, arg_offset: i32) -> Result<
                 handler.get_gas_costs().clone(),
             )
             .map_err(|err| WasmV1Error::RuntimeError(format!("Could not run function: {}", err)))?;
+            interface.decrement_recursion_counter().map_err(|e| {
+                WasmV1Error::RuntimeError(format!("Could not decrement recursion counter: {}", e))
+            })?;
             handler.set_remaining_gas(response.remaining_gas);
 
             #[cfg(feature = "execution-trace")]
@@ -1022,6 +1033,9 @@ fn abi_local_execution(
             let module = helper_get_tmp_module(handler, req.bytecode.clone(), remaining_gas)?;
 
             let interface = handler.exec_env.get_interface();
+            interface.increment_recursion_counter().map_err(|e| {
+                WasmV1Error::RuntimeError(format!("Could not increment recursion counter: {}", e))
+            })?;
             match crate::execution::run_function(
                 interface,
                 module,
@@ -1031,6 +1045,12 @@ fn abi_local_execution(
                 handler.get_gas_costs().clone(),
             ) {
                 Ok(response) => {
+                    interface.decrement_recursion_counter().map_err(|e| {
+                        WasmV1Error::RuntimeError(format!(
+                            "Could not decrement recursion counter: {}",
+                            e
+                        ))
+                    })?;
                     handler.set_remaining_gas(response.remaining_gas);
 
                     #[cfg(feature = "execution-trace")]