feat: vm.sign for scripts (foundry-rs#7454)

* feat: vm.sign for script wallets

* more tests

* clippy

* if let some else

* review fixes

crates/cheatcodes/spec/src/vm.rs

@@ -249,6 +249,22 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Safe)]
     function sign(uint256 privateKey, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);
 
+    /// Signs `digest` with signer provided to script using the secp256k1 curve.
+    ///
+    /// If `--sender` is provided, the signer with provided address is used, otherwise,
+    /// if exactly one signer is provided to the script, that signer is used.
+    ///
+    /// Raises error if signer passed through `--sender` does not match any unlocked signers or
+    /// if `--sender` is not provided and not exactly one signer is passed to the script.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function sign(bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);
+
+    /// Signs `digest` with signer provided to script using the secp256k1 curve.
+    ///
+    /// Raises error if none of the signers passed into the script have provided address.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function sign(address signer, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);
+
     /// Signs `digest` with `privateKey` using the secp256r1 curve.
     #[cheatcode(group = Evm, safety = Safe)]
     function signP256(uint256 privateKey, bytes32 digest) external pure returns (bytes32 r, bytes32 s);
@@ -1563,8 +1579,12 @@ interface Vm {
 
     // -------- Broadcasting Transactions --------
 
-    /// Using the address that calls the test contract, has the next call (at this call depth only)
-    /// create a transaction that can later be signed and sent onchain.
+    /// Has the next call (at this call depth only) create transactions that can later be signed and sent onchain.
+    ///
+    /// Broadcasting address is determined by checking the following in order:
+    /// 1. If `--sender` argument was provided, that address is used.
+    /// 2. If exactly one signer (e.g. private key, hw wallet, keystore) is set when `forge broadcast` is invoked, that signer is used.
+    /// 3. Otherwise, default foundry sender (1804c8AB1F12E6bbf3894d4083f33e07309d1f38) is used.
     #[cheatcode(group = Scripting)]
     function broadcast() external;
 
@@ -1578,8 +1598,12 @@ interface Vm {
     #[cheatcode(group = Scripting)]
     function broadcast(uint256 privateKey) external;
 
-    /// Using the address that calls the test contract, has all subsequent calls
-    /// (at this call depth only) create transactions that can later be signed and sent onchain.
+    /// Has all subsequent calls (at this call depth only) create transactions that can later be signed and sent onchain.
+    ///
+    /// Broadcasting address is determined by checking the following in order:
+    /// 1. If `--sender` argument was provided, that address is used.
+    /// 2. If exactly one signer (e.g. private key, hw wallet, keystore) is set when `forge broadcast` is invoked, that signer is used.
+    /// 3. Otherwise, default foundry sender (1804c8AB1F12E6bbf3894d4083f33e07309d1f38) is used.
     #[cheatcode(group = Scripting)]
     function startBroadcast() external;
 

crates/cheatcodes/src/evm.rs

@@ -145,6 +145,20 @@ impl Cheatcode for sign_0Call {
     }
 }
 
+impl Cheatcode for sign_1Call {
+    fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { digest } = self;
+        super::utils::sign_with_wallet(ccx, None, digest)
+    }
+}
+
+impl Cheatcode for sign_2Call {
+    fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { signer, digest } = self;
+        super::utils::sign_with_wallet(ccx, Some(*signer), digest)
+    }
+}
+
 impl Cheatcode for signP256Call {
     fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { privateKey, digest } = self;

crates/cheatcodes/src/utils.rs

@@ -1,7 +1,7 @@
 //! Implementations of [`Utils`](crate::Group::Utils) cheatcodes.
 
 use crate::{Cheatcode, Cheatcodes, CheatsCtxt, DatabaseExt, Result, Vm::*};
-use alloy_primitives::{keccak256, B256, U256};
+use alloy_primitives::{keccak256, Address, B256, U256};
 use alloy_signer::{
     coins_bip39::{
         ChineseSimplified, ChineseTraditional, Czech, English, French, Italian, Japanese, Korean,
@@ -10,7 +10,8 @@ use alloy_signer::{
     LocalWallet, MnemonicBuilder, Signer, SignerSync,
 };
 use alloy_sol_types::SolValue;
-use foundry_evm_core::constants::DEFAULT_CREATE2_DEPLOYER;
+use foundry_common::types::{ToAlloy, ToEthers};
+use foundry_evm_core::{constants::DEFAULT_CREATE2_DEPLOYER, utils::RuntimeOrHandle};
 use k256::{
     ecdsa::SigningKey,
     elliptic_curve::{sec1::ToEncodedPoint, Curve},
@@ -49,7 +50,7 @@ impl Cheatcode for getNonce_1Call {
     }
 }
 
-impl Cheatcode for sign_1Call {
+impl Cheatcode for sign_3Call {
     fn apply_full<DB: DatabaseExt>(&self, _: &mut CheatsCtxt<DB>) -> Result {
         let Self { wallet, digest } = self;
         sign(&wallet.privateKey, digest)
@@ -156,6 +157,10 @@ fn create_wallet(private_key: &U256, label: Option<&str>, state: &mut Cheatcodes
         .abi_encode())
 }
 
+fn encode_vrs(v: u8, r: U256, s: U256) -> Vec<u8> {
+    (U256::from(v), B256::from(r), B256::from(s)).abi_encode()
+}
+
 pub(super) fn sign(private_key: &U256, digest: &B256) -> Result {
     // The `ecrecover` precompile does not use EIP-155. No chain ID is needed.
     let wallet = parse_wallet(private_key)?;
@@ -165,11 +170,46 @@ pub(super) fn sign(private_key: &U256, digest: &B256) -> Result {
 
     assert_eq!(recovered, wallet.address());
 
-    let v = U256::from(sig.v().y_parity_byte_non_eip155().unwrap_or(sig.v().y_parity_byte()));
-    let r = B256::from(sig.r());
-    let s = B256::from(sig.s());
+    let v = sig.v().y_parity_byte_non_eip155().unwrap_or(sig.v().y_parity_byte());
+
+    Ok(encode_vrs(v, sig.r(), sig.s()))
+}
 
-    Ok((v, r, s).abi_encode())
+pub(super) fn sign_with_wallet<DB: DatabaseExt>(
+    ccx: &mut CheatsCtxt<DB>,
+    signer: Option<Address>,
+    digest: &B256,
+) -> Result {
+    let Some(script_wallets) = &ccx.state.script_wallets else {
+        return Err("no wallets are available".into());
+    };
+
+    let mut script_wallets = script_wallets.inner.lock();
+    let maybe_provided_sender = script_wallets.provided_sender;
+    let signers = script_wallets.multi_wallet.signers()?;
+
+    let signer = if let Some(signer) = signer {
+        signer
+    } else if let Some(provided_sender) = maybe_provided_sender {
+        provided_sender
+    } else if signers.len() == 1 {
+        *signers.keys().next().unwrap()
+    } else {
+        return Err("could not determine signer".into());
+    };
+
+    let wallet = signers
+        .get(&signer)
+        .ok_or_else(|| fmt_err!("signer with address {signer} is not available"))?;
+
+    let sig = RuntimeOrHandle::new()
+        .block_on(wallet.sign_hash(digest))
+        .map_err(|err| fmt_err!("{err}"))?;
+
+    let recovered = sig.recover(digest.to_ethers()).map_err(|err| fmt_err!("{err}"))?;
+    assert_eq!(recovered.to_alloy(), signer);
+
+    Ok(encode_vrs(sig.v as u8, sig.r.to_alloy(), sig.s.to_alloy()))
 }
 
 pub(super) fn sign_p256(private_key: &U256, digest: &B256, _state: &mut Cheatcodes) -> Result {

crates/forge/tests/cli/script.rs

@@ -1157,3 +1157,25 @@ contract ScriptC {}
     tester.cmd.forge_fuse().args(["script", "script/B.sol"]);
     tester.simulate(ScriptOutcome::OkNoEndpoint);
 });
+
+forgetest_async!(can_sign_with_script_wallet_single, |prj, cmd| {
+    foundry_test_utils::util::initialize(prj.root());
+
+    let mut tester = ScriptTester::new_broadcast_without_endpoint(cmd, prj.root());
+    tester
+        .add_sig("ScriptSign", "run()")
+        .load_private_keys(&[0])
+        .await
+        .simulate(ScriptOutcome::OkNoEndpoint);
+});
+
+forgetest_async!(can_sign_with_script_wallet_multiple, |prj, cmd| {
+    let mut tester = ScriptTester::new_broadcast_without_endpoint(cmd, prj.root());
+    let acc = tester.accounts_pub[0].to_checksum(None);
+    tester
+        .add_sig("ScriptSign", "run(address)")
+        .arg(&acc)
+        .load_private_keys(&[0, 1, 2])
+        .await
+        .simulate(ScriptOutcome::OkRun);
+});

crates/test-utils/src/script.rs

@@ -264,6 +264,7 @@ pub enum ScriptOutcome {
     ScriptFailed,
     UnsupportedLibraries,
     ErrorSelectForkOnBroadcast,
+    OkRun,
 }
 
 impl ScriptOutcome {
@@ -279,21 +280,23 @@ impl ScriptOutcome {
             Self::ScriptFailed => "script failed: ",
             Self::UnsupportedLibraries => "Multi chain deployment does not support library linking at the moment.",
             Self::ErrorSelectForkOnBroadcast => "cannot select forks during a broadcast",
+            Self::OkRun => "Script ran successfully",
         }
     }
 
     pub fn is_err(&self) -> bool {
         match self {
-            ScriptOutcome::OkNoEndpoint
-            | ScriptOutcome::OkSimulation
-            | ScriptOutcome::OkBroadcast
-            | ScriptOutcome::WarnSpecifyDeployer => false,
-            ScriptOutcome::MissingSender
-            | ScriptOutcome::MissingWallet
-            | ScriptOutcome::StaticCallNotAllowed
-            | ScriptOutcome::UnsupportedLibraries
-            | ScriptOutcome::ErrorSelectForkOnBroadcast
-            | ScriptOutcome::ScriptFailed => true,
+            ScriptOutcome::OkNoEndpoint |
+            ScriptOutcome::OkSimulation |
+            ScriptOutcome::OkBroadcast |
+            ScriptOutcome::WarnSpecifyDeployer |
+            ScriptOutcome::OkRun => false,
+            ScriptOutcome::MissingSender |
+            ScriptOutcome::MissingWallet |
+            ScriptOutcome::StaticCallNotAllowed |
+            ScriptOutcome::UnsupportedLibraries |
+            ScriptOutcome::ErrorSelectForkOnBroadcast |
+            ScriptOutcome::ScriptFailed => true,
         }
     }
 }

testdata/default/cheats/Broadcast.t.sol

@@ -528,3 +528,45 @@ contract ScriptAdditionalContracts is DSTest {
         new Parent();
     }
 }
+
+contract SignatureTester {
+    address public immutable owner;
+
+    constructor() {
+        owner = msg.sender;
+    }
+
+    function verifySignature(bytes32 digest, uint8 v, bytes32 r, bytes32 s) public view returns (bool) {
+        require(ecrecover(digest, v, r, s) == owner, "Invalid signature");
+    }
+}
+
+contract ScriptSign is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    bytes32 digest = keccak256("something");
+
+    function run() external {
+        vm.startBroadcast();
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(digest);
+
+        vm._expectCheatcodeRevert(
+            bytes(string.concat("signer with address ", vm.toString(address(this)), " is not available"))
+        );
+        vm.sign(address(this), digest);
+
+        SignatureTester tester = new SignatureTester();
+        (, address caller,) = vm.readCallers();
+        assertEq(tester.owner(), caller);
+        tester.verifySignature(digest, v, r, s);
+    }
+
+    function run(address sender) external {
+        vm._expectCheatcodeRevert(bytes("could not determine signer"));
+        vm.sign(digest);
+
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(sender, digest);
+        address actual = ecrecover(digest, v, r, s);
+
+        assertEq(actual, sender);
+    }
+}