From 827a0132b4e2eb197abf0138ce51a64852344186 Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Fri, 3 May 2024 15:39:19 +0200
Subject: [PATCH] fix: Save evm context after errors

---
 src/evm.rs             |  6 +++---
 src/executor.rs        | 36 +++++++++++++++++++-----------------
 tests/fixtures/min.bin |  6 ++++++
 tests/test_evm.py      | 17 +++++++++++++++++
 4 files changed, 45 insertions(+), 20 deletions(-)
 create mode 100644 tests/fixtures/min.bin

diff --git a/src/evm.rs b/src/evm.rs
index 318bf96..8d0ddfe 100644
--- a/src/evm.rs
+++ b/src/evm.rs
@@ -373,9 +373,9 @@ impl EVM {
         let evm_context: EvmContext<DB> =
             replace(&mut self.context, EvmContext::new(DB::new_memory()));
         let (result, evm_context) =
-            call_evm(evm_context, self.handler_cfg, self.tracing, is_static)?;
+            call_evm(evm_context, self.handler_cfg, self.tracing, is_static);
         self.context = evm_context;
-        self.result = Some(result.clone());
-        Ok(result)
+        self.result = result.as_ref().map(|r| r.clone()).ok();
+        result
     }
 }
diff --git a/src/executor.rs b/src/executor.rs
index 9666187..6c208e2 100644
--- a/src/executor.rs
+++ b/src/executor.rs
@@ -1,18 +1,20 @@
-use crate::database::DB;
-use crate::utils::pyerr;
+use std::mem::replace;
+
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::PyResult;
+use revm::{
+    Context, ContextWithHandlerCfg, Evm, EvmContext, FrameOrResult, FrameResult,
+    inspector_handle_register,
+};
 use revm::inspectors::TracerEip3155;
 use revm::precompile::Log;
-use revm::primitives::TransactTo;
 use revm::primitives::{ExecutionResult, ShanghaiSpec};
-use revm::{
-    inspector_handle_register, Context, ContextWithHandlerCfg, Evm, EvmContext, FrameOrResult,
-    FrameResult,
-};
+use revm::primitives::TransactTo;
+use revm_interpreter::{CallInputs, CreateInputs, gas, SuccessOrHalt};
 use revm_interpreter::primitives::HandlerCfg;
-use revm_interpreter::{gas, CallInputs, CreateInputs, SuccessOrHalt};
-use std::mem::replace;
+
+use crate::database::DB;
+use crate::utils::pyerr;
 
 /// Calls the EVM with the given context and handler configuration.
 pub(crate) fn call_evm(
@@ -20,10 +22,10 @@ pub(crate) fn call_evm(
     handler_cfg: HandlerCfg,
     tracing: bool,
     is_static: bool,
-) -> PyResult<(ExecutionResult, EvmContext<DB>)> {
+) -> (PyResult<ExecutionResult>, EvmContext<DB>) {
     if tracing {
         let tracer = TracerEip3155::new(Box::new(crate::pystdout::PySysStdout {}));
-        let evm = Evm::builder()
+        let mut evm = Evm::builder()
             .with_context_with_handler_cfg(ContextWithHandlerCfg {
                 cfg: handler_cfg,
                 context: Context {
@@ -33,9 +35,9 @@ pub(crate) fn call_evm(
             })
             .append_handler_register(inspector_handle_register)
             .build();
-        run_evm(evm, is_static)
+        (run_evm(&mut evm, is_static), evm.context.evm)
     } else {
-        let evm = Evm::builder()
+        let mut evm = Evm::builder()
             .with_context_with_handler_cfg(ContextWithHandlerCfg {
                 cfg: handler_cfg,
                 context: Context {
@@ -44,15 +46,15 @@ pub(crate) fn call_evm(
                 },
             })
             .build();
-        run_evm(evm, is_static)
+        (run_evm(&mut evm, is_static), evm.context.evm)
     }
 }
 
 /// Calls the given evm. This is originally a copy of revm::Evm::transact, but it calls our own output function
 fn run_evm<EXT>(
-    mut evm: Evm<'_, EXT, DB>,
+    evm: &mut Evm<'_, EXT, DB>,
     is_static: bool,
-) -> PyResult<(ExecutionResult, EvmContext<DB>)> {
+) -> PyResult<ExecutionResult> {
     let logs_i = evm.context.evm.journaled_state.logs.len();
 
     evm.handler
@@ -137,7 +139,7 @@ fn run_evm<EXT>(
     let logs = ctx.evm.journaled_state.logs[logs_i..].to_vec();
 
     // Returns output of transaction.
-    Ok((output(ctx, result, logs)?, evm.context.evm))
+    Ok(output(ctx, result, logs)?)
 }
 
 fn call_inputs<EXT>(
diff --git a/tests/fixtures/min.bin b/tests/fixtures/min.bin
new file mode 100644
index 0000000..7bd5459
--- /dev/null
+++ b/tests/fixtures/min.bin
@@ -0,0 +1,6 @@
+61003e61000f60003961003e6000f35f3560e01c637ae2b5c781186100365760443610341761003a5760043560243580820382811161003a579050905060405260206040f35b5f5ffd5b5f80fd84183e8000a16576797065728300030b0012
+
+@external
+@view
+def min(x: uint256, y: uint256) -> uint256:
+    return x - y
\ No newline at end of file
diff --git a/tests/test_evm.py b/tests/test_evm.py
index a643b27..ed0aef0 100644
--- a/tests/test_evm.py
+++ b/tests/test_evm.py
@@ -280,3 +280,20 @@ def test_get_blobhashes(blob_hashes):
 
     # the contract logs 6 blob hashes, so pad with 0s
     assert logged == blob_hashes + [b"\0" * 32] * (6 - len(blob_hashes))
+
+
+@pytest.mark.parametrize("kwargs", KWARG_CASES)
+def test_call_reverting(kwargs):
+    evm = EVM()
+    code = load_contract_bin("min.bin")
+    deploy_address = evm.deploy(address, code)
+
+    with pytest.raises(RuntimeError) as err:
+        evm.message_call(
+            caller=address,
+            to=deploy_address,
+            value=10,
+        )
+
+    assert evm.get_code(deploy_address), "The code should still be deployed after revert"
+    assert str(err.value).startswith("Transaction(LackOfFundForMaxFee")