initial support for custom paymaster

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;

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<Address>,
 
+    pub paymaster_params: Option<ZkPaymasterData>,
+
     /// 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

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<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> 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<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self {

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;

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"));
+}

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 {

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:

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<H256, Vec<u8>>>,
+    /// Paymaster data
+    pub paymaster_data: Option<ZkPaymasterData>,
 }
 
 /// Tracer result to return back to foundry.

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;
+    }
+}