Support EIP-7702 Delegations in Forge

crates/cheatcodes/spec/src/vm.rs

@@ -1973,6 +1973,18 @@ interface Vm {
     #[cheatcode(group = Scripting)]
     function broadcastRawTransaction(bytes calldata data) external;
 
+    /// Create an EIP-7702 authorization for delegation
+    #[cheatcode(group = Scripting)]
+    function createDelegation(address implementation, uint64 nonce) external returns (bytes32);
+
+    /// Sign an EIP-7702 authorization for delegation
+    #[cheatcode(group = Scripting)]
+    function signDelegation(bytes32 delegation, uint256 privateKey) external returns (uint8 v, bytes32 r, bytes32 s);
+
+    /// Designate the next call as an EIP-7702 transaction
+    #[cheatcode(group = Scripting)]
+    function attachDelegation(address implementation, uint64 nonce, uint8 v, bytes32 r, bytes32 s) external;
+
     /// Returns addresses of available unlocked wallets in the script environment.
     #[cheatcode(group = Scripting)]
     function getWallets() external returns (address[] memory wallets);

crates/cheatcodes/src/inspector.rs

@@ -46,7 +46,10 @@ use revm::{
         EOFCreateInputs, EOFCreateKind, Gas, InstructionResult, Interpreter, InterpreterAction,
         InterpreterResult,
     },
-    primitives::{BlockEnv, CreateScheme, EVMError, EvmStorageSlot, SpecId, EOF_MAGIC_BYTES},
+    primitives::{
+        BlockEnv, CreateScheme, EVMError, EvmStorageSlot, SignedAuthorization, SpecId,
+        EOF_MAGIC_BYTES,
+    },
     EvmContext, InnerEvmContext, Inspector,
 };
 use serde_json::Value;
@@ -373,6 +376,9 @@ pub struct Cheatcodes {
     /// execution block environment.
     pub block: Option<BlockEnv>,
 
+    /// EIP-7702 delegations
+    pub delegations: HashMap<Address, SignedAuthorization>,
+
     /// The gas price.
     ///
     /// Used in the cheatcode handler to overwrite the gas price separately from the gas price
@@ -497,6 +503,7 @@ impl Cheatcodes {
             labels: config.labels.clone(),
             config,
             block: Default::default(),
+            delegations: Default::default(),
             gas_price: Default::default(),
             prank: Default::default(),
             expected_revert: Default::default(),

crates/cheatcodes/src/script.rs

@@ -2,10 +2,13 @@
 
 use crate::{Cheatcode, CheatsCtxt, Result, Vm::*};
 use alloy_primitives::{Address, B256, U256};
+use alloy_rpc_types::Authorization;
+use alloy_signer::{Signature, SignerSync};
 use alloy_signer_local::PrivateKeySigner;
 use alloy_sol_types::SolValue;
 use foundry_wallets::{multi_wallet::MultiWallet, WalletSigner};
 use parking_lot::Mutex;
+use revm::primitives::SignedAuthorization;
 use std::sync::Arc;
 
 impl Cheatcode for broadcast_0Call {
@@ -29,6 +32,57 @@ impl Cheatcode for broadcast_2Call {
     }
 }
 
+impl Cheatcode for createDelegationCall {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { implementation, nonce } = self;
+        let auth = Authorization {
+            address: *implementation,
+            nonce: *nonce,
+            chain_id: U256::from(ccx.ecx.env.cfg.chain_id),
+        };
+        let hash = auth.signature_hash();
+        Ok(hash.to_vec())
+    }
+}
+
+impl Cheatcode for attachDelegationCall {
+    // @todo - change this to apply_full?
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { implementation, nonce, v, r, s } = self;
+        let auth = Authorization {
+            address: *implementation,
+            nonce: *nonce,
+            chain_id: U256::from(ccx.ecx.env.cfg.chain_id),
+        };
+        let addr = auth.address;
+        let sig = Signature::from_rs_and_parity(
+            U256::try_from(*r).unwrap(),
+            U256::try_from(*s).unwrap(),
+            alloy_primitives::Parity::from(*v == 1),
+        )?;
+        let signed_auth = auth.into_signed(sig);
+        // store in CheatcodeState
+        ccx.state.delegations.insert(addr, signed_auth);
+        Ok(Default::default())
+    }
+}
+
+impl Cheatcode for signDelegationCall {
+    fn apply_stateful(&self, _: &mut CheatsCtxt) -> Result {
+        let Self { delegation, privateKey } = self;
+        let signer = super::crypto::parse_wallet(privateKey)?;
+        let sig = signer.sign_hash_sync(delegation)?;
+        Ok(encode_delegation_sig(sig))
+    }
+}
+
+fn encode_delegation_sig(sig: alloy_primitives::Signature) -> Vec<u8> {
+    let v = U256::from(sig.v().y_parity() as u8);
+    let r = sig.r();
+    let s = sig.s();
+    (v, r, s).abi_encode()
+}
+
 impl Cheatcode for startBroadcast_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self {} = self;
@@ -80,6 +134,8 @@ pub struct Broadcast {
     pub depth: u64,
     /// Whether the prank stops by itself after the next call
     pub single_call: bool,
+    /// EIP-7702 delegation
+    pub delegation: Option<SignedAuthorization>,
 }
 
 /// Contains context for wallet management.
@@ -154,18 +210,24 @@ fn broadcast(ccx: &mut CheatsCtxt, new_origin: Option<&Address>, single_call: bo
     ensure!(ccx.state.broadcast.is_none(), "a broadcast is active already");
 
     let mut new_origin = new_origin.cloned();
+    let mut delegation = None;
 
     if new_origin.is_none() {
-        let mut wallets = ccx.state.wallets().inner.lock();
-        if let Some(provided_sender) = wallets.provided_sender {
-            new_origin = Some(provided_sender);
+            let mut wallets = ccx.state.wallets().inner.lock();
+            if let Some(provided_sender) = wallets.provided_sender {
+                new_origin = Some(provided_sender);
+                delegation = ccx.state.delegations.get(&provided_sender).cloned();
         } else {
             let signers = wallets.multi_wallet.signers()?;
             if signers.len() == 1 {
                 let address = signers.keys().next().unwrap();
                 new_origin = Some(*address);
+                delegation = ccx.state.delegations.get(address).cloned();
             }
         }
+    } else if let Some(addr) = new_origin {
+        // check delegation for explicitly provided address
+        delegation = ccx.state.delegations.get(&addr).cloned();
     }
 
     let broadcast = Broadcast {
@@ -174,6 +236,7 @@ fn broadcast(ccx: &mut CheatsCtxt, new_origin: Option<&Address>, single_call: bo
         original_origin: ccx.ecx.env.tx.caller,
         depth: ccx.ecx.journaled_state.depth(),
         single_call,
+        delegation,
     };
     debug!(target: "cheatcodes", ?broadcast, "started");
     ccx.state.broadcast = Some(broadcast);

crates/script/Cargo.toml

@@ -24,6 +24,7 @@ foundry-cheatcodes.workspace = true
 foundry-wallets.workspace = true
 foundry-linking.workspace = true
 forge-script-sequence.workspace = true
+revm = { workspace = true, default-features = false, features = ["std"] }
 
 serde.workspace = true
 eyre.workspace = true

testdata/default/cheats/AttachDelegation.t.sol

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity 0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract AttachDelegationTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    uint256 pk = 77814517325470205911140941194401928579557062014761831930645393041380819009408;
+    address public ACCOUNT_A = 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266;
+
+    function testCreateDelegation(address implementation, uint64 nonce) public {
+        bytes32 delegation = vm.createDelegation(implementation, nonce);
+        (uint8 v, bytes32 r, bytes32 s) = vm.signDelegation(delegation, pk);
+
+        vm.attachDelegation(implementation, nonce, v, r, s);
+        vm.broadcast();
+        // @todo - assert a transaction's msg.sender is implementation address
+    }
+}