add interceptInitcode cheatcode (#10242)

* add interceptInitcode cheatcode

* rename intercept_next to intercept_next_create_call

* move logic to create_common

Author: tushar994

Diff for: crates/cheatcodes/assets/cheatcodes.json

@@ -2838,6 +2838,16 @@ interface Vm {
     /// Randomly shuffles an array.
     #[cheatcode(group = Utilities)]
     function shuffle(uint256[] calldata array) external returns (uint256[] memory);
+
+    /// Causes the next contract creation (via new) to fail and return its initcode in the returndata buffer.
+    /// This allows type-safe access to the initcode payload that would be used for contract creation.
+    /// Example usage:
+    /// vm.interceptInitcode();
+    /// bytes memory initcode;
+    /// try new MyContract(param1, param2) { assert(false); }
+    /// catch (bytes memory interceptedInitcode) { initcode = interceptedInitcode; }
+    #[cheatcode(group = Utilities, safety = Unsafe)]
+    function interceptInitcode() external;
 }
 }
 

Diff for: crates/cheatcodes/spec/src/vm.rs

@@ -489,6 +489,9 @@ pub struct Cheatcodes {
     /// `char -> (address, pc)`
     pub breakpoints: Breakpoints,
 
+    /// Whether the next contract creation should be intercepted to return its initcode.
+    pub intercept_next_create_call: bool,
+
     /// Optional cheatcodes `TestRunner`. Used for generating random values from uint and int
     /// strategies.
     test_runner: Option<TestRunner>,
@@ -549,6 +552,7 @@ impl Cheatcodes {
             mapping_slots: Default::default(),
             pc: Default::default(),
             breakpoints: Default::default(),
+            intercept_next_create_call: Default::default(),
             test_runner: Default::default(),
             ignored_traces: Default::default(),
             arbitrary_storage: Default::default(),
@@ -656,6 +660,25 @@ impl Cheatcodes {
     where
         Input: CommonCreateInput,
     {
+        // Check if we should intercept this create
+        if self.intercept_next_create_call {
+            // Reset the flag
+            self.intercept_next_create_call = false;
+
+            // Get initcode from the input
+            let output = input.init_code();
+
+            // Return a revert with the initcode as error data
+            return Some(CreateOutcome {
+                result: InterpreterResult {
+                    result: InstructionResult::Revert,
+                    output,
+                    gas: Gas::new(input.gas_limit()),
+                },
+                address: None,
+            });
+        }
+
         let ecx = &mut ecx.inner;
         let gas = Gas::new(input.gas_limit());
         let curr_depth = ecx.journaled_state.depth();

Diff for: crates/cheatcodes/src/inspector.rs

@@ -193,6 +193,18 @@ impl Cheatcode for resumeTracingCall {
     }
 }
 
+impl Cheatcode for interceptInitcodeCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self {} = self;
+        if !state.intercept_next_create_call {
+            state.intercept_next_create_call = true;
+        } else {
+            bail!("vm.interceptInitcode() has already been called")
+        }
+        Ok(Default::default())
+    }
+}
+
 impl Cheatcode for setArbitraryStorage_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target } = self;

Diff for: crates/cheatcodes/src/utils.rs

@@ -3411,3 +3411,113 @@ Selectors successfully uploaded to OpenChain
 
 "#]]);
 });
+
+// tests `interceptInitcode` function
+forgetest_init!(intercept_initcode, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.insert_ds_test();
+    prj.insert_vm();
+    prj.clear();
+
+    prj.add_source(
+        "InterceptInitcode.t.sol",
+        r#"
+import {Vm} from "./Vm.sol";
+import {DSTest} from "./test.sol";
+
+contract SimpleContract {
+    uint256 public value;
+    constructor(uint256 _value) {
+        value = _value;
+    }
+}
+
+contract InterceptInitcodeTest is DSTest {
+    Vm vm = Vm(HEVM_ADDRESS);
+
+    function testInterceptRegularCreate() public {
+        // Set up interception
+        vm.interceptInitcode();
+
+        // Try to create a contract - this should revert with the initcode
+        bytes memory initcode;
+        try new SimpleContract(42) {
+            assert(false);
+        } catch (bytes memory interceptedInitcode) {
+            initcode = interceptedInitcode;
+        }
+
+        // Verify the initcode contains the constructor argument
+        assertTrue(initcode.length > 0, "initcode should not be empty");
+
+        // The constructor argument is encoded as a 32-byte value at the end of the initcode
+        // We need to convert the last 32 bytes to uint256
+        uint256 value;
+        assembly {
+            value := mload(add(add(initcode, 0x20), sub(mload(initcode), 32)))
+        }
+        assertEq(value, 42, "initcode should contain constructor arg");
+    }
+
+    function testInterceptCreate2() public {
+        // Set up interception
+        vm.interceptInitcode();
+
+        // Try to create a contract with CREATE2 - this should revert with the initcode
+        bytes memory initcode;
+        try new SimpleContract(1337) {
+            assert(false);
+        } catch (bytes memory interceptedInitcode) {
+            initcode = interceptedInitcode;
+        }
+
+        // Verify the initcode contains the constructor argument
+        assertTrue(initcode.length > 0, "initcode should not be empty");
+
+        // The constructor argument is encoded as a 32-byte value at the end of the initcode
+        uint256 value;
+        assembly {
+            value := mload(add(add(initcode, 0x20), sub(mload(initcode), 32)))
+        }
+        assertEq(value, 1337, "initcode should contain constructor arg");
+    }
+
+    function testInterceptMultiple() public {
+        // First interception
+        vm.interceptInitcode();
+        bytes memory initcode1;
+        try new SimpleContract(1) {
+            assert(false);
+        } catch (bytes memory interceptedInitcode) {
+            initcode1 = interceptedInitcode;
+        }
+
+        // Second interception
+        vm.interceptInitcode();
+        bytes memory initcode2;
+        try new SimpleContract(2) {
+            assert(false);
+        } catch (bytes memory interceptedInitcode) {
+            initcode2 = interceptedInitcode;
+        }
+
+        // Verify different initcodes
+        assertTrue(initcode1.length > 0, "first initcode should not be empty");
+        assertTrue(initcode2.length > 0, "second initcode should not be empty");
+
+        // Extract constructor arguments from both initcodes
+        uint256 value1;
+        uint256 value2;
+        assembly {
+            value1 := mload(add(add(initcode1, 0x20), sub(mload(initcode1), 32)))
+            value2 := mload(add(add(initcode2, 0x20), sub(mload(initcode2), 32)))
+        }
+        assertEq(value1, 1, "first initcode should contain first arg");
+        assertEq(value2, 2, "second initcode should contain second arg");
+    }
+}
+     "#,
+    )
+    .unwrap();
+    cmd.args(["test", "-vvvvv"]).assert_success();
+});