From 8b86ed13d7c0aab0f6fea866eea89e77b0ffcd06 Mon Sep 17 00:00:00 2001
From: Trevor Richard <trmidpoint@gmail.com>
Date: Fri, 5 Apr 2024 21:37:55 +0000
Subject: [PATCH] skip transfer if leaving rewards in prize pool

---
 src/PrizePool.sol    |  5 ++++-
 test/PrizePool.t.sol | 10 ++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/PrizePool.sol b/src/PrizePool.sol
index fc296e3..cce27ff 100644
--- 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
index 559529f..b36dec4 100644
--- 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);