fix: Keep referral amounts even after withdrawals

rabbitholegg · Apr 18, 2024 · 33a7089 · 33a7089
1 parent da7e09c
commit 33a7089
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/contracts/Quest.sol b/contracts/Quest.sol
@@ -170,7 +170,7 @@ contract Quest is ReentrancyGuardUpgradeable, PausableUpgradeable, Ownable, IQue
         uint256 protocolFeeForRecipient = (this.protocolFee() / 2) - referralClaimTotal;
         rewardToken.safeTransfer(protocolFeeRecipient, protocolFeeForRecipient);
 
-        uint256 remainingBalanceForOwner = rewardToken.balanceOf(address(this));
+        uint256 remainingBalanceForOwner = rewardToken.balanceOf(address(this)) - referralClaimTotal;
         rewardToken.safeTransfer(owner(), remainingBalanceForOwner);
 
         questFactoryContract.withdrawCallback(questId, protocolFeeRecipient, protocolPayout, address(owner()), ownerPayout);

diff --git a/test/Quest.t.sol b/test/Quest.t.sol
@@ -401,9 +401,18 @@ contract TestQuest is Test, TestUtils, Errors, Events {
             );
         }
 
-
         vm.warp(timestamp);
         vm.prank(referrer);
+
+        // verify that withdrawals can still work
+        quest.withdrawRemainingTokens();
+
+        assertEq(
+            SampleERC20(rewardTokenAddress).balanceOf(address(quest)),
+            quest.referralClaimTotal(),
+            "expected to have referralClaimTotal() amount left inside the contract"
+        );
+
         quest.claimReferralFees(referrer);
 
         assertEq(
@@ -412,9 +421,6 @@ contract TestQuest is Test, TestUtils, Errors, Events {
             "referrer should claim their allocated referral rewards"
         );
 
-        // verify that withdrawals can still work
-        quest.withdrawRemainingTokens();
-
         uint256 participantBalance = SampleERC20(rewardTokenAddress).balanceOf(participant);
         uint256 protocolFeeRecipientBalance = SampleERC20(rewardTokenAddress).balanceOf(quest.protocolFeeRecipient());
         uint256 ownerBalance = SampleERC20(rewardTokenAddress).balanceOf(owner);