Paymaster: handle deposit and stake based on entrypoint()

OpenZeppelin · Feb 4, 2025 · bc0c040 · bc0c040
1 parent 3b1a9fa
commit bc0c040
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 9 deletions.
diff --git 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
@@ -135,23 +135,29 @@ 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);
     });
   });
 
   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,20 +169,37 @@ 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
       await expect(this.paymaster.connect(this.admin).withdrawStake(this.receiver)).to.changeEtherBalances(
         [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 () {