From b1f1ee9b20a5cc8d0686b566267bb09221b9a23b Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Thu, 28 Mar 2024 17:18:45 +0400 Subject: [PATCH] feat: `vm.sign` for scripts (#7454) * feat: vm.sign for script wallets * more tests * clippy * if let some else * review fixes --- crates/cheatcodes/assets/cheatcodes.json | 44 ++++++++++++++++++- crates/cheatcodes/spec/src/vm.rs | 32 ++++++++++++-- crates/cheatcodes/src/evm.rs | 14 ++++++ crates/cheatcodes/src/utils.rs | 54 +++++++++++++++++++++--- crates/forge/tests/cli/script.rs | 22 ++++++++++ crates/test-utils/src/script.rs | 23 +++++----- testdata/cheats/Vm.sol | 2 + testdata/default/cheats/Broadcast.t.sol | 42 ++++++++++++++++++ 8 files changed, 210 insertions(+), 23 deletions(-) diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 9c69772b69f9..dea51432227f 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -2941,7 +2941,7 @@ { "func": { "id": "broadcast_0", - "description": "Using the address that calls the test contract, has the next call (at this call depth only)\ncreate a transaction that can later be signed and sent onchain.", + "description": "Has the next call (at this call depth only) create transactions that can later be signed and sent onchain.\nBroadcasting address is determined by checking the following in order:\n1. If `--sender` argument was provided, that address is used.\n2. If exactly one signer (e.g. private key, hw wallet, keystore) is set when `forge broadcast` is invoked, that signer is used.\n3. Otherwise, default foundry sender (1804c8AB1F12E6bbf3894d4083f33e07309d1f38) is used.", "declaration": "function broadcast() external;", "visibility": "external", "mutability": "", @@ -7081,6 +7081,46 @@ { "func": { "id": "sign_1", + "description": "Signs `digest` with signer provided to script using the secp256k1 curve.\nIf `--sender` is provided, the signer with provided address is used, otherwise,\nif exactly one signer is provided to the script, that signer is used.\nRaises error if signer passed through `--sender` does not match any unlocked signers or\nif `--sender` is not provided and not exactly one signer is passed to the script.", + "declaration": "function sign(bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);", + "visibility": "external", + "mutability": "pure", + "signature": "sign(bytes32)", + "selector": "0x799cd333", + "selectorBytes": [ + 121, + 156, + 211, + 51 + ] + }, + "group": "evm", + "status": "stable", + "safety": "safe" + }, + { + "func": { + "id": "sign_2", + "description": "Signs `digest` with signer provided to script using the secp256k1 curve.\nRaises error if none of the signers passed into the script have provided address.", + "declaration": "function sign(address signer, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);", + "visibility": "external", + "mutability": "pure", + "signature": "sign(address,bytes32)", + "selector": "0x8c1aa205", + "selectorBytes": [ + 140, + 26, + 162, + 5 + ] + }, + "group": "evm", + "status": "stable", + "safety": "safe" + }, + { + "func": { + "id": "sign_3", "description": "Signs data with a `Wallet`.", "declaration": "function sign(Wallet calldata wallet, bytes32 digest) external returns (uint8 v, bytes32 r, bytes32 s);", "visibility": "external", @@ -7181,7 +7221,7 @@ { "func": { "id": "startBroadcast_0", - "description": "Using the address that calls the test contract, has all subsequent calls\n(at this call depth only) create transactions that can later be signed and sent onchain.", + "description": "Has all subsequent calls (at this call depth only) create transactions that can later be signed and sent onchain.\nBroadcasting address is determined by checking the following in order:\n1. If `--sender` argument was provided, that address is used.\n2. If exactly one signer (e.g. private key, hw wallet, keystore) is set when `forge broadcast` is invoked, that signer is used.\n3. Otherwise, default foundry sender (1804c8AB1F12E6bbf3894d4083f33e07309d1f38) is used.", "declaration": "function startBroadcast() external;", "visibility": "external", "mutability": "", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index cdcca3331785..8f0de363a7fa 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/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; diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 2269c467cf91..b694a9f64da4 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -145,6 +145,20 @@ impl Cheatcode for sign_0Call { } } +impl Cheatcode for sign_1Call { + fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { digest } = self; + super::utils::sign_with_wallet(ccx, None, digest) + } +} + +impl Cheatcode for sign_2Call { + fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { signer, digest } = self; + super::utils::sign_with_wallet(ccx, Some(*signer), digest) + } +} + impl Cheatcode for signP256Call { fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { let Self { privateKey, digest } = self; diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs index 525120d7de8f..f9edc8e40d7c 100644 --- a/crates/cheatcodes/src/utils.rs +++ b/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(&self, _: &mut CheatsCtxt) -> 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 { + (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( + ccx: &mut CheatsCtxt, + signer: Option
, + 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 { diff --git a/crates/forge/tests/cli/script.rs b/crates/forge/tests/cli/script.rs index e212683c1a2b..ff03a6b57e88 100644 --- a/crates/forge/tests/cli/script.rs +++ b/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); +}); diff --git a/crates/test-utils/src/script.rs b/crates/test-utils/src/script.rs index 1a5f46283ae7..b60723d9d8d2 100644 --- a/crates/test-utils/src/script.rs +++ b/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, } } } diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index b5a604e14d96..2116acf2b8ff 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -351,6 +351,8 @@ interface Vm { function setNonceUnsafe(address account, uint64 newNonce) external; function signP256(uint256 privateKey, bytes32 digest) external pure returns (bytes32 r, bytes32 s); function sign(uint256 privateKey, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s); + function sign(bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s); + function sign(address signer, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s); function sign(Wallet calldata wallet, bytes32 digest) external returns (uint8 v, bytes32 r, bytes32 s); function skip(bool skipTest) external; function sleep(uint256 duration) external; diff --git a/testdata/default/cheats/Broadcast.t.sol b/testdata/default/cheats/Broadcast.t.sol index d44bc954006f..6a099dc6ed54 100644 --- a/testdata/default/cheats/Broadcast.t.sol +++ b/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); + } +}