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 974c3c8c0..9ba15ab11 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::{ @@ -360,6 +360,8 @@ pub struct Cheatcodes { /// to EVM. Alternatively, we'd need to add `vm.zkVmSkip()` to these calls manually. pub skip_zk_vm_addresses: HashSet
, + pub paymaster_params: Option, + /// Records the next create address for `skip_zk_vm_addresses`. pub record_next_create_address: bool, @@ -463,6 +465,7 @@ impl Cheatcodes { skip_zk_vm_addresses: Default::default(), record_next_create_address: Default::default(), persisted_factory_deps: Default::default(), + paymaster_params: None, } } @@ -959,6 +962,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), @@ -1550,6 +1554,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/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..7d3493ae7 --- /dev/null +++ b/crates/forge/tests/it/zk/paymaster.rs @@ -0,0 +1,37 @@ +//! Forge tests for zksync contracts. + +use foundry_config::fs_permissions::PathPermission; +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 mut config = cmd.config(); + config.fs_permissions.add(PathPermission::read("./zkout")); + prj.write_config(config); + + prj.add_source("MyPaymaster.sol", include_str!("../../../../../testdata/zk/MyPaymaster.sol")) + .unwrap(); + prj.add_source("Paymaster.t.sol", include_str!("../../../../../testdata/zk/Paymaster.t.sol")) + .unwrap(); + + cmd.args(["test", "--zk-startup", "--evm-version", "shanghai", "--via-ir"]); + 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. diff --git a/testdata/zk/MyPaymaster.sol b/testdata/zk/MyPaymaster.sol new file mode 100644 index 000000000..3cc330f88 --- /dev/null +++ b/testdata/zk/MyPaymaster.sol @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +import "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol"; +import "forge-std/console2.sol"; + +import {IPaymaster, ExecutionResult, PAYMASTER_VALIDATION_SUCCESS_MAGIC} from "../lib/zksync-contracts/zksync-contracts/l2/system-contracts/interfaces/IPaymaster.sol"; +import {IPaymasterFlow} from "../lib/zksync-contracts/zksync-contracts/l2/system-contracts/interfaces/IPaymasterFlow.sol"; +import {TransactionHelper, Transaction} from "../lib/zksync-contracts/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol"; +import "../lib/openzeppelin-contracts/contracts/token/ERC20/ERC20.sol"; +import "../lib/zksync-contracts/zksync-contracts/l2/system-contracts/Constants.sol"; + +contract MyPaymaster is IPaymaster { + uint256 constant PRICE_FOR_PAYING_FEES = 1; + + address public allowedToken; + + modifier onlyBootloader() { + require( + msg.sender == BOOTLOADER_FORMAL_ADDRESS, + "Only bootloader can call this method" + ); + // Continue execution if called from the bootloader. + _; + } + + constructor(address _erc20) { + allowedToken = _erc20; + } + + function validateAndPayForPaymasterTransaction( + bytes32, + bytes32, + Transaction calldata _transaction + ) + external + payable + onlyBootloader + returns (bytes4 magic, bytes memory context) + { + // By default we consider the transaction as accepted. + magic = PAYMASTER_VALIDATION_SUCCESS_MAGIC; + require( + _transaction.paymasterInput.length >= 4, + "The standard paymaster input must be at least 4 bytes long" + ); + + bytes4 paymasterInputSelector = bytes4( + _transaction.paymasterInput[0:4] + ); + if (paymasterInputSelector == IPaymasterFlow.approvalBased.selector) { + // While the transaction data consists of address, uint256 and bytes data, + // the data is not needed for this paymaster + (address token, uint256 amount, bytes memory data) = abi.decode( + _transaction.paymasterInput[4:], + (address, uint256, bytes) + ); + + // Verify if token is the correct one + require(token == allowedToken, "Invalid token"); + + // We verify that the user has provided enough allowance + address userAddress = address(uint160(_transaction.from)); + + address thisAddress = address(this); + + uint256 providedAllowance = IERC20(token).allowance( + userAddress, + thisAddress + ); + require( + providedAllowance >= PRICE_FOR_PAYING_FEES, + "Min allowance too low" + ); + + // Note, that while the minimal amount of ETH needed is tx.gasPrice * tx.gasLimit, + // neither paymaster nor account are allowed to access this context variable. + uint256 requiredETH = _transaction.gasLimit * + _transaction.maxFeePerGas; + try + IERC20(token).transferFrom(userAddress, thisAddress, amount) + {} catch (bytes memory revertReason) { + // If the revert reason is empty or represented by just a function selector, + // we replace the error with a more user-friendly message + if (revertReason.length <= 4) { + revert("Failed to transferFrom from users' account"); + } else { + assembly { + revert(add(0x20, revertReason), mload(revertReason)) + } + } + } + // The bootloader never returns any data, so it can safely be ignored here. + (bool success, ) = payable(BOOTLOADER_FORMAL_ADDRESS).call{ + value: 1 ether + }(""); + require( + success, + "Failed to transfer tx fee to the bootloader. Paymaster balance might not be enough." + ); + } else { + revert("Unsupported paymaster flow"); + } + } + + function postTransaction( + bytes calldata _context, + Transaction calldata _transaction, + bytes32, + bytes32, + ExecutionResult _txResult, + uint256 _maxRefundedGas + ) external payable override onlyBootloader { + } + + receive() external payable {} +} + +contract MyERC20 is ERC20 { + uint8 private _decimals; + + constructor( + string memory name_, + string memory symbol_, + uint8 decimals_ + ) ERC20(name_, symbol_) { + _decimals = decimals_; + } + + function mint(address _to, uint256 _amount) public returns (bool) { + _mint(_to, _amount); + return true; + } + + function decimals() public view override returns (uint8) { + return _decimals; + } +} diff --git a/testdata/zk/Paymaster.t.sol b/testdata/zk/Paymaster.t.sol new file mode 100644 index 000000000..2150061a4 --- /dev/null +++ b/testdata/zk/Paymaster.t.sol @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import {Test, console} from "forge-std/Test.sol"; +import "../lib/zksync-contracts/zksync-contracts/l2/system-contracts/Constants.sol"; +import {MyPaymaster, MyERC20} from "./MyPaymaster.sol"; + +contract TestPaymasterFlow is Test { + MyERC20 private erc20; + MyPaymaster private paymaster; + DoStuff private do_stuff; + address private alice; + bytes private paymaster_encoded_input; + + function setUp() public { + alice = makeAddr("Alice"); + do_stuff = new DoStuff(); + erc20 = new MyERC20("Test", "JR", 1); + paymaster = new MyPaymaster(address(erc20)); + + // Initial funding + vm.deal(address(do_stuff), 1 ether); + vm.deal(alice, 1 ether); + vm.deal(address(paymaster), 10 ether); + + // Mint and approve ERC20 tokens + erc20.mint(alice, 1); + vm.prank(alice, alice); + erc20.approve(address(paymaster), 1); + + // Encode paymaster input + paymaster_encoded_input = abi.encodeWithSelector( + bytes4(keccak256("approvalBased(address,uint256,bytes)")), + address(erc20), + uint256(1), + bytes("0x") + ); + } + + function testCallWithPaymaster() public { + 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(); + + require(address(do_stuff).balance == 0, "Balance is not 0 ether"); + require(address(alice).balance == alice_balance, "Balance is not the same"); + } + + 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(); + + 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() public { + (bool success, ) = payable(BOOTLOADER_FORMAL_ADDRESS).call{ + value: address(this).balance + }(""); + require(success, "Failed to transfer tx fee to the bootloader. Paymaster balance might not be enough."); + } +}