diff --git a/crates/revive-strategy/src/cheatcodes/mock_handler.rs b/crates/revive-strategy/src/cheatcodes/mock_handler.rs index e53480e805f16..c58d6d3448ff9 100644 --- a/crates/revive-strategy/src/cheatcodes/mock_handler.rs +++ b/crates/revive-strategy/src/cheatcodes/mock_handler.rs @@ -5,7 +5,7 @@ use std::{ }; use alloy_primitives::{Address, Bytes, map::foldhash::HashMap, ruint::aliases::U256}; -use foundry_cheatcodes::{Ecx, MockCallDataContext, MockCallReturnData}; +use foundry_cheatcodes::{DealRecord, Ecx, MockCallDataContext, MockCallReturnData}; use foundry_evm::constants::CHEATCODE_ADDRESS; use polkadot_sdk::{ frame_system, @@ -15,7 +15,7 @@ use polkadot_sdk::{ }, pallet_revive_uapi::ReturnFlags, polkadot_sdk_frame::prelude::OriginFor, - sp_core::H160, + sp_core::{H160, U256 as SpU256}, }; use revive_env::Runtime; @@ -60,17 +60,36 @@ impl MockHandlerImpl { state.mocked_functions = mock_inner.mocked_functions.clone(); } - pub(crate) fn fund_pranked_accounts(&self, account: Address) { - // Fuzzed prank addresses have no balance, so they won't exist in revive, and - // calls will fail, this is not a problem when running in REVM. - // TODO: Figure it out why this is still needed. - let balance = Pallet::::evm_balance(&H160::from_slice(account.as_slice())); - if balance == 0.into() { - Pallet::::set_evm_balance( - &H160::from_slice(account.as_slice()), - u128::MAX.into(), - ) - .expect("Could not fund pranked account"); + /// Syncs balances for pranked accounts between REVM and pallet-revive. + /// + /// If the account was explicitly dealt to via vm.deal(), sync that balance to pallet-revive. + /// This handles cases where vm.deal() was called in a callback and pallet-revive's balance + /// diverged from REVM's balance. + /// + /// If the account was NOT dealt to and has 0 balance, fund with u128::MAX so fuzzed + /// prank addresses can make calls in pallet-revive. + pub(crate) fn fund_pranked_accounts(account: Address, eth_deals: &[DealRecord]) { + let account_h160 = H160::from_slice(account.as_slice()); + + // Check if account was explicitly dealt to via vm.deal() + // Use the most recent deal record for this account + if let Some(deal) = eth_deals.iter().rev().find(|d| d.address == account) { + // Sync the dealt balance to pallet-revive + let target_balance = + SpU256::from_little_endian(&deal.new_balance.as_le_bytes()).min(u128::MAX.into()); + let pvm_balance = Pallet::::evm_balance(&account_h160); + if pvm_balance != target_balance { + Pallet::::set_evm_balance(&account_h160, target_balance) + .expect("Could not sync dealt account balance"); + } + return; + } + + // Account was not dealt to - fund with u128::MAX if balance is 0 + let pvm_balance = Pallet::::evm_balance(&account_h160); + if pvm_balance == 0.into() { + Pallet::::set_evm_balance(&account_h160, u128::MAX.into()) + .expect("Could not fund pranked account"); } } } diff --git a/crates/revive-strategy/src/cheatcodes/mod.rs b/crates/revive-strategy/src/cheatcodes/mod.rs index 5132f1fdaa81d..b65967f35113d 100644 --- a/crates/revive-strategy/src/cheatcodes/mod.rs +++ b/crates/revive-strategy/src/cheatcodes/mod.rs @@ -1001,6 +1001,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector sp_core::U256::from_little_endian(&U256::from(ecx.tx.gas_price).as_le_bytes()); let mut tracer = Tracer::new(state.expected_calls.clone(), state.expected_creates.clone()); let caller_h160 = H160::from_slice(input.caller().as_slice()); + let eth_deals = &state.eth_deals; let res = ctx.externalities.execute_with(|| { tracer.watch_address(&caller_h160); @@ -1011,7 +1012,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector AccountId32Mapper::::to_fallback_account_id(&caller_h160); let origin = OriginFor::::signed(origin_account_id.clone()); let evm_value = sp_core::U256::from_little_endian(&input.value().as_le_bytes()); - mock_handler.fund_pranked_accounts(input.caller()); + MockHandlerImpl::fund_pranked_accounts(input.caller(), eth_deals); if !exists { let nonce = ecx .journaled_state @@ -1185,6 +1186,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector let caller_h160 = H160::from_slice(call.caller.as_slice()); let mut tracer = Tracer::new(state.expected_calls.clone(), state.expected_creates.clone()); + let eth_deals = &state.eth_deals; let res = ctx.externalities.execute_with(|| { // Watch the caller's address so its nonce changes get tracked in prestate trace tracer.watch_address(&caller_h160); @@ -1193,7 +1195,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector let origin = OriginFor::::signed( AccountId32Mapper::::to_fallback_account_id(&caller_h160), ); - mock_handler.fund_pranked_accounts(call.caller); + MockHandlerImpl::fund_pranked_accounts(call.caller, eth_deals); let evm_value = sp_core::U256::from_little_endian(&call.call_value().as_le_bytes()); let target = H160::from_slice(call.target_address.as_slice()); diff --git a/testdata/default/revive/Prank.t.sol b/testdata/default/revive/Prank.t.sol index d41c5743e425d..54a1f806edb33 100644 --- a/testdata/default/revive/Prank.t.sol +++ b/testdata/default/revive/Prank.t.sol @@ -675,3 +675,40 @@ contract Counter { number++; } } + +contract FundPrankedAccountsReproTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + address alice = address(0xA11CE); + + function test_PrankDealtAccountWithZeroBalance() public { + vm.deal(alice, 3 ether); + assertEq(alice.balance, 3 ether, "alice should have 3 ether after deal"); + + vm.deal(alice, 0); + assertEq(alice.balance, 0, "alice should have 0 balance after second deal"); + + vm.startPrank(alice); + assertEq(alice.balance, 0, "alice balance should remain 0, not be reset to u128::MAX"); + vm.stopPrank(); + } + + function test_PrankFuzzedAddressGetsFunded() public { + address fuzzedAddr = address(0xF022ED); + uint256 initialBalance = fuzzedAddr.balance; + + vm.prank(fuzzedAddr); + + assertTrue(fuzzedAddr.balance >= initialBalance, "fuzzed address should be auto-funded"); + } + + function test_PrankDealtAccountWithBalance() public { + vm.deal(alice, 5 ether); + assertEq(alice.balance, 5 ether, "alice should have 5 ether"); + + vm.startPrank(alice); + assertEq(alice.balance, 5 ether, "alice balance should remain 5 ether"); + vm.stopPrank(); + } + +}