From 468f4eec244ce1719693c703f37a73d3c1e676bc Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 20 Mar 2024 03:39:10 +0400 Subject: [PATCH 1/5] feat: vm.sign for script wallets --- crates/cheatcodes/assets/cheatcodes.json | 44 +++++++++++++++++++++- crates/cheatcodes/spec/src/vm.rs | 32 ++++++++++++++-- crates/cheatcodes/src/evm.rs | 14 +++++++ crates/cheatcodes/src/utils.rs | 48 ++++++++++++++++++++++-- crates/forge/tests/cli/script.rs | 22 +++++++++++ crates/test-utils/src/script.rs | 5 ++- testdata/cheats/Vm.sol | 2 + testdata/default/cheats/Broadcast.t.sol | 34 +++++++++++++++++ 8 files changed, 191 insertions(+), 10 deletions(-) diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index b086072263d1..4fcd2451416f 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": "", @@ -7041,6 +7041,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", @@ -7141,7 +7181,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 bd6500fb7c81..a7e81ba4960e 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); @@ -1553,8 +1569,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; @@ -1568,8 +1588,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 601c97ad0046..2f3dafedefe6 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..e57bd3d051ac 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) @@ -172,6 +173,47 @@ pub(super) fn sign(private_key: &U256, digest: &B256) -> Result { Ok((v, r, s).abi_encode()) } +pub(super) fn sign_with_wallet( + ccx: &mut CheatsCtxt, + signer: Option
, + digest: &B256, +) -> Result { + if let Some(script_wallets) = &ccx.state.script_wallets { + 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); + + let v = U256::from(sig.v); + let r = B256::from(sig.r.to_alloy()); + let s = B256::from(sig.s.to_alloy()); + + Ok((v, r, s).abi_encode()) + } else { + Err("no wallets are available".into()) + } +} + pub(super) fn sign_p256(private_key: &U256, digest: &B256, _state: &mut Cheatcodes) -> Result { ensure!(*private_key != U256::ZERO, "private key cannot be 0"); let n = U256::from_limbs(*p256::NistP256::ORDER.as_words()); 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 80b06ae769b2..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,6 +280,7 @@ 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", } } @@ -287,7 +289,8 @@ impl ScriptOutcome { ScriptOutcome::OkNoEndpoint | ScriptOutcome::OkSimulation | ScriptOutcome::OkBroadcast | - ScriptOutcome::WarnSpecifyDeployer => false, + ScriptOutcome::WarnSpecifyDeployer | + ScriptOutcome::OkRun => false, ScriptOutcome::MissingSender | ScriptOutcome::MissingWallet | ScriptOutcome::StaticCallNotAllowed | diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 623ef254bf06..3eb4f36f4ecb 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -349,6 +349,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..df4a24d0044c 100644 --- a/testdata/default/cheats/Broadcast.t.sol +++ b/testdata/default/cheats/Broadcast.t.sol @@ -528,3 +528,37 @@ 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); + + SignatureTester tester = new SignatureTester(); + (, address caller,) = vm.readCallers(); + assertEq(tester.owner(), caller); + tester.verifySignature(digest, v, r, s); + } + + function run(address sender) external { + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sender, digest); + address actual = ecrecover(digest, v, r, s); + + assertEq(actual, sender); + } +} From 5930fa04870ec4b0306fb8b50c25568d4db02190 Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 20 Mar 2024 03:51:52 +0400 Subject: [PATCH 2/5] more tests --- testdata/default/cheats/Broadcast.t.sol | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/testdata/default/cheats/Broadcast.t.sol b/testdata/default/cheats/Broadcast.t.sol index df4a24d0044c..6a099dc6ed54 100644 --- a/testdata/default/cheats/Broadcast.t.sol +++ b/testdata/default/cheats/Broadcast.t.sol @@ -549,6 +549,11 @@ contract ScriptSign is DSTest { 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); @@ -556,6 +561,9 @@ contract ScriptSign is DSTest { } 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); From acbfe0a586235971c89c3e183ce1542771742287 Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 20 Mar 2024 04:17:28 +0400 Subject: [PATCH 3/5] clippy --- crates/cheatcodes/src/utils.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs index e57bd3d051ac..c9d9f806ca87 100644 --- a/crates/cheatcodes/src/utils.rs +++ b/crates/cheatcodes/src/utils.rs @@ -198,7 +198,7 @@ pub(super) fn sign_with_wallet( .ok_or_else(|| fmt_err!("signer with address {signer} is not available"))?; let sig = RuntimeOrHandle::new() - .block_on(wallet.sign_hash(&digest)) + .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}"))?; From 3161570f7aa38e6e28a42f1d537d5c9d193d4237 Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 20 Mar 2024 05:41:22 +0400 Subject: [PATCH 4/5] if let some else --- crates/cheatcodes/src/utils.rs | 58 +++++++++++++++++----------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs index c9d9f806ca87..aa024658312d 100644 --- a/crates/cheatcodes/src/utils.rs +++ b/crates/cheatcodes/src/utils.rs @@ -178,40 +178,40 @@ pub(super) fn sign_with_wallet( signer: Option
, digest: &B256, ) -> Result { - if let Some(script_wallets) = &ccx.state.script_wallets { - 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 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 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 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); + let recovered = sig.recover(digest.to_ethers()).map_err(|err| fmt_err!("{err}"))?; + assert_eq!(recovered.to_alloy(), signer); - let v = U256::from(sig.v); - let r = B256::from(sig.r.to_alloy()); - let s = B256::from(sig.s.to_alloy()); + let v = U256::from(sig.v); + let r = B256::from(sig.r.to_alloy()); + let s = B256::from(sig.s.to_alloy()); - Ok((v, r, s).abi_encode()) - } else { - Err("no wallets are available".into()) - } + Ok((v, r, s).abi_encode()) } pub(super) fn sign_p256(private_key: &U256, digest: &B256, _state: &mut Cheatcodes) -> Result { From 90759d6f6b56379bff3e2905b98be98d58e8504d Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 20 Mar 2024 05:47:18 +0400 Subject: [PATCH 5/5] review fixes --- crates/cheatcodes/src/utils.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs index aa024658312d..f9edc8e40d7c 100644 --- a/crates/cheatcodes/src/utils.rs +++ b/crates/cheatcodes/src/utils.rs @@ -157,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)?; @@ -166,11 +170,9 @@ 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((v, r, s).abi_encode()) + Ok(encode_vrs(v, sig.r(), sig.s())) } pub(super) fn sign_with_wallet( @@ -207,11 +209,7 @@ pub(super) fn sign_with_wallet( let recovered = sig.recover(digest.to_ethers()).map_err(|err| fmt_err!("{err}"))?; assert_eq!(recovered.to_alloy(), signer); - let v = U256::from(sig.v); - let r = B256::from(sig.r.to_alloy()); - let s = B256::from(sig.s.to_alloy()); - - Ok((v, r, s).abi_encode()) + 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 {