Add support for etch cheatcode

crates/forge/tests/it/revive/cheat_etch.rs

@@ -0,0 +1,16 @@
+use crate::{config::*, test_helpers::TEST_DATA_REVIVE};
+use foundry_test_utils::Filter;
+use revive_strategy::ReviveRuntimeMode;
+use revm::primitives::hardfork::SpecId;
+use rstest::rstest;
+
+#[rstest]
+#[case::pvm_mode_with_any_etched_evm_code(ReviveRuntimeMode::Pvm)]
+#[case::evm_mode_with_any_etched_evm_code(ReviveRuntimeMode::Evm)]
+#[tokio::test(flavor = "multi_thread")]
+async fn test_etch(#[case] runtime_mode: ReviveRuntimeMode) {
+    let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
+    let filter = Filter::new(".*", "EtchTest", ".*/revive/EtchTest.t.sol");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}

crates/forge/tests/it/revive/mod.rs

@@ -1,5 +1,6 @@
 //! Revive strategy tests
 
+pub mod cheat_etch;
 pub mod cheat_gas_metering;
 pub mod cheat_mock_call;
 pub mod cheat_mock_calls;

crates/revive-strategy/src/cheatcodes/mod.rs

@@ -8,8 +8,8 @@ use foundry_cheatcodes::{
     CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyRunner, CheatsConfig, CheatsCtxt,
     CommonCreateInput, DealRecord, Ecx, Error, EvmCheatcodeInspectorStrategyRunner, Result,
     Vm::{
-        dealCall, getNonce_0Call, loadCall, pvmCall, resetNonceCall, rollCall, setNonceCall,
-        setNonceUnsafeCall, storeCall, warpCall,
+        dealCall, etchCall, getNonce_0Call, loadCall, pvmCall, resetNonceCall, rollCall,
+        setNonceCall, setNonceUnsafeCall, storeCall, warpCall,
     },
     journaled_account, precompile_error,
 };
@@ -23,10 +23,11 @@ use std::{
 };
 use tracing::warn;
 
+use alloy_eips::eip7702::SignedAuthorization;
 use polkadot_sdk::{
     pallet_revive::{
         self, AccountInfo, AddressMapper, BalanceOf, BytecodeType, Code, ContractInfo, ExecConfig,
-        Pallet, evm::CallTrace,
+        Executable, Pallet, evm::CallTrace,
     },
     polkadot_sdk_frame::prelude::OriginFor,
     sp_core::{self, H160, H256},
@@ -41,7 +42,6 @@ use crate::{
 };
 use foundry_cheatcodes::Vm::{AccountAccess as FAccountAccess, ChainInfo};
 
-use alloy_eips::eip7702::SignedAuthorization;
 use revm::{
     bytecode::opcode as op,
     context::{CreateScheme, JournalTr},
@@ -201,6 +201,49 @@ fn set_block_number(new_height: U256, ecx: Ecx<'_, '_, '_>) {
     });
 }
 
+// Implements the `etch` cheatcode for PVM.
+fn etch_call(target: &Address, new_runtime_code: &Bytes, ecx: Ecx<'_, '_, '_>) -> Result {
+    let origin_address = H160::from_slice(ecx.tx.caller.as_slice());
+    let origin_account = AccountId::to_fallback_account_id(&origin_address);
+
+    execute_with_externalities(|externalities| {
+        externalities.execute_with(|| {
+            let code = new_runtime_code.to_vec();
+            let code_type =
+                if code.starts_with(b"PVM\0") { BytecodeType::Pvm } else { BytecodeType::Evm };
+            let contract_blob = Pallet::<Runtime>::try_upload_code(
+                origin_account.clone(),
+                code,
+                code_type,
+                BalanceOf::<Runtime>::MAX,
+                &ExecConfig::new_substrate_tx(),
+            )
+            .map_err(|_| <&str as Into<Error>>::into("Could not upload PVM code"))?
+            .0;
+
+            let mut contract_info = if let Some(contract_info) =
+                AccountInfo::<Runtime>::load_contract(&H160::from_slice(target.as_slice()))
+            {
+                contract_info
+            } else {
+                ContractInfo::<Runtime>::new(
+                    &origin_address,
+                    System::account_nonce(origin_account),
+                    *contract_blob.code_hash(),
+                )
+                .map_err(|_| <&str as Into<Error>>::into("Could not create contract info"))?
+            };
+            contract_info.code_hash = *contract_blob.code_hash();
+            AccountInfo::<Runtime>::insert_contract(
+                &H160::from_slice(target.as_slice()),
+                contract_info,
+            );
+            Ok::<(), Error>(())
+        })
+    })?;
+    Ok(Default::default())
+}
+
 fn set_timestamp(new_timestamp: U256, ecx: Ecx<'_, '_, '_>) {
     // Set timestamp in EVM context (seconds).
     ecx.block.timestamp = new_timestamp;
@@ -394,6 +437,12 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
 
                 Ok(Default::default())
             }
+            t if using_pvm && is::<etchCall>(t) => {
+                let etchCall { target, newRuntimeBytecode } =
+                    cheatcode.as_any().downcast_ref().unwrap();
+                etch_call(target, newRuntimeBytecode, ccx.ecx)?;
+                Ok(Default::default())
+            }
             t if using_pvm && is::<loadCall>(t) => {
                 tracing::info!(cheatcode = ?cheatcode.as_debug() , using_pvm = ?using_pvm);
                 let &loadCall { target, slot } = cheatcode.as_any().downcast_ref().unwrap();

testdata/default/revive/EtchTest.t.sol

@@ -0,0 +1,132 @@
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+import "../../default/logs/console.sol";
+
+contract Adder {
+    function add(uint256 a, uint256 b) public pure returns (uint256) {
+        return a * b; // Intentional bug to verify etch works
+    }
+}
+
+contract NestedAdder {
+    uint256 public inner_a;
+    uint256 public inner_b;
+
+    constructor(uint256 a, uint256 b) {
+        inner_a = a;
+        inner_b = b;
+    }
+
+    function nested_call(address target) public returns (uint256) {
+        // Perform the add call on the target contract address
+        (bool success, bytes memory data) =
+            target.call(abi.encodeWithSignature("add(uint256,uint256)", inner_a, inner_b));
+        require(success, "Nested call failed");
+        uint256 result = abi.decode(data, (uint256));
+        assert(success);
+        return result;
+    }
+}
+
+contract EtchTest is DSTest {
+    Vm constant vm = Vm(address(bytes20(uint160(uint256(keccak256("hevm cheat code"))))));
+    // This is the bytecode for the  correct Adder contract above compiled with resolc.
+    bytes constant code =
+        hex"50564d00008c0c000000000000010700c13000c00080047106000000000e0000001c0000002a000000350000004700000063616c6c5f646174615f636f707963616c6c5f646174615f6c6f616463616c6c5f646174615f73697a657365616c5f72657475726e7365745f696d6d757461626c655f6461746176616c75655f7472616e7366657272656405110287920463616c6c879e066465706c6f79068bec25028a531a001f004b007c00a500ae00c100cd00fa005e01630192019701bb0212036003b103c0033c048d04b2046c05ec05f5050106180627063d0660061a0769079e07ab07c607ca07080850099511f07b10087b158475010a02013d0700000251050750100209501004b3009511807b10787b15707b1668951580008411e04911384911304911208000330740951820491128501006e6084911584911504911484911408317400a0501821750821858821948821a40d49808d4a707d4870752072e6417501008ec008217188218108219088216d49707d48609d47909989920d48707977720d497075107090050100a61056467330850100c75098377330833090a283e076467330850100e62098378330733093300100a03019511a07b10587b15507b16489515608411e064164916384916304916208000330740956820491628501012370839070000025317045f9517e08477e07b67186471837733080a010182671880771c977720140800000000f7021677ab8736330014951120fe7b10d8017b15d0017b16c8019515e0018411e04921b8014921b0014921a8014921a0018317a0010a05288d02501016a3049511c07b10387b15307b16289515408411f0647664173300189511f07b10087b156475330820330740330048288f0850121a3a063200828910828a18828b088288d4ba0ad4980bd4ab0b98bb20d4a909979920d4b9095209449511c07b10387b15307b16289515408411e06476838883170a01821718821810821908821a7b67187b68107b69087b6a9551c08210388215308216289511403200009511e0fe7b1018017b1510017b160801951520018411f0828310828208829a08828c829b7b1b20829410d32a067b1638d82a06d8cb00821b38dab006828b188298187b1838c93409c969087b1828d869087b1330d8340664b4821938c9b909c96909c98909c92a08c908087b1c38821a20c9ca0a8e8b88aa4085aa01db8b0a8f98821c288ecbdb980bd49c08db8b0a510ace0064768217387b17c0007b127b12c8008217307b17d0009517e0009518c0007b14087b14d80050101ce3fe8217e8007b17288217e0007b17208217f8007b17188217f0007b1710821838958720d88708821ac88a09d8a90ada880a821830c88a0ad88a08821b08c88b0b7b1a707b19687b1760951780009518607b1b7850101e8cfe821790008218980082198000821a88007b67307b68387b69207b6a288217107b67108217187b67188217207b678217287b67089551e0fe821018018215100182160801951120013200501020a602828a10828b18828c088289d4cb0bd4a908d4b808988820d4ba0a97aa20d4a80852083f9511d07b10287b15209515308411f0827a18827810827b0882777b177b1b087b181064187b1a18649750102280059551d08210288215209511303200008218108217087b87088217187b877b861082177b87189551808210788215708216689511800032008217b0018218b8018219a801821aa001d49808d4a707d487075207e70138070000024921380149213001492128017b1720014921580149215001492140010495176001951840019519200149214801501026c7fd82126001821768018218700182197801821b8001821c880182169001821a98017b1ad8007b16d0007b1cc8007b1bc0007b19f8007b18f0007b17e800951700019518e0009519c0007b12e000501028d801821700017b1738821708017b1730821710017b1728821718017b17209517a00050102ab6fc8217b8007b17188216b0008218a8007b1810821aa0007b1a088219207b19588219287b19508219307b19488219387b19407b17787b16707b1868951780009518609519407b1a6033002c9511807b10787b15707b1668951580008411f07b171082878292828b08829308957a207b1a18d87a06c86b0a7b1a08d8ba0cda660c828a10828818829410829918c8ca06d8a60cc88c0c7b1c7b18387b1a307b1b287b17207b19587b14507b13489517409518207b124033002428f7fd821918821b10821008d49b07d46008d47808988820d46707977720d4870752075d646482178800821898007b183882138000821a9000d3b706d8b70cd80308da680cc94a06c9c602d8c606d84a0a821c38c99c0cc9ac0cc96c0cc9b707c98707c90306d4c707d42608d47808d42707988820977720d487075107080050102e19640764685010302e048378836933073300320a03019511f87b103308100002838833070133093300340a03013308491718491710491708490732004911184911104901113307046418491108501038f402390804000256183f0b200304000240013308100002838833070133092433003a0a03019511e87b10107b15087b16828b188294188282828c08829a088295828610829810c8ca09c82503d85305c85909d3a900d8a90ada050ac86805c85a0ad85a00c8b404d88508c84808c88000d86a05d3b008d8b00bda850bd3a606d46808d3c906d8c90cd82305db6c05db8b0552051b7b737b79087b7a107b70188210108215088216951118320033003c9511b07b10487b15409515508411f0491130491128491120140700000000717b484e9518207b173833073300362815029511807b10787b15707b1668951580008411f08282828308828410828818829a829b08829c1082991864767b19187b1c107b1b087b1a7b18387b14307b132895174095182064197b122050103efcfe821750821858821940821a487b67107b68187b697b6a089551808210788215708216689511800032009511f87b10330750104067f89511f87b103307015010425af89511c07b10387b15307b16289515408411f064766417501044a2f95012460632008217108218188219821a087b67107b68187b697b6a089551c0821038821530821628951140320239080800025108c0f8330730000383770a0428b3f87c78017c797c7a027c7b03978808d4980897aa1097bb18d4ba0ad4a8087c79057c7a047c7b067c7c07979908d4a90997bb1097cc18d4cb0bd4b909979920d489027c79097c7a087c7b0a7c7c0b979908d4a90997bb1097cc18d4cb0bd4b9097c7a0d7c7b0c7c7c0e7c780f97aa08d4ba0a97cc10978818d4c808d4a808978820d498037c78117c7a107c7b127c7c13978808d4a80897bb1097cc18d4cb0bd4b8087c7a157c7b147c7c167c791797aa08d4ba0a97cc10979918d4c909d4a909979920d4890a7c78197c79187c7b1a7c7c1b978808d4980897bb1097cc18d4cb0bd4b8087c791d7c7b1c7c7c1e7c771f979908d4b90997cc10977718d4c707d49707977720d487076f776fa86f396f2a7b5a187b59107b58087b57821008821595111032009511d87b10207b15187b161082897b19088289087b1982851082861833082050104ad3006f686f59821a6faa821b086fbb787b18787a10787908787898bc38787c1f98bc30787c1e98bc28787c1d98bc20787c1c98bc18787c1b98bc10787c1a98bb08787b1998ab38787b1798ab30787b1698ab28787b1598ab20787b1498ab18787b1398ab10787b1298aa08787a11989a38787a0f989a30787a0e989a28787a0d989a20787a0c989a18787a0b989a10787a0a98990878790998893878790798893078790698892878790598892078790498891878790398891078790298880878780182102082151882161095112832008b7910520931c8780883881f8488e05638000001253309040002390a040002ae8a093d080400020133081000028377c887073200004969488424892421494892344992a490a4244992a1423515aa3449929290248448523549004944442422224a4892a4925492849294244992244955929294246944442449254992244992a424499294a494244a1489884488442291242549922489244444244992a42422220a494a9224499224a9a41492a449524a924411114922444444848888101111111111111111112112894422442212499224491249922491549224499224494992244992244992a4884424499224499224499251a14a5328a525a9242921443094a4494a4a4a922449922449922449929224a52449920490942449a4242529494992129224294912a5424a4892a4429494a424490a112a294992244992244992244992244992244992244992244992244992244992244992aa24a5942425250955294949922449922449922449922449922449922449922449529224082161280500";
+    bytes constant code_evm =
+        hex"608060405234801561000f575f5ffd5b5060043610610029575f3560e01c8063771602f71461002d575b5f5ffd5b610047600480360381019061004291906100a9565b61005d565b60405161005491906100f6565b60405180910390f35b5f818361006a919061013c565b905092915050565b5f5ffd5b5f819050919050565b61008881610076565b8114610092575f5ffd5b50565b5f813590506100a38161007f565b92915050565b5f5f604083850312156100bf576100be610072565b5b5f6100cc85828601610095565b92505060206100dd85828601610095565b9150509250929050565b6100f081610076565b82525050565b5f6020820190506101095f8301846100e7565b92915050565b7f4e487b71000000000000000000000000000000000000000000000000000000005f52601160045260245ffd5b5f61014682610076565b915061015183610076565b92508282019050808211156101695761016861010f565b5b9291505056fea2646970667358221220bde9424751d367d702063695cb7d0afb42c0a83d370954296d07e2b1684208fb64736f6c634300081e0033";
+
+    // Test etching code into an existing contract instance works correctly.
+    function testEtchExistingContractPvmCode() public {
+        vm.pvm(true);
+        Adder adder = new Adder();
+
+        // Without etch, the add function is broken
+        uint256 buggy_result = adder.add(1, 2);
+        assertEq(buggy_result, 2);
+
+        // Etch the correct bytecode into the existing contract
+        vm.etch(address(adder), code);
+        uint256 result = adder.add(1, 2);
+        assertEq(result, 3);
+
+        // Verify that nested calls also work correctly after etch
+        uint256 nested_call_result = (new NestedAdder(1, 2)).nested_call(address(adder));
+        assertEq(nested_call_result, 3);
+    }
+
+    // Test etching code into an existing contract instance works correctly.
+    function testEtchExistingContractEvmCode() public {
+        vm.pvm(true);
+        Adder adder = new Adder();
+
+        // Without etch, the add function is broken
+        uint256 buggy_result = adder.add(1, 2);
+        assertEq(buggy_result, 2);
+
+        // Etch the correct bytecode into the existing contract
+        vm.etch(address(adder), code_evm);
+        uint256 result = adder.add(1, 2);
+        assertEq(result, 3);
+
+        // Verify that nested calls also work correctly after etch
+        uint256 nested_call_result = (new NestedAdder(1, 2)).nested_call(address(adder));
+        assertEq(nested_call_result, 3);
+    }
+
+    // Test etching code into any arbitrary address works correctly.
+    function testEtchAnyContractPvmCode() public {
+        vm.pvm(true);
+        // Etch the correct bytecode into an arbitrary address
+        address target = address(7070707);
+        vm.etch(target, code);
+        (bool success, bytes memory output) = target.call(abi.encodeWithSignature("add(uint256,uint256)", 1, 2));
+        uint256 result1 = abi.decode(output, (uint256));
+
+        assert(success);
+        assertEq(result1, 3);
+
+        uint256 nested_call_result = (new NestedAdder(1, 2)).nested_call(address(target));
+        assertEq(nested_call_result, 3);
+
+        // Etch into the zero address as well to verify it works for reserved addresses
+        address target2 = address(0);
+        vm.etch(target2, code);
+        (bool success2, bytes memory output2) = target2.call(abi.encodeWithSignature("add(uint256,uint256)", 1, 2));
+        uint256 result2 = abi.decode(output2, (uint256));
+
+        assert(success2);
+        assertEq(result2, 3);
+
+        uint256 nested_call_result2 = (new NestedAdder(1, 2)).nested_call(address(target2));
+        assertEq(nested_call_result2, 3);
+    }
+
+    // Test etching code into any arbitrary address works correctly.
+    function testEtchAnyContractEvmCode() public {
+        vm.pvm(true);
+        // Etch the correct bytecode into an arbitrary address
+        address target = address(7070707);
+        vm.etch(target, code);
+        (bool success, bytes memory output) = target.call(abi.encodeWithSignature("add(uint256,uint256)", 1, 2));
+        uint256 result1 = abi.decode(output, (uint256));
+
+        assert(success);
+        assertEq(result1, 3);
+
+        uint256 nested_call_result = (new NestedAdder(1, 2)).nested_call(address(target));
+        assertEq(nested_call_result, 3);
+
+        // Etch into the zero address as well to verify it works for reserved addresses
+        address target2 = address(0);
+        vm.etch(target2, code_evm);
+        (bool success2, bytes memory output2) = target2.call(abi.encodeWithSignature("add(uint256,uint256)", 1, 2));
+        uint256 result2 = abi.decode(output2, (uint256));
+
+        assert(success2);
+        assertEq(result2, 3);
+
+        uint256 nested_call_result2 = (new NestedAdder(1, 2)).nested_call(address(target2));
+        assertEq(nested_call_result2, 3);
+    }
+}