Prevent prize strategy from moving deposit token and added sweep func…

…tion
pooltogether · Oct 14, 2021 · 335dc47 · 335dc47
1 parent 293e224
commit 335dc47
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 3 deletions.
diff --git a/contracts/prize-pool/PrizePool.sol b/contracts/prize-pool/PrizePool.sol
@@ -176,9 +176,9 @@ abstract contract PrizePool is IPrizePool, Ownable, ReentrancyGuard, IERC721Rece
 
         ITicket _ticket = ticket;
 
-        _mint(_to, _amount, _ticket);
-
         _token().safeTransferFrom(_operator, address(this), _amount);
+
+        _mint(_to, _amount, _ticket);
         _supply(_amount);
 
         emit Deposited(_operator, _to, _ticket, _amount);

diff --git a/contracts/prize-pool/YieldSourcePrizePool.sol b/contracts/prize-pool/YieldSourcePrizePool.sol
@@ -26,6 +26,10 @@ contract YieldSourcePrizePool is PrizePool {
     /// @param yieldSource Address of the yield source.
     event Deployed(address indexed yieldSource);
 
+    /// @notice Emitted when stray deposit token balance in this contract is swept
+    /// @param amount The amount that was swept
+    event Swept(uint256 amount);
+
     /// @notice Deploy the Prize Pool and Yield Service with the required contract connections
     /// @param _owner Address of the Yield Source Prize Pool owner
     /// @param _yieldSource Address of the yield source
@@ -47,13 +51,25 @@ contract YieldSourcePrizePool is PrizePool {
         emit Deployed(address(_yieldSource));
     }
 
+    /// @notice Sweeps any stray balance of deposit tokens into the yield source.
+    /// @dev This becomes prize money
+    function sweep() external nonReentrant onlyOwner {
+        uint256 balance = _token().balanceOf(address(this));
+        _supply(balance);
+
+        emit Swept(balance);
+    }
+
     /// @notice Determines whether the passed token can be transferred out as an external award.
     /// @dev Different yield sources will hold the deposits as another kind of token: such a Compound's cToken.  The
     /// prize strategy should not be allowed to move those tokens.
     /// @param _externalToken The address of the token to check
     /// @return True if the token may be awarded, false otherwise
     function _canAwardExternal(address _externalToken) internal view override returns (bool) {
-        return _externalToken != address(yieldSource);
+        return (
+            _externalToken != address(yieldSource) &&
+            _externalToken != yieldSource.depositToken()
+        );
     }
 
     /// @notice Returns the total balance (in asset tokens).  This includes the deposits and interest.

diff --git a/test/prize-pool/YieldSourcePrizePool.test.ts b/test/prize-pool/YieldSourcePrizePool.test.ts
@@ -142,4 +142,18 @@ describe('YieldSourcePrizePool', function () {
             expect(await prizePool.canAwardExternal(yieldSource.address)).to.equal(false);
         });
     });
+
+    describe('sweep()', () => {
+        it('should sweep stray tokens', async () => {
+            await depositToken.mint(prizePool.address, toWei('100'))
+            await yieldSource.mock.supplyTokenTo.withArgs(toWei('100'), prizePool.address).returns()
+            await expect(prizePool.sweep())
+                .to.emit(prizePool, 'Swept')
+                .withArgs(toWei('100'))
+        })
+
+        it('should not allow a non-owner to call it', async () => {
+            await expect(prizePool.connect(wallet2).sweep()).to.be.revertedWith('Ownable/caller-not-owner')
+        })
+    })
 });