skip transfer if leaving rewards in prize pool

GenerationSoftware · Apr 5, 2024 · 8b86ed1 · 8b86ed1
1 parent 19ffd4d
commit 8b86ed1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/src/PrizePool.sol b/src/PrizePool.sol
@@ -581,7 +581,10 @@ contract PrizePool is TieredLiquidityDistributor {
     _totalWithdrawn = SafeCast.toUint128(_totalWithdrawn + _amount);
     _totalRewardsToBeClaimed = SafeCast.toUint104(_totalRewardsToBeClaimed - _amount);
 
-    prizeToken.safeTransfer(_to, _amount);
+    // skip transfer if recipient is the prize pool (tokens stay in this contract)
+    if (_to != address(this)) {
+      prizeToken.safeTransfer(_to, _amount);
+    }
 
     emit WithdrawRewards(msg.sender, _to, _amount, _available);
   }

diff --git a/test/PrizePool.t.sol b/test/PrizePool.t.sol
@@ -1579,6 +1579,16 @@ contract PrizePoolTest is Test {
     prizePool.withdrawRewards(address(1), 5e17);
   }
 
+  function testWithdrawRewards_transferToPrizePool() public {
+    contribute(100e18);
+    awardDraw(winningRandomNumber);
+    mockTwab(address(this), msg.sender, 0);
+    claimPrize(msg.sender, 0, 0, 1e18, address(this));
+    prizePool.withdrawRewards(address(prizePool), 1e18); // leave the tokens in the prize pool
+    assertEq(prizeToken.balanceOf(address(this)), 0);
+    assertEq(prizeToken.balanceOf(address(prizePool)) - prizePool.accountedBalance(), 1e18); // tokens are in prize pool
+  }
+
   function testDrawToAward_zeroDraw() public {
     // current time *is* lastAwardedDrawAwardedAt
     assertEq(prizePool.getDrawIdToAward(), 1);