fix re-entrancy

crates/cheatcodes/src/config.rs

@@ -1,6 +1,6 @@
 use super::Result;
 use crate::{
-    strategy::{CheatcodeInspectorStrategy, EvmCheatcodeInspectorStrategy},
+    strategy::{new_evm_strategy, Strategy},
     Vm::Rpc,
 };
 use alloy_primitives::{map::AddressHashMap, U256};
@@ -57,7 +57,7 @@ pub struct CheatsConfig {
     /// Version of the script/test contract which is currently running.
     pub running_version: Option<Version>,
     /// The behavior strategy.
-    pub strategy: Box<dyn CheatcodeInspectorStrategy>,
+    pub strategy: Strategy,
     /// Whether to enable legacy (non-reverting) assertions.
     pub assertions_revert: bool,
     /// Optional seed for the RNG algorithm.
@@ -73,7 +73,7 @@ impl CheatsConfig {
         available_artifacts: Option<ContractsByArtifact>,
         running_contract: Option<String>,
         running_version: Option<Version>,
-        strategy: Box<dyn CheatcodeInspectorStrategy>,
+        strategy: Strategy,
     ) -> Self {
         let mut allowed_paths = vec![config.root.0.clone()];
         allowed_paths.extend(config.libs.clone());
@@ -234,7 +234,7 @@ impl Default for CheatsConfig {
             available_artifacts: Default::default(),
             running_contract: Default::default(),
             running_version: Default::default(),
-            strategy: Box::new(EvmCheatcodeInspectorStrategy::default()),
+            strategy: new_evm_strategy(),
             assertions_revert: true,
             seed: None,
         }
@@ -253,7 +253,7 @@ mod tests {
             None,
             None,
             None,
-            Box::new(EvmCheatcodeInspectorStrategy::default()),
+            new_evm_strategy(),
         )
     }
 

crates/cheatcodes/src/evm.rs

@@ -64,7 +64,7 @@ impl Cheatcode for getNonce_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account } = self;
 
-        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_get_nonce(ccx, *account))
+        ccx.state.strategy.inner.clone().cheatcode_get_nonce(ccx, *account)
     }
 }
 
@@ -350,7 +350,7 @@ impl Cheatcode for getBlobhashesCall {
 impl Cheatcode for rollCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { newHeight } = self;
-        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_roll(ccx, *newHeight))
+        ccx.state.strategy.inner.clone().cheatcode_roll(ccx, *newHeight)
     }
 }
 
@@ -372,7 +372,7 @@ impl Cheatcode for txGasPriceCall {
 impl Cheatcode for warpCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { newTimestamp } = self;
-        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_warp(ccx, *newTimestamp))
+        ccx.state.strategy.inner.clone().cheatcode_warp(ccx, *newTimestamp)
     }
 }
 
@@ -407,40 +407,38 @@ impl Cheatcode for dealCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account: address, newBalance: new_balance } = *self;
 
-        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_deal(ccx, address, new_balance))
+        ccx.state.strategy.inner.clone().cheatcode_deal(ccx, address, new_balance)
     }
 }
 
 impl Cheatcode for etchCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target, newRuntimeBytecode } = self;
 
-        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_etch(ccx, *target, newRuntimeBytecode))
+        ccx.state.strategy.inner.clone().cheatcode_etch(ccx, *target, newRuntimeBytecode)
     }
 }
 
 impl Cheatcode for resetNonceCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account } = self;
-        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_reset_nonce(ccx, *account))
+        ccx.state.strategy.inner.clone().cheatcode_reset_nonce(ccx, *account)
     }
 }
 
 impl Cheatcode for setNonceCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account, newNonce } = *self;
 
-        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_set_nonce(ccx, account, newNonce))
+        ccx.state.strategy.inner.clone().cheatcode_set_nonce(ccx, account, newNonce)
     }
 }
 
 impl Cheatcode for setNonceUnsafeCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account, newNonce } = *self;
 
-        ccx.with_strategy(|strategy, ccx| {
-            strategy.cheatcode_set_nonce_unsafe(ccx, account, newNonce)
-        })
+        ccx.state.strategy.inner.clone().cheatcode_set_nonce_unsafe(ccx, account, newNonce)
     }
 }
 

crates/cheatcodes/src/evm/fork.rs

@@ -125,7 +125,11 @@ impl Cheatcode for selectForkCall {
         persist_caller(ccx);
         check_broadcast(ccx.state)?;
 
-        ccx.with_strategy(|strategy, ccx| strategy.zksync_select_fork_vm(ccx.ecx, *forkId));
+        ccx.state.strategy.inner.zksync_select_fork_vm(
+            ccx.state.strategy.context.as_mut(),
+            ccx.ecx,
+            *forkId,
+        );
         ccx.ecx.db.select_fork(*forkId, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?;
 
         Ok(Default::default())
@@ -281,7 +285,11 @@ fn create_select_fork(ccx: &mut CheatsCtxt, url_or_alias: &str, block: Option<u6
 
     let fork = create_fork_request(ccx, url_or_alias, block)?;
     let id = ccx.ecx.db.create_fork(fork)?;
-    ccx.with_strategy(|strategy, ccx| strategy.zksync_select_fork_vm(ccx.ecx, id));
+    ccx.state.strategy.inner.zksync_select_fork_vm(
+        ccx.state.strategy.context.as_mut(),
+        ccx.ecx,
+        id,
+    );
     ccx.ecx.db.select_fork(id, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?;
     Ok(id.abi_encode())
 }

crates/cheatcodes/src/evm/mock.rs

@@ -15,9 +15,7 @@ impl Cheatcode for clearMockedCallsCall {
 impl Cheatcode for mockCall_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { callee, data, returnData } = self;
-        ccx.with_strategy(|strategy, ccx| {
-            strategy.cheatcode_mock_call(ccx, *callee, data, returnData)
-        })
+        ccx.state.strategy.inner.clone().cheatcode_mock_call(ccx, *callee, data, returnData)
     }
 }
 
@@ -85,9 +83,7 @@ impl Cheatcode for mockCalls_1Call {
 impl Cheatcode for mockCallRevert_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { callee, data, revertData } = self;
-        ccx.with_strategy(|strategy, ccx| {
-            strategy.cheatcode_mock_call_revert(ccx, *callee, data, revertData)
-        })
+        ccx.state.strategy.inner.clone().cheatcode_mock_call_revert(ccx, *callee, data, revertData)
     }
 }
 

crates/cheatcodes/src/fs.rs

@@ -283,7 +283,7 @@ impl Cheatcode for getArtifactPathByDeployedCodeCall {
 impl Cheatcode for getCodeCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { artifactPath: path } = self;
-        state.with_strategy(|strategy, state| strategy.get_artifact_code(state, path, false))
+        state.strategy.inner.get_artifact_code(state, path, false)
     }
 }
 

crates/cheatcodes/src/inspector.rs

@@ -7,7 +7,7 @@ use crate::{
         DealRecord, GasRecord,
     },
     script::{Broadcast, Wallets},
-    strategy::CheatcodeInspectorStrategy,
+    strategy::Strategy,
     test::{
         assume::AssumeNoRevert,
         expect::{self, ExpectedEmit, ExpectedRevert, ExpectedRevertKind},
@@ -69,7 +69,6 @@ pub use utils::CommonCreateInput;
 
 pub type Ecx<'a, 'b, 'c> = &'a mut EvmContext<&'b mut (dyn DatabaseExt + 'c)>;
 pub type InnerEcx<'a, 'b, 'c> = &'a mut InnerEvmContext<&'b mut (dyn DatabaseExt + 'c)>;
-pub type Strategy<'a> = &'a mut dyn CheatcodeInspectorStrategy;
 
 /// Helper trait for obtaining complete [revm::Inspector] instance from mutable reference to
 /// [Cheatcodes].
@@ -528,7 +527,7 @@ pub struct Cheatcodes {
     pub wallets: Option<Wallets>,
 
     /// The behavior strategy.
-    pub strategy: Option<Box<dyn CheatcodeInspectorStrategy>>,
+    pub strategy: Strategy,
 }
 
 impl Clone for Cheatcodes {
@@ -568,7 +567,7 @@ impl Clone for Cheatcodes {
             arbitrary_storage: self.arbitrary_storage.clone(),
             deprecated: self.deprecated.clone(),
             wallets: self.wallets.clone(),
-            strategy: self.strategy.as_ref().map(|s| s.new_cloned()),
+            strategy: self.strategy.clone(),
         }
     }
 }
@@ -588,7 +587,7 @@ impl Cheatcodes {
         Self {
             fs_commit: true,
             labels: config.labels.clone(),
-            strategy: Some(config.strategy.clone()),
+            strategy: config.strategy.clone(),
             config,
             block: Default::default(),
             active_delegation: Default::default(),
@@ -763,7 +762,8 @@ impl Cheatcodes {
                 if ecx_inner.journaled_state.depth() == broadcast.depth {
                     input.set_caller(broadcast.new_origin);
 
-                    self.strategy.as_mut().unwrap().record_broadcastable_create_transactions(
+                    self.strategy.inner.record_broadcastable_create_transactions(
+                        self.strategy.context.as_mut(),
                         self.config.clone(),
                         &input,
                         ecx_inner,
@@ -801,9 +801,9 @@ impl Cheatcodes {
             }]);
         }
 
-        if let Some(result) = self.with_strategy(|strategy, cheatcodes| {
-            strategy.zksync_try_create(cheatcodes, ecx, &input, executor)
-        }) {
+        if let Some(result) =
+            self.strategy.inner.clone().zksync_try_create(self, ecx, &input, executor)
+        {
             return Some(result);
         }
 
@@ -917,10 +917,7 @@ where {
             }
         }
 
-        self.strategy
-            .as_mut()
-            .expect("failed acquiring strategy")
-            .zksync_record_create_address(&outcome);
+        self.strategy.inner.zksync_record_create_address(self.strategy.context.as_mut(), &outcome);
 
         outcome
     }
@@ -963,10 +960,12 @@ where {
             let prev = account.info.nonce;
             let nonce = prev.saturating_sub(1);
             account.info.nonce = nonce;
-            self.strategy
-                .as_mut()
-                .expect("failed acquiring strategy")
-                .zksync_sync_nonce(sender, nonce, ecx);
+            self.strategy.inner.zksync_sync_nonce(
+                self.strategy.context.as_mut(),
+                sender,
+                nonce,
+                ecx,
+            );
 
             trace!(target: "cheatcodes", %sender, nonce, prev, "corrected nonce");
         }
@@ -1000,10 +999,7 @@ where {
             return None;
         }
 
-        self.strategy
-            .as_mut()
-            .expect("failed acquiring strategy")
-            .zksync_set_deployer_call_input(call);
+        self.strategy.inner.zksync_set_deployer_call_input(self.strategy.context.as_mut(), call);
 
         // Handle expected calls
 
@@ -1137,17 +1133,15 @@ where {
                         })
                     }
 
-                    self.strategy
-                        .as_mut()
-                        .expect("failed acquiring strategy")
-                        .record_broadcastable_call_transactions(
-                            self.config.clone(),
-                            call,
-                            ecx_inner,
-                            broadcast,
-                            &mut self.broadcastable_transactions,
-                            &mut self.active_delegation,
-                        );
+                    self.strategy.inner.record_broadcastable_call_transactions(
+                        self.strategy.context.as_mut(),
+                        self.config.clone(),
+                        call,
+                        ecx_inner,
+                        broadcast,
+                        &mut self.broadcastable_transactions,
+                        &mut self.active_delegation,
+                    );
 
                     let account =
                         ecx_inner.journaled_state.state().get_mut(&broadcast.new_origin).unwrap();
@@ -1220,9 +1214,8 @@ where {
             }]);
         }
 
-        if let Some(result) = self.with_strategy(|strategy, cheatcodes| {
-            strategy.zksync_try_call(cheatcodes, ecx, call, executor)
-        }) {
+        if let Some(result) = self.strategy.inner.clone().zksync_try_call(self, ecx, call, executor)
+        {
             return Some(result);
         }
 
@@ -1264,17 +1257,6 @@ where {
             None => false,
         }
     }
-
-    pub fn with_strategy<F, R>(&mut self, mut f: F) -> R
-    where
-        F: FnMut(Strategy, &mut Self) -> R,
-    {
-        let mut strategy = self.strategy.take();
-        let result = f(strategy.as_mut().expect("failed acquiring strategy").as_mut(), self);
-        self.strategy = strategy;
-
-        result
-    }
 }
 
 impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
@@ -1294,10 +1276,11 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
             self.gas_metering.paused_frames.push(interpreter.gas);
         }
 
-        self.strategy
-            .as_mut()
-            .expect("failed acquiring strategy")
-            .post_initialize_interp(interpreter, ecx);
+        self.strategy.inner.post_initialize_interp(
+            self.strategy.context.as_mut(),
+            interpreter,
+            ecx,
+        );
     }
 
     #[inline]
@@ -1342,8 +1325,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
 
     #[inline]
     fn step_end(&mut self, interpreter: &mut Interpreter, ecx: Ecx) {
-        if self.strategy.as_mut().expect("failed acquiring strategy").pre_step_end(interpreter, ecx)
-        {
+        if self.strategy.inner.pre_step_end(self.strategy.context.as_mut(), interpreter, ecx) {
             return;
         }
 

crates/cheatcodes/src/lib.rs

@@ -17,7 +17,6 @@ extern crate tracing;
 
 use alloy_primitives::Address;
 use foundry_evm_core::backend::DatabaseExt;
-use inspector::Strategy;
 use revm::{ContextPrecompiles, InnerEvmContext};
 use spec::Status;
 
@@ -175,15 +174,4 @@ impl CheatsCtxt<'_, '_, '_, '_> {
     pub(crate) fn is_precompile(&self, address: &Address) -> bool {
         self.precompiles.contains(address)
     }
-
-    pub(crate) fn with_strategy<F, R>(&mut self, mut f: F) -> R
-    where
-        F: FnMut(Strategy, &mut Self) -> R,
-    {
-        let mut strategy = self.state.strategy.take();
-        let result = f(strategy.as_mut().expect("failed acquiring strategy").as_mut(), self);
-        self.state.strategy = strategy;
-
-        result
-    }
 }