Add stake management function to ERC4337Utils

OpenZeppelin · Jan 31, 2025 · 5bc48ae · 5bc48ae
1 parent 840c974
commit 5bc48ae
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 5 deletions.
diff --git a/contracts/account/utils/draft-ERC4337Utils.sol b/contracts/account/utils/draft-ERC4337Utils.sol
@@ -3,7 +3,7 @@
 
 pragma solidity ^0.8.20;
 
-import {PackedUserOperation} from "../../interfaces/draft-IERC4337.sol";
+import {IEntryPoint, PackedUserOperation} from "../../interfaces/draft-IERC4337.sol";
 import {Math} from "../../utils/math/Math.sol";
 import {Calldata} from "../../utils/Calldata.sol";
 import {Packing} from "../../utils/Packing.sol";
@@ -16,6 +16,9 @@ import {Packing} from "../../utils/Packing.sol";
 library ERC4337Utils {
     using Packing for *;
 
+    /// @dev Address of the entrypoint v0.7.0
+    IEntryPoint internal constant ENTRYPOINT = IEntryPoint(0x0000000071727De22E5E9d8BAf0edAc6f37da032);
+
     /// @dev For simulation purposes, validateUserOp (and validatePaymasterUserOp) return this value on success.
     uint256 internal constant SIG_VALIDATION_SUCCESS = 0;
 
@@ -160,4 +163,29 @@ library ERC4337Utils {
     function paymasterData(PackedUserOperation calldata self) internal pure returns (bytes calldata) {
         return self.paymasterAndData.length < 52 ? Calldata.emptyBytes() : self.paymasterAndData[52:];
     }
+
+    /// @dev Deposit ether into the entrypoint.
+    function depositTo(address to, uint256 value) internal {
+        ENTRYPOINT.depositTo{value: value}(to);
+    }
+
+    /// @dev Withdraw ether from the entrypoint.
+    function withdrawTo(address payable to, uint256 value) internal {
+        ENTRYPOINT.withdrawTo(to, value);
+    }
+
+    /// @dev Add stake to the entrypoint.
+    function addStake(uint256 value, uint32 unstakeDelaySec) internal {
+        ENTRYPOINT.addStake{value: value}(unstakeDelaySec);
+    }
+
+    /// @dev Unlock stake on the entrypoint.
+    function unlockStake() internal {
+        ENTRYPOINT.unlockStake();
+    }
+
+    /// @dev Withdraw unlocked stake from the entrypoint.
+    function withdrawStake(address payable to) internal {
+        ENTRYPOINT.withdrawStake(to);
+    }
 }
diff --git a/test/account/utils/draft-ERC4337Utils.test.js b/test/account/utils/draft-ERC4337Utils.test.js
@@ -4,15 +4,16 @@ const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers');
 
 const { packValidationData, UserOperation } = require('../../helpers/erc4337');
 const { MAX_UINT48 } = require('../../helpers/constants');
+const time = require('../../helpers/time');
 const ADDRESS_ONE = '0x0000000000000000000000000000000000000001';
 
 const fixture = async () => {
-  const [authorizer, sender, factory, paymaster] = await ethers.getSigners();
+  const [authorizer, sender, factory, paymaster, other] = await ethers.getSigners();
   const utils = await ethers.deployContract('$ERC4337Utils');
   const SIG_VALIDATION_SUCCESS = await utils.$SIG_VALIDATION_SUCCESS();
   const SIG_VALIDATION_FAILED = await utils.$SIG_VALIDATION_FAILED();
 
-  return { utils, authorizer, sender, factory, paymaster, SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED };
+  return { utils, authorizer, sender, factory, paymaster, other, SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED };
 };
 
 describe('ERC4337Utils', function () {
@@ -284,4 +285,68 @@ describe('ERC4337Utils', function () {
       });
     });
   });
+
+  describe('stake management', function () {
+    const unstakeDelaySec = 3600n;
+
+    beforeEach(async function () {
+      await this.authorizer.sendTransaction({ to: this.utils, value: ethers.parseEther('1') });
+    });
+
+    it('deposit & withdraw', async function () {
+      await expect(entrypoint.balanceOf(this.utils)).to.eventually.equal(0n);
+
+      // deposit
+      await expect(this.utils.$depositTo(this.utils, 42n)).to.changeEtherBalances(
+        [this.utils, entrypoint],
+        [-42n, 42n],
+      );
+
+      await expect(entrypoint.balanceOf(this.utils)).to.eventually.equal(42n);
+
+      // withdraw
+      await expect(this.utils.$withdrawTo(this.other, 17n)).to.changeEtherBalances(
+        [entrypoint, this.other],
+        [-17n, 17n],
+      );
+
+      await expect(entrypoint.balanceOf(this.utils)).to.eventually.equal(25n); // 42 - 17
+    });
+
+    it('stake, unlock & withdraw stake', async function () {
+      await expect(entrypoint.deposits(this.utils)).to.eventually.deep.equal([0n, false, 0n, 0n, 0n]);
+
+      // stake
+      await expect(this.utils.$addStake(42n, unstakeDelaySec)).to.changeEtherBalances(
+        [this.utils, entrypoint],
+        [-42n, 42n],
+      );
+
+      await expect(entrypoint.deposits(this.utils)).to.eventually.deep.equal([0n, true, 42n, unstakeDelaySec, 0n]);
+
+      // unlock
+      const unlockTx = this.utils.$unlockStake();
+      await expect(unlockTx).to.changeEtherBalances([this.utils, entrypoint], [0n, 0n]); // no ether movement
+
+      const timestamp = await time.clockFromReceipt.timestamp(unlockTx);
+      await expect(entrypoint.deposits(this.utils)).to.eventually.deep.equal([
+        0n,
+        false,
+        42n,
+        unstakeDelaySec,
+        timestamp + unstakeDelaySec,
+      ]);
+
+      // wait
+      await time.increaseBy.timestamp(unstakeDelaySec);
+
+      // withdraw stake
+      await expect(this.utils.$withdrawStake(this.other)).to.changeEtherBalances(
+        [this.utils, entrypoint, this.other],
+        [0n, -42n, 42n],
+      );
+
+      await expect(entrypoint.deposits(this.utils)).to.eventually.deep.equal([0n, false, 0n, 0n, 0n]);
+    });
+  });
 });
diff --git a/test/helpers/time.js b/test/helpers/time.js
@@ -7,8 +7,11 @@ const clock = {
   timestamp: () => time.latest().then(ethers.toBigInt),
 };
 const clockFromReceipt = {
-  blocknumber: receipt => Promise.resolve(ethers.toBigInt(receipt.blockNumber)),
-  timestamp: receipt => ethers.provider.getBlock(receipt.blockNumber).then(block => ethers.toBigInt(block.timestamp)),
+  blocknumber: receipt => Promise.resolve(receipt).then(({ blockNumber }) => ethers.toBigInt(blockNumber)),
+  timestamp: receipt =>
+    Promise.resolve(receipt)
+      .then(({ blockNumber }) => ethers.provider.getBlock(blockNumber))
+      .then(({ timestamp }) => ethers.toBigInt(timestamp)),
 };
 const increaseBy = {
   blockNumber: mine,