feat: 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,
         }
     }
 
@@ -976,6 +979,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 +1571,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/fixtures/zk/MyPaymaster.sol

@@ -0,0 +1,114 @@
+// 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;
+    }
+}

crates/forge/tests/fixtures/zk/Paymaster.t.sol

@@ -0,0 +1,87 @@
+// 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.");
+    }
+}

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,35 @@
+//! 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!("../../fixtures/zk/MyPaymaster.sol")).unwrap();
+    prj.add_source("Paymaster.t.sol", include_str!("../../fixtures/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: