updated exclusivity logic

Author: aalavandhan

contracts/ExclusiveGeyser.sol

@@ -2,12 +2,13 @@
 pragma solidity 0.7.6;
 pragma abicoder v2;
 
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {IUniversalVault} from "./UniversalVault.sol";
 import {Geyser} from "./Geyser.sol";
 
 /// @title ExclusiveGeyser
-/// @notice A special extension of GeyserV2 which allows staking in,
-///         at most one distribution program in any given time, for a given staking token.
+/// @notice A special extension of GeyserV2 which limits staking tokens,
+///         only distribution program at any given time.
 /// @dev Security contact: dev-support@ampleforth.org
 contract ExclusiveGeyser is Geyser {
     /// @inheritdoc Geyser
@@ -16,21 +17,31 @@ contract ExclusiveGeyser is Geyser {
         uint256 amount,
         bytes calldata permission
     ) public override {
-        // verify that vault has NOT locked staking token in other programs
-        _enforceExclusiveStake(IUniversalVault(vault));
+        // verify that vault isn't staking the same tokens in multiple programs
+        _enforceExclusiveStake(IUniversalVault(vault), amount);
 
         // continue with regular stake
         super.stake(vault, amount, permission);
     }
 
-    function _enforceExclusiveStake(IUniversalVault vault) private view {
+    /// @dev Enforces that the vault can't use tokens which have already been staked.
+    function _enforceExclusiveStake(IUniversalVault vault, uint256 amount) private view {
+        require(amount <= computeAvailableStakingBalance(vault), "ExclusiveGeyser: expected exclusive stake");
+    }
+
+    /// @notice Computes the amount of staking tokens in the vault available to be staked exclusively.
+    function computeAvailableStakingBalance(IUniversalVault vault) public view returns (uint256) {
+        // Iterates through the vault's locks to compute the total "stakingToken" balance staked across all geysers.
         address stakingToken = super.getGeyserData().stakingToken;
+        uint256 vaultBal = IERC20(stakingToken).balanceOf(address(vault));
+        uint256 totalStakedBal = 0;
         uint256 lockCount = vault.getLockSetCount();
         for (uint256 i = 0; i < lockCount; i++) {
             IUniversalVault.LockData memory lock = vault.getLockAt(i);
             if (lock.token == stakingToken) {
-                require(lock.delegate == address(this), "ExclusiveGeyser: expected exclusive stake");
+                totalStakedBal += lock.balance;
             }
         }
+        return vaultBal - totalStakedBal;
     }
 }

frontend/src/utils/stakingToken.ts

@@ -141,12 +141,10 @@ const uniswapV2Pair = async (
   const totalSupply: BigNumber = await contract.totalSupply()
   const totalSupplyNumber = parseFloat(formatUnits(totalSupply, decimals))
 
-  const tokenCompositions = await getTokenCompositions(
-    [token0Address, token1Address],
-    address,
-    signerOrProvider,
-    [0.5, 0.5],
-  )
+  const tokenCompositions = await getTokenCompositions([token0Address, token1Address], address, signerOrProvider, [
+    0.5,
+    0.5,
+  ])
   const [token0Symbol, token1Symbol] = tokenCompositions.map((c) => c.symbol)
   const marketCap = getMarketCap(tokenCompositions)
 
@@ -183,12 +181,10 @@ const getMooniswap = async (tokenAddress: string, signerOrProvider: SignerOrProv
 
   const totalSupplyNumber = parseFloat(formatUnits(totalSupply, decimals))
 
-  const tokenCompositions = await getTokenCompositions(
-    [token0Address, token1Address],
-    address,
-    signerOrProvider,
-    [0.5, 0.5],
-  )
+  const tokenCompositions = await getTokenCompositions([token0Address, token1Address], address, signerOrProvider, [
+    0.5,
+    0.5,
+  ])
   const marketCap = getMarketCap(tokenCompositions)
 
   return {

test/ExclusiveGeyser.ts

@@ -1031,8 +1031,9 @@ describe('ExclusiveGeyser', function () {
       })
       describe('with insufficient balance', function () {
         it('should fail', async function () {
+          // Exclusive stake condition fails first
           await expect(stake(user, geyser, vault, stakingToken, stakeAmount.mul(2))).to.be.revertedWith(
-            'UniversalVault: insufficient balance',
+            'ExclusiveGeyser: expected exclusive stake',
           )
         })
       })
@@ -1123,7 +1124,7 @@ describe('ExclusiveGeyser', function () {
             }
           })
           it('should fail', async function () {
-            await expect(stake(user, geyser, vault, stakingToken, stakeAmount.div(quantity))).to.be.revertedWith(
+            await expect(stake(user, geyser, vault, stakingToken, 1)).to.be.revertedWith(
               'Geyser: MAX_STAKES_PER_VAULT reached',
             )
           })
@@ -1176,18 +1177,31 @@ describe('ExclusiveGeyser', function () {
           ]
           otherGeyser = await deployGeyser(args, 'Geyser')
           await otherGeyser.connect(admin).registerVaultFactory(vaultFactory.address)
-          await stake(user, otherGeyser, vault, stakingToken, 1)
+          await stake(user, otherGeyser, vault, stakingToken, stakeAmount)
         })
         it('should fail', async function () {
           await expect(stake(user, geyser, vault, stakingToken, stakeAmount)).to.be.revertedWith(
             'ExclusiveGeyser: expected exclusive stake',
           )
+          expect(await geyser.computeAvailableStakingBalance(vault.address)).to.eq(0)
+          expect(await vault.checkBalances()).to.eq(true)
         })
-        it('should not fail when there are no locks', async function () {
-          await unstakeAndClaim(user, otherGeyser, vault, stakingToken, 1)
-          await expect(stake(user, geyser, vault, stakingToken, stakeAmount)).not.to.be.revertedWith(
-            'ExclusiveGeyser: expected exclusive stake',
-          )
+
+        it('should NOT fail when there is some unlocked amount', async function () {
+          await unstakeAndClaim(user, otherGeyser, vault, stakingToken, 15)
+          await expect(stake(user, geyser, vault, stakingToken, 14)).not.to.be.reverted
+          expect(await geyser.computeAvailableStakingBalance(vault.address)).to.eq(1)
+          expect(await vault.checkBalances()).to.eq(true)
+          await expect(stake(user, geyser, vault, stakingToken, 1)).not.to.be.reverted
+          expect(await geyser.computeAvailableStakingBalance(vault.address)).to.eq(0)
+          expect(await vault.checkBalances()).to.eq(true)
+        })
+
+        it('should NOT fail when there are no locks', async function () {
+          await unstakeAndClaim(user, otherGeyser, vault, stakingToken, stakeAmount)
+          await expect(stake(user, geyser, vault, stakingToken, stakeAmount)).not.to.be.reverted
+          expect(await geyser.computeAvailableStakingBalance(vault.address)).to.eq(0)
+          expect(await vault.checkBalances()).to.eq(true)
         })
       })
     })