From f34e90bb2df5857ee51ba777bcd60d8bb484e4f4 Mon Sep 17 00:00:00 2001
From: Ariel Barmat <abarmat@gmail.com>
Date: Wed, 22 Sep 2021 11:09:55 -0300
Subject: [PATCH] Allow unstake passing a larger amount than available (#487)

* staking: allow unstake with larger amount than available
* staking: remove unused check and add comment
---
 contracts/staking/Staking.sol | 17 ++++++++++-----
 test/staking/staking.test.ts  | 41 +++++++++++++++++++++++++++++------
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/contracts/staking/Staking.sol b/contracts/staking/Staking.sol
index 61614eeb1..1672d707c 100644
--- a/contracts/staking/Staking.sol
+++ b/contracts/staking/Staking.sol
@@ -707,27 +707,34 @@ contract Staking is StakingV2Storage, GraphUpgradeable, IStaking {
 
     /**
      * @dev Unstake tokens from the indexer stake, lock them until thawing period expires.
+     * NOTE: The function accepts an amount greater than the currently staked tokens.
+     * If that happens, it will try to unstake the max amount of tokens it can.
+     * The reason for this behaviour is to avoid time conditions while the transaction
+     * is in flight.
      * @param _tokens Amount of tokens to unstake
      */
     function unstake(uint256 _tokens) external override notPartialPaused {
         address indexer = msg.sender;
         Stakes.Indexer storage indexerStake = stakes[indexer];
 
-        require(_tokens > 0, "!tokens");
         require(indexerStake.tokensStaked > 0, "!stake");
-        require(indexerStake.tokensAvailable() >= _tokens, "!stake-avail");
+
+        // Tokens to lock is capped to the available tokens
+        uint256 tokensToLock = MathUtils.min(indexerStake.tokensAvailable(), _tokens);
+        require(tokensToLock > 0, "!stake-avail");
 
         // Ensure minimum stake
-        uint256 newStake = indexerStake.tokensSecureStake().sub(_tokens);
+        uint256 newStake = indexerStake.tokensSecureStake().sub(tokensToLock);
         require(newStake == 0 || newStake >= minimumIndexerStake, "!minimumIndexerStake");
 
-        // Before locking more tokens, withdraw any unlocked ones
+        // Before locking more tokens, withdraw any unlocked ones if possible
         uint256 tokensToWithdraw = indexerStake.tokensWithdrawable();
         if (tokensToWithdraw > 0) {
             _withdraw(indexer);
         }
 
-        indexerStake.lockTokens(_tokens, thawingPeriod);
+        // Update the indexer stake locking tokens
+        indexerStake.lockTokens(tokensToLock, thawingPeriod);
 
         emit StakeLocked(indexer, indexerStake.tokensLocked, indexerStake.tokensLockedUntil);
     }
diff --git a/test/staking/staking.test.ts b/test/staking/staking.test.ts
index 5fc359f0d..8db3e530c 100644
--- a/test/staking/staking.test.ts
+++ b/test/staking/staking.test.ts
@@ -14,10 +14,11 @@ import {
   latestBlock,
   toBN,
   toGRT,
+  provider,
   Account,
 } from '../lib/testHelpers'
 
-const { AddressZero } = constants
+const { AddressZero, MaxUint256 } = constants
 
 function weightedAverage(
   valueA: BigNumber,
@@ -269,14 +270,18 @@ describe('Staking:Stakes', () => {
         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('!tokens')
+      it('should unstake available tokens even if passed a higher amount', async function () {
+        // Try to unstake a bit more than currently staked
+        const tokensOverCapacity = tokensToStake.add(toGRT('1'))
+        await staking.connect(indexer.signer).unstake(tokensOverCapacity)
+
+        // Check state
+        const tokensLocked = (await staking.stakes(indexer.address)).tokensLocked
+        expect(tokensLocked).eq(tokensToStake)
       })
 
-      it('reject unstake more than available tokens', async function () {
-        const tokensOverCapacity = tokensToStake.add(toGRT('1'))
-        const tx = staking.connect(indexer.signer).unstake(tokensOverCapacity)
+      it('reject unstake zero tokens', async function () {
+        const tx = staking.connect(indexer.signer).unstake(toGRT('0'))
         await expect(tx).revertedWith('!stake-avail')
       })
 
@@ -305,6 +310,28 @@ describe('Staking:Stakes', () => {
         await staking.connect(indexer.signer).unstake(tokensToStake)
         expect(await staking.getIndexerCapacity(indexer.address)).eq(0)
       })
+
+      it('should allow unstake of full amount with no upper limits', async function () {
+        // Use manual mining
+        await provider().send('evm_setAutomine', [false])
+
+        // Setup
+        const newTokens = toGRT('2')
+        const stakedTokens = await staking.getIndexerStakedTokens(indexer.address)
+        const tokensToUnstake = stakedTokens.add(newTokens)
+
+        // StakeTo & Unstake
+        await staking.connect(indexer.signer).stakeTo(indexer.address, newTokens)
+        await staking.connect(indexer.signer).unstake(MaxUint256)
+        await provider().send('evm_mine', [])
+
+        // Check state
+        const tokensLocked = (await staking.stakes(indexer.address)).tokensLocked
+        expect(tokensLocked).eq(tokensToUnstake)
+
+        // Restore automine
+        await provider().send('evm_setAutomine', [true])
+      })
     })
 
     describe('withdraw', function () {