Merge pull request #65 from hats-finance/fix-54

Fix #54
hats-finance · Dec 4, 2023 · 19117c0 · 19117c0
2 parents fc0c464 + 933528f
commit 19117c0
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 3 deletions.
diff --git a/contracts/tokenlock/TokenLock.sol b/contracts/tokenlock/TokenLock.sol
@@ -42,6 +42,7 @@ abstract contract TokenLock is Ownable, ITokenLock {
     error AmountCannotBeZero();
     error AmountRequestedBiggerThanSurplus();
     error LockIsNonRevocable();
+    error LockIsAlreadyRevoked();
     error NoAvailableUnvestedAmount();
     error OnlySweeper();
     error CannotSweepVestedToken();
@@ -190,17 +191,27 @@ abstract contract TokenLock is Ownable, ITokenLock {
         if (!revocable)
             revert LockIsNonRevocable();
 
-        uint256 unvestedAmount = managedAmount - vestedAmount();
+        if (isRevoked)
+            revert LockIsAlreadyRevoked();
+
+        uint256 vestedAmount = vestedAmount();
+
+        uint256 unvestedAmount = managedAmount - vestedAmount;
         if (unvestedAmount == 0)
             revert NoAvailableUnvestedAmount();
 
         isRevoked = true;
 
+        managedAmount = vestedAmount;
+
+        // solhint-disable-next-line not-rely-on-time
+        endTime = block.timestamp;
+
         token.safeTransfer(owner(), unvestedAmount);
 
         emit TokensRevoked(beneficiary, unvestedAmount);
 
-        selfdestruct(payable(msg.sender));
+        trySelfDestruct();
     }
 
     /**
@@ -445,7 +456,7 @@ abstract contract TokenLock is Ownable, ITokenLock {
     }
 
     function trySelfDestruct() private {
-        if (currentTime() > endTime && currentBalance() == 0) {
+        if (currentTime() >= endTime && currentBalance() == 0) {
             selfdestruct(payable(msg.sender));
         }
     }

diff --git a/docs/dodoc/tokenlock/HATTokenLock.md b/docs/dodoc/tokenlock/HATTokenLock.md
@@ -849,6 +849,17 @@ error DelegateDisabled()
 
 
 
+### LockIsAlreadyRevoked
+
+```solidity
+error LockIsAlreadyRevoked()
+```
+
+
+
+
+
+
 ### LockIsNonRevocable
 
 ```solidity

diff --git a/docs/dodoc/tokenlock/TokenLock.md b/docs/dodoc/tokenlock/TokenLock.md
@@ -779,6 +779,17 @@ error CliffTimeMustBeBeforeEndTime()
 
 
 
+### LockIsAlreadyRevoked
+
+```solidity
+error LockIsAlreadyRevoked()
+```
+
+
+
+
+
+
 ### LockIsNonRevocable
 
 ```solidity

diff --git a/test/tokenlock.js b/test/tokenlock.js
@@ -295,6 +295,18 @@ contract("TokenLock", (accounts) => {
     }
   });
 
+  it("cannot revoke twice", async () => {
+    await setup(accounts);
+    await utils.increaseTime(300);
+    await tokenLock.revoke();
+    try {
+      await tokenLock.revoke();
+      assert(false, "cannot revoke after lock was already revoked");
+    } catch (ex) {
+      assertVMException(ex, "LockIsAlreadyRevoked");
+    }
+  });
+
   it("sinceStartTime", async () => {
     await setup(accounts, true, false, 100);
     assert.equal(await tokenLock.sinceStartTime(), 0);
@@ -305,6 +317,22 @@ contract("TokenLock", (accounts) => {
     assert.equal(await tokenLock.availableAmount(), web3.utils.toWei("0.2"));
   });
 
+  it("redeem and withdraw surplus after revoke", async () => {
+    await setup(accounts);
+    await stakingToken.mint(tokenLock.address, web3.utils.toWei("1"));
+    await utils.increaseTime(300);
+    await tokenLock.revoke();
+    assert.equal(await stakingToken.balanceOf(accounts[0]), web3.utils.toWei("0.8"));
+    assert.equal(await stakingToken.balanceOf(tokenLock.address), web3.utils.toWei("1.2"));
+    await tokenLock.release({ from: accounts[1] });
+    assert.equal(await stakingToken.balanceOf(accounts[1]), web3.utils.toWei("0.2"));
+    assert.equal(await stakingToken.balanceOf(tokenLock.address), web3.utils.toWei("1"));
+    await tokenLock.withdrawSurplus(web3.utils.toWei("1"), { from: accounts[1] });
+    assert.equal(await stakingToken.balanceOf(accounts[1]), web3.utils.toWei("1.2"));
+    assert.equal(await stakingToken.balanceOf(tokenLock.address), web3.utils.toWei("0"));
+    expect(await ethers.provider.getCode(tokenLock.address)).to.equal("0x");
+  });
+
   it("cannot revoke after renaunceOwnership revoke", async () => {
     await setup(accounts);
     await tokenLock.renounceOwnership();