From bc0c0401ecee0a27f0e41265863fd8df4d247fdc Mon Sep 17 00:00:00 2001
From: Hadrien Croubois <hadrien.croubois@gmail.com>
Date: Tue, 4 Feb 2025 14:29:29 +0100
Subject: [PATCH] Paymaster: handle deposit and stake based on entrypoint()

---
 contracts/account/paymaster/PaymasterCore.sol | 10 +++---
 test/account/paymaster/Paymaster.behavior.js  | 31 ++++++++++++++++---
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/contracts/account/paymaster/PaymasterCore.sol b/contracts/account/paymaster/PaymasterCore.sol
index 67ae0782..a5946ab4 100644
--- a/contracts/account/paymaster/PaymasterCore.sol
+++ b/contracts/account/paymaster/PaymasterCore.sol
@@ -93,27 +93,27 @@ abstract contract PaymasterCore is IPaymaster {
 
     /// @dev Calls {IEntryPointStake-depositTo}.
     function deposit() public payable virtual {
-        ERC4337Utils.depositTo(address(this), msg.value);
+        entryPoint().depositTo{value: msg.value}(address(this));
     }
 
     /// @dev Calls {IEntryPointStake-withdrawTo}.
     function withdraw(address payable to, uint256 value) public virtual onlyWithdrawer {
-        ERC4337Utils.withdrawTo(to, value);
+        entryPoint().withdrawTo(to, value);
     }
 
     /// @dev Calls {IEntryPointStake-addStake}.
     function addStake(uint32 unstakeDelaySec) public payable virtual {
-        ERC4337Utils.addStake(msg.value, unstakeDelaySec);
+        entryPoint().addStake{value: msg.value}(unstakeDelaySec);
     }
 
     /// @dev Calls {IEntryPointStake-unlockStake}.
     function unlockStake() public virtual onlyWithdrawer {
-        ERC4337Utils.unlockStake();
+        entryPoint().unlockStake();
     }
 
     /// @dev Calls {IEntryPointStake-withdrawStake}.
     function withdrawStake(address payable to) public virtual onlyWithdrawer {
-        ERC4337Utils.withdrawStake(to);
+        entryPoint().withdrawStake(to);
     }
 
     /// @dev Ensures the caller is the {entrypoint}.
diff --git a/test/account/paymaster/Paymaster.behavior.js b/test/account/paymaster/Paymaster.behavior.js
index 3f74256f..f25296bf 100644
--- a/test/account/paymaster/Paymaster.behavior.js
+++ b/test/account/paymaster/Paymaster.behavior.js
@@ -135,7 +135,7 @@ function shouldBehaveLikePaymaster({ postOp } = { postOp: false }) {
 
   describe('postOp', function () {
     it('reverts if the caller is not the entrypoint', async function () {
-      await expect(this.paymaster.connect(this.other).postOp(0, '0x', 0, 0))
+      await expect(this.paymaster.connect(this.other).postOp(0n, '0x', 0n, 0n))
         .to.be.revertedWithCustomError(this.paymaster, 'PaymasterUnauthorized')
         .withArgs(this.other);
     });
@@ -143,15 +143,21 @@ function shouldBehaveLikePaymaster({ postOp } = { postOp: false }) {
 
   describe('deposit lifecycle', function () {
     it('deposits and withdraws effectively', async function () {
+      await expect(entrypoint.balanceOf(this.paymaster)).to.eventually.equal(0n);
+
       await expect(this.paymaster.connect(this.other).deposit({ value })).to.changeEtherBalances(
         [this.other, entrypoint],
         [-value, value],
       );
 
-      await expect(this.paymaster.connect(this.admin).withdraw(this.receiver, value)).to.changeEtherBalances(
+      await expect(entrypoint.balanceOf(this.paymaster)).to.eventually.equal(value);
+
+      await expect(this.paymaster.connect(this.admin).withdraw(this.receiver, 1n)).to.changeEtherBalances(
         [entrypoint, this.receiver],
-        [-value, value],
+        [-1n, 1n],
       );
+
+      await expect(entrypoint.balanceOf(this.paymaster)).to.eventually.equal(value - 1n);
     });
 
     it('reverts when an unauthorized caller tries to withdraw', async function () {
@@ -163,13 +169,28 @@ function shouldBehaveLikePaymaster({ postOp } = { postOp: false }) {
 
   describe('stake lifecycle', function () {
     it('adds and removes stake effectively', async function () {
+      await expect(entrypoint.deposits(this.paymaster)).to.eventually.deep.equal([0n, false, 0n, 0n, 0n]);
+
       // stake
       await expect(this.paymaster.connect(this.other).addStake(delay, { value })).to.changeEtherBalances(
         [this.other, entrypoint],
         [-value, value],
       );
+
+      await expect(entrypoint.deposits(this.paymaster)).to.eventually.deep.equal([0n, true, 42n, delay, 0n]);
+
       // unlock
-      await this.paymaster.connect(this.admin).unlockStake();
+      const unlockTx = this.paymaster.connect(this.admin).unlockStake();
+
+      const timestamp = await time.clockFromReceipt.timestamp(unlockTx);
+      await expect(entrypoint.deposits(this.paymaster)).to.eventually.deep.equal([
+        0n,
+        false,
+        42n,
+        delay,
+        timestamp + delay,
+      ]);
+
       await time.increaseBy.timestamp(delay);
 
       // withdraw stake
@@ -177,6 +198,8 @@ function shouldBehaveLikePaymaster({ postOp } = { postOp: false }) {
         [entrypoint, this.receiver],
         [-value, value],
       );
+
+      await expect(entrypoint.deposits(this.paymaster)).to.eventually.deep.equal([0n, false, 0n, 0n, 0n]);
     });
 
     it('reverts when an unauthorized caller tries to unlock stake', async function () {