Skip to content

Commit 6234614

Browse files
nbaztecnbaztec
authored
feat: fix re-entrancy in strategies (#801)
* fix re-entrancy
1 parent 5353a10 commit 6234614

File tree

31 files changed

+1060
-542
lines changed

31 files changed

+1060
-542
lines changed

‎crates/cheatcodes/src/config.rs‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,5 @@
11
use super::Result;
2-
use crate::{
3-
strategy::{CheatcodeInspectorStrategy, EvmCheatcodeInspectorStrategy},
4-
Vm::Rpc,
5-
};
2+
use crate::{strategy::CheatcodeInspectorStrategy, Vm::Rpc};
63
use alloy_primitives::{map::AddressHashMap, U256};
74
use foundry_common::{fs::normalize_path, ContractsByArtifact};
85
use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
@@ -57,7 +54,7 @@ pub struct CheatsConfig {
5754
/// Version of the script/test contract which is currently running.
5855
pub running_version: Option<Version>,
5956
/// The behavior strategy.
60-
pub strategy: Box<dyn CheatcodeInspectorStrategy>,
57+
pub strategy: CheatcodeInspectorStrategy,
6158
/// Whether to enable legacy (non-reverting) assertions.
6259
pub assertions_revert: bool,
6360
/// Optional seed for the RNG algorithm.
@@ -73,7 +70,7 @@ impl CheatsConfig {
7370
available_artifacts: Option<ContractsByArtifact>,
7471
running_contract: Option<String>,
7572
running_version: Option<Version>,
76-
strategy: Box<dyn CheatcodeInspectorStrategy>,
73+
strategy: CheatcodeInspectorStrategy,
7774
) -> Self {
7875
let mut allowed_paths = vec![config.root.0.clone()];
7976
allowed_paths.extend(config.libs.clone());
@@ -234,7 +231,7 @@ impl Default for CheatsConfig {
234231
available_artifacts: Default::default(),
235232
running_contract: Default::default(),
236233
running_version: Default::default(),
237-
strategy: Box::new(EvmCheatcodeInspectorStrategy::default()),
234+
strategy: CheatcodeInspectorStrategy::new_evm(),
238235
assertions_revert: true,
239236
seed: None,
240237
}
@@ -253,7 +250,7 @@ mod tests {
253250
None,
254251
None,
255252
None,
256-
Box::new(EvmCheatcodeInspectorStrategy::default()),
253+
CheatcodeInspectorStrategy::new_evm(),
257254
)
258255
}
259256

‎crates/cheatcodes/src/evm.rs‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ impl Cheatcode for getNonce_0Call {
6464
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
6565
let Self { account } = self;
6666

67-
ccx.with_strategy(|strategy, ccx| strategy.cheatcode_get_nonce(ccx, *account))
67+
ccx.state.strategy.runner.clone().cheatcode_get_nonce(ccx, *account)
6868
}
6969
}
7070

@@ -350,7 +350,7 @@ impl Cheatcode for getBlobhashesCall {
350350
impl Cheatcode for rollCall {
351351
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
352352
let Self { newHeight } = self;
353-
ccx.with_strategy(|strategy, ccx| strategy.cheatcode_roll(ccx, *newHeight))
353+
ccx.state.strategy.runner.clone().cheatcode_roll(ccx, *newHeight)
354354
}
355355
}
356356

@@ -372,7 +372,7 @@ impl Cheatcode for txGasPriceCall {
372372
impl Cheatcode for warpCall {
373373
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
374374
let Self { newTimestamp } = self;
375-
ccx.with_strategy(|strategy, ccx| strategy.cheatcode_warp(ccx, *newTimestamp))
375+
ccx.state.strategy.runner.clone().cheatcode_warp(ccx, *newTimestamp)
376376
}
377377
}
378378

@@ -407,40 +407,38 @@ impl Cheatcode for dealCall {
407407
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
408408
let Self { account: address, newBalance: new_balance } = *self;
409409

410-
ccx.with_strategy(|strategy, ccx| strategy.cheatcode_deal(ccx, address, new_balance))
410+
ccx.state.strategy.runner.clone().cheatcode_deal(ccx, address, new_balance)
411411
}
412412
}
413413

414414
impl Cheatcode for etchCall {
415415
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
416416
let Self { target, newRuntimeBytecode } = self;
417417

418-
ccx.with_strategy(|strategy, ccx| strategy.cheatcode_etch(ccx, *target, newRuntimeBytecode))
418+
ccx.state.strategy.runner.clone().cheatcode_etch(ccx, *target, newRuntimeBytecode)
419419
}
420420
}
421421

422422
impl Cheatcode for resetNonceCall {
423423
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
424424
let Self { account } = self;
425-
ccx.with_strategy(|strategy, ccx| strategy.cheatcode_reset_nonce(ccx, *account))
425+
ccx.state.strategy.runner.clone().cheatcode_reset_nonce(ccx, *account)
426426
}
427427
}
428428

429429
impl Cheatcode for setNonceCall {
430430
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
431431
let Self { account, newNonce } = *self;
432432

433-
ccx.with_strategy(|strategy, ccx| strategy.cheatcode_set_nonce(ccx, account, newNonce))
433+
ccx.state.strategy.runner.clone().cheatcode_set_nonce(ccx, account, newNonce)
434434
}
435435
}
436436

437437
impl Cheatcode for setNonceUnsafeCall {
438438
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
439439
let Self { account, newNonce } = *self;
440440

441-
ccx.with_strategy(|strategy, ccx| {
442-
strategy.cheatcode_set_nonce_unsafe(ccx, account, newNonce)
443-
})
441+
ccx.state.strategy.runner.clone().cheatcode_set_nonce_unsafe(ccx, account, newNonce)
444442
}
445443
}
446444

‎crates/cheatcodes/src/evm/fork.rs‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,11 @@ impl Cheatcode for selectForkCall {
125125
persist_caller(ccx);
126126
check_broadcast(ccx.state)?;
127127

128-
ccx.with_strategy(|strategy, ccx| strategy.zksync_select_fork_vm(ccx.ecx, *forkId));
128+
ccx.state.strategy.runner.zksync_select_fork_vm(
129+
ccx.state.strategy.context.as_mut(),
130+
ccx.ecx,
131+
*forkId,
132+
);
129133
ccx.ecx.db.select_fork(*forkId, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?;
130134

131135
Ok(Default::default())
@@ -281,7 +285,11 @@ fn create_select_fork(ccx: &mut CheatsCtxt, url_or_alias: &str, block: Option<u6
281285

282286
let fork = create_fork_request(ccx, url_or_alias, block)?;
283287
let id = ccx.ecx.db.create_fork(fork)?;
284-
ccx.with_strategy(|strategy, ccx| strategy.zksync_select_fork_vm(ccx.ecx, id));
288+
ccx.state.strategy.runner.zksync_select_fork_vm(
289+
ccx.state.strategy.context.as_mut(),
290+
ccx.ecx,
291+
id,
292+
);
285293
ccx.ecx.db.select_fork(id, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?;
286294
Ok(id.abi_encode())
287295
}

‎crates/cheatcodes/src/evm/mock.rs‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,7 @@ impl Cheatcode for clearMockedCallsCall {
1515
impl Cheatcode for mockCall_0Call {
1616
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
1717
let Self { callee, data, returnData } = self;
18-
ccx.with_strategy(|strategy, ccx| {
19-
strategy.cheatcode_mock_call(ccx, *callee, data, returnData)
20-
})
18+
ccx.state.strategy.runner.clone().cheatcode_mock_call(ccx, *callee, data, returnData)
2119
}
2220
}
2321

@@ -85,9 +83,7 @@ impl Cheatcode for mockCalls_1Call {
8583
impl Cheatcode for mockCallRevert_0Call {
8684
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
8785
let Self { callee, data, revertData } = self;
88-
ccx.with_strategy(|strategy, ccx| {
89-
strategy.cheatcode_mock_call_revert(ccx, *callee, data, revertData)
90-
})
86+
ccx.state.strategy.runner.clone().cheatcode_mock_call_revert(ccx, *callee, data, revertData)
9187
}
9288
}
9389

‎crates/cheatcodes/src/fs.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -283,7 +283,7 @@ impl Cheatcode for getArtifactPathByDeployedCodeCall {
283283
impl Cheatcode for getCodeCall {
284284
fn apply(&self, state: &mut Cheatcodes) -> Result {
285285
let Self { artifactPath: path } = self;
286-
state.with_strategy(|strategy, state| strategy.get_artifact_code(state, path, false))
286+
state.strategy.runner.get_artifact_code(state, path, false)
287287
}
288288
}
289289

‎crates/cheatcodes/src/inspector.rs‎

Lines changed: 34 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,6 @@ pub use utils::CommonCreateInput;
6969

7070
pub type Ecx<'a, 'b, 'c> = &'a mut EvmContext<&'b mut (dyn DatabaseExt + 'c)>;
7171
pub type InnerEcx<'a, 'b, 'c> = &'a mut InnerEvmContext<&'b mut (dyn DatabaseExt + 'c)>;
72-
pub type Strategy<'a> = &'a mut dyn CheatcodeInspectorStrategy;
7372

7473
/// Helper trait for obtaining complete [revm::Inspector] instance from mutable reference to
7574
/// [Cheatcodes].
@@ -528,7 +527,7 @@ pub struct Cheatcodes {
528527
pub wallets: Option<Wallets>,
529528

530529
/// The behavior strategy.
531-
pub strategy: Option<Box<dyn CheatcodeInspectorStrategy>>,
530+
pub strategy: CheatcodeInspectorStrategy,
532531
}
533532

534533
impl Clone for Cheatcodes {
@@ -568,7 +567,7 @@ impl Clone for Cheatcodes {
568567
arbitrary_storage: self.arbitrary_storage.clone(),
569568
deprecated: self.deprecated.clone(),
570569
wallets: self.wallets.clone(),
571-
strategy: self.strategy.as_ref().map(|s| s.new_cloned()),
570+
strategy: self.strategy.clone(),
572571
}
573572
}
574573
}
@@ -588,7 +587,7 @@ impl Cheatcodes {
588587
Self {
589588
fs_commit: true,
590589
labels: config.labels.clone(),
591-
strategy: Some(config.strategy.clone()),
590+
strategy: config.strategy.clone(),
592591
config,
593592
block: Default::default(),
594593
active_delegation: Default::default(),
@@ -763,7 +762,8 @@ impl Cheatcodes {
763762
if ecx_inner.journaled_state.depth() == broadcast.depth {
764763
input.set_caller(broadcast.new_origin);
765764

766-
self.strategy.as_mut().unwrap().record_broadcastable_create_transactions(
765+
self.strategy.runner.record_broadcastable_create_transactions(
766+
self.strategy.context.as_mut(),
767767
self.config.clone(),
768768
&input,
769769
ecx_inner,
@@ -801,9 +801,9 @@ impl Cheatcodes {
801801
}]);
802802
}
803803

804-
if let Some(result) = self.with_strategy(|strategy, cheatcodes| {
805-
strategy.zksync_try_create(cheatcodes, ecx, &input, executor)
806-
}) {
804+
if let Some(result) =
805+
self.strategy.runner.clone().zksync_try_create(self, ecx, &input, executor)
806+
{
807807
return Some(result);
808808
}
809809

@@ -917,10 +917,7 @@ where {
917917
}
918918
}
919919

920-
self.strategy
921-
.as_mut()
922-
.expect("failed acquiring strategy")
923-
.zksync_record_create_address(&outcome);
920+
self.strategy.runner.zksync_record_create_address(self.strategy.context.as_mut(), &outcome);
924921

925922
outcome
926923
}
@@ -963,10 +960,12 @@ where {
963960
let prev = account.info.nonce;
964961
let nonce = prev.saturating_sub(1);
965962
account.info.nonce = nonce;
966-
self.strategy
967-
.as_mut()
968-
.expect("failed acquiring strategy")
969-
.zksync_sync_nonce(sender, nonce, ecx);
963+
self.strategy.runner.zksync_sync_nonce(
964+
self.strategy.context.as_mut(),
965+
sender,
966+
nonce,
967+
ecx,
968+
);
970969

971970
trace!(target: "cheatcodes", %sender, nonce, prev, "corrected nonce");
972971
}
@@ -1000,10 +999,7 @@ where {
1000999
return None;
10011000
}
10021001

1003-
self.strategy
1004-
.as_mut()
1005-
.expect("failed acquiring strategy")
1006-
.zksync_set_deployer_call_input(call);
1002+
self.strategy.runner.zksync_set_deployer_call_input(self.strategy.context.as_mut(), call);
10071003

10081004
// Handle expected calls
10091005

@@ -1137,17 +1133,15 @@ where {
11371133
})
11381134
}
11391135

1140-
self.strategy
1141-
.as_mut()
1142-
.expect("failed acquiring strategy")
1143-
.record_broadcastable_call_transactions(
1144-
self.config.clone(),
1145-
call,
1146-
ecx_inner,
1147-
broadcast,
1148-
&mut self.broadcastable_transactions,
1149-
&mut self.active_delegation,
1150-
);
1136+
self.strategy.runner.record_broadcastable_call_transactions(
1137+
self.strategy.context.as_mut(),
1138+
self.config.clone(),
1139+
call,
1140+
ecx_inner,
1141+
broadcast,
1142+
&mut self.broadcastable_transactions,
1143+
&mut self.active_delegation,
1144+
);
11511145

11521146
let account =
11531147
ecx_inner.journaled_state.state().get_mut(&broadcast.new_origin).unwrap();
@@ -1220,9 +1214,9 @@ where {
12201214
}]);
12211215
}
12221216

1223-
if let Some(result) = self.with_strategy(|strategy, cheatcodes| {
1224-
strategy.zksync_try_call(cheatcodes, ecx, call, executor)
1225-
}) {
1217+
if let Some(result) =
1218+
self.strategy.runner.clone().zksync_try_call(self, ecx, call, executor)
1219+
{
12261220
return Some(result);
12271221
}
12281222

@@ -1264,17 +1258,6 @@ where {
12641258
None => false,
12651259
}
12661260
}
1267-
1268-
pub fn with_strategy<F, R>(&mut self, mut f: F) -> R
1269-
where
1270-
F: FnMut(Strategy, &mut Self) -> R,
1271-
{
1272-
let mut strategy = self.strategy.take();
1273-
let result = f(strategy.as_mut().expect("failed acquiring strategy").as_mut(), self);
1274-
self.strategy = strategy;
1275-
1276-
result
1277-
}
12781261
}
12791262

12801263
impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
@@ -1294,10 +1277,11 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
12941277
self.gas_metering.paused_frames.push(interpreter.gas);
12951278
}
12961279

1297-
self.strategy
1298-
.as_mut()
1299-
.expect("failed acquiring strategy")
1300-
.post_initialize_interp(interpreter, ecx);
1280+
self.strategy.runner.post_initialize_interp(
1281+
self.strategy.context.as_mut(),
1282+
interpreter,
1283+
ecx,
1284+
);
13011285
}
13021286

13031287
#[inline]
@@ -1342,8 +1326,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
13421326

13431327
#[inline]
13441328
fn step_end(&mut self, interpreter: &mut Interpreter, ecx: Ecx) {
1345-
if self.strategy.as_mut().expect("failed acquiring strategy").pre_step_end(interpreter, ecx)
1346-
{
1329+
if self.strategy.runner.pre_step_end(self.strategy.context.as_mut(), interpreter, ecx) {
13471330
return;
13481331
}
13491332

‎crates/cheatcodes/src/lib.rs‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ extern crate tracing;
1717

1818
use alloy_primitives::Address;
1919
use foundry_evm_core::backend::DatabaseExt;
20-
use inspector::Strategy;
2120
use revm::{ContextPrecompiles, InnerEvmContext};
2221
use spec::Status;
2322

@@ -175,15 +174,4 @@ impl CheatsCtxt<'_, '_, '_, '_> {
175174
pub(crate) fn is_precompile(&self, address: &Address) -> bool {
176175
self.precompiles.contains(address)
177176
}
178-
179-
pub(crate) fn with_strategy<F, R>(&mut self, mut f: F) -> R
180-
where
181-
F: FnMut(Strategy, &mut Self) -> R,
182-
{
183-
let mut strategy = self.state.strategy.take();
184-
let result = f(strategy.as_mut().expect("failed acquiring strategy").as_mut(), self);
185-
self.state.strategy = strategy;
186-
187-
result
188-
}
189177
}

0 commit comments

Comments
 (0)