diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index 5e4d567f2..a7997690d 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -700,6 +700,10 @@ interface Vm { #[cheatcode(group = Testing, safety = Safe)] function zkVmSkip() external pure; + /// Enables/Disables use of a paymaster for ZK transactions. + #[cheatcode(group = Testing, safety = Safe)] + function zkUsePaymaster(address paymaster_address, bytes calldata paymaster_input) external pure; + /// Registers bytecodes for ZK-VM for transact/call and create instructions. #[cheatcode(group = Testing, safety = Safe)] function zkRegisterContract(string calldata name, bytes32 evmBytecodeHash, bytes calldata evmDeployedBytecode, bytes calldata evmBytecode, bytes32 zkBytecodeHash, bytes calldata zkDeployedBytecode) external pure; diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 95cb85356..ce7dc31ae 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -37,8 +37,8 @@ use foundry_evm_core::{ use foundry_zksync_compiler::{DualCompiledContract, DualCompiledContracts}; use foundry_zksync_core::{ convert::{ConvertH160, ConvertH256, ConvertRU256, ConvertU256}, - get_account_code_key, get_balance_key, get_nonce_key, Call, ZkTransactionMetadata, - DEFAULT_CREATE2_DEPLOYER_ZKSYNC, + get_account_code_key, get_balance_key, get_nonce_key, Call, ZkPaymasterData, + ZkTransactionMetadata, DEFAULT_CREATE2_DEPLOYER_ZKSYNC, }; use itertools::Itertools; use revm::{ @@ -363,6 +363,9 @@ pub struct Cheatcodes { /// Records the next create address for `skip_zk_vm_addresses`. pub record_next_create_address: bool, + /// Paymaster params + pub paymaster_params: Option, + /// Dual compiled contracts pub dual_compiled_contracts: DualCompiledContracts, @@ -463,6 +466,7 @@ impl Cheatcodes { skip_zk_vm_addresses: Default::default(), record_next_create_address: Default::default(), persisted_factory_deps: Default::default(), + paymaster_params: None, } } @@ -976,6 +980,7 @@ impl Cheatcodes { expected_calls: Some(&mut self.expected_calls), accesses: self.accesses.as_mut(), persisted_factory_deps: Some(&mut self.persisted_factory_deps), + paymaster_data: self.paymaster_params.take(), }; let create_inputs = CreateInputs { scheme: input.scheme().unwrap_or(CreateScheme::Create), @@ -1567,6 +1572,7 @@ impl Cheatcodes { expected_calls: Some(&mut self.expected_calls), accesses: self.accesses.as_mut(), persisted_factory_deps: Some(&mut self.persisted_factory_deps), + paymaster_data: self.paymaster_params.take(), }; // We currently exhaust the entire gas for the call as zkEVM returns a very high amount diff --git a/crates/cheatcodes/src/test.rs b/crates/cheatcodes/src/test.rs index deaaa71ba..064d2daba 100644 --- a/crates/cheatcodes/src/test.rs +++ b/crates/cheatcodes/src/test.rs @@ -5,6 +5,7 @@ use alloy_primitives::Address; use alloy_sol_types::SolValue; use foundry_evm_core::constants::{MAGIC_ASSUME, MAGIC_SKIP}; use foundry_zksync_compiler::DualCompiledContract; +use foundry_zksync_core::ZkPaymasterData; pub(crate) mod assert; pub(crate) mod expect; @@ -31,6 +32,15 @@ impl Cheatcode for zkVmSkipCall { } } +impl Cheatcode for zkUsePaymasterCall { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { paymaster_address, paymaster_input } = self; + ccx.state.paymaster_params = + Some(ZkPaymasterData { address: *paymaster_address, input: paymaster_input.clone() }); + Ok(Default::default()) + } +} + impl Cheatcode for zkRegisterContractCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { diff --git a/crates/forge/tests/fixtures/zk/MyPaymaster.sol b/crates/forge/tests/fixtures/zk/MyPaymaster.sol new file mode 100644 index 000000000..225547d35 --- /dev/null +++ b/crates/forge/tests/fixtures/zk/MyPaymaster.sol @@ -0,0 +1,49 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +import { + IPaymaster, + ExecutionResult, + PAYMASTER_VALIDATION_SUCCESS_MAGIC +} from "zksync-contracts/zksync-contracts/l2/system-contracts/interfaces/IPaymaster.sol"; +import {IPaymasterFlow} from "zksync-contracts/zksync-contracts/l2/system-contracts/interfaces/IPaymasterFlow.sol"; +import { + TransactionHelper, + Transaction +} from "zksync-contracts/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol"; +import "zksync-contracts/zksync-contracts/l2/system-contracts/Constants.sol"; + +contract MyPaymaster is IPaymaster { + modifier onlyBootloader() { + require(msg.sender == BOOTLOADER_FORMAL_ADDRESS, "Only bootloader can call this method"); + _; + } + + constructor() payable {} + + function validateAndPayForPaymasterTransaction(bytes32, bytes32, Transaction calldata _transaction) + external + payable + onlyBootloader + returns (bytes4 magic, bytes memory context) + { + // Always accept the transaction + magic = PAYMASTER_VALIDATION_SUCCESS_MAGIC; + + // Pay for the transaction fee + uint256 requiredETH = _transaction.gasLimit * _transaction.maxFeePerGas; + (bool success,) = payable(BOOTLOADER_FORMAL_ADDRESS).call{value: requiredETH}(""); + require(success, "Failed to transfer tx fee to the bootloader"); + } + + function postTransaction( + bytes calldata _context, + Transaction calldata _transaction, + bytes32, + bytes32, + ExecutionResult _txResult, + uint256 _maxRefundedGas + ) external payable override onlyBootloader {} + + receive() external payable {} +} diff --git a/crates/forge/tests/fixtures/zk/Paymaster.t.sol b/crates/forge/tests/fixtures/zk/Paymaster.t.sol new file mode 100644 index 000000000..8b661e1f4 --- /dev/null +++ b/crates/forge/tests/fixtures/zk/Paymaster.t.sol @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import {Test, console} from "forge-std/Test.sol"; +import "zksync-contracts/zksync-contracts/l2/system-contracts/Constants.sol"; +import {MyPaymaster} from "./MyPaymaster.sol"; + +contract TestPaymasterFlow is Test { + MyPaymaster private paymaster; + DoStuff private do_stuff; + address private alice; + address private bob; + bytes private paymaster_encoded_input; + + function setUp() public { + alice = makeAddr("Alice"); + bob = makeAddr("Bob"); + do_stuff = new DoStuff(); + paymaster = new MyPaymaster(); + + // A small amount is needed for initial tx processing + vm.deal(alice, 1 ether); + vm.deal(address(paymaster), 10 ether); + + // Encode paymaster input + paymaster_encoded_input = abi.encodeWithSelector(bytes4(keccak256("general(bytes)")), bytes("0x")); + } + + function testCallWithPaymaster() public { + vm.deal(address(do_stuff), 1 ether); + require(address(do_stuff).balance == 1 ether, "Balance is not 1 ether"); + + uint256 alice_balance = address(alice).balance; + (bool success,) = address(vm).call( + abi.encodeWithSignature("zkUsePaymaster(address,bytes)", address(paymaster), paymaster_encoded_input) + ); + require(success, "zkUsePaymaster call failed"); + + vm.prank(alice, alice); + do_stuff.do_stuff(bob); + + require(address(do_stuff).balance == 0, "Balance is not 0 ether"); + require(address(alice).balance == alice_balance, "Balance is not the same"); + require(address(bob).balance == 1 ether, "Balance is not 1 ether"); + } + + function testCreateWithPaymaster() public { + uint256 alice_balance = address(alice).balance; + (bool success,) = address(vm).call( + abi.encodeWithSignature("zkUsePaymaster(address,bytes)", address(paymaster), paymaster_encoded_input) + ); + require(success, "zkUsePaymaster call failed"); + + vm.prank(alice, alice); + DoStuff new_do_stuff = new DoStuff(); + + require(address(alice).balance == alice_balance, "Balance is not the same"); + } + + function testFailPaymasterBalanceDoesNotUpdate() public { + uint256 alice_balance = address(alice).balance; + uint256 paymaster_balance = address(paymaster).balance; + (bool success,) = address(vm).call( + abi.encodeWithSignature("zkUsePaymaster(address,bytes)", address(paymaster), paymaster_encoded_input) + ); + require(success, "zkUsePaymaster call failed"); + + vm.prank(alice, alice); + do_stuff.do_stuff(bob); + + require(address(alice).balance == alice_balance, "Balance is not the same"); + require(address(paymaster).balance < paymaster_balance, "Paymaster balance is not less"); + } +} + +contract DoStuff { + function do_stuff(address recipient) public { + uint256 amount = address(this).balance; + (bool success,) = payable(recipient).call{value: amount}(""); + require(success, "Failed to transfer balance to the recipient."); + } +} diff --git a/crates/forge/tests/it/zk/mod.rs b/crates/forge/tests/it/zk/mod.rs index d7337efb0..bbe6cda52 100644 --- a/crates/forge/tests/it/zk/mod.rs +++ b/crates/forge/tests/it/zk/mod.rs @@ -13,6 +13,7 @@ mod invariant; mod logs; mod nft; mod ownership; +mod paymaster; mod proxy; mod repros; mod traces; diff --git a/crates/forge/tests/it/zk/paymaster.rs b/crates/forge/tests/it/zk/paymaster.rs new file mode 100644 index 000000000..e21a48daf --- /dev/null +++ b/crates/forge/tests/it/zk/paymaster.rs @@ -0,0 +1,33 @@ +//! Forge tests for zksync contracts. + +use foundry_test_utils::util; + +#[tokio::test(flavor = "multi_thread")] +async fn test_zk_contract_paymaster() { + let (prj, mut cmd) = util::setup_forge( + "test_zk_contract_paymaster", + foundry_test_utils::foundry_compilers::PathStyle::Dapptools, + ); + util::initialize(prj.root()); + + cmd.args([ + "install", + "OpenZeppelin/openzeppelin-contracts", + "cyfrin/zksync-contracts", + "--no-commit", + "--shallow", + ]) + .ensure_execute_success() + .expect("able to install dependencies"); + + cmd.forge_fuse(); + + let config = cmd.config(); + prj.write_config(config); + + prj.add_source("MyPaymaster.sol", include_str!("../../fixtures/zk/MyPaymaster.sol")).unwrap(); + prj.add_source("Paymaster.t.sol", include_str!("../../fixtures/zk/Paymaster.t.sol")).unwrap(); + + cmd.args(["test", "--zk-startup", "--via-ir", "--match-contract", "TestPaymasterFlow"]); + assert!(cmd.stdout_lossy().contains("Suite result: ok")); +} diff --git a/crates/zksync/core/src/lib.rs b/crates/zksync/core/src/lib.rs index b93dc4013..b9fdd7432 100644 --- a/crates/zksync/core/src/lib.rs +++ b/crates/zksync/core/src/lib.rs @@ -74,6 +74,15 @@ pub fn get_nonce_key(address: Address) -> rU256 { zksync_types::get_nonce_key(&address.to_h160()).key().to_ru256() } +/// Represents additional data for ZK transactions that require a paymaster. +#[derive(Clone, Debug, Default)] +pub struct ZkPaymasterData { + /// Paymaster address. + pub address: Address, + /// Paymaster input. + pub input: Bytes, +} + /// Represents additional data for ZK transactions. #[derive(Clone, Debug, Default)] pub struct ZkTransactionMetadata { diff --git a/crates/zksync/core/src/vm/runner.rs b/crates/zksync/core/src/vm/runner.rs index 520f65f1f..3996107be 100644 --- a/crates/zksync/core/src/vm/runner.rs +++ b/crates/zksync/core/src/vm/runner.rs @@ -137,6 +137,15 @@ where let (gas_limit, max_fee_per_gas) = gas_params(ecx, caller); info!(?gas_limit, ?max_fee_per_gas, "tx gas parameters"); + let paymaster_params = if let Some(paymaster_data) = &ccx.paymaster_data { + PaymasterParams { + paymaster: paymaster_data.address.to_h160(), + paymaster_input: paymaster_data.input.to_vec(), + } + } else { + PaymasterParams::default() + }; + let tx = L2Tx::new( CONTRACT_DEPLOYER_ADDRESS, calldata, @@ -150,7 +159,7 @@ where caller.to_h160(), call.value.to_u256(), factory_deps, - PaymasterParams::default(), + paymaster_params, ); let call_ctx = CallContext { @@ -186,6 +195,16 @@ where let (gas_limit, max_fee_per_gas) = gas_params(ecx, caller); info!(?gas_limit, ?max_fee_per_gas, "tx gas parameters"); + + let paymaster_params = if let Some(paymaster_data) = &ccx.paymaster_data { + PaymasterParams { + paymaster: paymaster_data.address.to_h160(), + paymaster_input: paymaster_data.input.to_vec(), + } + } else { + PaymasterParams::default() + }; + let tx = L2Tx::new( call.bytecode_address.to_h160(), call.input.to_vec(), @@ -202,7 +221,7 @@ where _ => U256::zero(), }, factory_deps, - PaymasterParams::default(), + paymaster_params, ); // address and caller are specific to the type of call: diff --git a/crates/zksync/core/src/vm/tracer.rs b/crates/zksync/core/src/vm/tracer.rs index 0b765a539..661f1000f 100644 --- a/crates/zksync/core/src/vm/tracer.rs +++ b/crates/zksync/core/src/vm/tracer.rs @@ -28,7 +28,7 @@ use zksync_utils::bytecode::hash_bytecode; use crate::{ convert::{ConvertAddress, ConvertH160, ConvertH256, ConvertU256}, vm::farcall::{CallAction, CallDepth}, - EMPTY_CODE, + ZkPaymasterData, EMPTY_CODE, }; use super::farcall::FarCallHandler; @@ -82,6 +82,8 @@ pub struct CheatcodeTracerContext<'a> { pub accesses: Option<&'a mut RecordAccess>, /// Factory deps that were persisted across calls pub persisted_factory_deps: Option<&'a mut HashMap>>, + /// Paymaster data + pub paymaster_data: Option, } /// Tracer result to return back to foundry.