From fe2887d4ead78ebc122a0bc5493399e0933ce646 Mon Sep 17 00:00:00 2001
From: Li-Qing Wang <liqingnz@gmail.com>
Date: Thu, 23 Feb 2023 02:13:57 +0800
Subject: [PATCH] fix the storage type for DepositInfo (EntryPoint:withdrawTo)
 from memory to storage

---
 contracts/StakeManager.sol |  2 +-
 test/wallet.test.ts        | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/contracts/StakeManager.sol b/contracts/StakeManager.sol
index d1077ef..23952c6 100644
--- a/contracts/StakeManager.sol
+++ b/contracts/StakeManager.sol
@@ -144,7 +144,7 @@ abstract contract StakeManager {
      * @param withdrawAmount the amount to withdraw.
      */
     function withdrawTo(address payable withdrawAddress, uint256 withdrawAmount) external {
-        DepositInfo memory info = deposits[msg.sender];
+        DepositInfo storage info = deposits[msg.sender];
         require(withdrawAmount <= info.deposit, "Withdraw amount too large");
         info.deposit = uint112(info.deposit - withdrawAmount);
         emit Withdrawn(msg.sender, withdrawAddress, withdrawAmount);
diff --git a/test/wallet.test.ts b/test/wallet.test.ts
index 505954d..6fe95b7 100644
--- a/test/wallet.test.ts
+++ b/test/wallet.test.ts
@@ -149,6 +149,17 @@ describe("Wallet testing", () => {
       expect(await testNft.ownerOf(0)).eq(testSimpleWallet.address);
     });
 
+    it("test deposit/withdraw", async () => {
+      expect(await testSimpleWallet.getDeposit()).to.eq(BigNumber.from(0));
+      await testSimpleWallet.addDeposit({ value: TWO_ETH });
+      expect(await testSimpleWallet.getDeposit()).to.eq(TWO_ETH);
+      await testSimpleWallet.withdrawDepositTo(deployerAddress, ONE_ETH);
+      expect(await testSimpleWallet.getDeposit()).to.eq(ONE_ETH);
+      await expect(testSimpleWallet.withdrawDepositTo(deployerAddress, TWO_ETH)).to.be.revertedWith(
+        "Withdraw amount too large",
+      );
+    });
+
     it("test trusted paymaster", async () => {
       await deployer.sendTransaction({
         from: deployerAddress,
@@ -170,6 +181,9 @@ describe("Wallet testing", () => {
     it("test initialization/upgradeability with ownership", async () => {
       expect(await testSimpleWallet.owner()).to.eq(deployerAddress);
       expect(await maskToken.allowance(testSimpleWallet.address, beneficialAccountAddress)).to.eq(ONE_ETH);
+      await expect(testSimpleWallet.updateEntryPoint(AddressZero))
+        .to.emit(testSimpleWallet, "EntryPointChanged")
+        .withArgs(entryPoint.address, AddressZero);
       // only using "testSimpleWallet.address" for upgrade testing, could use any address here
       await expect(
         testProxy.connect(beneficialAccount).upgradeToAndCall(testSimpleWallet.address, "0x", false),
@@ -249,6 +263,7 @@ describe("Wallet testing", () => {
       await maskToken.transfer(walletProxyAddress, utils.parseEther("100"));
       // depositPaymaster setup
       depositPaymaster = await new DepositPaymaster__factory(deployer).deploy(entryPoint.address, maskToken.address);
+      await expect(depositPaymaster.withdrawStake(deployerAddress)).to.be.revertedWith("No stake to withdraw");
       await depositPaymaster.addStake(0, { value: TWO_ETH });
       await maskToken.approve(depositPaymaster.address, constants.MaxUint256);
       await depositPaymaster.connect(deployer).adjustAdmin(await deployer.getAddress(), true);
@@ -286,6 +301,23 @@ describe("Wallet testing", () => {
       await depositPaymaster.addDepositFor(walletProxyAddress, TWO_ETH);
     });
 
+    it("staking through paymaster", async () => {
+      expect((await entryPoint.getDepositInfo(depositPaymaster.address)).stake).to.be.eq(TWO_ETH);
+      const tempSigner = createWallet();
+      await expect(depositPaymaster.withdrawStake(tempSigner.address)).to.be.revertedWith(
+        "must call unlockStake() first",
+      );
+      await depositPaymaster.unlockStake();
+      await expect(depositPaymaster.unlockStake()).to.be.revertedWith("already unstaking");
+      await expect(depositPaymaster.withdrawStake(tempSigner.address)).to.be.revertedWith(
+        "Stake withdrawal is not due",
+      );
+      await network.provider.send("evm_increaseTime", [20]);
+      await depositPaymaster.withdrawStake(tempSigner.address);
+      expect((await entryPoint.getDepositInfo(depositPaymaster.address)).stake).to.be.eq(BigNumber.from(0));
+      expect(await ethers.provider.getBalance(tempSigner.address)).to.be.eq(TWO_ETH);
+    });
+
     it("wallet deployment through EntryPoint", async () => {
       let userOperation = createDefaultUserOp(walletProxyAddress);
       userOperation.nonce = salt;