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::{
@@ -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<ZkPaymasterData>,
+
     /// 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

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,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 {}
+}

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

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

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.