staking: withdraw any unlocked tokens before unstaking

This fixes: tokens ready for withdrawal can get locked again by future unstaking calls.

- Withdraw tokens that already passed the thawing period when unstake is called
- Add require to avoid unstaking zero tokens

Author: abarmat

contracts/staking/Staking.sol

@@ -551,6 +551,7 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
         address indexer = msg.sender;
         Stakes.Indexer storage indexerStake = stakes[indexer];
 
+        require(_tokens > 0, "Cannot unstake zero tokens");
         require(indexerStake.hasTokens(), "Indexer has no stakes");
         require(
             indexerStake.tokensAvailable() >= _tokens,
@@ -564,6 +565,12 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
             "Stake must be above minimum required"
         );
 
+        // Before locking more tokens, withdraw any unlocked ones
+        uint256 tokensToWithdraw = indexerStake.tokensWithdrawable();
+        if (tokensToWithdraw > 0) {
+            _withdraw(indexer);
+        }
+
         indexerStake.lockTokens(_tokens, thawingPeriod);
 
         emit StakeLocked(indexer, indexerStake.tokensLocked, indexerStake.tokensLockedUntil);
@@ -573,16 +580,7 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
      * @dev Withdraw indexer tokens once the thawing period has passed.
      */
     function withdraw() external override notPaused {
-        address indexer = msg.sender;
-
-        // Get tokens available for withdraw and update balance
-        uint256 tokensToWithdraw = stakes[indexer].withdrawTokens();
-        require(tokensToWithdraw > 0, "No tokens available to withdraw");
-
-        // Return tokens to the indexer
-        require(graphToken().transfer(indexer, tokensToWithdraw), "!transfer");
-
-        emit StakeWithdrawn(indexer, tokensToWithdraw);
+        _withdraw(msg.sender);
     }
 
     /**
@@ -881,6 +879,23 @@ contract Staking is StakingV1Storage, GraphUpgradeable, IStaking {
         emit StakeDeposited(_indexer, _tokens);
     }
 
+    /**
+     * @dev Withdraw indexer tokens once the thawing period has passed.
+     * @param _indexer Address of indexer to withdraw funds from
+     */
+    function _withdraw(address _indexer) internal {
+        Stakes.Indexer storage indexerStake = stakes[_indexer];
+
+        // Get tokens available for withdraw and update balance
+        uint256 tokensToWithdraw = indexerStake.withdrawTokens();
+        require(tokensToWithdraw > 0, "No tokens available to withdraw");
+
+        // Return tokens to the indexer
+        require(graphToken().transfer(_indexer, tokensToWithdraw), "!transfer");
+
+        emit StakeWithdrawn(_indexer, tokensToWithdraw);
+    }
+
     /**
      * @dev Allocate available tokens to a subgraph deployment.
      * @param _indexer Indexer address to allocate funds from.

contracts/staking/libs/Stakes.sol

@@ -56,20 +56,20 @@ library Stakes {
     }
 
     /**
-     * @dev Lock tokens until a thawing period expires.
+     * @dev Lock tokens until a thawing period pass.
      * @param stake Stake data
      * @param _tokens Amount of tokens to unstake
-     * @param _thawingPeriod Period in blocks that need to pass before withdrawal
+     * @param _period Period in blocks that need to pass before withdrawal
      */
     function lockTokens(
         Stakes.Indexer storage stake,
         uint256 _tokens,
-        uint256 _thawingPeriod
+        uint256 _period
     ) internal {
         // Take into account period averaging for multiple unstake requests
-        uint256 lockingPeriod = _thawingPeriod;
+        uint256 lockingPeriod = _period;
         if (stake.tokensLocked > 0) {
-            lockingPeriod = stake.getLockingPeriod(_tokens, _thawingPeriod);
+            lockingPeriod = stake.getLockingPeriod(_tokens, _period);
         }
 
         // Update balances

test/staking/staking.test.ts

@@ -197,10 +197,10 @@ describe('Staking:Stakes', () => {
         const tx1 = await staking.connect(indexer.signer).unstake(tokensToUnstake)
         const receipt1 = await tx1.wait()
         const event1: Event = receipt1.events.pop()
-        const tokensLockedUntil1 = event1.args[2]
+        const tokensLockedUntil1 = event1.args['until']
 
-        // Move forward
-        await advanceBlockTo(tokensLockedUntil1)
+        // Move forward before the tokens are unlocked for withdrawal
+        await advanceBlockTo(tokensLockedUntil1.sub(5))
 
         // Calculate locking time for tokens taking into account the previous unstake request
         const currentBlock = await latestBlock()
@@ -215,9 +215,57 @@ describe('Staking:Stakes', () => {
         // Unstake (2)
         const tx2 = await staking.connect(indexer.signer).unstake(tokensToUnstake)
         const receipt2 = await tx2.wait()
+
+        // Verify events
         const event2: Event = receipt2.events.pop()
-        const tokensLockedUntil2 = event2.args[2]
-        expect(expectedLockedUntil).eq(tokensLockedUntil2)
+        expect(event2.args['until']).eq(expectedLockedUntil)
+
+        // Verify state
+        const afterIndexerStake = await staking.stakes(indexer.address)
+        expect(afterIndexerStake.tokensLocked).eq(tokensToUnstake.mul(2)) // we unstaked two times
+        expect(afterIndexerStake.tokensLockedUntil).eq(expectedLockedUntil)
+      })
+
+      it('should unstake and withdraw if some tokens are unthawed', async function () {
+        const tokensToUnstake = toGRT('10')
+        const thawingPeriod = toBN(await staking.thawingPeriod())
+
+        const beforeIndexerBalance = await grt.balanceOf(indexer.address)
+
+        // Unstake (1)
+        const tx1 = await staking.connect(indexer.signer).unstake(tokensToUnstake)
+        const receipt1 = await tx1.wait()
+        const event1: Event = receipt1.events.pop()
+        const tokensLockedUntil1 = event1.args['until']
+
+        // Move forward after the tokens are unlocked for withdrawal
+        await advanceBlockTo(tokensLockedUntil1)
+
+        // Calculate locking time for tokens taking into account some tokens are withdraweable
+        const currentBlock = await latestBlock()
+        const expectedLockedUntil = currentBlock.add(thawingPeriod).add(toBN('1'))
+
+        // Unstake (2)
+        const tx2 = await staking.connect(indexer.signer).unstake(tokensToUnstake)
+        const receipt2 = await tx2.wait()
+
+        // Verify events
+        const unstakeEvent: Event = receipt2.events.pop()
+        const withdrawEvent: Event = receipt2.events.pop()
+        expect(unstakeEvent.args['until']).eq(expectedLockedUntil)
+        expect(withdrawEvent.args['tokens']).eq(tokensToUnstake)
+
+        // Verify state
+        const afterIndexerStake = await staking.stakes(indexer.address)
+        const afterIndexerBalance = await grt.balanceOf(indexer.address)
+        expect(afterIndexerStake.tokensLocked).eq(tokensToUnstake)
+        expect(afterIndexerStake.tokensLockedUntil).eq(expectedLockedUntil)
+        expect(afterIndexerBalance).eq(beforeIndexerBalance.add(tokensToUnstake))
+      })
+
+      it('reject unstake zero tokens', async function () {
+        const tx = staking.connect(indexer.signer).unstake(toGRT('0'))
+        await expect(tx).revertedWith('Cannot unstake zero tokens')
       })
 
       it('reject unstake more than available tokens', async function () {