feat(chisel): add support for setting msg.data (foundry-rs#4775)

* feat(chisel): add support for setting `msg.data`

* reformat the code

* fix clippy error

* add missing "calldata" in chisel tests

* add better input trimming

* use "first" instead of "get(0)"

chisel/src/bin/chisel.rs

@@ -83,6 +83,7 @@ async fn main() -> eyre::Result<()> {
         foundry_config: config,
         evm_opts,
         backend: None,
+        calldata: None,
     })?;
 
     // Check for chisel subcommands

chisel/src/cmd.rs

@@ -36,6 +36,8 @@ pub enum ChiselCommand {
     Fork,
     /// Enable / disable traces for the current session
     Traces,
+    /// Set calldata (`msg.data`) for the current session (appended after function selector)
+    Calldata,
     /// Dump the raw memory
     MemDump,
     /// Dump the raw stack
@@ -69,6 +71,7 @@ impl FromStr for ChiselCommand {
             "clearcache" | "cc" => Ok(ChiselCommand::ClearCache),
             "fork" | "f" => Ok(ChiselCommand::Fork),
             "traces" | "t" => Ok(ChiselCommand::Traces),
+            "calldata" | "cd" => Ok(ChiselCommand::Calldata),
             "memdump" | "md" => Ok(ChiselCommand::MemDump),
             "stackdump" | "sd" => Ok(ChiselCommand::StackDump),
             "export" | "ex" => Ok(ChiselCommand::Export),
@@ -132,6 +135,7 @@ impl From<ChiselCommand> for CmdDescriptor {
             // Environment
             ChiselCommand::Fork => (&["fork <url>", "f <url>"], "Fork an RPC for the current session. Supply 0 arguments to return to a local network", CmdCategory::Env),
             ChiselCommand::Traces => (&["traces", "t"], "Enable / disable traces for the current session", CmdCategory::Env),
+            ChiselCommand::Calldata => (&["calldata [data]", "cd [data]"], "Set calldata (`msg.data`) for the current session (appended after function selector). Clears it if no argument provided.", CmdCategory::Env),
             // Debug
             ChiselCommand::MemDump => (&["memdump", "md"], "Dump the raw memory of the current state", CmdCategory::Debug),
             ChiselCommand::StackDump => (&["stackdump", "sd"], "Dump the raw stack of the current state", CmdCategory::Debug),

chisel/src/dispatcher.rs

@@ -379,6 +379,40 @@ impl ChiselDispatcher {
                     DispatchResult::CommandFailed(Self::make_error("Session not present."))
                 }
             }
+            ChiselCommand::Calldata => {
+                if let Some(session_source) = self.session.session_source.as_mut() {
+                    // remove empty space, double quotes, and 0x prefix
+                    let arg = args
+                        .first()
+                        .map(|s| {
+                            s.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\'')
+                        })
+                        .map(|s| s.strip_prefix("0x").unwrap_or(s))
+                        .unwrap_or("");
+
+                    if arg.is_empty() {
+                        session_source.config.calldata = None;
+                        return DispatchResult::CommandSuccess(Some("Calldata cleared.".to_string()))
+                    }
+
+                    let calldata = hex::decode(arg);
+                    match calldata {
+                        Ok(calldata) => {
+                            session_source.config.calldata = Some(calldata);
+                            DispatchResult::CommandSuccess(Some(format!(
+                                "Set calldata to '{}'",
+                                Paint::yellow(arg)
+                            )))
+                        }
+                        Err(e) => DispatchResult::CommandFailed(Self::make_error(format!(
+                            "Invalid calldata: {}",
+                            e
+                        ))),
+                    }
+                } else {
+                    DispatchResult::CommandFailed(Self::make_error("Session not present."))
+                }
+            }
             ChiselCommand::MemDump | ChiselCommand::StackDump => {
                 if let Some(session_source) = self.session.session_source.as_mut() {
                     match session_source.execute().await {

chisel/src/executor.rs

@@ -258,7 +258,7 @@ impl SessionSource {
 
         // Create a [ChiselRunner] with a default balance of [U256::MAX] and
         // the sender [Address::zero].
-        ChiselRunner::new(executor, U256::MAX, Address::zero())
+        ChiselRunner::new(executor, U256::MAX, Address::zero(), self.config.calldata.clone())
     }
 }
 

chisel/src/runner.rs

@@ -30,6 +30,8 @@ pub struct ChiselRunner {
     pub initial_balance: U256,
     /// The sender
     pub sender: Address,
+    /// Input calldata appended to `RUN_SELECTOR`
+    pub input: Option<Vec<u8>>,
 }
 
 /// Represents the result of a Chisel REPL run
@@ -64,8 +66,13 @@ impl ChiselRunner {
     /// ### Returns
     ///
     /// A new [ChiselRunner]
-    pub fn new(executor: Executor, initial_balance: U256, sender: Address) -> Self {
-        Self { executor, initial_balance, sender }
+    pub fn new(
+        executor: Executor,
+        initial_balance: U256,
+        sender: Address,
+        input: Option<Vec<u8>>,
+    ) -> Self {
+        Self { executor, initial_balance, sender, input }
     }
 
     /// Run a contract as a REPL session
@@ -93,8 +100,14 @@ impl ChiselRunner {
         // Reset the sender's balance to the initial balance for calls.
         self.executor.set_balance(self.sender, self.initial_balance)?;
 
+        // Append the input to the `RUN_SELECTOR` to form the calldata
+        let mut calldata = RUN_SELECTOR.to_vec();
+        if let Some(mut input) = self.input.clone() {
+            calldata.append(&mut input);
+        }
+
         // Call the "run()" function of the REPL contract
-        let call_res = self.call(self.sender, address, Bytes::from(RUN_SELECTOR), 0.into(), true);
+        let call_res = self.call(self.sender, address, Bytes::from(calldata), 0.into(), true);
 
         call_res.map(|res| (address, res))
     }

chisel/src/session_source.rs

@@ -74,6 +74,8 @@ pub struct SessionSourceConfig {
     pub backend: Option<Backend>,
     /// Optionally enable traces for the REPL contract execution
     pub traces: bool,
+    /// Optionally set calldata for the REPL contract execution
+    pub calldata: Option<Vec<u8>>,
 }
 
 impl SessionSourceConfig {

chisel/tests/cache.rs

@@ -43,6 +43,7 @@ fn test_write_session() {
         evm_opts: EvmOpts::default(),
         backend: None,
         traces: false,
+        calldata: None,
     })
     .unwrap_or_else(|_| panic!("Failed to create ChiselSession!"));