Skip to content

Commit 7a7ad4e

Browse files
fenghaojianggithub-actions[bot]jes16jupytergrandizzy
authored
feat: support signDelegation and signAndAttachDelegation with nonce (#10182)
* feat: support signAndAttachDelegationWithNonce * chore: remove useless param * style: fmt * feat: support signDelegationWithNonce * style: add comment for signAndAttachDelegationWithNonce * feat: support write_delegation_skip_nonce * feat: support signDelegationWithNonce * style: add comment for signAndAttachDelegationWithNonce * feat: support write_delegation_skip_nonce * doc: fix function explanation * chore: use overloda * cargo fmt * fix: remove with nonce * feat: add delegation and attach delegation helper * cargo cheats * Simplify code and add test for new cheatcodes --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: jes16jupyter <jes163940@gmail.com> Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>
1 parent 3accf6b commit 7a7ad4e

File tree

5 files changed

+129
-50
lines changed

5 files changed

+129
-50
lines changed

‎crates/cheatcodes/assets/cheatcodes.json

Lines changed: 42 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/cheatcodes/spec/src/vm.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2200,6 +2200,10 @@ interface Vm {
22002200
#[cheatcode(group = Scripting)]
22012201
function signDelegation(address implementation, uint256 privateKey) external returns (SignedDelegation memory signedDelegation);
22022202

2203+
/// Sign an EIP-7702 authorization for delegation for specific nonce
2204+
#[cheatcode(group = Scripting)]
2205+
function signDelegation(address implementation, uint256 privateKey, uint64 nonce) external returns (SignedDelegation memory signedDelegation);
2206+
22032207
/// Designate the next call as an EIP-7702 transaction
22042208
#[cheatcode(group = Scripting)]
22052209
function attachDelegation(SignedDelegation calldata signedDelegation) external;
@@ -2208,6 +2212,10 @@ interface Vm {
22082212
#[cheatcode(group = Scripting)]
22092213
function signAndAttachDelegation(address implementation, uint256 privateKey) external returns (SignedDelegation memory signedDelegation);
22102214

2215+
/// Sign an EIP-7702 authorization and designate the next call as an EIP-7702 transaction for specific nonce
2216+
#[cheatcode(group = Scripting)]
2217+
function signAndAttachDelegation(address implementation, uint256 privateKey, uint64 nonce) external returns (SignedDelegation memory signedDelegation);
2218+
22112219
/// Returns addresses of available unlocked wallets in the script environment.
22122220
#[cheatcode(group = Scripting)]
22132221
function getWallets() external returns (address[] memory wallets);

‎crates/cheatcodes/src/script.rs

Lines changed: 55 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! Implementations of [`Scripting`](spec::Group::Scripting) cheatcodes.
22
33
use crate::{Cheatcode, CheatsCtxt, Result, Vm::*};
4-
use alloy_primitives::{Address, PrimitiveSignature, B256, U256};
4+
use alloy_primitives::{Address, Uint, B256, U256};
55
use alloy_rpc_types::Authorization;
66
use alloy_signer::SignerSync;
77
use alloy_signer_local::PrivateKeySigner;
@@ -54,46 +54,71 @@ impl Cheatcode for attachDelegationCall {
5454
}
5555
}
5656

57-
impl Cheatcode for signDelegationCall {
57+
impl Cheatcode for signDelegation_0Call {
5858
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
59-
let Self { implementation, privateKey } = self;
60-
let signer = PrivateKeySigner::from_bytes(&B256::from(*privateKey))?;
61-
let authority = signer.address();
62-
let (auth, nonce) = create_auth(ccx, *implementation, authority)?;
63-
let sig = signer.sign_hash_sync(&auth.signature_hash())?;
64-
Ok(sig_to_delegation(sig, nonce, *implementation).abi_encode())
59+
let Self { implementation, privateKey } = *self;
60+
sign_delegation(ccx, privateKey, implementation, None, false)
6561
}
6662
}
6763

68-
impl Cheatcode for signAndAttachDelegationCall {
64+
impl Cheatcode for signDelegation_1Call {
6965
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
70-
let Self { implementation, privateKey } = self;
71-
let signer = PrivateKeySigner::from_bytes(&B256::from(*privateKey))?;
72-
let authority = signer.address();
73-
let (auth, nonce) = create_auth(ccx, *implementation, authority)?;
74-
let sig = signer.sign_hash_sync(&auth.signature_hash())?;
75-
let signed_auth = sig_to_auth(sig, auth);
76-
write_delegation(ccx, signed_auth.clone())?;
77-
ccx.state.active_delegation = Some(signed_auth);
78-
Ok(sig_to_delegation(sig, nonce, *implementation).abi_encode())
66+
let Self { implementation, privateKey, nonce } = *self;
67+
sign_delegation(ccx, privateKey, implementation, Some(nonce), false)
7968
}
8069
}
8170

82-
fn create_auth(
71+
impl Cheatcode for signAndAttachDelegation_0Call {
72+
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
73+
let Self { implementation, privateKey } = *self;
74+
sign_delegation(ccx, privateKey, implementation, None, true)
75+
}
76+
}
77+
78+
impl Cheatcode for signAndAttachDelegation_1Call {
79+
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
80+
let Self { implementation, privateKey, nonce } = *self;
81+
sign_delegation(ccx, privateKey, implementation, Some(nonce), true)
82+
}
83+
}
84+
85+
/// Helper function to sign and attach (if needed) an EIP-7702 delegation.
86+
/// Uses the provided nonce, otherwise retrieves and increments the nonce of the EOA.
87+
fn sign_delegation(
8388
ccx: &mut CheatsCtxt,
89+
private_key: Uint<256, 4>,
8490
implementation: Address,
85-
authority: Address,
86-
) -> Result<(Authorization, u64)> {
87-
let authority_acc = ccx.ecx.journaled_state.load_account(authority, &mut ccx.ecx.db)?;
88-
let nonce = authority_acc.data.info.nonce;
89-
Ok((
90-
Authorization {
91-
address: implementation,
92-
nonce,
93-
chain_id: U256::from(ccx.ecx.env.cfg.chain_id),
94-
},
91+
nonce: Option<u64>,
92+
attach: bool,
93+
) -> Result<Vec<u8>> {
94+
let signer = PrivateKeySigner::from_bytes(&B256::from(private_key))?;
95+
let nonce = if let Some(nonce) = nonce {
96+
nonce
97+
} else {
98+
let authority_acc =
99+
ccx.ecx.journaled_state.load_account(signer.address(), &mut ccx.ecx.db)?;
100+
authority_acc.data.info.nonce
101+
};
102+
let auth = Authorization {
103+
address: implementation,
104+
nonce,
105+
chain_id: U256::from(ccx.ecx.env.cfg.chain_id),
106+
};
107+
let sig = signer.sign_hash_sync(&auth.signature_hash())?;
108+
// Attach delegation.
109+
if attach {
110+
let signed_auth = SignedAuthorization::new_unchecked(auth, sig.v() as u8, sig.r(), sig.s());
111+
write_delegation(ccx, signed_auth.clone())?;
112+
ccx.state.active_delegation = Some(signed_auth);
113+
}
114+
Ok(SignedDelegation {
115+
v: sig.v() as u8,
116+
r: sig.r().into(),
117+
s: sig.s().into(),
95118
nonce,
96-
))
119+
implementation,
120+
}
121+
.abi_encode())
97122
}
98123

99124
fn write_delegation(ccx: &mut CheatsCtxt, auth: SignedAuthorization) -> Result<()> {
@@ -108,24 +133,6 @@ fn write_delegation(ccx: &mut CheatsCtxt, auth: SignedAuthorization) -> Result<(
108133
Ok(())
109134
}
110135

111-
fn sig_to_delegation(
112-
sig: PrimitiveSignature,
113-
nonce: u64,
114-
implementation: Address,
115-
) -> SignedDelegation {
116-
SignedDelegation {
117-
v: sig.v() as u8,
118-
r: sig.r().into(),
119-
s: sig.s().into(),
120-
nonce,
121-
implementation,
122-
}
123-
}
124-
125-
fn sig_to_auth(sig: PrimitiveSignature, auth: Authorization) -> SignedAuthorization {
126-
SignedAuthorization::new_unchecked(auth, sig.v() as u8, sig.r(), sig.s())
127-
}
128-
129136
impl Cheatcode for startBroadcast_0Call {
130137
fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
131138
let Self {} = self;

‎testdata/cheats/Vm.sol

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎testdata/default/cheats/AttachDelegation.t.sol

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,19 @@ contract AttachDelegationTest is DSTest {
3939
assertEq(token.balanceOf(bob), 100);
4040
}
4141

42+
/// forge-config: default.allow_internal_expect_revert = true
43+
function testCallSingleAttachDelegationWithNonce() public {
44+
Vm.SignedDelegation memory signedDelegation = vm.signDelegation(address(implementation), alice_pk, 11);
45+
vm.broadcast(bob_pk);
46+
vm._expectCheatcodeRevert("vm.attachDelegation: invalid nonce");
47+
vm.attachDelegation(signedDelegation);
48+
49+
signedDelegation = vm.signDelegation(address(implementation), alice_pk, 0);
50+
vm.attachDelegation(signedDelegation);
51+
signedDelegation = vm.signDelegation(address(implementation), alice_pk, 1);
52+
vm.attachDelegation(signedDelegation);
53+
}
54+
4255
function testMultiCallAttachDelegation() public {
4356
Vm.SignedDelegation memory signedDelegation = vm.signDelegation(address(implementation), alice_pk);
4457
vm.broadcast(bob_pk);
@@ -126,6 +139,15 @@ contract AttachDelegationTest is DSTest {
126139

127140
assertEq(token.balanceOf(bob), 100);
128141
}
142+
143+
/// forge-config: default.allow_internal_expect_revert = true
144+
function testCallSingleSignAndAttachDelegationWithNonce() public {
145+
vm._expectCheatcodeRevert("vm.signAndAttachDelegation: invalid nonce");
146+
vm.signAndAttachDelegation(address(implementation), alice_pk, 11);
147+
148+
vm.signAndAttachDelegation(address(implementation), alice_pk, 0);
149+
vm.signAndAttachDelegation(address(implementation), alice_pk, 1);
150+
}
129151
}
130152

131153
contract SimpleDelegateContract {

0 commit comments

Comments
 (0)